基于轴-Transformer的医学图像分割模型Axial-TransUNet

摘 要：针对TransUNet中Transformer自注意力机制计算复杂度高、捕获位置信息能力不足的问题，提出一种基于轴向注意力机制的医学图像分割网络Axial-TransUNet。该网络在保留TransUNet网络编码器、解码器以及跳跃连接的基础上，使用基于轴向注意力机制的残差轴向注意力块代替TransUNet的Transformer层。实验结果表明，在多个医学数据集上，相较于TransUNet等其他医学图像分割网络，Axial-TransUNet的Dice系数、交并比IoU有更好的表现。与TransUNet相比，Axial-TransUNet网络的参数量与浮点运算数（FLOPs）分别降低14.9%和30.5%。可见，Axial-TransUNet有效降低了模型复杂度，并增强了模型捕获位置信息的能力。

关键词：医学图像分割；卷积神经网络；位置信息；计算复杂度；轴向注意力机制

中图分类号：TP183；TP39 文献标识码：A 文章编号：2096-4706（2024）16-0028-06

Axial-TransUNet of Medical Image Segmentation Model Based on Axis-Transformer

Abstract： A medical image segmentation network Axial-TransUNet based on Axial Attention Mechanism is proposed to address the issues of high computational complexity and insufficient ability to capture positional information in the Transformer Self-Attention Mechanism in TransUNet. On the basis of retaining the TransUNet network encoder， decoder， and skip connections， this network uses residual axial attention blocks based on Axial Attention Mechanism to replace the Transformer layer of TransUNet. The experimental results show that compared to other medical image segmentation networks such as TransUNet， Axial TransUNet performs better in Dice coefficient and intersection union ratio on multiple medical datasets. Compared with TransUNet， the parameter count and FLOPs of the Axial TransUNet network are reduced by 14.9% and 30.5%， respectively. It can be seen that Axial TransUNet effectively reduces model complexity and enhances the model's ability to capture positional information.

Keywords： medical image segmentation; Convolutional Neural Networks; positional information; computational complexity; Axial Attention Mechanism

0 引 言

传统的人工医学图像识别通常依赖于专业医生的经验和主观判断，且难适应大规模数据集和复杂的疾病模式。因此，深度学习在医学图像分析方面的应用具有重要意义。

在医学图像分割任务中，卷积神经网络（ConvolutionalNeuralNetworks， CNNs）特别是全卷积网络（FullyConvolutionalNetworks， FCNs）[1]，凭借其优异的分割性能占据了主导地位。以U-Net [2]为代表的经典U型结构是最为流行和广泛使用的结构之一。U-Net主要由对称的编码器和解码器组成，结合跳跃连接避免下采样过程中的信息丢失。U-Net的变体主要有UNet++ [3]、UNet 3+ [4]、Res-UNet [5]、ResUNet++ [6]、3D U-Net [7]、Attention U-Net [8]等，这些模型均在不同的医学图像分割任务上展示了优异的性能。

然而，由于卷积运算的内在局部性，它们捕获长距离依赖的能力不足，无法很好地利用全局信息。为了克服这一局限性，Wang [9]等人提出建立基于CNN特征的自我注意力机制，通过非局部操作捕获长期依赖，增大模型感受野。此外，为序列建模而设计的Transformer在自然语言处理领域（NLP）取得良好成效。受此启发，研究人员将其应用到计算机视觉领域。为使Transformer适用于计算机视觉任务，研究人员对其进行了一些修改。例如，Parmar等人[10]在每个查询像素的局部邻域中应用自关注；Child等人[11]提出稀疏变换器，采用全局自关注的可扩展近似；Dosovitski等人[12]提出的Vision Transformer将图像块的线性嵌入式序列作为Transformer的输入，使具有全局自关注的注意力机制直接应用于全尺寸图像，是当时ImageNet分类的最先进技术。然而，纯Transformer方法会带来特征损失等问题，所以建立基于CNN和Transformer相结合的网络架构受到了广泛的关注。例如，Yao等人[13]所提出的Transclaw U-Net在编码器部分将卷积和Transformer结合，分别用于提取浅层特征和全局信息；Xu等人[14]提出的LeViT-UNet结合了LeViT的全局感知能力和UNet的分割能力，生成精确的分割结果。其中，Chen [15]等人提出的TransUNet使用CNN和Transformer的混合编码器结构，结合CNN对局部细节信息以及Transformer对全局信息的提取优势来提升分割精度，取得了优异表现。TransUNet [15]编码器部分的Transformer建立在VisionTransformer（ViT）之上，Vision Transformer基于全局自注意力机制提取注意力，其计算的注意力矩阵大小与输入序列长度成平方关系，需要消耗大量的内存与计算资源，所以TransUNet在处理较大数据时存在一定的局限性。同时，ViT所采用的自注意力机制会将输入序列中的所有位置进行交互和整合，而不考虑位置之间的差异，这可能导致一些细节信息在编码过程中被模糊化或丢失。为解决上述问题，Wang [16]等人提出轴向注意力机制，在保持全局感受野的前提下，将2D注意力按高度轴和宽度轴顺序分解为两个1D注意力，有效降低了参数量，弥补了自注意力机制的不足。同时，轴向注意力机制引入更加精确的位置编码，使注意模型对位置更加敏感并具有边际成本。此外，研究人员还将轴向注意力模块引入残差结构中形成残差轴向注意力块，进一步提升分割精度。

