关联自注意力机制改进的UNet用于肺结节图像分割

Associative Self-Attention Mechanism Improved UNet for Lung Nodule Image Segmentation

摘要: 肺癌在各项肺部疾病中是非常致命的，而早期的诊断和治疗则能够有效的降低死亡率，对于早期的预防以及诊断，肺结节分割是一个非常重要的步骤，但结节与其领域之间的相似性难以处理。在这里本文提出了一个端到端的基于UNet的分割框架，称为PR-Net。为了保证标准卷积的可接受精度以及深层结构带来的信息损失，本文构建了残差结构的变体作为网络的主干。此外在解码过程中由于信息的冗余和语义差距导致难以有效的融合病变区域的空间及语义信息，为此本文采用了Transformer中的注意力机制，通过该机制的多头自注意力来有效提升空间细节和语义定位级别的特征识别。最终的模型IOU指标达到了97.71%，Dice相似系数达到了98.84%。大量的实验和结果表明，本文的模型在肺结节分割上有着非常不错的性能及稳定性。

Abstract: Lung cancer is very fatal among all lung diseases, while early diagnosis and treatment can effective-ly reduce the mortality rate. For early prevention as well as diagnosis, lung nodule segmentation is a very important step, but the similarity between the nodule and its domain is difficult to handle. Here in this paper we propose an end-to-end UNet-based segmentation framework called PR-Net. To ensure acceptable accuracy of standard convolution as well as loss of information due to deep structure, a variant of residual structure is constructed as the backbone of the network. In addition, in the decoding process, the redundancy of information and semantic gaps lead to the difficulty of effectively fusing the spatial and semantic information of the lesion region, for this reason, this pa-per adopts the attention mechanism in Transformer, and effectively improves the feature recogni-tion at the spatial detail and semantic localization level through the multi-head self-attention of this mechanism. The final model IOU index reaches 97.71% and Dice similarity coefficient reaches 98.84%. A large number of experiments and results show that the model in this paper has very good performance and stability in lung nodule segmentation.

1. 引言

肺部CT图像分割是疾病检测等许多应用中的关键过程。由于肺部结构中存在相似的图像密度、不同类型的扫描仪和扫描协议，这被认为是一个具有挑战性的问题 [1] 。在各种肺部疾病中，肺癌是最致命的，根据世界卫生组织(WHO)的数据，肺癌每年造成全球130万人死亡。早期治疗可降低死亡率，对于预防和诊断病变，肺结节分割 [2] 是一个重要步骤。诊断的第一部分是计算机辅助诊断(CAD)系统，该系统改进了结节的检测，有助于将结节分为恶性或良性。然而，随着时间的推移，数据不断扩大，CT图像的数量也在增加。随着图像数量的增加，采用手动肺结节分割程序变得具有挑战性。考虑到这一问题，一种高效的肺结节自动分割算法在临床上具有重要的意义。

最近，基于深度学习的框架已应用于许多医学图像处理领域，包括医学图像分割、病变检测和分类。根据这些成功的经验，已经提出了许多用于CT图像分割的深度学习系统 [3] [4] [5] [6] 。众所周知的UNet [7] 及其变体，包括UNet++ [8] 、V-Net [9] ，已被作为一些早期但有效的尝试，用于从CT扫描中分割肺部 [10] 、肺部阴影 [11] 和病变区域 [12] 。继它们工作之后Maqsood等人提出了一种端到端基于U-Net的分割框架(DA-Net)用于有效的肺结节分割，该方法通过集成与Atrous卷积块合并的紧密和密集链接的丰富卷积块来提取丰富的特征，以在不丢失和覆盖数据的情况下拓宽滤波器的视野 [13] 。Luo等人针对肺结节表面详细信息学习不足以及模型参数和计算过多的问题，提出了一种3D双注意力阴影网络(DAS-Net)来解决肺结节的准确分割问题 [14] 。提出了深度学习和形状驱动水平集的协同组合，以实现自动化和准确的肺结节分割 [15] 。

2. PR-Net分割模型

2.1. 模型结构

