YOLOX-PAI:加速YOLOX,比YOLOV6更快更强

XGundam05

作者：忻怡、周楼、谦言、临在
导言



目标检测（object detection）旨在定位并识别出图像中的目标物体，一直以来都是计算机视觉领域研究的热点问题，也是自动驾驶、目标追踪等任务的基础。近年来，优秀的目标检测算法不断涌现，其中单阶段的YOLO系列以其高效、简洁的优势，始终是目标检测算法领域的一个重要分支。2021年，旷视提出YOLOX[1]算法，在速度和精度上构建了新的基线，组件灵活可部署，深受工业界的喜爱。本文基于阿里云 PAI-EasyCV框架复现YOLOX算法，探究用以实际提升YOLOX精度的实用技巧，并进一步结合阿里巴巴计算平台PAI自研的PAI-Blade推理加速框架优化模型性能。经过我们对社区诸多YOLOX 改进技巧的复现和探索，进一步提升了YOLOX的性能，在速度和精度上都比现阶段的40~50mAP 的SOTA的YOLOv6更胜一筹。同时，PAI-EasyCV提供高效简洁的模型部署和端到端推理接口，供社区快速体验使用YOLOX-PAI的功能。
总结一下我们的工作贡献：

• 我们提供了一套Apache License 训练/优化/推理的代码库以及镜像，可以实现当前社区40+MAP 量级最快（相比  YOLOV6 +0.4mAP/加速13~20%）的目标检测模型。
  
• 我们调研了YOLOX相关的改进技术和消融实验，挑选了其中一些相对有帮助的改进，补齐了40/0.7ms(YOLOXS)~47.6/1.5ms(YOLOXM) 之间的模型，并以配置的方式提供出来。
  
• 我们对目标检测的端到端推理进行灵活封装及速度优化，在V100上的端到端推理为3.9ms，相对原版YOLOX的9.8ms，加速接近250%，供用户快速完成目标检测推理任务。
  

本文，我们将逐一介绍所探索的相关改进与消融实验结果，如何基于PAI-EasyCV使用PAI-Blade优化模型推理过程，及如何使用PAI-EasyCV进行模型训练、验证、部署和端到端推理。欢迎大家关注和使用PAI-EasyCV和PAI-Blade，进行简单高效的视觉算法开发及部署任务。
PAI-EasyCV项目地址：https://github.com/alibaba/EasyCV
PAI-BladeDISC项目地址：https://github.com/alibaba/BladeDISC
YOLOX-PAI-算法改进

