PyTorch学习（8）:模型保存和加载

PyTorch学习（8）:模型保存和加载

• Pytorch官方文档： https://pytorch-cn.readthedocs.io/zh/latest/
• Pytorch学习文档： https://github.com/tensor-yu/PyTorch_Tutorial
• 参考： https://mp.weixin.qq.com/s/iYuMcbnJqVUMDMvwDEPs6Q
      https://www.dtmao.cc/news_show_329038.shtml

---

前言

PyTorch框架在保存和加载模型时，需要熟悉三个核心功能：
（1）torch.save: 将序列化对象保存到磁盘。此函数使用Python的pickle模块进行系列化。使用此函数可以保存如模型、tensor、字典等各种对象。
（2）torch.load: 使用pickle的unpickling功能将pickle对象文件反序列化到内存。此功能还可以有助于设备加载数据。
（3）torch.nn.Module.load_state_dict: 使用反序列化函数state_dict来加载模型的参数字典。

---

1.保存和加载推理模型

官方链接： https://pytorch-cn.readthedocs.io/zh/latest/notes/serialization/

1）保存/加载state_dict（推荐使用：只保存和加载模型参数）

（1）保存

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

（2）加载

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

当保存好模型用来推断的时候，只需要保存模型学习到的参数，使用torch.save()函数来保存模型state_dict，它会给模型恢复提供最大的灵活性，这就是为什么要推荐它来保存的原因。

2）保存/加载完整模型

（1）保存

torch.save(model, PATH)

（2）加载

# 模型类必须在此之前被定义
model = torch.load(PATH)
model.eval()

此部分保存/加载过程使用最直观的语法并涉及最少量的代码。以Python‘pickle’模块的方式来保存模型。这种方法的缺点是序列化数据受限于某种特殊的类而且需要确切的字典结构。这是因为pickle无法保存模型类本身。相反，它保存包含类的文件的路径，该文件在加载时使用。因此，当在其他项目使用或者重构之后，您的代码可能会以各种方式中断。

注1： 在PyTorch中最常见的模型保存使 ‘.pt’或者是‘.pth’作为模型文件扩展名 。
注2： 在运行推理之前，务必调用model.eval()去设置dropout和batch normalization层为评估模式。如果不这么做，可能导致模型推断结果不一致。

---

2.保存和加载Checkpoint用于推理/继续训练

（1）保存

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

（2）加载

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也很重要，因为它包含作为模型训练更新的缓冲区和参数。也可以保存其他项目，比如最新记录的训练损失、迭代次数等等。

注1： 要保存多个组件，请在字典中组织它们并使用torch.save()来序列化字典。PyTorch中常见的保存Checkpoint是使用 .tar文件扩展名 。
注2： 如果你想要恢复训练，请调用 model.train() 以确保dropout、BN等层处于训练模式。

---

3.在一个文件中保存多个模型

（1）保存

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)

（2）加载

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(对抗生成网络)、sequence-to-sequence (序列到序列模型)，或者是多个模型融合，可以采用与保存常规检查点相同的方法。换句话说，保存每个模型的state_dict的字典和相对应的优化器。如前所述，可以通过将它们附加到字典的方式来保存任何其他项目，这样有助于恢复训练。

---

4.使用在不同模型条件下的热启动模式

（1）保存

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

（2）加载

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

在迁移学习或者训练新的复杂模型时，加载部分模型是常见的情况。利用训练好的参数，有助于热启动训练过程，并希望帮助你的模型比从头开始训练能够更快地收敛。
无论是从缺少某些键的state_dict加载还是从键的数目多于加载模型的state_dict, 都可以通过在load_state_dict()函数中将strict参数设置为False来忽略非匹配键的函数。

注： 如果要将参数从一个层加载到另一个层，但是某些键不匹配，主要修改正在加载的state_dict 中的参数键的名称以匹配要在加载到模型中的键即可。

---

5.通过设备保存/加载模型

1）保存到CPU、加载到CPU

（1）保存

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

（2）加载

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

当在CPU上加载GPU上训练的模型时, 将torch.device(‘cpu’)传递给torch.load()函数中的map_location参数。在这种情况下，使用map_location参数将张量下的存储器动态的重新映射到CPU设备。

2）保存到GPU、加载到GPU

（1）保存

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

（2）加载

device = torch.device("cuda")
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.to(device) # 确保在你提供给模型的任何输入张量上调用input = input.to(device)

当在GPU上训练并把模型保存到GPU上时，只需要使用model.to(torch.device(‘cuda’)) ，将初始化的model转换为CUDA优化模型。另外，请务必在所有模型输入上使用.to(torch.device(‘cuda’)) 函数来为模型准备数据。请注意，调用 my_tensor.to(device)会在GPU上返回my_tensor的副本。因此，请记住手动覆盖张量：my_tensor=my_tensor.to(torch.device(‘cuda’))。

3）保存到CPU，加载到GPU

（1）保存

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

（2）加载

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) # 确保提供给模型的任何输入张量上调用input = input.to(device)

在CPU上训练好并保存的模型加载到GPU时，将torch.load()函数中的map_location参数设置为cuda:device_id。这会将模型加载到指定的GPU设备。接下来，请务必调model.to(torch.device(‘cuda’)) 将模型的参数张量转换为CUDA张量。最后，确保在所有模型输入上使用.to(torch.device(‘cuda’))函数来为CUDA优化模型。请注意，调用my_tensor.to(device) 会在GPU上返回my_tensor的新副本。它不会覆盖my_tensor。因此，请手动覆盖张量my_tensor = my_tensor.to(torch.device(‘cuda’))。

4）保存torch.nn.DataParallel模型

（1）保存

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

（2）加载

# 加载任何你想要的设备

torch.nn.DataParallel 是一个模型封装，支持并行GPU使用。要普通保存DataParallel模型,
请保存model.module.state_dict()。这样，就可以非常灵活地以任何方式加载模型到你想要的设备中。

---

6.Torch.save的版本更新

官方文档： https://pytorch.org/docs/stable/generated/torch.save.html
在Torch官网查看了torch.save的介绍，可以看到在torch1.6版本中，对torch.save进行了更改，如下图：

为了兼容Torch1.6之后版本到之前的版本，有如下两种解决方法：

（1）解决方法1

torch.save(model.state_dict(), PATH, _use_new_zipfile_serialization=False)  # 训练所有数据后，保存网络的参数

（2）解决方法2

如果已经使用Torch 1.6之后的版本进行了训练，并且在训练训练过程中没有更改参数“_use_new_zipfile_serialization=False”，为了节省时间，大家可以在1.6之后版本中直接加载模型，然后再次使用torch.save 进行保存为非zip格式的。具体代码如下：

#在torch 1.6之后版本中重新加载一下网络参数
model = MyNetwork().to(device) #实例化模型并加载GPU中
model.load_state_dict(torch.load(PATH))  #加载模型参数
#重新保存网络参数，此时注意改为非zip格式
torch.save(model.state_dict(), model_cp,_use_new_zipfile_serialization=False)

总结

  至此，基于PyTorch的模型保存和加载已基本了解，可根据应用场景进行选择使用。