深度学习的进步推动了遥感语义分割领域的显著进展。虽然注意力机制能够进行全局建模并利用上下文信息，但它们面临着计算成本高的问题，并且需要基于窗口的操作，这削弱了捕捉长距离依赖的能力，限制了它们在遥感图像处理中的有效性。
在本文中，作者提出了AMMUNet，一个基于UNet的框架，该框架采用多尺度注意力图融合，包括两个关键创新：粒度多头自注意力（GMSA）模块和注意力图融合机制（AMMM）。
GMSA在获取全局信息的同时，显著降低了与全局多头自注意力机制相比的计算成本。这是通过策略性地利用维度对应来对齐粒度，并减少相对位置偏置参数来实现的，从而优化了计算效率。
所提出的AMMM有效地使用固定 Mask 模板将多尺度注意力图融合为统一表示，实现了全局注意力机制的建模。
实验评估突显了作者方法的高性能，在具有挑战性的 Vailingen 数据集上获得了75.48%的卓越平均交并比（mIoU）得分，在 Potsdam 数据集上更是达到了77.90%，证明了作者方法在精确遥感语义分割方面的优越性。
关注公众号，私信 『代码』获取。
I Introduction
远程感知图像的语义分割至关重要，因为它为城市规划、环境监测和资源管理等一系列应用奠定了基础。深度学习的出现，尤其是卷积神经网络（CNNs），带来了显著的范式转变，革新了计算机视觉领域，进而也影响了远程感知图像分析。
Long等人[1]开创了全卷积网络，成为首个端到端的可训练语义分割架构。此后，ResNet[2]通过引入残差连接，改变了深度神经网络，有效训练了非常深的架构，同时减轻了退化问题，保持了丰富的多级表示。在这些进展的基础上，DeepLabv3[3]及其增强版本DeepLabv3+[4]采用了扩张/孔洞卷积来扩大感受野并捕捉多尺度信息。然而，尽管这些增强，基于传统CNN的方法仍然由于其固有的局部感受野而难以建模全局上下文信息。
为了增强全局信息的利用，ViT[5]通过窗口分割在图像上应用了多头自注意力（MSA）[6]，Swin Transformer[7]则通过移动窗口交换窗口信息。对于大规模远程感知图像的密集语义分割任务[8, 9]，Transformer和CNN已经在U-Net[10]架构中协同使用，旨在有效地捕捉细粒度细节和全局上下文信息[11, 12]。这些努力包括通过线性注意力提高效率的MAResUNet[13]，以及通过融合卷积和注意力模块的高效框架组合全局和局部特征信息的UNeFormer[14]。
然而，尽管取得了成功，基于Transformer的方法有一个重大限制：为了减轻全局自注意力机制的计算需求，需要基于窗口的操作[15, 16]。首先，它削弱了 Patch 之间的上下文关系，可能阻碍了对长距离依赖的准确建模[17]。其次，窗口之间的相对位置关系变得隐含，进一步复杂化建模过程。即使是Swin Transformer采用的移动窗口方法也只能部分缓解这个问题，通过建模相邻窗口之间的关系。
为了充分发挥基于Transformer的架构在远程感知图像分割中的潜力，作者引入了AMMUNet，这是一个基于U-Net的框架，它使用ResNet编码器提取多尺度特征，并在解码器中采用创新性注意力机制。这个机制结合了CNN和Transformer的优势，复用高级注意力图以保持长距离建模能力，同时利用作者的粒状多头自注意力（GMSA）模块计算细粒度的局部注意力关系。此外，作者提出了一种注意力图合并机制（AMMM），它使用固定的 Mask 模板来合并在不同尺度生成的注意力图，使得能够在最小化计算需求的同时建模全局注意力机制，确保最佳准确度。作者工作的关键创新可以总结如下：
- 所提出的AMMUNet是一个利用基于CNN的ResNet编码器和解码器中新颖全局注意力机制的语义分割网络，能够合并多尺度注意力图。
- 作者设计了GMSA模块以在当前尺度提取注意力图，并使用AMMM合并多尺度注意力图。作者的方法包括粒度对齐、相对位置偏置减少以及用于注意力图合并的固定 Mask 模板，与全局多头自注意力相比，有效减少计算需求，同时实现实时性能。
- 对Vaihingen和Potsdam数据集的实验评估表明，在准确度和效率方面，作者的方法具有优越性。
II Methodology
Overview of AMMUNet
提出的语义分割网络如图1(a)所示。该网络采用UNet作为框架，这是一个整体的编码器-解码器结构。编码器使用ResNet50作为 Backbone 网络来提取多尺度特征。ResNet50由四个阶段组成，每个阶段执行空间下采样。这四个输出对应于不同大小的特征图，从浅层到深层分别为：、、和。
解码器将ResNet中最深层级的特征图作为其初始输入。这个输入首先通过MSA模块，生成相应的注意力图和特征图。注意力图作为作者提出的AMMM模块的输入，而特征图则通过上采样并与ResNet中的中间特征图通过跳跃连接进行拼接，共同输入到GMSA模块中。在解码器的中间两层，作者采用GMSA模块来计算和转换输入特征图，生成的注意力图传递给AMMM模块。AMMM模块合并不同尺度的注意力图以形成一个统一的全局注意力图，然后将其反馈到GMSA中以指导后续的特征图计算。最终，解码器输出的高分辨率特征图通过多层感知机（MLP）进行像素级语义预测。作者根据预测和真实值之间的差异计算损失函数，并使用反向传播以端到端的方式训练整个网络。
Global Multi-head Self-Attention
传统的计算机视觉中的Transformer方法通常采用块嵌入（patch embeddings）将输入图像划分为较小的窗口，然后在这些窗口上计算自注意力，这种策略旨在降低计算需求。然而，这样做以牺牲块之间的稀疏关联为代价。相比之下，作者的方法在解码器中最深层特征图上采用全局方法来计算自注意力，如图1(d)所示。这个决定基于两个关键考虑：首先，在解码器的最深层，窗口大小通常较小，限制了基于窗口操作的计算节省；其次，全局方法能够在单个像素区域之间进行细粒度的相关性计算。
注意力图（AM）和注意力输出（Attn）计算如下：
其中，Q、K和V分别表示 Query 、键和值，d表示输入维度，RPB表示相对位置偏置。值得注意的是，注意力图不仅有助于注意力输出计算，还作为AMMM模块的输入，促进注意力图的融合。这种方法能够建模浅层注意力图中远距离像素之间的长距离依赖。
Granular Multi-head Self-Attention
如图1(b)所示，与标准的多头自注意力模块相比，GMSA模块将整体特征图划分为2×2像素的粒度窗口。对于每个子区域，它计算 Query （Q）、键（K）和值（V）以获得相应的注意力图，记为。假设输入特征图的尺寸为，则有个子区域，产生个单独的子区域注意力图。
相对位置偏置扩展了自注意力的概念，以编码任意两个标记之间的距离。由于每个子区域内的相对位置是固定的，为了计算效率，可以在不同的子区域中应用相同的相对位置偏置。
然后，这些单独的子区域注意力图被插入到整体注意力图的对角区域中：
GMSA计算当前层每个子区域内的注意力，既有助于当前层的注意力计算，也代表了在更浅层对应的粒度上的相关性。全局MSA模块、窗口MSA [7]和GMSA的计算成本为：
其中hw表示输入特征图的空间尺寸，为窗口大小，表示解码器中最深层特征图的尺寸。与MSA模块和W-MSA相比，所提出的GMSA在输入规模增加时，计算复杂度的增长速率较慢，使其特别适合于大型遥感图像特征提取。
Attention Map Merging Mechanism
AMMM旨在将当前尺度下由GMSA计算的注意力图与从更深层次传播的注意力图合并，这些更深层次的注意力图可能来自基于 Patch 的MSA或之前的GMSA模块，如图1(c)所示。更深层次的注意力图首先乘以一个固定的 Mask 模板，为当前尺度的注意力图保留空间，然后进行上采样，如下方程所示：
其中表示作为单位矩阵实现的固定 Mask 模板，表示哈达玛积运算符，表示由GMSA计算的当前尺度的注意力图，表示来自前一个模块的注意力图，代表合并的更深层次的注意力图。
Dimension Correspondence Module
维度对应模块（DCM）有助于将内存中基于窗口格式的注意力映射转换成基于行格式。对于浅层特征，注意力映射以窗口方式存储。然而，由于AMMM在多个尺度上合并注意力映射，通过方程2计算的注意力输出以行格式存储，因此需要进行尺度对齐。这是通过以下维度转换实现的：
其中 表示Attn'的维度，子区域大小为22，对应的维度大小为 。可能需要多次DCM操作，这与架构中上采样层的数量一致。
III Experiment
以下是III实验部分的开始。
Dataset
分割模型使用两个数据集进行了评估：Vaihingen数据集和Potsdam数据集。Vaihingen数据集包含33张真实正射影像，其中17张用于训练，其余用于测试。图像尺寸在800到3000像素之间变化。Potsdam数据集由38张图像组成，每张图像尺寸为6000x6000像素。作者使用24张图像进行训练，14张进行测试。这两个数据集的地面采样距离为9厘米。标注包括六个类别的像素级标签：不透水表面、建筑物、低矮植被、树木、汽车和杂物。为了便于GPU训练，作者将图像和标签裁剪成了不重叠的512×512像素块。
Evaluation Metrics
为了评估作者提出的AMMUNet框架的性能，作者采用了两种广泛采用的指标：平均交并比（mIoU）和平均像素准确率（mAcc）18。这些指标的范围是从0到1，值越高表示模型性能越好。值得注意的是，对于不平衡数据集，mIoU是一个更严格的指标，在语义分割任务中更常使用。mIoU的计算方式如下：
mAcc定义如下：
其中，TP、FP、FN和TN分别表示真正例、假正例、假反例和真反例的像素计数，C表示数据集中的总类别数。
Implementation Details
所有实验都是在单个NVIDIA GTX 3090 GPU上使用PyTorch框架进行的。在训练阶段，作者使用了带有0.01权重衰减参数的AdamW优化器。最大训练周期数设置为200，训练批处理大小为8。初始学习率为1e-4，所使用的损失函数为交叉熵。作者采用“poly”学习率策略来更新学习率，其指数为0.9。
Comparison with State-of-the-Art Methods
为了验证作者提出的模型的优势，作者将其与五种最先进的方法进行了基准测试，包括DeeplabV3+[4]、DANet[19]、OCRNet[20]、SegFormer[21]和UNetFormer[14]。DeeplabV3+，一个纯粹的基于CNN的架构，被广泛认为是语义分割的 Baseline ，并通过使用扩张卷积增强模型的感受野。其余方法结合了CNN和Transformer架构，利用自注意力机制捕捉全局语义上下文。在遥感分割任务上的优势基础上，作者的方法以UNetFormer作为 Baseline ，并对其进行实验以评估作者提出的方法相对于这一 Baseline 网络的性能。为了确保公平评估，所有网络都使用了ResNet50作为主干网络，OCRNet和SegFormer分别使用了HRNet18s和MixTransformerB1，因为它们的模型大小与ResNet50相似。作者使用mIoU和mAcc作为评估指标，在Vaihingen和Potsdam数据集上评估这些模型。详细的包括类别特定IoU指标的结果展示在表1和表2中。
具体来说，对于Vaihingen数据集，作者的方法在mIoU上达到75.48%，在mAcc上达到81.23%；对于Potsdam数据集，分别达到77.90%和85.92%，优于像DeeplabV3+这样的高级CNN-based方法。得益于类别基础的语义信息以及空间细节上下文，作者方法的性能也超过了最初为遥感图像设计的最近的Transformer-based方法，如UNetFormer。作者的AMMUNet在两个数据集的三个子类别上均取得了最佳性能，并获得了最高的总体得分。视觉示例可以在图2和图3中找到。
在U-Net[10]和UNetFormer架构上，对不带 Patch 的MSA和GMSA+AMMM进行了比较。如表3所示的结果表明，作者提出的方法优于其他方法，在U-Net和UNetFormer框架上分别实现了0.39%和0.64%的mIoU提升。
此外，作者还评估了作者的方法的计算效率和推理速度，与其它方法在模型大小、浮点运算和推理时间方面进行了比较。作者的方法保持了与UNetFormer相当的推理速度和内存占用，同时将计算需求和推理时间减少了一半左右，与不带 Patch 的MSA方法相比，保持了相似的参数数量，并满足了实时推理的要求。作者提出的方法不仅在此方法上保持了准确性的优势，还提供了更小、更易于部署的参数占用。
IV Conclusion
在这封信中，作者提出了AMMUNet，一种新颖的基于U-Net的架构，该架构采用ResNet作为编码器，并在解码器中集成了粒度多 Head 自注意力模块与注意力图合并机制。通过避免窗口堆叠，作者的模型有效地平衡了全局和局部信息的提取与利用。在ISPRS Vaihingen和Potsdam数据集上的大量实验，与现有方法进行比较，证明了作者类别引导机制的优势以及所提出方法在分割精度方面的优越性。
参考
[1].AMMUNet: Multi-Scale Attention Map Merging for Remote Sensing Image Segmentation.