4、模型结构与权重加载：模型定义与预训练权重key不匹配、strict=False的隐患、多GPU加载的坑

模型加载这事儿，看着简单，其实坑特别多。我见过太多同学在这上面浪费一整天，最后发现就是key没对上。今天咱们就把这几个常见问题掰开揉碎了讲清楚。

4.1 模型定义与预训练权重的key不匹配

先说说最常见的情况。你从网上下载了一个预训练模型，比如ResNet50在ImageNet上的权重，然后自己改了一下模型结构——加了个头，或者改了中间层的通道数。一加载，报错：Missing key(s) in state_dict。

为什么会这样？说白了，PyTorch加载权重时，是按照key-value的方式匹配的。key就是每一层的名字，比如conv1.weight、bn1.bias。你的模型定义和预训练权重的key必须完全一致，才能加载成功。

我遇到过最离谱的一次，是一个同学把features.0.conv写成了features.0.Conv，大小写不对，结果加载失败。查了半天才发现。

常见的不匹配场景：

模型结构改了，但权重还是原来的
层命名不一致（比如自己写了个self.my_conv，但权重里叫self.conv）
分类头维度不同（比如你的数据集是10类，但预训练模型是1000类）

解决办法其实很简单：加载时打印出key，对比一下。

# 打印预训练权重的key
checkpoint = torch.load('pretrained.pth')
print(checkpoint.keys())

# 打印模型定义的key
model = MyModel()
print(model.state_dict().keys())

然后手动调整。我一般会写个脚本，把不匹配的key改掉，或者只加载匹配的部分。

4.2 strict=False的隐患

很多人遇到key不匹配，第一反应就是加个strict=False。嗯，这确实能跳过报错，但隐患很大。

strict=False的意思是：能匹配的就加载，匹配不上的就忽略。你想想看，如果模型定义里有个层叫fc.weight，但权重里没有这个key，那这个层就不会被初始化，保持随机值。

我曾经踩过的坑：

有一次我改了模型结构，加了一个新的注意力模块。加载预训练权重时用了strict=False，结果新加的模块没有被初始化，训练了好几个epoch才发现loss一直不降。后来一查，那个模块的权重全是随机的，等于没学到任何东西。

所以我的建议是：

尽量不用strict=False，除非你明确知道哪些层需要忽略
如果非要用，加载后手动检查一下哪些key没匹配上
对于新加的层，单独初始化，不要依赖随机值

# 加载后检查未匹配的key
model.load_state_dict(checkpoint['state_dict'], strict=False)
missing_keys, unexpected_keys = model.load_state_dict(checkpoint['state_dict'], strict=False)
print(f"Missing keys: {missing_keys}")
print(f"Unexpected keys: {unexpected_keys}")

这样你至少知道哪些层没加载上，心里有数。

4.3 多GPU加载的坑

多GPU训练现在很常见，但加载权重时容易出问题。最典型的就是key前面多了个module.前缀。

你想想看，如果用nn.DataParallel包装了模型，保存的权重里所有key都会变成module.conv1.weight、module.bn1.bias这种格式。但如果你在单GPU上加载，模型定义里是conv1.weight，没有module.前缀，自然就匹配不上了。

我的处理方式：

保存权重时，我习惯只保存model.module.state_dict()，而不是model.state_dict()。这样不管在单GPU还是多GPU上加载，key都是统一的。

# 保存时去掉module前缀
if isinstance(model, nn.DataParallel):
    torch.save(model.module.state_dict(), 'model.pth')
else:
    torch.save(model.state_dict(), 'model.pth')

如果已经保存了带module.前缀的权重，加载时可以用下面这个技巧：

# 加载时去掉module前缀
checkpoint = torch.load('model.pth')
new_state_dict = {}
for k, v in checkpoint.items():
    if k.startswith('module.'):
        new_state_dict[k[7:]] = v  # 去掉module.前缀
    else:
        new_state_dict[k] = v
model.load_state_dict(new_state_dict)

还有一个坑是batch size不一致。多GPU训练时，每个GPU上的batch size是总batch size除以GPU数量。如果你在单GPU上加载多GPU训练的权重，BN层的running mean和running var可能不匹配。不过这个影响不大，一般训练几个iteration就调整过来了。

4.4 知识体系总览

下面这张图总结了模型加载的核心逻辑和常见问题，你可以对照着排查：

4.5 避坑指南总结

最后，我把这些经验总结成几条实用建议：

加载前先打印key——花30秒看一眼，能省下半天排查时间
尽量不用strict=False——除非你明确知道自己在做什么
多GPU训练时统一保存格式——我习惯只保存model.module.state_dict()
加载后验证一下——跑几个batch看看loss是不是正常下降

一个小技巧：

我写了个工具函数，每次加载权重时自动打印匹配和不匹配的key，还会统计匹配率。这样一眼就能看出问题。你也可以试试。

嗯，模型加载这块其实不难，但细节决定成败。下次遇到加载报错，别慌，按上面的步骤一步步排查，很快就能解决。