前言

onnx模型作为中间模型,相较于pytorch直接推理,是有加速度效果的,且推理代码简单,不需要load各种网络。最近某些项目因为显存不够,onnxruntime推理时切换CPU/GPU,实现某些模型在CPU上推理,某些在GPU上推理。 查了一些别人的文章发现很多人都说onnxruntime推理没法像pytorch那样指定GPU,CPU,只能卸载一个GPU用CPU,卸载CPU用GPU。个人感觉不应该是这样的,点进去源码一看,明明有配置CPU,GPU的参数,而且很简单。这里把自己踩得一些坑给记录一些,分享给有需要的人。

onnxruntime CPU/GPU切换

点进去源码之后看到有CUDAExecutionProvider和CPUExecutionProvider这样的选择,开始使用的使用低版本onnxruntime-gpu的时候发现选了“CPUExecutionProvider”依旧是用GPU的资源,后来从一条警告看到1.10版本之后建立推理引擎的时候需要providers参数选择CPU/GPU,安装了onnxruntime-gpu==1.10.0之后果然可以在GPU版本的onnxruntime下使用CPU了。

import onnxruntime as rt
sess = rt.InferenceSession(MODEL_PATH,providers=['CPUExecutionProvider'])# providers=['CPUExecutionProvider','CUDAExecutionProvider']

注意:高版本ONNX动态尺寸推理的时候,如果inference输出的shape和你模型文件本身的outputshape不匹配会报警告,这时候会降低推理速度,这时候改一下模型本身设output shape就行。

修改onnx的输入/输出尺寸为动态

其实这个在转onnx模型的时候input/output就应该设置为静态或动态,但是如果没有原模型,只有一个onnx模型,可以使用下面给的方法快速转换模型为动态输入输出。就把对应的输入输出某一需要动态的维度改成“?”就行了。

import onnx
file_path = './my.onnx'
model = onnx.load(file_path)
#model.graph.input[0].type.tensor_type.shape.dim[0].dim_param = '?'
model.graph.output[0].type.tensor_type.shape.dim[2].dim_param = '?'
model.graph.output[0].type.tensor_type.shape.dim[3].dim_param = '?'
onnx.save(model, './my_dynamic.onnx')

后续:

每次我都要挑战一下这个发文助手,看起来有点搞笑,东西就这么点非说我字数不够,的检测为什么能这么无聊,非要把字数的权重弄得那么高,让人看着很无语,我这就看看到底多少字才能通过检测,看看评论区多少在吐槽这个烂检测系统,产品经理都装看不见。