pytochのメモリー漸増問題

最近モデルが大きくなってきて、GPUのメモリーぎりぎりなっています。
学習途中で使用メモリが増えて落ちるケースに困っていましたが、その対策が終わったので整理しておきます。

現象：
　交差検証のstep（1回の試行）毎に使用メモリーが増加し、途中でメモリ不足が発生する。

torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 66.00 MiB (GPU 0; 23.69 GiB total capacity; 22.67 GiB already allocated; 63.88 MiB free; 23.29 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

1-step目ではメモリ足りていたので安心していたのに、ｘｘ時間後にn-step目でメモリ不足で落ちていると、とても残念な気分になってしまいますよね

★結論：gc.collect()最高！

以下は対策毎のメモリ使用量の推移です。（交差は5回）

・case0：初期状態のメモリ推移

*on init gpu memory =0.2
ep=0.0, los=0.6705(-9.3295), auc=0.554(+0.554), val_los=0.6592(-9.3408), val_auc=0.554(+0.554) 0.0 hours
*** gpu mem= 19.2  on train
*** gpu mem =19.2  on step end
ep=1.0, los=0.6716(-9.3284), auc=0.557(+0.557), val_los=0.6529(-9.3471), val_auc=0.556(+0.556) 0.0 hours
*** gpu mem= 20.9  on train
*** gpu mem =20.9  on step end
ep=2.0, los=0.6716(-9.3284), auc=0.546(+0.546), val_los=0.6581(-9.3419), val_auc=0.553(+0.553) 0.0 hours
*** gpu mem= 23.7  on train
*** gpu mem =23.7  on step end
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 66.00 MiB (GPU 0; 23.69 GiB total capacity; 22.67 GiB already allocated; 59.88 MiB free; 23.29 GiB reserved in total by PyTorch)

注）ep={steps}.{epoch}　
　　on train：　　学習中（1epoch目）のメモリ使用量
　　on step end：学習ループが終わり交差検証のstepの最後

stepが進む毎にメモリ使用量が増えて、4step目でエラーになっています。

・case1：model、optimizerを削除した場合

*on init gpu memory =0.2
ep=0.0, los=0.6710(-9.3290), auc=0.556(+0.556)
*** gpu mem= 19.2  on train
*** gpu mem =19.2  on step end
ep=1.0, los=0.6713(-9.3287), auc=0.547(+0.547)
*** gpu mem= 20.9  on train
*** gpu mem =20.9  on step end
ep=2.0, los=0.6722(-9.3278), auc=0.553(+0.553)
*** gpu mem= 23.7  on train
*** gpu mem =23.7  on step end
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 66.00 MiB (GPU 0; 23.69 GiB total capacity; 22.67 GiB already allocated; 59.88 MiB free; 23.29 GiB reserved in total by PyTorch)

まず考え付く”オブジェクトの削除”をしてみました。
全く効果でていませんね・・・

・case2：case1＋torch.cuda.empty_cache()した場合

*on init gpu memory =0.2
ep=0.0, los=0.6704(-9.3296), auc=0.556(+0.556)
*** gpu mem= 19.2  on train
*** gpu mem =10.2  on step end
ep=1.0, los=0.6711(-9.3289), auc=0.552(+0.552)
*** gpu mem= 22.2  on train
*** gpu mem =14.5  on step end
torch.cuda.OutOfMemoryError: CUDA out of memory. Tried to allocate 66.00 MiB (GPU 0; 23.69 GiB total capacity; 22.67 GiB already allocated; 57.88 MiB free; 23.30 GiB reserved in total by PyTorch)

train終了後のstep endではempty_cache()の効果で使用メモリが減っています。
が、相変わらずstep endでのメモリ使用量が漸増しています。
また肝心のon trainでのメモリ使用量は変わらず漸増しています。
というか、むしろ増えています！
おまけにメモリ不足時の要求メモリが
256MB -> 1.75GB
に増えています！
これはcacheクリアした”つもり”になって、次のメモリ確保に行ってエラーになったと思われます。（pytorchの管理からは外れたが、未だにメモリを占拠している状態）

・case3：case2＋gc.collect()の場合

*on init gpu memory =0.2
ep=0.0, los=0.6731(-9.3269), auc=0.554(+0.554)
*** gpu mem= 19.2  on train
*** gpu mem =10.2  on step end
ep=1.0, los=0.6708(-9.3292), auc=0.554(+0.554)
*** gpu mem= 15.2  on train
*** gpu mem =9.2  on step end
ep=2.0, los=0.6719(-9.3281), auc=0.545(+0.545)
*** gpu mem= 15.3  on train
*** gpu mem =9.9  on step end
ep=3.0, los=0.6721(-9.3279), auc=0.553(+0.553)
*** gpu mem= 15.3  on train
*** gpu mem =9.2  on step end
ep=4.0, los=0.6702(-9.3298), auc=0.558(+0.558)
*** gpu mem= 15.3  on train
*** gpu mem =9.2  on step end

gc.collect()を追加してみました。（pythonのgarbage collectionを強制的に発動させる）
なんでしょう・・・、2step目が1step目より少なくなってしまいましたｗ
プログラム構造にまだ見直せる箇所があるのか？原因は不明。
2step目以降を見ると、漸増する傾向は”ほぼ”無くなったと言えると思います。
（まだ削除の取りこぼしがあるのかもしれませんね）

case4：model、optimizer削除＋gc.collect()（cuda_emputy_cache()無し）の場合

*on init gpu memory =0.2
ep=0.0, los=0.6725(-9.3275), auc=0.552(+0.552)
*** gpu mem= 19.2  on train
*** gpu mem =19.2  on step end
ep=1.0, los=0.6717(-9.3283), auc=0.550(+0.550)
*** gpu mem= 20.9  on train
*** gpu mem =20.9  on step end
ep=2.0, los=0.6720(-9.3280), auc=0.552(+0.552)
*** gpu mem= 20.9  on train
*** gpu mem =20.9  on step end
ep=3.0, los=0.6729(-9.3271), auc=0.542(+0.542)
*** gpu mem= 20.9  on train
*** gpu mem =20.9  on step end
ep=4.0, los=0.6703(-9.3297), auc=0.555(+0.555)
*** gpu mem= 20.9  on train
*** gpu mem =20.9  on step end

2step目は増えていますが、それ以降は増えていません。
とりあえず目的達成です！

case5：gc.collect()のみ（削除なし）

*on init gpu memory =0.2
ep=0.0, los=0.6710(-9.3290), auc=0.553(+0.553)
*** gpu mem= 19.2  on train
*** gpu mem =19.2  on step end
ep=1.0, los=0.6709(-9.3291), auc=0.558(+0.558)
*** gpu mem= 20.9  on train
*** gpu mem =20.9  on step end
ep=2.0, los=0.6703(-9.3297), auc=0.555(+0.555)
*** gpu mem= 20.9  on train
*** gpu mem =20.9  on step end
ep=3.0, los=0.6725(-9.3275), auc=0.548(+0.548)
*** gpu mem= 20.9  on train
*** gpu mem =20.9  on step end
ep=4.0, los=0.6716(-9.3284), auc=0.553(+0.553)
*** gpu mem= 20.9  on train
*** gpu mem =20.9  on step end

case4と比較すると衝撃の事実が判明します・・・

（modelとoptimizerの）削除関係ないじゃん！！

結論：GPUメモリ削減に重要なのは各要素のdelではなくgc.collect()である！！

かもしれない。

case3とcase4を見るとempty_cache()ありの方が2step目以降のメモリ使用量が少ないことがわかります。
ですが、1step目のメモリ使用量が同じなのでメモリ削減としてはあまり効果がありません。
（平行して別プログラムを走らせるなら効果がありますが）

また、case2を見ると、empty_cache()単独での使用は逆効果だと思われます。
（gcのタイミング次第）

おまけ

2step目に増加していた原因は　scheduler　でした。

case_ex1：Modle+scheduler削除＋gc.collect()

*on init gpu memory =0.2
ep=0.0, los=0.6701(-9.3299), auc=0.554(+0.554)
*** gpu mem= 19.2  on train
*** gpu mem =19.2  on step end
ep=1.0, los=0.6714(-9.3286), auc=0.554(+0.554)
*** gpu mem= 19.2  on train
*** gpu mem =19.2  on step end
ep=2.0, los=0.6703(-9.3297), auc=0.556(+0.556)
*** gpu mem= 19.2  on train
*** gpu mem =19.2  on step end
ep=3.0, los=0.6720(-9.3280), auc=0.550(+0.550)
*** gpu mem= 19.2  on train
*** gpu mem =19.2  on step end
ep=4.0, los=0.6718(-9.3282), auc=0.552(+0.552)
*** gpu mem= 19.2  on train
*** gpu mem =19.2  on step end

漸増傾向は完全に消えました！（小数点2位以下は見えないことにしています）

これで完全に目的達成です！＼(^_^)／

case_ex2：追加scheduler削除＋gc.collect()+empty_cache()

*on init gpu memory =0.2
ep=0.0, los=0.6713(-9.3287), auc=0.554(+0.554)
*** gpu mem= 19.2  on train
*** gpu mem =10.2  on step end
ep=1.0, los=0.6718(-9.3282), auc=0.551(+0.551)
*** gpu mem= 17.9  on train
*** gpu mem =10.1  on step end
ep=2.0, los=0.6721(-9.3279), auc=0.549(+0.549)
*** gpu mem= 17.9  on train
*** gpu mem =10.2  on step end
ep=3.0, los=0.6730(-9.3270), auc=0.556(+0.556)
*** gpu mem= 17.9  on train
*** gpu mem =10.2  on step end
ep=4.0, los=0.6718(-9.3282), auc=0.553(+0.553)
*** gpu mem= 17.9  on train
*** gpu mem =10.2  on step end

すっきりしました。

上で削除は効果ない風に書きましたが、削除が効果ある（うまく使いまわせていない？）オブジェクトもあるようです。
上の例では纏めていますが、正確に書くと
modelの削除　->　効果あり
optimizerの削除　->　効果なし
schedulerの削除　->　効果あり
dataloaderの削除　->　効果なし
になります。
（lossはtrain()内で削除しているので、今回は触れていません）

確実なのは全部セットでの採用でしょうか

gc.collect()　<–　これが一番大事
使用済みobjectの削除　<–　効果はあり
empty_cache()　<–　単独での使用は逆効果の場合あり

これって、java、python等でgc詳しくないと通常思いつかないですよね・・・
昔javaやってたけど、存在すっかり忘れてましたｗ

最近はメモリ24GBだと足りません。
しかし48GBに増やそうとすると、A6000旧世代で60万円、新世代だと130万円になってしまい、購入に躊躇してしまいます。

次はDeepDpeedに手を出しそうです。