如何在PyTorch中保存训练好的模型?

发现这个页面在他们的github回购:

保存模型的推荐方法

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

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

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官方教程中的保存并加载模型部分。

这取决于你想做什么。

案例# 1:保存模型以便自己使用它进行推理:你保存模型，你恢复它，然后你改变模型到评估模式。这样做是因为你通常有BatchNorm和Dropout层，它们在构造时默认处于火车模式:

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


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

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

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()。

torch.save(model, filepath)


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

这种方式仍然不是防弹的，因为pytorch仍然在经历很多变化，我不推荐它。

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

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

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

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 () < / p >

注意:不要忘记最后一行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: 拯救策略:< / p >

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: 保存< / p >

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

本地保存 .

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


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

# Save:
torch.save(old_model, PATH)


# Load:
new_model = torch.load(PATH)

保存到云- TorchHub

如果你希望你的模型是可移植的，你可以很容易地使用torch.hub导入它。如果你添加一个适当定义的hubconf.py文件到github repo，这可以很容易地从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安装火炬闪电

确保你的父模型使用pl.LightningModule而不是nn。模块

使用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模型的权重。

我用这个方法，希望对大家有用。

num_labels = len(test_label_cols)
robertaclassificationtrain = '/dbfs/FileStore/tables/PM/TC/roberta_model'
robertaclassificationpath = "/dbfs/FileStore/tables/PM/TC/ROBERTACLASSIFICATION"


model = RobertaForSequenceClassification.from_pretrained(robertaclassificationpath,
num_labels=num_labels)
model.cuda()


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

我保存我的火车模型已经在“roberta_model”路径。保存一个火车模型。

torch.save(model.state_dict(), '/dbfs/FileStore/tables/PM/TC/roberta_model')