详解PyTorch模型保存与加载


Posted in Python onApril 28, 2022
  • torch.save:保存序列化的对象到磁盘,使用了Python的pickle进行序列化,模型、张量、所有对象的字典。
  • torch.load:使用了pickle的unpacking将pickled的对象反序列化到内存中。
  • torch.nn.Module.load_state_dict:使用反序列化的state_dict加载模型的参数字典。

state_dict 是一个Python字典,将每一层映射成它的参数张量。注意只有带有可学习参数的层(卷积层、全连接层等),以及注册的缓存(batchnorm的运行平均值)在state_dict 中才有记录。state_dict同样包含优化器对象,存储了优化器的状态,所使用到的超参数。

一个简单的例子

# 定义模型
class TheModelClass(nn.Module):
    def __init__(self):
        super(TheModelClass, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

# 初始化模型
model = TheModelClass()

# 初始化优化器
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# 打印模型的 state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

# 打印优化器的 state_dict
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])

输出:

Model's state_dict:
conv1.weight     torch.Size([6, 3, 5, 5])
conv1.bias   torch.Size([6])
conv2.weight     torch.Size([16, 6, 5, 5])
conv2.bias   torch.Size([16])
fc1.weight   torch.Size([120, 400])
fc1.bias     torch.Size([120])
fc2.weight   torch.Size([84, 120])
fc2.bias     torch.Size([84])
fc3.weight   torch.Size([10, 84])
fc3.bias     torch.Size([10])

Optimizer's state_dict:
state    {}
param_groups     [{'lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [4675713712, 4675713784, 4675714000, 4675714072, 4675714216, 4675714288, 4675714432, 4675714504, 4675714648, 4675714720]}]

保存/加载 state_dict(推荐)

保存:

torch.save(model.state_dict(), PATH)

加载:

model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

要注意这个细节,如果使用nn.DataParallel在一台电脑上使用了多个GPU,那么加载模型的时候也必须先进行nn.DataParallel。

保存模型的推理过程的时候,只需要保存模型训练好的参数,使用torch.save()保存state_dict,能够方便模型的加载。因此推荐使用这种方式进行模型保存。

记住一定要使用model.eval()来固定dropout和归一化层,否则每次推理会生成不同的结果。

注意,load_state_dict()需要传入字典对象,因此需要先反序列化state_dict再传入load_state_dict()

保存/加载整个模型

保存:

torch.save(model, PATH)

加载:

# 模型类必须在别的地方定义
model = torch.load(PATH)
model.eval()

这种保存/加载模型的过程使用了最直观的语法,所用代码量少。这使用Python的pickle保存所有模块。这种方法的缺点是,保存模型的时候,序列化的数据被绑定到了特定的类和确切的目录。这是因为pickle不保存模型类本身,而是保存这个类的路径,并且在加载的时候会使用。因此,当在其他项目里使用或者重构的时候,加载模型的时候会出错。

一般来说,PyTorch的模型以.pt或者.pth文件格式保存。

一定要记住在评估模式的时候调用model.eval()来固定dropout和批次归一化。否则会产生不一致的推理结果。

保存加载用于推理的常规Checkpoint/或继续训练

保存:

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

加载:

model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# - 或者 -
model.train()

在保存用于推理或者继续训练的常规检查点的时候,除了模型的state_dict之外,还必须保存其他参数。保存优化器的state_dict也非常重要,因为它包含了模型在训练时候优化器的缓存和参数。除此之外,还可以保存停止训练时epoch数,最新的模型损失,额外的torch.nn.Embedding层等。

要保存多个组件,则将它们放到一个字典中,然后使用torch.save()序列化这个字典。一般来说,使用.tar文件格式来保存这些检查点。

加载各个组件,首先初始化模型和优化器,然后使用torch.load()加载保存的字典,然后可以直接查询字典中的值来获取保存的组件。

同样,评估模型的时候一定不要忘了调用model.eval()。

保存多个模型到一个文件

保存:

torch.save({
            'modelA_state_dict': modelA.state_dict(),
            'modelB_state_dict': modelB.state_dict(),
            'optimizerA_state_dict': optimizerA.state_dict(),
            'optimizerB_state_dict': optimizerB.state_dict(),
            ...
            }, PATH)

加载:

modelA = TheModelAClass(*args, **kwargs)
modelB = TheModelBClass(*args, **kwargs)
optimizerA = TheOptimizerAClass(*args, **kwargs)
optimizerB = TheOptimizerBClass(*args, **kwargs)

checkpoint = torch.load(PATH)
modelA.load_state_dict(checkpoint['modelA_state_dict'])
modelB.load_state_dict(checkpoint['modelB_state_dict'])
optimizerA.load_state_dict(checkpoint['optimizerA_state_dict'])
optimizerB.load_state_dict(checkpoint['optimizerB_state_dict'])

modelA.eval()
modelB.eval()
# - 或者 -
modelA.train()
modelB.train()

保存的模型包含多个torch.nn.Modules时,比如GAN,一个序列-序列模型,或者组合模型,使用与保存常规检查点的方式来保存模型。也就是说,保存每个模型的state_dict和对应的优化器到一个字典中。我们可以保存任何能帮助我们继续训练的东西到这个字典中。

使用其他模型来预热当前模型

保存:

torch.save(modelA.state_dict(), PATH)

加载:

modelB = TheModelBClass(*args, **kwargs)
modelB.load_state_dict(torch.load(PATH), strict=False)

在迁移学习或者训练新的复杂模型时,加载部分模型是很常见的。利用经过训练的参数,即使只有少数参数可用,也将有助于预热训练过程,并且使模型更快收敛。

在加载部分模型参数进行预训练的时候,很可能会碰到键不匹配的情况(模型权重都是按键值对的形式保存并加载回来的)。因此,无论是缺少键还是多出键的情况,都可以通过在load_state_dict()函数中设定strict参数为False来忽略不匹配的键。

如果想将某一层的参数加载到其他层,但是有些键不匹配,那么修改state_dict中参数的key可以解决这个问题。

跨设备保存与加载模型

GPU上保存,CPU上加载

保存:

torch.save(model.state_dict(), PATH)

加载:

device = torch.device('cpu')
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location=device))