为解决TransUNet参数量大、计算复杂度高以及对位置信息交互能力差的问题，本文提出Axial-TransUNet解决了这些问题带来的局限性。Axial-TransUNet将基于轴向注意力机制的残差轴向注意力块应用于轴-Transformer部分，建立CNN和轴-Transformer混合编码器架构，在保证较大感受野的前提下实现复杂度的降低，同时增强模型捕获位置信息的能力。在多个医学图像数据集上与有代表性的算法进行评估，Axial-TransUNet表现出了不错的性能。

1 相关工作

1.1 CNN与Transformer结合

以U-Net为代表的卷积神经网络存在感受野较小、捕捉长距离依赖不足的问题。受NLP中Transformer的启发，Dosovitski等人[12]提出的Vision Transformer将图像拆分为块，并提供这些块的线性嵌入序列作为Transformer的输入，利用其捕获长距离依赖的优势克服了这一局限性。然而，纯Transformer仍存在特征损失的问题。为解决上述问题，Chen等人[15]提出的TransUNet采用CNN与Transformer相结合的U型结构，将两者的优势结合起来并取得优异的性能。其中，编码器部分的Transformer建立在Vision Transformer之上，将来自CNN特征映射的标记化图像块编码处理成用于提取全局上下文的输入序列。然而，由于Vision Transformer中自注意力机制是全局的，注意力矩阵大小与输入序列长度成平方关系，所以TransUNet的计算复杂度较高。同时，自注意力机制在计算注意力权重时，只考虑了特征本身的相似性，而忽略了位置信息。这可能导致在某些任务中，模型无法准确地捕捉到物体的准确位置。

1.2 轴向注意力机制

Wang等人[16]提出的轴向注意力机制是一种具有全局感受野的独立自注意力机制，其核心思想是将2D注意力分解为两个1D注意力模块，第一个模块在特征图高度轴上进行自我关注，第二个模块在宽度轴上进行操作，有效地模拟了原始的自注意机制，具有更好的计算效率。同时，轴向注意力机制通过在计算轴向注意力权重时考虑位置信息，可以更好地保留物体的空间结构和位置关系，提高模型在位置感知任务中的性能。

2 方法介绍

Axial-TransUNet的整体结构如图1所示。输入图像首先被送入混合编码器模块，由CNN进行逐级特征提取，然后送入轴-Transformer层进行全局上下文信息提取。解码器模块将编码特征进行上采样并将其与高分辨率CNN特征图相结合，实现特征聚合。其中，轴-Transformer部分通过采用基于轴向注意力机制的残差轴向注意力块。从分割结果及参数量对比结果可以看出，Axial-TransUNet实现了更好的分割精度，且参数量相较TransUNet [15]有明显降低。

2.1 融合轴向注意力机制的编码器模块

在Axial-TransUNet中，编码器模块包括残差特征提取部分和轴-Transformer部分。

2.1.1 残差特征提取器

残差特征提取器负责对输入图像进行逐级特征提取，并将不同层次的高分辨率特征图输出至跳跃连接路径中。首先将输入图像Ho×Wo×Co通过卷积变换到指定大小H×W×C，其中H、W、C分别表示图像的高度、宽度和通道数，然后通过四级残差特征提取块进行特征提取。每一级残差特征提取块将图像通道数变为原来的2倍，将图像的分辨率降低为原来的1/2。其中，每个残差特征提取块基于3个堆叠的残差单元进行特征提取。每个残差单元包含3个卷积层，每个卷积层结构如图2所示。

每个残差单元可以用下面的通用形式表示：

