一种基于TransUNet的双分支并行医学图像分割模型

A Two-Branch Parallel Medical Image Segmentation Model Based on TransUNet

摘要: 医学图像分割任务中，为充分利用TransUNet模型能有效捕获全局和局部特征的优势，在其基础上，提出DouTransNet模型，编码器部分针对单分支Transformer模块学习角度单一、容易丢失细节特征的问题，将Transformer设计成双分支并行结构，来提取不同尺度的特征，融合两个分支的特征，实现特征互补；针对在融合两个分支的特征时可能存在冗余信息问题，添加多核并行池化模块，在保留多尺度特征的同时去除冗余信息；在解码器设计多尺度融合模块(USF)，融合来自编码器的三个尺度信息，有效弥补编码器与解码器之间的信息差距。在Synapse和ACDC数据集上进行了多次对比实验，Synapse数据集上平均DSC系数可达79.20%，较TransUNet模型提高3.34%，HD距离为25.24%，降低了11.67%；在ACDC数据集上平均DSC系数可达90.30%，提高1.67%。

Abstract: In the medical image segmentation task, in order to make full use of the advantages of TransUNet model which can effectively capture global and local features, the DouTransNet model is proposed based on it. The encoder part aims at the problems of single learning Angle and easy loss of detail features in single-branch Transformer module. Transformer is designed as a two-branch parallel structure to extract features of different scales and fuse features of two branches to achieve feature complementarity. To solve the problem of redundant information when fusing the features of two branches, a multi-core parallel pooling module is added to remove redundant information while retaining multi-scale features. In the decoder design multi-scale fusion module (USF), the information of three scales from the encoder is fused to effectively bridge the information gap between the encoder and the decoder. Several comparison experiments were conducted on Synapse and ACDC data sets. The average DSC coefficient on Synapse data set can reach 79.20%, which is 3.34% higher than TransUNet model, and the HD distance is 25.24%, which is 11.67% lower. The average DSC coefficient on ACDC dataset can reach 90.30%, an increase of 1.67%.

1. 引言

医学图像是一类很特殊的图像。医学图像通常具有高分辨率、低对比度且包含复杂的解剖结构和组织，因此，在医学图像分割领域，对于模型的要求就有针对性[1]。自UNet架构提出后，其由对称的编码器和解码器组成，通过跳跃连接来传递信息，以改善分割性能。2017年，Transformer模型的提出，引入了自注意力机制，处理长距离依赖关系，并且具有更高的并行性[2]。

在医学图像分割领域，Transformer模型通常与传统的卷积神经网络(CNN)结合使用，以充分利用两种模型的优势[3]。TransUNet [4]模型首次将Transformer用于医学图像分割领域，它将Transformer的自注意力机制和U-Net的编码器–解码器结构相结合，综合两个模型的优势，有效捕获全局和局部特征，以在图像分割任务中实现更好的性能。

为充分利用TransUNet模型的优势，在其基础上，提出DouTransNet模型，编码器部分保留CNN与Transformer的串联结构，继承TransUNet模型同时捕获全局和局部特征的优点，针对单分支Transformer模块学习角度单一、容易丢失细节特征的问题，将Transformer设计成双分支并行结构来提取不同尺度的特征，其中两个分支的区别在于输入的patch大小不同，将两个分支的编码特征进行融合，可以保证Transformer模块在提取全局上下文信息时也避免丢失细节特征，实现特征互补；其中虽然Transformer模块中添加了位置编码，但仍存在只有位置的相对信息没有语义信息的局限性，DouTransNet在原有位置编码的基础上，单独添加卷积操作，作为位置编码的一部分，解决没有考虑语义信息的问题；由于两个分支可能会提取到相似或重叠的信息，在融合时可能存在冗余信息和噪声，受CE-Net [5]模型的启发，添加并行池化模块，将Transformer的融合特征进行平均池化操作，抑制噪声，保留多尺度特征的同时去除冗余信息；解码器中设计了多尺度融合模块(USF)，融合来自编码器的三个尺度信息，分别为CNN层提取的浅层特征、Transformer的融合特征及解码器中上采样后的特征，通过捕捉不同层次特征之间的相关性，弥补编码器与解码器之间的信息差距和上采样过程中的信息丢失。

