Ubuntu20.04部署YOLOv5

前言

最近折腾了下双系统的安装为后续学习做准备，之前都是介绍的关于 Jetson 嵌入式上的模型部署工作，现在来尝试下 Ubuntu20.04 主机的部署工作，本次使用的 repo 是 tensorRT_Pro，部署的模型是 yolov5s.pt，部署流程其实和 Jetson 嵌入式一样(无非是 pytorch 训练 .pt => 转 onnx => tensorRT 生成 engine => 推理部署)，难点在于环境的安装和配置，博主使用的软件环境如下：

系统：Ubuntu20.04

软件环境：Diver Version：510.108.03；CUDA Version：11.6；cuDNN Version：8.4.0；tensorRT Version：8.4.1；OpenCV：4.6.0；protobuf：3.11.4

本篇博客主要分享环境配置、YOLOv5模型训练、YOLOv5模型部署三方面的内容。若有问题欢迎各位看官批评指正！！！

一、环境配置

环境配置可以查看 Ubuntu20.04软件安装大全

1 显卡驱动安装

参考自：拯救者Y9000K2021H在ubuntu18.04安装显卡驱动、Ubuntu卸载Nvidia驱动和安装最新驱动

描述：显卡驱动安装主要参考文章1，流程按照上面走没有问题，博主显卡为 RTX3060，切记在安装显卡驱动的时候要在 BIOS 中将安全启动关闭

1.1 卸载显卡驱动

进行该步目的是防止系统中已经存在旧版本的驱动，或者存在没有完全安装成功的驱动。

在终端输入如下指令：

sudo apt-get purge nvidia*

1.2 准备工作

点击系统中的设置(Settings)，点击最下面的关于(About)，点击 Software Updates

• 在第一栏 Ubuntu Software 页面中，前四个 main，universe，restricted，multiverse 都打上勾，找到 Download from，选择=>Others=>China=>mirrors.aliyun.com，点击确认，输入自己的密码，完成。(该步主要是进行 apt 换源)
• 在第二栏 Other Software 页面中，将 Canonical Partners 勾选上
• 在第三栏 Updates 页面中，将 Automatic check for updates 修改为 Never

换源成功后打开终端，执行如下语句：

cat /etc/apt/sources.list

在终端上会打印 apt 的源，可以看到原始的 Ubuntu 的源会被替换成之前选择的阿里源或者清华源

1.3 驱动安装

本次采用的方式是使用图形驱动程序 PPA 存储库完成 NVIDIA 驱动程序的快捷安装工作。

1.将 ppa 存储库添加到当前系统当中，指令如下：

sudo add-apt-repository ppa:graphics-drivers/ppa

2.更新软件列表，指令如下：

sudo apt-get update

3.查看检测到的驱动程序，指令如下，推荐下载带有 recommended 标志的驱动

sudo ubuntu-drivers devices

4.根据第 3 步检测到的驱动，选择一个合适的版本号(博主选择的是510)，安装指令如下：

sudo apt install nvidia-driver-510

执行完成之后执行命令 sudo reboot 重启

1.4 验证

重新启动之后在命令输入 nvidia-smi，如果出现对应的驱动版本和CUDA版本说明安装成功。

同时可以点击设置(Settings)，点击最下面的关于(About)，查看下 Graphics 这一栏是否已经变为 NVIDIA Corporation，如果是，代表安装全部完成。

2 CUDA安装

参考自linux安装tensorflow，cuda，cudnn安装，pytorch兼容，RTX30系列GPU兼容cuda

描述：本来想体验下 kiwi 一键安装 CUDA+cuDNN+TensorRT 的，但现在还处于内测阶段，还未公测，等 kiwi 公测吧，只能先自己来了。CUDA 的安装按照上面视频的操作来就行，没有问题，下面简单过一遍流程。切记！！！CUDA 的安装依赖于显卡驱动的安装，请务必安装完显卡驱动后再来安装 CUDA

2.1 准备工作

首先确定我们要安装的 CUDA 版本，它是根据你的驱动程序版本来安装的，打开终端输入如下指令：

cat /proc/driver/nvidia/version
# 或者
nvidia-smi

通过上面两种方法都可以获取 nvidia-driver 的版本，如下图所示，博主的驱动程序版本是 510.108.03，记住这个数字，后续会用到

我们知道了驱动版本后就需要确认 CUDA 下载的版本了，查看 CUDA 版本与驱动版本对应关系见下面的链接(查看表 3 即可)：

https://docs.nvidia.com/cuda/cuda-toolkit-release-notes/index.html

下面是表 3 的部分截图，左边一列是 CUDA 版本，后面两列分别对应 Linux 下和 Windows下 驱动版本的要求，GA = General Availability，通用版本，指软件的通用版本(可以理解为稳定版本吧)，后续我们就安装 GA 版本，那么怎么查看自己要下载的 CUDA 版本呢？看对应的驱动要求就像，比如说在下面的红色框内就是我想要下载的 CUDA 版本，即 CUDA 11.6 GA，驱动版本要求大于等于510.39.01，而我的驱动版本之前提到过是 510.108.03，显然是满足要求的。根据表中来看，我还可以选择 CUDA 11.5GA、CUDA 11.4 等 CUDA 版本，因为我的驱动版本是满足要求的。具体选择那个根据个人情况来就行。

2.2 CUDA下载

确认了 CUDA 下载的版本后我们就可以去 NVIDIA 官网进行下载了，博主需要下载的是 CUDA 11.6 GA

NVIDIA 官网 CUDA 下载链接：https://developer.nvidia.com/cuda-toolkit-archive

打开上面的链接后，会出现如下的界面，点击你想要下载的 CUDA Toolkit 就行，比如博主想要下载 CUDA 11.6 GA，点击下面红色框中的 CUDA Toolkit 11.6.0 即可。(为什么点击 11.6.0 而不是其它的 11.6 版本呢？其实大家可以从上面的版本对应图可知CUDA 11.6 Update 2对应CUDA Toolkit 11.6.2；CUDA 11.6 Update1对应CUDA Toolkit 11.6.1；CUDA 11.6 GA对应CUDA Toolkit 11.6.0)

点进去之后选择目标平台，如下图所示，博主这里选择 Linux 系统、x86_64 架构、Ubuntu、20.04 版本、runfile(local) 安装方式

选择完成后通过 wget 后面的地址就可以下载了，建议将网址复制到浏览器中下载，访问的是外网，下载慢，最好是开代理

