如何在 pytorch 模型中加载检查点文件？答案

【问题标题】：How to load a checkpoint file in a pytorch model?如何在 pytorch 模型中加载检查点文件？
【发布时间】：2019-07-07 17:29:51
【问题描述】：

在我的 pytorch 模型中，我正在像这样初始化我的模型和优化器。

model = MyModelClass(config, shape, x_tr_mean, x_tr,std)
optimizer = optim.SGD(model.parameters(), lr=config.learning_rate)

这是我的检查点文件的路径。

checkpoint_file = os.path.join(config.save_dir, "checkpoint.pth")

为了加载这个检查点文件，我检查检查点文件是否存在，然后我加载它以及模型和优化器。

if os.path.exists(checkpoint_file):
    if config.resume:
        torch.load(checkpoint_file)
        model.load_state_dict(torch.load(checkpoint_file))
        optimizer.load_state_dict(torch.load(checkpoint_file))

另外，这是我保存模型和优化器的方法。

 torch.save({'model': model.state_dict(), 'optimizer': optimizer.state_dict(), 'iter_idx': iter_idx, 'best_va_acc': best_va_acc}, checkpoint_file)

由于某种原因，每当我运行此代码时，我都会收到一个奇怪的错误。

model.load_state_dict(torch.load(checkpoint_file))
File "/home/Josh/.local/lib/python3.6/site-packages/torch/nn/modules/module.py", line 769, in load_state_dict
self.__class__.__name__, "\n\t".join(error_msgs)))
RuntimeError: Error(s) in loading state_dict for MyModelClass:
        Missing key(s) in state_dict: "mean", "std", "attribute.weight", "attribute.bias".
        Unexpected key(s) in state_dict: "model", "optimizer", "iter_idx", "best_va_acc"

有人知道我为什么会收到这个错误吗？

【问题讨论】：

标签： python python-3.x pytorch checkpointing

【解决方案1】：

您将模型参数保存在字典中。您应该使用之前保存时使用的密钥来加载模型检查点和state_dicts，如下所示：

if os.path.exists(checkpoint_file):
    if config.resume:
        checkpoint = torch.load(checkpoint_file)
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])

您可以在 PyTorch 网站上查看官方 tutorial 了解更多信息。

【讨论】：