提议的PR-Net模型结构主要由三个部分构成，分别是编码器部分、解码器部分以及特征增强部分。详细的模型结构如图1所示。编码器部分由改进的残差块和平均池化下采样构成，通过不断地压缩原始图像从而逐步提取富含高级语义信息的特征图，但随着网络层数的加深，可能会出现梯度消失甚至梯度爆炸的问题，残差块不仅能有限的解决这类问题还能缓解深层网络带来的原始信息丢失问题。解码器部分通过反卷积进行上采样，逐步恢复图像特征，并通过跳跃连接方式将浅层的位置信息与深层的语义信息相结合，结合的方式采用通道拼接，这样的结合方式能够保留更多的信息，此外，融合后的特征由双层卷积结构(即由两个3 × 3卷积，批归一化和ReLU激活函数堆叠而成)进一步提取特征，并恢复特征通道数，减少后续计算参数。特征增强部分由双层卷积结构和渐进注意模块构成，由双层卷积结构对来自编码器的特征进行进一步提取，后由渐进注意模块进行增强，该模块不断能够获取更大的感受野，而且能够加强特征之间的长距离依赖关系，即加强序列不同位置之间的联系，以方便模型能够正确处理序列中的信息。

Figure 1. General structure of the PR-Net model

图1. PR-Net模型总体结构

2.2. 改进的残差块

传统卷积神经网络在前向传播的时候，会将输入信号经过一系列的层次变换，这些变换通常表示为非线性映射(例如ReLU激活函数)和线性变换(例如卷积或全连接层)。而网络的深度增加时，梯度在反向传播过程中可能会逐渐消失，导致难以训练的问题。残差结构则旨在解决模型训练过程中的梯度消失和梯度爆炸等问题。通过引入残差连接(或者称之为跳跃连接skip connection)，这种结构使得信息能够直接从输入层传递到网络的较深层次，而不完全依赖于层次之间的变换。这样做的关键是网络学习到的是残差(差值)，而不是直接学习输出。从而可以更容易捕获到输入和输出之间的细微输出异。

本文改进的残差块结构如图2所示，与传统双层卷积结构堆叠的残差块不同的是该模块由三组Conv2d、InsatnceNorm2d、ReLU激活函数堆叠以及残差连接构成，除开首个卷积采用1 × 1的卷积核外，其他两个卷积的尺寸均为3 × 3，这样可以通过1 × 1的卷积来自适应特征通道并引入非线性变换，且与直接采用输入作为残差连接的残差结构不同的是该结构将1 × 1卷积后的特征作为残差连接的主体，这样更加有益于全局信息的聚合。在这里没用采用传统残差块中的归一化方式即Batch Normalization (批归一化，简称BN)，而是采用了Instance Normalization (实例归一化，简称IN)，与BN方法不同的是该方法以单个样本的实例为单位进行归一化，而不是以批次为单位。该方法主要有以下几个作用，一个是可以减少内部协变量偏移：这是一个训练网络时常见的问题。该方法有助于确保网络每一层的输入分布都在归一化范围内，从而加速模型的训练。另一个是增加模型的稳定性：当批次大小较小或输入分布差异较大时。它可以使每个样本都在相似的范围内归一化，减少了网络对批次中每个样本的依赖性。

Figure 2. Specific structure of the residual block

图2. 残差块的具体结构

2.3. 渐进注意模块

注意力机制(Attention Mechanism)是一种深度学习中的重要技术，主要用于处理序列数据(图像也属于这一类)和建立更复杂的模型。注意力机制的核心思想是模拟人类在处理信息时的注意力分配过程，使模型能够根据输入的不同部分来动态地分配不同的权重和关注度。这其中注意力机制(Self-Attention Mechanism)，也被称为自注意力模型或多头自注意力，该注意力机制可以捕获序列中所有元素之间的关联关系，而不受固定窗口大小或局部依赖的限制。这有助于模型捕获长距离依赖关系，从而在处理序列任务时表现出色。此外自注意力机制可以学习到不同元素之间的权重，以适应不同任务的需求。这增加了模型的表现力，使其能够更好地适应不同数据分布和模式。