2.3 CUDA安装

首先检查下前面下载的安装包，看下安装包名，博主下载的安装包名为 cuda_11.6.0_510.39.01_linux.run，首先看 11.6.0 是 CUDA 版本没有问题，其次看 510.39.01 是要求的驱动最小版本也没有问题，最后是 linux 即安装的系统也没有问题，检查完毕后下面我们开始安装工作‍

首先 cd 到安装包所在目录，执行如下指令：

sudo sh cuda_11.6.0_510.39.01_linux.run

输入密码后等待一段时间，出现下图，键盘按键上下移动，Enter 按键选择确认，我们选择 Continue

选择完成后跳转到如下界面，在最下面输入 accept

之后跳转到如下界面，选择你想要安装的选项

我们只考虑安装 CUDA Toolkit 11.6，其他均取消(空格键取消)，最后的界面如下，选择到 Install，确认(等待终端一段时间，这个时候终端没有任何显示，不要关闭！！！)

最后出现如下界面，表示安装成功

安装完成之后的路径：

• /usr/local/cuda-11.6：Toolkit
• /usr/local/cuda-11.6/include：头文件
• /usr/local/cuda-11.6/lib64：库文件

2.4 配置环境变量

CUDA 安装完成之后需要添加一下环境变量，输入如下指令打开配置文件：

sudo gedit ~/.bashrc

在最末尾添加如下内容(将 CUDA 版本替换成你自己的)

export PATH=/usr/local/cuda-11.6/bin${PATH:+:${PATH}}
export LD_LIBRARY_PATH=/usr/local/cuda-11.6/lib64${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}

点击保存退出，然后刷新一下

source ~/.bashrc

2.5 验证

环境变量配置完成后，在终端输入如下指令，有对应的输出代表安装成功：

nvcc --version

2.6 小结

CUDA 的安装流程其实并不复杂，难点在于两个步骤，一是通过你的显卡驱动找到对应的 CUDA 版本，二是能否将外网的 CUDA 安装包顺利下载下来(需要代理)

3 cuDNN安装

参考自linux安装tensorflow，cuda，cudnn安装，pytorch兼容，RTX30系列GPU兼容cuda

描述：cuDNN 的安装按照上面视频的操作来就行，没有问题，下面简单过一遍流程。切记！！！cuDNN 的安装依赖于 CUDA 的安装，请务必安装完 CUDA 后再来安装 cuDNN

3.1 cuDNN下载

cuDNN 的版本选择是根据 CUDA 版本来的，在前面 CUDA 的安装中，我们选择的是 CUDA 11.6.0 版本，依据此我们来安装 cuDNN

NVIDIA 官网 cuDNN 下载链接：https://developer.nvidia.com/rdp/cudnn-archive

打开上面的链接，会出现如下的界面(可能需要登陆，可以先注册一个 NVIDIA 账号，也不麻烦)

根据你的 CUDA 版本选择对应的 cuDNN 即可，CUDA 11.x 代表 CUDA11 版本的都支持，博主这里选择的是

Download cuDNN v8.4.0 (April 1st, 2022), for CUDA 11.x，如下所示，点击之后选择对应的平台安装包下载就行，如下面红色框所示，选择的是 Linux，Ubuntu，x86_64 的 Tar 安装包

注意！！！访问的是外网，下载慢，最好是开代理

3.2 cuDNN安装

首先检查下前面下载的安装包，看下安装包名，博主下载的安装包名为 cudnn-linux-x86_64-8.4.0.27_cuda11.6-archive.tar.xz，首先看 linux-x86_64 即安装的系统没有问题，其次看 cuda11.6 是依赖的 cuda 版本也没有问题，检查完毕后下面我们开始安装工作‍

首先 cd 到安装包所在目录，进行解压，执行如下指令：

tar -xf cudnn-linux-x86_64-8.4.0.27_cuda11.6-archive.tar.xz

等待解压完成，解压完成之后在目录下有一个 cudnn-linux-x86_64-8.4.0.27_cuda11.6-archive 文件夹，这个文件夹中又包含 include 和 lib 两个文件夹，分别代表 cuDNN 头文件和 cuDNN 库文件。

cuDNN 安装非常简单，就是将这两个文件中的所有内容复制到 CUDA 对应的文件夹中，我们来看下接下来的操作

首先输入如下指令，打开一个超级窗口(直接复制权限不够)：

sudo nautilus

• 在超级窗口中找到之前安装好的 CUDA 路径，在 /usr/local/cuda-11.6
• 将 cuDNN 解压包中 include 文件全部复制到 /usr/local/cuda-11.6/include中
• 将 cuDNN 解压包中 lib 文件全部复制到 /usr/local/cuda-11.6/lib64 中

至此，cuDNN 的安装完毕

3.3 小结

cuDNN 的安装流程其实并不复杂，其实都不能称之为安装，就是根据 CUDA 版本找到对应的 cuDNN 版本，下载后把其头文件和库文件复制到对应 CUDA 目录即可。难点在于两个步骤，一是通过你的 CUDA 版本找到对应的 cuDNN 版本，二是能否将外网的 cuDNN 安装包顺利下载下来(需要代理)

4 TensorRT安装

参考自Linux安装TensorRT、Ubuntu20.04安装TensorRT

描述：TensorRT 的安装按照上面视频的操作来就行，没有问题，下面简单过一遍流程。切记！！！TensorRT 安装的版本依赖于 CUDA 和 cuDNN 的版本，请先确认下 CUDA 和 cuDNN 版本

4.1 TensorRT下载

TensorRT 的版本是根据 CUDA 版本和 cuDNN 版本来的，在前面的 CUDA 的安装中，我们选择的是 CUDA 11.6.0 版本，依据此来安装 TensorRT

NVIDIA 官网 TensorRT 下载链接：https://developer.nvidia.com/nvidia-tensorrt-8x-download

打开上面的链接，会出现如下的界面(需要登陆，可以先注册一个 NVIDIA 账号，也不麻烦，可能需要开代理)

登陆完成之后，可以看到下面的界面，点击 I Agree To the Terms of the ... 即下面的红色框，就可以显示各种 TensorRT 版本(其中 GA 版本=>稳定版本 EA 版本=>测试版本)

