Jetson Xavier AGX/NX Jetpack 4.5.1 + DeepStream 5.1 + yolov5(6.0)转换教程

本文实验使用是Jetpack4.5， cuda10.2，tensorrt 7.1.3.0-1+cuda10.2

1 下载依赖项目

Yolov5-in-Deepstream-5.0 https://github.com/Abandon-ht/Yolov5-in-Deepstream-5.0
tensorrtx-yolov5-v6.0 https://github.com/wang-xinyu/tensorrtx.git
yolov5 https://github.com/ultralytics/yolov5

2 编译 tensorrtx-yolov5-v6.0

git clone https://github.com/wang-xinyu/tensorrtx.git
mv tensorrtx tensorrtx-yolov5-v6.0
cd tensorrtx-yolov5-v6.0/yolov5
mkdir build && cd build
cmake ..
make  -j2

注意：修改tensorrtx-yolov5-v6.0/yolov5/yololayer.h将，类别数量和输入尺寸换成实际的数量和尺寸。否则，转换时会出现下面的错误：kernel weights has count 2688 but 32640 was expected。

static constexpr int CLASS_NUM = 2;
static constexpr int INPUT_H = 640;
static constexpr int INPUT_W = 640;

在yolov5.cpp中可以修改FP16、FP32，默认是FP16。

3 编译 Yolov5-in-Deepstream-5.0

https://github.com/DanaHan/Yolov5-in-Deepstream-5.0 【fork的仓库，支持tensort7 https://github.com/Abandon-ht/Yolov5-in-Deepstream-5.0/tree/tensorrt7】

git clone https://github.com/Abandon-ht/Yolov5-in-Deepstream-5.0/tree/tensorrt7
cd Yolov5-in-Deepstream-5.0
mv Deepstream\ 5.0 Deepstream5.0
cd Deepstream5.0/nvdsinfer_custom_impl_Yolo
make  # 修改Makefile中的CUDA_VER?=10.2

修改Makefile中的CUDA_VER?=10.2

4 生成wts文件

wts就是模型的权重文件，yolov5s.wts文件，需要通过tensorrtx->yolov5目录下的gen_wts.py文。因此具体操作步骤如下：

① 下载yolov5代码yolov5, 【选择好对应的版本】；
② 在https://github.com/ultralytics/yolov5下载pt文件；
③ 安装yolov5运行依赖的三方包，pip3 install -r requirements.txt；
④ 拷贝tensortx/yolov5目录下的gen_wts.py文件到yolov5目录；
⑤ 在yolov5目录下执行 python3 gen_wts.py -w yolov5s.pt，生成对应的yolov5s.wts文件。

如果是自己训练的模型，则直接使用训练生成的pt文件即可。

5 生成engine文件

将生成的wts后缀的文件上传到/home/nvidia/george/yolo_models目录下。
执行下面脚本开始转换：

#!/bin/bash

function check_result(){
    if [ $1 -ne 0 ]; then
        echo -e "\031[32m $2 execute failed \031[0m \n"
        exit $1
    else
        echo -e "\033[32m $2 execute success \033[0m \n"
    fi
}

if [ $# -ne 2 ];then
    echo  "Usage: $0 SIZE MODEL_NAME"
    exit 0
fi

CURR_PATH=`pwd`
SIZE_TYPE=$1
MODEL_NAME=$2

YOLOV5_5_SRC_DIR=/root/taitong/yolov5-6.0
YOLOV5_BIN_DIR=/root/taitong/tensorrtx-yolov5-v6.0/yolov5/build
#YOLOV5_BIN_DIR=/data/george/code/1804/Yolov5-in-Deepstream-5.0-tensorrt7/build

if [ ! -f "$CURR_PATH/taitong_models/${MODEL_NAME}.wts" ];then
    cd $YOLOV5_5_SRC_DIR
    python3 gen_wts.py -w $CURR_PATH/taitong_models/${MODEL_NAME}.pt -o $CURR_PATH/taitong_models/${MODEL_NAME}.wts
    check_result $? "gen_wts"
else
    echo "$CURR_PATH/taitong_models/${MODEL_NAME}.wts exist"
fi


if [ ! -f "$CURR_PATH/taitong_models/${MODEL_NAME}.engine" ];then
    cd $YOLOV5_BIN_DIR
    echo " $YOLOV5_BIN_DIR/yolov5 -s $CURR_PATH/taitong_models/${MODEL_NAME}.pt $CURR_PATH/taitong_models/${MODEL_NAME}.engine $SIZE_TYPE"
    LD_PRELOAD=$YOLOV5_BIN_DIR/libmyplugins.so $YOLOV5_BIN_DIR/yolov5 -s $CURR_PATH/taitong_models/${MODEL_NAME}.wts $CURR_PATH/taitong_models/${MODEL_NAME}.engine $SIZE_TYPE
    check_result $? "gen_engine"
else
    echo "$CURR_PATH/taitong_models/${MODEL_NAME}.engine exist"
fi