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

导言

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

总结一下我们的工作贡献：

1. 我们提供了一套Apache License 训练/优化/推理的代码库以及镜像，可以实现当前社区40+MAP 量级最快（相比  YOLOV6 +0.4mAP/加速13~20%）的目标检测模型。
2. 我们调研了YOLOX相关的改进技术和消融实验，挑选了其中一些相对有帮助的改进，补齐了40/0.7ms(YOLOXS)~47.6/1.5ms(YOLOXM) 之间的模型，并以配置的方式提供出来。
3. 我们对目标检测的端到端推理进行灵活封装及速度优化，在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进行优化，

1. Backbone :  repvgg backbone
2. Neck  : gsconv / asff
3. Head :  toods / rtoods
4. Loss  : siou / giou

在算法改进的基础上，利用PAI-Blade对优化后的模型进行推理优化，开发了如下的PAI-YOLOX模型。筛选有效改进与现有主流算法的对比结果如下：

（ -ASFF 代表使用了 NeckASFF， -TOODN代表使用N个中间层的TOODHead取代原有的YOLOXHead）

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

有关测评需要注意以下几点：

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

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

Backbone

RepConv

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

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

Neck

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

1. ASSF[6]：基于对PAN输出特征维度变换后基于SE-Attention特征融合的特征增强，大幅提升参数量和精度，部分降低推理速度。

1. ASSF-Sim ： 我们选取了参数量更低的特征融合实现，用较少的参数量（ASFF:5M -> ASFF-Sim:380K）来保留了84%的精度精度提升（+0.98map->+0.85map）。然而，这一改进会让推理速度变慢，未来我们会针对这个OP实现对应的Plugin完成推理加速。

2. 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

针对不同的特征图，其融合机制如下：

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的提升。

Head

TOOD

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

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

此外，我们利用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次推理的平均值）：

下图，我们展示了由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未来规划：

1. 基于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.