点开之后就是对应 TensorRT 版本的各种形式的安装包(包名有提示适合的对应的 CUDA 版本)，比如博主选择的是 TensorRT 8.4 GA，选择的是 x86_64 架构，系统是 Linux 操作系统，下载形式选择压缩包形式(即TAR Package)下载，适用于 CUDA 11.6 版本，即TensorRT 8.4 GA for Linux x86_64 and CUDA 11.0, 11.1, 11.2, 11.3, 11.4, 11.5, 11.6 and 11.7 TAR Package，如下面的红色框所示。

注意！！！访问的是外网，下载慢，最好是开代理

4.2 TensorRT安装

首先检查下前面下载的安装包，看下安装包名，博主下载的安装包为 TensorRT-8.4.1.5.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz，首先看 Linux.x86_64 即安装的系统没有问题，其次看 cuda-11.6 是依赖的 cuda 版本也没有问题，最后看 cudnn8.4 是依赖的 cudnn 版本也没有问题，检查完毕后下面我们开始安装工作‍

首先 cd 到安装包所在目录，进行解压，执行如下指令：

tar zxf TensorRT-8.4.1.5.Linux.x86_64-gnu.cuda-11.6.cudnn8.4.tar.gz

等待解压完成。解压完成之后在目录下有一个 TensorRT-8.4.1.5 的文件夹

可以移动该文件夹存放在自己想放的目录下(也可以重命名)，如：

mv TensorRT-8.4.1.5 /opt

至此，TensorRT 安装完成

4.3 配置环境变量

TensorRT 安装完成之后需要添加一下环境变量，输入如下指令打开配置文件：

sudo gedit ~/.bashrc

在最末尾添加如下内容(将 TensorRT 版本替换成你自己的)

export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/opt/TensorRT-8.4.1.5/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
export PATH=/opt/TensorRT-8.4.1.5/bin${PATH:+:${PATH}}

点击保存退出，然后刷新一下

source ~/.bashrc

4.4 验证

测试代码：/opt/TensorRT-8.4.1.5/samples/sampleMNIST

在安装的 tensorRT 目录下 samples 文件夹下用 sampleMNIST 示例测试

首先进入该文件夹

cd /opt/TensorRT-8.4.1.5/samples/sampleMNIST

编译

make

如下图所示

编译成功后执行如下指令会有对应的结果输出：

./../../bin/sample_mnist

4.5 小结

TensorRT 的安装流程其实并不复杂，其实都不能称之为安装，就是根据 CUDA 版本找到对应的 TensorRT 版本，下载后添加下环境变量即可。难点在于两个步骤，一是通过你的 CUDA 版本找到对应的 TensorRT 版本，二是否能将外网的 TensorRT 安装包顺利下载下来(需要代理)

5 编译OpenCV-4.6.0

参考自Ubuntu 20.04搭建OpenCV 4.5.0 & C++环境、Ubuntu 安装 OpenCV 4.5.5、关于安装opencv编译报错：runtime library * hidden by 等若干问题、cmake编译出错hidden by files

描述：注意！！！如果你在编译 opencv 之前安装了 annaconda，请先看问题解决(与 Anaconda 冲突)，除了上述冲突问题之外，没有其他问题，按照上文的操作即可。博主主要参考上文1中的 opencv 安装(即前 3 项)以及上文2中的 opencv 环境配置和验证。下面简单过一遍整个编译流程

5.1 安装依赖项

1.安装 g++，cmake，make，wget，unzip，若已安装，此步跳过

sudo apt install g++
sudo apt install cmake
sudo apt install make
sudo apt install wget unzip

2.安装 opencv 依赖的库

sudo apt-get install build-essential libgtk2.0-dev libgtk-3-dev libavcodec-dev libavformat-dev libjpeg-dev libswscale-dev libtiff5-dev

3.安装一些可选的库

# python3支持（首次安装了python的库，但make报错了，之后删了这两个库，若不使用python，建议不安装）
sudo apt install python3-dev python3-numpy
# streamer支持
sudo apt install libgstreamer-plugins-base1.0-dev libgstreamer1.0-dev
# 可选的依赖
sudo apt install libpng-dev libopenexr-dev libtiff-dev libwebp-dev

5.2 下载OpenCV 4.6.0源文件

可以在官网下载相应版本的 OpenCV，主要有 Source 和 GitHub 两种方式下载

1.Source：https://opencv.org/releases/

点击 Source 进行下载，外网访问慢，可在 Windows 下通过代理下载后传到 Linux；也可通过我提供的下载链接 Baidu Driver[pwd:yolo] 进行下载

2.Github下载方式：

# 安装4.6.0版本
wget -O opencv.zip https://github.com/opencv/opencv/archive/4.6.0.zip
# 安装最新版本
wget -O opencv.zip https://github.com/opencv/opencv/archive/master.zip

下载完成后，解压到主目录，将解压后的文件重命名为 opencv

5.3 Cmake配置和编译OpenCV

1.在主目录下打开终端

2.进入到下载好的 opencv 目录，新建并进入目录 build

cd opencv
mkdir build
cd build

3.使用 cmake 配置 opencv

cmake -D CMAKE_BUILD_TYPE=Release -D OPENCV_GENERATE_PKGCONFIG=YES ..

4.用 make 进行编译

make -j24

5.用 make 安装

sudo make install

默认安装路径如下：

• /usr/local/bin - executable files
• /usr/local/bin - libraries (.so)
• /usr/local/cmake/opencv4 - cmake package
• /usr/local/include/opencv4 - headers
• /usr/local/share/opencv4 - other files (e.g. trained cascades in XML format)

5.4 环境配置

动态库配置

sudo gedit /etc/ld.so.conf.d/opencv.conf

添加

/usr/local/lib

保存后，终端内执行

sudo ldconfig 

更新PKG_CONFIG_PATH

sudo gedit /etc/bash.bashrc

文末添加

PKG_CONFIG_PATH=$PKG_CONFIG_PATH:/usr/local/lib/pkgconfig
export PKG_CONFIG_PATH

保存后，终端执行

source /etc/bash.bashrc
sudo updatedb

若出现sudo: updatedb: command not found，执行 sudo apt install mlocate 后再次操作即可。

5.5 验证

pkg-config

在终端输入如下指令，会显示对应的 opencv 版本 4.6.0

pkg-config --modversion opencv4

程序验证，代码如下

test.cpp

#include 
#include 
#include 
#include 

using namespace cv;
using namespace std;

/**
 * 读写与显示
 * @return
 */
