1. 网络保存
import torch
import torchvision
from torch import nnvgg16 = torchvision.models.vgg16(pretrained=False)#保存方式1 不仅保存了网络模型结构也保存了参数
torch.save(vgg16,'vgg16_method1.pth')#保存方式2 获取模型状态(参数)并且保存 (官方推荐 比较小)
torch.save(vgg16.state_dict(),'vgg16_method2.pth')#保存自定义的模型
class Tudui(nn.Module):def __init__(self):super(Tudui, self).__init__()self.conv1 = nn.Conv2d(3,64,5)def forward(self,x):x = self.conv1(x)return xtudui = Tudui()torch.save(tudui,'tudui_method.pt')
2. 网络加载
import torch# 加载方式1 --保存方式1
model1 = torch.load('vgg16_method1.pth')
print(model1)
#
# 加载方式2 --保存方式2
model2 = torch.load('vgg16_method2.pth')
print(model2) #打印出的是参数#如何获取模型结构?
import torchvisionvgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load('vgg16_method2.pth'))
print(vgg16)#加载自定义保存的模型
from model_save import tudui
tudui = torch.load('tudui_method.pt')
print(tudui) #失败的话需要引入包