在他们的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教程中的保存和加载模型部分。

保存在本地

如何保存模型取决于将来如何访问它。如果你可以调用模型类的一个新实例，那么你所需要做的就是用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导入它。如果你在github repo中添加了一个适当定义的hubconf.py文件，这可以很容易地在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)

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

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

我总是喜欢使用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')

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保存优化器的状态字典。