在PyTorch中保存经过训练的模型的最佳方法?
deep-learning
python
pytorch
tensor
11
0

我一直在寻找其他方法来在PyTorch中保存经过训练的模型。到目前为止,我发现了两种选择。

  1. 使用torch.save()保存模型,使用torch.load()加载模型。
  2. model.state_dict()保存训练的模型, model.load_state_dict()加载保存的模型。

我碰到过这种讨论 ,其中建议方法2优于方法1。

我的问题是,为什么选择第二种方法呢?仅仅是因为torch.nn模块具有这两个功能,我们被鼓励使用它们吗?

参考资料:
Stack Overflow
收藏
评论
共 4 个回答
高赞 时间 活跃

这取决于您想做什么。

案例1:保存模型以供您自己进行推断 :保存模型,还原模型,然后将模型更改为评估模式。这样做是因为您通常具有BatchNormDropout层,默认情况下它们在构造中处于训练模式:

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仍在进行大量更改,因此我不建议这样做。

收藏
评论

pickle 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

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()只是生成器对象。

另一方面, torch.save(model, filepath)保存模型对象本身,但是请记住,模型没有优化器的state_dict 。检查@Jadiel de Armas的其他出色答案,以保存优化程序的状态字典。

收藏
评论

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

保存/加载整个模型保存:

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

加载:

模型类必须在某处定义

model = torch.load(PATH)
model.eval()
收藏
评论

我在他们的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)

但是,在这种情况下,序列化的数据将绑定到所使用的特定类和确切的目录结构,因此在其他项目中使用或经过一些严重的重构后,它可能以各种方式中断。

收藏
评论
新手导航
  • 社区规范
  • 提出问题
  • 进行投票
  • 个人资料
  • 优化问题
  • 回答问题

关于我们

常见问题

内容许可

联系我们

@2020 AskGo
京ICP备20001863号