WTPose 框架：基于 Transformer 的 Waterfall 模块提升姿态估计能力！

作者提出了 Waterfall Transformer 架构用于人体姿态估计 (WTPose)，这是一个单次通过、端到端可训练的框架，旨在进行多人姿态估计。
作者的框架利用基于 Transformer 的 Waterfall 模块，从多种主干阶段生成多尺度特征图。
该模块在级联架构中进行筛选操作，以扩展感受野并捕获局部和全局上下文，从而提高网络的整体特征表示能力。
作者在 COCO 数据集上的实验表明，使用修改后的 Swin 主干和基于 Transformer 的 Waterfall 模块的 WTPose 架构，在多人姿态估计方面优于其他 Transformer 架构。

1. Introduction

人体姿态估计是一项具有广泛实际应用的挑战性计算机视觉任务[27], [19]。基于卷积神经网络(CNNs)的深度学习方法提高了最先进的性能[23], [2], [5]。近期，视觉 Transformer 在计算机视觉任务中，包括姿态估计方面表现出色。

在本文中，作者提出了一种名为 WTPose 的“瀑布 Transformer”架构，该架构在一个灵活的框架中运行，以改进基本模型的表现。图 1 展示了使用 WTPose 进行的姿态估计示例。作者架构的一个关键特性是将多尺度的瀑布 Transformer 模块（WTM）集成进来，以增强视觉 Transformer 模型（如 Shifted Window（Swin）Transformer [14]）的表现。作者通过 WTM 的瀑布分支处理主干网络从多个层次提取的特征图。该模块基于膨胀注意力机制执行滤波操作，增加视野（Field-of-View, FOV），并捕捉局部和全局语境，从而显著提高性能。本文的主要贡献包括：

作者提出了一种新颖的 Waterfall Transformer 架构用于姿态估计，这是一种单次通过、端到端可训练的多尺度方法，适用于自上而下的多人 2D 姿态估计。该方法包含一个多尺度注意力模块，并采用扩张注意力机制，能够捕捉更大的感受野，从而获取全局和局部上下文。

作者的实验在 COCO 数据集上的结果显示，与类似的 Transformer 方法相比，该方法在姿态估计性能上有所提升。

2. Related Work

2.1. CNNs for pose estimation

随着深度卷积神经网络的发展，人体姿态估计已经取得了卓越的结果。Convolutional Pose Machine (CPM) [23] 架构包括多个阶段，逐步生成更为精细的关节检测。OpenPose 方法 [5] 引入了部分关联场（Part Affinity Fields），以处理单张图像中多人的姿态。Stacked Hour-glass 网络 [15] 使用重复的自底向上和自顶向下的处理，并在中间加入监督，以处理所有尺度的数据并捕捉与身体相关的最佳空间关系，从而实现精确的人体姿态估计。在此基础上，多上下文注意方法 [6] 设计了 Hourglass 剩余单元（HRUs），旨在生成具有较大感受野且不同语义 Level 的注意力图。此外，通过条件随机场（Conditional Random Fields, CRFs）进行后处理，生成局部和全局一致的人体姿态估计。

High-分辨率网络（HRNet）架构[18],[22]将高分辨率子网络和低分辨率子网络并行连接，在整个过程中保持高分辨率表示，并生成更准确且空间上更精确的姿态估计。多阶段姿态网络[11]与 HRNet[22]的工作方式类似，但它采用跨阶段特征聚合策略从早期阶段传递信息到后期阶段，并配备了从粗到细的监督机制。

UniPose (+) [2]、OmniPose [1] 和 BAPose [3] 方法提出了多种 Waterfall Atrous 空间池化(WASP)模块的变种，用于单人、多人顶下式和多人底上式的姿态估计。WASP 模块启发了 WTPose 中的瀑布 Transformer 模块，因为它显著提高了网络的多尺度表示能力和视野（FOV），并提取了包含更多上下文信息的特征，从而在无需后续处理的情况下获得更精确的姿态估计。

Vision Transformers for Pose Estimati

近年来，使用 Transformer 架构进行人体姿态估计的兴趣有所激增。在早期的工作中，使用 CNN Backbone 网络作为特征提取器，而将 Transformer 视为更优的解码器[25], [12]。TransPose[25]架构结合了基于 CNN 的 Backbone 网络的初始部分，用于从图像中提取特征，并采用标准的 Transformer 架构[21]，利用注意力层学习依赖关系并预测 2D 人体姿态的关键点。然而，TransPose 在建模关键点之间的直接关系方面存在局限性。

TokenPose[12]则明确地将每个关键点嵌入为一个 token，并通过自我注意交互学习视觉线索和约束关系。HRFormer[26]受到 HRNet[22]的启发，采用了多分辨率并行设计。它在茎部和第一阶段使用卷积，随后是 Transformer 块。Transformer 块在不重叠划分的特征图上执行自我注意，并使用 3x3 深度卷积进行划分映射间的跨注意。ViTPose[24]采用简单的非分层视觉 Transformer[8]作为 Backbone ，用于提取特征图。该架构然后采用反卷积层或双线性上采样解码器进行 2D 姿态估计。PoseFormer[28]提出了一种基于视频帧中的 2D 姿态序列的纯 Transformer 架构，用于 3D 姿态估计。

3. Waterfall Transformer

提出的瀑布式 Transformer 架构，如图 2 所示，是一种单次通过、端到端可训练的网络，该网络结合了一个修改后的 Swin Transformer Backbone 和作者基于 Transformer 的多尺度瀑布模块，用于多人姿态估计。Swin 中的切片分割层被两个 3 × 3 卷积（Stem）和 ResNet-101 的第一个残差块所取代，从而改进了 Swin 的特征表示。

