Question

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

为什么我们不应该加载优化器的状态（亚当等）？

Answer 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()

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

Answer 2

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

在这种情况下，优化器的学习率将被重新调整到保存状态。

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

2 个答案: