如何在Python中利用多进程并行化嵌套循环？

问题描述 #

需要并行化计算函数func(a,b)在数组M和N的所有元素组合上的结果，已将嵌套索引扁平化存储为字典：

idx_map = {} 
count = 0
for i in range(n):
    for j in range(m):
        idx_map[count] = (i,j)
        count += 1 

当前串行执行方式为：

arr = []
for idx in range(n*m):
    i,j = idx_map[idx]
    arr.append( func(M[i], N[j]) )  

希望用Python内置的multiprocessing实现并行化，且无需聚合操作，只需得到所有组合结果的数组，无竞态条件。同时无法理解一段示例代码的含义，请求最优实现方案。

最优实现方案 #

完全可以用multiprocessing实现，而且不需要提前构建索引字典，直接生成所有(i,j)组合更高效，具体步骤如下：

1. 定义并行任务函数 #

先编写worker函数，接收(i,j)对，调用目标函数并返回结果：

def worker(args):
    i, j = args
    return func(M[i], N[j])

如果需要严格保持原串行循环的结果顺序（避免乱序问题），也可以返回索引和结果的元组：

def worker_with_idx(args):
    idx, (i, j) = args
    return (idx, func(M[i], N[j]))

2. 生成所有组合并并行执行 #

用itertools.product生成M和N的索引组合，再用multiprocessing.Pool批量处理：

import multiprocessing
from itertools import product

# 假设M、N、func已定义，n=len(M), m=len(N)
n = len(M)
m = len(N)

# 生成所有(i,j)组合
all_combinations = product(range(n), range(m))

# 创建进程池，默认使用CPU核心数，也可手动指定如processes=4
with multiprocessing.Pool() as pool:
    # map方法按顺序返回结果，和输入组合的顺序完全一致
    results = pool.map(worker, all_combinations)

# results就是所有func(M[i],N[j])的结果数组，顺序和串行循环一致

如果用带索引的worker函数，适合任务执行时间差异较大的场景：

# 生成带索引的组合：(0,(i0,j0)), (1,(i1,j1)), ...
indexed_combinations = enumerate(product(range(n), range(m)))

with multiprocessing.Pool() as pool:
    # imap_unordered可更快获取完成的任务结果
    raw_results = pool.imap_unordered(worker_with_idx, indexed_combinations)
    # 按索引排序恢复原顺序
    sorted_results = sorted(raw_results, key=lambda x: x[0])
    # 提取最终结果数组
    results = [res for idx, res in sorted_results]

3. 解释你提到的示例代码 #

那段代码的核心逻辑是：

• 创建包含28个进程的进程池（作者建议尝试CPU核心数的2倍，调整到CPU满载状态）
• 用imap_unordered批量调用worker函数，传入range(data_Y)的每个整数作为任务参数
• 循环遍历结果但不保存（pass），仅用于等待所有任务执行完毕
• 最后计算并打印并行执行的总耗时
  这段代码是一个极简的并行框架示例，没有处理实际业务逻辑（比如保存结果），所以看起来比较抽象。

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