利用pytorch自定义CNN网络(五):保存、加载自定义模型【转载】

这篇具有很好参考价值的文章主要介绍了利用pytorch自定义CNN网络(五):保存、加载自定义模型【转载】。希望对大家有所帮助。如果存在错误或未考虑完全的地方,请大家不吝赐教,您也可以点击"举报违法"按钮提交疑问。

本文转载自:PyTorch | 保存和加载模型

1. 简介

本文主要介绍如何加载和保存 PyTorch 的模型。这里主要有三个核心函数:

  1. torch.save :把序列化的对象保存到硬盘。它利用了 Python 的 pickle 来实现序列化。模型、张量以及字典都可以用该函数进行保存;
  2. torch.load:采用 pickle 将反序列化的对象从存储中加载进来。
  3. torch.nn.Module.load_state_dict:采用一个反序列化的 state_dict加载一个模型的参数字典。

本文主要内容如下:

  • 什么是状态字典(state_dict)?
  • 预测时加载和保存模型
  • 加载和保存一个通用的检查点(Checkpoint)
  • 在同一个文件保存多个模型
  • 采用另一个模型的参数来预热模型(Warmstaring Model)
  • 不同设备下保存和加载模型

2. 什么是状态字典(state_dict)

PyTorch 中,一个模型(torch.nn.Module)的可学习参数(也就是权重和偏置值)是包含在模型参数(model.parameters())中的,一个状态字典就是一个简单的 Python 的字典,其键值对是每个网络层和其对应的参数张量。模型的状态字典只包含带有可学习参数的网络层(比如卷积层、全连接层等)和注册的缓存(batchnorm的 running_mean)。优化器对象(torch.optim)同样也是有一个状态字典,包含的优化器状态的信息以及使用的超参数。
由于状态字典也是 Python 的字典,因此对 PyTorch 模型和优化器的保存、更新、替换、恢复等操作都很容易实现。
下面是一个简单的使用例子,例子来自:Training a Classifier — PyTorch Tutorials 2.0.1+cu117 documentation

# Define model
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

        # Initialize model
model = TheModelClass()

# Initialize optimizer
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)

# Print model's state_dict
print("Model's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

    # Print optimizer's state_dict
print("Optimizer's state_dict:")
for var_name in optimizer.state_dict():
    print(var_name, "\t", optimizer.state_dict()[var_name])

上述代码先是简单定义一个 5 层的 CNN,然后分别打印模型的参数和优化器参数。
输出结果:

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]}]

3. 预测时加载和保存模型

3.1. 加载/保存状态字典(推荐做法)

保存的代码:

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

加载的代码:

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

当需要为预测保存一个模型的时候,只需要保存训练模型的可学习参数即可。采用 torch.save() 来保存模型的状态字典的做法可以更方便加载模型,这也是推荐这种做法的原因。
通常会用 .pt 或者 .pth 后缀来保存模型。
记住

  1. 在进行预测之前,必须调用 model.eval() 方法来将 dropout 和 batch normalization 层设置为验证模型。否则,只会生成前后不一致的预测结果。
  2. load_state_dict() 方法必须传入一个字典对象,而不是对象的保存路径,也就是说必须先反序列化字典对象,然后再调用该方法,也是例子中先采用 torch.load() ,而不是直接 model.load_state_dict(PATH)

3.2. 加载/保存整个模型

保存

torch.save(model, PATH)

加载

# Model class must be defined somewhere
model = torch.load(PATH)
model.eval()

保存和加载模型都是采用非常直观的语法并且都只需要几行代码即可实现。这种实现保存模型的做法将是采用 Python 的 pickle 模块来保存整个模型,这种做法的缺点就是序列化后的数据是属于特定的类和指定的字典结构,原因就是 pickle 并没有保存模型类别,而是保存一个包含该类的文件路径,因此,当在其他项目或者在 refactors 后采用都可能出现错误。