通过在Synapse和ACDC数据集上进行一系列对比实验，所得结果证实了DouTransNet模型的高效性。此外，还通过实验考察了双分支结构中不同patch大小组合对模型性能的影响。

2. 模型结构

如图1所示为DouTransNet模型的整体架构。2.1部分详细介绍了整体架构的实现方法；2.2部分详细介绍了编码器的实现；2.3详细介绍了解码器的实现。

Figure 1. Overall architecture diagram of the model

图1. 模型整体架构图

2.1. DouTransNet整体架构

DouTransNet模型由CNN-Transformer模块组成编码器，上采样联合USF模块作为解码器。其中Transformer块由并行结构组成以此来提取多尺度特征，一张$H\times W\times C$ 的图像经过三层卷积后图像的尺寸会变小，同时通道数增加，根据两个Transformer分支所要求的patch大小，将图像分成多个patches，两个分支的patch大小不同，这样可以利用小patch来提取细节特征，两个分支加上各自的位置编码，再将其输入Transformer块中利用注意力机制对图像进行编码。两个分支独立编码。

为了更完整的保留不同尺度的特征，将两个分支输出的特征进行相加融合。之后将融合特征输入并行多核池化模块，输出的特征需要与池化前的特征进行拼接再经过1 × 1卷积恢复原来通道数，至此得到完整的编码特征。将编码后的特征进行reshape操作恢复三维形式，输入解码器。

解码时为了增加不同层次特征之间的依赖关系，同时为了避免上采样过程中信息的丢失，USF模块将浅层特征、多尺度融合特征、解码时的上采样特征进行拼接融合，其中多尺度融合特征为池化前的特征，是为了保证特征的完整性，可以对浅层特征及深层特征进行补充，弥合信息差距。解码器共由三个模块组成，第一个模块不需要融合多尺度融合特征，将第一个模块输出的结果通过卷积和上采样输入USF-1，第二个模块输出的结果通过卷积和上采样输入USF-2，通过USF-1、USF-2来融合三个不同层次特征。最后得到分割结果。

2.2. 编码器

2.2.1. CNN-Transformer编码

将一张图像$x_i\in {ℝ}^{H\times W\times C}$ 输入模型，首先经过三个CNN层图像尺寸变为H/8 × W/8 × 256，其中CNN层通过Resnet-50块[6]来实现，将CNN层得到的中间特征记为$\left\{f_i:H_i\times W_i\times C_i|i=1,2,3\right\}$ ；之后将卷积得到的特征图扁平化为二维patches，patches也就是输入transformer模块的特征长度，第一个分支为$\left\{x_i^k\in {ℝ}^{P_1\times P_1\times C}|k=1,\cdots ,N_1\right\}$ ，其中每个patch的大小为$P_1\times P_1$ ，共$N_1=HW/P_1^2$ 个patch，第二个分支为$\left\{x_j^k\in {ℝ}^{P_2\times P_2\times C}|k=1,\cdots ,N_2\right\}$ ，其中每个patch的大小为$P_2\times P_2$ ，共$N_2=HW/P_2^2$ 个patch ($P_1$ 和$P_2$ 不相同)。同时为了保留原图像的空间位置信息，将位置信息一同嵌入，得到图像特征的二维嵌入序列：

${\partial }^0=\left[x_i^1{F}_i;x_i^2{F}_i;\cdots ;x_i^{N_1}{F}_i\right]+F_{pos1}$ , (1)

