关于机器学习: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.