报错“can‘t pickle onnxruntime.capi.onnxruntime_pybind11_state.InferenceSession objects“问题解决

在使用onnx模型进行python的多进程推理的时候,遇到这个报错。搞明白报错原因之后,这里给出一个解决办法
使用python多进程进行ONNX和torch模型的的推理时候,习惯性把初始化构建推理引擎,初始化模型这个操作放在主进程,一般来讲都是这样做的,子进程只负责推理。在设置为multiprocessing.set_start_method(‘spawn’),torch可以正常多进程推理。onnx这样的话

import onnxruntime as rt
from multiprocessing import Process
multiprocessing.set_start_method('spawn')
def infer_process(sess,arg):
	pass
sess = rt.InferenceSession("./mine.onnx")
p1 = Process(target=infer_process,args=(sess,args1,))

是会报错"can’t pickle onnxruntime.capi.onnxruntime_pybind11_state.InferenceSession objects"。这是因为python的多进程在multiprocessing .Process在传参时需要用pickle对参数进行序列化操作,但是onnx的推理引擎不在可序列化类型中,属于复杂类型,但是pytorch的model是可以。这时候需要把onnx模型的推理引擎构建放到子进程中。类似的can’t pickle等错误,我猜测也是这个原因。
但这样就相当于每次都额外增加耗时。自己这边采用使用队列的方式,把这个子进程中的初始化onnx引擎放在while循环的外面,然后再循环内读取队列。

猜你喜欢

转载自blog.csdn.net/qq_36276587/article/details/119913628
今日推荐