YOLOX-PAI 是我们在阿里云机器学习平台PAI 的开源计算机视觉代码库EasyCV(https://github.com/alibaba/EasyCV)中集成的 YOLOX 算法。若读者不了解YOLOX算法，可以自行学习(可参考：链接)，本节主要介绍我们基于YOLOX算法的改进。
通过对YOLOX 算法的分析，结合检测技术的调研，我们从以下4个方向对原版的YOLOX进行优化，

• Backbone :  repvgg backbone
  
• Neck  : gsconv / asff
  
• Head :  toods / rtoods
  
• Loss  : siou / giou
  

在算法改进的基础上，利用PAI-Blade对优化后的模型进行推理优化，开发了如下的PAI-YOLOX模型。筛选有效改进与现有主流算法的对比结果如下：
（ -ASFF 代表使用了 NeckASFF， -TOODN代表使用N个中间层的TOODHead取代原有的YOLOXHead）

Model	Params (M)	Flops (G)	mAPval 0.5:0.95 （672）	mAPval 0.5:0.95 （640）	SpeedV100 fp16 b32(ms)	Download
YOLOXs	9.0	26.8	40.2	40.1	0.68	model - log
Meituan-YOLOv6-tiny	15.0	36.7	41.3	-	0.9	-
PAI-YOLOXs	15.9	36.8	41.5	41.4	0.7(+23%)	model - log
Meituan-YOLOv6-s	17.2	44.2	43.1	42.4	1.0	-
PAI-YOLOXs-ASFF	21.3	41.0	43.3	42.8	0.87(+13%)	model - log
PAI-YOLOXs-ASSF-TOOD3	23.7	49.9	44.0	43.9（+1.5）	1.15	model - log
YOLOXm	25.3	73.8	46.3	46.3	1.50	model - log

从结果中可以看到，相比目前同水平(1ms以内)SOTA的YOLO6模型，融合上述改进的YOLOX-PAI在同等精度/速度的条件下有一定的速度/精度优势。
有关测评需要注意以下几点：

• YOLOV6 release测试速度不包含 decode和nms，所以我们为了公平对比，也进行了相关测速设置的关闭。（上表所示结果计算了Blade优化后的对应模型在bs32下平均一张图像模型前向推理所用时间，关于端到端推理的时间（包含图像前、后处理的总用时）见PAI-EasyCV Export一节）
  
• YOLOV6 release的精度是在训练过程中测试的，会出现部分 shape=672的情况，然而测速是在导出到image_size=(640, 640) 的完成，实际上社区也有相关同学补充了YOLOV6在640下的测试精度，所以上表放了两个测试精度。
  
• 使用EasyCV的Predictor 接口加载相关模型预测进行从图片输入到结果的预测，由于包含了预处理和NMS后处理，相对应的时间会变慢一些，详细参考端到端优化结果。
  

下面我们将详细介绍每一个模块的改进和消融实验。
Backbone

RepConv

近期YOLO6 [2]，PP-YOLOE [3]等算法都改进了CSPNet[4]的骨干网络，基于RepVGG[5]的思想设计了可重参数化的骨干网络，让模型在推理上具有更高效的性能。我们沿用了这一思想，利用YOLO6的EfficientRep代替YOLOX原来的CSPDarkNet-53骨干网络。得到的实验结果与YOLO6相关模型对比如下（YOLOX-Rep表示使用了EfficientRep作为骨干网络的YOLOX模型）：

Model	Params (M)	Flops (G)	mAPval 0.5:0.95 （672）	mAPval 0.5:0.95 （640）	SpeedV100 fp16 b32 (ms)
YOLOXs	9.0	26.8	40.2	40.1	0.68
YOLOv6-tiny	15.0	36.7	41.3	-	0.9
PAI-YOLOXs	15.9	36.8	41.5	41.4	0.7

        RepVGG结构的网络设计确实会增大参数量和计算量，但实际推理速度都更有优势，所以我们选择YOLO6 EfficientRep 作为可以配置的Backbone。
Neck

在更换了骨干网络的基础上，我们对Neck部分分别进行了两方面的探究。

• ASSF[6]：基于对PAN输出特征维度变换后基于SE-Attention特征融合的特征增强，大幅提升参数量和精度，部分降低推理速度。
  
  
  • ASSF-Sim ： 我们选取了参数量更低的特征融合实现，用较少的参数量（ASFF:5M -> ASFF-Sim:380K）来保留了84%的精度精度提升（+0.98map->+0.85map）。然而，这一改进会让推理速度变慢，未来我们会针对这个OP实现对应的Plugin完成推理加速。
    
  
  
• GSNeck[7] ：基于DW Conv 对Neck信息融合，降低Neck参数量，轻微提升精度，也会会降低推理速度。
  

ASFF 信息融合

ASFF，通过进行不同PAN不同特征图之间的信息交互，利用attention机制完成Neck部分的信息融合和增强，具体思想如下图。


ASFF-SIM轻量版

参考YOLO5[8]中的Fcous层的设计，PAI-EasyCV利用切片操作进行特征通道的增加和特征图的缩小。同时，利用求平均操作进行通道的压缩，基于这种实现的ASFF，我们简单区分为ASFF-Sim。我们进行特征图统一的核心操作（通道扩展和通道压缩）如下：
def expand_channel(self, x):
    # [b,c,h,w]->[b,c*4,h/2,w/2]
    patch_top_left = x[..., ::2, ::2]
    patch_top_right = x[..., ::2, 1::2]
    patch_bot_left = x[..., 1::2, ::2]
    patch_bot_right = x[..., 1::2, 1::2]
    x = torch.cat(
        (
            patch_top_left,
            patch_bot_left,
            patch_top_right,
            patch_bot_right,
        ),
        dim=1,
    )
    return x

def mean_channel(self, x):
    # [b,c,h,w]->[b,c/2,h,w]
    x1 = x[:, ::2, :, :]
    x2 = x[:, 1::2, :, :]
    return (x1 + x2) / 2针对不同的特征图，其融合机制如下：

Model	Params (M)	Flops (G)	mAPval 0.5:0.95	SpeedV100 fp16 b32 (ms)
PAI-YOLOXs	15.9	36.8	41.4	0.70
YOLOX-s-ASFF	21.3	41.0	42.8	0.87
YOLOX-s-ASFFSim	16.3	37.1	42.6	1.13 (need speed up)

GSConvNeck

采用DWConv降低参数量是一种常用技巧，在YOLOX中，GSconv设计了一种新型的轻量级的卷积用来减少模型的参数和计算量。为了解决Depth-wise Separable Convolution (DSC)在计算时通道信息分离的弊端，GSConv（如下图所示）采用shuffle的方式将标准卷积（SC）和DSC得到的特征图进行融合，使得SC的输出完全融合到DSC中。


此外，GSConv原文指出，如果在整个网络都使用GSconv，则会大大加深网络的深度，降低模型的推理速度，而仅在通道信息维度最大，空间信息维度最小的Neck处使用GSConv是一种更优的选择。我们在YOLOX中利用GSConv优化模型，特别的我们采用了两种方案分别进行实验(a: 仅将Neck的部分用GSConv，b: Neck的所有模块均使用GSConv)：



  （a）gsconv_part   



(b) gsconv

所得到的实验结果如下（仅统计Neck部分的Params, Flops）。可以看到GSConv对参数量进行了优化，且提升了模型的性能，降低3%的推理速度可以换来0.3mAP的提升。

Model	Params (M)	Flops (G)	mAPval 0.5:0.95	SpeedV100 fp16 b32 (ms)
Yolox-original	2.83	2.67	41.35	0.70
gsconv	1.22	1.33	41.58	0.78
gsconv_part	2.39	2.39	41.68	0.72

Head

TOOD

参考PPYOLOE，我们同样考虑利用TOOD[9]算法中的Task-aligned predictor中的注意力机制（T-Head）分别对分类和回归特征进行增强。如下图所示，特征先通过解耦头的stem层(1x1)进行通道压缩，接着由通过堆叠卷积层得到中间的特征层，再分别对分类和回归分支利用注意力机制进行特征的增强，来解耦两个任务。


我们对堆叠的中间层个数进行消融实验，每堆叠可以进一步提升一定的精度，并带来速度上的一些损失。(下表的Params和Flops只统计了检测头部分。测速及精度基于的基线方法为以EfficientRep为backbone + ASFF进行特征增强。)

Model	Params (M)	Flops (G)	mAPval 0.5:0.95	SpeedV100 fp16 b32 (ms)
stack = 6	2.83M	12.66G	44.66	1.40
stack = 5	2.68M	11.42G	44.41	1.32
stack = 4	2.53M	10.18G	44.12	1.24
stack = 3	2.37M	8.94G	43.85	1.15
stack = 2	2.22M	7.71G	43.45	1.09
YOLOX Head	1.92M	5.23G	42.78	0.87

此外，我们利用RepVGGBlock分别优化inter_conv，cls_conv/reg_conv层。实验结果表明用RepVGGBlock实现inter_conv层会带来性能上的下降，而RepVGGBlock实现的cls_conv/reg_conv层与原始卷积层在stack较大时效果相当,在stack比较小时，RepVGGBlock能起到优化作用。
Loss function

S/G/D/E/CIou

PAI-EasyCV实现了社区常用的集中IOU损失函数，用户可以通过config自行选择，特别的，对于最新提出的SIoU[10]，在实验过程中发现原公式中的反三角函数会使得模型训练不稳定，且计算开销相对较高，我们对原公式利用三角函数公式化简可得（符号与论文原文一致，详见原论文）：


实验结果显示，在YOLOX上引入SIoU训练模型的确能加快模型的收敛速度，但在最终精度上使用GIoU[11]性能达到最优。
综合上述Backbone/Neck/Head/Loss 方向的改进，我们可以获得如上的YOLOX-PAI模型。进一步，我们采用PAI-Blade 对已经训练好的模型进行推理优化，实现高性能的端到端推理。
YOLOX-PAI-推理优化



PAI-EasyCV Predictor

针对使用PAI-EasyCV训练的YoloX-PAI 模型，用户可以使用EasyCV自带的导出(export)功能得到优化后的模型，并使用 EasyCV Predictor 进行端到端的推理。 该导出功能对检测模型进行了如下优化：

• 使用PAI-Blade优化模型推理速度，简化对模型的推理加速（TensorRT/编译优化）开发流程。
  
• 支持EasyCV配置TorchScript/PAI-Blade对图像前处理、模型推理、图像后处理分别优化，供用户灵活使用
  
• 支持Predictor 结构端到端的模型推理优化，简化图片预测过程。
  

也可以参考[EasyCV detector.py] 自行组织相应的图像前处理/后处理过程，或直接使用我们导出好的模型和接口)：我们这里提供一个已经导出好的检测模型，用户下载三个模型文件到本地 [preprocess, model, meta]
export_blade/
├── epoch_300_pre_notrt.pt.blade
├── epoch_300_pre_notrt.pt.blade.config.json
└── epoch_300_pre_notrt.pt.preprocess用户可以直接使用PAI-EasyCV提供的Predictor接口，通过如下简单的API调用，高效的进行图像的检测任务：
from easycv.predictors import TorchYoloXPredictor
from PIL import Image