int main() {
    cout << CV_VERSION << endl;
    Mat src = cv::imread("test.png", IMREAD_GRAYSCALE);
    if (src.empty()) {
        cout << "Could not load image\n" << endl;
        return -1;
    }
    namedWindow("gray", WINDOW_AUTOSIZE);
    imshow("gray", src);
    waitKey(0);

    return 0;
}

CMakeLists.txt

cmake_minimum_required(VERSION 3.16)
project(test)
find_package(OpenCV REQUIRED)
add_executable(test test.cpp)
target_link_libraries(test ${OpenCV_LIBS})

编译运行

cmake .
make
./test

运行后屏幕会显示目录下 test.png 的灰度图像

6 编译Protobuf-3.11.4

参考自Linux下编译protobuf、Linux下添加protobuf的环境变量

描述：由于要经常使用到 tensorRT_Pro 这个 repo，因此需要编译特定版本的 protobuf-3.11.4，编译流程看我下面的即可

6.1 安装

首先获取压缩包，通过我提供的软件安装包下载链接 Baidu Driver[pwd:yolo]，里面就有 protobuf-3.11.4

解压

先 cd 到安装包路径，然后解压，指令如下：

unzip protobuf-3.11.4.zip

编译

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

耐心等待编译完成(需要一段时间)

创建安装目录

我们先要创建一个文件用于存放安装后的 protobuf 的头文件和库文件，我们选择在 /home 目录下创建一个 protobuf 文件，指令如下：

mkdir protobuf

安装

安装的位置选择上面创建的 protobuf 路径即可，指令如下：

make install DESTDIR=/home/jarvis/protobuf

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

6.2 环境变量的配置

首先打开配置文件，指令如下：

sudo vim /etc/profile

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

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

source 生效

source /etc/profile

配置动态路径

sudo vim /etc/ld.so.conf

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

/home/jarvis/protobuf/lib

验证

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

二、YOLOv5模型训练

1. 项目的克隆和必要的环境依赖

1.1 项目的克隆

yolov5 的代码是开源的可直接从 github 官网上下载，源码下载地址是 https://github.com/ultralytics/yolov5，由于 yolov5 版本较多，本次采用 yolov5-v7.0 分支进行模型的训练和部署工作。Linux 下代码克隆指令如下：

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

1.2 环境安装

关于 Linux 下深度学习的环境安装可参考Ubuntu20.04软件安装大全

2. 数据集和预训练权重的准备

2.1 数据集

训练采用的 VOC 数据集，这里给出下载链接 Baidu Drive[pwd:yolo]，本次训练并没有用到所有的数据，博主将 train2007 和 val2007 作为训练集，将 test2007 作为验证集，整个数据集文件夹内容如下图所示：

其中，images 存放的内容和图片文件，labels 存放的内容是 YOLO 格式的 .txt 标签文件，所有文件都可以从我分享的链接下载，大家可以按照上述方式将数据集进行整合。

2.2 预训练权重准备

yolov5-7.0 预训练权重可以通过下载链接 Baidu Drive[pwd:yolo] 进行下载，将预训练权重放在 yolov5-7.0 主目录下，本次训练 VOC 数据集使用的预训练权重是 yolov5s.pt

3. 训练模型

代码和数据集准备好后就可以进行训练了，训练修改的文件主要是两个即 VOC.yaml 用于指定数据集的相关路径，yolov5s.yaml 用于指定训练的网络结构

VOC.yaml 位于 yolov5-7.0/data 下，其具体内容如下：

• 首先 path 路径指定为上面整合的数据集的绝对路径，路径中最好不要含有中文
• train、val、test 的内容就是 VOC 数据集下用于训练、验证以及测试的图片
• names 不用修改
• download 内容全部删除即可

# YOLOv5  by Ultralytics, GPL-3.0 license
# PASCAL VOC dataset http://host.robots.ox.ac.uk/pascal/VOC by University of Oxford
# Example usage: python train.py --data VOC.yaml
# parent
# ├── yolov5
# └── datasets
#     └── VOC  ← downloads here (2.8 GB)


# Train/val/test sets as 1) dir: path/to/imgs, 2) file: path/to/imgs.txt, or 3) list: [path/to/imgs1, path/to/imgs2, ..]
path: /home/jarvis/Learn/Datasets
train: # train images (relative to 'path')  16551 images
  - images/train2007
  - images/val2007
val: # val images (relative to 'path')  4952 images
  - images/test2007
test: # test images (optional)
  - images/test2007

# Classes
names:
  0: aeroplane
  1: bicycle
  2: bird
  3: boat
  4: bottle
  5: bus
  6: car
  7: cat
  8: chair
  9: cow
  10: diningtable
  11: dog
  12: horse
  13: motorbike
  14: person
  15: pottedplant
  16: sheep
  17: sofa
  18: train
  19: tvmonitor

yolov5s.yaml 位于 yolov5-7.0/models 下，将 nc：80 修改为 VOC 的 20 个类别，即 nc：20

进入到 yolov5-7.0 目录，在终端执行如下指令即可开始训练：

python train.py --weights=./yolov5s.pt --cfg=./models/yolov5s.yaml --data=./data/VOC.yaml --epochs=100 --batch-size=32

博主训练的模型是 yolov5s.pt 且使用的是单个 GPU 进行训练，显卡为 RTX3060，操作系统为 Ubuntu20.04，pytorch 版本为 1.12.0，训练时长大概 1.7 小时

训练参数的指定简要解释如下：

• –weights 预训练权重
• –cfg 模型配置文件路径
• –data 数据配置文件路径
• –epochs 训练轮数
• –batch-size 每次输入到网络的图片数

还有其它参数博主并未设置，如 --workers 最大工作核心数等。大家一定要根据自己的实际情况(如显卡算力)指定不同的参数，如果你之前训练过自己的数据集，那我相信这对你来说应该是小 case

训练完成后的模型权重保存在 runs/train/exp/weights 文件夹下，我们选用 best.pt 进行后续模型部署即可，这里提供博主训练好的权重文件下载链接 Baidu Driver[pwd:yolo]

4. 测试

利用项目中的 detect.py 文件进行测试，将需要推理的图片放入 执行指令如下：

python detect.py --weights=./best.pt

也可进行视频或者摄像头推理，执行指令如下，0 代表本地摄像头

python detect.py --source=0  --weights=./best.pt

至此，YOLOv5模型训练已经完毕，下面开始模型部署工作。

三、YOLOv5模型部署

