TX2之TensorRT加速TensorFlow目标检测模型

TX2之TensorRT加速TensorFlow目标检测模型

本文链接：https://blog.csdn.net/zong596568821xp/article/details/86622430

参考：https://github.com/jkjung-avt/tf_trt_models/tree/tf-1.8

在之前的文章中，详细介绍了tensorflow目标检测接口的使用方法，包括训练与推理流程，该工作主要是在服务器端完成的，而我们实际应用的时候，推理不一定发生在服务器上，而可能是TX2等终端设备，因此，在TX2上进行模型的推理加速，对于实际应用很有必要。

幸运的是，nvidia已经针对推理加速提供了一套框架tensorrt，关于tensorrt的介绍，我之前也提到过，这里不再赘述，本文主要讲解如何在TX2上应用tensorrt来对目标检测模型进行加速。

注意：本文介绍的tensorrt加速方法与官网介绍的有区别，不是在x86主机上生成uff文件然后导入到TX2上，而是直接在TX2上用tensorrt优化生成pb文件，然后按照传统的方法读入推理（关于第一种实现方法，有时间会尝试）

1 环境准备

TX2刷机JetPack3.3，安装时注意勾选TensorRT

tensorflow-gpu 1.8（1.8版本支持最友好，1.9/1.10等会有各种问题），安装方法参考《TX2之tensorflow环境部署（最简单）》

opencv3.4

2 下载代码

git clone --recursive https://github.com/jkjung-avt/tf_trt_modelscd tf_trt_models./install.sh

3 安装tensorrt

由于我刷机时没有留意，导致在TX2上没有安装tensorrt，不过没有关系，我们可以手动安装

刷机时，在host主机上生成一个jetpack_download的文件夹，我们把它拷贝到TX2上

然后读进该目录，运行安装以下包

sudo dpkg -i libnvinfer4_4.1.3-1+cuda9.0_arm64.debsudo dpkg -i libnvinfer-dev_4.1.3-1+cuda9.0_arm64.debsudo dpkg -i libnvinfer-samples_4.1.3-1+cuda9.0_arm64.debsudo dpkg -i tensorrt_4.0.2.0-1+cuda9.0_arm64.deb

4 运行示例

4.1 示例1

原作者提供了两个例子，第一个例子是运行摄像头进行检测，我给TX2插了一个USB摄像头，然后运行以下语句

python3 camera_tf_trt.py --usb --model ssd_mobilenet_v1_coco --build

表示构建ssd_mobilenet_v1_coco优化的Tensorrt模型，然后进行摄像头的实时检测，运行后结果如下

可以看到检测速度能达到10帧，还是很不错的

4.2 示例2

第二个例子是读取优化后的模型，然后进行图片检测

python3 camera_tf_trt.py --image --filename examples/detection/data/huskies.jpg --model ssd_mobilenet_v1_coco

运行结果如下

5 优化自己的模型

通过以上步骤，我们就将所需要的环境配置完毕了，接下来就开始优化我们自己的模型，测试一下实际效果

首先还是按照《TensorFlow之目标检测API接口调试（超详细）》里边介绍的方法，准备数据，训练模型，我这里使用的预训练模型是ssd_mobilenet_v1_coco，训练结束后，在结果保存文件里生成了以下三个文件，具体的数字根据训练的回合数可能不一样，总之选择一组模型导出

将上述3个文件导入到TX2的tf_trt_models/data/ssd_mobilenet_v1_egohands文件夹下（其他模型可仿照博主已给的模型方法构建），然后根据README.md文件里的提示，将数字统一改为20000，修改data文件夹下的ssd_mobilenet_v1_egohands.config文件，将num_classes值修改为自己的类别数，否则会报错

同时将训练时的object-detection.pbtxt文件导出，放到data文件夹下

接下来就可以进行模型转化了，运行如下代码，记得第一次运行要加–build选项

python3 camera_tf_trt.py --image --filename examples/detection/data/huskies.jpg --model ssd_mobilenet_v1_egohands --build --labelmap data/object-detection.pbtxt

运行成功之后，会生成ssd_mobilenet_v1_coco_trt.pb文件，接下来我们可以用该文件来验证（不需要再加build选项），比如说检测一个视频

python3 camera_tf_trt.py --file --filename 2018-09-10_162811.mp4 --model ssd_mobilenet_v1_egohands --labelmap data/object-detection.pbtxt

运行成功后，显示如下画面

我又测试了不用tensorrt优化的pb模型，画面如下

可以看到，tensorrt会有一定的速度提升，但是也损失了一部分精度，并不一定好用，我们在实际使用时要注意做对比，选择满足自己使用要求的模型