在 PyTorch 中保存训练模型的最佳方法是啥？ [关闭]

Posted 2023-03-29

【中文标题】在 PyTorch 中保存训练模型的最佳方法是啥？ [关闭]

【英文标题】：Best way to save a trained model in PyTorch? [closed]在 PyTorch 中保存训练模型的最佳方法是什么？ [关闭]

【发布时间】：2021-08-19 06:34:53

【问题描述】：

我一直在寻找在 PyTorch 中保存训练模型的替代方法。到目前为止，我已经找到了两种选择。

torch.save() 保存模型，torch.load() 加载模型。 model.state_dict() 保存经过训练的模型，model.load_state_dict() 加载保存的模型。

我遇到过这个discussion，建议使用方法 2 而不是方法 1。

我的问题是，为什么首选第二种方法？难道仅仅是因为torch.nn模块有这两个功能，我们被鼓励使用它们吗？

【问题讨论】：

我认为这是因为 torch.save() 也保存了所有中间变量，例如用于反向传播的中间输出。但是您只需要保存模型参数，例如权重/偏差等。有时前者可能比后者大得多。

我测试了torch.save(model, f) 和torch.save(model.state_dict(), f)。保存的文件大小相同。现在我很困惑。另外，我发现使用 pickle 保存 model.state_dict() 非常慢。我认为最好的方法是使用torch.save(model.state_dict(), f)，因为您处理模型的创建，而torch 处理模型权重的加载，从而消除了可能的问题。参考：discuss.pytorch.org/t/saving-torch-models/838/4

似乎 PyTorch 在他们的 tutorials section 中更明确地解决了这个问题——这里的答案中没有列出很多好的信息，包括一次保存多个模型和热启动模型。

使用pickle有什么问题？

@CharlieParker torch.save 基于泡菜。以下来自上面链接的教程：“[torch.save]将使用Python的pickle模块保存整个模块。这种方法的缺点是序列化数据绑定到特定的类和模型时使用的确切目录结构被保存。这是因为pickle不保存模型类本身。相反，它保存了包含该类的文件的路径，在加载时使用。因此，您的代码可能会以各种方式中断用于其他项目或重构后。”

【参考方案1】：

在他们的 github 仓库中找到 this page，我将在此处复制粘贴内容。

---

保存模型的推荐方法

序列化和恢复模型有两种主要方法。

第一个（推荐）只保存和加载模型参数：

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

然后：

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

第二个保存并加载整个模型：

torch.save(the_model, PATH)

然后：

the_model = torch.load(PATH)

但是在这种情况下，序列化的数据绑定到特定的类 以及使用的确切目录结构，因此它可以以各种方式中断 在其他项目中使用，或者经过一些严重的重构。

---

更新：另请参阅 PyTorch 教程中的 Save and Load the Model 部分

【讨论】：

根据@smth discuss.pytorch.org/t/saving-and-loading-a-model-in-pytorch/… 模型默认重新加载以训练模型。所以加载后需要手动调用the_model.eval()，如果你加载它是为了推理，而不是恢复训练。

第二种方法在 Windows 10 上出现***.com/questions/53798009/… 错误。无法解决

有什么选项可以保存而不需要访问模型类？

通过这种方法，您如何跟踪您需要为负载情况传递的 *args 和 **kwargs？

其实我只需要保存一个nn.Sequential 模型。你知道怎么做吗？我没有模型类定义。

【参考方案2】：

这取决于你想做什么。

案例 1：保存模型以供自己用于推理：您保存模型，恢复它，然后将模型更改为评估模式。这样做是因为您通常有 BatchNorm 和 Dropout 层，它们在构造时默认处于训练模式：

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

#Later to restore:
model.load_state_dict(torch.load(filepath))
model.eval()

案例#2：保存模型以便稍后恢复训练：如果您需要继续训练您将要保存的模型，您需要保存的不仅仅是模型。您还需要保存优化器的状态、时期、分数等。您可以这样做：

state = 
    'epoch': epoch,
    'state_dict': model.state_dict(),
    'optimizer': optimizer.state_dict(),
    ...

torch.save(state, filepath)

要恢复训练，您可以执行以下操作：state = torch.load(filepath)，然后恢复每个对象的状态，如下所示：

