BiUNet：基于双层路由注意力的轻量化医学分割网络

BiUNet：Lightweight Medical Segmentation Network Based on Bi⁃Level Routing Attention

摘要

针对视觉Transformer骨干提取网络计算开销大， 模型训练缓慢的问题， 同时为了进一步提升Transformer结构在医学图像领域的分割性能， 提出一种名为BiUNet的新型轻量级U型架构的医学图像分割网络。将输入医学图像整切成若干图像块后， 送入一种基于双层路由动态稀疏注意力机制的BiFormer转换器中， 通过组合下采样和特定块数的BiFormer模块， 构建多级金字塔结构实现特征提取。随后通过组合上采样和卷积模块， 相应构建多级金字塔结构进行特征解码， 进而实现像素级语义分割。该模型在3个医学数据集上依次取得了90.2%， 93.7%和85.6%的mIoU值以及5.55 G的Flops和28.10 M的参数量。结果表明， BiUNet能够以轻量化的效果有效提升医学图像分割的精度。

Abstract

To solve the problems of high computing cost and slow model training in vision Transformer backbone extraction network， and to further improve the performance of Transformer structure in the field of medical image segmentation， a new lightweight U-architecture medical image segmentation network named BiUNet was proposed. The input medical image was cut into several blocks， and then the blocks were fed into the BiFormer based on the dynamic sparse attention mechanism of Bi-level routing. By combining downsampling and BiFormer modules with a specific number of blocks， a multi-level pyramid structure was constructed to achieve feature extraction. Subsequently， the feature map output from the encoder was decoded by a multi-level pyramid structure which was constructed by combining the upsampling and convolution modules， and pixel-level semantic segmentation was realized. The model achieved 90.2%， 93.7% and 85.6% mIoU values as well as 5.55 G Flops and 28.10 M parameters on the three medical datasets sequentially. The results show that BiUNet can effectively improve the accuracy of medical image segmentation with a lightweight effect.

Graphical abstract

关键词

双层路由注意力机制 / Transformer结构 / 医学图像分割 / 轻量级 / U型结构

Key words

Bi-level routed attention mechanism / transformer structure / medical image segmentation / lightweight / U-shaped structure

至今为止， 黑色素瘤， 器官息肉等疾病仍然是导致人类痛苦甚至死亡的棘手问题， 对于上述疾病的预防和治疗， 相当依赖于先进的医学图像处理技术^［1］。自深度学习算法迅速发展以来， 基于人工智能的计算机工作方式已经在各种领域广泛应用^［2］。而在医学领域， 依托于医学图像分割结果的CAD（Computer-aided Diagnosis）^［3］方式也借助深度学习模型发生了重大变革。相较于分类和检测任务， 通常医疗专业及相关人员更关注于语义分割的结果^［4］。

在语义分割任务上， 尤其是医学图像分割任务， 现存的分割模型结构重度依赖于UNet的U型网络结构^［5］。而U型网络结构是由编码器、 解码器、 跳跃连接3个核心模块组成。在ViT^［6］（Vision Transformer）结构出现之前， 在计算机视觉领域， 由CNN（Convolutional Neural Networks）模块主导实现特征提取的功能。UNet本身的编码器和解码器模块核心部件由特定大小卷积核的卷积算子实现。研究者们利用CNN结合U型框架， 提出了种种改进UNet的网络模型。这一类结构主要有UNet++^［7］， UNet+++^［8］， AACN-UNet^［9］， DCSAU-Net^［10］， Attention-UNet^［11］等。而ViT结构出现之后， 将图片切割成图像块送入Transformer结构进行注意力计算的思想激励了研究者们针对这一新的特征提取骨干模块的探索。在TransUNet^［12］中， 作者先利用CNN生成特征图， 将其送入堆叠好的12个连续的Transformer块实现特征提取， 之后， 利用CNN实现特征解码， 在多器官分割任务上取得了较好的分割结果。而在SwinUNet^［13］中， 作者设计了独特的补丁扩展层， 使用纯Transformer模块搭建的U型结构取得了比TransUNet在同一数据集上更好的分割结果。此外， 利用Transformer结构进行医学图像分割的模型还有TransFuse^［14］， CoTr^［15］， DS-TransUNet^［16］， TransClaw UNet^［17］等。虽然Transformer结构带来了优秀的特征提取能力， 但是它所依赖的传统自注意力机制却带来了很大的计算开销。针对这一问题， 有研究者开始探索轻量化的Transformer结构， 如LeViT^［18］， Swin Transformer等^［19］。而Zhu等^［20］提出的BiFormer， 在保证自身轻量化的同时， 在图像分类、 目标检测、 语义分割等下游视觉任务上取得了超越Swin Transformer的效果。

