基于RT-Unet的医学图像分割算法

赵佳美，吴迪康，王志芳

(黑龙江大学 电子工程学院，黑龙江 哈尔滨 150080)

0 引言

随着计算机技术和人工智能的快速发展，深度学习在自然语言处理、机器视觉任务和多模态分析等各个方面都取得了突破性的进展[1-3]。深度学习能够自动且准确地提取特征并具有强大的建模能力，相较于传统算法在医学图像分割中得到了更广泛的应用。但由于不同患者间器官的结构、纹理和大小等方面存在差异，医学图像分割精度高度依赖于成像设备与医生的专业知识，快速且准确地完成分割任务仍然具有挑战性[4]。

目前，U-Net[5]是基于深度学习的医学图像分割算法的主流方法，由编码器与解码器构成的U型对称结构，通过跳过连接使用简单的拼接操作完成多尺度特征的融合，达到减少空间信息丢失的目的。这种U型结构使得U-Net在各种医学影像处理方面都取得了巨大的成功，许多算法都是在此基础上进行的改进，诸如U-Net++[6]，3D U-Net[7]和Res-UNet[8]等。但是由于卷积运算固有的局限性，图像上的对角信息很难直接进行卷积运算，无法有效地提取全局信息，使U-Net在远程关系建模方面存在一定限制。为了解决这个问题，此前方法大多是引入图像金字塔[9]、深度卷积层[10]或注意力机制[11]等，但会造成算法计算量的直线上升。

近两年，许多研究者将自然语言处理领域的Transformer引入到机器视觉领域[12]。2020年Dosovitskiy等[13]首次提出基于Transformer的ViT算法完成图像识别任务，取得了与卷积神经网络(Convolutional Neural Network，CNN)相当的性能，但是却需要较大的数据集来完成模型的训练。随着研究的深入，CNN与Transformer二者结合在图像处理领域上取得了较大的突破。2021年Chen等[14]开创性地提出了TransUNet，利用CNN提取细节信息，将其送入到Transformer中提取全局信息，通过跳过连接融合不同分辨率的多尺度特征。该算法在腹部多器官分割、心脏分割数据集中取得了优异的成绩，充分证明了CNN与Transformer组合的有效性，但却忽略了跨尺度特征之间的依赖性。2021年Cao等[15]提出了基于Transformer的Swin-Unet算法应用于医学图像分割任务，极大地减少了算法的计算量并取得了较高的分割准确率。但Swin-Unet在训练过程中只能学习单尺度的上下文特征，缺乏处理局部信息与多尺度信息的内在归纳偏置[16]。针对上述问题，本文在Swin-Unet的基础上提出了RT-Unet算法，可以充分地学习多尺度特征信息，提高算法的训练精度与分割准确率。

1 研究方法

1.1 整体结构

RT-Unet是一种由编码器、解码器和跳过连接组成的U型对称拓扑结构，在残差连接的过程中加入局部归纳偏置模块构建RESwin Transformer模块，协助完成特征信息的学习，通过叠加聚合不同感受野的卷积层来提取各个层间或层内的多尺度特征。然后，在位置嵌入与编码模块使用4层连续卷积层，并选择GELU作为激活函数。RT-Unet不但可以有效地学习局部与远程信息，而且可以获得更丰富的多尺度特征，增强网络的泛化能力，RT-Unet的总体结构如图1所示。

图1 RT-Unet的结构Fig.1 Structure of RT-Unet

首先，编码部分将医学影像数据输入图像分块处理层(Patch Partition)中，将其切分成大小相等的非重叠图像块，并对图像块进行位置嵌入与编码(Linear Embedding)，完成转换特征维度等操作；然后，将转换后的图像块输入RESwin Transformer模块进行特征学习，通过图像块合并层(Patch Merging)实现下采样功能。与之对称的解码部分，采用图像块扩展层(Patch Expanding)完成上采样操作，逐步恢复图像的大小。同时，为了减少空间信息的丢失，编-解码部分利用跳过连接完成多尺度上下文的特征融合；最后，将得到的特征信息应用线性投影层完成最终的分割预测。

