YOLO系列模型的部署、精度对齐与int8无损失量化加速

本文来自社区投稿，作者：海滨（邮箱 [e-mail]）

在目标检测模型的部署与量化加速过程中，我们往往会遇到训练侧和实际推理侧不能完全对齐的问题，比如即使是以 FP32 形式推理，由于 TensorRT 算子参数的一些限制和 TensorRT 和 Torch 内部实现的不同，导致 torch 推理结果会和 TensorRT 推理结果天然的不统一。

基于此，本人于 2023 年初基于 MMYOLO 完成了一项工作——YOLO 系列模型在 TensorRT 上的部署与量化加速。本项目麻雀虽小但五脏俱全，系统介绍了 YOLO 系列模型在 TensorRT 上的量化方案，工程型较强，我们给出的工具可以实现不同量化方案在 YOLO 系列模型的量化部署，无论是工程实践还是学术实验，相信都会对大家带来一定的帮助。

目前项目已开源，对应的讲解视频也已在 B 站发布，欢迎大家关注！

B站地址（求关注和三连）：

https://www.bilibili.com/video/BV1Ds4y1k7yr/

GitHub 开源地址（求star）：

https://github.com/thb1314/mmyolo_tensorrt/

接下来，就带大家详细了解下项目包含的主要内容。

YOLO 系列模型在 TensorRT 

上的部署与精度对齐

该项目详细介绍了 YOLO 系列模型在 TensorRT 上的 FP32 的精度部署，基于 MMYOLO 框架导出各种 YOLO 模型的onnx，在 coco val 数据集上对齐 Torch 版本与 TensorRT 版本的精度。

在此过程中我们发现，由于 TopK 算子限制和 NMS 算子实现上的不同，我们无法完全对齐 Torch 和 YOLO 模型的精度，不过这种风险是可解释且可控的。

详解 TensorRT 量化的三种实现方式

TensorRT 量化的三种实现方式包括 trt7 自带量化、dynamic range api，trt8 引入的 QDQ 算子。

Dynamic range api 会在采用基于 MQbench 框架做 PTQ 时讲解。

TensorRT 引入的 QDQ 算子方式在针对 YOLO 模型的 PTQ 和 QAT 方式时都有详细的阐述，当然这个过程也没有那么顺利。

在基于 PytorchQuantization 导出的含有 QDQ 节点的 onnx 时，我们发现尽管量化版本的 torch 模型精度很高，但是在 TensorRT 部署时精度却很低，TRT 部署时精度损失很严重，通过对可视化其他量化形式的 engine 和问题 engine 进行对比，我们发现是一些层的 int8 量化会出问题，由此找出问题量化节点解决。

详解 MQBench 量化工具包

在 TensorRT 上的应用

我们研究了基于 MQbench 框架的普通 PTQ 算法，包括 Adaround 的高阶 PTQ 算法，以及启发于 Adaround 的高阶 PTQ 算法。

我们将 torch 版本中的 HistogramObserver 引入到 MQBench 中，activation 采用 HistogramObserver，weight 采用 MinMaxObserver，在 PTQ 过程中，weight 的校准前向传播一次，activation 的校准需要多次。因此我们将 weight 的 PTQ 过程和 activation 的 PTQ 过程分开进行，加速 PTQ 量化。

实践证明，我们采用上述配置的分离 PTQ 量化在 YOLOV8 上可以取得基本不掉点的 int8 量化精度。

model_mqbench = prepare_by_platform(torch_model.train(), BackendType.Tensorrt, prepare_custom_config_dict)model_mqbench.to(device)
model_mqbench.eval()
enable_calibration_woquantization(model_mqbench, quantizer_type='weight_fake_quant')model_mqbench(torch.rand(1,3,640,640).to(device))
# do activation and weight calibration seperatelyenable_calibration_woquantization(model_mqbench, quantizer_type='act_fake_quant')
# calibration loop PTQ processfrom torch.utils.data import Dataset, DataLoaderclass QuantDataset(Dataset):    def __init__(self) -> None:        super().__init__()        self.calib_files, _ = get_file_list(args.calib_imgs)    def __getitem__(self, index):        file = self.calib_files[index]        bgr = mmcv.imread(file, channel_order='bgr')        data, samples = test_pipeline(dict(img=bgr, img_id=index)).values()        data = pre_pipeline(data)[0]        return data    def __len__(self):        return len(self.calib_files)        quant_dataset = QuantDataset()calib_batch_size = 32quant_dataloader = DataLoader(quant_dataset, batch_size=calib_batch_size, num_workers=os.cpu_count() // 8 * 8, shuffle=False)
max_index = 100for i, data in tqdm.tqdm(enumerate(quant_dataloader), total=max_index):    if i >= max_index:        break    data = data.to(device)            with torch.no_grad():        model_mqbench(data)

针对 YOLOV6 这种难量化模型，

分别采用部分量化和 QAT 来弥补量化精度损失

在部分量化阶段，我们采用量化敏感层分析技术来判断哪些层最需要恢复原始精度，给出各种 metric 的量化敏感层实现。

在 QAT 阶段，不同于原始 YOLOV6 论文中蒸馏 + RepOPT 的方式，我们直接采用上述部分量化后的模型做出初始模型进行 finetune ，结果发现 finetune 后的模型依然取得不错效果。

我们最终部署在 YOLOV6-S int8 QAT 版本的精度如下：

Average Precision  (AP) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.430 Average Precision  (AP) @[ IoU=0.50      | area=   all | maxDets=100 ] = 0.602 Average Precision  (AP) @[ IoU=0.75      | area=   all | maxDets=100 ] = 0.464 Average Precision  (AP) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.232 Average Precision  (AP) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.479 Average Precision  (AP) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.587 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=  1 ] = 0.346 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets= 10 ] = 0.575 Average Recall     (AR) @[ IoU=0.50:0.95 | area=   all | maxDets=100 ] = 0.628 Average Recall     (AR) @[ IoU=0.50:0.95 | area= small | maxDets=100 ] = 0.414 Average Recall     (AR) @[ IoU=0.50:0.95 | area=medium | maxDets=100 ] = 0.694 Average Recall     (AR) @[ IoU=0.50:0.95 | area= large | maxDets=100 ] = 0.808

针对旋转目标检测，

我们同样给出一种端到端方案，

最后的输出就是 NMS 后的结果

通过将 TensorRT 中的 EfficientNMS Plugin 和 mmcv 中旋转框 iou 计算的 cuda 实现相结合，给出 EfficientNMS For Rotated Box 版本，经过简单验证我们的 TRT 版本与 Torch 版本模型输出基本对齐。

以上就是我们这个项目做的事情，主要是针对 YOLO 这种 Single-Stage 目标检测模型在 TensorRT 上的部署和量化做一个总结。

欢迎大家观看 B 站视频了解项目详情，如果各位有更好的想法也欢迎给我们提 PR。

B站地址（求关注和三连）：

https://www.bilibili.com/video/BV1Ds4y1k7yr/

GitHub 开源地址（求star）：

https://github.com/thb1314/mmyolo_tensorrt/

点击下方“阅读原文”直达项目 GitHub 页面