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

Question

有没有办法并行运行多个 ONNX 模型并使用多个可用内核？
目前，我已经训练了两个 ONNX 模型并想使用它们进行推断。我使用了 Python 中的线程，但这并没有真正使用多核。
之后，我尝试了多处理，但这给了我以下错误:

can't pickle onnxruntime.capi.onnxruntime_pybind11_state.SessionOptions objects

请让我知道是否有任何解决方法？

Answer 1

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

import multiprocessing

CUDA 运行时不支持 fork start 方法，因此在多进程调用之前使用 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”作为参数传递给进程调用，然后可以在所有多个进程中访问该参数。