为了解决现存医学图像分割网络计算开销大、 训练缓慢的问题， 更出于对实现更高分割精度的目标， 本文在BiFormer的启发下， 针对二维医学图像分割任务， 提出了BiUNet。BiUNet在编码器阶段充分利用BiFormer优秀的特征提取能力， 构建了金字塔式分级结构， 实现分级特征图提取， 利用双层路由机制保证模型轻量化， 在解码器阶段利用传统的卷积操作和独特的上采样模块， 实现特征解码和逐级分辨率恢复， 输出分割结果。

1 方　法

设计的BiUNet整体结构如图 1 所示。受典型U型架构的影响， BiUNet也由编码器、 解码器、 跳跃连接构成。编码器的核心特征提取模块由基于双层路由注意力机制的BiFormer构成。

如图 1 所示， 对于输入到网络里的一张皮肤病图片， 送入Patch Embedding模块， 处理成图像序列。将处理后的图像序列送入特定深度的BiFormer块与Patch Merging模块构成的层级结构中学习多尺度的特征表示； 随后送入由卷积模块和Patch Expanding， 将已学到的高级语义表示进行解析和分辨率逐级恢复； 在解码之前， 将编码器多级结构里输出的不同分辨率特征图与对应解码器上采样模块输出的同分辨率大小特征图进行跳跃连接， 最后恢复至输入大小分辨率得到预测分割图。

1.1 编码器

在编码器阶段， 构建了4级金字塔式结构， 实现特征提取。在第1层的Patch Embedding中， 首先使用2个3*3大小的卷积核， 步长为2的卷积层将3通道的输入（H\times W\times C）映射至通道维度C（本文设置为64）， 同时分辨率大小变至\frac{H}{\mathrm{4}}\times \frac{W}{\mathrm{4}}， 再送入N1个连续的BiFormer模块。在之后的3层里， 均为Patch Merging模块和BiFormer模块的组合， 第2层、 第3层、 第4层的BiFormer深度分别为N2， N3和N4（N1， N2， N3， N4的取值在后文BiUNet-T和BiUNet-S配置中给出）。Patch Merging模块本质上是1个卷积核大小为3\times3， 步长为2的卷积操作， 从而对特征图分辨率进行2倍下采样， 再通过对输出通道数的设置， 将通道数扩展为原先的2倍。经过这样的分层结构， 实现多尺度特征提取。

1.1.1 BiFormer

BiFormer的具体结构由图 2 所示。

对于每一个具体的BiFormer结构， 均由一个3×3的深度卷积模块、 一个层归一化模块和BRA（Bi-level Routing Attention）模块的组合， 以及一个层归一化模块和MLP（）的组合以残差连接的方式组合而成。对于输入的特征图， 通过深度堆叠的BiFormer模块， 分辨率和通道数均维持不变。

1.1.2 双层路由注意力模块（BRA）

BRA模块是一种动态的， 具有查询感知特性的稀疏注意力机制。它的核心思想是只保留一小部分路由区域， 而在粗粒度上过滤掉最不相关的键值对。在细粒度上， 在所保留路由区域的并集中进行标记到标记的注意力的计算。以单个输入的单头自注意力为例， 如式（1）， 首先将一个X\in R^{H\times W\times C}的输入特征图划分成S\times S个互不重叠的区域， 每个区域包含\frac{HW}{S^{\mathrm{2}}}个特征向量。这一步也可表述为将X重构成X_r\in R^{s^{\mathrm{2}}\times \frac{HW}{s^{\mathrm{2}}}\times C}。再对X_r线性投影得到Q,K,V\in R^{s^{\mathrm{2}}\times \frac{HW}{s^{\mathrm{2}}}\times C}。

