使用tensorrt_inference库导出RetinaFace-R50模型到TensorRT Engine

前言

TensorRT是NVIDIA推出的一款高效深度学习模型推理框架，其包括了深度学习推理优化器和运行时，能够让深度学习推理应用拥有低时延和高吞吐的优点。

RetinaFace 作为一款InsightFace出品的较为优秀的单阶段人脸检测器，使用TensorRT进行模型转化来提升推理速度很有必要。

RetinaFace TensorRT加速

本文主要参考tensorrt_inference开源库，将RetinaFace -R50模型转为TensorRT engine并使用自带C++代码运行测试。

模型转化

根据开源库中带的说明文档中的步骤执行，对于部分需要注意的步骤会做特别提示。

1. 下载Retinaface源码和模型（InsightFace）

拉取源码：

git clone https://github.com/deepinsight/insightface.git

下载模型文件：

下载好模型文件后解压并放进指定目录：

cd insightface/detection/retinaface
mkdir model
cp R50-0000.params R50-symbol.json model/

2. 安装retinaface所需依赖环境

安装cython，供pycocotools/setup.py调用

pip install cython

在 insightface/detection/retinaface/目录，编译 cxx tools

make

3. 下载tensorrt_inference源码

拉取源码：

git clone https://github.com/linghu8812/tensorrt_inference.git

4. 生成onnx模型文件

拷贝tensorrt_inference的export_onnx.py脚本到原始retinaface目录下：

cd tensorrt_inference/RetinaFace/
cp export_onnx.py insightface/detection/retinaface/

运行export_onnx.py脚本生成onnx模型文件：

cd insightface/detection/retinaface/
python export_onnx.py

运行完成后model目录如下：

5. 拷贝生成的onnx文件到tensorrt_inference/RetinaFace/目录下

cp model/R50.onnx ./tensorrt_inference/RetinaFace/

6. 编译tensorrt_inference的RetinaFace_trt工程

cd ./tensorrt_inference/RetinaFace/

需要注意修改 目录下CMakeLists.txt中TensorRT lib的路径为实际机器的路径：

vi ../CMakeLists.txt

编译工程：

mkdir build && cd build
make -j8

生成可执行文件RetinaFace_trt：

• 注：安装环境

----------软硬件信息------------
	操作系统：Ubuntu20.04
	显卡：GeForce RTX 3090
	TensorRT版本：TensorRT-7.2.1.6
	CUDA版本：cuda-11.1
------------------------------

效果测试

运行RetinaFace_trt对图片进行推理：

cd build
./RetinaFace_trt ../config.yaml ../samples

第一次运行时，会对onnx模型文件进行解析并生成TensorRT engine文件：

其他说明

onnx报错

在运行推理文件时，可能会遇到onnx相关报错。报错信息如下：

$ ./RetinaFace_trt ../config.yaml ../samples
----------------------------------------------------------------
Input filename: ../R50.onnx
ONNX IR version: 0.0.8
Opset version: 15
Producer name:
Producer version:
Domain:
Model version: 0
Doc string:
----------------------------------------------------------------
[01/08/2022-17:09:42] [W] [TRT] /data/installs/TensorRT-7.2.1.6-OSS/parsers/onnx/onnx2tr
t_utils.cpp:220: Your ONNX model has been generated with INT64 weights, while TensorRT doe
s not natively support INT64. Attempting to cast down to INT32.
While parsing node number 200 [Resize]:
ERROR: /data/installs/TensorRT-7.2.1.6-OSS/parsers/onnx/builtin_op_importers.cpp:2594 In
function importResize:
[8] Assertion failed: (mode != "nearest" || nearest_mode == "floor") && "This version of T
ensorRT only supports floor nearest_mode!"
[01/08/2022-17:09:42] [E] Failure while parsing ONNX file
start building engine
[01/08/2022-17:09:42] [E] [TRT] Network must have at least one output
[01/08/2022-17:09:42] [E] [TRT] Network validation failed.
build engine done
RetinaFace_trt: ./tensorrt_inference/RetinaFace/../includes/common/commo
n.hpp:138: void onnxToTRTModel(const string&, const string&, nvinfer1::ICudaEngine*&, cons
t int&): Assertion `engine' failed.
Aborted (core dumped)

经查是onnx版本问题，onnx有版本要求，如果版本过高运行时会报错。需要对onnx降级，需从1.10.0降级到1.5.0。

pip install onnx==1.5.0

NMS操作支持

tensorrt_inference库仅实现了模型的推理，事实上NMS的操作均用C++实现（见代码中的RetinaFace::postProcess函数），在实际使用过程中常常需要将NMS操作也放进TensorRT Engine中。

而TensorRT已开源了NMS op的支持（BatchedNMSDynamic_TRT，详见参考资料[8]），可以通过对onnx运行图进行添加op来支持导出到TRT engine。

import onnx_graphsurgeon as gs

nms = gs.Node(op="BatchedNMSDynamic_TRT",
                  attrs=nms_attrs,
                  inputs=out_tensors,
                  outputs=[
                      nms_num_detections,
                      nms_boxes,
                      nms_scores,
                      nms_classes
                  ])

另外值得一提的是，Retinaface是由bbox分支和landmark分支多分支输出的，所以在做NMS的时候需要保证取舍的bbox和landmark对应。而TensorRT自带的NMS op显然无法满足这一需求，这个需求就需要用户自定义TensorRT的NMS Plugin来满足了。此处不做过多说明。

版权说明

本文为原创文章，独家发布在blog.csdn.net/TracelessLe。未经个人允许不得转载。如需帮助请email至tracelessle@163.com。

参考资料

[1] linghu8812/tensorrt_inference
[2] tensorrt_inference/RetinaFace at master · linghu8812/tensorrt_inference
[3] tensorrt_inference/INSTALL.md at master · linghu8812/tensorrt_inference
[4] InsightFace力作：RetinaFace单阶段人脸检测器 - 知乎
[5] retinaface-R50.zip-深度学习文档类资源-CSDN文库
[6] insightface/detection/retinaface at master · deepinsight/insightface
[7] RetinaFace mnet.25模型的问题 · Issue #75 · linghu8812/tensorrt_inference
[8] TensorRT/plugin/batchedNMSPlugin at master · NVIDIA/TensorRT