python deploy/TensorRT/onnx_to_trt.py -m ./yolov6s_base.onnx \
-d int8 \
--batch-size 32 \
--num-calib-batch 4 \
--calib-img-dir ../coco/images/train2017/ \
--calib-cache ./yolov6s_base-int8-4x32.cache
进阶版模型大量采用重参数化结构,但在提升模型精度的同时,也为模型的量化带来了困难,该模型直接采用 PTQ 部署一般很难获得可接受的量化精度。为此,我们针对进阶版模型准备了一套基于 RepOPT 算法的量化流程。该方案需要首先采用 RepOPT 结构重新训练网络,然后再基于该网络预训练权重进行模型量化。该网络的部署模型结构与原网络基于 RepVGG 版本的 deploy 模型完全一致。
python tools/train.py --batch 32 --conf configs/repopt/yolov6s_hs.py --data data/coco.yaml --device 0
已发布的scale.pt见:yolov6n_scale.pt、yolov6t_scale.pt、yolov6s_scale.pt、yolov6m_scale.pt.
scale.pt
文件,config
中配置flag training_mode='hyper_search'
,训练代码如下,获得预训练模型yolov6s_reopt.pt
.python -m torch.distributed.launch --nproc_per_node 8 tools/train.py \
--batch 256 \
--conf configs/repopt/yolov6s_opt.py \
--data data/coco.yaml \
--epoch 300 \
--device 0,1,2,3,4,5,6,7 \
--name yolov6s_coco_repopt # yolov6l_coco
pytorch_quantization
,请先参考教程安装。该过程需要首先通过calib过程获得一个calib.pt文件,然后传入第二步进行QAT训练.CUDA_VISIBLE_DEVICES=0 python tools/train.py \
--data ./data/coco.yaml \
--output-dir ./runs/opt_train_v6s_ptq \
--conf configs/repopt/yolov6s_opt_qat.py \
--quant \
--calib \
--batch 32 \
--workers 0
步骤二: 量化感知训练(QAT)。该过程结合了CWD蒸馏,加载calib.pt
文件,训练10个epoch即可获得量化精度的提升,获得量化模型yolov6s_reopt_qat.pt
。
CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 python -m torch.distributed.launch --nproc_per_node=8 \
tools/train.py \
--data ./data/coco.yaml \
--output-dir ./runs/opt_train_v6s_qat \
--conf configs/repopt/yolov6s_opt_qat.py \
--quant \
--distill \
--distill_feat \
--batch 128 \
--epochs 10 \
--workers 32 \
--teacher_model_path ./assets/yolov6s_v2_reopt_43.1.pt \
--device 0,1,2,3,4,5,6,7
python3 qat_export.py --weights yolov6s_reopt.pt --quant-weights yolov6s_reopt_qat.pt --graph-opt --export-batch-size 1
执行该步骤可以得到带QDQ算子的ONNX模型,和不带QDQ算子的ONNX模型及对应的cache文件(带_remove_qdq_后缀),然后在目标部署环境下使用trtexec工具来生成trt模型进行部署,例如:
trtexec --workspace=1024 --percentile=99 --streams=1 --onnx=***_remove_qdq.onnx --calib=***_remove_qdq_add_insert_
YOLOv6提供了TRT模型精度测试的脚本,使用如下:
python deploy/TensorRT/eval_yolo_trt.py \
--imgs-dir ../data/coco/images/val2017/ \
--annotations ../data/coco/annotations/instances_val2017.json \
-m ***_remove_qdq.trt
pip install onnx>=1.10.0
python ./deploy/ONNX/export_onnx.py \
--weights yolov6s.pt \
--img 640 \
--batch 1 \
--simplify
--weights : yolov6 模型权重路径
--img : 模型输入图片尺寸,默认640
--batch : 模型输入的批大小
--half : 是否导出半精度(fp16)模型
--inplace : 是否需要设置Detect()类inplace为True
--simplify : 是否采用onnx-sim简化模型,端到端导出模型不支持简化
--end2end : 是否需要导出端到端的onnx模型,仅支持 onnxruntime 和 TensorRT >= 8.0.0
--trt-version : TensorRT 版本,支持7或8
--ort : 是否为 onnxruntime 后端导出模型
--with-preprocess : 是否需要预处理操作(bgr2rgb和归一化)
--topk-all : 保留每张图像的topK个目标
--iou-thres : NMS算法使用的IOU阈值
--conf-thres : NMS算法使用的置信度阈值
--device : 导出时用的环境设备,如显卡0或CPU
下载
端到端导出模型
onnxruntime 后端
使用方式
python ./deploy/ONNX/export_onnx.py \
--weights yolov6s.pt \
--img 640 \
--batch 1 \
--end2end \
--ort
您将获得带有 NMS 操作的 onnx 模型。
使用方式
python ./deploy/ONNX/export_onnx.py \
--weights yolov6s.pt \
--img 640 \
--batch 1 \
--end2end \
--trt-version 7
您将获得带有 BatchedNMSDynamic_TRT 插件的 onnx 模型。
使用方式
python ./deploy/ONNX/export_onnx.py \
--weights yolov6s.pt \
--img 640 \
--batch 1 \
--end2end \
--trt-version 8
您将获得带有 BatchedNMSDynamic_TRT 插件的 onnx 模型。
num_dets 表示其批次中每个图像中的目标数
det_boxes 表示 topk(100) 目标的坐标信息 [x0,y0,x1,y1] .
det_scores 表示每个 topk(100) 个对象的置信度分数
det_classes 表示每个 topk(100) 个对象的类别
您可以使用 trtexec 工具导出 TensorRT 引擎。
trtexec --onnx=yolov6s.onnx \
--saveEngine=yolov6s.engine \
--workspace=8192 # 8GB
--fp16 # if export TensorRT fp16 model
python deploy/ONNX/eval_trt.py --weights yolov6s.engine --batch-size=1 --data data/coco.yaml
动态批量推理,YOLOv6支持动态批量导出和推理,请参考以下教程:
export ONNX model with dynamic batch
export TensorRT model with dynamic batch
TensorRT-8.2.3.0
OpenCV-4.1.0
yolov6n.onnx
.python ./deploy/ONNX/export_onnx.py \
--weights yolov6n.pt \
--img 640 \
--batch 1
yolov6.engine
.python3 onnx_to_tensorrt.py --fp16 --int8 -v \
--max_calibration_size=${MAX_CALIBRATION_SIZE} \
--calibration-data=${CALIBRATION_DATA} \
--calibration-cache=${CACHE_FILENAME} \
--preprocess_func=${PREPROCESS_FUNC} \
--explicit-batch \
--onnx ${ONNX_MODEL} -o ${OUTPUT}
第三步:构建 demo,按照指引TensorRT Installation Guide 安装 TensorRT.
并且您应该在 CMakeLists.txt 中设置 TensorRT 路径和 CUDA 路径。
如果您训练自定义数据集,您可能需要修改 num_class
、image width,
image height
和 class name
的值。
const int num_class = 80;
static const int INPUT_W = 640;
static const int INPUT_H = 640;
static const char* class_names[] = {
"person", "bicycle", "car", "motorcycle", "airplane", "bus", "train", "truck", "boat", "traffic light",
"fire hydrant", "stop sign", "parking meter", "bench", "bird", "cat", "dog", "horse", "sheep", "cow",
"elephant", "bear", "zebra", "giraffe", "backpack", "umbrella", "handbag", "tie", "suitcase", "frisbee",
"skis", "snowboard", "sports ball", "kite", "baseball bat", "baseball glove", "skateboard", "surfboard",
"tennis racket", "bottle", "wine glass", "cup", "fork", "knife", "spoon", "bowl", "banana", "apple",
"sandwich", "orange", "broccoli", "carrot", "hot dog", "pizza", "donut", "cake", "chair", "couch",
"potted plant", "bed", "dining table", "toilet", "tv", "laptop", "mouse", "remote", "keyboard", "cell phone",
"microwave", "oven", "toaster", "sink", "refrigerator", "book", "clock", "vase", "scissors", "teddy bear",
"hair drier", "toothbrush"
};
mkdir build
cd build
cmake ..
make
./yolov6 ../you.engine -i image_path
python3 deploy/TensorRT/eval_yolo_trt.py -v -m model.trt \
--imgs-dir /workdir/datasets/coco/images/val2017 \
--annotations /workdir/datasets/coco/annotations/instances_val2017.json \
--conf-thres 0.40 --iou-thres 0.45
pip install --upgrade pip
pip install openvino-dev
python deploy/OpenVINO/export_openvino.py --weights yolov6s.pt --img 640 --batch 1
benchmark_app -m yolov6s_openvino/yolov6s.xml -i data/images/image1.jpg -d CPU -niter 100 -progress