Jetson嵌入式系列模型部署-3

前言

给大家安利的第二个仓库是tensorRT_Pro。该仓库通过TensorRT的ONNX parser解析ONNX文件来完成模型的构建工作。对模型部署有疑问的可以参考Jetson嵌入式系列模型部署-1，想了解通过TensorRT的Layer API一层层完成模型的搭建工作可参考Jetson嵌入式系列模型部署-2。本文主要是针对于tensorRT_Pro项目中的yolov5完成嵌入式模型部署，本文参考自tensorRT_Pro的README.md，具体操作流程作者描述非常详细，这里再简单过一遍，本次训练的模型使用yolov5s-6.0，类别数为2，为口罩识别。

1. 源码下载

使用如下指令

$ git clone https://github.com/shouxieai/tensorRT_Pro.git

文件较大下载可能比较慢，给出下载好的源码链接Baidu Drive[password:yolo]，若有改动请参考最新

删除不必要的文件，给出简化后的源码链接Baidu Drive[password:yolo]，若有改动请参考最新

2. 环境配置

需要使用的软件环境有TensorRT、CUDA、CUDNN、OpenCV、Protobuf。前四个软件环境在JetPack镜像中已经安装完成，故只需要配置protobuf即可。博主使用的jetpack版本为JetPack4.6.1(PS:关于jetson nano刷机就不再赘述了，需要各位看官自行配置好相关环境，外网访问较慢，这里提供Jetson nano的JetPack镜像下载链接Baidu Drive[password:nano]【更新完毕!!!】(PS:提供4.6和4.6.1两个版本，注意4GB和2GB的区别，不要刷错了)，关于Jetson nano 2GB和4GB的区别可参考链接Jetson NANO是什么？如何选？。(吐槽下这玩意上传忒慢了，超级会员不顶用呀，终于上传完了，折磨!!!)

2.1 Jtop(option)

可使用如下指令查看自己的JetPack版本简单信息

$ cat /etc/nv_tegra_release

使用Jtop可查看JetPack详细信息。Jtop是一个由第三方开发，用于显示Jetson开发板信息的包，可以查询当前板子CPU，GPU使用率，实时功耗，Jetpack软件包信息等，参考自Jetson nano安装jtop，Jetson nano安装pip并换源

2.1.1 配置pip

$ sudo apt install python-pip python3-pip
$ pip3 install --upgrade pip
$ pip install --upgrade pip

pip换源，指令如下

$ sudo mkdir .pip && cd .pip
$ sudo touch pip.conf
$ sudo vim pip.conf

添加如下内容

[global]
timeout = 6000
index-url = https://pypi.tuna.tsinghua.edu.cn/simple
trusted-host = pypi.tuna.tsinghua.edu.cn

2.1.2 安装jtop

$ sudo pip3 install -U jetson-stats

2.1.3 使用jtop

$ sudo jtop
The jetson_stats.service is not active. Please run:
sudo systemctl restart jetson_stats.service

需要启动相关服务，指令如下

$ sudo systemctl restart jetson_stats.service
$ jtop

博主Jtop显示的jetson nano软件包信息页面如下

2.2 编译Protobuf

tensorRT_Pro需要Protobuf用于ONNX解析器，需要下载并编译protobuf源码。这里使用的protobuf版本为3.11.4，若需要修改为其他版本，请参照README/环境配置/适配Protobuf版本。关于protobuf的更多相关介绍请参考protubuf简介，给出protobuf的安装包下载链接地址Baidu Drive[password:yolo]。参考自Linux下编译protobuf，Linux下添加protobuf环境变量

2.2.1 解压

$ mkdir protobuf-3.11.4 && cd protobuf-3.11.4    // 创建protobuf编译的文件夹
$ uzip protobuf-3.11.4.zip    // 解压protobuf压缩包

2.2.2 编译

$ cd protobuf-3.11.4/cmake
$ cmake . -Dprotobuf_BUILD_TESTS=OFF
$ cmake --build .

2.2.3 安装

$ mkdir protobuf  // 创建protobuf安装的文件夹
$ make install DESTDIR=/home/nvidia/protobuf // 指定protobuf安装的路径

注：编译完成之后protobuf文件夹下仅仅只有user一个文件夹，需要将编译好的protobuf/user/local下的bin、include、lib文件夹复制到protobuf当前文件夹下，方便后续tensorRT_Pro项目CMakeLists.txt的指定。

2.2.4 环境配置(option)

• 配置环境变量

  $ sudo vim /etc/profile
  

  添加如下内容保存并退出，注意路径修改为自己的路径

  export PATH=$PATH:/home/nvidia/protobuf/bin
  export PKG_CONFIG_PATH=/home/nvidia/protobuf/lib/pkgconfig/
  

  source生效

  $ source /etc/profile
  

• 配置动态路径

  $ sudo vim /etc/ld.so.conf
  

  追加如下内容，注意路径修改为自己的路径

  /home/nvidia/protobuf/lib
  

• 验证

  protoc --version输出对应版本信息说明安装成功

2.3 配置CMakeLists.txt

主要修改三处

• 1. 修改第10行，选择不支持python(也可选择支持)

  set(HAS_PYTHON OFF)
  

• 2. 修改第20行，修改CUDA路径

  set(CUDA_TOOLKIT_ROOT_DIR  "usr/local/cuda-10.2")
  

• 3. 修改第33行，修改自编译的protobuf的路径

  set(PROTOBUF_DIR "/home/zhlab/protobuf")
  

3. ONNX导出

• 训练的模型使用yolov5s-6.0，torch版本1.7
• torch版本>=1.7时，导出的ONNX模型可以被tensorRT_Pro框架所使用
• 参考自tensorRT_Pro/README/各项任务支持/YoloV5支持，导出细节的说明已经非常详细了，这里简单过一遍

3.1 源码下载

$ git clone -b v6.0 https://github.com/ultralytics/yolov5.git

将训练好的权重文件复制到yolov5文件夹中，给出权重下载链接Baidu Drive[password:yolo]

3.2 修改代码

主要修改两个文件的内容

• 1. yolov5-6.0/models/yolo.py
• 2. yolov5-6.0/export.py

# yolov5/models/yolo.py第55行，forward函数 
# bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
# x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
# 修改为:

bs, _, ny, nx = x[i].shape  # x(bs,255,20,20) to x(bs,3,20,20,85)
bs = -1
ny = int(ny)
nx = int(nx)
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()

# yolov5/models/yolo.py第70行
#  z.append(y.view(bs, -1, self.no))
# 修改为：
z.append(y.view(bs, self.na * ny * nx, self.no))


############# 对于 yolov5-6.0 #####################
# yolov5/models/yolo.py第65行
# if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
#    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
# 修改为:
if self.grid[i].shape[2:4] != x[i].shape[2:4] or self.onnx_dynamic:
    self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)

