小土堆pytorch学习笔记（十七、网络模型的保存与读取）

一、模型的保存

import torch
import torchvision

vgg16 = torchvision.models.vgg16(pretrained=16)
# 保存方式1, 模型结构+模型参数
torch.save(vgg16, "vgg16_method1.pth")

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

再搭建一个神经网络进行保存：

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

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


cow = Cow()
torch.save(cow, "cow_method1.pth")

二、模型的加载

import torch
import torchvision

# 保存方式1，加载模型
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)

加载搭建好的神经网络模型

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

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


model = torch.load("cow_method1.pth")
print(model)

运行结果：

Cow(
  (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))
)

用方法1保存的模型进行加载时必须将网络模型也复制过去，不过不用创建一个子类。
或者直接导入保存的模型：

from model_save import *

这样就不用将搭建的网络模型复制过去，可以直接加载。