ONNX 运行时推理 | session.run() 多处理

Question

目标：运行在多个 CPU 核上并行推理

我正在使用 simple_onnxruntime_inference.ipynb 进行推理试验。

个人：

outputs = session.run([output_name], {input_name: x})

很多：

outputs = session.run(["output1", "output2"], {"input1": indata1, "input2": indata2})

依次：

%%time
outputs = [session.run([output_name], {input_name: inputs[i]})[0] for i in range(test_data_num)]

这个多处理 tutorial 提供了许多并行化任何任务的方法。

但是，我想知道哪种方法最适合 session.run()，无论是否通过 outputs。

我如何将所有输出和输入一起并行推理？

代码：

import onnxruntime
import multiprocessing as mp

session = onnxruntime.InferenceSession('bert.opt.quant.onnx')

i = 0
# First Input
input_name = session.get_inputs()[i].name
print("Input Name  :", input_name)

# First Output
output_name = session.get_outputs()[i].name
print("Output Name  :", output_name)  

pool = mp.Pool(mp.cpu_count())

# PARALLELISE THIS LINE
outputs = [session.run([], {input_name: inputs[i]})[0] for i in range(test_data_num)]
# outputs = pool.starmap(func, zip(iter_1, iter_2))

pool.close()

print(results)

Update：这个建议使用 starmap() 和 zip() 来传递一个函数名和 2 个独立的迭代器。

用这个替换行：

outputs = pool.starmap(session.run, zip([output_name], [ {input_name: inputs[i]}[0] for i in range(test_data_num) ]))

回溯：

---------------------------------------------------------------------------
KeyError                                  Traceback (most recent call last)
<ipython-input-45-0aab302a55eb> in <module>
     25 #%%time
     26 #outputs = [session.run([output_name], {input_name: inputs[i]})[0] for i in range(test_data_num)]
---> 27 outputs = pool.starmap(session.run, zip([output_name], [ {input_name: inputs[i]}[0] for i in range(test_data_num) ]))
     28 
     29 pool.close()

<ipython-input-45-0aab302a55eb> in <listcomp>(.0)
     25 #%%time
     26 #outputs = [session.run([output_name], {input_name: inputs[i]})[0] for i in range(test_data_num)]
---> 27 outputs = pool.starmap(session.run, zip([output_name], [ {input_name: inputs[i]}[0] for i in range(test_data_num) ]))
     28 
     29 pool.close()

KeyError: 0

Answer 1

def run_inference(i):
    output_name = session.get_outputs()[0].name
    return session.run([output_name], {input_name: inputs[i]})[0]  # [0] bc array in list

outputs = pool.map(run_inference, [i for i in range(test_data_num)])

大家可以批评指正

ONNX 运行时推理 | session.run() 多处理

ONNX Runtime Inference | session.run() multiprocessing

parallel-processing

inference

multiprocessing

onnx

onnxruntime