在 Python 中并行运行多个 ONNX 模型进行推理

Question

有没有办法并行运行多个 ONNX 模型并使用多个可用内核？

\n

目前，我已经训练了两个 ONNX 模型，并想使用它们进行推断。我已经使用了 Python 的线程，但这并没有\xe2\x80\x99 真正使用多核。

\n

之后我尝试了多重处理，但这给了我以下错误：

\n

\n

无法pickle onnxruntime.capi.onnxruntime_pybind11_state.SessionOptions对象

\n

\n

请告诉我是否有任何解决方法？

\n

Answer 1

步骤 1：如果您正在 GPU 上运行应用程序，以下解决方案将会有所帮助。

import multiprocessing

CUDA运行时不支持fork启动方法，因此在多进程调用之前使用spawn，如下所示：

multiprocessing.set_start_method('spawn')

更多理解请参考： https://github.com/microsoft/onnxruntime/issues/7846

第2步：您需要在单独的进程内存中传递引用的对象，或者也可以使用共享内存方法，如下所示：

 from multiprocessing.managers import BaseManager
 from PythonFile import ClassName

 BaseManager.register('LabelName', ClassName)
 manager = BaseManager()
 manager.start()
 obj = manager.ClassName()

现在，您可以将此“obj”作为参数传递给进程调用，然后可以在所有多个进程中访问该参数。