如何基于现有多进程代码实现单Inference Session多线程并行推理？

我想要编写多线程代码，创建一个推理会话（Inference Session）并同时处理多张图像。目前找到的代码是用多进程实现的，想问下multiprocessing.Process类有没有相关参数能实现多线程需求，还是需要在其中引入多线程类？基于下面给出的代码，具体该怎么操作？

import onnxruntime as ort
import numpy as np
import multiprocessing as mp


def init_session(model_path):
    EP_list = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    sess = ort.InferenceSession(model_path, providers=EP_list)
    return sess


class PickableInferenceSession:  # 这是一个包装类，用于让当前的InferenceSession类可被序列化
    def __init__(self, model_path):
        self.model_path = model_path
        self.sess = init_session(self.model_path)

    def run(self, *args):
        return self.sess.run(*args)

    def __getstate__(self):
        return {'model_path': self.model_path}

    def __setstate__(self, values):
        self.model_path = values['model_path']
        self.sess = init_session(self.model_path)


class IOProcess(mp.Process):
    def __init__(self):
        super(IOProcess, self).__init__()
        self.sess = PickableInferenceSession('model.onnx')

    def run(self):
        print("calling run")
        print(
            self.sess.run({}, {
                'a': np.zeros((3, 4), dtype=np.float32),
                'b': np.zeros((4, 3), dtype=np.float32)
            }))
        #print(self.sess)


if __name__ == '__main__':
    mp.set_start_method(
        'spawn')  # 这一点很重要，必须放在name==main代码块内
    io_process = IOProcess()
    io_process.start()
    io_process.join()

解答 #

核心结论 #

multiprocessing.Process本身是多进程实现，没有参数可以直接转成多线程。要实现多线程需求，要么直接改用threading.Thread类，要么在进程内部结合多线程使用。另外要注意：ONNX Runtime的InferenceSession本身并非线程安全，多线程共享会话时需要通过配置控制并行策略，避免资源冲突。

方案1：直接改用多线程（推荐单会话多场景） #

不需要原代码中的序列化包装类，多线程共享同一进程内存，可直接复用推理会话实例：

import onnxruntime as ort
import numpy as np
import threading

def init_session(model_path):
    EP_list = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    # 配置会话并行参数，避免多线程冲突
    sess_options = ort.SessionOptions()
    sess_options.inter_op_num_threads = 1  # 算子间并行设为1，避免多线程抢占资源
    sess_options.intra_op_num_threads = 4  # 算子内并行数，根据CPU核心数调整
    return ort.InferenceSession(model_path, sess_options=sess_options, providers=EP_list)

# 单张图像处理逻辑
def process_image(sess, image_data):
    result = sess.run({}, image_data)
    print(f"线程{threading.current_thread().name}处理完成，结果：{result}")

if __name__ == '__main__':
    # 初始化共享的推理会话
    model_path = 'model.onnx'
    sess = init_session(model_path)

    # 模拟多张待处理图像数据
    image_list = [
        {'a': np.zeros((3, 4), dtype=np.float32), 'b': np.zeros((4, 3), dtype=np.float32)},
        {'a': np.ones((3, 4), dtype=np.float32), 'b': np.ones((4, 3), dtype=np.float32)},
        {'a': np.random.rand(3,4).astype(np.float32), 'b': np.random.rand(4,3).astype(np.float32)}
    ]

    # 创建并启动多线程
    threads = []
    for idx, img_data in enumerate(image_list):
        thread = threading.Thread(target=process_image, args=(sess, img_data), name=f"Thread-{idx+1}")
        threads.append(thread)
        thread.start()

    # 等待所有线程执行完毕
    for thread in threads:
        thread.join()

    print("所有图像处理完成")

方案2：进程内结合多线程（适合多CPU核心场景） #

如果需要保留多进程架构（比如利用独立CPU核心），可以在每个进程内部启动多线程处理图像批次：

import onnxruntime as ort
import numpy as np
import multiprocessing as mp
import threading

def init_session(model_path):
    EP_list = ['CUDAExecutionProvider', 'CPUExecutionProvider']
    sess_options = ort.SessionOptions()
    sess_options.inter_op_num_threads = 1
    sess_options.intra_op_num_threads = 2
    return ort.InferenceSession(model_path, sess_options=sess_options, providers=EP_list)

# 进程内的线程处理逻辑
def process_image_in_thread(sess, img_data):
    result = sess.run({}, img_data)
    print(f"进程{mp.current_process().name}的线程{threading.current_thread().name}处理完成：{result}")

class IOProcess(mp.Process):
    def __init__(self, image_batch):
        super().__init__()
        self.image_batch = image_batch
        self.sess = None

    def run(self):
        # 进程内初始化推理会话
        self.sess = init_session('model.onnx')
        # 启动多线程处理当前批次图像
        threads = []
        for idx, img_data in enumerate(self.image_batch):
            thread = threading.Thread(target=process_image_in_thread, args=(self.sess, img_data), name=f"Thread-{idx+1}")
            threads.append(thread)
            thread.start()
        for thread in threads:
            thread.join()

if __name__ == '__main__':
    mp.set_start_method('spawn')
    # 拆分图像为两个批次，分配给不同进程
    image_list = [
        {'a': np.zeros((3,4), np.float32), 'b': np.zeros((4,3), np.float32)},
        {'a': np.ones((3,4), np.float32), 'b': np.ones((4,3), np.float32)},
        {'a': np.random.rand(3,4).astype(np.float32), 'b': np.random.rand(4,3).astype(np.float32)},
        {'a': np.full((3,4), 2.0, np.float32), 'b': np.full((4,3), 2.0, np.float32)}
    ]
    batch1 = image_list[:2]
    batch2 = image_list[2:]

    # 启动两个进程
    p1 = IOProcess(batch1)
    p2 = IOProcess(batch2)
    p1.start()
    p2.start()
    p1.join()
    p2.join()
    print("所有进程的线程处理完成")

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