神经网络---网络模型的保存、加载

方式1：结构+参数

保存

import torch
import torchvision
from torch import nn
from torchvision.models import vgg16, VGG16_Weights

vgg16 = torchvision.models.vgg16(weights=VGG16_Weights.DEFAULT)

# 保存方式1, 模型解构+模型参数
torch.save(vgg16, 'vgg16_1.pth')

加载

from p26_model_svae import *

# 方式1 -》 保存方式1，加载模型
model = torch.load('vgg16.pth')
print(model)

方式1的陷阱
自定义网络结构如下：

import torch
import torchvision
from torch import nn

class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)

    def forward(self, x):
        x = self.conv1(x)
        return x
torch.save(Tudui(), 'tudui_1.pth')

在另一个文件加载该模型，会报错
正确的调用格式需要复制原模型的类定义

class Tudui(nn.Module):
    def __init__(self):
        super(Tudui, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)

    def forward(self, x):
        x = self.conv1(x)
        return x
model = torch.load('tudui_1.pth')
print(model)

或者用import

from p26_model_svae import *

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

方式2 模型参数（官方推荐）

import torch
import torchvision
from torch import nn

vgg16 = torchvision.models.vgg16(pretrained=False)

# 保存方式2 模型参数（官方推荐）
torch.save(vgg16.state_dict(), 'vgg16_2.pth')

模型加载（在另一个文件加载）

# 方式2 ，加载模型
vgg16 = torchvision.models.vgg16(weights=None)
vgg16.load_state_dict(torch.load('vgg16_2.pth'))
# model2 = torch.load('vgg16_2.pth')  #字典形式
print(vgg16)