基于改进U-Net网络的神经元分割算法

程维东，叶 曦*，王 芳，平晶晶，钱同惠，张志玮

（江汉大学 a.智能制造学院；b.人工智能学院；c.设计学院，湖北 武汉 430056）

0 引言

图像分割在生物医学图像分析中有着重要作用，也是医学图像处理方法的重要组成之一。随着计算机硬件性能的提高，基于U-Net［1］的图像分割技术已经成为医学图像处理的有力工具，被广泛应用于各种医学图像分割任务中。目前，U-Net对视网膜、脑组织、肺部和血流图像等都有较好的分割效果，也在生物图像分析中起到了重要的作用。

神经元结构分割是目前生物医学图像分析的重要课题之一［2］。超高分辨率的电镜图像越来越多，但对电镜图像进行有效分割依然存在挑战。采用人工分割的方式虽然能获得很好的分割效果，但这一过程十分繁琐且耗时甚高［3］。2015年开始，U-Net在医学图像分割中的应用取得了突破性进展，从此获得了广泛的关注和应用，进而基于深度学习的分割算法成为医学图像分割算法的主流方法［1，4］。因此，建立合适的深度学习模型对图像分割进行精确、快速、有效的神经元结构分割，对于生物医学图像分析有着重大意义。

深度学习图像分割模型从2012年开始，Hinton和学生Krizhevsky设计的AlexNet在Image Net比赛上取得了突破性的识别准确率后［5］，卷积神经网络（convolutional neural network，CNN）开始成为图像处理的主要工具。2015年，网络结构经历了从AlexNet到Oxford Visual Geometry Group提出的VGG模型［6］的转变。该模型的亮点在于用3*3、1*1的卷积核替代了5*5卷积，在相同的感受野下，3*3的卷积参数更少。之后，U-Net在MICCAI 2016被扩展为3D U-Net应用于三维图像分割［7］，该深度学习网络在2018年仍然能在Multimodal Brain Tumor Segmentation Challenge 2018中取得第二名的好成绩。Top2的网络模型［8］在U-Net的基础上做出了一些改动，该模型与U-Net同样采用4次下采样，通过Dice loss可以提升模型的Dice的得分，并使用Leaky Re LU激活函数提高性能。因此，在医学图像分割中，损失函数的设计与选用尤为重要［9］，其中用得较多的损失函数为Dice损失函数（Dice coefficient loss）［10］、交叉熵损失函数（cross entropy loss）、相对熵损失函数（KL loss）等。目前，一些基于CNN的算法对神经元重要细节分割上有所缺失，分割图像中像素间的关联度不强，导致线条也存在断裂。因此，如何利用注意力机制，对电镜神经元图像中神经元的某些特征进行重点学习、怎样针对神经元图像选用合适的损失函数或者组合损失函数，都是值得研究的重要问题。

综上所述，本文提出一种用于电镜图像中神经元结构分割的深度学习模型：Attention-KL U-Net（A-KLU-Net）。该模型主要由4个下采样、4个上采样、一个自注意力模块组成，并使用组合损失函数来评估模型。本文以含有BN层和卷积时进行边缘补零的U-Net结构为骨架，经过大量实验发现自注意力机制的使用可以有效地避免神经元分割结果图中的断裂和多余部分的出现，并进一步提出了改进的注意力机制模型来加强对细节的学习；其次，在网络模型中使用了Dice与KL散度叠加的损失函数，使得模型可以更好地学习。

1 模型建立

1.1 U-Net模型

U-Net网络的核心思想是利用较少的训练集进行端对端的训练［11］，U-Net模型是建立在全连接卷积神经网络（fully connected network，FCN）［12］架构上，对FCN结构进行修改和更为精细设计。但该网络不能很好地恢复经池化层之后的图像细节结构信息，易造成分割断裂的现象［11］。基于此，本研究将在U-Net模型的基础上添加边缘补零和BN（batch normalization）层，通过在卷积过程中对图像进行边缘补零操作，使得网络更好地考虑到图像的边缘像素点；同时，图像的尺寸在卷积后保持原来大小，只在用最大池化层对图像特征进行池化时，改变图像分辨率。添加BN层主要是将每个通道输出数据的均值和方差进行归一化处理，使得其均值接近0、方差接近1，从而提升关键性指标F1，更好地提升网络分割性能。

1.2 采样模型

模型中下采样和上采样过程中每层的卷积结构都含有卷积层（Conv）、BN层以及线性整流单元（rectified linear unit，Re LU）激活层。与原始U-Net不同，本文模型每次卷积会添加zero padding，通过在卷积过程中对图像进行边缘补零操作，使得图像的尺寸在采样过程中不变，同时进一步让网络考虑边缘像素增强对细节的敏感度。下采样对图像特征进行提取，上采样对图像进行还原，用转置卷积（transpose convolution，TSConv）进行上采样，再与下采样对应的特征图拼接，对其卷积缩小通道数。模型如图1所示（图中卷积核大小为3*3）。