${\partial }^1=\left[x_j^1{F}_j;x_j^2{F}_j;\cdots ;x_j^{N_2}{F}_j\right]+F_{pos2}$ , (2)

其中，$F_{i,j}\in {ℝ}^{\left(P_{1,2}^2\times C\right)\times D}$ 表示每个patch在D维空间的映射，$F_{pos1,2}\in {ℝ}^{N_{1,2}\times D}$ 为位置信息。

两个分支的Transformer编码器分别由六个模块组成，六个模块之间串联，两个分支之间并联，每个注意力模块的结构如图2，由多头自注意力(MSA)和多层感知机(MLP)块组成，所以第层的输出为：

${\partial }_l^{°}=\text{MSA}\left(\text{LN}\left({\partial }_{l-1}\right)\right)+{\partial }_{l-1}$ , (3)

Figure 2. Attention module structure diagram

图2. 注意力模块结构图

${\partial }_l=\text{MLP}\left(\text{LN}\left({\partial }_l^{°}\right)\right)+{\partial }_l^{°}$ , (4)

其中，$\text{LN}(·)$ 表示层归一化，${\partial }_l$ 为特征的编码表示。

2.2.2. 位置编码

由于Transformer模型原有的位置编码只提供了关于输入序列中位置的相对信息，使得模型在处理输入序列时可能会忽略一些重要的语义信息。卷积操作可以在位置编码中进行特征提取，从而使得位置编码不仅仅包含位置信息，还可能包含一定程度上的语义信息。有助于模型更好地理解位置与语义之间的关系。

因此，通过二维空间卷积[7]来添加关于位置编码的语义信息。将得到的位置编码进行卷积操作得到新的位置编码，加到特征图序列中。(1)和(2)式中的$F_{pos1}$ 和$F_{pos2}$ 即为新的位置编码。

2.2.3. 并行多核池化模块

将Transformer模块输出的融合特征进行池化操作，具体模块结构如图3，由四个不同接受域大小的并行池化操作组成，分别为2 × 2 pooling、3 × 3 pooling、4 × 4 pooling及5 × 5 pooling，具体为自适应平均池化，四个分支输出不同大小的特征图，以此来去除冗余信息同时保留多尺度信息；为了不增加模型的复杂度，将pooling后的特征图进行1 × 1卷积及上采样恢复pooling前模型的尺寸，每个分支的通道数为1，最后将四个分支的输出与pooling前的特征进行拼接融合，再经过一个1 × 1卷积恢复pooling前特征的通道数输出。

Figure 3. Parallel multi-core pooling module structure diagram

图3. 并行多核池化模块结构图

2.3. 解码器

将经过池化后的特征${\partial }_p$ 输入解码器，解码器由三个上采样模块组成，在上采样的过程中第一个模块不需要融合池化前的特征，由此可得经过第一个模块后的特征：

$m_1=\text{Conv}\left(\text{Concat}\left({\partial }_p,f_3\right)\right)\times r$ , (5)

其中$\text{Concat}(·)$ 表示对两个特征图进行拼接操作，$\text{Conv}(·)$ 表示对其进行卷积操作，$f_3$ 卷积层特征，r为上采样的放大倍数。

将$m_1$ 输入USF-1进行下一步解码。

2.3.1. USF-1

详细结构如图4所示。首先Transformer输出的融合特征$\partial$ 需要经过一次上采样图像尺寸变为(H/4, W/4)，解码器第一个模块的输出$m_1$ 图像尺寸也为(H/4, W/4)，将这两部分相加融合，再与CNN第二层的特征进行拼接，经过一次3 × 3卷积及上采样，输出$m_2$ ，具体如下：

$m_2=\text{Conv}\left(\text{Concat}\left(m_1+\partial \times r,f_2\right)\right)\times r$ , (6)

其图像尺寸为(H/2, W/2)；

Figure 4. USF-1 structure

图4. USF-1结构图

2.3.2. USF-2

