当前位置: 萬仟网 > IT编程>软件设计>面向对象 > 在Jetson TX2开发套件上使用TensorRT7.1.0加速YOLOv4

在Jetson TX2开发套件上使用TensorRT7.1.0加速YOLOv4

2020年07月26日  | 萬仟网IT编程  | 我要评论
目录前言前言

前言

记录一下最近一两天的简单工作,我在把TX2开发套件刷机、安装最新的JetPack后,跑了跑TensorRT官方给的UFF-SSD示例,然后又把SSD的骨架网络换成轻量级的MobileNetv2,运行发现速度倒是非常不错,平均在20FPS以上,但精度嘛就比较抱歉了,在行人、车辆检测这种常规任务下性能表现得不如人意。当然主要原因是由于没有用特定场景的数据集进行训练微调,不过精度低的小体量模型确实不适合这类任务。于是我开始尝试在TX2上运行YOLOv3和YOLOv4,图个方便,我没有自己从头写代码,而是尝试在GitHub上找直接能用的代码进行测试。简单看了一下,GitHub上开源的支持在TX2上跑的仓库主要分为两大类:一是直接使用onnx转换darknet的YOLO模型为onnx格式,然后通过TensorRT加速后运行,无需TensorFlow、PyTorch等深度学习框架,典型的有:

二是使用PyTorch、TensorFlow等框架,先将darknet的YOLO模型进行转换,接着导出到onnx或者uff格式,最后经过TensorRT加速后运行,典型的有:

我本着几个项目总有一两个能直接跑起来的想法,从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<input.nbDims; ++d ) {
    if (d == 0) {
	  output.type[d] = nvinfer1::DimensionType::kCHANNEL;
	  output.d[d] = input.d[d];
	}
    else {
	  output.type[d] = input.type[d];
	  output.d[d] = int(input.d[d] * _scale[s++]);
	}
  }
  return output;
}

原因是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<input.nbDims; ++d ) {
        output.type[d] = input.type[d];
        output.d[d] = input.d[d];
    }
	output.type[0] = nvinfer1::DimensionType::kCHANNEL;
    output.d[0] = 7;
    return output;
}

最后,由于原项目中的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例程来说,精度倒是高了不少。如果再结合模型压缩的一些技巧,并在特定场景的应用数据集上做训练微调,那么是能达到实时性与精度的双重要求的,这也是我接下要做得事情。行吧,今天就到这里。

本文地址:https://blog.csdn.net/hlld__/article/details/107551084

如您对本文有疑问或者有任何想说的,请点击进行留言回复,万千网友为您解惑!

相关文章:

验证码:
Copyright © 2017-2020  萬仟网 保留所有权利. 粤ICP备17035492号-1