PIL图像处理多进程并行优化失效及AssertionError报错排查

咱们先拆解你遇到的问题，从实现细节到多进程的适配性，一步步来解决：

一、你的并行实现里的核心坑点 #

你的代码逻辑本身是通顺的（按列拆分独立处理），但进程间通信（IPC）的开销完全盖过了并行计算的收益，甚至小尺寸分割时触发错误，主要原因有这几个：

1. 大对象的重复序列化/反序列化开销 #

你通过partial把neigh（sklearn近邻模型）、img_train_rot（训练图像集合）这些大对象传给每个子进程，这会导致：

• 每个任务都要把这些大对象通过pickle序列化后传给子进程，子进程再反序列化，这个过程的耗时可能比处理单个切片的时间还长。
• PIL图像的pickle效率极低，因为它内部包含C级别的资源，序列化时需要额外转换，数据量越大开销越夸张。

2. 并行颗粒度过细（小尺寸分割时更严重） #

当你减小lenx时，图像被拆成更多列切片，任务数量暴增。进程池的调度开销（比如进程切换、任务分配）会大幅增加，而每个任务的处理时间很短，完全抵消了并行的优势。

3. 小尺寸分割时的AssertionError #

这个错误大概率是pickle大对象时的内存溢出或数据结构异常导致的——切片越多，需要传递的图像数据总和越大，加上每个任务附带的大模型/训练图像，超过了pickle的处理阈值。

二、为什么单切片工作量足够时还没提速？ #

哪怕单切片工作量不小，你传递的neigh和img_train_rot这些大对象的序列化开销，可能已经超过了并行处理节省的时间。举个例子：如果处理一个切片需要0.1秒，但传递大对象给子进程需要0.2秒，那并行反而比串行慢。

三、优化方案：从减少IPC开销入手 #

1. 让子进程共享大对象，避免重复传递 #

用进程池的initializer参数，在每个子进程启动时一次性初始化模型和训练数据，而不是每个任务都传递一次：

import multiprocessing
import numpy as np
from PIL import Image

# 子进程初始化函数：全局变量存储共享资源
def init_worker(train_imgs, split_dict, img_train_rot_np):
    global neigh, split_dict_global, img_train_rot_global
    neigh = setup_nearest_neighbour(train_imgs, n_neighbors=1)
    split_dict_global = split_dict
    # 把PIL图像转成numpy数组存储，避免pickle PIL对象
    img_train_rot_global = img_train_rot_np

def work_image_parallel(img_slice_np, leny):
    # 把numpy数组转回PIL图像
    img_slice = Image.fromarray(img_slice_np)
    constructed_img_slice = Image.new(mode='L', size=img_slice.size)
    
    for y in range(0, img_slice.size[1], leny):
        box = (0, y, img_slice.size[0], y+leny)
        img_part = img_slice.crop(box)
        res = neigh.kneighbors(np.asarray(img_part).ravel().reshape((1,-1)))
        # 注意：这里i是近邻查询的结果索引，比如i = res[1][0][0]
        i = res[1][0][0]
        # 从共享的numpy数组转回PIL图像并裁剪
        matched_img = Image.fromarray(img_train_rot_global[i])
        constructed_img_slice.paste(matched_img.crop(new_box), (0, y))
    
    # 返回numpy数组，减少pickle开销
    return np.asarray(constructed_img_slice)

if __name__ == '__main__':
    lenx, leny = 16, 16
    # 准备你的训练数据、测试图像等
    train_imgs = ...
    split_dict = ...
    img_train_rot = ...
    test_img = test_imgs[0]
    
    # 把测试图像切片转成numpy数组，避免传递PIL对象
    img_slices_np = [
        np.asarray(test_img.crop((x, 0, x+lenx, test_img.size[1])))
        for x in range(0, test_img.size[0], lenx)
    ]
    # 把训练图像转成numpy数组列表，用于子进程共享
    img_train_rot_np = [np.asarray(img) for img in img_train_rot]
    
    # 初始化进程池，传入共享资源
    with multiprocessing.Pool(
        initializer=init_worker,
        initargs=(train_imgs, split_dict, img_train_rot_np)
    ) as pool:
        # 用starmap传递切片和leny参数
        res_np = pool.starmap(work_image_parallel, [(slice, leny) for slice in img_slices_np])
    
    # 拼接结果
    test_result2 = Image.new(mode='L', size=test_img.size)
    for i, slice_np in enumerate(res_np):
        slice_img = Image.fromarray(slice_np)
        test_result2.paste(slice_img, box=(i*lenx, 0, i*lenx+lenx, test_result2.size[1]))

2. 调整并行粒度，减少任务数量 #

如果lenx很小，不要按单列拆分，而是按块组拆分（比如一次处理4列或8列），让每个任务的工作量足够大，抵消进程调度开销。比如把图像分成10个大的垂直块，每个块包含多列，这样任务数量减少，每个任务的处理时间更长。

3. 用更高效的并行库：joblib #

joblib对numpy数组和sklearn模型的序列化效率远高于原生multiprocessing，还支持内存映射，减少数据拷贝：

from joblib import Parallel, delayed
import numpy as np
from PIL import Image

def work_image_parallel(x, lenx, leny, neigh, split_dict, img_train_rot, test_img):
    img_slice = test_img.crop((x, 0, x+lenx, test_img.size[1]))
    constructed_img_slice = Image.new(mode='L', size=img_slice.size)
    
    for y in range(0, img_slice.size[1], leny):
        box = (0, y, img_slice.size[0], y+leny)
        img_part = img_slice.crop(box)
        res = neigh.kneighbors(np.asarray(img_part).ravel().reshape((1,-1)))
        i = res[1][0][0]
        constructed_img_slice.paste(img_train_rot[i].crop(new_box), (0, y))
    
    return (x, constructed_img_slice)

if __name__ == '__main__':
    lenx, leny = 16, 16
    neigh = setup_nearest_neighbour(train_imgs, n_neighbors=1)
    test_img = test_imgs[0]
    
    # 用joblib并行，n_jobs=-1表示用全部CPU核心
    results = Parallel(n_jobs=-1, verbose=10)(
        delayed(work_image_parallel)(x, lenx, leny, neigh, split_dict, img_train_rot, test_img)
        for x in range(0, test_img.size[0], lenx)
    )
    
    # 拼接结果
    test_result2 = Image.new(mode='L', size=test_img.size)
    for x, slice_img in results:
        test_result2.paste(slice_img, box=(x, 0, x+lenx, test_result2.size[1]))

四、总结 #

你的并行逻辑本身没问题，但大对象的IPC开销是核心瓶颈。通过让子进程共享资源、调整并行粒度、用更高效的序列化工具，就能大幅提升并行效率，同时解决小尺寸分割时的AssertionError问题。

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