Pytorch笔记19 model的保存与读取

19.1 model的两种保存方式

19.1.1 代码实现

import torch
import torchvision
from torch import nn

vgg16 = torchvision.models.vgg16()
# 保存方式1,保存结构及参数
torch.save(vgg16, "vgg16_method1.pth")

# 保存方式2,将参数保存成字典(官方推荐)
torch.save(vgg16.state_dict(), "vgg16_method2.pth")


# 陷阱
class Avlon(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

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


avlon = Avlon()
torch.save(avlon, "avlon_method1.pth")

19.1.2 两种方式对比

  在当前文件目录下进行操作


  win10系统下执行"ls"命令


  可见方法2 保存的文件更小

19.2 model的两种读取方式

  分别对应两种保存方式

19.2.1 代码实现

import torch
from model_save import *    # 导入 model_save 模块中定义的所有函数和变量
import torchvision
from torch import nn

# 方式1-> 保存方式1，加载模型
model = torch.load("vgg16_method1.pth")
print(model)

# 方式2，加载模型
vgg16 = torchvision.models.vgg16()
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))    # 将字典形式的参数加载到模型中
print(vgg16)


# 陷阱1
# class Avlon(nn.Module):
#     def __init__(self):
#         super().__init__()
#         self.conv1 = nn.Conv2d(3, 64, kernel_size=3)
#
#     def forward(self, x):
#         x = self.conv1(x)
#         return x


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