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

Z时代
2024-02-29
分类：IT

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

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

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

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

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

回答：

你可以用 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)