onnx 部署模型，如何实现类似 pytorch 的 DataLoader 来实现并行预处理图片？

Question

训练用 pytorch；生产的时候，会把 pytorch 的模型转成 onnx 格式的，然后用 onnxruntime 来推理

但是有一些问题，因为 python 有 GIL，所以调用 numpy+PIL 预处理图片的时候，最多只能使用一个 CPU core。这样就太慢了，瓶颈变成了图片的预处理

像 pytorch 提供了 DataLoader ，用多进程来实现数据的预处理

但是现在转成了 onnx 之后，onnx 没有提供 DataLoader 怎么办？

总不能生产环境推理的时候，图片预处理用 pytorch、推理用 onnxruntime 吧，感觉很不伦不类

Answer 1

你可以用 Python 的多进程库（multiprocessing）在部署 ONNX 模型时实现类似于 PyTorch DataLoader 的功能。先定义一个预处理函数，接受图像路径列表作为输入，并返回预处理后的图像列表：

import numpy as np
from PIL import Image

def preprocess_images(image_paths):
    preprocessed_images = []
    for image_path in image_paths:
        # 使用 PIL 和 numpy 完成图像预处理
        image = Image.open(image_path)
        # ... 图像预处理代码 ...
        preprocessed_image = np.array(image)
        preprocessed_images.append(preprocessed_image)
    return preprocessed_images

使用 multiprocessing 库并行预处理图像：

import os
import glob
from multiprocessing import Pool

def parallel_preprocess_images(image_paths, num_workers):
    chunk_size = len(image_paths) // num_workers
    chunks = [image_paths[i:i + chunk_size] for i in range(0, len(image_paths), chunk_size)]

    with Pool(num_workers) as p:
        preprocessed_images = p.map(preprocess_images, chunks)

    # 合并所有预处理后的图像
    preprocessed_images = [image for chunk in preprocessed_images for image in chunk]

    return preprocessed_images

if __name__ == "__main__":
    image_folder = "path/to/your/image/folder"
    image_paths = glob.glob(os.path.join(image_folder, "*.jpg"))

    num_workers = 4
    preprocessed_images = parallel_preprocess_images(image_paths, num_workers)