Pytorch如何部分加载模型参数

Pytorch如何加载部分模型参数

在处理深度学习相关的任务时，有些情况下是在别人已有的模型上做了少量的更改，此时如果从头开始训练此模型，感觉很浪费时间，因为可以充分利用别人已经训练出来的模型的参数，供自己更改的模型使用。因此本文便简单分享一种方法，来加载模型的部分参数。

方法

经过查阅相关博客，发现加载模型部分参数的方法如下：

model.load_state_dict(related_params, strict=True)

strict是False时，当权重中的key和网络中匹配就加载，不匹配就跳过。如果strict是True，那必须完全匹配，不然就报错。所以如果只想加载部分的模型参数，只需要另strict参数为False即可，这样就可以将相匹配的部分加载进来

踩坑

笔者直接进行尝试后，发现会报如下错误

'model' object has no attribute 'copy'

笔者所写的代码如下所示：

newModel.load_state_dict(torch.load("model.pth"), strict=False)

这里的"model.pth"，笔者在进行保存时，是使用torch.save(model, “model.pth”)直接进行保存的，因此无法进行加载。如果在保存模型时使用如下命令：

torch.save(model.state_dict(), "model.pth")

模型就可以正确地被加载进去了

原因

这其中的原因其实很简单，就是模型的保存和加载要配套使用，笔者也是在找到此篇文章后才知道的：[错误汇总]‘model’ object has no attribute 'copy’加载模型报错
简而言之，就是保存模型和加载模型要配套使用

1 直接进行模型的保存和加载：

torch.save(model, "model.pth")
newModel = torch.load("model.pth")

2 分别进行加载

torch.save(model.state_dict(), "model.pth")
newModel.load_state_dict(torch.load("model.pth"))

所以注意到配套使用，模型部分参数的加载就很容易了