基于自注意力机制的思想，本文构建了一个渐进注意力模块(Progressive attention module，简称为PAM)，该模块主要由两部分组成，即扩张自注意力模块(Dilated Self Attention，简称为DSA)和门控卷积模块(Gated Convolution，简称为GC)。如图3所示，使用一个3 × 3卷积运算和两个扩张的自注意力模块，扩张的自注意力模块的扩张率分别为2和3，以获得来自不同感受野的特征。然后，将通过3 × 3卷积运算和扩张率为2的扩张自注意力模块获得的特征馈送到GC中，并使较大的感受野特征指导原始特征的判别提取过程。因此，来自第一GC模块的特征和来自扩张率为3的扩张自注意卷积模块的特征被再次馈送到GC中，并且判别特征被进一步提取。最后，将原始的3 × 3卷积特征与两个GC的输出特征相结合，作为最终输出。

DSA建立在Transformer的多头自注意力基础上，允许模型只关心来自全局表示子空间的信息。使用卷积嵌入而不是线性嵌入，这样DSA不仅可以聚合全局上下文信息，还可以考虑局部空间信息。与传统卷积相比，扩张卷积可以通过改变扩张率灵活地改变感受野，同时确保特征大小的一致性。DSA可以选择性地将全局上下文聚合到学习的特征中，并使用卷积嵌入和矩阵乘法将更广泛的上下文位置信息编码到局部特征中，这可以提高类内紧凑性并优化特征表示。

DSA的细节如图4所示。首先，对经过双层卷积结构的编码器特征x应用三个扩张卷积运算来生成特征图q、k和v。随后，对q和k进行整形，并使用softmax归一化进行矩阵乘法，得到位置相关性注意力图。上述操作可以定义如下：

$\begin{array}{l}M=reshape\left(q\right)\\ N=reshape\left(k\right)\\ T=reshape\left(v\right)\\ B_{i,j}=\frac{\exp \left(M_i\cdot N_j\right)}{{\displaystyle {\sum }_{k=1}^n\exp \left(M_k\cdot N_j\right)}}\\ \text{DSA}\left(B,T\right)=B\cdot T\end{array}$ (1)

其中， $B_{i,j}$ 测量第i个位置对第j个位置的影响， $n=h\times w$ 是像素数，M、N和T表示重塑的特征。B表示位置相关性注意力图。然后，将T乘以B，对优化后的特征图进行整形，以获得DSA的输出。

Figure 3. General structure of the progressive attention module

图3. 渐进注意模块的总体结构

Figure 4. Specific structure of DSA

图4. 扩张自注意力的具体结构

门控卷积模块(GC)由两个输入组成，指示一个大的和一个小的感受野特征，如图3所示。然后，将两种不同的卷积运算应用于输入特征以生成门映射。最后，执行乘法运算以获得最终输出。计算过程可以公式化如下：

$\begin{array}{l}Gate=W_g\cdot F_{high}\\ F=W_f\cdot F_{low}\\ G=\delta \left(F\right)\ast \sigma \left(Gate\right)\end{array}$ (2)

其中 $W_g$ 和 $W_f$ 是不同卷积投影的嵌入矩阵， $F_{high}$ 和 $F_{low}$ 表示两个输入。Gate是注意力映射， $\sigma$ 是sigmoid函数，它将所有值映射到0和1之间的区间。最后，F是特征嵌入， $\delta$ 表示ReLU激活。

3. 实验和结果

3.1. 数据集与评估指标

为了评估模型的有效性，本文采用一个公开的肺结节分割数据集，该分割数据集共包含264张带有分割标签的肺结节CT图像。这264张图像随后被分为211张用于训练，26张用于验证，27张用于测试。

为了进一步评估提出的模型的功效，使用以下指标对最终的分割结果进行了定量分析：交并比(IOU)、骰子相似系数(DSC)、豪斯多夫距离(HD)、准确率(ACC)、曲线下面积(AUC)、特异性(SPE)和平均绝对误差(MAE)。IOU通过比较模型的预测区域和真实区域之间的交集和并集来度量重叠程度，IOU的值越接近1，表示预测结果与真实结果的重叠越好，分割效果越准确。DSC通过计算预测图和地面实况图之间的重叠率来衡量它们之间的相似性。HD用于确定两个集合之间的距离，值越小表示两个集合之间的相似度越高。ACC是正确标注像素与像素总数的比率，准确率越高，表示分割性能越好。SPE衡量正确识别非目标样本的概率，特异性越高表示模型性能越好。MAE用于计算预测图像与地面实况之间的误差，而AUC则反映了模型正确预测阳性或阴性样本的概率，数值越接近1表示模型性能越好。上述指标的公式表示如下：

