Python多进程结合tqdm：进度条无法正常更新问题排查

问题描述 #

使用Python multiprocessing.Pool 加速3D LIDAR单点特征计算，结合tqdm实现进度条，但运行时进度条显示后无更新，看似计算未启动。原代码如下：

import numpy as np
from multiprocessing import Pool
from tqdm import tqdm
from functools import partial

# np.random.seed(0)
lidar_data = np.random.uniform(low=0.0, high=100.0, size=(1000, 3))

def compute_feature(i, lidar_data, radius):
    x_center, y_center, z_center = lidar_data[i]
    height = z_center
    bounding_box_x_min = x_center - radius - 5
    bounding_box_x_max = x_center + radius + 5
    bounding_box_y_min = y_center - radius - 5
    bounding_box_y_max = y_center + radius + 5
    points_in_cylinder = []
    z_height = []

    for point in lidar_data:
        x, y, z = point
        if bounding_box_x_min <= x <= bounding_box_x_max and bounding_box_y_min <= y <= bounding_box_y_max:
            if z <= z_center:
                dist_to_center = np.sqrt((x - x_center)**2 + (y - y_center)**2)
                if dist_to_center <= radius and z_center - height <= z:
                    points_in_cylinder.append(point)
                    z_height.append(z)

    points_in_cylinder = np.array(points_in_cylinder)
    z = [round(float(point[2]), 3) for point in points_in_cylinder]
    minimum_z = min(z) if points_in_cylinder.size > 0 else z_center
    feature = z_center - 2 * minimum_z
    return feature

# lidar_data = np.random.uniform(low=0.0, high=100.0, size=(1000, 3))
radius = 3.0

with Pool() as pool:
    func = partial(compute_feature, lidar_data=lidar_data, radius=radius)
    features = list(tqdm(pool.imap(func, range(lidar_data.shape[0])), total= lidar_data.shape[0], desc="Computing feature"))

features = np.array(features)

问题成因 #

1. 大数组重复复制开销：通过partial传递lidar_data时，每个子进程都会复制一份完整的LIDAR数组，内存占用暴增，导致进程启动和数据传输耗时极长，看起来像计算未开始。
2. 任务调度通信开销：pool.imap默认块大小为1，每个任务单独和主进程通信，1000个任务的频繁通信掩盖了计算进度，进度条无法及时获取更新。
3. 单任务计算效率极低：原函数用Python循环遍历所有点，单任务耗时过长，进一步延缓了进度条的更新频率。

解决方案 #

1. 共享全局数据避免重复复制 #

使用Pool的initializer参数，在子进程启动时一次性加载lidar_data为全局变量，避免每个任务重复传递大数组。

2. 优化任务分块减少通信开销 #

给imap设置合适的chunksize，将多个任务打包成块传递给子进程，降低通信次数。

3. 矢量化运算提升单任务速度 #

用Numpy矢量化操作替代Python循环，大幅减少单任务计算时间。

修改后的完整代码 #

import numpy as np
from multiprocessing import Pool
from tqdm import tqdm
from functools import partial

# 全局变量，用于子进程共享LIDAR数据
global_lidar_data = None

def init_worker(lidar_data):
    """子进程初始化函数，加载共享数据"""
    global global_lidar_data
    global_lidar_data = lidar_data

def compute_feature(i, radius):
    """优化后的特征计算函数，使用矢量化操作"""
    x_center, y_center, z_center = global_lidar_data[i]
    height = z_center
    
    # 生成边界掩码
    mask_x = (global_lidar_data[:, 0] >= x_center - radius -5) & (global_lidar_data[:, 0] <= x_center + radius +5)
    mask_y = (global_lidar_data[:, 1] >= y_center - radius -5) & (global_lidar_data[:, 1] <= y_center + radius +5)
    mask_z = (global_lidar_data[:, 2] <= z_center) & (global_lidar_data[:, 2] >= z_center - height)
    
    # 计算距离中心的距离掩码
    dx = global_lidar_data[:, 0] - x_center
    dy = global_lidar_data[:, 1] - y_center
    dist_mask = np.sqrt(dx**2 + dy**2) <= radius
    
    # 合并所有掩码筛选点
    total_mask = mask_x & mask_y & mask_z & dist_mask
    points_in_cylinder = global_lidar_data[total_mask]
    
    minimum_z = np.round(points_in_cylinder[:, 2], 3).min() if points_in_cylinder.size > 0 else z_center
    return z_center - 2 * minimum_z

# 生成测试数据
lidar_data = np.random.uniform(low=0.0, high=100.0, size=(1000, 3))
radius = 3.0

with Pool(initializer=init_worker, initargs=(lidar_data,)) as pool:
    func = partial(compute_feature, radius=radius)
    # 按进程数均分任务块大小，平衡负载
    chunksize = max(1, lidar_data.shape[0] // pool._processes)
    # 使用imap并设置chunksize，进度条可实时更新
    features = list(tqdm(pool.imap(func, range(lidar_data.shape[0]), chunksize=chunksize),
                         total=lidar_data.shape[0], desc="Computing feature"))

features = np.array(features)

额外优化建议 #

如果不需要严格保持特征的顺序，可以改用pool.imap_unordered，它会在任务完成后立即返回结果，进度条更新更及时：

features = list(tqdm(pool.imap_unordered(func, range(lidar_data.shape[0]), chunksize=chunksize),
                     total=lidar_data.shape[0], desc="Computing feature"))

内容的提问来源于stack exchange，提问作者Purple_Ad