如何用multiprocessing Pool结合scikit-learn Pipeline批量预测图像？

针对你遇到的内存不足、多进程使用错误等问题，以下是几种高效可行的解决方案，从简单到进阶排列：

方案1：分块批量预测（最直接，无需多进程） #

无需复杂并行逻辑，利用HDF5支持切片读取的特性，分批次加载数据并预测，完全控制内存占用：

import h5py
import numpy as np

# 打开HDF5文件（无需一次性加载全部数据）
with h5py.File('test_images.h5', 'r') as hf:
    X_test = hf['X_test']  # HDF5 Dataset对象，支持切片
    total_samples = X_test.shape[0]
    batch_size = 1000  # 根据内存调整，比如每次处理1000张
    y_preds = []
    
    # 分批次读取并预测
    for start in range(0, total_samples, batch_size):
        end = min(start + batch_size, total_samples)
        batch = X_test[start:end]  # 读取当前批次的展平图像(shape=(batch_size,22500))
        batch_preds = model.predict(batch)
        y_preds.append(batch_preds)

# 合并所有预测结果
y_preds = np.concatenate(y_preds)

方案2：正确使用multiprocessing并行预测 #

你之前的错误在于：Pool.map会遍历传入的可迭代对象，直接传model.predict会导致单样本预测（效率极低），打包元组的方式也错误（把Pipeline和整个数据集作为两个独立任务传入）。以下是正确实现：

方法A：用functools.partial固定模型参数 #

import multiprocessing as mp
from functools import partial
import numpy as np
import h5py

def predict_batch(model, batch):
    # 若批次是150x150格式，先展平为模型需要的22500维向量
    if batch.ndim == 3:
        batch = batch.reshape(batch.shape[0], -1)
    return model.predict(batch)

# 从HDF5分块生成批量任务
with h5py.File('test_images.h5', 'r') as hf:
    X_test = hf['X_test']
    total_samples = X_test.shape[0]
    batch_size = 1000
    batches = []
    for start in range(0, total_samples, batch_size):
        end = min(start + batch_size, total_samples)
        batches.append(X_test[start:end])

# 固定模型参数，生成仅接收batch的函数
predict_with_model = partial(predict_batch, model)

# 多进程并行预测
n_cores = mp.cpu_count()
with mp.Pool(n_cores) as pool:
    results = pool.map(predict_with_model, batches)

y_preds = np.concatenate(results)

方法B：打包模型与批次为任务元组 #

import multiprocessing as mp
import numpy as np
import h5py

def predict_wrapper(task_args):
    model, batch = task_args
    if batch.ndim == 3:
        batch = batch.reshape(batch.shape[0], -1)
    return model.predict(batch)

# 生成(模型, 数据批次)的任务列表
with h5py.File('test_images.h5', 'r') as hf:
    X_test = hf['X_test']
    total_samples = X_test.shape[0]
    batch_size = 1000
    tasks = []
    for start in range(0, total_samples, batch_size):
        end = min(start + batch_size, total_samples)
        tasks.append( (model, X_test[start:end]) )

# 多进程执行任务
n_cores = mp.cpu_count()
with mp.Pool(n_cores) as pool:
    results = pool.map(predict_wrapper, tasks)

y_preds = np.concatenate(results)

方案3：用Dask处理超大规模数据 #

针对千万级以上图像，Dask支持延迟计算和分布式并行，可无缝对接scikit-learn：

import dask.array as da
import h5py
from dask_ml.wrappers import ParallelPostFit

# 将Pipeline包装为支持Dask的并行模型
parallel_model = ParallelPostFit(model)

# 用Dask加载HDF5数据（不占用内存，仅生成计算图）
with h5py.File('test_images.h5', 'r') as hf:
    X_dask = da.from_array(hf['X_test'], chunks=(1000, 22500))  # chunks大小适配内存

# 自动分块并行预测
y_preds_dask = parallel_model.predict(X_dask)

# 计算并获取结果（或直接保存到磁盘）
y_preds = y_preds_dask.compute()

关键注意事项 #

• 始终批量处理数据：单样本预测效率极低，sklearn的predict本身支持批量输入，务必分批次处理。
• 利用HDF5的切片特性：无需一次性加载全部数据，这是处理海量数据的核心前提。
• 多进程任务粒度：每个任务对应一个数据批次，而非单个样本，避免进程调度开销。

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