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

PyTorchを使っていて，

RuntimeError: cannot join current thread

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

原因となる箇所

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

import torch
torch.backends.cudnn.benchmark = True

# 学習用関数
def train(train_loader, model_obj, optimizer, loss_fn, device, total_epoch, epoch):
    
    model_obj.train() # モデルを学習モードに変更

    # ミニバッチごとに学習
    for mini_batch in tqdm(train_loader):

        # データをDoubleTensorからFloatTensorに変換した上で，to()で明示的にGPUを使用する
        images = mini_batch['image'].to(device, dtype=torch.float)
        labels = mini_batch['labels'].to(device, dtype=torch.float)

        optimizer.zero_grad() # 勾配を初期化
        outputs = model_obj(images) # 順伝播の計算
        loss = loss_fn(outputs, labels) # 誤差を計算

        loss.backward() # 誤差を逆伝播させる
        optimizer.step() # 重みを更新する
    
    print ('Epoch [%d/%d], Loss: %.4f' % (epoch, total_epoch, loss.item()))

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

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

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

調査結果

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

benchmark=True かつ tqdmを使用

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

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

benchmark＝True かつ tqdmは不使用

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

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

まとめ

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

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