\mathit{Q}={\mathit{X}}^r{\mathit{W}}^q， \mathit{K}={\mathit{X}}^r{\mathit{W}}^k， \mathit{V}={\mathit{X}}^r{\mathit{W}}^v，（1）

式中： W^q， W^k， W^v\in R^{C\times C}分别代表Q,K,V的映射权重。

第二步是先构造出一个有向图， 找到每个给定区域该关注的区域。具体实现是先通过平均每个区域的Q和K， 来得出区域级的Q_r,K_r。如式（2）， 通过Q^r和K^{r\mathrm{T}}矩阵相乘得到区域到区域亲和力图的邻接矩阵{\mathit{A}}^r。{\mathit{A}}^r中的元素值实际上衡量了两个区域语义上的相关程度。

{\mathit{A}}^r={\mathit{Q}}^r({\mathit{K}}^r{)}^{\mathrm{T}}\mathrm{。}（2）

接下来在每个区域仅保留top-k个连接来裁剪亲和力图， 通过行向topk算子构建一个路由索引矩阵{\mathit{I}}_r实现。如式（3）

{\mathit{I}}^r=topkIndex({\mathit{A}}^r)\mathrm{。}（3）

{\mathit{I}}^r的第i行代表了第i个区域关联度最高区域的k个索引值。

最后一步根据索引矩阵{\mathit{I}}^r，进行细粒度的注意力计算。即对于第i个区域中每个查询标记， 将以I_{(i,\mathrm{1})}^r,I_{(i,\mathrm{2})}^r,\cdots ,I_{(i,k)}^r的索引方式， 来关注k个路由区域并集区域中的所有键值对， 如式（4）， 聚集键张量和值张量得到聚集后的K^g和V^g。如式（5）， 对K^g和V^g作注意力计算。LCE（）函数实际上是通过一个5*5核大小的深度卷积进行参数化。

{\mathit{K}}^g=\mathrm{g}\mathrm{a}\mathrm{t}\mathrm{h}\mathrm{e}\mathrm{r}(\mathit{K},{\mathit{I}}^r),{\mathit{V}}^g=\mathrm{g}\mathrm{a}\mathrm{t}\mathrm{h}\mathrm{e}\mathrm{r}(\mathit{V},{\mathit{I}}^r),（4）

\mathit{O}=\mathrm{A}\mathrm{t}\mathrm{t}\mathrm{e}\mathrm{n}\mathrm{t}\mathrm{i}\mathrm{o}\mathrm{n}(\mathit{Q},{\mathit{K}}^g,{\mathit{V}}^g)+\mathrm{L}\mathrm{C}\mathrm{E}(\mathit{V})\mathrm{。}（5）

整个BRA计算复杂度的计算如式（6）

\begin{array}{c}FLOPs=FLOP{s}_{proj}+FLOP{s}_{routing}+FLOP{s}_{attn}=\\ \mathrm{3}HW{C}^{\mathrm{2}}+\mathrm{2}(s^{\mathrm{2}})C+\mathrm{2}HWk\frac{HW}{s^{\mathrm{2}}}C=\\ \mathrm{3}HW{C}^{\mathrm{2}}+C\left(\mathrm{2}{S}^{\mathrm{4}}+\frac{k{(HW)}^{\mathrm{2}}}{S^{\mathrm{2}}}+\frac{k{(HW)}^{\mathrm{2}}}{S^{\mathrm{2}}}\right)\ge \\ \mathrm{3}HW{C}^{\mathrm{2}}+\mathrm{3}C{\left(\mathrm{2}{S}^{\mathrm{4}}\frac{k{(HW)}^{\mathrm{2}}}{S^{\mathrm{2}}}\frac{k{(HW)}^{\mathrm{2}}}{S^{\mathrm{2}}}\right)}^{\frac{\mathrm{1}}{\mathrm{3}}}=\\ \mathrm{3}HW{C}^{\mathrm{2}}+\mathrm{3}C{k}^{\frac{\mathrm{2}}{\mathrm{3}}}{(\mathrm{2}HW)}^{\frac{\mathrm{4}}{\mathrm{3}}},\end{array}（6）

