【TIP2024】HAFormer: Unleashing the Power of Hierarchy-Aware Features for Lightweight Semantic Segment

【TIP2024】HAFormer: Unleashing the Power of Hierarchy-Aware Features for Lightweight Semantic Segmentation

论文链接：http://arxiv.org/abs/2407.07441
代码链接：https://github.com/XU-GITHUB-curry/HAFormer

一、摘要

 文中引入了HAFormer模型，该模型结合了CNN层次化特征提取能力与Transformer全局依赖建模能力，以应对轻量级语义分割的挑战。具体来说，设计了一个层次感知像素激励（HAPE）模块，用于自适应的多尺度局部特征提取。在全局感知建模过程中，设计了一个高效的Transformer （ET）模块，简化了传统Transformer中的二次计算。此外，还设计了一个相关加权融合（cwF）模块，该模块能够选择性地融合不同的特征表征，显著提高了预测精度。HAFormer在保持高性能的同时，具有最小的计算开销和紧凑的模型尺寸，在Cityscapes测试数据集上达到了74.2%的mIoU，在CamVid测试数据集上达到了71.1%的mIoU，在单个2080Ti GPU上的帧率分别为105FPS和118FPS。

二、创新点

• 提出了一种新颖的层次感知像素激励（HAPE）模块，该模块利用层次结构和内容感知注意力机制来减少计算负载，同时能够在各种感受野下提取更深层次的像素语义信息。

• 开发了一种有效的特征融合机制，称为相关加权融合（cwF），用于协同集成CNN和Transformer学习到的局部和全局上下文特征，有效地提高了准确性。

• 提出了一种有效的Transformer，用于分解Q、K和V 矩阵，有效地解决了传统Transformer模型中存在的二次计算复杂性挑战。

三、原理

A. Overall Architecture

 提出的HAFormer的整体架构如图2所示，它包含三个组成部分：一个带有层次感知像素激励的CNN编码器、一个高效的Transformer编码器和一个轻量级解码器。

 对于给定的输入图像$I\in R^{3\times H\times W}$其尺寸为$H\times \times W$，模型首先通过CNN编码器生成特征$F\in R^{C_f\times H_f\times W_f}$ （其中$H_f=\frac{H}{8}$，$W_f=\frac{W}{8}$）。同时，输入$I$在Transformer Stem块之后的Transformer编码器中进行处理，生成特征嵌入$T\in R^{N\times D}$，其中$N=\frac{H_t}{P}\times \frac{W_t}{P}$（$H_t=\frac{H}{16}$，$W_t=\frac{W}{16}$）表示token数量，$D=C_t\times P^2$表示每个token的维度，$P$表示token大小。随后，新设计的关联加权融合（cwF）模块有效地将两种不同类型的上下文特征$F$和$T$结合起来。这种融合相关联的CNN和Transformer特征，通过轻量级解码器分割头增强了边界信息和恢复。

 具体来说，为了优化CNN编码器，在CNNStem块中使用三个$3\times 3$卷积层。在这种配置中，最后一层的步长为$2$，产生特征图的大小为$C_1\times \frac{H}{2}\times \frac{W}{2}$，其中$C_1$表示输出通道数量。

 相反，Transformer stem在Transformer Encoder中降低了分辨率，同时提取特征表示，通过最小化计算负载，因为更高的分辨率意味着更多的计算，从而为模型设计提供了轻量级。因此，在Transformer Stem块中，使用四个$3\times 3$卷积层，步长为$2$，产生输出特征的大小为$C_t\times \frac{H}{16}\times \frac{W}{16}$。

B. Hierarchy-Aware Pixel-Excitation (HAPE) Module

 采用同一层中不同大小的卷积核，结合像素激励，有助于从不同大小的对象中提取特征。在此基础上，受到如 [30], [44]等工作的启发，采用了多尺度策略来捕捉不同感受野层次上的独特像素特征。与ESPNet[12]中的层间合并和Inception [45]中的串联不同，所提模块避免了冗余计算，使得网络更加精简，同时保持了特征的有效性。此外，为了进一步提高不同尺度上的像素表征，在本研究中引入了创新的层次感知像素激励（HAPE）模块。该模块增强了模型识别图像中各种大小对象的能力，最终降低了像素误分类率。

 如图3所示，给定特征输入$X_{in}\in R^{N_c\times H_c\times W_c}$，首先将其输入到一个$1\times 1$卷积层，以减少其通道维度至$\frac{N_c}{4}$，即输出特征图$\tilde{X}$被表示为