详细结构如图5所示。将Transformer输出的融合特征经过两次上采样图像尺寸变为(H/2, W/2)，与USF-1模块的输出$m_2$ 相加融合，再与CNN第一层的特征进行拼接，经过一次3 × 3卷积及上采样，输出$m_3$ 。具体如下：

$m_3=\text{Conv}\left(\text{Concat}\left(m_2+\partial \times r\times r,f_1\right)\right)\times r$ , (7)

其图像尺寸为(H, W)。

2.4. 损失函数

本模型的损失是通过计算交叉熵损失[8] (CE Loss)和Dice损失[9] (Dice Loss)并将它们加权求和得到的。

对于每个像素点，交叉熵损失函数计算模型预测的概率分布与真实标签概率分布之间的差异，公式如下：

$L_{ce}=-\frac{1}{N}{\displaystyle {\sum }_{i=1}^N\left(y_i\ln \left(p_i\right)+\left(1-y_i\right)\ln \left(1-p_i\right)\right)}$ , (8)

Dice损失函数过计算预测分割结果与真实分割结果之间的相似度来度量模型的性能，公式如下：

$L_{dice}=1-\frac{2\times {\displaystyle {\sum }_{i=1}^N{p}_i{y}_i}}{{\displaystyle {\sum }_{i=1}^N{p}_i+{\displaystyle {\sum }_{i=1}^N{y}_i}}}$ , (9)

Figure 5. USF-2 structure

图5. USF-2结构图

其中，$y_i$ 是第i个像素的真实标签，$p_i$ 是第i个像素的预测值，N为图像中的像素点数。

为了结合两个损失函数的优势，选择0.5的权重来得到最终损失，公式如下：

$L=0.5\times L_{ce}+0.5\times L_{dice}$ , (10)

3. 实验

3.1. 实验准备

3.1.1. 数据集

为了充分证明本模型的有效性，在两个数据集上进行了实验：腹部多器官分割(Synapse数据集) [10]、自动心脏诊断挑战(ACDC数据集) [11]。

Synapse数据集由30名受试者的腹部CT扫描组成，共有3779张腹部轴向增强CT图像。按照[6]的拆分方法，随机使用18例(2212个轴向切片)来训练模型和12例(1567个轴向切片)来验证。在八个器官上使用平均DSC和平均豪斯多夫距离(HD)来显示分割效果：主动脉(Aorta)、胆囊(Gallbladder)、脾脏(Spleen)、左肾(Kidney (L))、右肾(Kidney (R))、肝(Liver)、胰腺(Pancreas)、胃(Stomach)。ACDC数据集共包含100名患者的心脏MRI图像，每例图像都手动标注了左心室(LV)、右心室(RV)和心肌(Myo)的真实标签。其中70例(1930个轴向切片)训练模型，10例验证和20例测试，在LV、RV和Myo上使用平均DSC来评价分割效果。

3.1.2. 评价指标

本文实验主要使用平均Dice相似系数(DSC)和平均豪斯多夫距离(HD)这两个指标来衡量模型的性能。Dice系数于衡量预测分割结果与真实分割结果的相似度，平均Dice系数计算的是所有类别Dice系数的平均值，其公式为：

$\text{Dice}=\frac{1}{C}{\displaystyle {\sum }_{i=1}^C\frac{2\times \left|X_i\cap Y_i\right|}{\left|X_i\right|+\left|Y_i\right|}}$ , (11)

其中，C是类别总数，$X_i$ 表示第i个类别的预测分割结果，$Y_i$ 表示第i个类别的真实分割结果，$\left|X_i\cap Y_i\right|$ 表示预测结果与真实结果之间交集的元素数量，$\left|X_i\right|$ 和$\left|Y_i\right|$ 分别表示预测结果和真实结果各自的元素数量。

平均HD系数，可以用来衡量预测分割边界与真实分割边界之间的差异，其公式为：

