YOLO-TensorRT8 是基于 TensorRT8 和 Efficient NMS Plugin 实现的端到端检测架构。
利用 PyTorch 中的 torch.autograd.Function 和 symbolic 实现了 Efficient NMS Plugin 注册到 onnx 中。
效果图如图所示:
图中的输出共计 4 个，分别表示:
num_dets: Batch 中每张图片检测的目标数量, 类型为int32;det_boxes: Batch 中每张图片中每个目标的检测框, 目标框格式为左上角和右下角坐标(x0,y0,x1,y1), 类型为float32;det_scores: Batch 中每张图片中每个目标的置信度, 类型为float32;det_classes: Batch 中每张图片中每个目标的类别id,类型为int32;
# 安装环境
pip install -r requirements.txt
# 导出普通的 ONNX
python ./export.py \
--weights ./yolov7.pt \
--version 7 \
--img-size 640 \
--batch-size 1 \
--device cpu \
--simplify \
--opset 12
# 导出带有 Efficient NMS 的 ONNX
python ./export.py \
--weights ./yolov7.pt \
--version 7 \
--img-size 640 \
--batch-size 1 \
--device cpu \
--simplify \
--opset 12 \
--end2end \
--max-obj 100 \
--iou-thres 0.65 \
--score-thres 0.35
# 转换 ONNX 为 TensorRT engine
trtexec --onnx=./yolov7.onnx \
--saveEngine=./yolov7.engine \
--fp16 # 如果使用 fp16--weights: 使用最新版本开源仓库训练好的模型权重。--version: 指定模型版本。目前支持 5:yolov5, 6:yolov6, 7:yolov7; 未来考虑增加 8:airdet, -1:ppyoloe(需要安装PaddleDetection)。--img-size: 输入图片尺寸。--batch-size: 输入批量大小。--device: 指定使用 CPU 或者 CUDA:0 CUDA:1 ...导出 ONNX。--simplify: 使用 onnx-simplifier 简化模型。需要安装最新版本的 onnx-simplifier。--opset: 导出支持 opset 版本算子的 ONNX。--end2end: 导出端到端的带有 Efficient NMS 的 ONNX。--max-obj: 每张图片的最大检测框数量。--iou-thres: NMS 算法中的 IOU 阈值。--conf-thres: NMS 算法中的置信度阈值。
准备模型训练时使用的图片放到 ./calib_data 下，路径格式如下：
# 文件名随意，文件后缀最好是 .jpg
./calib_data/0001.jpg
./calib_data/0002.jpg
./calib_data/0003.jpg
./calib_data/0004.jpg
./calib_data/0005.jpg
....本仓库支持 pytorch 和 cuda-python 作为量化后端喂给 TensorRT 校准数据。
用法：
# pytorch + TensorRT int8量化
python ./build_engine.py \
--onnx yolov7.onnx \
--engine yolov7.engine \
--batch-size 1 \ # 需要和前文导出 onnx 时指定一致
--imgsz 640 \
--device 0 \
--verbose \ # 打印详细日志(未来会使用它增加绘制 TensorRT engine 结构的功能)
--workspace 8 \
--fp16 \ # 如果需要 fp16 + int8 可以同时开启
--int8 \
--calib-data ./calib_data \ # 校准图片路径可以自定义
--cache ./calib.cache \ # 校准集缓存,方便下次使用
--method torch