如何量化加速 YOLOv6

首先请参考选型指导,选择适合于自己的模型。对于量化新手同学,我们推荐选择基础版(base)模型,并直接采用 PTQ 的方式进行模型量化并部署。对于量化有一定基础的同学,可以选择进阶版模型来进行完整的量化模型生产。

基础版模型PTQ量化(推荐)

模型训练完成后,首先参考ONNX模型导出生成对应的ONNX模型,然后根据以下命令进行 PTQ生成 TRT 量化模型

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

进阶版RepOPT

进阶版模型大量采用重参数化结构,但在提升模型精度的同时,也为模型的量化带来了困难,该模型直接采用 PTQ 部署一般很难获得可接受的量化精度。为此,我们针对进阶版模型准备了一套基于 RepOPT 算法的量化流程。该方案需要首先采用 RepOPT 结构重新训练网络,然后再基于该网络预训练权重进行模型量化。该网络的部署模型结构与原网络基于 RepVGG 版本的 deploy 模型完全一致。

RepOpt两阶段训练

RepOptimizer给出了一种量化友好的重参数化网络训练范式。该方法采用优化器Optimizer的重参数化操作来代替结构的重参数化过程,从而保持网络结构在train和deploy阶段保持一致。该方法分为两个步骤: 步骤一: 超参数搜索(Hyper-Search)。该阶段通过训练LinearAddBlock获得超参数scale(通过 scale.pt 的形式存储),来初始化RepOPT结构的Optimizer。您可以直接使用我们所发布的超参数预训练模型,跳过该步骤;或者选择重新搜索获得自己的超参数scale

python tools/train.py --batch 32 --conf configs/repopt/yolov6s_hs.py --data data/coco.yaml --device 0

已发布的scale.pt见:yolov6n_scale.ptyolov6t_scale.ptyolov6s_scale.ptyolov6m_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

量化感知训练

对生产的RepOPT预训练模型进行QAT提升量化精度,该步骤依赖于pytorch_quantization,请先参考教程安装。该过程需要首先通过calib过程获得一个calib.pt文件,然后传入第二步进行QAT训练. 步骤一: 数据校准(Calibration)。

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

量化模型导出及转换

导入前两步所得pt模型,导出成ONNX模型。

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_qdq_calibration.cache --int8 --fp16 --saveEngine=***_remove_qdq.trt

TRT模型精度测试

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