如何在 PyTorch 中保存经过训练的模型?我读过:
torch.save()/torch.load() 用于保存/加载可序列化的对象。 model.state_dict()/model.load_state_dict() 用于保存/加载模型状态。
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
pickle
有什么问题?
在他们的 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 部分。
这取决于你想做什么。
案例 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 仍在经历很多变化,我不推荐它。
torch.load
中只返回一个 OrderedDict。您如何获得模型以进行预测?
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_dict
。
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 的另一个出色答案以保存优化器的状态字典。
一个常见的 PyTorch 约定是使用 .pt 或 .pth 文件扩展名保存模型。
保存/加载整个模型
节省:
path = "username/directory/lstmmodelgpu.pth"
torch.save(trainer, path)
加载:
(模型类必须在某处定义)
model.load_state_dict(torch.load(PATH))
model.eval()
如果您想保存模型并希望稍后恢复训练:
单 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
本地保存
保存模型的方式取决于将来要如何访问它。如果您可以调用 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)
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")
我总是喜欢使用 Torch7 (.t7) 或 Pickle (.pth, .pt) 来保存 pytorch 模型的权重。
这些天一切都写在官方教程中:https://pytorch.org/tutorials/beginner/saving_loading_models.html
您有几个关于如何保存和保存什么的选项,所有这些都在该教程中进行了解释。