本文介绍 JetsonNano2GB编译Yolov5并用TensorRT加速教程

JetsonNano2GB编译Yolov5并用TensorRT加速教程

This article was original written by Jin Tian, welcome re-post, first come with https://jinfagang.github.io . but please keep this copyright info, thanks, any question could be asked via wechat: jintianiloveu

上一期我们教大家如何给新的JetsonNano 2GB烧录系统。这一期我们将教大家如何在JetsonNano上部署最新的Yolov5检测模型,并且采用TensorRT加速,看看我们的模型能否在JetsonNano这样的小设备上跑到实时。

image-20201107222319190.png

首先我们来确认一下系统里面opencv对应的版本:

image-20201107223801733.png

是最新的4.1,不错。

今天这是我们要实现的效果:

d

image.png

是的,你没有看错,我们将在JetsonNano上跑一个经过TensorRT加速的YoloV5算法,算法的功能是实现口罩佩戴的检测。从图中你可以看到,对于口罩佩戴的检测还是很准确的,而且基本上可以达到实时!这里要说明的是,这个模型的输入尺寸不大,并且在Yolov5的基础上经过了裁剪,所以能够做到这么快,但是如果你的模型直接拿官方的部署,可能无法跑到这个速度,最后速度带来的当然是精度的降低,不过在可以接受的范围之内!

还等什么,快来依照我们的教程来实现吧~

技术栈

本篇文章不涉及到Python或者任何训练方面的东西,纯部署,而且纯C++,因此教程里面可能会出现一些你不熟悉的库或者术语,没关心,你看到最后我们跑出来的结果就可以了,至于过程你可以通过加入我们的社群来获取想要的材料。大部分代码都在github开源。核心技术包括:

  • C++;
  • TensorRT相关基础知识;
  • 对ONNX了解;
  • 对Yolov5的基本网络结构了解。

准备工作

先做一些准备工作,什么准备工作呢?我们需要先安装一个C++下面比较常用的可视化库,thor:

git clone https://github.com/jinfagang/thor

大家可以进入到主页依照readme进行编译。为什么要做这个准备干工作?原因有两个:

  • 它依赖一些库,如果你能够成功编译它,说明你的环境没有问题,如果不能,你的环境可能有问题了;
  • 它可以很方便的可视化我们的检测输出结果。

image-20201107222409982.png

引用一张thor官方的图:

画检测结果既好看又方便!

编译onnx-tensorrt

我们给大家准备好了一个转好的onnx文件,如果你需要导出自己的onnx,过程有点复杂,你可以通过加入社群交流如何导出,我们没有现成的代码提供。

git clone --recursive https://github.com/onnx/onnx-tensorrt

编译onnx-tensorrt 在JetsonNano上可能会遇到一些问题。为什么呢?

第一个错误:

CMake Error at CMakeLists.txt:21 (cmake_minimum_required):
  CMake 3.13 or higher is required.  You are running version 3.10.2


-- Configuring incomplete, errors occurred!

然后你就需要重新编译一下CMake了:

git clone https://github.com/Kitware/CMake

说实话这个CMake也是异常的庞大了。clone下来100多M。然后开始编译CMake了:

./bootstrap.sh
make
sudo make install
cmake -V

编译过程还是挺慢的.

然后就可以接着编译onnx-tensorrt了。

编译完onnx-tensorrt之后,就可以用 onnx2trt 将模型转出到tensorrt 的engine了。

在JetsonNano上将ONNX转到engine文件

想必大家已经安装好了必备的工具,如果你在安装工具的时候遇到问题,不要灰心丧气,因为最终是可以work的。我们已经走到这一步了。

接下来就是拿一个onnx来做试验了。这里我推荐一个onnx文件,大小只有6M,我们将用它来实现前面的检测效果。

至于如何训练对应的模型,大家可以在神力平台找到对应的训练代码。

http://manaai.cn

我们成功的将onnx转到了engine:

onnx2trt best_sim.onnx -o a.trt -w 10909920

image-20201110220437019

这里需要注意的是这个-w, 它指的是我们的workspace多大,你不要默认,也不要设置的太小,大概按照我这样的设置就可以了,默认的话就死机了,太小会报错。

对TensorRT Engine进行推理

接下来就要对这个Engine文件进行推理了。推理部分我们要实现load上面的tensorrt engine并实现口罩的后处理检测。对应的代码也可以在神力平台找到:

http://manaai.cn

最后我们这里贴出核心的推理代码:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
YoloV5 model(inputParams, trtParams, v5Params);
  if (!model.initSession(0)) {
    std::cerr << "model init failed!\n";
    exit(0);
  }

  string data_f = FLAGS_data_f;
  if (thor::os::isfile(data_f)) {
    if (thor::os::suffix(data_f) == "mp4" ||
        thor::os::suffix(data_f) == "avi") {
      std::cout << "inference on video.\n";
      VideoCapture cap(data_f);
      if (!cap.isOpened()) {
        cout << "Error opening video stream or file" << endl;
        return -1;
      }
      cap.set(CAP_PROP_POS_FRAMES, 60);

      VideoWriter video("yolov5_result.mp4", CV_FOURCC('M', 'J', 'P', 'G'), 25,
                        Size(cap.get(3), cap.get(4)));

      while (1) {
        Mat image;
        cap >> image;

        if (image.empty()) break;

        auto bboxes = model.predOneImage(image);

        std::vector<thor::Detection> dets = toThorDetections(bboxes);
        auto res =
            thor::vis::VisualizeDetections(image, dets, CLASSES, &COLORS);
        cv::imshow("yolov5", res);
        video.write(res);

        char c = (char)waitKey(25);
        if (c == 27) break;
      }
      video.release();

    } 

由此可以得到我们上面给大家展示的效果。