PyTorch | torch.save()函数的使用

Pytorch保存模型等相关参数，利用torch.save()，以及读取保存之后的文件。

函数信息

torch.save(obj, f, pickle_module=pickle, pickle_protocol=DEFAULT_PROTOCOL,_use_new_zipfile_serialization=True)

作用：将对象保存到磁盘文件。

参数：

• obj (object) – 保存的对象
• f (Union[str, PathLike, BinaryIO, IO[bytes]]) – 类似文件的对象（必须实现写入和刷新）或包含文件名的字符串或 os.PathLike 对象
• pickle_module (Any) – 用于pickling元数据和对象的模块
• pickle_protocol (int) – 可以指定覆盖默认协议

PyTorch 1.6 版本改用torch.save新的基于 zipfile 的文件格式。torch.load仍然保留加载旧格式文件的能力。如果出于任何原因您想torch.save 使用旧格式，请传递 kwarg _use_new_zipfile_serialization=False。

注意：

• 一种常见的PyTorch约定是使用.pt或.pth文件扩展名保存张量。

例子：

# Save to file
x = torch.tensor([0, 1, 2, 3, 4])
torch.save(x, 'tensor.pt')
# Save to io.BytesIO buffer
buffer = io.BytesIO()
torch.save(x, buffer)

模型的保存与加载

1、保存

• 建立一个字典：

state = {
    "step": step,
    "epoch": epoch,
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict()
}

• 调用torch.save(state,dir)：

checkpoint_dir = os.path.join(params.output,"model-final.pt")
torch.save(state, checkpoint_dir)

其中dir表示保存文件的绝对路径+保存文件名，如'/home/zlq/fine-turn/out/model-final.pt'。

2、加载

用来加载模型。torch.load() 使用 Python 的 解压工具（unpickling）来反序列化 pickled object 到对应存储设备上。首先在 CPU 上对压缩对象进行反序列化并且移动到它们保存的存储设备上，如果失败了（如：由于系统中没有相应的存储设备），就会抛出一个异常。用户可以通过 register_package进行扩展，使用自己定义的标记和反序列化方法。

当你想恢复某一阶段的训练（或者进行测试）时，那么就可以读取之前保存的网络模型参数等。

checkpoint = torch.load(checkpoint_dir)
model.load_state_dict(checkpoint[‘model’])
optimizer.load_state_dict(checkpoint[‘optimizer’])
start_epoch = checkpoint[‘epoch’] + 1

模型的保存与加载的使用情况

1、state_dict（推荐）

• 保存：

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

• 加载：

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

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

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

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

2、整个模型

• 保存

torch.save(model, PATH)

• 加载

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

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

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

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

3、保存加载用于推理的常规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()。

4、保存多个模型到一个文件

• 保存

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和对应的优化器到一个字典中。我们可以保存任何能帮助我们继续训练的东西到这个字典中。

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

• 保存

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可以解决这个问题。