模型保存与提取
- 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.csdn.net/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>)
只有net3的输出是固定的,因为在模型导入的时候执行了net3.eval().
结论:无论采用 方式1 还是 方式2 导入的模型, 在模型测试时,都需要用.eval()方法固定一下网络在训练过程中的随机项目,如dropout 等,避免网络在同一个输入下产生不一样的结果。