img = Image.open(img_path = 'demo.jpg')
pred = TorchYoloXPredictor(export_model_name='epoch_300_pre_notrt.pt.blade',
                           use_trt_efficientnms=False)
res = pred.predict([img])PAI-EasyCV Export

下面我们简单介绍如何通过PAI-EasyCV的配置文件，导出不同的模型（具体的模型部署流程即相应的配置文件说明介绍见链接），并展示导出的不同模型进行端到端图像推理的性能。
为导出不同的模型，用户需要对配置文件进行修改，配置文件的说明如下：
export = dict(export_type='ori', # 导出的模型类型['ori','jit','blade']
              preprocess_jit=False, # 是否用jit对前处理进行加速
              static_opt=True, # 是否使用static shape优化，默认True
              batch_size=1,  # 静态图的输入batch_size
              blade_config=dict(
                  enable_fp16=True,
                  fp16_fallback_op_ratio=0.05 # 所有的layer都会针对转fp16前后的输出
                  # 的偏移进行排序，会对数值变化最大的层回退到fp32，该参数用于控制回退的比例，
                  # 如果出现模型输出漂移太大，影响相关测试结果，可以手动调整该参数。
              ),
              use_trt_efficientnms=False) # 是否使用trt优化的efficientnms根据不同的模型配置，我们在单卡V100上测试YOLOX-s所有配置下模型的端到端推理性能 （1000次推理的平均值）：