当在CPU上加载一个GPU上训练的模型时,在torch.load()中指定map_location=torch.device('cpu'),此时,map_location动态地将tensors的底层存储重新映射到CPU设备上。

上述代码只有在模型是在一块GPU上训练时才有效,如果模型在多个GPU上训练,那么在CPU上加载时,会得到类似如下错误:

KeyError: ‘unexpected key “module.conv1.weight” in state_dict'

原因是在使用多GPU训练并保存模型时,模型的参数名都带上了module前缀,因此可以在加载模型时,把key中的这个前缀去掉:

# 原始通过DataParallel保存的文件
state_dict = torch.load('myfile.pth.tar')
# 创建一个不包含`module.`的新OrderedDict
from collections import OrderedDict
new_state_dict = OrderedDict()
for k, v in state_dict.items():
    name = k[7:] # 去掉 `module.`
    new_state_dict[name] = v
# 加载参数
model.load_state_dict(new_state_dict)

GPU上保存,GPU上加载

保存:

torch.save(model.state_dict(), PATH)

加载:

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device)
# 往模型中输入数据的时候不要忘记在任意tensor上调用input = input.to(device)

在把GPU上训练的模型加载到GPU上时,只需要使用model.to(torch.devie('cuda'))将初始化的模型转换为CUDA优化模型。同时确保在模型所有的输入上使用.to(torch.device('cuda'))。注意,调用my_tensor.to(device)会返回一份在GPU上的my_tensor的拷贝。不会覆盖原本的my_tensor,因此要记得手动将tensor重写:my_tensor = my_tensor.to(torch.device('cuda'))。

CPU上保存,GPU上加载

保存:

torch.save(model.state_dict(), PATH)

加载:

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH, map_location="cuda:0"))  # 选择希望使用的GPU
model.to(device)

保存torch.nn.DataParallel模型

