关于加载模型后恢复训练的问题答案

【问题标题】：Question on restoring training after loading model关于加载模型后恢复训练的问题
【发布时间】：2020-07-25 13:47:55
【问题描述】：

经过 24 小时的训练，训练过程通过torch.save 保存了模型文件。关闭电源或其他问题导致进程退出。通常，我们可以加载模型并从最后一步开始继续训练。

为什么我们不应该加载优化器（Adam 等）的状态，有必要吗？

【问题讨论】：

【解决方案1】：

是的，您可以从最后一步加载模型并从该步骤重新训练。

如果您只想将其用于推理，则将模型的 state_dict 保存为

torch.save(model, PATH)

并将其加载为

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

但是，出于您的考虑，您还需要保存优化器状态字典。为此，您需要将其另存为

torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,
            ...
            }, PATH)

并将模型加载为进一步训练：

model = TheModelClass(*args, **kwargs)
optimizer = TheOptimizerClass(*args, **kwargs)

checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
epoch = checkpoint['epoch']
loss = checkpoint['loss']

model.eval()
# - or -
model.train()

有必要保存优化器状态字典，因为它包含随着模型训练而更新的缓冲区和参数。

【讨论】：

【解决方案2】：

在某些情况下需要加载优化器的状态，例如正在使用学习率调度器的情况。

在这种特殊情况下，优化器的学习率将重新调整到它处于保存状态的点。

【讨论】：