【PyTorch】tqdmとベンチマークモードの組合せ可否

PyTorchを使っていて,

というエラーが出力されるときがあって,最初は原因がよく分からなかったのですが,ソースを触っていると原因となる箇所が分かったので残しておきます。

原因となる箇所

Pytorchで学習させる部分の関数は以下のように書いていたのですが,このままだと上述したエラーが出てしまいます。

この原因となる箇所が以下の2か所の関係です。

  • 2行目:ベンチマークモードの設定
  • 10行目:tqdmの使用有無

ここの組み合わせによってはエラーが出てしまうようなので調べた結果を整理しておきます。

調査結果

組み合わせの結果は以下の通りです。

benchmark tqdm 実行結果
True 使用 RuntimeError: cannot join current thread
False 使用 正常完了
True 不使用 RuntimeError: CUDA error: out of memory
False 不使用 正常完了

benchmark=True かつ tqdmを使用

この場合は,完全にエラーが出て止まってしまいます。

パッと調べてみたところ,tqdm側の問題かもしれないとか何とか。

benchmark=True かつ tqdmは不使用

今回調査した環境ではGPUのメモリ不足で落ちてしまいました。
メモリはかなりかつかつの状態で検証していたため,何かしらの追加の処理でメモリを食ってしまってメモリ不足になったのだと思います。

ベンチマークモードをオン(torch.backends.cudnn.benchmark = True)にすると,オートチューナーがネットワーク構成に対して最適なアルゴリズムを見つけるため,計算速度が向上するらしいです。
恐らくその過程で,追加でメモリを使用しているんでしょう。

まとめ

以下のような感じでしょうか。

  • ベンチマークモードとtqdmの併用は使用不可
  • ベンチマークモードを使用する場合は,メモリが追加で消費されるため要注意
スポンサーリンク