tqdmでメモリリークにハマった話（機械学習）

事件

PytorchのResNetモデルを使って画像の訓練をしていたら、エポックのループ毎に、CPU / GPUメモリ両方の使用率がガンガン上がっていく事件に遭遇した。

解決方法（tqdmの使い方）

tqdmは、進捗をプログレスバーで表示してくれる便利なライブラリだが、そこに落とし穴があった。

import tqdm

# 訓練画像の data loader をループを回す
for x in tqdm(loader):
    y = model(x)

このような書き方をすると、tqdmがloaderを参照したままになり、メモリが解放されない。

正解はこれ。

tqdm(loader)             # ダメ
enumerate(tqdm(loader))  # ダメ
tqdm(enumerate(loader))  # OK

enumerateでラッパーしてからtqdmを噛ませるのが正解。

更にイテレータの長さを渡して、

tqdm(enumerate(loader), total=len(loader))

とすれば、完璧。

こちらの記事が参考になった。
https://github.com/tqdm/tqdm/issues/746

束縛されている変数を発見しよう

参照が残ったままで、メモリから解放されない変数がないかと、疑問を持ったときは、

sys.getrefcount()

を試して見よう。

import sys

a = "aiueo"

# getrefcount 自体も a を参照するため、
# 参照数は常に+1で表示されることに注意。

print(sys.getrefcount(a))
>> 2 

b = a
print(sys.getrefcount(a))
>> 3

del b
print(sys.getrefcount(a))
>> 2

このように、参照されている変数が存在しない場合は「2」、参照が残っている場合は「2」より大きい値を返す。

参照を残さずに関数などのスコープを抜けると、Pythonはガベージコレクションの機能で、自動的にオブジェクト自体もメモリから掃除してくれる。

メモリリークを疑ったときは参照数カウントを試してみては。