$$\tilde{X}=f_{1\times 1}(X_{in}),\tilde{X}\in R^{\frac{N_c}{4}\times H_c\times W_c}.\text{\hspace{1em}(1)}$$
其中$f_{1\times 1}$ 表示一个核大小为 $1\times 1$ 的卷积操作。这种尺寸的减少有利于后续层次卷积层中的通道操作。

 随后，执行四个并行的卷积操作，包括因子化卷积和深度可分离卷积，核大小分别为3、3、5和7。此外，最后三个卷积层使用膨胀卷积来增强感受野，如图4所示。这一策略使得模型能够捕捉不同尺度下的图像特征，确保了全面和详细的特征提取。上述过程可以表示为
$$l_1=f_{1\times 3}(f_{3\times 1}(\tilde{X})),\text{\hspace{1em}(2)}$$
$$l_i=f_{1\times k_i}^{d_c}(f_{k_i\times 1}^{d_c}(\tilde{X})),k_i=3,5,7;i=2,3,4,\text{\hspace{1em}(3)}$$

其中$l_i$表示中间特征，$f_{1\times k_i}$ 是一个核大小为$k_i$的1-D卷积操作，而$d_c$表示膨胀率。为了简化，等式中省略了一些激活和批量归一化操作。

在像素激发模块（PEM）中，一个关键要素在于通过一种内容感知的空间注意力机制来增强特征的可表示性。如图3所示，这一过程首先将输入$x\in R^{c\times h\times w}$送入全局平均池化（GAP）层，生成$x_1\in R^{1\times h\times w}$。随后，$x_1$经过reshape和flatten操作，然后被输入到Softmax函数中，以计算权重矩阵$A\in R^{1\times h\times w}$。这个权重矩阵随后与输入特征相乘，产生一个内容感知注意力增强的输出$x^{\prime }$。

 这一过程可以表示为
$$x_1=Reshape(GAP(x)),\text{\hspace{1em}(4)}$$
$$A=Reshap{e}^{-1}(Softmax(x_1)),\text{\hspace{1em}(5)}$$
$$x^{\prime }=\delta (x\odot A+x).\text{\hspace{1em}(6)}$$
这里，Reshape 和 $Reshap{e}^{-1}$ 表示reshape操作及其逆操作，$\delta$是一个激活函数，而$\odot$表示逐元素相乘。

 最后，使用残差结构来保留原始特征，得到最终输出$Y\in R^{N_c\times H_c\times W_c}$。这四个卷积层被联合添加到一个$1\times 1$ 卷积中进行特征融合和通道恢复。模块内部保持了残差连接，通道混洗Channel Shuffle操作有效地促进了通道之间的信息交互，表达式为
$$Y=Shuffle(f_{1\times 1}(\delta (\sum _{i=1}^{4}PEM(l_i))+X_{in})),\text{\hspace{1em}(7)}$$
其中，Shuffle 表示通道混洗操作，而$\delta$是一个激活函数。

C. Efficient Transformer

 所提方法专注于通过减少特征维度而不显著丢失图像细节来降低计算成本。为了实现这一点，引入了一种空间降维线性投影方法，该方法首先将特征映射到一个低维的潜在嵌入空间，然后将其用于多头自注意力计算。这种方法，称为高效的多头自注意力（eMHSA），包括学习的投影和分割操作，如图5所示。

 设输入特征为$X_t\in R^{C_t\times H_t\times W_t}$，其中$C_t$，$H_t$和$W——t$分别表示通道数、高度和宽度的特征图。在Reshape操作之后，得到一系列平铺的非重叠patch的序列，结果为$X_t\in R^{N\times (C_t\cdot P^2)}$，其中$N=\frac{H_t{W}_t}{P^2}$表示patch的数量（即输入序列长度），每个patch的大小为$P\times P$。随后，patch通过一个可学习的线性投影层$E\in R^{(P^2\cdot C_t)\times D}$映射到一个潜在的$D$维嵌入空间，表示为$Z\in R^{N\times D}$。这个过程可以表示为
