1.模型保存与读取方式一
(1)模型保存 保存模型结构+模型参数
vgg16=torchvision.models.vgg16(pretrained=False)
torch.save(vgg16,"vgg16_method1.pth")#指定模型保存路径
(2)加载模型
model1=torch.load("vgg16_method1.pth") print(model1)
结果如下:
2.模型保存与读取方式二
(1) 模型保存 保存模型参数(官方推荐)占用空间较少
vgg16=torchvision.models.vgg16(pretrained=False)
torch.save(vgg16.state_dict(),"vgg16_method2.pth") #将vgg16中模型参数保存为字典形式
(2)加载模型
a.保存的是字典形式
model2=torch.load("vgg16_method2.pth")#输出为字典形式 print(model2)#输出为字典形式
b.保存的是网络模型结构
vgg16=torchvision.models.vgg16(pretrained=False)#定义网络模型结构 vgg16.load_state_dict(torch.load("vgg16_method2.pth"))#参数要求是字典形式,把加载的模型放进去,加载的模型保存的是字典形式 print(vgg16)#输出网络模型结构
model_save.py
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")#指定路径
#将vgg16的状态保存为字典形式
#将vgg16中网络模型的参数保存为字典
#陷阱
class Tudui(nn.Module):
def __init__(self) -> None:
super().__init__()
self.conv1=nn.Conv2d(3,64,3)
def forward(self,x):
x=self.conv1(x)
return x
tudui=Tudui()
torch.save(tudui,'tuidui_method1.pth')
model_load.py
import torch
from model_save import *
import torchvision
# 保存方式1,加载模型
model1=torch.load("vgg16_method1.pth")
print(model1)
# 保存方式2,加载模型
model2=torch.load("vgg16_method2.pth")#输出为字典形式
print(model2)#输出为字典形式
vgg16=torchvision.models.vgg16(pretrained=False)#定义网络模型结构
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))#参数要求是字典形式
print(vgg16)#输出网络模型结构
#陷阱1
from torch import nn
model3=torch.load('tuidui_method1.pth')
print(model3)