1.2 RESwin Transformer 模块

图2 RESwin Transformer模块的结构Fig.2 Structure of RESwin Transformer block

RESwin Transformer模块加入LIB模块后，一方面可以在尺度不变的情况下，通过叠加聚合具有不同感受野的卷积层获取丰富的多尺度上下文特征。另一方面，RESwin Transformer模块自身的注意力机制在远程关系建模上具有一定优势，又通过模块中的卷积计算相邻像素的相关性获得边缘与角落的局部细节信息。因此，RESwin Transformer模块可以更有效地建模局部和远程依赖关系，较好地提升网络的鲁棒性与分割准确率。RESwin Transformer模块的计算过程总结如下：

(1)

(2)

(3)

(4)

LIB模块嵌入到RESwin Transformer模块的残差连接过程中,具体结构如图3所示。首先，Seq2Img层完成简单的重塑作用，将输入的一维序列转换为二维特征图像块。然后，依次将图像块输入感受野为1×1,3×3,1×1组成的3个卷积层(Conv)中提取局部多尺度特征。随后，每层再进行批量归一化(BN)和ReLU激活函数处理。最后，由Img2Seq层将具有多尺度特征的二维图像块转换回一维序列并入到MLP中，协助RESwin Transformer模块学习局部特征与远程依赖关系。

图3 LIB模块的结构Fig.3 Structure of the LIB block

1.3 位置嵌入与编码模块

位置嵌入与编码模块的主要作用是将图片转换为RESwin Transformer模块可以处理的一维序列，并对序列进行相对位置编码。与自然语言处理领域中的词汇不同，图像与图像之间是没有位置顺序的。因此，RT-Unet首先将一张医学影像划分为若干个不重叠的图像块，每个图像块充当一个词汇。然后，通过位置嵌入与编码模块将得到的图像块展平，再利用线性映射将二维图像转换为一维向量并对其逐一进行编码，使图像块可以嵌入位置信息。

位置嵌入与编码模块由4个连续的卷积层完成像素级空间信息的编码，每个卷积层后跟随一个GELU激活函数和归一化层，结构如图4所示。

图4 位置嵌入与编码模块的结构Fig.4 Structure of position embedding and encoding block

激活函数对网络上一层输出的神经元进行处理，将结果传递到下一神经元中，给神经元引入非线性因素，使网络可以应用到众多的非线性模型中。激活函数ReLU输入值为正时，网络可以正常收敛；但为负时，ReLU的输出值将始终为0，导致网络参数的更新出现停滞，进而出现梯度消失问题。而GELU激活函数在ReLU的基础上加入统计特性完成对输入神经元的概率描述。它是一种符合预期的随机正则变换方式，能够有效地避免梯度消失问题，增加算法的非线性。GELU激活函数的实验效果要优于ReLU，尤其是在Transformer模型中表现最好[17]。本文采用GELU激活函数，计算公式如下：

GELU(x)=xP(X≤x)=xφ(x)，

(5)

式中，φ(x)为正态分布的概率函数，当为标准正态分布时：

(6)

2 实验结果分析

2.1 评价指标与Synapse数据集

本文采用2个评价指标：平均Dice系数(DSC)和Hausdorff距离(HD)[18]。二者都是描述2个样本集合间的相似性度量，但DSC更加侧重于内部填充部分的分割准确性，而HD则对分割的边界更为敏感。它们的定义如下：

(7)

HD(A,B)=max(h(A,B),h(B,A))，

(8)

(9)

(10)

式中，A代表医学影像的标签图；B代表算法分割的预测图。式(10)为双向HD，式(11)和式(12)分别为从A到B和从B到A的单向HD。

