運行多個 ONNX Model 以在 Python 中並行推理

Question

有沒有辦法並行運行多個 ONNX 模型並使用多個可用內核？

目前，我已經訓練了兩個 ONNX 模型並想使用它們進行推斷。 我使用了 Python 的線程，但它並沒有真正使用多核。

之后我嘗試了 multiprocessing 但這給了我以下錯誤：

無法腌制 onnxruntime.capi.onnxruntime_pybind11_state.SessionOptions 對象

請讓我知道是否有任何解決方法？

Answer 1

第 1 步：如果您在 GPU 上運行應用程序，則以下解決方案將有所幫助。

import multiprocessing

CUDA runtime 不支持 fork start 方法所以使用 spawn，在多進程調用之前如下所示：

multiprocessing.set_start_method('spawn')

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

第 2 步：您需要在個人流程 memory 中傳遞您的參考 object 或者您可以使用共享 memory 方法，如下所示：

 from multiprocessing.managers import BaseManager
 from PythonFile import ClassName

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

現在您可以將此“obj”作為參數傳遞給進程調用，然后可以在所有多個進程中訪問該參數。