对于深度学习使用pytorch框架中有时候需要涉及到模型保存和读取，保存方式一般有两种方法，第一种是保存网络模型的模型结构+模型参数(占的内存比方式2大一点点)，第二种（官方推荐）只保存了模型参数。针对保存方式不一样的模型，其读取的方式和所得到的结果都不一样。

以下以VGG16为例：

1.对于方式一：

1.1保存模型：

import torch
import torchvision

# vgg16没有预训练模型的参数
vgg16 = torchvision.models.vgg16(pretrained=False)

# 保存方式1,模型结构+模型参数(占的内存比方式2大一点点)
torch.save(vgg16, "vgg16_method1.pth")    
# 保存路径为当前文件夹下，会生成vgg16_method1.pt权重文件

1.2读取模型：

import torch
model = torch.load("vgg16_method1.pth")
print(model)        # 得到的是模型的网络结构


结果：
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)
  )
)

进程已结束，退出代码为 0

注意：在使用方式一保存模型的时候有时也会出现加载模型的错误

例如：

正常保存了一个模型：

该文件命名为：model_test

import torch
from torch import nn

# 方式一是有陷阱的
class Luguo(nn.Module):
    def __init__(self):
        super(Luguo, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3)

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


luguo = Luguo()
torch.save(luguo, "luguo_method1.pth")

读取模型时候会报错，因为此时读取的模型找不到了模型的结构了，

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

结果：
AttributeError: Can't get attribute 'Luguo' on <module '__main__'

解决办法是，得将网络结构导过来才能运行，否则报错，或者直接在上面直接将model_test整个文件通过from model_test import * 导入。

解决方式一：

import torch
from torch import nn

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

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

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


结果：
Luguo(
  (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))
)

解决方式二：

import torch
# 保存模型文件命名为model_test，*表示将该文件中的所有内容全部导入到此
from model_test import *   

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


结果：
Luguo(
  (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))
)

测试过程中发现，只需将以上的模型加载一次过后似乎该文件已经记得此模型的结构此时再用原来的加载方式也可以加载出来。

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

结果：
Luguo(
  (conv1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))
)

2.对于方式二：

2.1保存模型：

import torch
import torchvision

# vgg16没有预训练模型的参数
vgg16 = torchvision.models.vgg16(pretrained=False)

# 保存方式2,模型参数
torch.save(vgg16, "vgg16_method2.pth")    
# 保存路径为当前文件夹下，会生成vgg16_method2.pt权重文件

2.2读取模型：

该方式的读取模型分为两种，第一种如下。

import torch
# 利用保存方式二加载模型1
model = torch.load("vgg16_method2.pth")
print(model)        # 得到的是模型的字典形式


结果：
OrderedDict([('features.0.weight', tensor([[[[-0.0276, -0.0197, -0.0358],
          [ 0.0672, -0.0233,  0.1106],
          [ 0.0983,  0.0539, -0.1141]],

         [[ 0.0238,  0.0743, -0.0127],
          [ 0.0178,  0.0504, -0.0228],
          [-0.0052,  0.0331,  0.1067]],

         

        ...,

该方式读取模型第二种：

import torch
# 利用保存方式二加载模型2
# 注意此时用到了load_state_dict
vgg16.load_state_dict(torch.load("gg16_method2.pth"))     
print(vgg16)        # 得到的是模型的网络结构（同方式一样）



结果：
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)
  )
)

进程已结束，退出代码为 0

3.总结：

在使用方式一保存模型时，保存的是网络结构+网络参数，在加载过程中可能会由于新的读取环境发生变化，识别不出网络结构，此时需要将之前的网络结构导入到新的读取环境当中。
在使用方式二保存模型时，保存的是字典存储的网络参数，在加载过程中注意此时用到了load_state_dict，才可以得到网络结构。

文章出处登录后可见！

已经登录？立即刷新

迁移学习–网络模型的保存和读取