PyTorch中增减修改模型层结构后能否加载预训练权重开展训练？

基础结论 #

对原有模型做层结构增减、维度调整后，可以加载对应预训练权重开展迁移训练，但你给出的示例代码无法直接运行，会触发报错。
PyTorch的load_state_dict()方法默认会严格校验权重字典的key名称、张量形状和当前模型完全匹配，一旦你修改了层结构，要么新增的层在预训练权重中无对应key，要么修改后的层张量形状和预训练权重不匹配，直接调用默认参数会抛出RuntimeError。同时你示例中加载原有优化器状态的操作完全不可行，优化器状态和原模型的参数是一一绑定的，模型结构修改后旧的优化器状态没有任何复用价值，加载反而会导致训练异常。

你给出的ResNet34修改场景的可行解决方案 #

方案1：非严格加载匹配权重（适合仅新增层、未修改原有层维度的场景） #

你提到只是在conv4_x模块新增了若干层，没有修改原有层的结构，直接在加载权重时添加strict=False参数即可自动跳过不匹配的key，原有未修改层的预训练权重会正常加载，新增层自动走默认初始化：

import torch
import torch.optim as optim

model = Resnet34() # 已修改层结构的自定义ResNet34
pretrained_ckpt = torch.load('Resnet.pt')
# 非严格加载，返回缺失的key和多余的key供校验
missing_keys, unexpected_keys = model.load_state_dict(pretrained_ckpt['state_dict'], strict=False)
# 建议打印确认：missing_keys应为你新增的层的key，unexpected_keys应为被你删掉的层的key，无异常再继续
print("未加载到权重的层（新增层）：", missing_keys)
print("预训练中多余的权重（被删除层）：", unexpected_keys)

# 初始化新的优化器，不要加载旧的优化器状态
optimizer = optim.Adam(model.parameters(), lr=0.00005)

方案2：手动过滤权重（适合同时修改了原有层维度的场景） #

如果你修改了原有层的输出/输入维度，key还存在但张量形状不匹配，用strict=False仍会报错，此时可以手动筛选出完全匹配的权重再加载：

model = Resnet34()
pretrained_state = torch.load('Resnet.pt')['state_dict']
model_state = model.state_dict()
# 仅保留key存在、且形状完全匹配的预训练权重
filtered_state = {k: v for k, v in pretrained_state.items() if k in model_state and v.shape == model_state[k].shape}
model_state.update(filtered_state)
model.load_state_dict(model_state)

方案3：分层冻结微调（小数据集训练推荐） #

如果你的自定义数据集规模较小，加载完匹配的预训练权重后，可以先冻结所有已经加载了预训练权重的层，仅训练新增的层，等训练到初步收敛后再解冻全网络微调，能避免预训练的通用特征被破坏，提升训练效果。

注意事项 #

1. 无论哪种场景，都不要加载原有预训练模型的优化器状态，模型结构修改后旧的优化器缓存参数完全不适用。
2. 如果你只是修改了模型最后的分类头（比如把ResNet默认的1000类全连接层改成自定义数据集的类别数），上述方案就是工业界最常用的迁移训练实现，效果非常稳定。
3. 如果在模型中间层新增了大量层，建议给预训练层设置更低的学习率，新增层使用更高的学习率，训练效率会更高。

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