export_type	preprocess_jit	use_trt_efficientnms	Infer time (end2end) /ms
raw	-	-	24.58
jit	False	False	18.30
jit	False	True	18.38
jit	True	False	13.44
jit	True	True	13.04
blade	False	False	8.72
blade	False	True	9.39
blade	True	False	3.93
blade	True	True	4.53

下图，我们展示了由PAI-EasyCV中集成的使用PAI-Blade/JIT优化的模型端到端推理速度与YOLOX官方原版的 不同模型(s/m/l/x)的推理速度对比：


可以看到PAI-EasyCV导出的模型，极大程度的优化了原模型的端到端推理速度，未来我们将进一步优化blade接入trt_efficientnms的速度，提升端到端性能。
PAI-Blade 推理优化

PAI-Blade是由阿里云机器学习平台PAI开发的模型优化工具，可以针对不同的设备不同模型进行推理加速优化。PAI-Blade遵循易用性，鲁棒性和高性能为原则，将模型的部署优化进行高度封装，设计了统一简单的API，在完成Blade环境安装后，用户可以在不了解ONNX、TensorRT、编译优化等技术细节的条件下，通过简单的代码调用方便的实现对模型的高性能部署。更多PAI-Blade相关技术介绍可以参考 [PAI-Blade介绍]。
PAI-EasyCV中对Blade进行了支持，用户可以通过PAI-EasyCV的训练config 中配置相关export 参数，从而对训练得到的模型进行导出。
这里我们提供一个 PAI-Blade + PAI-EasyCV 社区版 V100对应的镜像（cuda11.1/TensorRT8/cudnn8）：用户也可以基于Blade每日发布的镜像自行搭建推理环境  [PAI-Blade社区镜像发布]
registry.cn-shanghai.aliyuncs.com/pai-ai-test/eas-service:blade_cu111_easycv用户执行如下导出命令即可
cd ${EASYCV_ROOT}
export LD_LIBRARY_PATH=/usr/loca/cuda/lib64/:${LD_LIBRARY_PATH}
export CUDA_VISIBLE_DEVICES=0
export PYTHONPATH='./'
python tools/export.py yolox_pai_trainconfig.py input.pth export.pth值得注意的是上文所有的模型的推理速度都限定在 V100 BatchSize=32 静态Shape (end2end=False)的PAI-Blade优化设置结果。Blade中已经集成了常见OP的优化，针对用户自定义的op可以参考PAI-EasyCV中的easycv/toolkit/blade/trt_plugin_utils.py 自行实现。
YOLOX-PAI-训练与复现

