如何在PyTorch中保存训练好的模型?我读到过:

Torch.save ()/torch.load()用于保存/加载可序列化对象。 model.state_dict()/model.load_state_dict()用于保存/加载模型状态。


当前回答

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

其他回答

保存在本地

如何保存模型取决于将来如何访问它。如果你可以调用模型类的一个新实例,那么你所需要做的就是用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)

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

这取决于你想做什么。

案例# 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仍然在经历很多变化,我不推荐它。

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

保存/加载整个模型

拯救策略:

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

负载:

(模型类必须在某处定义)

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

现在所有内容都写在官方教程中: https://pytorch.org/tutorials/beginner/saving_loading_models.html

关于如何保存和保存什么,您有几个选项,所有这些都在本教程中进行了解释。