Pytorch自定义加载预训练权重

Posted Rogn

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了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

以上是关于Pytorch自定义加载预训练权重的主要内容,如果未能解决你的问题,请参考以下文章

在 pytorch 中修改预训练模型时,旧的权重会重新初始化吗?

如何加载部分预训练的 pytorch 模型?

如何使用 Pytorch 中的预训练权重修改具有 4 个通道作为输入的 resnet 50?

在预训练的 pytorch 网络中加载单个图像

如何使用 PyTorch 在预训练模型上添加新层? (给出了 Keras 示例。)

利用Pytorch的C++前端(libtorch)读取预训练权重并进行预测