model.load_state_dict(state['state_dict'])
optimizer.load_state_dict(state['optimizer'])

由于您正在恢复训练，一旦您在加载时恢复状态，请勿致电model.eval()。

案例 3：其他人无法访问您的代码的模型： 在 Tensorflow 中，您可以创建一个 .pb 文件，该文件定义模型的架构和权重。这非常方便，特别是在使用Tensorflow serve 时。在 Pytorch 中执行此操作的等效方法是：

torch.save(model, filepath)

# Then later:
model = torch.load(filepath)

这种方式仍然不是防弹的，而且由于 pytorch 仍在经历很多变化，我不会推荐它。

【讨论】：

这3个案例有推荐的文件结尾吗？还是总是 .pth？

在案例#3 torch.load 只返回一个 OrderedDict。您如何获得模型以进行预测？

您好，请问上面提到的“案例#2：保存模型以便稍后恢复训练”怎么做？我设法将检查点加载到模型中，然后我无法运行或恢复训练模型，例如“model.to(device) model = train_model_epoch(model, criteria, optimizer, sched, epochs)”

嗨，对于用于推理的情况，在官方 pytorch 文档中说必须保存优化器 state_dict 以进行推理或完成训练。 “保存一般检查点时，用于推理或恢复训练，您必须保存的不仅仅是模型的 state_dict。重要的是还保存优化器的 state_dict，因为它包含随着模型训练而更新的缓冲区和参数. "

情况#3，模型类应该定义在某处。

【参考方案3】：

pickle Python 库实现了用于序列化和反序列化 Python 对象的二进制协议。

当你import torch（或者当你使用PyTorch）它会为你import pickle，你不需要直接调用pickle.dump()和pickle.load()，这是保存和加载对象的方法.

事实上，torch.save() 和 torch.load() 将为您包装 pickle.dump() 和 pickle.load()。

state_dict 提到的另一个答案值得多加注意。

我们在 PyTorch 中有什么 state_dict？ 实际上有两个state_dicts。

PyTorch 模型是 torch.nn.Module，它调用 model.parameters() 来获取可学习的参数（w 和 b）。 这些可学习的参数一旦随机设置，就会随着我们的学习而随着时间的推移而更新。 可学习的参数是第一个state_dict。

第二个state_dict 是优化器状态字典。您还记得优化器用于改进我们的可学习参数。但是优化器state_dict 是固定的。那里没什么可学的。

因为state_dict 对象是 Python 字典，所以它们可以轻松保存、更新、更改和恢复，从而为 PyTorch 模型和优化器增加了大量的模块化。

让我们创建一个超级简单的模型来解释这一点：

import torch
import torch.optim as optim

model = torch.nn.Linear(5, 2)

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

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

print("Model weight:")    
print(model.weight)

print("Model bias:")    
print(model.bias)

print("---")
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:
weight      torch.Size([2, 5])
bias      torch.Size([2])
Model weight:
Parameter containing:
tensor([[ 0.1328,  0.1360,  0.1553, -0.1838, -0.0316],
        [ 0.0479,  0.1760,  0.1712,  0.2244,  0.1408]], requires_grad=True)
Model bias:
Parameter containing:
tensor([ 0.4112, -0.0733], requires_grad=True)
---
Optimizer's state_dict:
state      
param_groups      ['lr': 0.001, 'momentum': 0.9, 'dampening': 0, 'weight_decay': 0, 'nesterov': False, 'params': [140695321443856, 140695321443928]]

请注意，这是一个最小模型。您可以尝试添加顺序堆栈

model = torch.nn.Sequential(
          torch.nn.Linear(D_in, H),
          torch.nn.Conv2d(A, B, C)
          torch.nn.Linear(H, D_out),
        )

请注意，只有具有可学习参数的层（卷积层、线性层等）和已注册的缓冲区（batchnorm 层）在模型的state_dict 中有条目。

不可学习的东西属于优化器对象state_dict，其中包含有关优化器状态的信息，以及使用的超参数。

故事的其余部分是一样的；在推理阶段（这是我们在训练后使用模型的阶段）进行预测；我们确实根据我们学到的参数进行预测。所以为了推理，我们只需要保存参数model.state_dict()即可。

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

