PyTorch深度学习笔记（十八）网络模型的保存与读取

课程学习笔记，课程链接

在搭建自己的神经网络模型之后，需要将模型进行保存，同时也需要读取或加载现有的神经网络模型。

加载代码

# File : model_load.py
import torch
​
# 方式1 -> 保存方式1，加载模型
import torchvision
from torch import nn
​
model = torch.load("vgg16_method1.pth")
print(model)
​
# 方式2 ，加载模型
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
# model = torch.load("vgg16_method2.pth")
print(vgg16)
​
# 方式 1 陷阱
class Jiaolong(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
​
jiaolong = Jiaolong()
torch.save(jiaolong, 'jiaolong_method1.pth')

保存代码

# File : model_save.py
import torch
import torchvision
​
vgg16 = torchvision.models.vgg16(pretrained=False)
# 保存方式 1，模型结构+模型参数
torch.save(vgg16, "vgg16_method1.pth")
​
# 保存方式2，模型参数（官方推荐）
torch.save(vgg16.state_dict(), 'vgg16_method2.pth')  # 将模型的状态（参数）保存成字典形式
​
# 方式 1 陷阱，需要将模型写过来，但无需创建实例
# 可以引入，from model_save import *
model = torch.load('jiaolong_method1.pth')
print(model)  # 报错