Pytorch(4)-模型保存-载入-eval()

原创已于 2022-02-26 13:41:32 修改 · 4.3k 阅读

9 ·

CC 4.0 BY-SA版权

文章标签：

#pytorch #深度学习 #python

于 2019-07-20 17:49:32 首次发布

MLTools 专栏收录该内容

13 篇文章

订阅专栏

模型保存与提取

1. 整个模型保存-载入
2. 仅模型参数保存-载入
3. GPU/CPU模型保存与导入
4. net.eval()--固定模型随机项

神经网络模型在线训练完之后需要保存下来，以便下次使用时可以直接导入已经训练好的模型。pytorch 提供两种方式保存模型:

方式1：保存整个网络，载入时直接载入整个网络，优点：代码简单，缺点需要的存储空间大

方式2：只保存网络参数，载入时需要先建立与原来网络一样结构的网络，然后将网络参数导入到该网络中，方式2的优缺点与方式1相反。

1. 整个模型保存-载入

模型的结构参数都保存下来了

# 保存模型：设置 保存目录 和 保存文件名.扩展名，常用扩展名： .pkl .pth （扩展名只要好辨识就即可）
PATH="./model/mynet1.pkl"
# 导入官方提供的预训练模型
net1=torchvision.models.alexnet(pretrainend=True)
# 用数据集训练网络
.....
# 保存训练好的网络
torch.save(net1, PATH)
-----------------------------------------------------------
# 载入模型：设置载入路径，即模型保存的路径
PATH="./model/mynet1.pkl"
net1_1=torch.load(PATH)

2. 仅模型参数保存-载入

保存时–只保存网络中的参数 (速度快, 占内存少)，载入时–需要提前创建好结构和net2是一样的

# 保存模型：设置 保存目录 和 保存文件名.扩展名，常用扩展名： .pkl .pth （扩展名只要好辨识就即可）
PATH="./model/mynet2.pkl"
# 导入官方提供的预训练模型
net2=torchvision.models.alexnet(pretrainend=True)
# 用数据集训练网络
.....
# 保存训练好的网络
torch.save(net1.state_dict(), PATH)
-----------------------------------------------------------
# 载入模型：设置载入路径，即模型保存的路径
PATH="./model/net2.pkl"
# 新建一个网络
net2_2=torchvision.models.alexnet(pretrained=True)
# 载入模型参数
net2_2.load_state_dict(torch.load(PATH))

迷糊的现象

在使用莫烦的文档做实验时，保存的两个文件：net.pkl，net_params.pkl大小差异比较大。保证在导入模型是比较快。
在这里插入图片描述
但是使用torchvision.models.模块中的一系列网络时，因为网络的参数很大，所以实验过程中用两种方法保存模型的文件大小是一致的。（猜测是内置模型使用torch.save(net1, ‘net.pkl’)时默认保存的是模型参数）

提供一个神经网络模型占用空间大小的计算方法：
在这里插入图片描述
参考文档：经典CNN模型计算量与内存需求分析

3. GPU/CPU模型保存与导入

在训练是模型是GPU/CPU，决定了模型载入时的模型原型。可以分为下面三种情况
(只展示导入整个网络模型的情况,具体实验还没做过)：

1.CPU(原型)->CPU, GPU(原型)->GPU

torch.load( ‘net.pkl’)

2.GPU(原型)->CPU

torch.load(‘model_dict.pkl’, map_location=lambda storage, loc: storage)

3.CPU(模型文件)->GPU

torch.load(‘model_dic.pkl’, map_location=lambda storage, loc: storage.cuda)

参考文档：https://blog.youkuaiyun.com/u012135425/article/details/85217542

4. net.eval()–固定模型随机项

两种模型载入方式、.eval() 作用实验demo

step1: 载入模型

# 20191204 pytorch 模型载入测试
import torchvision as tvt
import torch
net1=tvt.models.alexnet(pretrained=True)  # 1.自动从网上下载的预先训练模型
net2=torch.load("./model/mynet1.pkl")     # 2.导入事先训练好的保存的整个网络

net3=tvt.models.alexnet(pretrained=True)  # 3.导入只保存模型参数的网络,需要新建一个网络
net3.load_state_dict(torch.load("../model/mynet2.pkl"))
net3.eval（）                              #   固定dropout和归一化层，否则每次推理会生成不同的结果。

step2：输出三个网络同一层参数的和，net2 和net3 对应参数相等。可以看出来，两种模型保存和导入方式是等价的。

net1 tensor(-21257.7656, grad_fn=<SumBackward0>)
net2 tensor(-21253.9473, device='cuda:0', grad_fn=<SumBackward0>)
net3 tensor(-21253.9551, grad_fn=<SumBackward0>)

step3: 产生一个随机输入a,输入到网络1,2,3，打印输出结果。

a=torch.randn([1,3,224,224])
y1=net1(a)
y2=net2(a)
y3=net3(a)
# 第二次输入
y11=net1(a)
y22=net2(a)
y33=net3(a)
# 打印y1,y2,y3,y11,y22,y33（1000维的和）
y1: tensor(-5.2689, grad_fn=<SumBackward0>)
y2: tensor(-1.6695, device='cuda:0', grad_fn=<SumBackward0>)
y3: tensor(-4.4349, device='cuda:0', grad_fn=<SumBackward0>)

y11: tensor(-4.4205, grad_fn=<SumBackward0>)
y22: tensor(-5.9475, device='cuda:0', grad_fn=<SumBackward0>)
y33: tensor(-4.4349, device='cuda:0', grad_fn=<SumBackward0>)