$$Z=[x_p^{1}E;x_p^{2}E;...;x_p^{N}E],\text{\hspace{1em}(8)}$$

其中，$x_p^i$表示第$i$个补丁。请注意，省略位置嵌入是有意为之，以便为不同大小的输入提供更大的适应性。

 随后，Transformer中的三个矩阵，即查询$Q$、键$K$和值$V$ ，是通过它们的线性投影$W^Q$、$W^K$ 和 $W^V\in R^{D\times D_h}$ 得出的。这可以表示为
$$Q,K,V=Z{W}^Q,Z{W}^K,Z{W}^V\in R^{N\times D_h}\text{\hspace{1em}(9)}$$

 此外，多头自注意力中的头数 $h$ 也是一个定义的参数，确保每个头的维度等于 $d=\frac{D_h}{h}$ 。因此，第$i$ 个头中 $q$、$k$ 和 $v$ 的维度是 $N\times d$。在第 $i$ 个头中，$k$和 $v$通过一个缩放因子 $r$（缩放比例设置为$2$）进行空间缩减。然后，特征分割操作产生的子token与仅代表原始感知 $\frac{1}{s}$ 的场进行矩阵乘法运算，其中 $s$ 表示特征分割的数量，设置为$4$。这个过程可以描述为
$$(q_1,...,q_s),(k_1,...,k_s),(v_1,...,v_s)=Feature\_Split(q,k,v).\text{\hspace{1em}(10)}$$

 因此，空间分布变为$q_i\in R^{N\times \frac{d}{s}}$，$k_i\in R^{\frac{N}{r}\times \frac{d}{s}}$，和$v_i\in R^{\frac{N}{r}\times \frac{d}{s}}$。这个想法与group卷积的概念有相似之处，并且能够有效地减少内存消耗。因此，第$n$个头的自注意力计算如下
$$o_i(q_i,k_i,v_i)=Softmax(\frac{q_i(k_i{)}^T}{\sqrt{d}})v_i,i\in [1,s],\text{\hspace{1em}(11)}$$
$$hea{d}^n=Concat[o_1,o_2,...,o_s],n\in [1,h]\text{\hspace{1em}(12)}$$

其中，$Concat[\cdot ,\cdot ]$ 表示拼接操作。

 因此，eMHSA的最终输出表示为:

$$eMHSA=Concat[hea{d}^1,hea{d}^2,...,hea{d}^h]W^O,\text{\hspace{1em}(13)}$$

其中，$h$表示 eMHSA 中的注意力头的数量，而$W^O\in R^{D_h\times D}$ 作为一个线性投影，用于恢复维度。因此，通过上述结构设计，将复杂度从 $O(N^2)$ 降低到了 $O(\frac{N^2}{sr})$。

 Transformer 系列[16], [17], [23] 也使用了一种自注意力机制，包括多头注意力。然而，所提方法在计算上对于捕捉特征之间的详细关系是密集型的，这与我们的目标相悖。

 至于 MLP 层，作者遵循[31], [48] 中描述的方法，用零填充位置编码替换固定大小的位置编码。此外，输入张量中引入了一个深度卷积，填充大小为 $1$，以捕捉全连接（FC）层和 MLP 中的GELU 之间的局部连续性。通过消除固定大小的位置嵌入，模型在处理不同分辨率的输入时变得通用。因此，高效 MLP 层的输出，记为“eMLP”，可以写为：
$$eMLP=\rho ({\xi }_{GELU}(f_{DWConv}(\rho (x_e)))),\text{\hspace{1em}(14)}$$
其中$\rho$表示FC层的操作，${\xi }_{GELU}$表示GELU激活函数，$f_{DWConv}$表示深度卷积，而$x_e$是eMLP的输入。

