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() 以
確保這些層處於訓練模式。