pytorch学习（十三）：保存模型，加载模型

 保存方式一：

保存网络结构和参数

vgg16=torchvision.models.vgg16(pretrained=False)
#保存方式1 保存网络的结构和参数
torch.save(vgg16,"./models/vgg16_method1.pth")

保存方式二;

以字典形式保存网络参数（官方推荐），只保存参数，不保存结构；

vgg16=torchvision.models.vgg16(pretrained=False)
#保存方式2 以字典形式保存模型参数（官方推荐）只保存参数
torch.save(vgg16.state_dict(),"./models/vgg16_method2.pth")

 加载方式一：

#加载方式一
model=torch.load("./models/vgg16_method1.pth")
print(model)

加载方式二：

#加载方式二
vgg16=torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("./models/vgg16_method2.pth"))

陷阱 1：

#陷阱
class Wang(nn.Module):def __init__(self):super(Wang, self).__init__()self.model1=nn.Sequential(nn.Conv2d(3, 32, 5, padding=2),nn.MaxPool2d(2),nn.Conv2d(32, 32, 5, padding=2),nn.MaxPool2d(2),nn.Conv2d(32, 64, 5, padding=2),nn.MaxPool2d(2),nn.Flatten(),nn.Linear(1024, 64),nn.Linear(64, 10),)def forward(self,x):x=self.model1(x)return x
wang=Wang()
torch.save(wang,"./models/wang_method1.pth")

上述代码保存模型后，加载模型时候会报错

错误加载模型代码：

#陷阱1
model=torch.load("./models/wang_method1.pth")
print(model)

报错信息：

Traceback (most recent call last):
  File "D:/pytorchlen/19_model_load.py", line 21, in <module>
    model=torch.load("./models/wang_method1.pth")
  File "D:\miniconda\envs\gpu_torch\lib\site-packages\torch\serialization.py", line 592, in load
    return _load(opened_zipfile, map_location, pickle_module, **pickle_load_args)
  File "D:\miniconda\envs\gpu_torch\lib\site-packages\torch\serialization.py", line 851, in _load
    result = unpickler.load()
AttributeError: Can't get attribute 'Wang' on <module '__main__' from 'D:/pytorchlen/19_model_load.py'>

解决办法，将自己写的网络结构赋值过来，加载模型的方式仅仅加载了参数以及省略了实例化类这一步。

还可以不复制 ，直接 from model_save import *