本文实验使用是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
