YOLOX-PAI：YOLOXを加速し、YOLOV6よりも速く、強く

著者: Xin Yi, Zhou Lou, Qian Yan, Presence

序文

物体検出は、画像内の対象物体を見つけて識別することを目的としています. これは、コンピュータービジョンの分野で常にホットな研究トピックであり、自動運転や対象追跡などのタスクの基礎でもあります. 近年、優れたターゲット検出アルゴリズムが継続的に登場しており、その中でもシングルステージのYOLOシリーズは、その高効率とシンプルさにより、ターゲット検出アルゴリズムの分野で常に重要な分野でした。2021 年、Megvii は YOLOX[1] アルゴリズムを提案しました。これは、柔軟で展開可能なコンポーネントを使用して、速度と精度の面で新しいベースラインを構築し、業界に深く愛されました。このホワイトペーパーでは、Alibaba Cloud PAI-EasyCV フレームワークに基づいて YOLOX アルゴリズムを再現し、実際に YOLOX の精度を向上させるための実用的な手法を探り、アリババコンピューティングプラットフォーム PAI によって開発された PAI-Blade 推論高速化フレームワークを使用してモデルのパフォーマンスをさらに最適化します。コミュニティでの多くの YOLOX 改善技術の再現と調査の後、YOLOX のパフォーマンスはさらに改善され、その速度と精度は現在の 40~50mAP SOTA の YOLOv6 よりも優れています。同時に、PAI-EasyCV は、コミュニティが YOLOX-PAI の機能をすばやく体験できるように、効率的で簡潔なモデル展開とエンドツーエンドの推論インターフェイスを提供します。

私たちの仕事の貢献を要約すると：

Apache ライセンスのトレーニング/最適化/推論コードベースとミラーのセットを提供します。これにより、現在のコミュニティで 40+MAP のオーダーで最速のターゲット検出モデルを実現できます (YOLOV6 と比較して +0.4mAP/加速度 13~20%) .
YOLOX に関連する改善手法とアブレーション実験を調査し、比較的有用なものをいくつか選択し、40/0.7ms (YOLOXS) ~ 47.6/1.5ms (YOLOXM) の間でモデルを埋め、提供された方法でそれらを構成しました。
ターゲット検出のエンドツーエンド推論の速度を柔軟にカプセル化し、最適化しました. V100 のエンドツーエンド推論は、元の YOLOX の 9.8ms と比較して 3.9ms であり、ほぼ 250% 高速です。ターゲットの検出と推論のタスクをすばやく完了します。

この記事では、関連する改善点とアブレーションの実験結果を 1 つずつ紹介し、PAI-Blade を使用して PAI-EasyCV に基づいてモデルの推論プロセスを最適化する方法と、PAI-EasyCV をモデルのトレーニング、検証、展開とエンドツーエンドの推論。シンプルで効率的なビジョンアルゴリズムの開発と展開タスクのために、PAI-EasyCV と PAI-Blade に注意を払い、使用してください。

PAI-EasyCV プロジェクトのアドレス: https://github.com/alibaba/EasyCV

PAI-BladeDISC プロジェクトのアドレス: https://github.com/alibaba/BladeDISC

YOLOX-PAI-アルゴリズムの改善

YOLOX-PAI は、Alibaba Cloud 機械学習プラットフォーム PAI のオープンソースコンピュータービジョンコードベースであるEasyCV ( https://github.com/alibaba/EasyCV ) に統合された YOLOX アルゴリズムです。読者がYOLOXアルゴリズムを知らない場合は、自分でそれを学ぶことができます (参照:リンク). このセクションでは、主にYOLOXアルゴリズムに基づく私たちの改善を紹介します.

YOLOX アルゴリズムの分析と検出技術の調査を通じて、次の 4 つの方向からオリジナルの YOLOX を最適化しました。

バックボーン: repvgg バックボーン
ネック：gsconv / asff
頭: toods / rtoods
Loss : siou / giou

アルゴリズムの改善に基づいて、PAI-Blade を使用して最適化されたモデルを推論および最適化し、次の PAI-YOLOX モデルが開発されました。スクリーニングの効果的な改善と既存の主流アルゴリズムとの比較結果は次のとおりです。

(-ASFF は NeckASFF の使用を表し、-TOODN は元の YOLOXHead を置き換えるために TOODHead の N 中間層を使用することを表します)

モデル	パラメータ (M)	フロップ (G)	mAP 値0.5 :0.95 （672）	mAP 値0.5 :0.95 （640）	速度V100 fp16 b32(ms)	ダウンロード
YOLOX	9.0	26.8	40.2	40.1	0.68	モデル-ログ
Meituan-YOLOv6-tiny	15.0	36.7	41.3	-	0.9	-
パイヨロックス	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	モデル-ログ

この結果からわかるように、SOTA が同じレベル (1ms 以内) の現行の YOLO6 モデルと比較すると、上記の改善を統合した YOLOX-PAI は、同じ精度/速度の条件下で、速度/精度において一定の利点があります。

評価については、次の点に注意する必要があります。

YOLOV6 リリースのテスト速度にはデコードと nms が含まれていないため、公正な比較のために関連する速度測定設定もオフにしました。(上記の表に示されている結果は、Blade によって最適化された対応するモデルについて、bs32 での画像モデルの平均前方推論時間を計算したものです。エンドツーエンドの推論時間については (画像の前処理と後処理の合計時間を含みます) )、「PAI-EasyCV エクスポート」セクションを参照)
YOLOV6 リリースの精度は、トレーニングの過程でテストされ、shape=672 の場合もありますが、速度測定は image_size=(640, 640) にエクスポートされた後に完了します. 実際には、関連する学生もいます。 640 未満の YOLOV6 を補足しているコミュニティ。テストの精度なので、上の表に 2 つのテストの精度を示します。
EasyCV の Predictor インターフェイスを使用して、画像入力から結果までの予測に関連するモデル予測を読み込みます. 前処理と NMS 後処理が含まれているため、対応する時間は遅くなります. 詳細については、エンドツーエンドの最適化結果を参照してください.

