广告

扫码关注官方微信

扫码下载APP
返回顶部
首页 > 资讯 > 后端开发 > Python >PyTorch模型保存与加载实例详解
  • 572

0
分享到

PyTorch模型保存与加载实例详解

2024-04-02 19:04:59 572人浏览 独家记忆

Python 官方文档：入门教程 => 点击学习

摘要

目录一个简单的例子保存/加载 state_dict（推荐）保存/加载整个模型保存加载用于推理的常规Checkpoint/或继续训练保存多个模型到一个文件使用其他模型来预热当前模型跨设

  • 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模型保存与加载的文章就介绍到这了,更多相关PyTorch模型保存加载内容请搜索编程网以前的文章或继续浏览下面的相关文章希望大家以后多多支持编程网！

您可能感兴趣的文档:
点击免费下载>>软考高级考试备考技巧/历年真题/备考精华资料

--结束END--

本文标题: PyTorch模型保存与加载实例详解

本文链接: https://www.lsjlt.com/news/117578.html(转载时请注明来源链接)

有问题或投稿请发送至: 邮箱/279061341@qq.com    QQ/279061341

本篇文章演示代码以及资料文档资料下载

下载Word文档到电脑，方便收藏和打印～

下载Word文档
去做题
猜你喜欢
  • PyTorch模型保存与加载实例详解
    目录一个简单的例子保存/加载 state_dict（推荐）保存/加载整个模型保存加载用于推理的常规Checkpoint/或继续训练保存多个模型到一个文件使用其他模型来预热当前模型跨设...
    99+
    2022-11-10
  • PyTorch模型的保存与加载方法实例
    目录模型的保存与加载保存和加载模型参数保存和加载模型参数与结构总结模型的保存与加载 首先，需要导入两个包 import torch import torchvision.models...
    99+
    2022-11-11
  • pytorch模型的保存加载与续训练详解
    目录前面模型保存与加载方式1方式2方式3总结前面 最近，看到不少小伙伴问pytorch如何保存和加载模型，其实这部分pytorch官网介绍的也是很清楚的，感兴趣的点击了解详情 但是肯...
    99+
    2022-11-13
    pytorch模型保存加载训练 pytorch 模型训练
  • PyTorch模型保存与加载的方法
    这篇文章主要介绍了PyTorch模型保存与加载的方法的相关知识，内容详细易懂，操作简单快捷，具有一定借鉴价值，相信大家阅读完这篇PyTorch模型保存与加载的方法文章都会有所收获，下面我们一起来看看吧。state_dict 是一个Pytho...
    99+
    2023-06-30
  • pytorch模型保存与加载问题怎么解决
    这篇“pytorch模型保存与加载问题怎么解决”文章的知识点大部分人都不太理解，所以小编给大家总结了以下内容，内容详细，步骤清晰，具有一定的借鉴价值，希望大家阅读完这篇文章能有所收获，下面我们一起来看看这篇“pytorch模型保存与加载问题...
    99+
    2023-07-04
  • Pytorch模型参数的保存和加载
    目录一、前言二、参数保存三、参数的加载四、保存和加载整个模型五、总结一、前言 在模型训练完成后，我们需要保存模型参数值用于后续的测试过程。由于保存整个模型将耗费大量的存储，故推荐的做...
    99+
    2023-03-11
    Pytorch模型参数保存 Pytorch模型参数加载
  • PyTorch深度学习模型的保存和加载流程详解
    一、模型参数的保存和加载  torch.save(module.state_dict(), path)：使用module.state_dict()函数获取各层已经...
    99+
    2022-11-12
  • pytorch模型的保存和加载、checkpoint操作
    其实之前笔者写代码的时候用到模型的保存和加载，需要用的时候就去度娘搜一下大致代码，现在有时间就来整理下整个pytorch模型的保存和加载，开始学习把~ pytorch的模型和参数是分...
    99+
    2022-11-12
  • Pytorch模型参数如何保存和加载
    这篇文章主要介绍“Pytorch模型参数如何保存和加载”，在日常操作中，相信很多人在Pytorch模型参数如何保存和加载问题上存在疑惑，小编查阅了各式资料，整理出简单好用的操作方法，希望对大家解答”Pytorch模型参数如何保存和加载”的疑...
    99+
    2023-07-05
  • pytorch模型保存与加载中的一些问题实战记录
    目录前言一、torch中模型保存和加载的方式1、模型参数和模型结构保存和加载2、只保存模型的参数和加载——这种方式比较安全，但是比较稍微麻烦一点点二、torc...
    99+
    2022-11-11
  • Three.jsGLTF模型加载实现示例详解
    目录引言1. 下载GLTF模型2. 加载GLTF模型3. 添加光源4. 动画控制总结引言 在Three.js中，要加载三维模型文件，可以使用GLTF格式。GLTF是一种基于JSON...
    99+
    2023-05-16
    Three.js GLTF模型加载 Three.js GLTF
  • 模型的保存加载、模型微调、GPU使用及Pytorch常见报错
    序列化与反序列化 序列化就是说内存中的某一个对象保存到硬盘当中，以二进制序列的形式存储下来，这就是一个序列化的过程。 而反序列化，就是将硬盘中存储的二进制的数，反序列化到内存当中，得到一个相应的对象，这样就可以再次使用这个模型了。 序列化和...
    99+
    2023-08-30
    pytorch 人工智能 python
  • API模型的保存与加载方法是什么
    本篇内容介绍了“API模型的保存与加载方法是什么”的有关知识，在实际案例的操作过程中，不少人都会遇到这样的困境，接下来就让小编带领大家学习一下如何处理这些情况吧！希望大家仔细阅读，能够学有所成！　　1.目的：　　将训练好的模型保存下来，已备...
    99+
    2023-06-02
  • pytorch加载预训练模型与自己模型不匹配的解决方案
    pytorch中如果自己搭建网络并且加载别人的与训练模型的话，如果模型和参数不严格匹配，就可能会出问题，接下来记录一下我的解决方法。 两个有序字典找不同 模型的参数和pth文件的参数...
    99+
    2022-11-12
  • Tensorflow加载模型实现图像分类识别流程详解
    目录前言正文VGG19网络介绍总结前言 深度学习框架在市面上有很多。比如Theano、Caffe、CNTK、MXnet 、Tensorflow等。今天讲解的就是主角Tensorflo...
    99+
    2022-11-11
  • Spark中的数据读取保存和累加器实例详解
    目录数据读取与保存Text文件Sequence文件Object对象文件累加器累加器概念系统累加器数据读取与保存 Text文件 对于 Text文件的读取和保存 ，其语法和实现是最简单的...
    99+
    2022-11-13
    Spark数据读取保存累加器 Spark数据读取保存
软考高级职称资格查询
软考职称历年真题下载
软考职称资料下载
热门wiki
近期文章
推荐阅读
热门问答
热门标签
编程网，编程工程师的家园，是目前国内优秀的开源技术社区之一，形成了由开源软件库、代码分享、资讯、协作翻译、讨论区和博客等几大频道内容，为IT开发者提供了一个发现、使用、并交流开源技术的平台。

  • 官方手机版

  • 微信公众号

  • 商务合作