式中： H， W为输入的特征图分辨率； C为特征图通道数； k为参与注意力计算的数目。等号仅在式（7）成立时实现

S={\left(\frac{k}{\mathrm{2}}{(HW)}^{\mathrm{2}}\right)}^{\frac{\mathrm{1}}{\mathrm{6}}}\mathrm{。}（7）

最终， BRA实现了O{\left(HW\right)}^{\frac{\mathrm{4}}{\mathrm{3}}}的计算复杂度。

1.2 解码器

对于编码器输出的多尺度特征图。在解码器阶段相应构建4级金字塔式结构。每一层级均由两个3×3核大小，步长为1的卷积模块和上采样模块组成。这里的上采样模块， 本文主要参考了SwinUNet中的Patch Expanding模块， 对于前3层的上采样， 它的核心实现是先将特征图通道图扩展为之前的2倍， 然后经过矩阵重排的操作实现分辨率扩展和通道数减半， 在最后一个上采样块， 将倍率调整为4。并且将跳跃连接设置在1/16， 1/8， 1/4分辨率处。最后通过1×1核大小的卷积层将特征通道数映射成类别数， 得到输出分割图。

2 实　验

2.1 数据集

本文所使用的数据集为CVC-ClinicDB， 2018 Data Science Bowl， ISIC-2018数据集。CVC-ClinicDB数据集是从结肠镜检查视频中提取的帧数据库， 主要用于息肉分割任务， 曾作为MICCAI 2015自动息肉分割的训练数据库。 2018 Data Science Bowl数据集主要用于自动化细胞核检测， 来源于2018数据科学碗竞赛中细胞显微镜图像。ISIC-2018数据集来源于国际皮肤成像协作组在2018年发布的一个皮肤镜图像数据集，曾作为2018皮肤病分类挑战赛训练数据库， 本文采用了task1皮肤病分割数据集。相关细节规模设置如表 1 所示。

2.2 实现细节

本文所有实验均在pytorch1.10和python3.8下完成， 所有网络遵循同样的实验配置， 所使用的显卡型号为Nvidia RTX 3090， 显存大小为24 GB。在训练过程中， 对３个数据集都统一裁剪到224×224的大小， batch size设置为16， 学习率设置为1e-4， 使用交叉熵损失函数搭配Adam优化器进行优化， 以５折交叉验证的方式共训练150个epoch。其中BiUNet， TransUNet， LeViT-UNet^［21］， HiFormer^［22］还需要加载各自预训练权重。

2.3 实验结果及分析

2.3.1 评估指标

采用F1分数（F1）、 准确率（Acc）、 召回率（Recall）、　平均交并比（mIoU）、　精确率（Precision）５个评价指标评估网络性能。在医学图像分割领域， 通常重点关注mIoU和F1值。相关公式为

Acc=\frac{TP+TN}{TP+FP+TN+FN},（8）

Precision=\frac{TP}{TP+FP},（9）

Recall=\frac{TP}{TP+FN},（10）

F_{\mathrm{1}}=\mathrm{2}\mathrm{*}\frac{Precision\mathrm{*}Recall}{Precision+Recall},（11）

mIoU=\frac{TP}{TP+FP+FN},（12）

式中： TP代表实际为正样本， 预测也为正样本的样本数； TN代表实际为负样本， 预测也为负样本的样本数； FP代表实际为负样本， 预测为正样本的样本数； FN代表实际为正样本， 预测为负样本的样本数。mIoU值和F1值越大， 代表预测的图像与标签图像重合度越高， 准确率越高。对于计算开销的评估， 本文采用计算量GFlops和参数量（Param）进行衡量。

2.3.2 实验结果分析

