使用PyTorch加载模型部分参数方法

前言

在深度学习领域，经常需要使用其他人已训练好的模型进行改进或微调，这个时候我们通常会希望加载预训练模型文件的参数，如果网络结构不变，只需要使用load_state_dict方法即可。而当我们改动网络结构后，由于load_state_dict方法要求读入的state_dict的key和net.state_dict()的key对应相等，如果有缺少就会报错。这个时候我们通常希望加载未改动部分结构的参数用来初始化网络。

方法

方法一：strict=False

直接使用load_state_dict提供的参数strict=False，网络结构名字一致的会被导入，不一致的会被舍弃：

net.load_state_dict(checkpoint['net'], strict=False)

方法二：update

使用net.state_dict().update()方法来更新筛选出来的key-value。

model_dict=net.state_dict()
pretrained_dict = {k: v for k, v in checkpoint['net'].items() if k in model_dict}
model_dict.update(pretrained_dict)
net.load_state_dict(model_dict)

注意事项

Optimizer的加载

在加载优化器时，需要注意如果改动了网络结构后，优化器中的参数长度可能会对应不上，这时候就会报错，例如"param 'initial_lr' is not specified in param_groups[*] when resuming an optimizer"。

而优化器的load_state_dict方法没有参数strict，此时可以选择只加载优化器的基本信息，例如初始学习率initial_lr。

optimizer_net= torch.optim.Adam([{'params': net.parameters(), 'initial_lr': 0.002}], lr=2.0e-3, betas=(0.5, 0.999))

参考资料

[1] PyTorch Docs - load_state_dict(state_dict, strict=True)
[2] PyTorch Docs - update(modules)
[3] csdn - pytorch 模型部分参数的加载
[4] 知乎 - PyTorch模型加载/参数初始化/Finetune
[5] csdn - Pytorch中，只导入部分模型参数的做法
[6] csdn - Pytorch：lr_schedule恢复训练的注意事项
[7] PyTorch Docs - torch.optim.Optimizer.load_state_dict