本文以VGG16模型为例,来学习神经网络模型的保存与读取,其中网络模型的读取有两种方法,对应的网络模型的读取也有两种方法,下面来一一学习。
文章目录
一、网络模型的保存方法
1 保存方法一
2 保存方法二
二、网络模型的读取方法
1 读取方法一
2 读取方法二
三、两种方法的对比及注意事项
1 解决方法一
2 解决方法二
网络模型的保存有两种方法
以vgg16模型为例,代码如下:
import torch
import torchvisionvgg16 = torchvision.models.vgg16(pretrained=False)
#保存方法一
torch.save(vgg16, "vgg16_model.pth")
其中,
pretrained=False表示使用初始化的参数,没有经过数据集训练,
vgg16,是要保存的网络模型
"vgg16_model.pth"是给这个保存的网络模型一个名字
.pth是常用的后缀格式,
点击运行,就会在左边创建一个网络模型vgg16_model.pth文件
方法一保存,不仅保存了网络模型的结构,也保存了网络模型的参数
方法二,保存的是模型参数(官方推荐)因为这种保存方式的内存占比较小
import torch
import torchvisionvgg16 = 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文件
同上图一样。二、网络模型的读取方法
对应上面的保存方法一,代码如下:
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的模型结构输出了
对应上面的保存方法二, 代码如下:
方法二对应的网络模型的读取
import torchmodel2 = 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 xtest = Test()
torch.save(test,"test.pth")
我们来看看直接读取这个test.pth会发生什么
新创建一个py文件,输入:
import torch
test1 = torch.load("test.pth")
print()
结果直接报错: AttributeError: Can't get attribute 'Test' on
大概意思就是:你不能得到Test这个属性,因为你没有这个类
解决这种报错一般有两种方法
方法一:将Test这个网络模型给复制过来,但是不需要创建神经网络的,也就是不需要test = Test() ,这一步的。
只需要将模型的定义写过来,目的应该就是确保,你加载的这个网络模型是你想要的网络模型
正确的应该这样写
import torch
from torch import nnclass 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 xtest1 = torch.load("test.pth")
print(test1)输出结果:
Test(
(conv): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1))
)
将写有这模型的文件直接引用过来就行了
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定义的网络模型引用过来
这样也是可以的,具体怎么操作要看实际情况来.