在Jetson TX2开发套件上使用TensorRT7.1.0加速YOLOv4
目录
- 前言
- 编译TensorRT-YOLOv4
- 修改TensorRT-YOLOv4源码
- 结语
前言
记录一下最近一两天的简单工作，我在把TX2开发套件刷机、安装最新的JetPack后，跑了跑TensorRT官方给的UFF-SSD示例，然后又把SSD的骨架网络换成轻量级的MobileNetv2，运行发现速度倒是非常不错，平均在20FPS以上，但精度嘛就比较抱歉了，在行人、车辆检测这种常规任务下性能表现得不如人意。当然主要原因是由于没有用特定场景的数据集进行训练微调，不过精度低的小体量模型确实不适合这类任务。于是我开始尝试在TX2上运行YOLOv3和YOLOv4，图个方便，我没有自己从头写代码，而是尝试在GitHub上找直接能用的代码进行测试。简单看了一下，GitHub上开源的支持在TX2上跑的仓库主要分为两大类：一是直接使用onnx转换darknet的YOLO模型为onnx格式，然后通过TensorRT加速后运行，无需TensorFlow、PyTorch等深度学习框架，典型的有：
- TensorRT-YOLOv4
- TensorRTx
- YOLO-TensorRT
二是使用PyTorch、TensorFlow等框架，先将darknet的YOLO模型进行转换，接着导出到onnx或者uff格式，最后经过TensorRT加速后运行，典型的有：
- PyTorch-YOLOv4
- TensorRT-Demos
我本着几个项目总有一两个能直接跑起来的想法，从TensorRT-YOLOv4开始尝试，直接就在2080Ti的机子上跑起来了，也是比较嗨心，于是暂时没有去尝试其他开源项目。接着我把TensorRT-YOLOv4放到TX2上编译，结果出现了TensorRT版本不兼容的错误，但错误的修改看起来不是特别复杂，于是我把原项目fork之后简单进行了一些修改，在安装有JetPack4.4DP的TX2板子上成功运行了起来，并且结果也是正确的。仓库地址是tensorrt-yolov4，有类似环境需求的读者可以试试，不过修改后的代码应该不能在TensorRT5.x/6.x的环境下编译通过，因此在低版本TensorRT的环境下建议使用原作者的仓库，尽管修改很简单。接下来就对tensorrt-yolov4改动的一些内容进行简要的说明。
编译TensorRT-YOLOv4
如README所描述的那样，首先需要安装所需的环境依赖，如我使用的TX2上的环境如下：
Ubuntu 18.04
TensorRT 7.1.0
CUDA 10.2
cuDNN 8.0.0
onnx 1.4.1
OpenCV 4.1.1
CMake 3.10.2
除此之外，在cmake检查环境的过程中如果没有安装protobuf也会报错，因此需要执行下面的指令进行安装：
sudo apt-get install libprotobuf-dev protobuf-compiler
为了让编译顺利通过，需要对项目顶层目录的CMakeLists.txt做简单的修改，修改其中的GPU_ARCHS为当前机器GPU的数字，如TX2的GPU_ARCHS为62，RTX 2080Ti则为75。具体的架构信息可在NVIDIA官网查询，或者在装有TensorFlow-GPU的机器上打开Python，创建一个TensorFlow会话查看显示信息，以及可编译/usr/local/cuda/samples/1_Utilities/deviceQuery/下的例子，执行例子会输出相应的GPU_ARCHS信息。
此外，原项目中在编译的时候会提示找不到opencv2/opencv.hpp头文件，原因是source/目录下的CMakeLists.txt未包含相关头文件，导致编译的时候报错。在CMakeLists.txt中添加如下两条内容即可：
find_package(OpenCV REQUIRED)
include_directories(${OpenCV_INCLUDE_DIRS})
修改TensorRT-YOLOv4源码
要在安装有TensorRT7.1.0的TX2板子上编译通过，需要对项目的源码进行简单的修改。修改的地方有三处，其中第一处是注释掉onnx-tensorrt/builtin_plugins.cpp:80处的REGISTER_BUILTIN_NVPLUGIN("Concat", ConcatPlugin);，由于TensorRT更新的版本去掉了对应的API，并且将Concat插件层集成到了的框架中，因此需要将这行代码注释掉。
第二处要修改的地方是onnx-tensorrt/ResizeNearest.cu:35，替换定义的getOutputDimensions()函数如下：
nvinfer1::Dims ResizeNearestPlugin::getOutputDimensions(
int index, const nvinfer1::Dims* inputDims, int nbInputs) {
assert(nbInputs == 1);
nvinfer1::Dims const& input = inputDims[0];
assert(input.nbDims == 3);
assert(_ndims == 2);
assert(index == 0);
nvinfer1::Dims output;
output.nbDims = input.nbDims;
int s = 0;
for( int d=0; d原因是TensorRT更新的版本在解析onnx模型时，传入插件层的CHW张量维度的类型定义有变化，导致针对旧版本的is_CHW()宏不起作用，但其它未改变。因此应去掉is_CHW()宏中判断CHW维度类型的部分。
第三处要修改的地方是onnx-tensorrt/yolo.cu:102，修改的原因类似第二处，需要替换文件的getOutputDimensions()函数如下：
nvinfer1::Dims YOLOPlugin::getOutputDimensions(
int index, const nvinfer1::Dims *inputDims, int nbInputs) {
assert(index == 0);
assert(inputDims);
assert(nbInputs == 1);
assert(index == 0);
nvinfer1::Dims const& input = inputDims[0];
assert(input.nbDims == 3);
nvinfer1::Dims output;
output.nbDims = input.nbDims;
for( int d=0; d最后，由于原项目中的onnx版本为1.3.0，而Python使用onnx版本却是1.4.1。因此我将onnx-tensorrt/third_party/下的onnx替换到了1.4.1版本，并且替换后的版本编译也能顺利通过。
结语
在编译通过后，我在TX2板子上试了试运行YOLOv3和YOLOv4。TX2在MAXN模式下使用float16进行量化能跑到7-8FPS，实时性满足不了要求。不过相对于TensorRT的UFF-SSD例程来说，精度倒是高了不少。如果再结合模型压缩的一些技巧，并在特定场景的应用数据集上做训练微调，那么是能达到实时性与精度的双重要求的，这也是我接下要做得事情。行吧，今天就到这里。