Ubuntu20.04 主机上 yolov5 模型部署流程和 Jetosn 嵌入式基本一致，大家可以参考我之前发的模型部署文章，部署使用的 Github 仓库是 tensorRT_Pro。该 repo 通过 tensorRT 的 ONNX parser 解析 ONNX 文件来完成模型的构建工作。对模型部署有疑问的可以参考 Jetson嵌入式系列模型部署-1，想了解通过 tensorRT 的 Layer API 一层层完成模型的搭建工作可参考 Jetson嵌入式系列模型部署-2，想了解通过 tensorRT 的 trtexec 工具构建模型的可参考 Jetson nano部署YOLOv8。本文主要是针对 tensorRT_Pro 项目中的 yolov5 完成模型部署，本文参考自 tensorRT_Pro的README.md，具体操作流程作者描述非常详细，下面再简单过一遍，本次训练的模型使用 yolov5s.pt，数据集为 VOC，类别数为20.

1. 源码下载

使用如下指令

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

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

2. 环境配置

需要使用的软件环境有 tensorRT、CUDA、cuDNN、OpenCV、Protobuf。之前已经详细介绍了各软件的安装情况，这里不再赘述。编译可采用 CMakeLists.txt 和 Makefile 两种方式，二者选一即可。

2.1 配置CMakeLists.txt

主要修改五处

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

set(HAS_PYTHON OFF)

2.修改第 18 行，修改 OpenCV 路径

set(OpenCV_DIR   "/usr/local/include/opencv4/")

3.修改第 20 行，修改 CUDA 路径

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

4.修改第 22 行，修改 tensorRT 路径

set(TENSORRT_DIR "/opt/TensorRT-8.4.1.5")

5.修改第 33 行，修改 protobuf 路径

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

完整的 CMakeLists.txt 的内容如下：

cmake_minimum_required(VERSION 2.6)
project(pro)

option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)
set(EXECUTABLE_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/workspace)

# 如果要支持python则设置python路径
set(HAS_PYTHON OFF)                                         # ===== 修改 1 =====
set(PythonRoot "/datav/software/anaconda3")
set(PythonName "python3.9")

# 如果你是不同显卡，请设置为显卡对应的号码参考这里：https://developer.nvidia.com/zh-cn/cuda-gpus#compute
#set(CUDA_GEN_CODE "-gencode=arch=compute_75,code=sm_75")

# 如果你的opencv找不到，可以自己指定目录
set(OpenCV_DIR   "/usr/local/include/opencv4/")             # ===== 修改 2 =====


set(CUDA_TOOLKIT_ROOT_DIR     "/usr/local/cuda-11.6")       # ===== 修改 3 =====
set(CUDNN_DIR    "/datav/lean/cudnn8.2.4.15-cuda11.4")
set(TENSORRT_DIR "/opt/TensorRT-8.4.1.5")                   # ===== 修改 4 =====

# set(CUDA_TOOLKIT_ROOT_DIR     "/data/sxai/lean/cuda-10.2")
# set(CUDNN_DIR    "/data/sxai/lean/cudnn7.6.5.32-cuda10.2")
# set(TENSORRT_DIR "/data/sxai/lean/TensorRT-7.0.0.11")

# set(CUDA_TOOLKIT_ROOT_DIR  "/data/sxai/lean/cuda-11.1")
# set(CUDNN_DIR    "/data/sxai/lean/cudnn8.2.2.26")
# set(TENSORRT_DIR "/data/sxai/lean/TensorRT-7.2.1.6")

# 因为protobuf，需要用特定版本，所以这里指定路径
set(PROTOBUF_DIR "/home/jarvis/protobuf")                   # ===== 修改 5 ======


find_package(CUDA REQUIRED)
find_package(OpenCV)

include_directories(
    ${PROJECT_SOURCE_DIR}/src
    ${PROJECT_SOURCE_DIR}/src/application
    ${PROJECT_SOURCE_DIR}/src/tensorRT
    ${PROJECT_SOURCE_DIR}/src/tensorRT/common
    ${OpenCV_INCLUDE_DIRS}
    ${CUDA_TOOLKIT_ROOT_DIR}/include
    ${PROTOBUF_DIR}/include
    ${TENSORRT_DIR}/include
    ${CUDNN_DIR}/include
)

# 切记，protobuf的lib目录一定要比tensorRT目录前面，因为tensorRTlib下带有protobuf的so文件
# 这可能带来错误
link_directories(
    ${PROTOBUF_DIR}/lib
    ${TENSORRT_DIR}/lib
    ${CUDA_TOOLKIT_ROOT_DIR}/lib64
    ${CUDNN_DIR}/lib
)

if("${HAS_PYTHON}" STREQUAL "ON")
    message("Usage Python ${PythonRoot}")
    include_directories(${PythonRoot}/include/${PythonName})
    link_directories(${PythonRoot}/lib)
    set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -DHAS_PYTHON")
endif()

