多尺度高分辨率视觉Transformer用于语义分割的研究及应用

120940多尺度高分辨率视觉Transformer用于语义分割0Jiaqi Gu 1 * , Hyoukjun Kwon 2 , Dilin Wang 2 , Wei Ye 2 , Meng Li 2 , Yu-Hsin Chen 2 , Liangzhen Lai 2 , VikasChandra 2 , David Z. Pan 101 德克萨斯大学奥斯汀分校，2 Meta平台公司0jqgu@utexas.edu, dpan@ece.utexas.edu0{ hyoukjun, wdilin, weiye, meng.li, yhchen, liangzhen, vchandra } @fb.com0摘要0与基于卷积神经网络（CNN）的模型相比，视觉Transformer（ViTs）在计算机视觉任务上表现出更优越的性能。然而，ViTs主要设计用于图像分类，会生成单尺度低分辨率的表示，这使得密集预测任务，如语义分割，对ViTs来说具有挑战性。因此，我们提出了HRViT，通过将高分辨率多分支架构与ViTs集成，增强ViTs学习语义丰富和空间精确的多尺度表示。我们通过各种分支块协同优化技术平衡了HRViT的模型性能和效率。具体而言，我们探索了异构分支设计，减少了线性层中的冗余，并增强了注意力块的表达能力。这些方法使HRViT能够将语义分割的性能和效率的Pareto前沿推向一个新的水平，正如我们在ADE20K和Cityscapes上的评估结果所示。HRViT在ADE20K上实现了50.20%的mIoU，在Cityscapes上实现了83.16%的mIoU，超过了MiT和CSWin骨干网络的最新成绩，平均提高了1.78个mIoU，参数节省了28%，FLOPs减少了21%，展示了HRViT作为强大的语义分割视觉骨干的潜力。我们的代码公开可用。01. 引言0密集预测任务，如语义分割，在新兴的智能计算平台上是重要的计算机视觉工作负载，例如AR/VR设备。卷积神经网络（CNNs）在语义分割方面取得了显著的性能提升[1, 4, 19,21, 25,29]。除了传统的CNNs，视觉Transformer（ViTs）在计算机视觉任务中表现出有竞争力的性能。0* 本工作在Meta平台公司实习期间完成。1https://github.com/facebookresearch/HRViT018, 20, 28, 31, 32, 35, 36, 39, 43].凭借自注意力操作，ViTs具有与远距离信息交互和动态特征聚合的强大表达能力。然而，ViT[13]产生的是单尺度和低分辨率的表示，这对于需要高位置敏感性和细粒度图像细节的语义分割并不友好。为了应对这一挑战，提出了各种能产生多尺度表示的ViT骨干网络用于语义分割[6, 12, 20, 30, 31, 35,38]。然而，它们仍然遵循一种类似分类的网络拓扑结构，具有顺序或系列架构。基于复杂性考虑，它们逐渐对特征图进行下采样，提取更高级别的低分辨率（LR）表示，并直接将每个阶段的输出馈送到下游分割头。这样的顺序结构缺乏足够的跨尺度交互，因此无法产生高质量的高分辨率（HR）表示。HRNet[29]是在ViT背景之外提出的解决这个问题的方法，它通过保持网络中的所有分辨率来增强跨分辨率交互。HRNet并行提取多分辨率特征，并重复融合它们，以生成具有丰富语义信息的高质量HR表示。这种设计理念在各种密集预测任务中取得了巨大成功。然而，它的表达能力受到了小感受野和级联卷积操作的强归纳偏差的限制。为了应对这一挑战，提出了一些HRNet变体，如Lite-HRNet [37]和HR-NAS[11]。然而，这些改进的HRNet设计仍然主要基于卷积构建模块，并且它们在语义分割上的表现仍远远落后于ViT对应模型的最新成绩。因此，将HRNet与ViTs协同集成是一种可以进一步提高性能的方法。通过结合这两种方法，ViTs可以从HR架构中获得丰富的多尺度可表示性，而HRNet可以获得更大的感受野。2.2. Efficient HRViT component designA straightforward choice to fuse HRNet and ViTs is toreplace convolutions in HRNet with self-attentions. How-ever, given the high complexity of multi-branch HRNet andself-attentions, this brute-force combining can cause an ex-plosion in memory footprint, parameter size, and computa-tional cost. In this section, we will discuss how to designHRViT blocks with balanced efficiency and performance.Augmented cross-shaped local self-attention. To achievehigh performance with improved efficiency, a hardware-efficient self-attention operator is necessary. We adopt oneof the SoTA efficient attention designs, cross-shaped self-attention [12], as our baseline attention operator.Basedon that, we design our augmented cross-shaped local self-attention HRViTAttn illustrated in Figure 2, which pro-vides the following benefits: (1) Fine-grained attention:Compared to globally-downsampled attentions [30,35], thisone has fine-grained feature aggregation that preserves de-tailed information. (2) Approximate global view: By using120950然而，将HRNet的成功迁移到ViT骨干中并不容易。由于多分支HR架构和自注意操作的高复杂性，简单地用Transformer块替换HRNet中的所有卷积残差块将遇到严重的可扩展性问题。继承自多尺度的高表示能力可能会被硬件上的严重延迟和能量成本所压倒，而没有经过精心的架构块协同优化。0因此，我们提出了HRViT，一种专门针对语义分割进行优化的高效多尺度高分辨率视觉Transformer骨干。HRViT实现了ViT中的多尺度表示学习，并通过以下方法提高了效率：（1）HRViT的多分支HR架构以并行方式提取多尺度特征，并通过跨分辨率融合增强了ViT的多尺度可表示性；（2）HRViT的增强局部自注意去除了冗余的键和值，提高了模型的表达能力，并通过额外的并行卷积路径、附加的非线性单元和特征多样性增强的辅助快捷方式；（3）HRViT采用混合尺度卷积前馈网络来加强多尺度特征提取；（4）HRViT的HR卷积干部和高效的补丁嵌入层以较低的硬件成本保留更多的低级细粒度特征。此外，与HRNet系列不同，HRViT采用了独特的异构分支设计，以平衡效率和性能，它不仅仅是改进的HRNet或HRNet和自注意的直接集成，而是由自注意构成的纯ViT的新拓扑结构，通过精心的分支块协同优化。0基于HRViT中的方法，我们做出了以下贡献：0•我们深入研究了视觉Transformer中的多尺度表示学习，并提出了HRViT，将多分支高分辨率架构与视觉Transformer集成。0•为了增强HRViT的可扩展HR-ViT集成的效率，我们提出了一系列方法，如利用Transformer块中的冗余信息，开发性能-效率协同优化的构建块，并采用异构分支设计。0•我们在ADE20K和Cityscapes上评估了HRViT，并呈现了推动性能和效率Pareto前沿的结果，具体如下：HRViT在ADE20K的val上实现了50.20%的mIoU，在Cityscapes的val上实现了83.16%的mIoU，用于语义分割任务，相对于SoTA的MiT和CSWin骨干，mIoU提高了1.78倍，参数减少了28%，FLOPs降低了21%。02. 提出的HRViT架构0近年来，视觉Transformer骨干设计的最新进展主要集中在注意力操作符的创新上。一种新的拓扑设计可以为ViT提供更强的视觉表达能力。将ViT的顺序拓扑扩展为多分支结构，受到HRNet的启发，是一种提高性能的有希望的方法。一个尚未回答的重要问题是HRNet的成功是否可以有效地迁移到ViT骨干中，以巩固它们在语义分割等密集预测任务中的领先地位。在本节中，我们深入探讨了ViT中的多尺度表示学习，并介绍了HR架构和Transformer的高效集成。02.1. 架构概述0如图1所示，HRViT的第一部分由一个卷积干部组成，用于降低空间维度并提取低级特征。在卷积干部之后，HRViT部署了四个渐进的Transformer阶段，其中第n个阶段包含n个并行的多尺度Transformer分支。每个阶段可以有一个或多个模块。每个模块以一个轻量级的密集融合层开始，以实现跨分辨率交互，并使用高效的补丁嵌入块进行局部特征提取，然后是重复的增强局部自注意块（HRViTAttn）和混合尺度卷积前馈网络（MixCFN）。与逐步减小空间维度以生成金字塔特征的顺序ViT骨干不同，我们在整个网络中保持HR特征，通过跨分辨率融合来增强HR表示的质量。StemPatch EmbedUpsampleDownsampleTransformer BlocksStage 3Module 1Stage 1Stage 2Stage 4Skip Transformer BlocksLayerNormLayerNormFigure 1. The overall architecture of our proposed HRViT. It progressively expands to 4 branches. Each stage has multiple modules. Eachmodule contains multiple Transformer blocks.two parallel orthogonal local attentions, this attention cancollect global information. (3) Scalable complexity: onedimension of the window is fixed, which avoids quadraticcomplexity to image sizes.To balance the performance and efficiency, we introduceour augmented version, denoted as HRViTAttn, with sev-eral key optimizations. In Figure 2a, we follow the cross-shaped window partitioning approach in CSWin that sepa-rates the input x ∈ RH×W ×C into two parts {xH, xV ∈RH×W ×C/2}.xH is partitioned into disjoint horizontalwindows, and the other half xV is chunked into vertical win-dows. The window is set to s × W or H × s. Within eachwindow, the patch is chunked into K dk-dimensional heads,then a local self-attention is applied,HRViTAttn(x) = BN�σ(W O[y1, · · ·, yk, · · ·, yK])�yk = zk + DWConv�σ(W Vk x)�[z1k, · · · , zMk ] = zk =�H-Attnk(x),1 ≤ k < K/2V-Attnk(x),K/2 ≤ k ≤ Kzmk = MHSA(W Qk xm, W Kk xm, W Vk xm)[x1, · · ·, xm, · · ·, xM] = x,xm ∈ R(H/s)×W ×C,(1)where W Qk , W Kk , W Vk ∈ Rdk×C are projection matrices togenerate query Qk, key Kk, and value Vk tensors for thek-th head, W O ∈ RC×C is the output projection matrix,and σ is Hardswish activation. If the image sizes are nota multiple of window size, e.g., s⌈H/s⌉ > H, we applyzero-padding to inputs xH or xV to allow a complete K-thwindow, shown in Figure 2b. Then the padded attentionsare masked to 0 to avoid incoherent semantic correlation.The original QKV linear layers are quite costly in com-putation and parameters. We share the linear projections forkey and value tensors in HRViTAttn to save computationand parameters as follows,MHSA(W Qk xm, W Vk xm, W Vk xm)=softmax�Qmk (V mk )T√dk�V mk ,(2)QKAVMM+Vertical MHSA[ | ]Horizontal MHSAMM+LNHswish+DWConvHswish +Linear +BNWHCsC/2WsWdksoftmax(a)12K= H/s-infSoftmax0MasksWsW0(b)120960共享0多样性增强快捷方式（DES）0零填充0图2.（a）HRViTAttn：带有并行CONV路径和高效多样性增强快捷方式的增强型十字形自注意力。（b）使用注意力映射掩蔽的窗口零填充。0此外，我们引入了一个辅助路径，其中包含并行的深度卷积，以注入归纳偏差以便于训练。与CSWin中的局部位置编码不同，我们的并行路径是非线性的，并且应用于整个4-D特征图W Vx，而不是窗口分区。该路径可以被视为一个倒置残差模块，与自注意力中的线性投影层共享逐点卷积。这个共享路径可以有效地注入归纳偏差，并通过边际硬件开销增强局部特征聚合。作为对上述键值共享的性能补偿，我们引入了额外的Hardswish函数来提高非线性性能。我们还附加了一个BatchNorm（BN）HWCHWrC/2LinearLinear33DWConv55DWConvHWrC[ | ]HWCGELU +Linear+LNFigure 3. MixCFN with multiple depth-wise convolution paths toextract multi-scale local information.DES(x) = A · Hardswish(˜xBT ).(3)Mixed-scale convolutional feedforward network.In-spired by the MixFFN in MiT [35] and multi-branch in-verted residual blocks in HR-NAS [11], we design a mixed-scale convolutional FFN (MixCFN) by inserting two multi-scale depth-wise convolution paths between two linear lay-ers, shown in Figure 3.After LayerNorm, we expand thechannel by a ratio of r, then split it into two branches. The3×3 and 5×5 depth-wise convolutions (DWConvs) are usedto increase the multi-scale local information extraction ofHRViT. For efficiency consideration, we exploit the chan-nel redundancy by reducing the MixCFN expansion ratio rfrom 4 [20,35] to 2 or 3 with marginal performance loss onmedium to large models.Downsampling stem. In semantic segmentation tasks, im-ages are high resolution, e.g., 1024×1024. Self-attentionoperators are expensive as their complexity is quadratic toimage sizes.To address the scalability issue when pro-cessing large images, we down-sample the inputs by 4×before feeding into the main body of HRViT. We do notuse attention operations in the stem since early convolu-tions are more effective to extract low-level features thanself-attentions [15, 34].As early convolutions, we fol-low the design in HRNet and use two stride-2 CONV-BN-ReLU blocks as a stronger downsampling stem to extractC-channel features with more information maintained, un-like prior ViTs [6,20,35] that used a stride-4 convolution.Efficient patch embedding.Before Transformer blocks+++GELUEffPatchEmbed(x) = LN DWConv(PWConv(x)) .(4)120970层初始化为恒等投影，以稳定分布以获得更好的可训练性。受到ViTs中快捷方式的重要性的最新研究的启发[24]，我们添加了一个通道级投影器作为多样性增强的快捷方式（DES）。与增强型快捷方式[27]不同，我们的快捷方式具有更高的非线性性能，并且不依赖于硬件不友好的傅里叶变换。我们的DES中的投影矩阵P C × C通过Kronecker分解P =A√来近似。0C以减少参数成本。然后，我们将x折叠为˜x∈RHW×√C，并将（A�B）x转换为（A˜xB0DWConv 步长=2,4,8 BatchNorm 1 1 Conv2d 最近邻0上采样0图4. 具有通道匹配、上采样和下采样的跨分辨率融合层。0在每个模块中，我们在每个分支上添加一个补丁嵌入块（CONV-LayerNorm），用于匹配通道并提取增强的补丁间通信的补丁信息。然而，由于HR架构中每个阶段的每个模块都将有n个嵌入块，补丁嵌入层具有非常复杂的硬件成本。因此，我们将补丁嵌入简化为一个点卷积后跟一个深度可分离卷积[16]，0跨分辨率融合层。跨分辨率融合层对于HRViT学习高质量的HR表示非常关键，如图4所示。为了增强跨分辨率交互，我们在每个模块的开头插入了重复的跨分辨率融合层，这是根据HRNet [29,37]的方法进行的。为了帮助LR特征保持更多的图像细节和精确的位置信息，我们将它们与下采样的HR特征合并。我们不使用渐进的基于卷积的下采样路径来匹配张量形状[29,37]，而是采用直接的下采样路径来减少硬件开销。在第i个输入和第j个输出（j >i）之间的下采样路径中，我们使用步长为2j -i的深度可分离卷积来缩小空间维度并匹配输出通道。DWConv中使用的内核大小为（2j - i +1），以创建补丁重叠。这些HR路径将更多的图像信息注入到LR路径中，以减轻信息损失并加强反向传播中的梯度流，从而促进深度LR分支的训练。另一方面，由于我们在HR路径上最小化了窗口大小和分支深度，HR块中的感受野通常是有限的。因此，我们将LR表示合并到HR路径中，以帮助它们获得具有更大感受野的更高级特征。具体而言，在上采样路径（j 

下载后可阅读完整内容，剩余1页未读，立即下载