PyTorch学习笔记（十四）——网络模型的保存与读取

两种方式保存和加载模型

方式一

保存模型

不仅保存了网络模型的结构，也保存了网络模型的参数

import torch
import torchvision

vgg16 = torchvision.models.vgg16(weights=False)
torch.save(vgg16,"vgg16_method1.pth")

加载模型

打印出的是网络模型的结构 

import torch

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

方式二

保存模型

网络模型的参数保存为字典，不保存网络模型的结构（官方推荐的保存方式，用的空间小）

import torch
import torchvision

vgg16 = torchvision.models.vgg16(weights=False)
torch.save(vgg16.state_dict(),"vgg16_method2.pth")

 加载模型

打印出的是参数的字典形式

import torch

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

 如何恢复网络模型结构？

import torchvision.models
 
vgg16 = torchvision.models.vgg16(pretrained=False)  # 预训练设置为False
vgg16.load_state_dict(torch.load("vgg16_method2.pth"))  # vgg16通过字典形式，加载状态即参数
print(vgg16)

 陷阱

问题描述：

首先在 model_save.py 中写以下代码并运行

import torch
from torch import nn


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

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

mynn = Mynn()
torch.save(mynn,"mynn_method.pth")

再在 model_load.py 中写以下代码

import torch

# 陷阱
model = torch.load("mynn_method.pth")
print(model)

运行后报错

 解决办法1：

需要将 model_save.py 中的网络结构复制到 model_load.py 中（不需要写mynn = Mynn()），即

import torch
from torch import nn

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

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

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

 解决办法2：

实际写项目过程中，直接定义在一个单独的文件中（如model_save.py），再在 model_load.py 中

from model_save import *