关于机器学习:YOLOXPAI加速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/Ea…

PAI-BladeDISC 我的项目地址：https://github.com/alibaba/Bl…

YOLOX-PAI- 算法改良

YOLOX-PAI 是咱们在阿里云机器学习平台 PAI 的开源计算机视觉代码库 EasyCV(https://github.com/alibaba/Ea…)中集成的 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. 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)：

<p align=center><img src=”https://p6-juejin.byteimg.com/tos-cn-i-k3u1fbpfcp/4a88fa8edb0340b9809984c9c4e734ff~tplv-k3u1fbpfcp-watermark.image?” alt=”（a）gsconv_part” /></p>

<p align=center><img src=”https://p6-juejin.byteimg.com/tos-cn-i-k3u1fbpfcp/2fe36ae518b341788bf8b1d0df598ef6~tplv-k3u1fbpfcp-watermark.image?” alt=”image.png” /></p>
所失去的试验后果如下（仅统计 Neck 局部的 Params, Flops）。能够看到 GSConv 对参数量进行了优化，且晋升了模型的性能，升高 3% 的推理速度能够换来 0.3mAP 的晋升。

Head

TOOD

参考 PPYOLOE，咱们同样思考利用 TOOD[9]算法中的 Task-aligned predictor 中的注意力机制（T-Head）别离对分类和回归特色进行加强。如下图所示，特色先通过解耦头的 stem 层 (1×1) 进行通道压缩，接着由通过重叠卷积层失去两头的特色层，再别离对分类和回归分支利用注意力机制进行特色的加强，来解耦两个工作。

咱们对重叠的中间层个数进行融化试验，每重叠能够进一步晋升肯定的精度，并带来速度上的一些损失。(下表的 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]](https://github.com/alibaba/Ea…) 自行组织相应的图像前解决 / 后处理过程，或间接应用咱们导出好的模型和接口)：咱们这里提供一个曾经导出好的检测模型，用户下载三个模型文件到本地 [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 介绍]](https://zhuanlan.zhihu.com/p/…)。

PAI-EasyCV 中对 Blade 进行了反对，用户能够通过 PAI-EasyCV 的训练 config 中配置相干 export 参数，从而对训练失去的模型进行导出。

这里咱们提供一个 PAI-Blade + PAI-EasyCV 社区版 V100 对应的镜像（cuda11.1/TensorRT8/cudnn8）：用户也能够基于 Blade 每日公布的镜像自行搭建推理环境 [[PAI-Blade 社区镜像公布]](https://github.com/alibaba/Bl…)

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/Ea…）是阿里云机器学习平台深耕一年多的计算机视觉算法框架，已在团体内外多个业务场景获得相干业务落地成绩，次要聚焦在自监督学习 /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/YO….

[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/ultralytic….

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