在Jupyter Notebook中并行化绘图循环的技术问询

问题场景 #

使用Python 3.5.1在Jupyter Notebook中处理2500+个二维数组，通过imshow串行绘制效率极低，尝试用multiprocessing并行提速时遇到TypeError: Invalid dimensions for image data错误——原本维度为(3,2)的数组被传入后变成了(2,)，最初误以为multiprocessing无法处理二维数组传参，实际是对pool.map的传参逻辑理解错误。

错误原因 #

pool.map要求目标函数只能接受单个参数，如果直接传入二维数组，map会将数组的每一行作为独立参数传递给函数，导致函数接收到的是一维数组，从而触发图像维度错误。

正确实现代码 #

import multiprocessing
import numpy as np
import matplotlib.pyplot as plt

# 绘图函数，接收打包后的参数（数组+索引）
def plot_single_array(args):
    arr, idx = args
    # 子进程中切换matplotlib后端，避免Jupyter环境冲突
    plt.switch_backend('Agg')
    fig, ax = plt.subplots(figsize=(4,3))
    ax.imshow(arr, cmap='viridis')
    ax.set_title(f'Array {idx}')
    # 保存绘制结果（若需在主进程显示，可返回图像对象或二进制数据）
    plt.savefig(f'array_plot_{idx}.png', bbox_inches='tight')
    plt.close(fig)
    return f'Completed plot {idx}'

if __name__ == '__main__':
    # 模拟生成2500+个二维数组
    array_collection = [np.random.rand(3, 2) for _ in range(2500)]
    # 将数组与索引打包，满足pool.map的单参数要求
    task_args = [(arr, idx) for idx, arr in enumerate(array_collection)]
    
    # 创建进程池，使用全部可用CPU核心
    with multiprocessing.Pool() as pool:
        # 并行执行绘图任务
        results = pool.map(plot_single_array, task_args)
    
    # 打印部分执行结果（可选）
    for res in results[:5]:
        print(res)

关键注意点 #

• 传参方式：通过将数组与额外参数（如索引）打包成元组，确保pool.map传递的是单个参数，避免数组被拆分。
• Jupyter环境适配：子进程中需切换matplotlib后端为Agg，防止与主进程的交互式绘图环境冲突，同时绘图后要关闭画布释放资源。
• 进程保护：必须用if __name__ == '__main__'包裹主逻辑，这是multiprocessing在Windows和Jupyter环境中正常运行的必要条件，避免无限递归创建子进程。

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