加载DynamicUnet权重时state_dict缺失与尺寸不匹配错误求助

问题场景 #

从其他训练环境导出序列化模型export.pkl与权重文件，在本地指定环境加载后用于新数据集推理时，出现权重加载错误。

环境版本 #

• torch==1.7.1
• fastai==2.7.7
• fastcore==1.5.6
• torchvision==0.8.2

执行代码 #

from fastai.vision.all import *
learn = load_learner('export.pkl', cpu=True)
learn.load('model_3C_34_CELW_V_1.1')

报错信息 #

加载DynamicUnet的state_dict时出现错误，存在缺失键（如layers.0.4.0.conv3.weight等），且layers.12.0.weight参数尺寸不匹配（checkpoint中为torch.Size([3,99,1,1])，当前模型为torch.Size([3,291,1,1])）

解决方案 #

1. 处理参数尺寸不匹配 #

尺寸不匹配的核心原因是当前模型的类别数/特征通道数与训练权重时的模型不一致：

• 原权重对应模型的最后一层输出通道为99，而当前加载的export.pkl模型对应291类，导致layers.12.0.weight的尺寸差异。
• 对应解决方法：
  • 如果要在291类的新数据集上复用原权重，需修改模型最后一层的输出通道，并过滤加载匹配的参数：

    learn = load_learner('export.pkl', cpu=True)
    # 替换最后一层ConvLayer的输入通道为291（匹配当前数据集类别数）
    learn.model[1][-1] = ConvLayer(291, 3, ks=1)  # 需根据实际模型结构调整层的索引
    # 加载权重文件并过滤匹配的参数
    state_dict = torch.load('model_3C_34_CELW_V_1.1.pth', map_location='cpu')
    filtered_dict = {k:v for k,v in state_dict.items() if k in learn.model.state_dict() and v.shape == learn.model.state_dict()[k].shape}
    learn.model.load_state_dict(filtered_dict, strict=False)
    

  • 如果原模型与当前数据集类别数应一致，检查是否加载了错误的export.pkl或权重文件，确保两者来自同一训练任务。

2. 处理缺失键问题 #

缺失键通常由训练环境与当前环境的FastAI/Torch版本差异导致，版本不同会使DynamicUnet的层命名结构变化；也可能是原模型使用了自定义组件，当前环境未定义：

• 优先保证训练环境与当前环境的FastAI、Torch版本完全一致，版本匹配是避免结构差异的核心。
• 若版本一致仍有缺失键，加载权重时添加strict=False参数跳过不匹配的键：

  learn.load('model_3C_34_CELW_V_1.1', strict=False)
  

  注意：此方法会跳过未找到的层权重，若缺失的是关键层，可能影响模型性能，需确认缺失层的作用。

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