Pytorch模型保存

Pytorch 模型保存和加载

  • 保存和加载整个模型
  • 保存加载模型参数

保存和加载整个模型

torch.save(modelname, 'model.pkl')
model = torch.load('model.pkl')

该方式无需自定义网络模型,保存时已把网络结构保存,不方便调整网络结构 。

保存加载模型参数

torch.save(modelname.state_dict(), 'params.pkl')
modelname.load_state_dict(torch.load('params.pkl'))

该方式需要自己定义网络,并且其中的参数名称与结构要与保存的模型中的一致(可以定义部分网络,比如只使用VGG的前几层),相对灵活,便于对网络进行修改。

[1] https://blog.csdn.net/u012436149/article/details/68948816
[2] https://blog.csdn.net/u011276025/article/details/78507950

你可能感兴趣的:(Pytorch)