Python函数顺序组合并行化：数组迭代应用实现问询

我有一个维度为(N, d)的数组A，一个输入输出均为(d,)数组的函数f，以及正整数M。需要生成维度为(N, M+1, d)的数组B，其中B[i, j, :]等于f对A[i, :]迭代应用j次的结果（f^0表示恒等函数，即j=0时直接取A[i, :]）。

循环版代码可以简单实现这个逻辑，我也写了优化后的compute_B_slice函数，但在用multiprocessing做并行化时遇到了问题，求可行的解决方案。

循环版实现示例 #

import numpy as np

def f(x):
    # 示例函数：对输入数组做线性变换
    return x * 0.5 + 1

def compute_B_loop(A, f, M):
    N, d = A.shape
    B = np.zeros((N, M+1, d))
    for i in range(N):
        current = A[i, :]
        B[i, 0, :] = current
        for j in range(1, M+1):
            current = f(current)
            B[i, j, :] = current
    return B

# 测试用例
A = np.random.rand(1000, 3)
M = 5
B_loop = compute_B_loop(A, f, M)

并行化实现方案 #

方案1：单样本并行（适合小维度样本） #

将每个样本独立交给进程处理，最后合并结果：

import numpy as np
from multiprocessing import Pool

def f(x):
    return x * 0.5 + 1

def compute_single_sample(x, f, M):
    # 处理单个样本，生成(M+1, d)的结果数组
    d = x.shape[0]
    result = np.zeros((M+1, d))
    result[0] = x
    current = x
    for j in range(1, M+1):
        current = f(current)
        result[j] = current
    return result

def compute_B_parallel(A, f, M, num_workers=None):
    N, _ = A.shape
    # 将数组拆分为单个样本的列表
    samples = [A[i, :] for i in range(N)]
    # 启动进程池处理
    with Pool(num_workers) as pool:
        results = pool.starmap(compute_single_sample, [(sample, f, M) for sample in samples])
    # 合并结果为目标维度数组
    return np.array(results)

# 测试验证
B_parallel = compute_B_parallel(A, f, M)
assert np.allclose(B_loop, B_parallel)

方案2：切片批量并行（减少进程通信开销） #

当N很大时，将数组拆分为大切片批量处理，降低进程间数据传递的开销：

def compute_slice(A_slice, f, M):
    slice_N, d = A_slice.shape
    B_slice = np.zeros((slice_N, M+1, d))
    for i in range(slice_N):
        current = A_slice[i, :]
        B_slice[i, 0, :] = current
        for j in range(1, M+1):
            current = f(current)
            B_slice[i, j, :] = current
    return B_slice

def compute_B_parallel_slice(A, f, M, num_workers=None):
    N, _ = A.shape
    # 自动适配进程数（默认用CPU核心数）
    if num_workers is None:
        num_workers = Pool()._processes
    # 拆分数组为对应数量的切片
    slice_sizes = [N // num_workers] * num_workers
    # 处理余数，将多出来的样本分配到前几个切片
    for i in range(N % num_workers):
        slice_sizes[i] += 1
    # 生成切片列表
    slices = []
    start_idx = 0
    for size in slice_sizes:
        end_idx = start_idx + size
        slices.append(A[start_idx:end_idx, :])
        start_idx = end_idx
    # 进程池批量处理切片
    with Pool(num_workers) as pool:
        results = pool.starmap(compute_slice, [(s, f, M) for s in slices])
    # 合并所有切片结果
    return np.concatenate(results, axis=0)

# 测试验证
B_parallel_slice = compute_B_parallel_slice(A, f, M)
assert np.allclose(B_loop, B_parallel_slice)

关键注意事项 #

• 函数序列化：f必须能被pickle序列化，避免用闭包、未在顶层模块定义的函数。如果函数复杂，可改用pathos.multiprocessing（支持更多序列化方式）替代标准库的multiprocessing。
• 内存控制：当N和M较大时，生成的B数组会占用大量内存，需确保系统有足够资源。
• 进程数选择：num_workers建议设为CPU核心数，过多进程会增加上下文切换开销，降低效率。

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