$H\left(A,B\right)=\max \left(h\left(A,B\right),h\left(B,A\right)\right)$ , (12)

其中，$h\left(A,B\right)$ 定义为集合A中每个点到集合B中最近点的距离的最大值，$h\left(B,A\right)$ 为集合B中每个点到集合A中最近点的距离的最大值，两者的最大值为豪斯多夫距离；

在计算豪斯多夫距离时，可以使用以下公式来计算两个点集之间的距离：

$d\left(p,q\right)=\sqrt{{\displaystyle \sum {\left(p_i-q_i\right)}^2}}$ , (13)

其中，p为模型的预测值，q为真实值，i为像素点的编号。

3.2. 实验细节及比较

3.2.1. 实验细节

本实验模型由PyTorch1.7.1框架搭建，根据文献[12]使用以上所有数据集时，都对三维体积进行了切片处理，并将预测的二维切片堆叠在一起重建成三维进行评估。对于所有的实验，应用了随机旋转、翻转等简单的数据扩充。输入模型的图像尺寸为224 × 224，两个Transformer分支默认的输入patch大小分别为16和14。使用SGD优化器，动量为0.9，学习率为0.01，权重衰减为1e−4的参数来训练模型，patch size设置为24。使用在ImageNet上进行预训练的ResNet-50和ViT模型来初始化CNN层和Transformer块。Synapse数据集的训练迭代次数为30 k，ACDC数据集为20 k。所有实验均使用有11 GB内存的单个NVIDIA 2080 GPU来训练。

3.2.2. 比较

Synapse数据集表1列出了部分模型在Synapse数据集上对八个器官进行分割的DSC系数和HD

Table 1. Average DSC (%) coefficients and average HD distances for different models on the Synapse dataset

表1. 不同模型在Synapse数据集上的平均DSC (%)系数和平均HD距离

模型	DSC↑ HD↓	主动脉	胆囊	左肾	右肾	肝	胰腺	脾脏	胃
V-Net [10]	68.81 -	75.34	51.87	77.10	80.75	87.84	40.05	80.56	56.98
R50 UNet [13]	74.68 36.87	87.74	63.66	80.60	78.19	93.74	56.90	85.87	74.16
UNet	76.85 39.70	89.07	69.72	77.77	68.60	93.43	53.98	86.67	75.58
R50 Att-UNet [14]	75.57 36.97	55.92	63.91	79.20	72.71	93.56	49.37	87.19	74.95
Att-UNet	77.77 36.02	89.55	68.88	77.98	71.11	93.57	58.04	87.30	75.75
R50 ViT [15]	71.29 32.87	73.73	55.13	75.80	72.20	91.51	45.99	81.99	73.95
TransUnet	75.86 36.91	85.96	59.96	80.73	74.02	93.25	57.71	84.77	70.50
SwinUnet [6]	79.13 21.55	85.47	66.53	83.28	79.61	94.29	56.58	90.66	76.60
DouTransNet	79.20 25.24	87.74	62.51	81.97	78.50	93.83	64.36	86.63	78.02

Table 2. Average DSC of different models on ACDC dataset (%)

表2. 不同模型在ACDC数据集上的平均DSC (%)

模型	平均DSC↑	右心室(RV)	心肌(Myo)	左心室(LV)
R50-U-Net	87.55	87.10	80.63	94.92
R50-Att-UNet	86.75	87.58	79.20	93.47
TransUNet	88.63	87.56	84.54	94.43
MISSFormer [16]	87.90	86.36	85.75	91.59
Swin-Unet	90.00	88.55	85.62	95.83
UNETR [17]	86.61	85.29	86.52	94.02
DouTransNet	90.30	89.02	85.66	96.24

距离及总体的平均DSC和平均HD。重新训练TransUNet模型得到的平均DSC系数和HD距离分别为75.86%和36.91%。DouTransNet模型相较于UNet模型DSC系数提高了2.35%，HD下降了14.46%。相较于TransUnet模型DSC系数提高了3.34%，HD下降了11.67%，尤其是对于胃和胰腺两个器官的分割，DSC系数分别提高了7.52%和6.65%。