保存:

torch.save(model.module.state_dict(), PATH)

 总结

到此这篇关于PyTorch模型保存与加载的文章就介绍到这了!


Tags in this post...

Python 相关文章推荐
Python中用altzone()方法处理时区的教程
May 22 Python
Python的网络编程库Gevent的安装及使用技巧
Jun 24 Python
python批量制作雷达图的实现方法
Jul 26 Python
用python实现简单EXCEL数据统计的实例
Jan 24 Python
python模拟登录并且保持cookie的方法详解
Apr 04 Python
Python基于hashlib模块的文件MD5一致性加密验证示例
Feb 10 Python
Python爬虫常用小技巧之设置代理IP
Sep 13 Python
Python 运行 shell 获取输出结果的实例
Jan 07 Python
Pycharm之快速定位到某行快捷键的方法
Jan 20 Python
python 获取微信好友列表的方法(微信web)
Feb 21 Python
如何理解python对象
Jun 21 Python
PyTorch 中的傅里叶卷积实现示例
Dec 11 Python
Python 图片添加美颜效果
Python 视频画质增强
python 单机五子棋对战游戏
python井字棋游戏实现人机对战
Apr 28 #Python
Python开发五子棋小游戏
Python简易开发之制作计算器
Apr 28 #Python
Python实现对齐打印 format函数的用法
Apr 28 #Python
You might like
PHP的面试题集
2006/11/19 PHP
php ftp文件上传函数(基础版)
2010/06/03 PHP
Yii使用DeleteAll连表删除出现报错问题的解决方法
2016/07/14 PHP
php中对象引用和复制实例分析
2019/08/14 PHP
AutoSave/自动存储功能实现
2007/03/24 Javascript
javascript css styleFloat和cssFloat
2010/03/15 Javascript
node.js中的buffer.write方法使用说明
2014/12/10 Javascript
JS给超链接加确认对话框的方法
2015/02/24 Javascript
jQuery结合ajax实现动态加载文本内容
2015/05/19 Javascript
jquery实现的简单二级菜单效果代码
2015/09/22 Javascript
JS获取地址栏参数的两种方法(简单实用)
2016/06/14 Javascript
100多个基础常用JS函数和语法集合大全
2017/02/16 Javascript
Angularjs 双向绑定时字符串的转换成数字类型的问题
2017/06/12 Javascript
Angular中的$watch方法详解
2017/09/18 Javascript
在微信小程序中使用vant的方法
2019/06/07 Javascript
vue element-ui读取pdf文件的方法
2019/11/26 Javascript
JavaScript事件冒泡机制原理实例解析
2020/01/14 Javascript
JS JQuery获取data-*属性值方法解析
2020/09/01 jQuery
详解Vue.js 可拖放文本框组件的使用
2021/03/03 Vue.js
对python中的try、except、finally 执行顺序详解
2019/02/18 Python
python BlockingScheduler定时任务及其他方式的实现
2019/09/19 Python
tensorflow 报错unitialized value的解决方法
2020/02/06 Python
翻转数列python实现,求前n项和,并能输出整个数列的案例
2020/05/03 Python
Python3实现英文字母转换哥特式字体实例代码
2020/09/01 Python
Python 多进程、多线程效率对比
2020/11/19 Python
html5 canvas的绘制文本自动换行的示例代码
2018/09/17 HTML / CSS
APM Monaco中国官网:来自摩纳哥珠宝品牌
2017/12/27 全球购物
Muziker英国:中欧最大的音乐家商店
2020/02/05 全球购物
普通大学毕业生自荐信
2013/11/04 职场文书
口头翻译求职人自荐信
2013/12/07 职场文书
党的群众路线教育实践活动总结报告
2014/04/28 职场文书
机关班子查摆问题及整改措施
2014/10/28 职场文书
2015年高三班主任工作总结
2015/05/21 职场文书
2016年党员公开承诺书范文
2016/03/24 职场文书
JavaScript原型链详解
2021/11/07 Javascript
MySQL详细讲解变量variables的用法
2022/06/21 MySQL