Xavier上使用TensorRT加速MaskRCNN

环境

Jetson AGX Xavier Linux Jetpack 4.5.1
cuda 10.2
cudnn 8.0
TensorRT 7.1.3

代码

Mask RCNN 版本 https://github.com/matterport/Mask_RCNN

将h5文件转为uff文件

1.Clone github 的TensorRT库，对应版本Xavier上的版本

git clone https://github.com/nvidia/TensorRT
cd TensorRT
git submodule update --init --recursive

2.Modify the conv2d_transpose conversion function in UFF

/usr/lib/python3.6/dist-packages/uff/converters/tensorflow/converter_functions.py

这里要注意的是你当前使用的环境下的dist-packages，根据自己的python环境去修改，比如我这里使用的conda py36 环境，在下面目录修改

/home/user/archiconda3/envs/py36/lib/python3.6/site-packages/uff/converters/tensorflow/converter_functions.py

uff_graph.conv_transpose(
	# inputs[2], inputs[1], inputs[0],
    inputs[0], inputs[2], inputs[1],
    strides, padding,
    dilation=None, number_groups=number_groups,
    left_format=lhs_fmt, right_format=rhs_fmt,
    name=name, fields=fields
    )

3.下载Mask R-CNN库并设置PYTHONPATH

git clone https://github.com/matterport/Mask_RCNN.git
export PYTHONPATH=$PYTHONPATH:$PWD/Mask_RCNN

4.Apply the patch into Mask R-CNN repo to update the model from NHWC to NCHW

cd Mask_RCNN
git checkout 3deaec5
patch -p1 < ../0001-Update-the-Mask_RCNN-model-from-NHWC-to-NCHW.patch
cd -

5.下载keras权重文件

wget https://github.com/matterport/Mask_RCNN/releases/download/v2.0/mask_rcnn_coco.h5

使用自己训练的权重文件可以跳过这一步

6.将h5文件转换为uff

python3 mrcnn_to_trt_single.py -w mask_rcnn_coco.h5 -o mrcnn_nchw.uff -p ./config.py

转换后得到.uff文件和.pbtxt文件（mrcnn_nchw.uff和mrcnn_nchw.pbtx）

使用uff文件加速

方案一 直接使用编译好的例子，模型对应coco数据类别

1.Xavier如果在刷机时选择了安装SDK其实就已经安装好了TensorRT，可以在usr/src/tensorrt下找到已经编译好的 trtexec，切换到samples/sampleUffMaskRCNN目录下，里面包含了Makeflie，可以直接执行make命令

make

2.make后会在tensorrt/bin目录下生成sample_Uff_mask_rcnn和sample_Uff_mask_rcnn_debug两个可执行文件，这里编译好的模型架构是对应coco上训练好的权重，可以通过下面命令下载，然后转换

wget https://github.com/matterport/Mask_RCNN/releases/download/v2.0/mask_rcnn_coco.h5

3.在data下新建maskrcnn文件夹，将转好的uff文件，pbtx文件与测试的ppm格式图片放入，执行命令

./sample_Uff_mask_rcnn -d data/maskrcnn/

输出

方案二 在自己训练的数据上，需要调整类别，重新编译

1.修改关键文件上的类别

TensorRT\samples\opensource\sampleUffMaskRCNN\mrcnn_config.h
TensorRT\plugin\multilevelProposeROI\tlt_mrcnn_config.h
TensorRT\plugin\proposalLayerPlugin\mrcnn_config.h

这三个文件修改的地方基本差不多

2.修改uff转换配置文件上的类别信息，这一步需要在模型转换成uff之前修改

TensorRT/samples/opensource/sampleUffMaskRCNN/converted/config.py

3.编译

cd TensorRT # tensorrt 目录
mkdir -p build && cd build
# DTRT_LIB_DIR Xavier上刷机会自动安装TensorRT,所以 lib 目录应该为/usr/lib/aarch64-linux-gnu
# DTRT_OUT_DIR 输出目录
# CMAKE_C_COMPILER 指定编译器,否者可能会遇到 No CMAKE_C_COMPILER could be found
cmake .. -DTRT_LIB_DIR=/usr/lib/aarch64-linux-gnu -DTRT_OUT_DIR=`pwd`/out -DTRT_PLATFORM_ID=aarch64 -DCUDA_VERSION=10.2 -D CMAKE_C_COMPILER="/usr/bin/gcc"
make -j$(nproc)

有时候make -j$(nproc)会出错，可能是编译顺序问题，直接make就好，等待编译结束。。。时间有点长（我这里把整个TensorRT编译了）

4.测试

进入编译输出目录，执行以下命令运行

./out/sample_uff_maskRCNN --datadir data/maskrcnn