(備忘録) Python で並列処理しつつ進捗を確認する

はじめに

並列処理を行う際、単純にループを回すだけでは全体の進捗が分かりにくいため、tqdm と組み合わせる手法の備忘録です。

前提 tqdm

from tqdm import tqdm
import time

# iterable なオブジェクトを tqdm() でラップするだけ
for i in tqdm(range(100)):
    time.sleep(0.01) # 何らかの処理

最低限の使い方は、 tqdm(イテレーター, total=int)。上限が自動的に取得できないときはtotalを指定しておくとわかりやすい（なくても動く、現在何％かはわからなくなるが）

やりかた

import multiprocessing
from tqdm import tqdm

def task(x):
    # 何らかの重い処理
    return x * 2

dataset = range(100)

with multiprocessing.Pool() as pool:
    # imap_unordered は完了した順に結果を返すため、進捗表示と相性が良い
    imap = pool.imap_unordered(task, dataset)
    # tqdm でラップすることで進捗を表示
    results = list(tqdm(imap, total=len(dataset)))

resultsを得る際のイテレーターにtqdmを噛ませるとよい

ほかでも似たような感じで使える

ProcessPoolExecutor

from concurrent.futures import ProcessPoolExecutor, as_completed
from tqdm import tqdm

def task(n):
    return n * 2

data = range(100)

with ProcessPoolExecutor() as executor:
    futures = [executor.submit(task, x) for x in data]

    results = [
        f.result() for f in tqdm(as_completed(futures), total=len(data))
    ]

joblib

from joblib import Parallel, delayed
from tqdm import tqdm

def task(i):
    return i * 2

results = Parallel(n_jobs=-1)(
    delayed(task)(i) for i in tqdm(range(100))
)