4. 加载和保存一个通用的检查点(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()
# - or -
model.train()

当保存一个通用的检查点(checkpoint)时,无论是用于继续训练还是预测,都需要保存更多的信息,不仅仅是 state_dict ,比如说优化器的 state_dict 也是非常重要的,它包含了用于模型训练时需要更新的参数和缓存信息,还可以保存的信息包括 epoch,即中断训练的批次,最后一次的训练 loss,额外的 torch.nn.Embedding 层等等。
上述保存代码就是介绍了如何保存这么多种信息,通过用一个字典来进行组织,然后继续调用 torch.save 方法,一般保存的文件后缀名是 .tar 。
加载代码也如上述代码所示,首先需要初始化模型和优化器,然后加载模型时分别调用 torch.load 加载对应的 state_dict 。然后通过不同的键来获取对应的数值。
加载完后,根据后续步骤,调用 model.eval() 用于预测,model.train() 用于恢复训练。

5. 在同一个文件保存多个模型

保存模型的示例代码

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()
# - or -
modelA.train()
modelB.train()

当我们希望保存的是一个包含多个网络模型 torch.nn.Modules 的时候,比如 GAN、一个序列化模型,或者多个模型融合,实现的方法其实和保存一个通用的检查点的做法是一样的,同样采用一个字典来保持模型的 state_dict 和对应优化器的 state_dict 。除此之外,还可以继续保存其他相同的信息。
加载模型的示例代码如上述所示,和加载一个通用的检查点也是一样的,同样需要先初始化对应的模型和优化器。同样,保存的模型文件通常是以 .tar 作为后缀名。

6. 采用另一个模型的参数来预热模型(Warmstaring Model)

保存模型的示例代码

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

加载模型的示例代码

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

在之前迁移学习教程中也介绍了可以通过预训练模型来微调,加快模型训练速度和提高模型的精度。
这种做法通常是加载预训练模型的部分网络参数作为模型的初始化参数,然后可以加快模型的收敛速度。
加载预训练模型的代码如上述所示,其中设置参数 strict=False 表示忽略不匹配的网络层参数,因为通常我们都不会完全采用和预训练模型完全一样的网络,通常输出层的参数就会不一样。
当然,如果希望加载参数名不一样的参数,可以通过修改加载的模型对应的参数名字,这样参数名字匹配了就可以成功加载。

7. 不同设备下保存和加载模型

7.1. 在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'),这个做法会将张量都重新映射到 CPU 上。

7.2. 在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)
                      # Make sure to call input = input.to(device) on any input tensors that you feed to the model

在 GPU 上训练和加载模型,调用 torch.load() 加载模型后,还需要采用 model.to(torch.device('cuda')),将模型调用到 GPU 上,并且后续输入的张量都需要确保是在 GPU 上使用的,即也需要采用 my_tensor.to(device)。

7.3. 在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"))  # Choose whatever GPU device number you want
model.to(device)
# Make sure to call input = input.to(device) on any input tensors that you feed to the model

这次是 CPU 上训练模型,但在 GPU 上加载模型使用,那么就需要通过参数 map_location 指定设备。然后继续记得调用 model.to(torch.device('cuda'))。

7.4. 保存 torch.nn.DataParallel 模型

保存模型的示例代码

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

torch.nn.DataParallel 是用于实现多 GPU 并行的操作,保存模型的时候,是采用 model.module.state_dict()。
加载模型的代码也是一样的,采用 torch.load() ,并可以放到指定的 GPU 显卡上。文章来源地址https://www.toymoban.com/news/detail-638398.html

到了这里,关于利用pytorch自定义CNN网络(五):保存、加载自定义模型【转载】的文章就介绍完了。如果您还想了解更多内容,请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章,希望大家以后多多支持TOY模板网!

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处: 如若内容造成侵权/违法违规/事实不符,请点击违法举报进行投诉反馈,一经查实,立即删除!

领支付宝红包 赞助服务器费用

