如何利用 python 多处理和 concurrent.futures 的所有内核？

Question

我正在编写一个简单的脚本来对涉及调整大小和添加过滤器的图像数据集进行一些预处理。

这是我的代码：

def preprocessing(tar_ratio, img_paths, label_paths, 
                  save_dir="output", resampling_mode=None):
    # with concurrent.futures.ThreadPoolExecutor() as executor:
    with concurrent.futures.ProcessPoolExecutor() as executor:
        for img_path, label_path in zip(img_paths, label_paths):
            src_ratio = get_ratio(label_path)
            if src_ratio is not np.nan:
                executor.submit(
                    process_single(src_ratio, tar_ratio, img_path, label_path, 
                                   save_dir=save_dir, resampling_mode=resampling_mode)
                )
            else:
                pass

我认为它更受 CPU 限制，因此multiprocessing比multithreading更合适。 但是在尝试了这两种方法之后，在只使用两个 CPU 内核的情况下，两者都没有按预期工作。

我已经阅读了以下帖子，我想知道是否有使用concurrent.futures的更新版本？ 如何利用 python 多处理的所有内核

Answer 1

Executor.submit方法接受一个可调用对象作为第一个参数，但你调用了 function，试试这个：

executor.submit(
    process_single,
    src_ratio,
    tar_ratio,
    img_path,
    label_path,
    save_dir=save_dir,
    resampling_mode=resampling_mode,
)

一个说明正确用法的简单示例：

测试.py：

import random
import time

from concurrent.futures import ProcessPoolExecutor


def worker(i):
    t = random.uniform(1, 5)
    print(f"START: {i} ({t:.2f}s)")
    time.sleep(t)
    print(f"END: {i}")
    return i * 2


def main():
    futures = []

    with ProcessPoolExecutor() as executor:
        for i in range(5):
            futures.append(executor.submit(worker, i))

        print([f.result() for f in futures])


if __name__ == "__main__":
    main()

例子：

$ python test.py
START: 0 (3.16s)
START: 1 (1.68s)
START: 2 (2.76s)
START: 3 (1.53s)
START: 4 (4.05s)
END: 3
END: 1
END: 2
END: 0
END: 4
[0, 2, 4, 6, 8]