TPAMI 2024 | HIRI-ViT:具有高分辨率输入的缩放视觉Transformer
HIRI-ViT:具有高分辨率输入的缩放视觉Transformer
作者:Ting Yao, Yehao Li, Yingwei Pan,Tao Mei
摘要
视觉变换器 (ViT) 和卷积神经网络 (CNN) 的混合深度模型已成为视觉任务的强大主干类别。扩大此类混合主干的输入分辨率自然会增强模型容量,但不可避免地会遭受二次方增长的沉重计算成本。相反,我们提出了一种具有高分辨率输入的新型混合主干 (即 HIRI-ViT),它将流行的四级 ViT 升级为针对高分辨率输入量身定制的五级 ViT。HIRI-ViT 建立在将典型的 CNN 操作以经济高效的方式分解为两个并行的 CNN 分支的开创性思想之上。一个高分辨率分支直接将主要的高分辨率特征作为输入,但使用较少的卷积运算。另一个低分辨率分支首先执行下采样,然后对这些低分辨率特征进行更多的卷积运算。在识别任务(ImageNet-1K 数据集)和密集预测任务(COCO 和 ADE20K 数据集)上的实验证明了 HIRI-ViT 的优越性。更值得注意的是,在相当的计算成本(∼5.0 GFLOPs)下,HIRI-ViT 在 ImageNet 上以 448×448 输入实现了迄今为止已公布的最佳 Top-1 准确率 84.3%,这比以 224×224 输入获得的 iFormer-S 的 83.4% 提高了 0.9%。
关键字
图像识别、自注意力学习、视觉转换器。
引言
受自然语言处理 (NLP) [1] 中占主导地位的 Transformer 结构的启发,计算机视觉 (CV) 领域见证了用于设计视觉主干的 Vision Transformer (ViT) 的兴起。这种趋势在图像/动作识别 [2]、[3]、[4]、[5] 和物体检测 [6] 等密集预测任务中最为明显。这些成功中的许多可以归因于通过传统 Transformer 块中的自注意机制对输入视觉标记之间的长距离交互进行灵活建模。最近,几项并发研究 [7]、[8]、[9]、[10]、[11] 指出,直接在视觉标记序列上使用纯 Transformer 块并不是最优的。这种设计不可避免地缺乏 2D 区域结构建模的正确归纳偏差。为了缓解这一限制,他们引领了将卷积神经网络(CNN)的二维归纳偏差注入 ViT 的新浪潮,产生了 CNN+ViT 混合主干。
CNN 主干设计中的一种常见做法是扩大网络深度/宽度/输入分辨率 [12]、[13],从而通过在输入中捕获更细粒度的模式来增强模型容量。在精神上类似,我们的工作希望深入研究使用高分辨率输入扩展 CNN+ViT 混合主干的过程。然而,与 CNN 主干的扩展类似,简单地扩大流行的 ViT 主干的输入分辨率将带来实际挑战,尤其是计算成本的急剧增加。以广泛采用的 ViT 主干 Swin Transformer [14] 为例,当直接将输入分辨率从 224×224 扩大到 384×384 时,ImageNet-1K 上的 top-1 准确率明显从 83.5% 提高到 84.5%。然而,如图 1 所示,具有 384×384 个输入(GFLOPs:47.0,推理时间:3.95 毫秒)的 Swin Transformer 的计算成本明显高于具有 224×224 个输入(GFLOPs:15.4,推理时间:1.17 毫秒)的 Swin Transformer。
鉴于此问题,我们的核心问题是——是否有一种原则性的方法可以在保持相当的计算开销的同时扩展具有高分辨率输入的 CNN+ViT 混合主干?为此,我们设计了一系列针对高分辨率输入的五阶段 Vision Transformers,其中包含早期阶段的两个分支构建块,以寻求性能和计算成本之间的更好平衡。具体而言,重新塑造的 stem/CNN 块中的关键组件是并行组合高分辨率分支(对高分辨率输入的卷积运算较少)和低分辨率分支(对低分辨率输入的卷积运算较多)。这种双分支设计取代了 stem/CNN 块中具有标准卷积运算的单个分支。通过这样做,我们不仅保留了高分辨率输入的增强模型容量,而且还通过每个分支的轻量级设计显着降低了计算成本。如图 1 所示,通过将输入分辨率从 224×224 扩大到 384×384,我们的 HIRI-ViT 获得了明显的性能提升,但计算成本仅略有增加(GFLOP:8.2 到 9.3,推理时间:0.84 毫秒到 1.04 毫秒)。即使将输入分辨率扩大到 768×768,HIRI-ViT 也能带来显着的性能提升,同时所需的计算成本低于 Swin Transformer。
通过将这种双分支设计集成到 CNN+ViT 混合主干中,我们提出了一种新的原则性五阶段视觉主干,即 HIRI-ViT,它有效地扩展了具有高分辨率输入的 Vision Transformer。具体来说,我们首先通过将单个 CNN 分支分解为两个并行分支(即高分辨率和低分辨率分支)来升级典型的 Conv-stem 块,从而产生高分辨率 stem (HRstem) 块。接下来,通过使用所提出的双分支设计替换 CNN 分支来重塑早期阶段的 CNN 块。这种新的高分辨率块(HR-block)触发了高分辨率输入的经济高效的编码。
这项工作的主要贡献是提出在保持有利的计算成本的同时扩展具有高分辨率输入的 CNN+ViT 混合主干。这也导致了将高分辨率输入上的典型 CNN 操作分解为两个并行的轻量级 CNN 分支的优雅设计。通过对多个视觉任务(例如,图像识别、对象检测和实例/语义分割)进行的大量实验,我们证明了我们的新 HIRI-ViT 主干与最先进的 ViT 和 CNN 主干相比在可比计算成本下的优势。
预备知识
传统的多阶段视觉变换器(M-ViT)[14],[38],[39],[40]通常由一个主干加四个阶段组成,如ConvNets [17],[21],[23](见图2(a))。具体而言,首先利用主干层将输入图像(分辨率:224×224)分割成块。每个块被视为“视觉标记”,并将进一步输入到以下阶段。每个阶段包含多个Transformer块,每个Transformer块由一个多头自注意模块(MHA)和一个前馈网络(FFN)组成。通常,在每两个阶段之间插入一个下采样层(DS)以合并输入的“视觉标记”(即降低特征图的分辨率)并同时扩大其通道维度。最后,采用分类器层根据最后一个特征图预测概率分布。
多头自注意力:多头自注意力旨在捕获视觉标记之间的长距离依赖关系。从技术上讲,让 X ∈ R n × D X \in \mathbb{R}^{n \times D} X∈Rn×D 表示视觉标记的特征,其中 n = H × W n = H \times W n=H×W 是视觉标记的数量, H / W / D H/W/D H/W/D 分别表示高度/宽度/通道数。输入 X X X 首先线性转换为查询 Q ∈ R n × D Q \in \mathbb{R}^{n \times D} Q∈Rn×D、键 K ∈ R n × D K \in \mathbb{R}^{n \times D} K∈Rn×D 和值 V ∈ R n × D V \in \mathbb{R}^{n \times D} V∈Rn×D,然后进一步沿通道维度分解为 N h N_h Nh 个头/部分。通过将第 j j j 个头的查询、键和值表示为 Q j ∈ R n × D h Q_j \in \mathbb{R}^{n \times D_h} Qj∈Rn×Dh、 K j ∈ R n × D h K_j \in \mathbb{R}^{n \times D_h} Kj∈Rn×Dh 和 V j ∈ R n × D h V_j \in \mathbb{R}^{n \times D_h} Vj∈Rn×Dh( D h D_h Dh:每个头的维度),自注意力模块的运行方式如下:
M u l t i H e a d ( Q , K , V ) = C o n c a t ( h e a d 0 , . . . , h e a d N h ) W O , h e a d j = A t t e n t i o n ( Q j , K j , V j ) , A t t e n t i o n ( Q j , K j , V j ) = S o f t m a x ( Q j K j T D h ) V j , \begin{align*} \mathbf{MultiHead}(Q, K, V) &= \mathbf{Concat}(head_0, ..., head_{N_h})W^O, \\ head_j &= \mathbf{Attention}(Q_j, K_j, V_j), \\ \mathbf{Attention}(Q_j, K_j, V_j) &= \mathbf{Softmax}(\frac{Q_jK^T_j}{\sqrt{D_h}})V_j, \tag{1} \end{align*} MultiHead(Q,K,V)headjAttention(Qj,Kj,Vj)=Concat(head0,...,headNh)WO,=Attention(Qj,Kj,Vj),=Softmax(DhQjKjT)Vj,(1)
其中 W O W^O WO 是权重矩阵, C o n c a t ( ⋅ ) \mathbf{Concat}(\cdot) Concat(⋅) 是连接操作。考虑到自注意力的计算成本与 token 数量成二次方关系,通常对键/值应用空间缩减以减少计算/内存开销 [8],[40]。
前馈神经网络:原始前馈神经网络 [1],[2] 由两个全连接( F C \mathbf{FC} FC)层组成,中间有一个非线性激活层:
FFN ( X ) = FC ( σ ( FC ( X ) ) ) , (2) \textbf{FFN}(X) = \textbf{FC}(\sigma(\textbf{FC}(X))), \tag{2} FFN(X)=FC(σ(FC(X))),(2)
其中 σ \sigma σ 表示非线性激活。受 [8]、[40] 的启发,我们通过额外的卷积运算升级 F F N \mathbf{FFN} FFN以施加 2D 归纳偏差,从而产生卷积前馈网络 ( C F F N \mathbf{CFFN} CFFN)。此 C F F N \mathbf{CFFN} CFFN 的整体操作总结如下
C F F N ( X ) = F C ( D W C o n v ( z ) + z ) , z = σ ( F C ( X ) ) , (3) \mathbf{CFFN}(X) = \mathbf{FC}(\mathbf{DWConv}(z) + z), z = \sigma(\mathbf{FC}(X)), \tag{3} CFFN(X)=FC(DWConv(z)+z),z=σ(FC(X)),(3)
其中 D W C o n v \mathbf{DWConv} DWConv 表示深度卷积。
方法
在本文中,我们的目标是设计一个原则性的 Transformer 结构(即 HIRI-ViT),从而能够以低成本的方式扩展具有高分辨率输入的 Vision Transformer。为此,我们将典型的四阶段 M-ViT 升级为一个新的五阶段 ViT 系列,其中包含早期阶段的两个分支构建块,将单分支 CNN 操作分解为两个并行的 CNN 分支。这种方式可以带来针对高分辨率输入量身定制的有利计算开销。图 2(b) 展示了我们的 HIRI-ViT 的整体架构。
A. 高分辨率Stem
传统 Vision Transformer 中的 stem 层设计可以简单地分为两个维度:ViT-stem 和 Convstem [41], [42]。如图 3(a) 所示,ViT-stem 实现为单个步幅卷积层(例如,步幅 = 4,内核大小 = 7 [40]),旨在将输入图像划分为块。最近,[10], [42] 表明,用几个堆叠的 3×3 卷积(即图 3(b) 中所示的 Conv-stem)替换 ViT-stem 可以稳定网络优化过程,同时提高峰值性能。对于典型的输入分辨率(例如 224×224),Conv-stem 会导致参数和 GFLOP 略有增加。然而,当输入分辨率显著增加(例如 448×448)时,Convstem 的 GFLOP 会比 ViT-stem 重得多。为了缓解这些问题,我们通过将单分支 Conv-stem 重塑为两个并行的 CNN 分支,设计了一个新的高分辨率 stem 层(图 3© 中的 HR-stem)。这样的设计不仅保留了 Conv-stem 的高模型容量,而且在高分辨率输入下消耗了有利的计算成本。
从技术上讲,HR-stem 首先利用步幅卷积(步幅 = 2,内核大小 = 3)对输入图像进行下采样,就像在 Conv-stem 中一样。之后,下采样的特征图被馈送到两个并行分支(即高分辨率和低分辨率分支)。高分辨率分支包含轻量级深度卷积,后跟步幅卷积。对于低分辨率分支,首先使用步幅卷积对特征图进行下采样。然后应用两个卷积(3×3 和 1×1 卷积)来施加归纳偏差。最后,通过聚合两个分支实现 HR-stem 的输出,并通过批量归一化(BN)进一步对总和进行归一化。
B. 高分辨率块
鉴于混合主干中前两个阶段的输入分辨率较大,Transformer 块的计算成本相对较高。为了解决这一限制,我们用新的高分辨率块(HR 块)替换了前两个阶段的 Transformer 块,从而能够在早期阶段以低成本对高分辨率输入进行编码。具体而言,与 HR-stem 类似,HR 块由两个并行的分支组成。轻量级高分辨率分支捕获高分辨率输入上的粗级信息,而低分辨率分支利用更多卷积运算来提取低分辨率输入上的高级语义。图 4 描绘了 HR 块的详细架构。
具体来说,高分辨率分支实现为轻量级深度卷积。对于低分辨率分支,首先使用具有 BN 的步幅深度卷积(步幅 = 2,内核大小 = 3)对输入特征图进行下采样。然后,在低分辨率特征图上应用前馈操作(即,两个 FC 之间有一个激活)。之后,通过重复对低分辨率输出进行上采样,然后与高分辨率输出进一步融合。
C. 反向残差下采样
在传统的 M-ViT 中,空间下采样是通过单个步幅卷积(例如,步幅 = 2,内核大小 = 3 [40])执行的,如图 5(a)所示。受到 ConvNets [17]、[43] 的启发,我们设计了一个具有两个并行分支的更强大的下采样层,即倒置残差下采样 (IRDS)。具体来说,对于具有高分辨率输入的前两个阶段,我们采用 IRDS-a(图 5(b))进行下采样。IRDS-a 首先使用步幅 3×3 卷积来扩展维度并减小空间大小,然后使用 1×1 卷积来缩小通道维度。对于最后两个下采样层,我们利用 IRDS-b(图 5(c)),它类似于倒置残差块 [44]。不同之处在于,我们仅在第一次卷积后应用归一化和激活操作。请注意,我们添加了额外的下采样捷径来稳定训练过程。
D. 块规范化
ConvNets [17], [21] 通常使用 BN 来稳定训练过程。BN 也可以合并到卷积操作中以加速推理。相比之下,Vision Transformer 主干倾向于使用层规范化 (LN) 对特征进行规范化。LN 对于训练批次大小较小的密集预测任务(例如,对象检测和语义分割)更为友好,因为它与批次大小无关。与 BN 相比,LN 也可以带来略好的性能 [23]。尽管如此,LN 会导致高分辨率输入的计算成本更高。因此,我们在前三个阶段使用高分辨率输入的 BN,而在最后两个阶段使用 LN 低分辨率输入。此外,我们还将 CFFN 块的 LN 替换为 BN。通过这样做,推理过程可以加速 7.6%,同时保持性能。
E. EMA 蒸馏
在训练过程中,指数移动平均 (EMA) [58] 已被广泛采用来稳定和改进 ConvNets [13] 和 ViT [25] 的训练程序。然而,传统 EMA 中的消息传递是单向的,即教师网络由 EMA 根据学生网络的参数进行更新,从而导致次优解决方案。为了触发教师和学生网络之间的双向消息交互,我们提出了一种新的 EMA 蒸馏策略来训练 HIRI-ViT。EMA 蒸馏还利用从教师网络中学习到的概率分布来指导学生网络的训练。与传统知识蒸馏 [25] 相比,我们的 EMA 蒸馏不依赖任何额外的大规模预训练网络。
从技术上讲,给定一对训练样本 ( x a , y a ) (x_a, y_a) (xa,ya) 和 ( x b , y b ) (x_b, y_b) (xb,yb),通过 Cutmix/Mixup 生成新的训练样本 x ~ , y ~ \tilde{x}, \tilde{y} x~,y~,并将其进一步输入到学生网络 F s \mathcal{F}^s Fs 中进行网络优化。以 Cutmix 为例,新的样本 x ~ , y ~ \tilde{x}, \tilde{y} x~,y~ 可以生成为
x ~ = M ⊙ x a + ( 1 − M ) ⊙ x b , y ~ = λ y a + ( 1 − λ ) y b , λ = ∑ M H W , \begin{align*} \tilde{x} = M \odot x_a + (1-M) \odot x_b, \\ \tilde{y} =\lambda y_a + (1 - \lambda)y_b, \lambda = \frac{\sum M}{HW}, \tag{4} \end{align*} x~=M⊙xa+(1−M)⊙xb,y~=λya+(1−λ)yb,λ=HW∑M,(4)
其中 M ∈ { 0 , 1 } H × W M \in \{0,1\}^{H \times W} M∈{0,1}H×W 表示矩形掩码, ⊙ \odot ⊙ 表示元素乘法。在 EMA 蒸馏中,我们将原始样本 x a x_a xa 和 x b x_b xb 输入教师网络 F t \mathcal{F}_t Ft(删除了最后一个特征图上的平均池化操作),分别得到概率分布图 P a P_a Pa 和 P b P_b Pb。接下来,我们为生成的样本 x ~ \tilde{x} x~ 分配一个混合目标标签:
P ~ = M ⊙ P a + ( 1 − M ) ⊙ P b , y ^ = A v g P o o l ( P ~ ) , (5) \tilde{P} = M \odot P_a + (1-M) \odot P_b, \hat{y} = \mathbf{AvgPool}(\tilde{P}), \tag{5} P~=M⊙Pa+(1−M)⊙Pb,y^=AvgPool(P~),(5)
其中 A v g P o o l \mathbf{AvgPool} AvgPool 表示沿空间维度使用平均池化操作。之后,我们将 Cutmix 和教师网络生成的两个目标标签组合起来: y ˉ = α y ~ + ( 1 − α ) y ^ \bar{y} = \alpha \tilde{y} + (1-\alpha)\hat{y} yˉ=αy~+(1−α)y^,其中 α \alpha α 是权衡参数。最后,我们在训练期间将混合样本 x ~ , y ˉ \tilde{x}, \bar{y} x~,yˉ 馈入学生网络 F s \mathcal{F}^s Fs。这样,从教师网络获得的知识被额外利用,以促进学生网络的学习。
F. 架构细节
表 I 详细介绍了我们的 HIRI-ViT 系列的架构。遵循现有 CNN+ViT 混合主干 [14]、[26] 的基本网络配置,我们构建了三种不同模型大小的 HIRI-ViT 变体,即 HIRI-ViT-S(小尺寸)、HIRI-ViT-B(基础尺寸)和 HIRI-ViT-L(大尺寸)。具体来说,HIRI-ViT 的整个架构由一个 HR-stem 层和五个阶段组成。对于具有高分辨率输入的前两个阶段,我们用轻量级高分辨率块替换传统的 Transformer 块,以避免巨大的计算开销。对于第三阶段,我们仅利用 CFFN 块来处理中等分辨率特征图。与传统的 Vision Transformer 类似,我们在最后两个阶段使用具有低分辨率输入的 Transformer 块。对于每个阶段 i i i, E i E_i Ei、 C i C_i Ci 和 H D i HD_i HDi 分别表示前馈层的扩展率、通道维度和头数量。
实验
我们在四个视觉任务上评估了我们的 HIRI-ViT:图像分类、对象检测、实例分割和语义分割。具体来说,HIRI-ViT 首先在 ImageNet-1K [59] 上从头开始训练,用于图像分类任务。接下来,我们对预训练的 HIRI-ViT 进行微调,以完成其余三个下游任务:在 COCO [60] 上进行对象检测和实例分割,在 ADE20K [61] 上进行语义分割。
ImageNet-1K 上的图像分类
设置:ImageNet-1K 数据集包含 1,000 个对象类的 128 万张训练图像和 50,000 张验证图像。在训练期间,我们采用 [14]、[40] 中常见的数据增强策略:随机裁剪、随机水平翻转、Cutmix [65]、Mixup [66]、随机擦除 [67]、RandAugment [68]。整个网络通过 AdamW [69] 在 8 个 V100 GPU 上进行优化,包括 300 个带有余弦衰减学习率调度程序 [70] 的 epoch 和 5 个线性预热 epoch。批量大小、初始学习率和权重衰减分别设置为 1,024、0.001 和 0.05。我们报告了 ImageNet-1K 验证集上的 Top-1/5 准确率,以及 ImageNet V2 匹配频率测试集上的 Top-1 准确率(即 V2 Top-1)[71],如 [25] 所示。
性能比较:表 II 显示了我们的 HIRI-ViT 系列与现有 CNN/ViT 主干之间的性能比较。值得注意的是,所有基线都采用典型分辨率输入(224×224),而我们的 HIRIViT 系列使用高分辨率输入(448×448)扩展了 Vision Transformer。总体而言,在每组可比的计算成本下,我们的 HIRI-ViT(448×448)在所有模型尺寸上都实现了与最先进的主干相比一致的性能改进。值得注意的是,对于大模型尺寸(GFLOP 超过 11.7 的主干),我们的 HIRI-ViT-L(448×448)的 Top-1 准确率为 85.7%,这比最佳竞争对手 MaxViT-L(85.2%)的绝对性能提升了 0.5%。尽管 HIRI-ViT-L 的输入分辨率明显大于 MaxViT-L,但我们的 HIRI-ViT-L(448×448)所需的 GFLOP 比 MaxViT-L 要少,并且表明推理速度更快,吞吐量几乎翻倍。这样的结果清楚地表明,我们的 HIRI-ViT 在性能和计算成本之间实现了更好的平衡,尤其是针对高分辨率输入。值得注意的是,当输入典型分辨率输入(224×224)时,我们的 HIRI-ViT 在每种模型尺寸下都实现了与最先进的主干网络相当的性能,同时所需的计算成本却要低得多。例如,在大模型尺寸组中,我们的 HIRI-ViT-L(224×224)的 Top-1 准确率为 85.3%,相应的吞吐量为每秒 660 张图像,这比最佳竞争对手 MaxViT-L(Top-1 准确率:85.2%,吞吐量:每秒 241 张图像)快得多。结果再次证实了我们的 HIRI-ViT 具有成本效益的设计。
更高分辨率下的性能比较:表 III 说明了我们的 HIRI-ViT 系列与具有更大输入图像尺寸的其他最先进的视觉主干之间的比较。对于具有更高分辨率输入(768 × 768)的升级版 HIRI-ViT,我们在 8 个 V100 GPU 上采用 AdamW 优化器 [69],动量为 0.9,初始学习率为 1.0 e − 5 1.0e^{-5} 1.0e−5,权重衰减为 1.0 e − 8 1.0e^{-8} 1.0e−8。优化过程包括 30 个epoch,使用余弦衰减学习率调度程序 [70]。同样,对于每个具有可比计算成本的组,我们的 HIRI-ViT 与其他高分辨率视觉主干相比始终获得性能提升。这些结果清楚地验证了我们提出的针对高分辨率输入量身定制的五级 ViT 主干的有效性。这种设计新颖地将典型的 CNN 操作并行分解为高分辨率和低分辨率分支,因此即使在具有较高分辨率输入的设置下也能保持良好的计算成本。
B. COCO 上的目标检测和实例分割
设置:我们在 COCO 数据集上执行对象检测和实例分割任务。我们采用 [14]、[40] 中的标准设置,并在 COCO-2017 训练集(约 118 K 张图像)上训练所有模型。最后在 COCO-2017 验证集(5 K 张图像)上对学习到的模型进行评估。我们使用两个主流检测器(RetinaNet [72] 和 Mask R-CNN [73])进行对象检测和实例分割。每个检测器中的主要 CNN 主干都被我们的 HIRI-ViT 系列(最初在 ImageNet-1 K 上进行预训练)替换。所有其他新添加的层都用 Xavier [74] 初始化。我们通过 AdamW 优化器 [69](批量大小:16)在 8 个 V100 GPU 上对检测器进行微调。对于 RetinaNet 和 Mask R-CNN,我们采用标准的 1× 训练计划(12 个时期)。每幅图像的短边调整为 1,600 像素,而长边不超过 2,666 像素。对于物体检测,我们对另外四种物体检测方法进行了实验:Cascade Mask R-CNN [75]、ATSS [76]、GFL [77] 和 Sparse RCNN [78]。按照 [14]、[40],采用多尺度策略的 3× 计划(36 个epoch)进行训练。输入图像的大小随机调整,将短边保持在 [960, 1,600] 的范围内,而长边则被强制小于 2,666 像素。我们报告了不同 IoU 阈值和三种不同物体大小(即小( A P S AP_S APS)、中( A P M AP_M APM)、大( A P L AP_L APL))的平均精度得分 ( A P AP AP)。例如,在分割任务中,报告了边界框和掩模 A P AP AP 分数( A P b AP^b APb、 A P m AP^m APm)。
性能比较:表 IV 总结了在 COCO 基准上具有不同主干的 RetinaNet 和 Mask R-CNN 的对象检测和实例分割性能。我们的 HIRI-ViT-S 和 HIRI-ViT-B 在所有指标上的表现始终优于每个组中具有可比计算成本的其他主干。具体而言,在 RetinaNet 和 Mask R-CNN 检测器的基础上,HIRI-ViT-S 分别比最佳竞争对手 ScalableViT-S 高出 1.4% ( A P AP AP) 和 Ortho-S 高出 0.5% ( A P m AP^m APm)。同时,HIRI-ViT-S 的输入分辨率是 ScalableViT-S 的两倍,而它们都需要相似的参数数量和 GFLOP。这清楚地验证了通过我们的设计扩展 Vision Transformer 在下游任务中的卓越通用性。表 V 进一步展示了 COCO 上不同主干下另外四个物体检测器的性能。类似地,HIRI-ViT-S 可以使每个物体检测器的性能与其他基线相比获得一致的提升。
C.ADE20K 上的语义分割
设置:接下来,我们在 ADE20 K 上评估 HIRI-ViT 的语义分割。该数据集涵盖 150 个语义类别,包含 20,000 张用于训练的图像、2,000 张用于验证的图像和 3,000 张用于测试的图像。在这里,我们遵循 [14] 并使用 UPerNet [79] 作为基础模型,其中 CNN 主干被我们的 HIRI-ViT 替换。整个网络通过 AdamW 优化器 [69] 在 8 个 V100 GPU 上进行了 160 K 次迭代训练。我们利用线性学习率衰减调度程序和 1,500 次迭代线性预热进行优化。批量大小、权重衰减和初始学习率设置为 16、0.01 和 0.00006。我们采用标准数据增强:随机水平翻转、随机光度失真和在 [0.5, 2.0] 比率范围内的随机重新缩放。为了公平比较,所有其他超参数和检测头的设置均与 Swin [14] 相同。
性能比较:表 VI 详细说明了不同主干在 ADE20 K 验证集上进行语义分割的性能。与对象检测和实例分割下游任务中的观察结果类似,我们的 HIRI-ViT-S 和 HIRI-ViT-B 在计算成本相当的情况下在每个组中获得了最佳 mIoU 分数。具体来说,在相同的 UPerNet 基础模型下,HIRI-ViT-S 将 HRViT-b2 的 IoU 分数提高了 1.2%,这再次证明了我们提案的有效性。
D. 消融研究
在本节中,我们首先说明如何构建一个具有四个阶段的强大的多阶段视觉变换器 (M-ViT),作为基础模型。然后,我们将基础模型扩展为五个阶段,并研究 HIRI-ViT 中的每个设计如何影响 ImageNet-1K 图像分类任务的整体性能。表 VII 通过将每个设计逐步纳入基础模型来总结不同消融运行的性能。
M-ViT:我们从基础模型开始,即具有四个阶段的多阶段 Vision Transformer (M-ViT)。M-ViT 的整体架构类似于 PVT [26],它由一个 ViT-stem 层和四个阶段组成。每个阶段都包含一个 Transformer 块堆栈。每个 Transformer 块由 MHA 和 FFN 组成。在前两个阶段,M-ViT 利用步幅卷积作为空间缩减来对键和值进行下采样,而后两个阶段不使用任何空间缩减操作。利用单个步幅卷积进行空间下采样,同时扩大每两个阶段之间的通道维度。采用 LN 进行特征归一化。如表 VII(第 1 行)所示,M-ViT 的 top-1 准确率达到 82.7%。
FFN → CFFN:第 2 行通过将每个 FFN 与额外的深度卷积(即 CFFN)集成来升级基础模型。通过这种方式,CFFN 能够利用归纳偏差,从而提高性能(82.9%)。
移除 MHA:接下来,我们借鉴 [53] 的灵感,在前两个阶段移除多头自注意力 (MHA)。CFFN 中的通道扩展维度也扩大了。如表 VII(第 3 行)所示,此消融运行保持了性能,同时 GFLOP 和参数数量均有所减少。
LN → BN:然后,我们在前两个阶段和每个 CFFN 块中用 BN 替换 LN(第 4 行)。top-1 准确率略有提高至 83.0%,表明 BN 比 LN 更适合具有卷积操作的块。
ViT-stem → Conv-stem:当我们用 Conv-stem 替换 ViT-stem(第 5 行)时,top-1 准确率进一步提高到 83.4%。这一观察证明了 Conv-stem 的优点,它在早期视觉处理中注入了少量的归纳偏差,从而稳定了优化并提高了峰值性能。
IRDS:之后,我们在每两个阶段之间应用反向残差下采样(第 6 行)。M-ViT 的最终版本具有四个阶段,成功实现了 83.6% 的 top-1 准确率,这具有竞争力,甚至超越了大多数现有的 ViT 主干。
分辨率 224→448:接下来,我们直接将这个四阶段结构的输入分辨率从 224×224 扩大到 448×448(第 7 行)。这种简化的运行可显著提高性能,同时将 GFLOP 从 4.7 大幅提高到 21.5。这一观察结果与现有的四阶段 ViT 架构(例如 PVT 和 SwinTransformer)一致,其中计算成本随着输入分辨率的扩大而呈二次方增长。
四阶段 → 五阶段:此外,我们通过利用附加阶段进一步对高分辨率输入进行下采样,扩展了四阶段消融运行(第 7 行)。这种五阶段结构(第 8 行)显著降低了高分辨率输入(448×448)的 GFLOP,Top-1 准确率下降到 84.0%,但仍优于 224×224 输入下的另一个四阶段消融运行(第 6 行)。结果基本证实了我们的五阶段结构的有效性,该结构寻求高分辨率输入的更好的性价比平衡。
Conv-stem → HR-stem:然后我们用我们的 HR-stem(在双分支设计中重新塑造)替换 Conv-stem。如表 VII(第 9 行)所示,GFLOP 明显下降(5.1),但 top-1 准确率仍然保持不变(84.0%)。
HR 块:之后,我们在前两个阶段(第 10 行)用两个堆叠的 HR 块替换每个 CFFN。请注意,对于小尺寸(GFLOP:∼5.0),ImageNet 上的 Top-1 准确率几乎饱和(∼84.0%),并且相对难以引入较大的改进幅度。然而,我们的 HR 块仍然设法在 Top-1 和 Top-5 准确率上实现了 0.1% 的性能提升,同时略微降低了计算成本,这再次验证了通过我们的高分辨率和低分辨率分支并行进行成本高效编码的优势。
EMA 提炼:最后,第 11 行是我们的 HIRI-ViT-S 的完整版本,通过额外的 EMA 提炼进行了优化,从而获得了最佳的 top-1 准确率 (84.3%)。
E. 每种区块的影响分为五个阶段
在这里,我们进一步进行了额外的消融研究,以全面检查每种块(HR/CFFN/Transformer 块)在五个阶段的影响。表 VIII 展示了通过构建具有不同块的五个阶段,HIRI-ViT 的不同消融变体的性能。具体而言,第 1-3 行的结果表明,在最后三个阶段用 Transformer 块替换更多的 HR 块通常可以提高性能,但计算成本会更高。例如,在第三阶段(第 3 行)使用 Transformer 块获得的性能增益很小(Top-1 准确率提高 0.1%),而 GFLOP 明显增加。结果基本验证了 Transformer 块在后期阶段对低分辨率输入的有效性,但 Transformer 块自然会导致巨大的计算开销,尤其是对于早期阶段的高分辨率输入。此外,我们在第三阶段(第 4 行)用 CFFN 块替换 Transformer 块来处理中分辨率输入。这样的设计很好地减少了 GFLOP,同时保持了 Top-1 和 Top-5 的准确率,体现了 HIRI-ViT 在计算成本和性能之间的最佳平衡。
F.计算成本与准确度
图 6 进一步说明了我们的 HIRI-ViT 和其他最先进的视觉主干的计算成本(即 (a) GFLOP、(b) 模型参数数量和 © 推理时间)的准确度曲线。如图所示,我们的 HIRI-ViT 主干的曲线始终高于其他视觉主干的曲线。也就是说,我们的 HIRI-ViT 主干比现有的视觉主干寻求更好的计算成本-准确度权衡。
G. 扩展到其他主干
在这里,我们利用我们的提案来扩展三个不同的视觉主干(HRNet、PVTv2、DaViT),使其具有更高分辨率的输入(224×224 到 448×448),从而报告性能/计算成本。如表 IX 所示,我们在 HIRI-ViT 中采用双分支设计的五阶段结构始终可以显着提高每个视觉主干的性能,同时保持有利的计算成本。结果进一步验证了我们的采用双分支设计的五阶段结构对于扩展具有高分辨率输入的视觉主干的通用性。
总结
在这项工作中,我们为高分辨率输入设计了一种新的五阶段 ViT 主干,即 HIRI-ViT,它新颖地将典型的 CNN 操作并行分解为高分辨率和低分辨率分支。凭借这种原则性的五阶段和双分支设计,我们的 HIRI-ViT 能够以经济高效的方式扩展具有高分辨率输入的 Vision Transformer 主干。在 ImageNet-1K(图像分类)、COCO(对象检测和实例分割)和 ADE20K 数据集(语义分割)上进行了广泛的实验,以验证我们的 HIRI-ViT 与竞争性 CNN 或 ViT 主干相比的有效性。
尽管有这些观察结果,但仍存在未解决的问题。虽然我们的双分支设计的五阶段结构明显提高了扩展 Vision Transformer 的效率,但我们观察到采用六阶段结构时性能/计算成本的改进较少。此外,如何利用高分辨率输入来扩展视频视觉变换器仍然是一个重大挑战。