グレートなtqdmの使い方

本記事は Python3.6 以上を想定しています．

Pythonで何かしら時間のかかる処理をする際にプログレスバーを表示するのに便利なライブラリとして tqdm というものが存在します．

tqdm/tqdm: A Fast, Extensible Progress Bar for Python and CLI

これの良さ気な使い方を紹介します．

導入

pip install tqdm

とりあえず使ってみる

from tqdm import tqdm

for i in tqdm(range(1000000)):
    pass

100%|████████████████████████████| 1000000/1000000 [00:00<00:00, 5139706.81it/s]

恐らくすぐ終わるでしょう．

rangeでなくても イテレータ に対してなら使用可能です(それはそうですね)．

from time import sleep

lines = ["a", "b", "c", "d", "e"]    # "abcde"でも動作します
for line in tqdm(lines):
    sleep(1)

100%|████████████████████████████| 5/5 [00:05<00:00,  1.00s/it]

また enumerateとの併用も可能です．
その場合，外側に使用します．

for i, line in enumerate(tqdm(lines)):
    sleep(1)

ジェネレータに使う場合

ジェネレータに対しても tqdm は使えます．

"""
適当なジェネレータを作成
"""
def gen():
    text = "abcde"
    for ch in text:
        yield ch

for ch in tqdm(gen()):
    sleep(1)

5it [00:05,  1.00s/it]

ただしこのような表示になってしまい，肝心のプログレスバーが出ていませんね．
ジェネレータは __len__ を持たないため，このような結果になります．

ジェネレータといえど，サイズがあらかじめわかる場合はそこそこあると思います．
例えばデータセットからミニバッチを作成する際に (データセットのサイズ) // (バッチサイズ) から計算できますね．

このような場合は以下のようにします．

"""
適当なジェネレータを作成
"""
text = "abcde"

def gen():
    for ch in text:
        yield ch

for ch in tqdm(gen(), total=len(text)):
    sleep(1)

100%|████████████████████████████| 5/5 [00:05<00:00,  1.00s/it]

totalに対してサイズを指定することでイテレータの場合と同じような見た目になります．

別の情報をprintしたい場合

プログレスバーの前後に情報を付与することも可能です．

prefix

機械学習をする場合，通常はtrainとvalidで分けて誤差を計算したりすることが多いと思います．
その際に今どちらのループが回っているのかを知ることで精神的な安心を得られます．

for ch in tqdm(text, desc="[train]"):
    sleep(0.1)

このように desc を指定することでprefixを設定できます．

ループの中で動的に変更したい場合は以下のようにします．
今回は train/valid の話をしておきながら説明の都合上 Epoch を表示しています．

with tqdm(text) as pbar:
    for i, ch in enumerate(pbar):
        pbar.set_description("[Epoch %d]" % (i + 1))
        sleep(0.1)

[Epoch 5]: 100%|████████████████████████████| 5/5 [00:00<00:00,  9.98it/s]

いい感じですね．

postfix(suffix)

ちなみに postfix というのは正しい英語ではないらしいです．
正しくはsuffixといいます．(接尾語，ですね)

さて，本題です．
プログレスバーの後ろに誤差とかAccuracyとかを表示したい場合もありますよね．
その場合以下のようにします．

for ch in tqdm(text, postfix="loss=1.0"):
    sleep(0.1)

100%|████████████████████████████| 5/5 [00:00<00:00,  9.72it/s, loss=1.0]

まあこういう使い方はしませんね．
おそらく後ろに情報を付けたい場合，なにか値が推移するものに対して使うと思います．

というわけで動的に変動させたい場合は以下のようにします．

from collections import OrderedDict

with tqdm(text) as pbar:
    for i, ch in enumerate(pbar):
        pbar.set_postfix(OrderedDict(loss=1-i/5, acc=i/10))
        sleep(0.1)

100%|████████████████████████████| 5/5 [00:00<00:00,  9.71it/s, loss=0.2, acc=0.4]

OrderedDict を使用することで順序をキープしたまま表示させることができます．

なお，Python3.6以降は dictが標準で順序をキープするようになったそうです．
が，どうやら実装によって違うとのことなので順序をキープしたいという強い意志がある方は OrderedDictを使用すると良いでしょう．

自分の環境では pbar.set_postfix({"loss": 1-i/5, "acc": i/10}) と pbar.set_postfix({"acc": i/10, "loss": 1-i/5}) で結果が変わり，順序が確かにキープされていました．

OrderedDictを使用するには import する必要があるので自分ならめんどくさくてそのまま dict を使うでしょう．

最終的に

with tqdm(text) as pbar:
    for i, ch in enumerate(pbar):
        pbar.set_description("[train] Epoch %d" % i)
        pbar.set_postfix(OrderedDict(loss=1-i/5, acc=i/10))
        sleep(0.1)

全部入りさせると．．．．．．

[train] Epoch 4: 100%|████████████████████████████| 5/5 [00:00<00:00,  9.71it/s, loss=0.2, acc=0.4]

おお！素晴らしい！
これで安心したPythonライフを送れます．

他にもプログレスバーの幅を指定したり，更新頻度を指定したり，手動でプログレスバーを動かしたりすることもできますが，個人的にあまり使わないので割愛します．