set(CMAKE_CXX_FLAGS  "${CMAKE_CXX_FLAGS} -std=c++11 -Wall -O0 -Wfatal-errors -pthread -w -g")
set(CUDA_NVCC_FLAGS "${CUDA_NVCC_FLAGS} -std=c++11 -O0 -Xcompiler -fPIC -g -w ${CUDA_GEN_CODE}")
file(GLOB_RECURSE cpp_srcs ${PROJECT_SOURCE_DIR}/src/*.cpp)
file(GLOB_RECURSE cuda_srcs ${PROJECT_SOURCE_DIR}/src/*.cu)
cuda_add_library(plugin_list SHARED ${cuda_srcs})
target_link_libraries(plugin_list nvinfer nvinfer_plugin)
target_link_libraries(plugin_list cuda cublas cudart cudnn)
target_link_libraries(plugin_list protobuf pthread)
target_link_libraries(plugin_list ${OpenCV_LIBS})

add_executable(pro ${cpp_srcs})

# 如果提示插件找不到，请使用dlopen(xxx.so, NOW)的方式手动加载可以解决插件找不到问题
target_link_libraries(pro nvinfer nvinfer_plugin)
target_link_libraries(pro cuda cublas cudart cudnn)
target_link_libraries(pro protobuf pthread plugin_list)
target_link_libraries(pro ${OpenCV_LIBS})

if("${HAS_PYTHON}" STREQUAL "ON")
    set(LIBRARY_OUTPUT_PATH ${PROJECT_SOURCE_DIR}/example-python/pytrt)
    add_library(pytrtc SHARED ${cpp_srcs})
    target_link_libraries(pytrtc nvinfer nvinfer_plugin)
    target_link_libraries(pytrtc cuda cublas cudart cudnn)
    target_link_libraries(pytrtc protobuf pthread plugin_list)
    target_link_libraries(pytrtc ${OpenCV_LIBS})
    target_link_libraries(pytrtc "${PythonName}")
    target_link_libraries(pro "${PythonName}")
endif()

add_custom_target(
    yolo
    DEPENDS pro
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
    COMMAND ./pro yolo
)

add_custom_target(
    yolo_gpuptr
    DEPENDS pro
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
    COMMAND ./pro yolo_gpuptr
)

add_custom_target(
    yolo_fast
    DEPENDS pro
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
    COMMAND ./pro yolo_fast
)

add_custom_target(
    centernet
    DEPENDS pro
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
    COMMAND ./pro centernet
)

add_custom_target(
    alphapose 
    DEPENDS pro
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
    COMMAND ./pro alphapose
)

add_custom_target(
    retinaface
    DEPENDS pro
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
    COMMAND ./pro retinaface
)

add_custom_target(
    dbface
    DEPENDS pro
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
    COMMAND ./pro dbface
)

add_custom_target(
    arcface 
    DEPENDS pro
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
    COMMAND ./pro arcface
)

add_custom_target(
    bert 
    DEPENDS pro
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
    COMMAND ./pro bert
)

add_custom_target(
    fall
    DEPENDS pro
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
    COMMAND ./pro fall_recognize
)

add_custom_target(
    scrfd
    DEPENDS pro
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
    COMMAND ./pro scrfd
)

add_custom_target(
    lesson
    DEPENDS pro
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/workspace
    COMMAND ./pro lesson
)

add_custom_target(
    pyscrfd
    DEPENDS pytrtc
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/example-python
    COMMAND python test_scrfd.py
)

add_custom_target(
    pyinstall
    DEPENDS pytrtc
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/example-python
    COMMAND python setup.py install
)

add_custom_target(
    pytorch
    DEPENDS pytrtc
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/example-python
    COMMAND python test_torch.py
)

add_custom_target(
    pyyolov5
    DEPENDS pytrtc
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/example-python
    COMMAND python test_yolov5.py
)

add_custom_target(
    pycenternet
    DEPENDS pytrtc
    WORKING_DIRECTORY ${PROJECT_SOURCE_DIR}/example-python
    COMMAND python test_centernet.py
)

2.2 配置Makefile

主要修改五处

1.修改第 4 行，修改 protobuf 路径

lean_protobuf  := /home/jarvis/protobuf

2.修改第 5 行，修改 tensorRT 路径

lean_tensor_rt := /opt/TensorRT-8.4.1.5

3.修改第 7 行，修改 OpenCV 路径

lean_opencv    := /usr/local

4.修改第 8 行，修改 CUDA 路径

lean_cuda      := /usr/local/cuda-11.6

5.修改第 9 行，选择不支持 python (也可选择支持)

use_python     := false

完整的 Makefile 的内容如下：

cc        := g++
nvcc      = ${lean_cuda}/bin/nvcc

lean_protobuf  := /home/jarvis/protobuf		# ===== 修改 1 =====
lean_tensor_rt := /opt/TensorRT-8.4.1.5		# ===== 修改 2 =====
lean_cudnn     := /datav/lean/cudnn8.2.4.15-cuda11.4
lean_opencv    := /usr/local				# ===== 修改 3 =====
lean_cuda      := /usr/local/cuda-11.6		# ===== 修改 4 =====
use_python     := false						# ===== 修改 5 =====
python_root    := /datav/software/anaconda3

# python_root指向的lib目录下有个libpython3.9.so，因此这里写python3.9
# 对于有些版本，so名字是libpython3.7m.so，你需要填写python3.7m
# /datav/software/anaconda3/lib/libpython3.9.so
python_name    := python3.9

# 如果是其他显卡，请修改-gencode=arch=compute_75,code=sm_75为对应显卡的能力
# 显卡对应的号码参考这里：https://developer.nvidia.com/zh-cn/cuda-gpus#compute
cuda_arch := # -gencode=arch=compute_75,code=sm_75

cpp_srcs  := $(shell find src -name "*.cpp")
cpp_objs  := $(cpp_srcs:.cpp=.cpp.o)
cpp_objs  := $(cpp_objs:src/%=objs/%)
cpp_mk    := $(cpp_objs:.cpp.o=.cpp.mk)

cu_srcs  := $(shell find src -name "*.cu")
cu_objs  := $(cu_srcs:.cu=.cu.o)
cu_objs  := $(cu_objs:src/%=objs/%)
cu_mk    := $(cu_objs:.cu.o=.cu.mk)

include_paths := src        \
			src/application \
			src/tensorRT	\
			src/tensorRT/common  \
			$(lean_protobuf)/include \
			$(lean_opencv)/include/opencv4 \
			$(lean_tensor_rt)/include \
			$(lean_cuda)/include  \
			$(lean_cudnn)/include 

library_paths := $(lean_protobuf)/lib \
			$(lean_opencv)/lib    \
			$(lean_tensor_rt)/lib \
			$(lean_cuda)/lib64  \
			$(lean_cudnn)/lib

link_librarys := opencv_core opencv_imgproc opencv_videoio opencv_imgcodecs \
			nvinfer nvinfer_plugin \
			cuda cublas cudart cudnn \
			stdc++ protobuf dl


# HAS_PYTHON表示是否编译python支持
support_define    := 

ifeq ($(use_python), true) 
include_paths  += $(python_root)/include/$(python_name)
library_paths  += $(python_root)/lib
link_librarys  += $(python_name)
support_define += -DHAS_PYTHON
endif

empty         :=
export_path   := $(subst $(empty) $(empty),:,$(library_paths))

run_paths     := $(foreach item,$(library_paths),-Wl,-rpath=$(item))
include_paths := $(foreach item,$(include_paths),-I$(item))
library_paths := $(foreach item,$(library_paths),-L$(item))
link_librarys := $(foreach item,$(link_librarys),-l$(item))

cpp_compile_flags := -std=c++11 -g -w -O0 -fPIC -pthread -fopenmp $(support_define)
cu_compile_flags  := -std=c++11 -g -w -O0 -Xcompiler "$(cpp_compile_flags)" $(cuda_arch) $(support_define)
link_flags        := -pthread -fopenmp -Wl,-rpath='$$ORIGIN'

cpp_compile_flags += $(include_paths)
cu_compile_flags  += $(include_paths)
link_flags        += $(library_paths) $(link_librarys) $(run_paths)

ifneq ($(MAKECMDGOALS), clean)
-include $(cpp_mk) $(cu_mk)
endif

pro    : workspace/pro
pytrtc : example-python/pytrt/libpytrtc.so
expath : library_path.txt

library_path.txt : 
	@echo LD_LIBRARY_PATH=$(export_path):"$$"LD_LIBRARY_PATH > $@

workspace/pro : $(cpp_objs) $(cu_objs)
	@echo Link $@
	@mkdir -p $(dir $@)
	@$(cc) $^ -o $@ $(link_flags)

example-python/pytrt/libpytrtc.so : $(cpp_objs) $(cu_objs)
	@echo Link $@
	@mkdir -p $(dir $@)
	@$(cc) -shared $^ -o $@ $(link_flags)

objs/%.cpp.o : src/%.cpp
	@echo Compile CXX $<
	@mkdir -p $(dir $@)
	@$(cc) -c $< -o $@ $(cpp_compile_flags)

objs/%.cu.o : src/%.cu
	@echo Compile CUDA $<
	@mkdir -p $(dir $@)
	@$(nvcc) -c $< -o $@ $(cu_compile_flags)

objs/%.cpp.mk : src/%.cpp
	@echo Compile depends CXX $<
	@mkdir -p $(dir $@)
	@$(cc) -M $< -MF $@ -MT $(@:.cpp.mk=.cpp.o) $(cpp_compile_flags)
	
objs/%.cu.mk : src/%.cu
	@echo Compile depends CUDA $<
	@mkdir -p $(dir $@)
	@$(nvcc) -M $< -MF $@ -MT $(@:.cu.mk=.cu.o) $(cu_compile_flags)

yolo : workspace/pro
	@cd workspace && ./pro yolo

yolo_gpuptr : workspace/pro
	@cd workspace && ./pro yolo_gpuptr

dyolo : workspace/pro
	@cd workspace && ./pro dyolo

dunet : workspace/pro
	@cd workspace && ./pro dunet

dmae : workspace/pro
	@cd workspace && ./pro dmae

dclassifier : workspace/pro
	@cd workspace && ./pro dclassifier

yolo_fast : workspace/pro
	@cd workspace && ./pro yolo_fast

bert : workspace/pro
	@cd workspace && ./pro bert

alphapose : workspace/pro
	@cd workspace && ./pro alphapose

fall : workspace/pro
	@cd workspace && ./pro fall_recognize

retinaface : workspace/pro
	@cd workspace && ./pro retinaface

arcface    : workspace/pro
	@cd workspace && ./pro arcface

test_warpaffine    : workspace/pro
	@cd workspace && ./pro test_warpaffine

test_yolo_map    : workspace/pro
	@cd workspace && ./pro test_yolo_map

arcface_video    : workspace/pro
	@cd workspace && ./pro arcface_video

arcface_tracker    : workspace/pro
	@cd workspace && ./pro arcface_tracker

test_all : workspace/pro
	@cd workspace && ./pro test_all

scrfd : workspace/pro
	@cd workspace && ./pro scrfd

centernet : workspace/pro
	@cd workspace && ./pro centernet

dbface : workspace/pro
	@cd workspace && ./pro dbface

high_perf : workspace/pro
	@cd workspace && ./pro high_perf

lesson : workspace/pro
	@cd workspace && ./pro lesson

plugin : workspace/pro
	@cd workspace && ./pro plugin

pytorch : pytrtc
	@cd example-python && python test_torch.py

pyscrfd : pytrtc
	@cd example-python && python test_scrfd.py

pyretinaface : pytrtc
	@cd example-python && python test_retinaface.py

pycenternet : pytrtc
	@cd example-python && python test_centernet.py

pyyolov5 : pytrtc
	@cd example-python && python test_yolov5.py

pyyolov7 : pytrtc
	@cd example-python && python test_yolov7.py

pyyolox : pytrtc
	@cd example-python && python test_yolox.py

pyarcface : pytrtc
	@cd example-python && python test_arcface.py

pyinstall : pytrtc
	@cd example-python && python setup.py install

clean :
	@rm -rf objs workspace/pro example-python/pytrt/libpytrtc.so example-python/build example-python/dist example-python/pytrt.egg-info example-python/pytrt/__pycache__
	@rm -rf workspace/single_inference
	@rm -rf workspace/scrfd_result workspace/retinaface_result
	@rm -rf workspace/YoloV5_result workspace/YoloX_result
	@rm -rf workspace/face/library_draw workspace/face/result
	@rm -rf build
	@rm -rf example-python/pytrt/libplugin_list.so
	@rm -rf library_path.txt

.PHONY : clean yolo alphapose fall debug

# 导出符号，使得运行时能够链接上
export LD_LIBRARY_PATH:=$(export_path):$(LD_LIBRARY_PATH)

3. ONNX导出

• 训练的模型使用 yolov8s.pt，torch 版本1.12.0，onnx 版本1.13.1

• ONNX 导出参考自 YoloV5案例第一部分，导出ONNX

关于静态 batch 和动态 batch 有以下几点说明，更多细节请查看视频

静态batch

• 导出的 onnx 指定所有维度均为明确的数字，是静态 shape 模型
• 在推理的时候，它永远都是同样的 batch 推理，即使你目前只有一个图推理，它也需要 n 和 batch 的耗时
• 适用于大部分场景，整个代码逻辑非常简单

动态batch

• 导出的时候指定特定维度为 dynamic，也就是不确定状态
• 模型推理时才决定所需推理的 batch 大小，耗时最优，但 onnx 复杂度提高了
• 适用于如 server 有大量不均匀的请求时的场景

3.1 静态batch导出

静态 batch 的 ONNX 模型导出指令如下

cd yolov5-7.0
python export.py --weights=./best.pt --include=onnx --opset=11

3.2 动态batch导出

修改代码，保证动态 batchsize，主要修改两个文件的内容

• yolov5-7.0/models/yolo.py
• yolov5-7.0/export.py

# yolov5-7.0/models/yolo.py第60行，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-7.0/export.py第141行
# output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output0']
# if dynamic:
#     dynamic = {'images': {0: 'batch', 2: 'height', 3: 'width'}}  # shape(1,3,640,640)
#     if isinstance(model, SegmentationModel):
#         dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
#         dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'}  # shape(1,32,160,160)
#         elif isinstance(model, DetectionModel):
#             dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
output_names = ['output0', 'output1'] if isinstance(model, SegmentationModel) else ['output']            
if dynamic:
    dynamic = {'images': {0: 'batch'}}  # shape(1,3,640,640)
    if isinstance(model, SegmentationModel):
        dynamic['output0'] = {0: 'batch', 1: 'anchors'}  # shape(1,25200,85)
        dynamic['output1'] = {0: 'batch', 2: 'mask_height', 3: 'mask_width'}  # shape(1,32,160,160)
        elif isinstance(model, DetectionModel):
            dynamic['output'] = {0: 'batch'}  # shape(1,25200,85)

动态 batch 的 ONNX 模型导出指令如下

cd yolov5-7.0
python export.py --weights=./best.pt --dynamic --include=onnx --opset=11

3.3 Resize节点解析的问题

先剧透下，当构建 engine 时会出现错误，我们一并解决，到时候可直接生成 engine，错误信息如下图所示，大概意思就是说 Resize_143 这个节点的 scales 没有初始化(应该是这样理解的吧)

更多细节的描述可查看Jetson nano部署YOLOv8，这里只提供解决方案即通过 onnxsim 来进行优化。新建一个 v5onnxsim.py 文件，用于优化 onnx 文件，具体内容如下：

import onnx
from onnxsim import simplify

onnx_model = onnx.load("best.onnx")
model_simp, check = simplify(onnx_model)
assert check, "Simplified ONNX model could not be Validated"
onnx.save(model_simp, "best.sim.onnx")

运行后会在当前文件夹生成一个 best.sim.onnx 模型，现在可以查看对应的 Resize_143 节点发生了变化。

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高性能部署视频，这些做法的必要性体现在，简化过程的复杂度，去掉 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，修改 yolov5s 为 best.sim，构建的模型为 best.sim.onnx
• 3.app_yolo.cpp 100行 cocolabels 修改为 voclabels
• 4.app_yolo.cpp 25行 新增 voclabels 数组，添加 VOC 的类别名称

具体修改如下：

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

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

static const char *voclabels[] = {"aeroplane",   "bicycle", "bird",   "boat",       "bottle",
                                  "bus",         "car",     "cat",    "chair",      "cow",
                                  "diningtable", "dog",     "horse",  "motorbike",  "person",
                                  "pottedplant",  "sheep",  "sofa",   "train",      "tvmonitor"};		 //修改4 25行新增代码，为自训练模型的类别名称

4.2 编译

OK！源码修改好了，编译文件也搞定了，可以编译运行了，编译方式有两种，根据自己喜好更改即可

1.采用 CMakeLists.txt 文件编译，指令如下：

cd tensorRT_Pro-main
mkdir build && cd build
cmake .. && make -j24

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

2.采用 Makefile 文件编译，指令如下：

cd tensorRT_Pro-main
make pro

4.3 模型构建和推理

编译完成后的可执行文件 .pro 存放在 workspace 文件夹下，模型推理指令如下：

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

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

模型推理效果如下图所示：

4.4 拓展-摄像头检测

简单写了一个摄像头检测的 demo，默认打开的是 USB 摄像头，如果有其它需求如 CSI 摄像头或者网络摄像头，大家自行修改即可，主要修改以下几点：

• 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 = voclabels[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("best.sim.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 文件夹下运行即可调用摄像头进行检测，指令如下(以 CMakeLists.txt 编译为例)：

cd build
make -j24
cd ../workspace
./pro yolo

图解如下所示：

5. CUDA-Tips

在 VSCode 中，要让编辑器正确识别 CUDA 代码并进行语法高亮等操作，需要安装相应的插件。可以在 VSCode 的插件商店中搜索 CUDA，安装名为 vscode-cudacpp 的插件，即可以在编辑器中获得对 CUDA 代码的支持。(from chatGPT)

安装完成后，如果打开的 CUDA 文件仍然存在语法高亮等问题，可以尝试执行以下操作：

1. 点击 VSCode 编辑器底部的选择语言模式按钮，选择CUDA C/C++作为语言模式

2. ctrl+shift+p快捷键选择Open Workspace Settings(JSON)/打开工作区设置(JSON)

3. 在打开的设置页面中，搜索files.associations设置，找到以 .cu 为键的配置项，确保其值为cuda-cpp(即将 .cu 文件与 CUDA C++ 关联)。如果该配置项不存在，可以手动添加以下配置：

   "files.associations": {
       "*.cu": "cuda-cpp"
   }
   

4. 修改完成后，保存配置文件，关闭并重新打开 VSCode，再次打开 CUDA 文件时就可以看到语法高亮等功能生效了。

结语

本篇博客介绍了 Ubuntu20.04 主机部署 YOLOv5 模型，其难点在于环境的安装和配置，部署流程和我之前介绍的 Jetson 嵌入式模型部署如出一辙，本篇博客从环境配置、模型训练、模型部署三方面进行了介绍，让大家对整体流程有一个基本的把握。博主对于 tensorRT_Pro 只做了最基础的演示，如果有更多的需求需要各位看官自己去挖掘啦。感谢各位看到最后，创作不易，读后有收获的看官请帮忙⭐️

下载链接

• 软件安装包[pwd:yolo]
• VOC数据集[pwd:yolo]
• yolov5-7.0预训练权重[pwd:yolo]
• 训练好的权重[pwd:yolo]
• tensorRT_Pro源文件[pwd:yolo]

参考

• tensorRT_Pro

• YOLOv5

• Ubuntu20.04软件安装大全

• 拯救者Y9000K2021H在ubuntu18.04安装显卡驱动

• Ubuntu卸载Nvidia驱动和安装最新驱动

• linux安装tensorflow，cuda，cudnn安装，pytorch兼容，RTX30系列GPU兼容cuda

• Linux安装TensorRT

• Ubuntu20.04安装TensorRT

• Ubuntu 20.04搭建OpenCV 4.5.0 & C++环境

• Ubuntu 安装 OpenCV 4.5.5

• 关于安装opencv编译报错：runtime library * hidden by 等若干问题

• cmake编译出错hidden by files

• Linux下编译protobuf

• Linux下添加protobuf的环境变量

• YoloV5案例第一部分，导出ONNX

• YoloV8的动态静态batch如何理解和使用

• Jetson nano部署YOLOv8

• 手写AI的详解TensorRT高性能部署视频

• make -j参数简介