1. 网络保存
import torch
import torchvision
from torch import nn
vgg16 = 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 x
tudui = 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 torchvision
vgg16 = 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) #失败的话需要引入包