D. Correlation-weighted Fusion

 考虑到Transformer和CNN的特性以及计算机制的差异，传统的逐元素相加或拼接操作可能无法达到最优效果。因此，设计一种能够充分利用两者互补优势的方案对于最大化提取特征的表示能力以及在解码过程中恢复信息至关重要。

 文中提出了一种有效的策略来弥合这一差距。所提方法通过相关性加权整合，无缝地将Transformer和CNN提取的不同类型的特征结合起来。通过融合与CNN特征高度相关的Transformer特征，开发了一种新的相关性加权融合（cwF）模块。

 如图6所示，$T$和$F$分别表示Transformer和CNN的中间特征。首先，将Transformer特征$T$reshape以匹配CNN特征$F$的形状，然后对两组特征进行后拼接操作。为了减少计算成本，使用深度可分离卷积进行通道维度的减少。经过全局平均池化（GAP）和Sigmoid操作后，计算出一个相关系数矩阵，记为$M$。然后将这个矩阵与原始特征相乘，得到 ̃$\tilde{F}$和$\tilde{T}$，将它们相加得到最终输出$Z$。

 这个过程可以表示为：
$$G=Concat[Reshape(T|F),F],\text{\hspace{1em}(15)}$$

其中$G\in R^{C_g\times H_g\times W_g}$ （$C_g=C_f+C_t$），$Concat$表示拼接操作，而$a|b$表示大小为$a$的特征图恢复到大小为$b$。然后，相关系数矩阵$M$可以计算为
$$M=\delta (f_{1\times 1}(GAP(f_{1\times 1}(f_{3\times 3}(G))))),\text{\hspace{1em}(16)}$$
其中，$M\in R^{C_m\times 1\times 1}$，$\delta$是 Sigmoid 函数，GAP表示全局平均池化操作，而 $f_{k_1\times k_2}$ 表示核大小为 $k_1\times k_2$ 的卷积操作。

 因此，得到的$cwF$特征，用 $Z_{cwF}$ 表示，可以表示为
$$Z_{cwf}=\varphi (\tilde{T}+\tilde{F}),\tilde{T}=T\odot M,\tilde{F}=F\odot M,\text{\hspace{1em}(17)}$$
其中$Z_{cwF}\in R^{C_z\times H_z\times W_z}$，$\varphi$是ReLU激活函数，并且$\odot$表示逐元素乘法。

四、实验

 为了展示HAFormer及其各个模块的质量和数量上的有效性，在基准数据集上进行了比较实验，并与最先进(SOTA)的方法进行了比较。在本节中，首先概述了在实验中使用的数据集、损失函数、硬件平台配置和参数设置。然后，展示了为了验证各个模块的有效性而进行的一系列消融实验。最后，进行了比较实验，以展示所提方法相对于SOTA方法的优越性。

A. Datasets

 HAFormer模型旨在应对街道场景中尺度
变化和上下文信息的挑战。Cityscapes [49]和CamVid[50]这两个数据集是街道场景分割研究中广泛使用的两个重要基准。因此，为了展示所提模型的有效性，在这两个数据集上进行了一系列全面的经验评估。

 Cityscapes。该数据集包含5,000张高质量图像，这些图像在50个城市的各种城市环境中拍摄，分辨率为$2,048\times 1,024$，主要描绘驾驶场景。数据集分为三个子集：2,975张图像用于训练，500张用于验证，1,525张用于测试。虽然数据集包括34个类别标签，但研究主要关注19个关键的语义类别。使用Cityscapes的内置工具来调整标签以满足研究需求。

 CamVid。这是由剑桥大学发布的一个公共城市道路场景数据集。这些分辨率为$960\times 720$的图像是从驾驶视角拍摄的，增加了观察目标的多样性。该数据集包含超过700张标注图像，适合于监督学习。CamVid数据集通常使用11个常见类别来评估分割精度。这些类别全面代表了城市道路场景中的对象，使其成为研究的宝贵资源。

B. Implementation Details

 本研究中的HAFormer模型在配备12GB内存的单个RTX 2080 Ti GPU卡上运行，使用CUDA10.1和PyTorch 1.8.1。该架构是从头开始训练的，没有使用任何预训练模型。实验采用随机梯度下降（SGD）优化算法，动量为0.9，权重衰减为$1e-5$，并采用
