负荷模型后恢复训练的问题

Question

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

为什么我们不应该加载优化器(Adam等)的状态，这是必要的吗？

Answer 1

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

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

torch.save(model, PATH)

并将其载入

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

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

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

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

在这种情况下，优化器的学习速度将被重新调整到处于保存状态的位置。