其中，Xl和Xl+1为第l个单元的输入和输出，Wl为可学习的权重矩阵，F为残差函数，f为激活函数，h为恒等映射函数，通常取h（Xl） = Xl。

2.1.2 轴-Transformer

轴-Transformer由若干残差轴向注意力块堆叠而成，残差轴向注意力块基于轴向注意力机制分别在高度轴及宽度轴上计算注意力，并采用残差结构提取特征。

1）轴向注意力机制。一个轴向注意力层沿着一个特定的轴传播信息。如图3所示，输入图像X分别与可学习的权重矩阵WQ、Wk、WV相乘得到查询q、键k、值v，查询q与键v相乘得到注意力矩阵，同时，为q、k添加可学习的位置矩阵，二者的输出与注意力矩阵拼接并进行Softmax操作后得到最终的注意力矩阵。注意力矩阵与添加位置编码信息的值v进行矩阵乘法操作得到最终的输出Y。对于一张H×W×C的输入图像，轴向注意力机制将图像宽度轴上的轴向注意力层定义为简单的1D位置敏感的自我注意力，并对高度轴使用类似定义。沿宽度轴的轴向注意层定义如下：

其中，m为区域大小，查询q = WQ X，键k = Wk X，值v = WV X，WQ、Wk、WV均为可学习的权重矩阵，rq、rk、rv分别为q、k、v的可学习的位置编码。

2）残差轴向注意力块。如图4所示，将ResNet中残差瓶颈块的3×3卷积替换为两个轴向层，其中一个用于高度轴，另一个用于宽度轴，形成残差轴向注意力块。

2.2 解码器模块

解码器部分主要通过多层解码模块对图像信息进行恢复。首先通过上采样操作将低分辨率图像转换成高分辨图像。然后将上采样后的特征图与对应的编码器层的特征图进行跳跃连接，通过逐层特征融合和特征提取，逐渐将特征图恢复到与原始输入图像相同的分辨率。最后，解码器模块使用一个卷积层将恢复到原始分辨率的特征图映射到与输出图像大小相同的特征空间中，并使用激活函数将其映射到0～1之间的数值范围内，得到预测结果[17]。

2.3 评价指标

对于医学图像分割任务，采用交并比IoU（Intersection

over Union）和相似系数Dice（Dice Similarity Coefficient）来评估模型的分割性能[18]。

2.3.1 Dice系数

Dice系数是一种集合相似度度量函数，通常用于计算两个样本的相似度，取值范围在[0，1]，值越大，分割效果越好。其计算式如下：

其中，X和Y分别表示真实标签和预测结果。

2.3.2 交并比IoU

IoU是预测结果和真实标签之间的重叠区域除以两者之间的联合区域，取值范围为[0，1]，值越大，分割效果越好。其计算式如下：

其中，X和Y分别为真实标签和预测结果。

3 实验分析

3.1 数据集

实验采用6个公开的医学图像分割数据集，如表1所示，包括5个息肉数据集Kvasir-SEG、Kvasir-Sessile、CVC-ClinicDB、CVC-ColonDB和ETIS-LaribPolypDB以及1个细胞核数据集Cell-Nuclei。每个数据集均以8：2的比例划分为训练集和测试集。

3.2 实验细节及环境

实验是在单个11 GB的RTX 2080 GPU上运行的，实验环境为Python 3.8和PyTorch 1.11.0。数据的预处理主要包括归一化、随机旋转、随机翻转、色调饱和度、亮度及对比度调节等。对于5个息肉数据集Kvasir-SEG、Kvasir-Sessile、CVC-ClinicDB、CVC-ColonDB和ETIS-LaribPolypDB，其输入图像大小为256×256，细胞核数据集Cell-Nuclei的输入图像大小为96×96。残差轴向注意力块默认堆叠1层，batch_size为8，使用BCEDiceLoss作为损失函数，学习率为0.01，动量为0.9，权重衰减为1×10-3的SGD作为优化器，迭代次数为100。对于模型，使用二分类交叉熵损失（BCELoss）和Dice损失（DiceLoss）的组合来训练Axial-TransUNet，如下式所示：

3.3 实验结果及分析

实验在6个医学图像分割数据集上对Axial-TransUNet进行训练，并与模型U-Net [2]、UNet++ [3]、

Attention U-Net [8]、Res-UNet [5]、ResUNNet++ [6]、TransUNet [15]进行分割精度以及参数量的对比。由实验结果可知，Axial-TransUNet的分割精度优于TransUNet [15]，并在CVC-ClinicDB、CVC-ColonDB、ETIS-LaribPolypDB三个数据集上达到最佳分割性能。除此之外，Axial-TransUNet的参数量相对TransUNet [15]也有显著降低。

