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"))
所以注意到配套使用,模型部分参数的加载就很容易了