pytorch学习笔记11-网络模型的读取与保存

模型的保存

先引入vgg16模型（没有经过预训练的）

import torchvision

vgg16_false = torchvision.models.vgg16(pretained=False)

方式一

保存网络模型的结构和其中的参数

torch.save(vgg16_false, "vgg16_method1.pth")

在左侧可以看到已经保存的文件

方式二

把模型的参数保存成字典形式，不保存网络结构，官方推荐的保存方式，因为这种保存方式占用空间小

torch.save(vgg16_false.state_dict(), "vgg16_method2.pth")

在terminal中输入dir查看文件，可以看到方式二比方式一要小一些。

方式1用pth后缀,方式2用pkl后缀,便于区分

模型的读取

对应方式一

import torch

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

输出结果如下，可以debug对比一下方式一保存时的参数，都是一样的。

VGG(
  (features): Sequential(
    (0): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(inplace=True)
    (2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(inplace=True)
    (4): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(inplace=True)
    (7): Conv2d(128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(inplace=True)
    (9): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(inplace=True)
    (12): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(inplace=True)
    (14): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(inplace=True)
    (16): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(inplace=True)
    (19): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(inplace=True)
    (21): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(inplace=True)
    (23): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(inplace=True)
    (26): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(inplace=True)
    (28): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(inplace=True)
    (30): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(output_size=(7, 7))
  (classifier): Sequential(
    (0): Linear(in_features=25088, out_features=4096, bias=True)
    (1): ReLU(inplace=True)
    (2): Dropout(p=0.5, inplace=False)
    (3): Linear(in_features=4096, out_features=4096, bias=True)
    (4): ReLU(inplace=True)
    (5): Dropout(p=0.5, inplace=False)
    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

对应方式二

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

输出是字典

如果想要输出网络结构需要这么写

vgg16_false = torchvision.models.vgg16(pretrained=False)
vgg16_false.load_state_dict(torch.load("vgg16_method2.pth"))

print(vgg16_false)

需要注意的问题

对自己创建的网络模型使用方式一保存时，读取时会出现问题。

比如自己创建一个模型

import torch
from torch import nn


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

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


model = Model()
torch.save(model, "model_method1.pth")

按方式一的读取方法

import torch
model = torch.load("model_method1.pth")
print(model)

输出会报错

这时需要把自己创建的类导入到当前文件中，不需要进行实例化

import torch
from torch import nn


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

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


# model = Model()
model = torch.load("model_method1.pth")
print(model)

或者

import torch
from torch import nn
from P26_model_save import *

# class Model(nn.Module):
#     def __init__(self):
#         super(Model, self).__init__()
#         self.conv1 = nn.Conv2d(3, 64, kernel_size=3)
#
#     def forward(self, x):
#         output = self.conv1(x)
#         return output


# model = Model()
model = torch.load("model_method1.pth")
print(model)

这时程序可以正常输出