[VGG模型] 神经网络模型的保存与读取

本文以VGG16模型为例，来学习神经网络模型的保存与读取，其中网络模型的读取有两种方法，对应的网络模型的读取也有两种方法，下面来一一学习。

---

文章目录

一、网络模型的保存方法

1 保存方法一

2 保存方法二

二、网络模型的读取方法

1 读取方法一

2 读取方法二

三、两种方法的对比及注意事项

1 解决方法一

2 解决方法二

---

一、网络模型的保存方法

网络模型的保存有两种方法

1 保存方法一

以vgg16模型为例，代码如下：

import torch
import torchvision

vgg16 = torchvision.models.vgg16(pretrained=False)

#保存方法一
torch.save(vgg16, "vgg16_model.pth") 

其中，

pretrained=False表示使用初始化的参数，没有经过数据集训练，

vgg16,是要保存的网络模型

 "vgg16_model.pth"是给这个保存的网络模型一个名字

.pth是常用的后缀格式，

点击运行，就会在左边创建一个网络模型vgg16_model.pth文件

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

2 保存方法二

方法二，保存的是模型参数（官方推荐）因为这种保存方式的内存占比较小

import torch
import torchvision

vgg16 = torchvision.models.vgg16(pretrained=False)  
#保存方法二

torch.save(vgg16.state_dict()，"vgg16_model2.pth")   
相当于把vgg16当中网络模型中的参数保存为字典形式，

其中：
vgg16.state_dict()是将vgg16中的参数保存成字典格式，

"vgg16_model2.pth"是给这个保存的网络模型一个名字
.pth是常用的后缀格式，

点击运行，就会在左边创建一个网络模型vgg16_model2.pth文件
同上图一样。二、网络模型的读取方法

二、网络模型的读取方法

1 读取方法一

对应上面的保存方法一，代码如下：

import torch

model = torch.load("vgg16_model.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)
  )
)

直接就将这个vgg16的模型结构输出了

2 读取方法二

对应上面的保存方法二， 代码如下：

方法二对应的网络模型的读取
import torch

model2 = torch.load("vgg16_model2.pth")

print(model2)

输出结构如下：

输出：
OrderedDict([('features.0.weight', tensor([[[[ 0.1032, -0.0300, -0.0190],
          [-0.0724, -0.0699, -0.0791],
          [ 0.0720, -0.0908,  0.0272]],

         [[ 0.0286,  0.0567, -0.0162],
          [ 0.0267, -0.0782,  0.0447],
          [ 0.0556,  0.0668,  0.0825]],

         [[ 0.0641, -0.0478,  0.0052],
          [-0.0937, -0.0636,  0.0672],
          [ 0.0316,  0.0376,  0.1007]]],
   ... ... ...

它保存的是一种字典形式的

如果我们要把这个恢复成网络模型

需要新建网络模型结构

代码如下：

import torch
import torchvision

vgg16 = torchvision.models.vgg16(pretrained=False)  
vgg16.load_state_dict(torch.load("vgg16_model2.pth"))
print(vgg16)

注：

vgg16 = torchvision.models.vgg16(pretrained=False)  是新建网络模型结构
vgg16.load_state_dict(torch.load("vgg16_model2.pth"))，表示vgg16要load_state加载一个状态（参数），通过dict字典

输出的结果如下：

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))

...........

    (6): Linear(in_features=4096, out_features=1000, bias=True)
  )
)

可以看见，新建网络模型结构后，输出结构与方法一一样。

三、两种方法的对比及注意事项

第二种方法进行网络模型的保存和读取是官方推荐的,

虽然看起来方法一比较简单

但是方法一有时候这样用会报错的,下面我们来看一下

代码如下：

首先,搭建一个Test神经网络,并保存网络模型test.pth
import torch
from torch import nn

#搭建神经网络
class Test(nn.Module):
    def __init__(self):              #初始化
        super(Test, self).__init__()       #继承父类
        self.conv = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3)
    def forward(self, x):
        x = self.conv(x)
        return x

test = Test()

torch.save(test,"test.pth")

我们来看看直接读取这个test.pth会发生什么

新创建一个py文件，输入：

import torch

test1 = torch.load("test.pth")
print()

结果直接报错: AttributeError: Can't get attribute 'Test' on <module '__main__' from 'C:/Users/Administrator/Desktop/learn_torch1/TEST.py'>

大概意思就是:你不能得到Test这个属性,因为你没有这个类

解决这种报错一般有两种方法

1 解决方法一

方法一:将Test这个网络模型给复制过来,但是不需要创建神经网络的,也就是不需要test = Test() ,这一步的。

只需要将模型的定义写过来,目的应该就是确保,你加载的这个网络模型是你想要的网络模型

正确的应该这样写

import torch
from torch import nn

class Test(nn.Module):
    def __init__(self):              
        super(Test, self).__init__()      
        self.conv = nn.Conv2d(in_channels=3, out_channels=64, kernel_size=3)
    def forward(self, x):
        x = self.conv(x)
        return x

test1 = torch.load("test.pth")
print(test1)

输出结果:
Test(
  (conv): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))
)

2 解决方法二

将写有这模型的文件直接引用过来就行了

import torch
from model_save import *                   
test1 = torch.load("test.pth")
print(test1)

输出结果:
Test(
  (conv): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))
)
 

其中,
from model_save import * 
的目的就是将model_save定义的网络模型引用过来

这样也是可以的,具体怎么操作要看实际情况来.