作者:忻怡、周楼、谦言、临在
导言
目标检测(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 |
mAPval |
SpeedV100 |
Download |
YOLOXs |
9.0 |
26.8 |
40.2 |
40.1 |
0.68 |
|
Meituan-YOLOv6-tiny |
15.0 |
36.7 |
41.3 |
– |
0.9 |
– |
PAI-YOLOXs |
15.9 |
36.8 |
41.5 |
41.4 |
0.7(+23%) |
|
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%) |
|
PAI-YOLOXs-ASSF-TOOD3 |
23.7 |
49.9 |
44.0 |
43.9(+1.5) |
1.15 |
|
YOLOXm |
25.3 |
73.8 |
46.3 |
46.3 |
1.50 |
从结果中可以看到,相比目前同水平(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 |
mAPval |
SpeedV100 |
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 |
SpeedV100 |
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 |
SpeedV100 |
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层(1×1)进行通道压缩,接着由通过堆叠卷积层得到中间的特征层,再分别对分类和回归分支利用注意力机制进行特征的增强,来解耦两个任务。
我们对堆叠的中间层个数进行消融实验,每堆叠可以进一步提升一定的精度,并带来速度上的一些损失。(下表的Params和Flops只统计了检测头部分。测速及精度基于的基线方法为以EfficientRep为backbone + ASFF进行特征增强。)
Model |
Params (M) |
Flops (G) |
mAPval |
SpeedV100 |
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.