课程学习笔记，课程链接
在搭建自己的神经网络模型之后，需要将模型进行保存，同时也需要读取或加载现有的神经网络模型。
加载代码
# 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) # 报错
