PyTorch快速入门教程【小土堆】-神经网络-网络模型的保存与读取

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)