以下では、各モジュールの改善とアブレーション実験について詳しく説明します。

背骨

RepConv

最近では、YOLO6 [2] や PP-YOLOE [3] などのアルゴリズムによって、CSPNet [4] のバックボーンネットワークが改良され、RepVGG [5] のアイデアに基づいて、モデルを作成するために再パラメータ化可能なバックボーンネットワークが設計されています。より効率的な推論、パフォーマンス。この考え方に従い、YOLO6 の EfficientRep を使用して、YOLOX の元の CSPDarkNet-53 バックボーンネットワークを置き換えます。得られた実験結果は、次のように YOLO6 関連モデルと比較されます (YOLOX-Rep は、バックボーンネットワークとして EfficientRep を使用する YOLOX モデルを表します)。

モデル	パラメータ (M)	フロップ (G)	mAP 値0.5 :0.95 （672）	mAP 値0.5 :0.95 （640）	速度V100 fp16 b32 (ミリ秒)
YOLOX	9.0	26.8	40.2	40.1	0.68
YOLOv6-小さな	15.0	36.7	41.3	-	0.9
パイヨロックス	15.9	36.8	41.5	41.4	0.7

RepVGG 構造のネットワーク設計はパラメーターと計算の量を増加させますが、実際の推論速度はより有利であるため、構成可能なバックボーンとして YOLO6 EfficientRep を選択します。

首

バックボーンネットワークの置き換えに基づいて、ネック部分を 2 つの側面から検討しました。

ASSF[6]: PAN 出力特徴次元を変換した後の SE-Attention 特徴融合に基づく特徴強化に基づいて、パラメーターの量と精度が大幅に改善され、推論速度が部分的に低下します。

1. ASSF-Sim: パラメーターの数が少ない機能融合の実装を選択し、より少ない数のパラメーター (ASFF: 5M -> ASFF-Sim: 380K) を使用して、84% の精度向上 (+0.98map -> + 0.85 マップ)。ただし、この改善により推論速度が遅くなります. 将来的には、このOPに対応するプラグインを実装して、推論の高速化を完了します.

GSNeck[7]: Neck 情報融合の DW Conv に基づいて、Neck パラメーターの量を減らし、精度をわずかに向上させ、推論速度も低下させます。

ASFF情報融合

ASFF は、異なる PAN と異なる特徴マップ間の情報相互作用を通じて、アテンションメカニズムを使用して情報の融合とネック部分の強化を完了します。

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

さまざまな機能マップの融合メカニズムは次のとおりです。

モデル	パラメータ (M)	フロップ (G)	mAP 値0.5 :0.95	速度V100 fp16 b32 (ミリ秒)
パイヨロックス	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 (スピードアップが必要)

GSConvNeck

DWConv を使用してパラメーターの量を削減することは一般的な手法ですが、YOLOX では、GSconv は新しいタイプの軽量畳み込みを設計して、モデルのパラメーターと計算を削減します。Depth-wise Separable Convolution (DSC) の計算におけるチャネル情報の分離の欠点を解決するために、GSConv (下の図に示すように) はシャッフルを使用して、標準畳み込み (SC) と DSC によって得られた特徴マップを融合します。 SCの出力がDSCに完全に統合されるようにします。

また、GSConv の原文では、GSconv をネットワーク全体で使用すると、ネットワークの深さが大幅に深まり、モデルの推論速度が低下することが指摘されています。チャネル情報の次元が最大で、空間情報の次元が最小です。YOLOX では GSConv を使用してモデルを最適化します. 特に、実験を別々に行うために 2 つのスキームを使用します (a: Neck の一部のみが GSConv を使用、b: Neck のすべてのモジュールが GSConv を使用):

（a）gsconv_part (b)gsconv

得られた実験結果は次のとおりです (Neck 部分の Params と Flops のみがカウントされます)。GSConv がパラメーターの量を最適化し、モデルのパフォーマンスを向上させることがわかります.3% の推論速度の低下は、0.3mAP の向上と交換できます.

モデル	パラメータ (M)	フロップ (G)	mAP 値0.5 :0.95	速度V100 fp16 b32 (ミリ秒)
Yoloxオリジナル	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.