模型的保存和读取

模型的保存

import torch
import torchvision
from torch import nn

vgg16 = torchvision.models.vgg16(pretrained=False)
# 保存方式1
torch.save(vgg16, "vgg16_method1.pth")

# 保存方式2,模型参数（官方推荐）
# 将vgg16的状态保存成字典形式
torch.save(vgg16.state_dict(), "vgg16_method2.pth")

模型的加载

import torch
import torchvision
from torch import nn

vgg16 = torchvision.models.vgg16(pretrained=False)
# 方式1-》保存方式1，加载模型
model = torch.load("vgg16_method1.pth")
print(model)

# 方式2-》加载模型
# ①字典形式
model = torch.load("vgg16_method2.pth")
print(model)
# ②将字典形式恢复成网络模型
vgg16 = torchvision.models.vgg16(pretrained=False)
# 通过Python字典形式加载网络模型
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))
print(vgg16)

陷阱

模型保存

# 保存方式1 的陷阱
class Peipei(nn.Module):

    def __init__(self) -> None:
        super(Peipei, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

    def forward(self, x):
        x = self.conv1(x)
        return x


peipei = Peipei()
model = torch.save(peipei, "peipei_method1.pth")
print(model)

模型加载： 加载的时候还需要复制网络结构，但不需要创建了

# 陷阱1  加载的时候还需要复制网络结构，但不需要创建了
class Peipei(nn.Module):

    def __init__(self) -> None:
        super(Peipei, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

    def forward(self, x):
        x = self.conv1(x)
        return x


model = torch.load('peipei_method1.pth')
print(model)