Goal: run Inference in parallel on multiple CPU cores
I'm experimenting with Inference using simple_onnxruntime_inference.ipynb.
Individually:
outputs = session.run([output_name], {input_name: x})
Many:
outputs = session.run(["output1", "output2"], {"input1": indata1, "input2": indata2})
Sequentially:
%%time
outputs = [session.run([output_name], {input_name: inputs[i]})[0] for i in range(test_data_num)]
This Multiprocessing tutorial offers many approaches for parallelising any tasks.
However, I want to know which approach would be best for session.run(), with or without outputs being passed.
How do I Inference all outputs and inputs together, in parallel?
Code:
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: this solution suggests using starmap() and zip() in order to pass a function name and 2 separate iterables.
Replacing line with this:
outputs = pool.starmap(session.run, zip([output_name], [ {input_name: inputs[i]}[0] for i in range(test_data_num) ]))
Traceback:
---------------------------------------------------------------------------
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
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)])
Anyone feel free to critique