本文在Synapse腹部多器官分割数据集上验证RT-Unet的分割效果。Synapse数据集共包含30张腹部扫描CT，每张CT里含有8个器官：主动脉、胆囊、脾脏、左肾、右肾、肝脏、胰腺和胃。为了避免由于数据集较小而导致过拟合现象的发生，将3D医学影像转换成大小为224 pixel×224 pixel的2D切片，共获得3 779张2D切片，其中2 212张用于算法的训练，剩余的完成算法的测试。

2.2 实验结果分析

RT-Unet在Python3.6，Pytorch1.5.0和NVIDIA 2080Ti GPU的基础上完成。损失函数为交叉熵损失函数，优化器选择随机梯度下降(SGD)。在相同的实验条件下，将RT-Unet在Synapse数据集上分割8个腹部器官得到的DSC，HD参数与经典的分割网络ViT，V-Net，U-Net，Swin-Unet，TU-Net进行对比分析，实验结果如表1所示。

表1 不同方法在Synapse数据集上的分割结果Tab.1 Segmentation results of different algorithms on Synapse dataset

由表1可以看出，与基础的ViT，V-Net相比，RT-Unet的DSC与HD参数都得到了较大的提升，证明CNN与Transformer组合的有效性。其次，与经典的医学图像分割网络U-Net相比，RT-Unet 的DSC，HD参数分别提升了2.23%与16.27 mm，说明RT-Unet对于较小器官与分割边界更为敏感。由于实验设备与参数设定等诸多因素，Swin-Unet与TU-Net的DSC，HD参数分别为：77.28%，77.95%和26.93，23.43 mm。RT-Unet较二者分别提升了1.8%，1.13%和3.5，1.27 mm，取得了DSC为79.08%，HD为23.43 mm较为准确的分割结果，表明RT-Unet同时注重了局部细节信息与远程依赖关系，又促进了多尺度上下文信息的提取，有效地提升了算法的分割精度。

本文主要对Swin-Unet，TU-Net，RT-Unet算法的分割结果进行可视化对比分析，结果如图5所示。由图5的第1行可以看出，RT-Unet分割的边界相较于其他算法更为清晰光滑；从第2行可以看出，Swin-Unet和TU-Net对于肝脏都存在漏分割的现象，并且TU-Net也并未分割出胰腺，相比之下RT-Unet更加注重细节信息的提取，对器官的分割也更为精准；从第3行可以看出Swin-Unet和TU-Net对于胃部都出现了过分割的现象，RT-Unet综合丰富的多尺度上下文特征，此类现象出现的概率较小。

图5 不同算法的分割结果Fig.5 Segmentation results of different algorithms

2.3 消融实验

通过分别移除算法中的局部归纳偏置模块或替换为原有的位置嵌入与编码模块来验证它们的有效性，并在Synapse数据集上完成了消融实验，实验结果如表2所示。从表2可以看出，当算法中仅有位置嵌入与编码模块作用时，取得了DSC为77.70%的效果，在不影响整个算法复杂度的情况下，较小地提升了分割的准确性。当算法中仅有局部归纳偏置模块作用时，取得了DSC为78.44%，提升了近1.2%，有效地促进了不同分辨率信息的融合。由此可证，二者对于算法性能的提升是不可缺少的，删除其中任何一个都可能导致性能的损失。

表2 基于位置嵌入与编码模块和LIB模块的消融实验Tab.2 Ablation studies on linear embedding block and LIB block

3 结束语

本文提出了一种基于RESwin Transformer的U型对称结构算法——RT-Unet。该算法在RESwin Transformer模块中引入了内在的LIB，由具有不同感受野的卷积层获取丰富的多尺度信息。卷积与Transformer二者组合，既可以关注局部细节信息，又可以有效地建模远程依赖关系。相较于其他算法，RT-Unet更加注重器官边界与小器官分割的准确性，较少会出现过分割与欠分割现象。而且，本文在位置嵌入与编码模块中，选择GELU激活函数增强算法的泛化能力，极大地避免了梯度消失问题的出现。在Synapse腹部多器官分割实验中，RT-Unet取得了DSC为79.08%，HD为23.43 mm的分割结果，证明了其有效性。