Pytorch自定义加载预训练权重


pytorch保存模型权重非常方便

 

保存模型可以分为两种

一种是保存整个网络(网络结构+权重参数)

torch.save(model, 'net.pth.tar')

一种是只保存模型的权重参数(速度快,内存占用少)

torch.save(model.state_dict(), 'net.pth.tar')

标准的加载还可以做一些设置

 

如果加载模型只是为了进行推理测试,则将每一层的 requires_grad 置为 False,即固定这些权重参数;还需要调用 model.eval() 将模型置为测试模式,主要是将 dropout 和 batch normalization 层进行固定,否则模型的预测结果每次都会不同。

如果希望继续训练,则调用 model.train(),以确保网络模型处于训练模式。

 

然后,想使用预训练权重有非常严格的要求,要求每一层一模一样,命名都要一样(不然dict的key就不一样了

如果两个模型实际是一样,既然是字典,可不可以手动赋值呢?

比如,我把cnn命名成conv了,导致load失败,因此我们来手动赋值

 

虽然load成功了,但是实际效果有点问题,直接用于生成,得到的图像颜色不对。

按道理两种方法的权重应该是一模一样啊,

 

前两个是手动赋值,颜色都有偏差,第三个直接load的正常。。。

难道model.state_dict() 没有包含所有的权重信息?

 

其实还有一个简便方法:使用strict=False 参数

model.load_state_dict(checkpoint["state_dict"], strict=False)
optimizer.load_state_dict(checkpoint["optimizer"])

奇怪的是,model的load_state_dict有strict参数,optimizer没有

 

 

参考链接:https://zhuanlan.zhihu.com/p/73893187


免责声明!

本站转载的文章为个人学习借鉴使用,本站对版权不负任何法律责任。如果侵犯了您的隐私权益,请联系本站邮箱yoyou2525@163.com删除。



 
粤ICP备18138465号  © 2018-2025 CODEPRJ.COM