3.3.1 模型定量分析

Axial-TransUNet与U-Net [2]、UNet++ [3]、Res-UNet [5]、ResUNet++ [6]、Attention UNet [8]、TransUNet [15]

在6个公开医学图像分割数据集上的定量分析结果如表2所示。由表2可知，Axial-TransUNet的分割精度相较TransUNet [15]均有不同程度的提升。在CVC-ClinicDB、CVC-ColonDB、ETIS-LaribPolypDB上，Axial-TransUNet获得了最好的分割效果。其中，在Kavsir-SEG数据集上，Axial-TransUNet相对TransUNet在Dice上提升了约5.5%，在IoU上提升了约8.8%。在CVC-ClinicDB数据集上，Axial-TransUNet相对TransUNet在Dice上提升了约10.1%，在IoU上提升了约16.0%。在CVC-ColonDB数据集上，Axial-TransUNet相对TransUNet在Dice上提升了约27.9%，在IoU上提升了约45.6%。在Kvasir-sessile数据集上，Axial-TransUNet相对TransUNet在Dice上提升了约10.9%，在IoU上提升了约11.9%。在ETIS-LaribPolypDB数据集上，Axial-TransUNet相对TransUNet在Dice上提升了约63.4%，在IoU上提升了约92.2%。在Cellnuclei数据集上，Axial-TransUNet相对TransUNet在Dice上提升了约1.3%，在IoU上提升了约1.8%。

3.3.2 模型参数量对比

如表3所示，与其他两种引入注意力机制的网络Attention UNet [8]、TransUNet [15]相比，Axial-TransUNet参数量更少，计算复杂度更低。其中，相比于TransUNet，Axial-TransUNet参数量降低了约14.9%，计算复杂度降低约30.5%。

3.3.3 模型可视化结果

不同模型在不同数据集上的分割结果如图5所示。从图中可以看出，相较于纯CNN的方法，Axial-TransUNet对边界信息分割更为精准。例如，第一行中，UNet [2]、UNet++ [3]等对图像的边界信息不够敏感，出现了特征丢失的问题，而Axial-TranUNet保留了很好的边界信息。

相较于TransUNet [15]，Axial-TransUNet能够较好地避免特征损失、过度分割的问题。例如，第三行中，TransUNet对白色区域左半部分进行分割时出现了特征损失的问题，而Axial-TransUNet则较好地保留了图像的特征。同样，Axial-TransUNet也实现了更好的边缘预测。

4 结 论

TransUNet模型的Transformer基于自注意力机制进行注意力计算。由于自注意力机制在计算注意力权重时需要考虑所有位置之间的关系，注意力矩阵大小与输入序列长度成平方关系，所以Transformer的计算复杂度较高。同时，由于自注意力机制仅对查询向量采用位置编码，所以Transformer在捕获位置信息方面能力较低。针对上述问题，本文提出Axial-TransUNet网络。Axial-TransUNet将轴向注意力机制融入Transformer中，轴向注意力机制仅考虑处于同一轴向位置之间的关系，因此可以通过降低注意力权重计算的复杂度来减少模型参数量，并且可以更好地捕捉输入数据中不同方向上的特征和关系。这种方向性信息的建模有助于提升模型的表示能力和泛化性能。此外，轴向注意力机制引入的更丰富的可学习的位置矩阵，使模型可以根据不同位置的需要自适应地调整注意力权重，从而更好地利用位置信息进行建模和预测。基于轴向注意力机制的Axial-TransUNet有效降低了TransUNet的计算复杂度，同时有效弥补了TransUNet捕捉位置信息能力的不足，提高了分割精度。如表2所示，Axial-TransUNet网络在6个医学图像数据集上表现出比TransUNet更优的分割性能，在CVC-ClinicDB、CVC-ColonDB、ETIS-LaribPolypDB这3个数据集其分割性能上达到最佳。如表3所示，与其他两种引入注意力机制的网络AttentionUNet、TransUNet相比，Axial-TransUNet参数量与计算复杂度均较低。这也表明，轴Transformer在未来仍有很高的研究价值。

然而，轴向注意力机制并不是对自注意力机制的全面替代，而是在特定任务或场景下的一种改进选择。在某些情况下，自注意力机制可能仍然是更合适的选择，特别是当全局信息之间的关系对任务至关重要时。

参考文献：

