Pytorchで学習再開

Pytorchでモデルをsave/loadして学習を途中から続ける方法です

学習の途中で、（コンピューターや人間の都合で）一旦中断して、後で学習を再開したい時があります。
Colabのような連続使用時間に制限があるような環境や、学習エポック数の初期設定に追加して学習したいときなどに、モデルの重みをファイルとして保存して、後でロードして再開したい。

モデルを保存するだけでは、同じ正解率から学習再開できない

モデルの保存と読み込みは以下の通りですが、これは、復元したモデルをevalモードで単に推論に使うときは機能しますが、学習を再開しようとすると、ロスや正解率が保存前から継続せず、初期値に戻ってしまっていることに気づきます。

# 保存
torch.save(model.state_dict(), PATH)

# 読み込み
model = TheModelClass(*args, **kwargs)
model.load_state_dict(torch.load(PATH))
model.eval()

Optimizerも保存すべし

学習を継続して再開するには、モデルの重みの状態に加えて、Optimizerの状態も保存／ロードする必要があります。

# 保存
save_path = "my_model_training_state.pt"
torch.save({'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': loss,},
           save_path)

#　読み込み
model = TheModelClass(*args, **kwargs)
optimizer = optim.SGD(model_ft.parameters(), lr=0.001, momentum=0.9)

PATH = "my_model_training_state.pt"
checkpoint = torch.load(PATH)
model.load_state_dict(checkpoint['model_state_dict'])
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

# optimizerのstateを現在のdeviceに移す。これをしないと、保存前後でdeviceの不整合が起こる可能性がある。
for state in optimizer.state.values():
    for k, v in state.items():
        if isinstance(v, torch.Tensor):
            state[k] = v.to(device)
epoch = checkpoint['epoch']
loss = checkpoint['loss']

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

model = model.to(device)
criterion = nn.CrossEntropyLoss()
exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=7, gamma=0.1)

これで、保存前のロスと正解率から再開できる。

🐣

---

フリーランスエンジニアです。
お仕事のご相談こちらまで
[e-mail]

Core MLやARKitを使ったアプリを作っています。
機械学習／AR関連の情報を発信しています。

Twitter
Medium
GitHub