# disconnect for pytorch trace
anchor_grid = (self.anchors[i].clone() * self.stride[i]).view(1, -1, 1, 1, 2)

# yolov5/models/yolo.py第70行
# y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
# 修改为:
y[..., 2:4] = (y[..., 2:4] * 2) ** 2 * anchor_grid  # wh

# yolov5/models/yolo.py第73行
# wh = (y[..., 2:4] * 2) ** 2 * self.anchor_grid[i]  # wh
# 修改为:
wh = (y[..., 2:4] * 2) ** 2 * anchor_grid  # wh

############# 对于 yolov5-6.0 #####################


# yolov5/export.py第52行
#torch.onnx.export(dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'},  # shape(1,3,640,640)
#                                'output': {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)  修改为
torch.onnx.export(dynamic_axes={'images': {0: 'batch'},  # shape(1,3,640,640)
                                'output': {0: 'batch'}  # shape(1,25200,85) 

3.3 导出ONNX模型

ONNX模型导出指令如下

$ cd yolov5-6.0
$ python export.py --weights=yolov5s.pt --dynamic --include=onnx --opset=11

导出的ONNX模型可使用Netron可视化工具查看，给出ONNX文件下载链接Baidu Drive[password:yolo]。

下图对比展示了原始onnx输出(不加修改直接导出)和简化后onnx输出(按照以上要求修改后导出)的部分差别(第一张未修改直接导出，第二张修改后导出)。主要体现在以下几点：

• 修改后的onnx导出文件更加简洁
• 不必要的节点如Shape、Gather、Unsqueeze等去除
• batch维度指定为动态

3.4 拓展-正确导出ONNX文件

如何正确导出ONNX文件？主要包含以下几条：

1. 对于任何用到shape、size返回值的参数时，例如：tensor.view(tensor.size(0),-1)这类操作，避免直接使用tensor.size的返回值，而是加上int转换如tensor.view(int(tensor(0)),-1)，断开跟踪

2. 对于nn.Unsample或nn.functional.interpolate函数，使用scale_factor指定倍率，而不是使用size参数指定大小

3. 对于reshape、view操作时，-1的指定需放到batch维度。其他维度计算出来即可。batch维度禁止指定为大于-1的明确数字

4. torch.onnx.export指定dynamic_axes参数，并且只指定batch维度，禁止其他动态

5. 使用opset_version=11，不要低于11

6. 避免使用inplace操作，如y[...,0:2] = y[..., 0:2] * 2 - 0.5

7. 尽量少的出现5个维度，例如ShuffleNet Module，可用考虑合并wh避免出现5维

8. 尽量将后处理部分在onnx模型中实现，降低后处理复杂度

注：参考自手写AI的详解TensorRT的C++/Python高性能部署视频，这些做法的必要性体现在，简化过程的复杂度，去掉Gather、Shape类节点，很多时候不这么改看似也可以成功，但是需求复杂后，依旧存在各类问题。按照上述要求修改后，基本总能成，就不需要使用onnx-simplifer了。具体更多细节描述请观看视频。

4. 运行

4.1 源码修改

yolo模型的推理代码主要在src/application/app_yolo.cpp文件中，需要推理的图片放在workspace/inference文件夹中，将上述修改后导出的ONNX文件放在workspace/文件夹下。源码修改较简单主要有以下几点：

• 1. app_yolo.cpp 177行 注释，测试yolov5而非yolov7
• 2. app_yolo.cpp 178行 取消注释，测试yolov5
• 3. app_yolo.cpp 100行 cocolabels修改为mylabels
• 4. app_yolo.cpp 25行 新增mylabels数组，添加自训练模型的类别名称

具体修改如下

//test(Yolo::Type::V7, TRT::Mode::FP32, "yolov7");   //修改1 注释177行
test(Yolo::Type::V5, TRT::Mode::FP32, "yolov5s");    //修改2 取消注释178行

for(auto& obj : boxes){
     ...
     auto name    = cocolabels[obj.class_label];	 //修改3 101行cocolabels修改为mylabels
	 ...
}

static const char* mylabels[] = {"have_mask", "no_mask"};		 //修改4 25行新增代码，为自训练模型的类别名称

4.2 编译

编译生成可执行文件.pro，保存在workspace/文件夹下，指令如下：

$ cd tensorRT_Pro-main
$ mkdir build && cd build
$ cmake .. && make -j8

耐心等待编译完成(PS:需要一段时间)，make -j参数的选取一般时以CPU核心数两倍为宜，参考自make -j参数简介，Linux下CPU核心数可通过lscpu指令查看，jetson nano的cpu核心数为4。

$ lscpu
Architecture:        aarch64
Byte Order:          Little Endian
CPU(s):              4
On-line CPU(s) list: 0-3
Thread(s) per core:  1
Core(s) per socket:  4
Socket(s):           1
Vendor ID:           ARM
Model:               1
Model name:          Cortex-A57
Stepping:            r1p1
CPU max MHz:         1479.0000
CPU min MHz:         102.0000
BogoMIPS:            38.40
L1d cache:           32K
L1i cache:           48K
L2 cache:            2048K
Flags:               fp asimd evtstrm aes pmull sha1 sha2 crc32

编译图解如下所示

4.3 模型构建和推理

编译完成后的可执行文件.pro存放在workspace/文件夹下，故进入workspace/文件夹下执行以下指令

$ cd workspace    // 进入可执行文件目录下
$ ./pro yolo	  // 构建模型并推理 

推理完成后在workspace/文件夹下会生成yolov5s.FP32.trtmodel引擎文件用于模型推理，会生成yolov5s_Yolov5_FP32_result文件夹，该文件夹下保存了推理的图片。模型构建和推理图解如下所示。

模型推理效果如下图所示

4.4 拓展-摄像头检测

简单写了一个摄像头检测的demo，主要修改以下几点:

• 1. app_yolo.cpp 新增app_yolo_video_demo()函数，具体内容参考下面
• 2. app_yolo.cpp 177行 注释
• 3. app_yolo.cpp 176行 新增调用app_yolo_video_demo()函数代码，具体内容参考下面

static void app_yolo_video_demo(const string& engine_file, TRT::Mode mode){  // 修改1
    auto yolo = Yolo::create_infer(
        engine_file,                    // engine file
        Yolo::Type::V5,                 // yolo type, Yolo::Type::V5 / Yolo::Type::X
        0,                              // gpu_id
        0.5f,                           // confidence threshold
        0.5f,                           // nms threshold
        Yolo::NMSMethod::FastGPU,       // NMS method, fast GPU / CPU
        1024,                           // max objects
        false                           // preprocess use multi stream
        );      
    if (yolo == nullptr){
        INFO("Engine is nullptr");
        return;
    }

    cv::Mat frame;
    cv::VideoCapture cap(0);
    if (!cap.isOpened()){
        INFO("Engine is nullptr");
        return;
    }
    
    while (true){
        cap.read(frame);
        auto t0 = iLogger::timestamp_now_float();
        time_t now = time(0);
        auto boxes = yolo->commit(frame).get();
        for (auto &obj : boxes){
            uint8_t b, g, r;
            tie(r, g, b) = iLogger::random_color(obj.class_label);
            cv::rectangle(frame, cv::Point(obj.left, obj.top), cv::Point(obj.right, obj.bottom), cv::Scalar(b, g, r), 5);

            auto name = mylabels[obj.class_label];
            auto caption = iLogger::format("%s %.2f", name, obj.confidence);

            int width = cv::getTextSize(caption, 0, 1, 2, nullptr).width + 10;
            cv::rectangle(frame, cv::Point(obj.left - 3, obj.top - 33), cv::Point(obj.left + width, obj.top), cv::Scalar(b, g, r), -1);
            cv::putText(frame, caption, cv::Point(obj.left, obj.top - 5), 0, 1, cv::Scalar::all(0), 2, 16);
        }
        imshow("frame", frame);
        auto fee = iLogger::timestamp_now_float() - t0;
        INFO("fee %.2f ms, fps = %.2f", fee, 1 / fee * 1000);
        int key = cv::waitKey(1);
        if (key == 27)
            break;
    }
    cap.release();
    cv::destroyAllWindows();
    INFO("Done");
    yolo.reset();
    return;
}

int app_yolo(){
    app_yolo_video_demo("yolov5s.FP32.trtmodel", TRT::Mode::FP32);		// 修改3
    // test(Yolo::Type::V7, TRT::Mode::FP32, "yolov7");					// 修改2
    // test(Yolo::Type::V5, TRT::Mode::FP32, "yolov5s");
    // test(Yolo::Type::V3, TRT::Mode::FP32, "yolov3");
}

进入build/文件夹下重新编译，然后进行workspace/文件夹下运行即可调用摄像头检测，指令如下

$ cd build
$ make -j8
$ cd ../workspace
$ ./pro yolo

图解如下所示

5. 结语

本篇博客只是一个引子，带大家认识到这个项目，并做了最基础的演示，你所看到的只是冰山一角!!!，更多的功能需要各位看官自己去挖掘啦。嵌入式模型部署到这里就告一段落了！

• Jetson嵌入式系列模型部署-1主要讲解了部署的相关知识以及驾驭tensorRT的几种方案讨论；
• Jetson嵌入式系列模型部署-2主要讲解了tensorrtx项目的简单使用，该项目通过tensorRT的Layer API一层层搭建模型，自定义权重加载并通过tensorRT序列化生成engine文件，完成高性能推理。最后带大家在Jetson nano上实现了自定义yolov5模型的简单推理；
• Jetson嵌入式系列模型部署-3即本篇文章主要讲解了tensorRT_Pro项目的简单使用，该项目通过tensorRT的ONNX parser解析ONNX文件搭建模型，最后带大家在Jetson nano上实现了自定义yolov5模型的简单推理。

这两个仓库的剩余潜能需要各位看官自行挖掘了，如果各位觉得有帮助不妨点个⭐️支持下吧。后续看情况如果有需要再更新吧。

6. 下载链接

• 源码下载链接Baidu Drive[password:yolo]，若有改动请参考最新
• 简化后的源码下载链接Baidu Drive[password:yolo]，若有改动请参考最新
  • 只保留了src CMakeLists.txt workspace三个文件
  • scr/application/app_yolo.cpp已经完成上述更改(包括新增摄像头调用demo)
  • CMakeLists.txt已经完成修改，可直接编译
  • workspace/inference存放着需要推理的图片
  • workspace/下存放着导出好的ONNX文件
• JetPack镜像下载链接Baidu Drive[password:nano]【更新完毕!!!】，上传忒慢
• protobuf安装包下载链接Baidu Drive[password:yolo]
• 权重文件下载链接Baidu Drive[password:yolo]
  • weights文件夹下共包含yolov5s.pt yolov5s.onnx yolov5s-origin.onnx三个文件
  • yolov5s.pt是使用yolov5s-6.0训练的口罩识别模型，共一个类别
  • yolov5s.onnx是修改yolov5s-6.0代码后导出的口罩识别ONNX模型，具体修改请参考here
  • yolov5s-origin.onnx是未修改yolov5s-6.0代码直接导出的口罩识别ONNX模型

7. 参考

• tensorRT_Pro
• tensorRT_Pro的README.md
• Jetson嵌入式系列模型部署-1
• Jetson嵌入式系列模型部署-2
• Jetson NANO是什么？如何选？
• Jetson nano安装jtop
• Jetson nano安装pip并换源
• protobuf简介
• Linux下编译protobuf
• Linux下添加protobuf环境变量
• Netron可视化工具
• 手写AI的详解TensorRT的C++/Python高性能部署视频 博主强烈推荐!!!
• make -j参数简介

感谢各位看到最后，创作不易，读后有收获的看官请帮忙点个⭐️