相关文章

  • 【AI】《动手学-深度学习-PyTorch版》笔记(十六):自定义网络层、保存/加载参数、使用GPU

    自定义网络层很简单,三步即可完成 继承类:nn.Module 定义初始化函数:__init__中定义需要初始化的代码 定义向前传播函数:forward 1)定义网络层

    2024年02月13日
    浏览(47)
  • PyTorch模型的保存与加载

    载入muti-GPU模型: 载入muti-GPU权重: 载入CPU权重: 模型保存的格式: pytorch中最常见的模型保存使用 .pt 或者是 .pth 作为模型文件扩展名,其他方式还有.t7/.pkl格式,t7文件是沿用torch7中读取模型权重的方式,而在keras中则是使用.h5文件 .pth 文件基本信息 四个键值: model(Ord

    2023年04月21日
    浏览(38)
  • pytorch保存、加载和解析模型权重

    1、模型保存和加载          主要有两种情况:一是仅保存参数,二是保存参数及模型结构。 保存参数:          torch.save(net.state_dict()) 加载参数(加载参数前需要先实例化模型):          param = torch.load(\\\'param.pth\\\')          net.load_state_dict(param) 保存模型结构

    2024年02月16日
    浏览(43)
  • 现有模型的保存与加载(PyTorch版)

    我们以VGG16网络为例,来说明现有模型的保存与加载操作。 保存与加载方式均有两种,接下来我们分别来学习这两种方式。 注意:保存与加载不在同一个py文件中,我们设定保存操作在save.py文件中,而加载操作在load.py文件中。 保存模型的两种方式如下代码所示,第一种为既

    2024年02月09日
    浏览(45)
  • pytorch11:模型加载与保存、finetune迁移训练

    往期回顾 pytorch01:概念、张量操作、线性回归与逻辑回归 pytorch02:数据读取DataLoader与Dataset、数据预处理transform pytorch03:transforms常见数据增强操作 pytorch04:网络模型创建 pytorch05:卷积、池化、激活 pytorch06:权重初始化 pytorch07:损失函数与优化器 pytorch08:学习率调整策略

    2024年02月01日
    浏览(48)
  • TensorFlow 高级技巧:自定义模型保存、加载和分布式训练

    本篇文章将涵盖 TensorFlow 的高级应用,包括如何自定义模型的保存和加载过程,以及如何进行分布式训练。 在 TensorFlow 中,我们可以通过继承 tf.train.Checkpoint 来自定义模型的保存和加载过程。 以下是一个例子: TensorFlow 提供了 tf.distribute.Strategy API,让我们可以在不同的设备

    2024年02月15日
    浏览(37)
  • 【模型部署】利用opencv Python API 加载与运行Pytorch模型

    OpenCV可以调用PyTorch模型。通过将PyTorch模型转换为ONNX格式,然后使用OpenCV中的dnn模块加载和运行模型。 以下是一个使用OpenCV加载和运行PyTorch模型的示例代码: 说明: 首先,使用torch.load函数加载PyTorch模型。 将模型转换为ONNX格式,使用torch.onnx.export函数将模型转换为ONNX格式

    2024年02月13日
    浏览(35)
  • 如何用pytorch做文本摘要生成任务(加载数据集、T5 模型参数、微调、保存和测试模型,以及ROUGE分数计算)

    摘要 :如何使用 Pytorch(或Pytorchlightning) 和 huggingface Transformers 做文本摘要生成任务,包括数据集的加载、模型的加载、模型的微调、模型的验证、模型的保存、ROUGE指标分数的计算、loss的可视化。 ✅ NLP 研 0 选手的学习笔记 ● python 需要 3.8+ ● 文件相对地址 : mian.py 和 tra

    2024年02月05日
    浏览(73)
  • PyTorch高级教程:自定义模型、数据加载及设备间数据移动

    在深入理解了PyTorch的核心组件之后,我们将进一步学习一些高级主题,包括如何自定义模型、加载自定义数据集,以及如何在设备(例如CPU和GPU)之间移动数据。 虽然PyTorch提供了许多预构建的模型层,但在某些情况下,你可能需要自定义模型层。这可以通过继承 torch.nn.Mo

    2024年02月14日
    浏览(32)
  • 【Pytorch】计算机视觉项目——卷积神经网络CNN模型识别图像分类

    在上一篇笔记《【Pytorch】整体工作流程代码详解(新手入门)》中介绍了Pytorch的整体工作流程,本文继续说明如何使用Pytorch搭建卷积神经网络(CNN模型)来给图像分类。 其他相关文章: 深度学习入门笔记:总结了一些神经网络的基础概念。 TensorFlow专栏:《计算机视觉入门

    2024年02月05日
    浏览(57)

觉得文章有用就打赏一下文章作者

支付宝扫一扫打赏

博客赞助

微信扫一扫打赏

请作者喝杯咖啡吧~博客赞助

支付宝扫一扫领取红包,优惠每天领

二维码1

领取红包

二维码2

领红包