pytorch 保存和加载 Checkpoint 模型，实现断点训练

保存和加载 Checkpoint 用于推理/继续训练

保存

# 模型类必须在此之前被定义
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()

 

当保存成 Checkpoint 的时候，可用于推理或者是继续训练，保存的不仅仅是模型的 state_dict。

保存优化器的 state_dict 也很重要, 因为它包含作为模型训练更新的缓冲区和参数。你也许
想保存其他项目，比如最新记录的训练损失，外部的 torch.nn.Embedding 层等等。

要保存多个组件，请在字典中组织它们并使用 torch.save() 来序列化字典。

PyTorch 中常见的保存checkpoint 是使用 .tar 文件扩展名。

要加载项目，首先需要初始化模型和优化器，然后使用 torch.load() 来加载本地字典。

这里,你可以非常容易的通过简单查询字典来访问你所保存的项目。

请记住在运行推理之前，务必调用 model.eval() 去设置 dropout 和 batch normalization 为评估。
如果不这样做，有可能得到不一致的推断结果。 如果你想要恢复训练，请调用 model.train() 以
确保这些层处于训练模式。