我们在PAI-EasyCV框架中复现了原版的YOLOX，及改进的YOLOX-PAI，并利用PAI-Blade对模型进行推理加速。为了更好的方便用户快速体验基于PAI-EasyCV和PAI-Blade的YOLOX，接下来，我们提供利用镜像对YOLOX-PAI进行模型的训练、测试、及部署工作。更多的关于如何在本地开发环境运行，可以参考该链接安装环境。若使用PAI-DSW进行实验则无需安装相关依赖，在PAI-DSW docker中已内置相关环境。
拉取镜像
sudo docker pull registry.cn-shanghai.aliyuncs.com/pai-ai-test/eas-service:blade_cu111_easycv启动容器
sudo nvidia-docker run -it -v path:path --name easycv_yolox_pai --shm-size=10g --network=host registry.cn-shanghai.aliyuncs.com/pai-ai-test/eas-service:blade_cu111_easycv数据代码准备
# 数据准备参考 https://github.com/alibaba/EasyCV/blob/master/docs/source/prepare_data.md
git clone https://github.com/alibaba/EasyCV.git
cd EasyCV模型训练
export PYTHONPATH=./ && python -m torch.distributed.launch  --nproc_per_node=8 --master_port=29500 tools/train.py config.py --work_dir workdir --launcher pytorch模型测试
python tools/eval.py config.py pretrain_model.pth --eval模型导出
python tools/export.py config.py pretrain_model.pth export.pth写在最后

YOLOX-PAI 是PAI-EasyCV团队基于旷视YOLOX 复现并优化的在V100BS32的1000fps量级下的SOTA检测模型。整体工作上集成和对比了很多社区已有的工作：通过对YOLOX的替换基于RepVGG的高性能Backbone， 在Neck中添加基于特征图融合的ASFF/GSConv增强，在检测头中加入了任务相关的注意力机制TOOD结构。结合PAI-Blade编译优化技术，在V100Batchsize32 1000FPS的速度下达到了SOTA的精度mAP=43.9，同等精度下比美团YOLOV6 加速13%，并提供了配套一系列算法/训练/推理优化代码和环境。
PAI-EasyCV（https://github.com/alibaba/EasyCV）是阿里云机器学习平台深耕一年多的计算机视觉算法框架，已在集团内外多个业务场景取得相关业务落地成果，主要聚焦在自监督学习/VisionTransformer等前沿视觉领域，并结合PAI-Blade等自研技术不断优化。欢迎大家参与进来一同进步。
YOLOX-PAI未来规划：

• 基于CustomOP（ASFFSim, EfficientNMS (fp16)）实现的加速推理
  

[1] Ge Z, Liu S, Wang F, et al. Yolox: Exceeding yolo series in 2021[J]. arXiv preprint arXiv:2107.08430, 2021.
[2] YOLOv6, https://github.com/meituan/YOLOv6.
[3] Xu S, Wang X, Lv W, et al. PP-YOLOE: An evolved version of YOLO[J]. arXiv preprint arXiv:2203.16250, 2022.
[4]  Wang C Y, Liao H Y M, Wu Y H, et al. CSPNet: A new backbone that can enhance learning capability of CNN[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition workshops. 2020: 390-391.
[5] Ding X, Zhang X, Ma N, et al. Repvgg: Making vgg-style convnets great again[C]//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2021: 13733-13742.
[6] Liu S, Huang D, Wang Y. Learning spatial fusion for single-shot object detection[J]. arXiv preprint arXiv:1911.09516, 2019.
[7] Li H, Li J, Wei H, et al. Slim-neck by GSConv: A better design paradigm of detector architectures for autonomous vehicles[J]. arXiv preprint arXiv:2206.02424, 2022.
[8] YOLOv5, https://github.com/ultralytics/yolov5.
[9] Feng C, Zhong Y, Gao Y, et al. Tood: Task-aligned one-stage object detection[C]//2021 IEEE/CVF International Conference on Computer Vision (ICCV). IEEE Computer Society, 2021: 3490-3499.
[10] Gevorgyan Z. SIoU Loss: More Powerful Learning for Bounding Box Regression[J]. arXiv preprint arXiv:2205.12740, 2022.
[11] Rezatofighi H, Tsoi N, Gwak J Y, et al. Generalized intersection over union: A metric and a loss for bounding box regression[C]//Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2019: 658-666.

yukamu

牛啊牛啊[抱抱][抱抱][抱抱]