PyTorch Lightning加载预训练模型权重未更新问题及临时修复方案

问题现象 #

• 基于预训练模型改造下游任务模型，仅替换最终fc层适配任务需求后，需要加载预训练权重
• 直接调用self.model.load_from_checkpoint(self.pretrained_model_path)加载时，无任何报错、警告输出，但对比加载前后的层权重数值完全没有变化，权重未实际更新

临时验证可行方案 #

经测试，以下写法可成功加载权重：

self.model.backbone = self.model.load_from_checkpoint(self.pretrained_model_path).backbone

该写法可将骨干网络参数更新为预训练权重值，属于临时修复方案，存在更规范的实现方式。

问题根因 #

load_from_checkpoint是类方法，不是实例方法。在已初始化的模型实例上调用该方法时，它不会原地修改当前实例的权重，而是会重新初始化一个全新的对应类的模型实例，把checkpoint权重加载到这个新实例后返回。如果不手动接收这个返回值，原实例的权重自然不会发生任何变化，流程也不会抛出异常。

更优实现方式 #

方案1：类方法调用加载后复用骨干权重 #

直接通过模型类调用load_from_checkpoint拿到完整的加载好权重的预训练实例，再提取需要的骨干部分赋值给下游模型，逻辑更清晰：

# 替换成你实际定义的Lightning模型类名，加载完整预训练模型实例
pretrained_instance = YourDefinedModelClass.load_from_checkpoint(self.pretrained_model_path)
# 复用加载好权重的骨干网络
self.model.backbone = pretrained_instance.backbone
# 自定义的下游fc层保持随机初始化即可

方案2：用原生state_dict接口灵活加载 #

如果需要更细粒度的加载控制，可以直接用PyTorch原生的状态字典加载接口，设置strict=False自动跳过维度不匹配的fc层参数，不需要额外实例化新模型：

import torch
# 读取checkpoint文件
ckpt = torch.load(self.pretrained_model_path, map_location="cpu")
# 加载权重，自动忽略shape不匹配的fc层参数
self.model.load_state_dict(ckpt["state_dict"], strict=False)

内容的提问来源于stack exchange，提问作者Mohbat Tharani