表 2 展示了BiUNet与其他分割算法在ISIC-2018上的对比结果。其中BiUNet-T是指BiFormer块数组合为［N₁=2， N₂=2， N₃=8， N₄=2］， 并且使用BiFormer-T的预训练权重。BiUNet-S是指BiFormer块数组合为［N₁=4， N₂=4， N₃=18， N₄=4］， 并且使用BiFormer-S的预训练权重。BiUNet-S以90.2%的mIoU值和94.6%的F1分数值取得最佳的表现。BiUNet-S超越了TransUNet约3%以及LeViT-UNet约1.5%的mIoU值。

表3是BiUNet与其他分割算法在CVC-ClinicDB上的对比结果。BiUNet-S也以93.7%的mIoU值和96.7%的F1分数值取得最佳的分割精度。BiUNet-S超越了HiFormer2.6%以及DCSA-UNet约3.7%的mIoU值。

表 4 是BiUNet与其他分割算法在2018 Data Science Bowl上的对比结果。BiUNet-S依旧以85.6%的mIoU值以及91.9%的F1分数值取得了最高的分割成绩。BiUNet-S超越了TransUNet约0.6%以及UNet约4%的mIoU值。

综合表 2、 表 3、 表 4 的结果， 除上述表述外， 事实上， BiUNet-S在3个数据集上5项指标均达到了最佳的表现， 充分说明了BiUNet优越的性能。部分主流分割算法的对比可视化结果如图 3 所示。图 3 中， 在CVC-ClinicDB数据集上， BiUNet-S的优势相对突出， 对息肉缺口部分信息预测准确率明显优于U-Net， LeViT-UNet以及HiFormer。在2018 Data Science Bowl数据集上， BiUNet-S对于细胞核边界信息的预测精度也明显优于Attention-UNet和LeViT-UNet。

表 5 是各个模型计算开销的统计结果。其中， BiUNet-T以3.44G的Flops取得了最小的计算复杂度。而BiUNet-S以5.55 G的Flops和28.10 M的参数量在ViT结构（HiFormer， LeViT-UNet， TransUNet）里取得了最小的计算开销。

3 消融实验

为了探索BiUNet网络性能细节， 本文在3个数据集上设置了2组消融实验， 结果如表 6、 表 7 所示。为简化数据集名称表述， 将2018 Data Science Bowl简记为cell， ISIC-2018简记为skin， CVC-ClinicDB简记为cvc。

CZ-BiUNet是指将BiUNet-S中上采样层由 Patch Excpanding替换为双线性插值， 表 6 结果表明Patch Expanding层作为上采样能够取得更好的分割效果。

DE-BiUNet是指使用Dice Loss损失函数替代交叉熵损失函数对BiUNet进行优化， 在skin， cvc， cell 3个数据集上的实验结果表明， 交叉熵损失函数更加契合BiUNet-S。

4 结　语

为了提升现有以Transformer为主干特征提取网络的分割算法的精度和解决Transformer结构计算开销大的问题， 本文提出了一个以BiFormer为骨干结构的U型医学图像分割网络。充分利用双层路由稀疏注意力机制， 不仅有效地减小了计算代价， 还更高效地提升了分割性能。但是在实验中未能以更小的计算代价超越快速CNN模型DCSAU-Net的性能， 所以在未来的研究中， 将继续探索更加高效的轻量化视觉Transformer分割模型， 实现以更小的计算代价获取更好的分割能力。

参考文献

[1] MA X， NIU Y， GU L， et al. Understanding adversarial attacks on deep learning based medical image analysis systems［J］. Pattern Recognition， 2021， 110： 107332.

[2] NARAYAN V， AWASTHI S， FATIMA N， et al. Deep learning approaches for human gait recognition： a review［C］//2023 International Conference on Artificial Intelligence and Smart Communication （AISC）. IEEE， 2023： 763-768.

[3] HE S， MINN K T， SOLNICA-KREZEL L， et al. Deeply-supervised density regression for automatic cell counting in microscopy images［J］. Medical Image Analysis， 2021， 68： 101892.

