pytorchで評価時のみOut of memory

trainingは問題なくできるのが、評価の段階になるとメモリー不足で落ちることがあった。

データセットのサイズとしてはメモリ不足になるような大きさではないため、ロジックを調べてみるとtensorを加算する処理をtorch.no_gradなしにやっていることが原因だった。

実際にやっていたことは異なるが以下のような形の処理をtorch.no_grad()をせずにやっていた。

model = Model()
sums = torch.zeros(1, 50)

for _ in range(100):
      time.sleep(0.5)
      data = torch.rand(1, 1, 24000)
      output = model(data)
      sums += output

modelでの処理結果がrequires_grad=Trueとなるため、勾配の計算結果がbackwardに備えて保存されていくため加算する毎にメモリが消費されていた。

以下のgistが試すのに使ったコード。
https://gist.github.com/y-kamiya/4f50185acd040dea4dda655cb24df372

topコマンドで見ながら試しに上記を動かすとわかりやすいが、torch.no_grad()をコメントアウトして実行すると処理が進むにつれて消費メモリが増えていくことがわかる。また、printしているrequires_gradの値もTrueとなる。

torch.no_grad()のブロック内で実行することにより、requires_grad=Falseとなり消費メモリも増えないことが確認できる。

ちなみに検証中に試した際に気になったが、torchvision.modelsに含まれる学習済みモデルを使って同様のことを試すと、no_gradを指定してもrequires_grad=Falseとならなかった。pretrainedモデルを利用した際にハマりそうなのでメモしておく。