YOLOv5-Lite TensorRT Deployment

Environment

TensorRT: 8.6.1.6 CUDA: 12.6 CUDNN: 8.9.0

配置参数修改

在开始之前，需要修改 include/yololayer.h 文件中的相关参数，确保它们与您的训练过程配置一致：

static constexpr int MAX_OUTPUT_BBOX_COUNT = 1000; 
static constexpr int CLASS_NUM = 80;  // 类别数
static constexpr int INPUT_H = 640;   // yolov5-lite的输入高度，必须能被32整除
static constexpr int INPUT_W = 640;   // yolov5-lite的输入宽度，必须能被32整除
static constexpr int DEVICE = 0;
static constexpr float NMS_THRESH = 0.4;
static constexpr float CONF_THRESH = 0.45;
static constexpr int BATCH_SIZE = 1;
const char* INPUT_BLOB_NAME = "data";
const char* OUTPUT_BLOB_NAME = "prob";

1. 由pt文件生成wts文件

此过程必须在yolov5-lite文件夹中完成：

python gen_wts.py -w v5lite-s.pt -o v5lite-s.wts
python gen_wts.py -w v5lite-e.pt -o v5lite-e.wts
python gen_wts.py -w v5lite-g.pt -o v5lite-g.wts

2. 编译engine文件和推理

编译步骤

a. 首先在include文件夹中找到yololayer.h文件，将CLASS_NUM改为您的数据集对应的类别数，这一步非常重要，否则会报错！

b. 在终端中执行以下命令：

mkdir build
cd build
cmake ..
make

生成engine文件

./v5lite -s ../v5lite-s.wts v5lite-s.engine s
./v5lite -s ../v5lite-g.wts v5lite-g.engine g
./v5lite -s ../v5lite-e.wts v5lite-e.engine e
./v5lite -s ../v5lite-c.wts v5lite-c.engine c

使用engine文件进行推理

(samples是您存放图片的文件夹)：

./v5lite -d v5lite-s.engine ../samples

您也可以使用yolov5-lite-trt.py进行推理。

3. INT8量化

准备工作

1. 获取训练集的图片，数量建议1000张左右

2. 将图片放入指定的校准数据文件夹中（例如：tensorrtx-int8calib-data/coco_calib）

3. 修改v5lite.cpp中的宏定义：

   将：

   // #define USE_FP16  // set USE_INT8 or USE_FP16 or USE_FP32
   // #define USE_INT8  // set USE_INT8 or USE_FP16 or USE_FP32

   改为：

   // #define USE_FP16  // set USE_INT8 or USE_FP16 or USE_FP32
   #define USE_INT8  // set USE_INT8 or USE_FP16 or USE_FP32

4. 更新代码中的数据路径，指向您的校准图像文件夹

5. 重新编译并生成engine进行推理(重复步骤2)

注意事项

• 实际测试中发现Python调用engine的推理效果可能更好