CUBE SUGAR CONTAINER

技術系のこと書きます。

Python: tqdm で処理の進捗状況をプログレスバーとして表示する

最近は Python がデータ分析や機械学習の分野でも使われるようになってきた。 その影響もあって REPL や Jupyter Notebook 上でインタラクティブに作業することも増えたように感じる。 そんなとき、重い処理を走らせると一体いつ終わるのか分からず途方に暮れることもある。 今回紹介する tqdm は、走らせた処理の進捗状況をプログレスバーとして表示するためのパッケージ。 このパッケージ自体はかなり昔からあるんだけど、前述した通り利用環境の変化や連携するパッケージの増加によって便利さが増してきてる感じ。

使った環境は次の通り。

$ sw_vers
ProductName:    Mac OS X
ProductVersion: 10.13.5
BuildVersion:   17F77
$ python -V    
Python 3.6.5

もくじ

下準備

まずは tqdm をインストールしておく。

$ pip install tqdm
$ pip list --format=columns | grep tqdm 
tqdm       4.23.4

終わったら Python の REPL を起動する。

$ python

基本的な使い方

ここからは tqdm の基本的な使い方を紹介する。

その前に、まずは tqdm がない場合から考えてみる。 次のサンプルコードでは 100 回のループを 100 ミリ秒の間隔を空けて実行している。 実行すると、何もレスポンスがない状態で 10 秒間待たされることになる。

>>> import time
>>> 
>>> for _ in range(100):
...     time.sleep(0.1)
... 

実際に実行してみると 10 秒間とはいえ長く感じる。

それでは、続いて上記の処理に tqdm を導入してみる。 変更点は一箇所だけで、上記の range() 関数の結果を tqdm() 関数に渡すだけ。 これだけで tqdm は渡された内容を読み取って全体の処理と現在の進捗をプログレスバーとして表示してくれる。

>>> from tqdm import tqdm
>>> 
>>> for _ in tqdm(range(100)):
...     time.sleep(0.1)
... 
 63%|█████████████████████████████▌                 | 63/100 [00:06<00:03,  9.69it/s]

プログレスバーがあるだけで同じ待ち時間でも感じ方はだいぶ変わるはず。

上記を見ると、どういう仕組みなのか結構気になる。 そもそも tqdm に渡せるオブジェクトは一体なんなのか? 答えから言ってしまうと tqdm にはイテラブルなオブジェクトなら何でも渡せる。 イテラブルなオブジェクトというのは、具体的には iter() 関数を使ってイテレータが返ってくるもの。

そもそもイテレータって何?っていう話については以下の記事に書いた。

blog.amedama.jp

なので、もちろんリストを渡すこともできるし。

>>> for _ in tqdm(list([1, 2, 3, 4])):
...     time.sleep(1)
... 
100%|██████████████████████████████████████████████████| 4/4 [00:04<00:00,  1.00s/it]

何なら文字列だって渡すことができる。

>>> for _ in tqdm('Hello, World!'):
...     time.sleep(0.1)
... 
100%|████████████████████████████████████████████████| 13/13 [00:01<00:00,  9.73it/s]

イテレータ自体には終わりを設ける必要はない。 なので、次のように無限に値を返し続けるイテレータを渡しても良い。 ただし、この場合はイテレータからオブジェクトを取り出した回数や、経過時間やスループットだけが表示される。

>>> from itertools import count
>>> 
>>> for _ in tqdm(count()):
...     time.sleep(0.01)
... 
417it [00:04, 85.10it/s]

ひとしきり満足したら Ctrl-C で止めよう。

pandas と連携させる

tqdm は pandas と連携させることもできる。

まずは pandas をインストールしよう。

$ pip install pandas
$ pip list --format=columns | grep pandas
pandas            0.23.3 

サンプルとなる DataFrame オブジェクトを用意しておく。

>>> import pandas as pd
>>> df = pd.DataFrame(list(range(10000)))

この状態では DataFrame には pregress_apply() というメソッドは存在しない。

>>> df.progress_apply
Traceback (most recent call last):
...(省略)...
AttributeError: 'DataFrame' object has no attribute 'progress_apply'

そこで、おもむろに tqdm をインポートしたら pandas() 関数を呼び出してみよう。

>>> from tqdm import tqdm
>>> tqdm.pandas()

すると DataFrame オブジェクトに progress_apply() メソッドが生えてきて使えるようになる。 これは単純に DataFrame#apply() メソッドの進捗表示ありバージョンと考えれば良い。

>>> df.progress_apply(lambda x: x ** 2, axis=1)
 96%|███████████████████████████████████████▎ | 9577/10000 [00:01<00:00, 5944.60it/s]

DataFrame#apply() 関数は結構重い処理をすることも多い (特に axis=1 のとき) ので、これは意外とありがたい。 ただし、それ以外のメソッドについてはこれまで通り何も表示されない。

Jupyter Notebook と連携させる

また、Jupyter Notebook と連携させることもできる。

まずは Jupyter Notebook 本体と ipyqidgets をインストールしておこう。

$ pip install notebook ipywidgets
$ pip list --format=columns | grep notebook
notebook            5.6.0

ノートブックのサーバを起動する。

$ jupyter notebook

適当なノートブックを新たに作ったら、次のコードをセルに入力して実行してみよう。 ターミナルとの違いはインポートするものが tqdm.tqdm から tqdm.tqdm_notebook に変わるだけ。

from tqdm import tqdm_notebook as tqdm
import time

for _ in tqdm(range(100)):
    time.sleep(0.1)

すると、次のようにプログレスバーが表示される。

f:id:momijiame:20180721130000p:plain

ちなみに、別に普通の tqdm.tqdm が使えないというわけではない。 試しに、最初に示した例を入力して実行してみよう。

from tqdm import tqdm
import time

for _ in tqdm(range(100)):
    time.sleep(0.1)

上記ほどしっかりとした表示ではないものの、次のようにちゃんと表示してくれる。

f:id:momijiame:20180721130420p:plain

めでたしめでたし。