Diffusion Planner数据处理优化：多进程与向量化实战

1. Diffusion Planner数据处理优化实战：从单线程到多进程的蜕变

最近在复现清华大学的Diffusion Planner项目时，遇到了数据处理速度慢的问题。作为一个长期从事自动驾驶算法开发的工程师，我深知数据处理效率对研发进度的影响。原项目的data_process.py采用单线程处理，对于nuPlan这种包含30多万场景的大规模数据集来说，处理时间长得令人难以接受。经过深入分析和优化，我将处理时间从原来的数十小时缩短到2.5小时（在Xeon Gold 6459C CPU上）。下面分享我的完整优化方案。

1.1 问题定位与分析

原项目的README.md给出的数据处理方法是执行data_process.sh脚本，但实际查看代码会发现两个关键问题：

1. 缺乏并行处理：data_process.py虽然使用了NuPlan提供的SingleMachineParallelExecutor，但并未真正启用多进程池，导致CPU资源利用率极低
2. 地图处理未优化：map_process.py中的地图数据处理逻辑存在重复计算和冗余操作，特别是车道边界处理部分

在Xeon Gold 6459C这样的服务器级CPU上（36核72线程），单线程处理显然是对计算资源的巨大浪费。我们的目标是充分利用多核优势，同时优化地图处理的关键路径。

2. 多进程优化方案详解

2.1 数据处理的并行化改造

原始data_process.py的主要处理流程是单线程顺序执行，我们的优化方案是引入Python的multiprocessing.Pool实现真正的并行处理。关键修改点包括：

python复制# 原始单线程处理方式
processor = DataProcessor(args)
results = [processor.process_single_scenario(scenario) for scenario in scenarios]

# 优化后的多进程处理
def process_scenario_wrapper(scenario):
    """每个worker进程内部创建processor实例，避免共享资源冲突"""
    processor = DataProcessor(args)
    return processor.process_single_scenario(scenario)

num_workers = min(cpu_count(), 32)  # 限制最大worker数量
with Pool(num_workers) as pool:
    results = list(tqdm(pool.imap(process_scenario_wrapper, scenarios), total=len(scenarios)))

2.1.1 多进程设计考量

1. 进程池大小：设置为min(cpu_count(), 32)是为了避免过多进程导致的内存溢出问题。实践中发现，对于内存密集型任务，32个worker在128GB内存的服务器上是最佳平衡点
2. worker内创建对象：每个worker内部创建DataProcessor实例，避免多进程共享对象带来的同步问题
3. 进度显示：使用tqdm包装pool.imap，方便观察处理进度和预估剩余时间

重要提示：不要使用多线程(ThreadPool)替代多进程，因为Python的GIL会导致CPU密集型任务无法真正并行。这也是我们选择multiprocessing.Pool的原因。

2.2 地图处理的关键优化

map_process.py的优化主要集中在以下几个方面：

1. 车道边界处理优化：原代码对左右车道边界的方向判断存在重复计算
2. 内存预分配：对固定大小的输出数组预先分配内存，避免append操作带来的频繁内存重分配
3. 向量化计算：将循环操作改为numpy向量化运算

优化后的车道处理核心逻辑：

python复制def _lane_polyline_process(polylines, left_boun