ACDC数据集表2列出了部分模型在ACDC数据集上对RV、Myo、LV分割的DSC系数。重新训练的TransUNet模型的平均DSC系数为88.63%。DouTransNet模型的DSC系数可以达到90.30%，相较于TransUNet模型提升了1.67%。相较于UNETR模型提升了3.69%，尤其是对RV和LV这两部分的分割，分别提升了3.73%和2.22%。

图6为不同模型在Synapse数据集上的分割结果可视化，从分割效果来看本文提出的DouTransNet

Figure 6. Segmentation results of different models on the Synapse dataset

图6. 不同模型在Synapse数据集上的分割结果

模型得到了更好的分割结果。从图6(c)可以看出Swin-UNet模型由于对局部特征关注较少导致了分割边缘非常不平滑，会出现各分割器官之间的分割错误的问题；图6(d)和图6(e)为TransUNet和UNet模型的分割实例，由于对细节特征的不敏感，会导致分割过度的问题。DouTransNet模型增加对细节特征的关注，可缓解上述问题。

3.3. 双分支不同patch组合对比

为了弥补串行的Transformer块提取特征尺度单一、学习角度单一的问题，在本文模型中将其设计成了双分支并行的结构，根据输入两个分支的patch大小来融合不同尺度的特征，但是patch大小的不同组合形式同样会影响模型的性能，如表3列出了部分不同patch大小组合模型的性能，使用的是Synapse数据集。通常情况下较小的patch大小有更好的分割效果[6]，如下表所示当patch1都为16时，patch2为8时DSC系数更高，但是一味减小patch的大小反而会忽略特征之间的相关性使模型性能下降，如表所示当patch1为14，patch2为8时，模型的DSC系数只有77.88%，所以在设定patch的大小时也要根据图像的特点及需要提取的特征来选择合适的patch。

Table 3. Average DSC (%) coefficient of different patch size combinations

表3. 不同patch大小组合的平均DSC (%)系数

patch1 patch2	DSC↑
16 14	79.20
16 8	79.46
14 8	77.88

3.4. 消融实验

为验证所提模块的有效性，对双分支结构，并行池化模块和USF模块进行了消融实验，如表4展示了双分支与单分支、池化与无池化、USF模块和级联上采样之间的消融实验。其中级联上采样为TransUNet模型中的解码器部分。使用Synapse数据集来进行实验。

可以看出，在只使用单分支时，由于没有双分支提取多尺度的特征，采用池化反而会使得信息丢失严重，模型的DSC系数只有77.88%；无池化模块时由于存在冗余信息，模型DSC系数为78.73%，提升效果不明显，无USF模块时模型DSC系数为78.34%，同样提升效果不明显。由此可证明模块的有效性。

Table 4. Ablation experiments of different modules

表4. 不同模块的消融实验

模型结构	DSC (%)
TransUNet+ 池化+ USF	77.88
TransUNet+ 双分支+ USF	78.73
TransUNet+ 双分支+ 池化	78.34
TransUNet+ 双分支+ 池化+ USF	79.20

4. 结论

本文在TransUNet模型的基础上，提出了DouTransNet模型，将串行的Transformer块设计成双分支并行结构，以此来提取不同尺度的特征，避免丢失细节信息，使得模型可以从不同角度理解图像，提高提取特征的多样性；添加并行池化操作，有效捕捉到图像不同尺寸的上下文信息，有效抑制噪声，去除冗余信息；解码器使用USF模块，融合来自编码器的三个不同层次的特征，有效缓解了编码器与解码器之间的特征差异，提高分割性能。实验结果表明，本文模型在不同数据集上相较于一些模型都有了明显的提升，证明了模型的有效性。