供以后使用 model.load_state_dict（torch.load（文件路径）） 模型.eval()

注意：不要忘记最后一行model.eval()，这是加载模型后的关键。

也不要尝试保存torch.save(model.parameters(), filepath)。 model.parameters() 只是生成器对象。

另一方面，torch.save(model, filepath) 保存模型对象本身，但请记住，模型没有优化器的 state_dict。检查@Jadiel de Armas 的另一个出色答案以保存优化器的状态字典。

【讨论】：

虽然不是直截了当的解决方案，但问题的本质分析得很深刻！点赞。

【参考方案4】：

一个常见的 PyTorch 约定是使用 .pt 或 .pth 文件扩展名保存模型。

保存/加载整个模型

保存：

path = "username/directory/lstmmodelgpu.pth"
torch.save(trainer, path)

加载：

（模型类必须在某处定义）

model.load_state_dict(torch.load(PATH))
model.eval()

【讨论】：

它引发了：AttributeError: 'dict' object has no attribute 'eval'

【参考方案5】：

如果您想保存模型并希望稍后恢复训练：

单 GPU： 保存：

state = 
        'epoch': epoch,
        'state_dict': model.state_dict(),
        'optimizer': optimizer.state_dict(),

savepath='checkpoint.t7'
torch.save(state,savepath)

加载：

checkpoint = torch.load('checkpoint.t7')
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch = checkpoint['epoch']

多 GPU： 保存

state = 
        'epoch': epoch,
        'state_dict': model.module.state_dict(),
        'optimizer': optimizer.state_dict(),

savepath='checkpoint.t7'
torch.save(state,savepath)

加载：

checkpoint = torch.load('checkpoint.t7')
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
epoch = checkpoint['epoch']

#Don't call DataParallel before loading the model otherwise you will get an error

model = nn.DataParallel(model) #ignore the line if you want to load on Single GPU

【参考方案6】：

Saving locally

您如何保存模型取决于您希望将来如何访问它。如果您可以调用model 类的新实例，那么您需要做的就是使用model.state_dict() 保存/加载模型的权重：

# Save:
torch.save(old_model.state_dict(), PATH)

# Load:
new_model = TheModelClass(*args, **kwargs)
new_model.load_state_dict(torch.load(PATH))

如果您因任何原因不能（或者更喜欢更简单的语法），那么您可以使用torch.save() 保存整个模型（实际上是对定义模型的文件及其 state_dict 的引用）：

# Save:
torch.save(old_model, PATH)

# Load:
new_model = torch.load(PATH)

但是由于这是对定义模型类的文件位置的引用，因此除非这些文件也被移植到相同的目录结构中，否则这段代码是不可移植的。

保存到云端 - TorchHub

如果您希望您的模型是可移植的，您可以使用torch.hub 轻松地将其导入。如果您将适当定义的 hubconf.py 文件添加到 github 存储库，则可以在 PyTorch 中轻松调用该文件，以使用户能够加载带/不带权重的模型：

hubconf.py (github.com/repo_owner/repo_name)

dependencies = ['torch']
from my_module import mymodel as _mymodel

def mymodel(pretrained=False, **kwargs):
    return _mymodel(pretrained=pretrained, **kwargs)

加载模型：

new_model = torch.hub.load('repo_owner/repo_name', 'mymodel')
new_model_pretrained = torch.hub.load('repo_owner/repo_name', 'mymodel', pretrained=True)

【参考方案7】：

pip install pytorch-lightning

确保您的父模型使用 pl.LightningModule 而不是 nn.Module

Saving and loading checkpoints using pytorch lightning

import pytorch_lightning as pl

model = MyLightningModule(hparams)
trainer.fit(model)
trainer.save_checkpoint("example.ckpt")
new_model = MyModel.load_from_checkpoint(checkpoint_path="example.ckpt")

【参考方案8】：

这几天一切都写在官方教程中： https://pytorch.org/tutorials/beginner/saving_loading_models.html

关于如何保存和保存什么，您有多种选择，所有这些都在该教程中进行了说明。

【参考方案9】：

我总是更喜欢使用 Torch7 (.t7) 或 Pickle (.pth, .pt) 来节省 pytorch 模型的权重。