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

单一的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

PIP安装火炬闪电

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

使用pytorch闪电保存和加载检查点

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")

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

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')

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

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

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

上面提到的另一个答案值得再多加注释。

PyTorch中有什么state_dict ? 实际上有两个state_dicts。

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

第二个state_dict是优化器状态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()只是生成器对象。

另一方面，火炬。Save (model, filepath)保存模型对象本身，但请记住，模型没有优化器的state_dict。检查其他优秀的回答@Jadiel de Armas保存优化器的状态字典。

这取决于你想做什么。

案例# 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()。

案例# 3:模型被其他人使用，而不能访问您的代码: 在Tensorflow中，你可以创建一个.pb文件来定义模型的架构和权重。这非常方便，特别是在使用Tensorflow服务时。在Pytorch中这样做的等效方法是:

torch.save(model, filepath)

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

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

在他们的github回购上找到了这个页面:

Recommended approach for saving a model There are two main approaches for serializing and restoring a model. The first (recommended) saves and loads only the model parameters: torch.save(the_model.state_dict(), PATH) Then later: the_model = TheModelClass(*args, **kwargs) the_model.load_state_dict(torch.load(PATH)) The second saves and loads the entire model: torch.save(the_model, PATH) Then later: the_model = torch.load(PATH) However in this case, the serialized data is bound to the specific classes and the exact directory structure used, so it can break in various ways when used in other projects, or after some serious refactors.

请参见官方PyTorch教程中的保存和加载模型部分。