基于joblib的批处理并行加速

python进阶

本文介绍一行命令加速python中for循环速度的方法。

这里以单反相机RAW格式的文件转换为JPG图像为例，介绍如何使用joblib 提供的库来加速图像批量转换。

Pasted-image-20241109093119.png-50b4bbccf3.png

现在计算机的CPU都有多个计算核心，但python的for循环默认是使用单个核心来处理。所以要加速 for 循环，一个最简单的思路就是将多个计算核心都利用起来，这就是并行处理。python中的 joblib 提供了非常方便地并行处理的方式。核心命令如下：

from joblib import Parallel, delayed






results = Parallel(n_jobs=-1)(delayed(func)(item) for item in box)

上面这段代码，首先是从 joblib中导入必要的对象和函数，然后具体使用时，把用户自定义的函数 func 填入 delayed 中，并且使用列表推导式的方式，对 box 中的 item 进行遍历。结合此案例中批量将 raw 文件转换为 jpg 的需求，还需要额外使用 rawpy 模块，具体的代码如下：

from glob import glob
from PIL import Image
import numpy as np
import rawpy
from joblib import Parallel, delayed
import os, time








def raw2jpg(fp):
    fp2 = fp.replace("ARW", "jpg")
    if not os.path.exists(fp2):
        file = rawpy.imread(fp)
        data = file.postprocess(use_camera_wb=True,
                                half_size=False,
                                no_auto_bright=True,
                                output_bps=16)
        # 直接调用postprocess可能出现偏色问题，所以要指定一些参数
        # https://blog.csdn.net/weixin_38342946/article/details/105789291
        rgb = np.float32(data / 65535 * 255)
        rgb = np.asarray(rgb, np.uint8)
        output = Image.fromarray(rgb)
        fp2 = fp.replace("ARW", "jpg")
        output.save(fp2)
        file.close()








def timer(func):
  '''函数装饰器，用于统计函数执行耗时'''
    def func_wrapper(*args, **kwargs):
        time_start = time.time()
        result = func(*args, **kwargs)
        time_end = time.time()
        time_spend = time_end - time_start
        print(f"{func.__name__} cost time: {time_spend:.3f} s")
        return result
    return func_wrapper




@timer
def sequential_run():
    fps = glob('*/*.ARW')
    for fp in fps:
        raw2jpg(fp)
    t1 = time.time()




@timer
def parallel_run():
    fps = glob('*/*.ARW')
    Parallel(n_jobs=-1)(delayed(raw2jpg)(fp) for fp in fps)
    # n_jobs=-1时，使用最大数量的CPU








if __name__=='__main__':
    # 处理7张ARW图像耗时比较
    # sequential_run()
    # sequential_run cost time: 27.351 s
    parallel_run()
    # parallel_run cost time: 4.428 s

测试机器CPU有16核心，所以并行处理时速度提升还是非常明显的。