使用multiprocessing的pool.starmap返回大数组时遇MaybeEncodingError

MaybeEncodingError本质是进程间传递大对象时，序列化/反序列化过程出错，或是结果堆积导致内存缓冲区溢出。结合你需要保留任务顺序的需求，以下是几种实用解决方案：

1. 分批次处理任务 #

将大任务列表拆分为多个小批次，每次仅处理一个批次的任务，完成后立即获取结果并释放内存，避免一次性堆积所有结果。这种方式改动最小，且天然保留任务顺序。

import multiprocessing as mp
import numpy as np

def process_image(img):
    # 模拟你的图像处理逻辑，返回大型数组
    return np.random.rand(1000, 1000)

if __name__ == "__main__":
    # 模拟输入图像列表
    input_images = [np.random.rand(256, 256) for _ in range(1000)]
    batch_size = 50  # 根据内存情况调整批次大小
    final_results = []

    with mp.Pool(processes=80) as pool:
        # 按批次迭代处理
        for start_idx in range(0, len(input_images), batch_size):
            batch = input_images[start_idx:start_idx+batch_size]
            # starmap保证批次内的结果顺序
            batch_results = pool.starmap(process_image, [(img,) for img in batch])
            final_results.extend(batch_results)

2. 用imap结合索引排序 #

imap返回迭代器，可逐个获取完成的任务结果，避免一次性加载所有数据。通过给任务添加索引，收集结果后按索引排序即可恢复原始顺序。

import multiprocessing as mp
import numpy as np

def process_with_index(task_idx, img):
    # 处理图像并返回索引+结果
    processed_arr = np.random.rand(1000, 1000)
    return (task_idx, processed_arr)

if __name__ == "__main__":
    input_images = [np.random.rand(256, 256) for _ in range(1000)]
    # 给每个任务添加原始索引
    indexed_tasks = [(idx, img) for idx, img in enumerate(input_images)]
    final_results = []

    with mp.Pool(processes=80) as pool:
        # 迭代获取结果
        result_iterator = pool.imap(process_with_index, indexed_tasks)
        # 按索引排序恢复原始顺序
        sorted_results = sorted(result_iterator, key=lambda x: x[0])
        # 提取最终结果数组
        final_results = [res for idx, res in sorted_results]

3. 使用共享内存直接写入 #

通过共享内存让子进程直接写入结果，完全避免进程间传递大型数组，内存效率最高。适合超大规模的数组输出场景。

import multiprocessing as mp
from multiprocessing import shared_memory
import numpy as np

def process_to_shared(img, shm_name, result_shape, dtype, task_idx):
    # 连接到已创建的共享内存
    existing_shm = shared_memory.SharedMemory(name=shm_name)
    # 创建共享数组视图
    shared_result_arr = np.ndarray(result_shape, dtype=dtype, buffer=existing_shm.buf)
    # 处理图像并写入对应位置
    shared_result_arr[task_idx] = np.random.rand(1000, 1000)
    # 关闭共享内存连接（主进程负责释放）
    existing_shm.close()

if __name__ == "__main__":
    total_tasks = 1000
    # 定义结果数组的形状和数据类型
    result_shape = (total_tasks, 1000, 1000)
    dtype = np.float64

    # 创建共享内存，大小匹配结果数组
    shm = shared_memory.SharedMemory(create=True, size=np.prod(result_shape)*dtype.itemsize)
    # 创建共享数组的本地视图
    shared_results = np.ndarray(result_shape, dtype=dtype, buffer=shm.buf)

    input_images = [np.random.rand(256, 256) for _ in range(total_tasks)]
    tasks = [(img, shm.name, result_shape, dtype, idx) for idx, img in enumerate(input_images)]

    with mp.Pool(processes=80) as pool:
        pool.starmap(process_to_shared, tasks)

    # 复制共享内存中的结果到本地数组（可选，若直接使用共享内存可跳过）
    final_results = shared_results.copy()

    # 释放共享内存资源
    shm.close()
    shm.unlink()

4. 替换序列化器为dill #

默认的pickle序列化器对某些大型数组的支持有限，改用dill可提升序列化效率并解决编码错误，代码改动极小。

import multiprocessing as mp
import dill
import numpy as np

def init_worker():
    # 将进程池的序列化器替换为dill
    mp.connection._converter = dill

def process_image(img):
    return np.random.rand(1000, 1000)

if __name__ == "__main__":
    input_images = [np.random.rand(256, 256) for _ in range(1000)]

    # 初始化进程池时传入自定义初始化函数
    with mp.Pool(processes=80, initializer=init_worker) as pool:
        final_results = pool.starmap(process_image, [(img,) for img in input_images])

额外注意事项 #

• 调整进程数：80核服务器无需满负载运行，可尝试减少进程数（如40），降低内存竞争压力。
• 监控内存：使用工具实时监控内存使用，调整批次大小或共享内存配置。
• 清理资源：确保子进程中临时创建的数组及时释放，避免内存泄漏。

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