$\text{H <msup> D <mo> ′ </mo> </msup>}\left(S_a,S_b\right)=\underset{i\in S_a}{\max }\left(\underset{j\in S_b}{\min }{\left|\left|i-j\right|\right|}_2\right)$ (3)

$\text{HD}\left(S_a,S_b\right)=\max \left(\text{H <msup> D <mo> ′ </mo> </msup>}\left(S_a,S_b\right),\text{H <msup> D <mo> ′ </mo> </msup>}\left(S_b,S_a\right)\right)$ (4)

其中，S_a和S_b分别代表模型预测结果和地面真实分割的目标表面点集合。

$\text{DSC}=\frac{2TP}{2TP+FP+FN}$ (5)

$\text{Acc}=\frac{TP+TN}{TP+FP+FN+TN}$ (6)

$\text{SPE}=\frac{TN}{FP+TN}$ (7)

其中，TP, TN分别代表正确分类并标记为真阳性和真阴性的像素数量。相反，FP, FN分别代表被错误分类并标记为假阳性和假阴性的像素数量。

$\text{MAE}=\frac{1}{w\times h}{\displaystyle \underset{x=1}{\overset{w}{\sum }}{\displaystyle \underset{y=1}{\overset{h}{\sum }}\left|P\left(x,y\right)-G\left(x,y\right)\right|}}$ (8)

其中h和w分别为输入图像的高度和宽度，(x, y)表示预测图和地面实况中每个像素的坐标。

3.2. 与其他模型的对比实验

为了证明提出的模型的有效性，本文将其与UNet [7] 、UNet++ [8] 、Attention-UNet [16] 等几种经典分割模型以及AnamNet [17] 、LightUNet [18] 、Inf-Net [19] 和AwsNet [20] 等其他四种前沿分割方法进行了比较。所有模型都在相同的配置下进行评估，使用相同的参数设置重新训练，并在相同的数据集上进行测试。

表1显示了提议的方法与其他方法的比较(表中红色的标注表示最好，蓝色的标注表示次好)，可以看出提出的PR-Net在所对比的指标上都取得了最佳结果。与一些经典方法(UNet [7] 、UNet++ [8] 、Attention-UNet [16] )相比，更先进的方法表现出更强的竞争力。作为医疗分割领域常用的网络，UNet [7] 在IOU上取得了97.16%的结果，而与UNet [7] 相比，Attention-UNet [16] 在IOU指标上提高了0.31，这无疑得益于其采用的注意力机制有助于网络更好地捕获重要的图像区域，而本文的方法达到了0.55。UNet++ [8] 作为UNet [7] 的改进，引入了多尺度路径和密集连接，其在DSC指标也是上取得了98.56%的结果，相较之下，Inf-Net [19] 采用信息传播单元和多尺度注意力模块，有着和其接近的效果。AwsNet [20] 采用由粗到细的逐步分割模型，并结合像素关注策略与自动加权监督，其在各项指标上也展示出了优秀的性能。LightUNet [18] 通过在U-net架构中引入群卷积，使得其在大多指标上取得了第二的好成绩。总体而言，在肺结节分割方面，本文提出的PR-Net与表现出了卓越的性能，在定量结果方面优于其他方法。

Table 1. Comparison with different models

表1. 和不同模型的比较

4. 结论

本文聚焦于肺结节胸部CT的病变区域分割任务，并为此提出了端到端的分割学习框架，称为残差渐进注意力网络(PR-Net)，该框架遵循编解码器结构，通过改进的残差块提取丰富的语义特征并减少丢失信息和梯度爆炸的可能。渐进注意力模块(PAM)旨在从重要的多尺度信息和长距离依赖关系中选择更有意义的编码特征，并将扩张卷积与自注意力相结合，使得该模块能够兼具不同感受野下的上下文特征。最后，本文将该模型与一些其他分割模型在肺结节的分割数据集上进行了实验验证，提议的模型在IOU和DSC指标上分别取得了97.71%和98.84%的结果，且在所有评价指标上优于其他对比网络，这些实验验证了本文提出的PR-Net模型在肺结节图像分割方面的优秀性能。

NOTES

^*通讯作者。