Poly学习率策略。

 对于Cityscapes数据集，初始学习率为$4.5e-2$，批量大小设置为5以最大化GPU内存使用。对于CamVid数据集，初始学习率为$1e-3$，批量大小设置为8。遵循现有做法，应用数据增强技术，包括水平翻转、随机缩放和随机裁剪，以增加训练数据的多样性，对于Cityscapes数据集，随机缩放范围从0.25到2.0，裁剪尺寸为$512\times 1024$，总共训练1,000轮次。为了公平比较，没有进行后处理。

 最后，遵循现有做法，性能通过所有类别平均的交并比（mIoU）进行定量评估，以及参数数量、FLOPs、GPU使用情况和处理速度。

C. Ablation Studies

  本部分进行了一系列的消融实验来验证所提方法中每个模块的有效性。

Ablation Study of the HAPE Module.
 HAPE模块提出了四种并行卷积操作，以全面捕捉图像特征的各个层次。然后是PEM，设计用来通过一个内容感知的空间注意力机制增强特征的可表示性。本节将展示HAPE模块中层次方法（记为“HM”）和PEM方法的有效性。

 用于比较的基线模型的结构是一种单行类型（如图7所示），它包含了标准的残差模块（RMs）。为了展示HM和PEM带来的性能提升，首先用HM模块替换基线模型中的RM，省略PEM部分，然后包含HM和PEM模块来测试整个HAPE模块的有效性。

 表I展示了HM的优越性能，它在mIoU上比RM高出1.47%和1.53%。HM在提取鲁棒特征方面表现出色，有效地促进了深层语义信息的提取。此外，多尺度结构显著增强了模型在特征提取和小对象识别方面的性能。引入PEM 进一步提高了两个数据集上的分割精度，提高了2.13%和2.74%。

 在整个实验过程中，在HM和HAPE中设置的膨胀卷积率都设置为1，以确保公平比较。图8也验证了当HAPE模块与Transformer模块集成时的有效性。

Ablation Study of the Dilation Rates.

 本节将探讨所选膨胀率对分割性能的影响。在模块数量保持一致的情况下，更大的膨胀率扩展感受野，允许模型感知更广泛的范围。

 表II显示，将膨胀卷积率从所有1s转换为所有2s（前两行），mIoU提高了约0.5%。进一步，通过逐步增加阶段2和3中的膨胀卷积率，在两个数据集上观察到性能提升1.21%和1.73%。因此，为了保留空间细节，在所提方法在阶段1和4分配了三个模块，而在阶段2和3使用了六个模块来捕捉网络深度内的精细语义信息。这种策略优化了Transformer的编码器计算，改进了长距离依赖建模

Ablation Study of the Efficient Transformer.
 在 HAFormer 中的另一个关键贡献是高效 Transformer （ET）模块，该模块在计算自注意力之前通过将其投影到最优的潜在嵌入空间中来减少特征的维度。表III展示了ET模块在Cityscapes和CamVid数据集上的分割准确性和计算复杂性方面相对于传统Transformer（标记为“TT”）的性能提升。

 ET设计在效率和准确性之间展现了出色的平衡。与传统Transformer“TT”相比，ET实
现了参数数量减少18%和计算负载减少17%，仅轻微的mIoU损失0.4%。这导致了一个效率更高的模型，对性能的影响极小，并且它甚至提供了更快的推理速度。此外，表 IV中的结果也揭示了通过Transformer学习到的特征的显著增强，mIoU有了显著的2.16%提升。这强调了Transformer捕捉长距离依赖关系的非凡能力，这是CNN单独无法实现的。

 此外，在提出的HAFormer中，ET层的数量L被限制为2，考虑到计算硬件的限制，并旨在在约束下实现最佳平衡。尽管堆叠更多的ET层可能带来更好的准确性结果，如图8所示，当L大于2时，性能提升会急剧放缓。此外，在像Cityscapes这样高分辨率的数据集上添加过多的ET层可能会对参数、计算和推理速度产生负面影响，并可能导致过拟合。

