保存方式一:
保存网络结构和参数
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 *