WTPose 的处理 Pipeline 如图 2 所示。输入图像被送入由作者修改后的 Swin Transformer 构建的 Backbone 网络。来自多个阶段的 Swin 的多尺度特征图通过作者的瀑布流 Transformer 模块（WTM）进行处理，并送入解码器生成 K 个 Heatmap ，每个关节一个 Heatmap 。多尺度的 WTM 维持了特征图的高分辨率，并且能够对可见和遮挡的关节生成准确的预测。

作者瀑布 Transformer 模块的架构如图 3 所示。WTM 从 Disentangled Waterfall Atrous Spatial Pooling (D-WASP)模块[3], [4]中获得灵感，该模块利用空洞块和瀑布架构来增强多尺度表示。然而，与 D-WASP[3], [4]通过空洞卷积扩大视野不同，作者提出的方法采用了一个具有扩张和非扩张邻域注意力的空洞 Transformer 块来扩大视野。这个空洞 Transformer 基于 DiNAT[9]架构构建，包括扩张和非扩张邻域注意力。扩张的邻域注意力通过增加扩张率来扩大局部感受野，并执行稀疏全局注意力；而非扩张的邻域注意力则将每个像素的 Self-Attention 限制在其最近的邻居范围内。

为了应对由于分层 Backbone 结构导致的上下文和空间信息丢失问题，WTM 通过瀑布分支处理 Swin Backbone 网四个阶段的多尺度特征图。首先，WTM 模块利用双线性插值对来自第二、三、四阶段的低分辨率特征图进行上采样操作，使其与第一阶段的高分辨率特征图匹配，然后将所有特征图结合起来生成增强联合估计所需的多尺度特征表示。

输出特征图 = 0 被送入瀑布型 Transformer 块(Waterfall Transformer Blocks, WTB)，通过逐级增强的过滤级 cascade 来扩展视场。每个 WTB 包含两种类型的注意力机制：膨胀多头邻域自注意力(Dilated Multi-Head Neighborhood Self-Attention, D-MHSA)，其后接多层感知机(Multi-Layer Perceptron, MLP)，用于捕捉全局上下文；非膨胀多头邻域自注意力(Nondilated Multi-Head Neighborhood Self-Attention, N-MHSA)，其后接 MLP，用于捕捉局部上下文。

瀑布模块旨在创建一个瀑布 Stream 处理过程，首先对输入进行初步处理，然后生成新的分支。与级联方法不同，WTM 通过结合所有 WTB 分支的所有流以及多尺度表示中的深度可分离池化（DWP）层，实现了更为全面的集成。

4. Experiments

作者在 Common Objects in Context（COCO）[13]数据集上进行了多人姿态估计实验。COCO 数据集[13]包含了超过 20 万张野外图片，并且包含了 25 万个实例的人类目标。作者使用 COCO 训练集 2017 版本，包含 5.7 万张图片和 15 万个实例的人目标进行 WTPose 的训练，并在包含 5 千张图片的 val 2017 验证集上进行验证。标注的姿态包含 17 个关键点。

作者采用了 Swin Base（Swin-B）Transformer 作为主干网络，并使用来自[14]预训练权重进行初始化。默认情况下，Swin-B 架构采用窗口大小为 7。对于 WTM 模块，作者尝试了多种膨胀率，并发现交替使用膨胀窗口带来的大感受野和非膨胀窗口带来的小感受野，可以提高预测效果。作者将 WTB 块的膨胀率设置为(2,1)，(4,1)，(4,1)，(8,1)，并保持窗口大小为 7。

作者的模型在包含 4 块 A100 GPU 的系统上使用 mmpose 代码库[7]进行训练，批量大小设为 32。作者采用 mmpose 中的默认训练设置来训练 WTPose，并使用 AdamW [16]优化器，学习率设为 5e-4。作者的模型总共训练了 210 个 epoch，在第 170 个和第 200 个 epoch 时学习率降低 10 倍进行衰减。

4.1. Experimental results on the COCO dataset

作者在 COCO 数据集上进行了训练和测试，并将 WTPose 与 Swin 框架进行了比较，详见表 1。作者的 WTPose 模型参数量比 Swin-B 多 130 万，平均精度和平均召回率分别提高了 1.2%和 0.9%。与 Swin-L 相比，WTPose 大约小 54%，但在平均精度和平均召回率方面仍分别优于 Swin-L 0.8%和 0.6%。瀑布 Transformer 模块改进了特征图，提高了关键点检测的准确性。

作者在 WTPose 中进行了消融研究以探究各个组件的效果。表 2 展示了使用 Swin-B 主干和分辨率为 384×288 的输入图像的各种配置结果。作者设置窗口大小为 7×7，在每个注意力层中分配 8 个头，并选择扩张率 2、4、4、8 来增加不同 WTB 块的感知区域大小。扩张率为 1、2、4、8 时的感知区域大小分别为 7×7、13×13、25×25 和 49×49。首先，作者分别针对每个 WTB 尝试了扩张率 2、4、4、8 的情况，在每个 WTB 中进行指定扩张率的一次扩张多头自注意力操作。接着，作者对每个 WTB 分别使用了一次扩张和未扩张的多头自注意力机制，并将扩张率设为(2, 1)、(4, 1)、(4, 1)、(8, 1)。作者的主要观察结果是：(i) 将瀑布流 Transformer 模块与修改后的 Swin 主干结合能够改进特征表示；(ii) 在 Swin-B 的开始添加一个 Stem 和 ResNet Bottleneck 可以进一步增强主干的能力。