[4] ZHOU S， NIE D， ADELI E， et al. High-resolution encoder-decoder networks for low-contrast medical image segmentation［J］. IEEE Transactions on Image Processing： a Publication of the IEEE Signal Processing Society， 2019.doi： 10.119/TIP.2019.2919937 .

[5] RONNEBERGER O， FISCHER P， BROX T. U-net： Convolutional networks for biomedical image segmentation［C］//NAVAB N， HORNEGGER J， WELLS W M， et al. Lecture Notes in Computer Science. Cham： Springer International Publishing， 2015： 234-241.

[6] DOSOVITSKIV A， BEYER L， Kolesnikov A， et al. An image is worth 16x16 words： transformers for image recognition at scale［EB/OL］. arXiv： 2020.

[7] ZHOU Z， RAHMAN SIDDIQUEE M M， TAJBAKHSH N， et al. Unet++： A nested U-net architecture for medical image segmentation［C］//Deep Learning in Medical Image Analysis and Multimodal Learning for Clinical Decision Support. Cham： Springer International Publishing， 2018： 3-11.

[8] HUANG H， LIN L， TONG R， 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）， 2020： 1055-1059.

[9] WU H， WANG J， WANG X， et al. Attention augmented ConvNeXt UNet for rectal tumour segmentation［EB/OL］. http：arxiv.org/abs/2210.002227v 2.

[10] XU Q， MA Z， HE N， et al. DCSAU-Net： a deeper and more compact split-attention U-Net for medical image segmentation［J］. Computers in Biology and Medicine， 2023， 154： 106626.

[11] SCHLEMPER J， OKTAY O， SCHAAP M， et al. Attention gated networks： Learning to leverage salient regions in medical images［J］. Medical Image Analysis， 2019， 53： 197-207.

[12] CHEN J， LU Y， YU Q， et al. TransUNet： transformers make strong encoders for medical image segmentation［EB/OL］.

[13] CAO H， WANG Y， CHEN J， et al. Swin-UNet： Unet-like pure transformer for medical image segmentation［C］//European Conference on Computer Vision. Cham： Springer Nature Switzerland， 2023： 205-218.

[14] ZHANG Y， LIU H， HU Q. TransFuse： fusing transformers and CNNs for medical image segmentation［C］//Medical Image Computing and Computer Assisted Intervention-MICCAI 2021. Cham： Springer International Publishing， 2021： 14-24.

[15] XIE Y， ZHANG J， SHEN C， et al. CoTr： efficiently bridging CNN and transformer for 3D medical image segmentation［C］//Medical Image Computing and Computer Assisted Intervention-MICCAI 2021. Cham： Springer International Publishing， 2021： 171-180.

[16] LIN A， CHEN B， XU J， et al. DS-TransUNet： dual swin transformer U-net for medical image segmentation［J］. IEEE Transactions on Instrumentation and Measurement， 2022， 71： 4005615.

[17] YAO C， HU M， ZHAI G， et al. TransClaw U-net： claw U-net with transformers for medical image segmentation［EB/OL］. httP//arxiv.org/abs/2017.05188v1.

[18] GRAHAM B， EL-NOUBY A， TOUVRON H， et al. LeViT： a vision transformer in ConvNet’s clothing for faster inference［C］//Proceedings of the IEEE/CVF International Conference on Computer Vision， 2021： 12239-12249.

[19] LIU Z， LIN Y， CAO Y， et al. Swin transformer： hierarchical vision transformer using shifted windows［C］//Proceedings of the IEEE/CVF international conference on computer vision. 2021： 19992-10002.

[20] ZHU L， WANG X， KE Z， et al. BiFormer： vision transformer with bi-level routing attention［C］//Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023： 10323-10333.

[21] XU G， WU X， ZHANG X， et al. LeViT-UNet： make faster encoders with transformer for medical image segmentation［EB/OL］.

[22] HEIDARI M， KAZEROUNI A， SOLTANY M， et al. HiFormer： hierarchical multi-scale representations using transformers for medical image segmentation［C］//Proceedings of the IEEE/CVF Winter Conference on Applications of Computer Vision. 2023： 6191-6201.