Pythonのmultiprocessingの進捗をtqdmで取りたい

概要

Pythonのmultiprocessing.Poolで並列処理をする際にtqdmで進捗を表示する方法の確認。

バージョン情報

• Python 3.7.1
• tqdm==4.29.1

導入

pip install tqdm

サンプルコード

Pool.imap、もしくはimap_unorderedを使えば進捗が出る。

import time, random
from tqdm import tqdm
from multiprocessing import Pool

# random時間sleepして引数に1を足して返す
def plus_one(x):
    time.sleep(random.random() * 10)
    return x + 1

# 4倍
with Pool(4) as pool:
    imap = pool.imap(plus_one, range(10))
    result = list(tqdm(imap))

しかしこれだとtqdm側はiterationの回数がわからず割合や終了予測時間等が表示されない。

事前にサイズがわかっていればtqdmの引数にtotalを渡せば解決。

with Pool(4) as pool:
    r = range(10)
    imap = pool.imap(plus_one, r)
    result = list(tqdm(imap, total=len(r)))

おまけ - ファイルを1行ずつ並列処理

個人的によく使うファイルを1行ずつ読み込んで処理するのを並列に実行するコードも貼っておく。

事前に1〜2000までの数字が書いてある2000行のファイル、test.txtを作る。

$ vi test.txt

i1
esc
qa V y p Ctr+a q 1998@a
:wq

1度ファイルを読み込んで行数をカウントし、imap_unorderedで各行に対して処理を実行。

import time
from tqdm import tqdm
from multiprocessing import Pool
from functools import reduce

# ファイルを読んで行数を取得するコード
with open('test.txt') as f:
    line_count = reduce((lambda x, y: y[0]), enumerate(f, 1)) 

# 1秒眠ってから+1して返す関数
def plus_one(x):
    time.sleep(1)
    return int(x.strip()) + 1

# ファイルを読みながら+1して結果を返す
with open('test.txt') as f, Pool(8) as pool:
    imap = pool.imap_unordered(plus_one, f)
    result = list(tqdm(imap, total=line_count))

4コア8スレッドのCPUを使った際の実行時間が下記。1秒に7.99回処理ができているとうことでほぼ理想的に並列処理が行われている。

2000/2000 [04:10<00:00,  7.99it/s

改定履歴

Author: Masato Watanabe, Date: 2019-01-25, 記事投稿