图1 上采样和下采样模型Fig.1 Upsampling and downsampling models

线性整流单元ReLU为

1.3 注意力机制及其改进模型

注意力机制选择性地关注整体信息的一部分，获得需要重点关注的区域，有效地提高了人类识别和认知的效率［13-15］。同理，在深度学习中提取到的信息流以同等重要性向后流动，通过某些先验信息，抑制某些无效信息的流动，使重要信息保留。因此，本文选取自注意力（self-attention）机制来对图像中的特征进行针对性学习，并对注意力模块进行了改进。模型如图2（a）所示，图中注意力机制由卷积层（Conv2d：卷积核尺寸为1*1卷积核）、BN层、激活层（激活函数为Sigmoid函数）组成。通过卷积降低特征图的通道数，使得特征维度为1；再使用BN层做归一化；最后，通过Sigmoid函数来进行激活，由于Sigmoid函数值域在（0，1）使得数据波动小，会生成一个合适的模板。

Sigmoid激活函数为

本文还通过实验发现在注意力机制中，把原图与最后一次上采样的结果进行拼接，用拼接的特征图卷积得到单通道的特征图，再与最后一次上采样的输出融合，可以很好地提高分割图像中需要的细节，像素点关联度也更好，故本文将使用改进后的自注意力模型。具体结构如图2（b）所示：将原图在第二维度与其拼接，给网络提供更多的图像细节和像素间的关联度，该方法理论上具有一定的通用性。

图2 改进前和改进后注意力模型Fig.2 Attention model before and after improvement

1.4 损失函数的组合

损失函数是机器学习里最为基础和关键的要素，决定了反向传播更新各层神经元连接之间的权重和偏差。本文结合训练数据，给出一个较优的损失函数。Dice本质上属于二分类，用于衡量两个样本的重叠部分，该指标属于［0，1］，越接近1重叠度越高。KL散度用于度量两个分布之间相似性。网络在选取损失函数时需要考虑数据集图像的类别、细致度、分割目标大小等多方面因素。本文分别对3种损失函数进行了实验，选取最为契合的一种。实验研究发现损失函数采用Dice loss和KL loss的组合能提高F1指标，从而提升网络性能。

1.5 A-KLU-Net模型

本文以U-Net为骨架与本研究提出的改进注意力模型结合构建了A-KLU-Net模型。如图3所示，该模型主要由4个下采样（池化层的卷积核为2*2）、4个上采样、一个Self-Attention模块组成。损失函数采用Dice loss与KL loss叠加的形式。图3中Conv2d为卷积层，Maxpool2d为最大池化层。

图3 A-KLU-Net模型示意图Fig.3 Schematic diagram of A-KLU-Net model

1.6 算法框架

算法框架主要分两大部分：网络训练和测试（见图4）。网络训练部分首先将图像进行几何变换（旋转和翻转），在不改变原始图像的条件下进行了数据增广，提升网络模型的鲁棒性。把训练集输入A-KLU-Net模型，对网络的输出用Softmax函数对图像的每个像素进行分类得到背景和细胞膜两种图，之后用损失函数对预测图和标签进行评估的同时通过Validation模块测试网络得分，通过Adam优化器优化，直到（5）式中l o s s值降到最低，模型训练完成，保存在Validation模块中得分最高的网络。在训练过程中，为了提高效率需要给网络设置一个耐性值patience，即在patience个epoch网络性能没有提高就结束训练。测试过程：将测试集的电镜神经元图像输入保存的训练效果最佳网络，对图像进行预测，从而达到分割的效果。

图4 网络结构框图Fig.4 Network structure block diagram

2 实验结果与分析

2.1 实验数据与基本参数

本研究采用的数据集为ISBI 2012提供的数据集［16-17］，训练数据是来自果蝇第一龄幼虫腹侧神经索（VNC）的连续部分透射电子显微镜（ssTEM）数据集的30个部分，共30张图片，每张图片像素大小为512*512，标签为二值图像，即分割对象的像素为黑色（细胞膜），其余像素为白色。本文采用旋转和翻转对数据集进行增广以减小过拟合的发生。网络模型实现的软件和模块分别是Python3.7.10、Torch1.8.1和CUDA10.0，并且使用NVIDIA GeForce RTX 3090 GPU进行训练。网络基本参数：优化器为Adam，学习率为0.001，训练Batch为4，测试Batch为9。

2.2 模型评估指标

