如何在Python多进程中传入半径列表批量计算激光雷达平面度

问题背景 #

现有激光雷达平面度计算函数，通过多进程对不同邻域半径值重复调用，但每次都要手动创建队列、进程变量，代码冗余且维护麻烦。希望能直接传入半径列表实现批量并行处理，简化流程。

优化方案 #

可以用multiprocessing.Pool替代手动管理进程和队列，它能自动处理任务分发、结果收集和进程生命周期管理，完美适配批量参数的并行计算需求。

步骤1：调整计算函数

去掉原函数中的output_queue参数，让函数直接返回计算结果，简化逻辑：

import numpy as np
from tqdm import tqdm

def calculate_planarity(lidar_data, neighborhood_radius, tree):
    planarity_values = np.zeros(len(lidar_data))
    try:
        for point_idx, point in tqdm(enumerate(lidar_data), total=len(lidar_data), desc=f"Radius={neighborhood_radius}"):
            neighbor_indices = tree.query_ball_point(point, neighborhood_radius)
            neighbors = lidar_data[neighbor_indices]

            if neighbors.shape[0] >= 3:
                cov_matrix = np.cov(neighbors, rowvar=False)
                eigenvalues, _ = np.linalg.eig(cov_matrix)
                sorted_eigenvalues = np.sort(eigenvalues)
                planarity = (sorted_eigenvalues[1] - sorted_eigenvalues[0]) / (sorted_eigenvalues.sum())
                planarity_values[point_idx] = planarity
            else:
                planarity_values[point_idx] = np.nan
    except Exception as e:
        print(f"Error in calculate_planarity (radius={neighborhood_radius}): {e}")
        raise
    return planarity_values

步骤2：用Pool实现批量并行处理

在主程序中，定义半径列表，通过Pool.starmap批量执行任务，自动收集结果：

from multiprocessing import Pool

if __name__ == "__main__":  
    # 假设inPoints和tree已经提前初始化
    radius_list = [1, 3, 10]
    
    # 创建进程池，进程数可根据CPU核心数调整，默认用全部核心
    with Pool() as pool:
        # 构建参数元组列表，每个元组对应一个任务的参数
        task_args = [(inPoints, radius, tree) for radius in radius_list]
        # 批量执行任务，获取结果列表
        planarity_results = pool.starmap(calculate_planarity, task_args)
    
    # 结果按半径顺序对应：planarity_results[0]是radius=1的结果，以此类推
    planarity_result = planarity_results[0]
    planarity_result3 = planarity_results[1]
    planarity_result10 = planarity_results[2]

进阶优化：减少大对象传递开销

如果lidar_data和tree是超大对象，每次传递会带来额外开销，可以用Pool的initializer参数把它们设为全局变量，避免重复序列化传递：

def init_worker(lidar_data_global, tree_global):
    global lidar_data, tree
    lidar_data = lidar_data_global
    tree = tree_global

def calculate_planarity_optimized(neighborhood_radius):
    # 直接使用全局变量lidar_data和tree，参数只保留radius
    planarity_values = np.zeros(len(lidar_data))
    try:
        for point_idx, point in tqdm(enumerate(lidar_data), total=len(lidar_data), desc=f"Radius={neighborhood_radius}"):
            neighbor_indices = tree.query_ball_point(point, neighborhood_radius)
            neighbors = lidar_data[neighbor_indices]

            if neighbors.shape[0] >= 3:
                cov_matrix = np.cov(neighbors, rowvar=False)
                eigenvalues, _ = np.linalg.eig(cov_matrix)
                sorted_eigenvalues = np.sort(eigenvalues)
                planarity = (sorted_eigenvalues[1] - sorted_eigenvalues[0]) / (sorted_eigenvalues.sum())
                planarity_values[point_idx] = planarity
            else:
                planarity_values[point_idx] = np.nan
    except Exception as e:
        print(f"Error in calculate_planarity (radius={neighborhood_radius}): {e}")
        raise
    return planarity_values

if __name__ == "__main__":  
    radius_list = [1, 3, 10]
    
    with Pool(initializer=init_worker, initargs=(inPoints, tree)) as pool:
        planarity_results = pool.map(calculate_planarity_optimized, radius_list)
    
    # 结果对应关系不变
    planarity_result, planarity_result3, planarity_result10 = planarity_results

优势说明 #

• 无需手动创建队列和进程变量，代码简洁易维护
• 自动处理进程的启动、执行、等待和结果收集
• 支持批量传入任意长度的半径列表，扩展性强
• 进阶优化方案减少大对象的重复传递，提升运行效率

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