在PyTorch中保存训练模型的最佳方法？

Question

我正在寻找在PyTorch中保存训练模型的替代方法.到目前为止,我找到了两种选择.

1. torch.save()保存模型,torch.load()加载模型.
2. model.state_dict()用于保存训练有素的模型,而model.load_state_dict()用于加载已保存的模型.

我已经讨论过这种讨论,其中建议方法2优于方法1.

我的问题是,为什么第二种方法更受欢迎？是否因为torch.nn模块具有这两个功能而我们被鼓励使用它们？

Answer 1

我在他们的github repo上找到了这个页面,我只是在这里粘贴内容.

保存模型的推荐方法

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

第一个(推荐)保存并仅加载模型参数:

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)

但是在这种情况下,序列化数据绑定到特定的类和使用的确切目录结构,因此当在其他项目中使用时,或者在一些严重的重构之后,它可以以各种方式中断.

Answer 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仍然经历了很多变化,我不推荐它.

Answer 3

将泡菜的Python库实现二进制协议的序列化和反序列化Python对象。

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

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

一state_dict提到对方的回答值得只是多了一些注意事项。

什么state_dict我们有内部PyTorch？实际上有两个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(filepath)) model.eval()

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

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

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

Answer 4

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

保存/加载整个模型

节省：

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

加载：

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

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

Answer 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

Answer 6

本地保存

保存模型的方式取决于您将来想要如何访问它。如果您可以调用该类的新实例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))

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

# 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）

# Save:
torch.save(old_model, PATH)

# Load:
new_model = torch.load(PATH)

加载模型：

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

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