本文为了更好地评估网络的分割效果，采用以下4个评估指标：Accuracy（A）表示正确性，即分类背景和细胞膜的像素在整个图像像素中的百分比；Precision（P）表示精确度，即测试样本中准确预测为细胞膜的数除以准确预测的总数；Recall（R）表示召回率又称为敏感度；医学图像分割常用指标Dice通用式：F1表示预测图与真实图的重叠率。

式中，T P表示预测为细胞膜，标签为细胞膜；T N表示预测为背景，标签为背景；F P表示预测为细胞膜，标签为背景；F N表示预测为背景，标签为细胞膜。

2.3 结果分析

2.3.1 自注意力机制与损失函数对模型的影响 为了进一步分析本文算法中的自注意力机制与损失函数对模型的影响，本研究通过对损失函数和自注意力机制的组合、不同损失函数的组合分析可知：①KL散度可以有效地提升正确性，自注意力机制可以有效提高正确性和精确度。②如表1所示，前三行数据与后三行都对不同的损失函数进行了比较，发现单独使用Dice损失函数比交叉熵损失函数（Ce loss）准确度上低，但Dice可以有效地提高F1指标和召回率，且与自注意力机制模型结合时，在关键性指标F1上也有一定的提升。基于此，本文采用组合损失函数方式，把Dice和交叉熵与KL散度进行适当的叠加，通过实验研究发现，当Dice与0.1倍的KL散度结合时，分割效果较之单独的Dice有所提升，具体对比表1的第二行和第五行，发现其对A、F1指标有一定提升。为提高效率，本文添加了Validation模块，在20代内网络无法提升就结束训练。为了避免测试集和验证集可能对实验的影响，采用不同的两种数据集比例进行实验。具体数据见表1（n=0.1）。

由表1可知，表中第四行和第五行带有自注意力机制，使得模型对图像细节分割的指标P有所提升，从而很好地提升了本文关键性指标F1。图5（a）为经过self-attention模块之后的效果图，图5（b）和图5（c）从左到右分别是原始图、标准图、预测图。通过对比图5（a）、图5（b）和图5（c）中红色方框的分割情况，可以看出有self-attention模块对图像的复杂区域分割的效果更好，给后续像素分类提供了更多的细节；上述差别反映到网络模型评价指标中，可以发现图像的像素分类精确度（P）有很好的提升。所以，在最后一次上采样结束后，使用Self-Attention可以有效减少图像信息的丢失，在恢复图像分辨率时更好地分割出图像的细节，从而提高网络模型的分割效果。

表1 自注意力与损失函数效果对比Tab.1 Effect comparison of self-attention and loss function

图5 Self-attention效果对比图Fig.5 Self-attention effect comparison graph

2.3.2 实验分析 为了与现有算法U-Net、文献［11］、文献［18］比较，实验采用25张图作为训练集，5张图作为测试集。如表2所示，本文和现有算法在分割效果上较之原始U-Net都有较好的提高。同时，本文与文献［11］相比，在正确性（A）、召回率（R）上分别提高了0.018 33和0.003 04，并在其基础上添加了精确度（P）对分割的精确度进行评估，相对U-Net提高了0.013 99；相对于文献［18］的方法用于本数据集，在精确度上稍有不足，但在其他3个指标上都有明显优势。文献［11］将分组卷积与改进的U-Net网络结构相结合，其F1指标与本文算法基本持平，但本文算法在准确度和召回率上都高于该算法。映射到分割图像上，即图像分割得更为准确和细致。

表2 不同方法比较Tab.2 Comparison of different methods

图6为本文算法与原始U-Net在ISBI 2012数据集中的部分图像的分割效果图。通过对图像彩色化可使图像可视效果更好。其中，图6从左到右分别是原始图像、标准图和3种方法的预测图。通过观察并将图像彩色化可以看出本文模型在细胞膜分割细节方面更好且大部分地方无明显断裂现象。通过对比可以看出本文在细胞膜分割的粗细上，更加与标准图匹配，在细节上也更加匹配且与label对比多余部分更少。本文算法更好地避免了细胞膜分割的断裂和多余部分的出现，具体如图（6）中的红色标记点所示。

图6 分割结果图Fig.6 Segmentation result graph

3 结语

本文提出了对神经元图像的深度学习分割模型A-KLU-Net，该网络在ISBI 2012数据集上进行了实验，其正确率、F1指标、精确度和召回率分别达到0.930 43、0.956 79、0.953 26、0.960 34。通过损失函数和自注意力机制的组合、不同损失函数的组合分析，验证了改进的自注意力机制可以有效提升图像分割的准确度，Dice与KL散度的结合对本研究的数据集的效果更好。后续可以考虑通过加入多尺度输入输出改善网络高层特征逐步丢失初始信息的问题，并且多个输出也有利于网络的梯度反向传播、改善训练效率，也可以在多尺度特征融合。