[1] LONG J，SHELHAMER E，DARRELL T. Fully Convolutional Networks for Semantic Segmentation[C]//2015 IEEE Conference on Computer Vision and Pattern Recognition （CVPR）.Boston：IEEE，2015：3431-3440.

[2] RONNEBERGER O，FISCHER P，BROX T. U-Net： Convolutional Networks for Biomedical Image Segmentation [C]//Proceedings of the Lecture Notes in Computer Science.Munich：Springer，2015：234-241.

[3] ZHOU Z W，SIDDIQUEE M M R，TAJBAKHSH N，et al. UNet++： A Nested U-Net Architecture for Medical Image Segmentation [C]//Medical Image Computing and Computer-Assisted Intervention - MICCAI 2015.Granada：Springer，2018：3-11.

[4] HUANG H M，LIN L F，TONG R F，et al. UNet 3+： A Full-Scale Connected UNet for Medical Image Segmentation [C]//ICASSP 2020 - 2020 IEEE International Conference on Acoustics， Speech and Signal Processing （ICASSP）.Barcelona：IEEE，2020：1055-1059.

[5] XIAO X，LIAN S，LUO Z M，et al. Weighted Res-UNet for High-Quality Retina Vessel Segmentation [C]//2018 9th International Conference on Information Technology in Medicine and Education （ITME）.Hangzhou：IEEE，2018：327-331.

[6] JHA D，SMEDSRUD P H，RIEGLER M A，et al. ResUNet++： An Advanced Architecture for Medical Image Segmentation [C]//2019 IEEE International Symposium on Multimedia （ISM）.San Diego：IEEE，2019：225-230.

[7] ÇIÇEK Ö，ABDULKADIR A，LIENKAMP S S，et al. 3D U-Net： Learning Dense Volumetric Segmentation from Sparse Annotation [C]//Medical Image Computing and Computer-Assisted Intervention - MICCAI 2016.Athens：Springer， 2016：424-432.

[8] OKTAY O，SCHLEMPER J，FOLGOC L L，et al. Attention U-Net： Learning Where to Look for the Pancreas [J/OL].arXiv：1804.03999 [cs.CV].[2023-11-15].https：//arxiv.org/abs/1804.03999.

[9] WANG X L，GIRSHICK R，GUPTA A，et al. Non-local Neural Networks [C]//2018 IEEE/CVF Conference on Computer Vision and Pattern Recognition.Salt Lake City：IEEE，2018：7794-7803.

[10] PARMAR N，VASWANI A，USZKOREIT J，et al. Image Transformer [J/OL].arXiv：1802.05751 [cs.CV].[2023-10-22].https：//arxiv.org/abs/1802.05751v3.

[11] CHILD R，GRAY S，RADFORD A，et al. Generating Long Sequences with Sparse Transformers [J/OL].arXiv：1904.10509 [cs.LG].[2023-10-20].https：//arxiv.org/abs/1904.10509.

[12] DOSOVITSKIY A，BEYER L，KOLESNIKOV A，et al. An Image is Worth 16×16 Words： Transformers for Image Recognition at Scale [J/OL].arXiv：2010.11929 [cs.CV].[2023-10-20].https：//arxiv.org/abs/2010.11929.

[13] YAO C，HU M H，LI Q L，et al. Transclaw U-Net： Claw U-Net With Transformers for Medical Image Segmentation [C]//2022 5th International Conference on Information Communication and Signal Processing （ICICSP）.Shenzhen：IEEE，2022：280-284.

[14] XU G P，ZHANG X，FANG Y，et al. LeViT-UNet： Make Faster Encoders with Transformer for Medical Image Segmentation [J/OL].arXiv：2107.08623 [cs.CV].[2023-09-19].https：//arxiv.org/abs/2107.08623.

[15] CHEN J N，LU Y Q，YU Q H，et al. TransUNet： Transformers Make Strong Encoders for Medical Image Segmentation [J/OL].arXiv：2102.04306 [cs.CV].[2023-09-19].https：//arxiv.org/abs/2102.04306.

[16] WANG H Y，ZHU Y K，GREEN B，et al. Axial-DeepLab： Stand-Alone Axial-Attention for Panoptic Segmentation [C]//Computer Vision - ECCV 2020.Glasgow：Springer，2020：108-126.

[17] HE K M，ZHANG X Y，REN S Q，et al. Deep Residual Learning for Image Recognition [C]//2016 IEEE Conference on Computer Vision and Pattern Recognition （CVPR）.Las Vegas：IEEE，2016：770-778.

[18] 丘文峰.基于Python的医学图像处理框架及其应用 [D].广州：华南师范大学，2010.