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

目标：在多个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)]

本多处理器教程提供了许多并行处理任何任务的方法。

然而，我想知道哪种方法最适合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()来传递一个函数名和两个单独的迭代。

将行替换为：

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