GPUで使用したoptimizerをsave & load する時の注意

#ブログ概況

• PyTorch の学習を別ジョブに引き継ぐ場合における注意点として、model の weight だけでなく、optimizer や scheduler なども save & load する必要があります。
• optimizer については特に注意が必要だと感じていて、うっかり前の情報を引き継がずに optimizer を初期化して学習開始してしまった場合、モデルを引き継いだのにも関わらす、急に精度が落ちてしまうこともあります。
• Optimizer の Load & Save について、デバイスを意識していないと、学習時に思わぬエラーを吐いてしまっていたので、その対処法の紹介。

#代表的な２つの最適化アルゴリズム
（一応、Optimizer を Load する必要性の簡単な説明から書いていますが、結論から先に知りたいって方は、このブログのGPU 環境で使用した Optimizer を再度 GPU 環境で学習する際の注意点に対処法を記述しています。）

##SGD
　確率的勾配降下法と呼ばれるものです。計算される勾配情報（損失関数の一階微分）を元に新たな weight に更新されます。式で表すと以下の形です。

\mathbf{w}^{t + 1} \gets \mathbf{w}^{t} - \eta \frac{\partial E(\mathbf{w}^{t})}{\partial \mathbf{w}^{t}}

見ての通り、学習率は勾配や前回の weight 更新量などに関係せず一定です。なので SGD に関しては optimizer の save&load に関して、特に気をつける必要はありません。

##Adam
　 Adam(Adaptive moment estimation)は AdaGrad や RMSProp、AdaDelta を改良したものです。
以下の式のように weight を更新していきます。

m_{t+1} = \beta_{1} m_{t} + (1 - \beta_{1}) \nabla E(\mathbf{w}^{t}) \ \ \ \cdots\ \ \ (1)\\
v_{t+1} = \beta_{2} v_{t} + (1 - \beta_{2}) \nabla E(\mathbf{w}^{t})^{2} \ \ \ \cdots\ \ \ (2)\\
\hat{m} = \frac{m_{t+1}}{1 - \beta_{1}^{t}}\\
\hat{v} = \frac{v_{t+1}}{1 - \beta_{2}^{t}}\\
\mathbf{w}^{t+1} \gets \mathbf{w}^{t} - \alpha \frac{\hat{m}}{\sqrt{\hat{v}} + \epsilon}

初期値は$m_0=0$、$\nu_0=0$となっています。またハイパーパラメータの推奨値は以下の通りです。

\alpha=0.001\\
\beta_{1}=0.9\\
\beta_{2}=0.999\\
\epsilon=10^{−8}\\

(1)式をみて分かる通り、Adam は更新前の勾配が式に組み込まれることで、結果的に勾配の移動平均を取っています。また(2)式をみて分かる通り、勾配の二乗の移動平均も式に組み込まれています。
　結果的に前回の勾配情報を元に学習率を逐次調整しながら traing が進んでいくので、SGD と違い学習率は一定にならないというわけです。

#PyTorch での save
　ここからは実際に、PyTorch での Optimizer のセーブやロードを見ていきます。まずはデモ用に簡単なモデルクラスや optimizer をインスタンス化していきます。尚、下記コードは PyTorch の公式リファレンスを参考に一部追記・削除しています。

import torch.nn as nn
import torch.optim as optim
from torch.optim.lr_scheduler import ReduceLROnPlateau


class SampleModel(nn.Module):
    def __init__(self):
        super(SampleModel, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

model = SampleModel()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
scheduler = ReduceLROnPlateau(optimizer, factor=0.1, patience=5)

• PyTorch の model クラスに存在する学習可能な weight はmodel.parameters()でアクセス可能です。
• 各レイヤーと model のモデルの weight をマッピングするオブジェクトをstate_dict()メソッドで生成可能です。非常に便利！
• 試しに model と optimizer にstate_dict()メソッドを適用してみます。

print('lr: ', optimizer.state_dict()['param_groups'][0]['lr'])
print('betas: ', optimizer.state_dict()['param_groups'][0]['betas'])
print('eps: ', optimizer.state_dict()['param_groups'][0]['eps'])
print('weight_decay: ', optimizer.state_dict()['param_groups'][0]['weight_decay'])
print('amsgrad: ', optimizer.state_dict()['param_groups'][0]['amsgrad'])
print('params: ', optimizer.state_dict()['param_groups'][0]['params'])

print("\nModel's state_dict:")
for param_tensor in model.state_dict():
    print(param_tensor, "\t", model.state_dict()[param_tensor].size())

Out

lr:  0.001
betas:  (0.9, 0.999)
eps:  1e-08
weight_decay:  0
amsgrad:  False
params:  [140713907680744, 140713907680888, 140713907680960, 140713907681032, 140713907681104, 140713907681176, 140713907681248, 140713907681320, 140713907681392, 140713907681464]

Model's state_dict:
conv1.weight 	 torch.Size([6, 3, 5, 5])
conv1.bias 	 torch.Size([6])
conv2.weight 	 torch.Size([16, 6, 5, 5])
conv2.bias 	 torch.Size([16])
fc1.weight 	 torch.Size([120, 400])
fc1.bias 	 torch.Size([120])
fc2.weight 	 torch.Size([84, 120])
fc2.bias 	 torch.Size([84])
fc3.weight 	 torch.Size([10, 84])
fc3.bias 	 torch.Size([10])

これらを保存、ロードすることで学習&推論が可能となっています。保存とロードは以下のように行います。

import torch

# save
torch.save(model.state_dict(), PATH1)
torch.save(optimizer.state_dict(), PATH2)

# initialize
model = SampleModel()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# load
model.load_state_dict(torch.load(PATH1))
optimizer.load_state_dict(torch.load(PATH2))

あるいは下記のようにまとめてシリアル化することも可能です。こちらの方が小分けにする必要がないのでバグに繋がりにくいです。

state = {
    'model': model.state_dict(),
    'optimizer': optimizer.state_dict(),
}

torch.save(state, PATH3)

# initialize
model = SampleModel()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# load
checkpoint = torch.load(PATH3)
model.load_state_dict(checkpoint['model'])
optimizer.load_state_dict(checkpoint['optimizer'])

#GPU 環境で使用した Optimizer を再度 GPU 環境で学習する際の注意点
　というわけでいざ学習を再開しようとすると、学習中にエラーが発生。メッセージを見てみると、どうやらtorch.FloatTensorとtorch.cuda.FloatTensorで整合性が取れてないようです。

ErrorMessage

RuntimeError: expected device cpu but got device cuda:0

Loading a saved model for continue trainingでも議論されているのですが、どうやら継続して optimizer を使用する場合、以下の様に逐一.cuda()をしないといけません。

for state in optimizer.state.values():
    for k, v in state.items():
        if isinstance(v, torch.Tensor):
            state[k] = v.to('cuda')

結構、初期の頃にハマりがちなポイントだと思うので、改めてブログ記事にさせて頂きました。