Ablation Study of the Correlation-weighted Fusion.

 为了解决CNN和Transformer之间的特征不匹配问题，并在解码过程中确保有效的特征恢复，引入了cwF机制。表 IV比较了使用所提cwF方法以及两种其他融合技术，即元素加法和串联，所获得的结果。表中显示了使用所有三种融合方法集成CNN和Transformer特征的增强分割准确性。值得注意的是，cwF在仅有一个ET层的情况下比基线提高了2.38%，在堆叠两层的情况下提高了4.06%。

 此外，从表 IV中观察到（a）与简单的元素加法融合方案相比，cwF在两种情况下分别实现了1.27%和1.90%的性能提升，仅在参数数量和FLOPs上有轻微的增加；（b）cwF在计算密集型的串联操作上分别实现了mIoU的0.84%和1.01%的提升，同时实现了大约5%的参数数量减少和15%的计算负载下降。这些实验结果进一步证明了cwF的有效性。

D. Comparisons with SOTA Methods

 本节全面评估和比较了所提方法与一些最先进的方法的性能和效率，以展示我所提
方法的优势。评估集中在三个方面：分割精度、模型参数和浮点运算（FLOPs）。

Evaluation Results on Cityscapes.

 表格V 展示了在 Cityscapes测试集上与先进语义分割方法的数量比较。表VI 详细列出了每类的结果，而图9-B显示了可视化结果。为了确保公平性，测试期间未使用增强技术，并且其他网络的数据参考了相关来源。当代语义分割模型主要分为两类：一类侧重于更大的规模和更高的精度，另一类则优先考虑实时实用性，在准确性和效率之间取得平衡。

 虽然大型模型实现了高精度，但它们FLOPs 和速度落后于轻量级模型，这使得它们不适合在资源有限的设备上进行实时处理。相比之下，轻量级模型如ENet [10]、ESPNet [12]、CGNet [53] 和 FPENet [59] 计算效率高。尽管它们的参数数量减少，但它们的整体性能，特别是在准确性方面，却有所欠缺。在准确性方面，EFRNet-16 [66] 显示了与我们结果的相似性。然而，值得注意的是，它的参数数量和 GFlops 是所提方法的两倍。显然，所提模型需要的参数和计算量更少，这突显了所提方法的有效性。

Evaluation Results on CamVid.
 为了进一步验证我们模型有效性和泛化能力，我们在CamVid数据集上与其他轻量级方法进行了比较，如表VII所示。虽然MGSeg [64]在准确性方面表现出色，超过了所提方法1.6个百分点，但它在参数数量上比所提方法多了22倍，这表明了一个不利的权衡。另一方面，SGCPNet [51]表现出了显著的速度，但缺乏准确性。相比之下，HAFormer在这些方面取得了更好的平衡。在CamVid数据集上的整体性能较低，与Cityscapes相比，这是由于其尺寸较小和分辨率较低，这突显了所提方法强大的泛化能力。图10中的可视化结果进一步展示了HAFormer的优势。

Speed Comparison.
 为了确保公平比较，所有方法都在相同的平台上执行，因为计算负载直接影响推理速度，这可能会根据设备的不同而有所不同。在控制评估中，使用单个 NVIDIA RTX 2080Ti GPU 来测量模型的执行时间。表VIII 详细比较了我们HAFormer与其他轻量级方法的速率和运行时间。实验涉及空间分辨率为 $512\times 1024$ 的评估，与具有官方代码的方法保持一致，以确保公平性。表VIII 展示了HAFormer的惊人速度，在处理大小为 512 ×1024 的图像流时，达到了105 fps的帧率，使其成为最快的几种方法之一。虽然DABNet在139 fps下运行，但HAFormer的74.2%的竞争准确性对于自动驾驶等现实世界应用来说意义重大。HAFormer在速度（105 fps）和准确性之间取得了良好的平衡，成为

五、总结

 文中引入了HAFormer，这是一种新的轻量级语义分割方法。设计了层次感知像素激励模块（HAPE）来提取增强的层次局部特征。此外，一个高效的Transformer模块在有限的计算负载下有效地捕捉了广泛的全球特征。然后，融入了一个相关性加权的融合（cwF）机制，将高度相关的CNN和Transformer特征结合起来，以进行改进的表征学习。通过在基准数据集上进行广泛的实验，所提方法已经显示出有效性和泛化能力，HAFormer实现分割准确性和计算效率之间的平衡。