本文介绍 ONNXRuntime采用TensorRT后端快速推理

ONNXRuntime采用TensorRT后端快速推理

This article was original written by Jin Tian, welcome re-post, first come with https://jinfagang.github.io . but please keep this copyright info, thanks, any question could be asked via wechat: jintianiloveu

onnxruntime是一个专门推理onnx模型的前向推理框架,它有一个好处就是模型统一、接口统一但是加速的框架不同。比如你如果在GPU上做推理,那么就用GPU后端,要更快一点可以切换到TensorRT后端,如果你在CPU上做推理,那么就采用openvino来做后端。

由于这些第三方加速很专业,比如TensorRT,比如OpenVINO,onnxruntime的聪明之处在于,集总价之长与一身。但是这样也带来了一些问题,那就是并非所有的模型都能很好的支持,比如CPU支持的op应该是最多的,有一些层TensorRT不支持,这种情况下onnxruntime会自动把不支持的操作,转到CUDA上来,这样就很方便了。其实onnruntime推理,也很简单,你只需要写一次代码,就可以在不同的平台上以最快的速度运行。我们接下来尝试使用onnxruntime来加速一下maskrcnn的onnx模型。

目前来说,有两个maskrcnn模型用的比较多,一个是maskrcnn-benmark, 另外一个是mmdetection, 这两个框架都可以将maskrcnn模型导出到onnx。我们将采用onnxruntime,来尝试推理,在推理之前,我们的思路是,先用CPU后端,然后用TensorRT后端。这样看看CPU结果是否正确,看看是否支持Tensorrt,但是大概率,是不支持tensorrt,至少有个别层是不支持的,这个时候就需要采用onnxruntime的fallback策略,将不支持的操作自动fallback到CUDA上。

在写这篇文章的时候,onnxruntime刚刚增加了TensorRT6.0的支持,这使得我们有可能对一些动态输入的模型在tensorrt中得到支持。比如我们要测试的maskrcnn。

经测试,最新的onnxruntime可以在GPU上推理maskrcnn的onnx模型。可以调用TensorRT的op,但是貌似,有一些不支持的自动在CPU上执行了,导致运行速度很慢。一些op比如:Gather, Concate 等还是在CPU上跑。