本文介绍 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这样的小设备上跑到实时。

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

是最新的4.1，不错。

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

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

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

技术栈

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

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

准备工作

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

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

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

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

引用一张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

这里需要注意的是这个-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();

    } 

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