当前位置:   article > 正文

python-pytorch 关于torch.load()和torch.load_state_dict()_torch.load stat

torch.load stat


最近在使用pytorch训练和加载模型时遇到了一些很玄学的问题,研究了一下,总结如下:

1、关于模型保存和加载

1、如果保存时使用了torch.save( model, “model.pth”),那么模型加载时只能使用torch.load(),而不能使用torch.load_state_dict(torch.load(model_path)),如果强行使用torch.load_state_dict(torch.load(model_path)),会报Net对象没有copy属性。

import torch

torch.save(model, 'model.pth')
model = torch.load('model.pth')
  • 1
  • 2
  • 3
  • 4

2、如果保存时使用了torch.save( model.state_dict(),“model.pth”),那么模型加载时只能使用model.load_state_dict(torch.Load(“model.pth”))。

import torch

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

state_dict = torch.load('model.pth')
model.load_state_dict(state_dict)
  • 1
  • 2
  • 3
  • 4
  • 5
  • 6

3、两种方式在模型加载时,都需要引入或定义模型结构,且模型结构的名称需要与训练时一致,不能使用别名。
训练时:class Net,测试时:定义或加载仍然是class Net,而不能from ··· import Net as resNet

# 两个Net的结构不同
from model1 import Net
from model2 import Net
  • 1
  • 2
  • 3

2、关于加载模型结构

关于在一个脚本中同时加载两个模型,在同时导入两个不同结构的Net的情况下,即使Net是同名,但加载的模型会自动找到自己训练时对应的那个结构。

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/酷酷是懒虫/article/detail/739378
推荐阅读
相关标签
  

闽ICP备14008679号