PyTorch学习笔记（14）--神经网络模型的保存与读取_神经网络保存模型

PyTorch学习笔记（14）–神经网络模型的保存与读取

    本博文是PyTorch的学习笔记，第14次内容记录，主要是讲解如何进行神经网络模型的保存和读取。

1.网络模型保存和加载–方法1

1.1网络模型保存方法1

    在搭建自己的神经网络模型之后，需要将模型进行保存，同时也需要读取或加载现有的神经网络模型，这里就介绍网络模型保存的第1种方法。以VGG16模型为例，代码如下：

import torch
import torchvision

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

# 保存方式1:不仅保存了网络模型，也保存了网络模型中的相关参数
torch.save(vgg16, "vgg16_model1.pth")
1
2
3
4
5
6
7

    运行上述代码会发现在其同路径下保存了神经网络模型文件：vgg16_model1.pth。

1.2网络模型加载方法1

    在保存神经网络模型之后，可以在其他文件中调用该模型，代码如下：

import torch
# 方式1-> 保存方式1，加载模型
import torchvision

model = torch.load("vgg16_model1.pth")
print(model)
1
2
3
4
5
6

    输出的网络模型结构信息如下：

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)
  )
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

2.网络模型保存和加载–方法2

2.1网络模型保存方法2

    在神经网络模型保存方法1中，不仅保存了模型的结构，同时也保存了模型的相关参数信息。此外，模型保存的方法2代码如下：

import torch
import torchvision

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

# 保存方式1:不仅保存了网络模型，也保存了网络模型中的相关参数
torch.save(vgg16, "vgg16_model1.pth")

# 保存方式2：只保存了模型的参数，占用空间更小，官方推荐方式
torch.save(vgg16.state_dict(), "vgg16_model2.pth")
1
2
3
4
5
6
7
8
9
10

    运行上述代码会发现在其同路径下保存了神经网络模型文件：vgg16_model2.pth。

2.2网络模型加载方法2

    采用方法2保存网络模型之后，需要采用相应的方法加载网络模型，具体代码如下：

import torch

# 方式1-> 保存方式1，加载模型
import torchvision

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

# 方式2-> 保存方式2，加载模型
model = torch.load("vgg16_model2.pth")  # 加载出来的是字典类型的数据
print(model)
1
2
3
4
5
6
7
8
9
10
11

    由于保存方法2保存的是模型结构，所以加载出来的模型是以“字典”的方式展示出来的，效果如下：

OrderedDict([('features.0.weight', tensor([[[[ 0.0814, -0.0414, -0.0450],
          [ 0.1337,  0.0163,  0.0327],
          [-0.0058,  0.0628, -0.0551]],

         [[ 0.0760,  0.0257,  0.0238],
          [ 0.0218,  0.0477, -0.0491],
          [ 0.0528, -0.0924,  0.0769]],

         [[ 0.0652, -0.0108,  0.0140],
          [ 0.0697,  0.0678,  0.0608],
          [-0.0158, -0.0084, -0.0335]]],


        [[[-0.0349,  0.0595,  0.1497],
          [-0.0481, -0.0335,  0.1096],
          [ 0.0062, -0.0172,  0.0420]],

     ...
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

    从上述输出结果中得到的结果是字典类型，其中参数的值也一起输出来了，如果想要查看具体的网络结构，则需要增加下述代码：

# 方式2-> 保存方式2，加载模型结构
vgg16 = torchvision.models.vgg16(pretrained=False)
vgg16.load_state_dict(torch.load("vgg16_model2.pth"))  # 输出完整的模型结构，与第一种方式输出的模型结构相同
print(vgg16)
1
2
3
4

    上述代码输出的网络结构如下所示：

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)
  )
)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

3.学习小结

    在本文重点介绍了如何保存和读取神经网络模型，这对日常中搭建自己的网络模型十分有用的。在下一篇博文中，将介绍神经网络模型的训练方法。