U-Net:用于生物医学图像分割的卷积神经网络
摘要 :大多数人认为,训练一个有效的神经网络需要海量带标签的数据。但在这篇文章里我们提出了一种网络和训练策略,它可依赖于充分使用数据增强技术使带标签的数据使用的更有效率,从而减少对大量带标签数据的依赖。该体系结构包括一个获得不同尺寸特征图定位信息的下采样收缩路径和一个与之对称上采样的扩展路径,该网络的输出图像具有高分辨率精确定位的特性。 我们证明这样一个端到端的网络可以被少量图片训练且表现优于会先前在ISBI的挑战赛中的最佳方法(滑动窗口卷积网络)对于电子显微镜堆栈中神经元结构的分割表现。使用同样的光学显微图像训练该网络(相位对比度和DIC),我们以较大优势赢得了2015年的ISBI细胞跟踪挑战赛中这两种类型的比赛。并且我们网络的测试花费时间短,在最新的GPU上分割一个512x512的图像花费不到一秒钟。
关键词:卷积神经网络 医学图像 U-Net 语义分割
第1章 绪论
在过去的两三年里,深卷积网络的性能超过了许多视觉识别任务中的其它类型神经网络的表现如。现在卷积网络已经存在了相当长一段时间,但受限于医学上可用训练集的大小和当下神经网络规模的大小,卷积神经网络的发展遇到了瓶颈。Krizhevsky等人所做出的创新归功于在ImageNet数据集上用100万张带标签图像数据有监督的训练一个深度为8层的带有数百万个参数的卷积神经网络,从那时起规模更大网络、程度更深的卷积网络不断地被有效搭建并训练出来[1]。
卷积神经网络的经典应用是图像的多分类任务,其网络的输出是图像所属的标签。在许多图像处理任务当中,我们期望包括目标位置信息的输出,即:算法应当为单个像素分配类别属性。然而在相当多的医学图像任务中,同以往收集大量数据来训练神经网络所不同,能够收集到数千张有效的标注的医学图像常常是我们无法想象的。所以Ciresan等人提出图像的滑动窗口法,并通过提供围绕着某点一个小区域的像素作为输入训练网络以预测整幅图的类别标签,这个网络具有的优势在于它可以标注出目标位置;然后,因其训练数据在滑动窗口的角度上讲其数量是远大于原训练图像数据量的,所以其最终在ISBI2012年挑战赛的EM细分组中以较大的优势取得了胜利[2]。
很明显,Ciresan等人的网络存在两个劣势。首先,因为滑动窗口的算法设计其网络需要为数量与大小远多于原图像的数据做出处理,导致网络会重复计算图像的某些重叠区域而导致整个单幅图像的运行速度会变慢。其次,这种网络在目标定位的准确性和上下文语境信息的使用之间有一个信息丢失的权衡。当运行较大的补丁时需要更多的MaxPo-oling层,这样会导致定位的精度降低;但当运行较小的补丁时模型只能得到原图中很少的上下文信息导致网络的定位精度下降。目前最新的方法[11,4]提出了考虑汇总来自多层网络的特征来作为一个分类器输出,如此使得良好的定位精度和上下文信息的兼顾在网络搭建中成为了可能。
在本文中,我们在一个更为简洁优美的结构,即所谓的“全卷积神经网络“[9]上建立了一个网络。我们修改、扩展了这个体系结构使它可以经较少的图像训练后能生成表现优异精确的语义分割结果;网络结构见图1-1。提出这个架构的论文的主要思想是通过连续上采样层补充通常情况下缩小的网络,其中池化操作被上采样操作所替换。因而本网络提高了输出结果的分辨率。为了图像目标定位的准确性,高分辨率的特征图将会从下采样的收缩路径中结合到上采样的输出操作中,一个连续卷积层可以去学习如何基于这种信息生成更精确的输出。
我们在结构上所做的一个重要创新是在上采样部分结合原先下采样过程中得到大量的位置特征信息,这样网络便可以在更高分辨率的网络层中融合到精确的定位信息。也因如此,我们网络模型的上采样部分与下采样路径一定是对称的,并因此产生一个u形结构的网络架构。我们网络没有任何全连接层并且仅使用每个卷积中有效的一部分,即分割图中仅包含像素,对此完整的上下文信息可以被用在输入图像当中。此策略允许通过重叠平铺策略(见图1-2)无缝分割任意大的输入图像。预测在图像中边界区域的像素,那些被丢失的上下文信息可以通过镜像输入图像来推断出。这种重叠平铺的网络对于将网络应用于大型图像处理很关键,不然输出图像的分辨率将会受到限制。
对于我们的任务,可用的训练数据很少,我们通过对现有训练图像应用弹性变形来大量扩充我们的训练数据。这使得网络能够学习对这种弹性形变的不变性,而不必需要在标注的图像库中看到这些转换。这在医学图像分割中很关键,因为变形是动物组织变化中最常见并且可以被有效地模拟地。Dosovitskiy等人在中证明了在无监督特征学习的范围内数据增强对学习不变性的价值。
在许多细胞分割任务中的另一个难点是分离相互接触的同一类对象;参见图3-1。为此,我们计划使用加权损失,在损失函数中细胞之间的接触边界会因此获得较大的损失权值。
本文所提网络适用于多种医学图像分割任务。在这里,我们分析了在EM stacks挑战赛中神经元网络结构分割的结果,我们网络的表现超过了Ciresan等人。并且,我们同样分析了ISBI细胞追踪挑战赛的光学显微镜图像细胞分割比赛。在这两个最具难度的2D光学数据集上中,我们都以较大的优势获胜。
第2章 网络架构
文本所提出的网络架构如图1-1所示。它包括一个收缩路径(左侧)和扩展路径(右侧)。收缩路径如同典型的卷积网络的下载样结构一般,它包括两个重复的3x3卷积(未加padding)每一个之后都应用一次线性矫正单元(ReLU)和一个步长为2的2x2最大池化层用于下采样,在每个下采样步骤中,我们将特征通道数加倍。在扩展路径中的每一步都包括特征映射后接2x2上卷积,该卷积将特征通道的数量减半,并且与来自收缩路径的相应裁剪的特征映射的级联,以及两个3x3卷积操作,每个卷积后接一个ReLU。由由于每次卷积都会丢失边界像素,因此必须进行裁剪。在最后一层,使用1x1卷积将每个64分量特征向量映射到到所需的类别数。
要允许输出分段映射的无缝平铺(参见图1-2),选择输入平铺大小非常重要,以便所有2x2-MaxPooling操作应用于x和y大小相等的层。
第3章 U-Net神经网络的训练
利用输入图像及其相应的标签图在随机梯度下降方法下的Caffe模型实现[6]。由于unpadding卷积层的存在,输出图像以不变的边框宽度小于输入图像。为减少开销并最大限度地利用对于GPU内存,我们更倾向于更大的数据输入而不是大的批处理,因此将批处理缩减为单个图像。因此,我们使用高动量参数0.99,使得大量先前看到的训练样本确定当前优化步骤中的更新。
网络的损失值由最终的特征图计算像素级soft-max函数与交叉熵损失函数相结合。Soft-max函数的定义为,其中表示在像素位置x上第k个特征通道的激活值。K表示总的类别数,估计的最大函数。即:≈1表示第k类有最大的激活值,并且≈0对于的k类。交叉熵然后惩罚每一个位置上的偏差当 使用
(1)
其中:Ω→{1,…,K}表示每一个像素的真是标签,ω:Ω→R是我们引进的一个加权图,在训练中提供给一些更重要的像素。
在网络训练过程中,我们线下计算每个标注图像的分割权值图,补偿某类像素不平衡的输入图像,即加大网络学习细胞间的接触轮廓的损失(见图3-1 c和d)。
权重边界图需要经形态学处理计算。这个重值图计算如下
其中,:Ω→R是平衡训练数据不平衡的权重图,:Ω→R表示到像素到最近单元格的边界的距离,:Ω→R表示像素到最近的第二个细胞边界的距离。在实验中,我们将设为10,σ≈5。
在大多数的神经网络的训练过程来说,一个有代表性的网络初始权重是极为关键的。否则,网络中的部分神经元可能会被赋予过大的权重,而其它部分则对网络的输出影响微乎其微。理想的条件下,网络的初始权重应调整应该与网络中的特征图拥有近似的方差。对于我们网络的架构(交替的卷积层和ReLU激活层)可以从标准差为的高斯分布中设置一个初始权值,其中N表示一个神经元输入的节点数[5]。例如,对于前一层中的3x3卷积和64个特征通道N=9 * 64=576。
3.1 数据增强
数据增强对于教导网络所需的不变性和鲁棒性,当只有很少的训练样本可用时。对于光学图像,我们尤其看重图像的旋转和平移变化后的不变性以及对变形和灰度值变化的抗干扰性。尤其是带标签样本的随机弹性变形,这在我们看来似乎是训练一个训练数据很少的语义分割网络的关键思想。实验中采用3×3粗糙的随机位移向量生成平滑弹性形变,位移变量是从10像素标准差的高斯分布中采样的,然后使用双三次插值计算每个像素位移。下采样收缩路径的末端采用Drop-out层可以进一步隐式执行数据扩充。
第4章 实验
文章分析了U-Net模型在三种不同图像分割任务中的表现。第一个任务是对电子显微镜所记录的图像做神经元结构分析。数据集的一个样例和我们分析得到的图像如图1-2所示。数据集由EM segmentation challenge[14]提供,该挑战始于ISBI 2012年。训练数据是序列切片透射电子显微镜30幅图像(512x512像素)果蝇一龄幼虫腹神经元。每个原始训练图像都会有与之相应的完全标注的细胞分割图和细胞边界。测试集的原始数据是公开的,但是它的语义标注图像是不公开的。可以通过发送预测图像到组织者的网站获得预测的准确率。通过对输出图像进行阈值化并在10个不同级别分别计算他们的“Warping误差”、“Rand误差”和“Pixel误差”[14]来完成评估。U-Net(输入图像的7个变化样本的平均值)在没有进一步预处理或后处理的情况下实现了0.0003529的warping误差(至今最好的分数,见表1)和0.0382的Rand误差。
这明显优于滑动卷积神经网络Ciresan等人的研究结果,其最佳分数有0.000420的Warping误差,且Rand误差为0.0504。就Rand误差而言,此数据集上唯一性能较好的算法使用了经Ciresan等人的提出的概率图后处理方法处理后的特定数据集 [1],代表性不强。
我们还将U-Net应用于光显微镜图像中的细胞分割任务。此分离任务是2014年ISBI小区跟踪挑战赛的一部分,并且2015年[10,13]。第一组数据集用相差显微镜记录的聚丙烯酰亚胺基质上的胶质母细胞瘤星形细胞瘤细胞(见图4a、b和补充材料)。它包含35个部分标注的图像。在这组数据集上,我们实现了92%的平均IOU(“交集对并集”),这明显优于83%的次优算法(参见-表2)。第二组数据集“DIC HeLa”是平板玻璃上的HeLa细胞,通过差分干涉对比显微镜所记录(见图3,图4c,d和补充材料)。它包含20个部分标注的训练数据。在这里我们得到了77.5%的平均IOU,这也极大的优于46%的次优算法。
5 结论
U-Net架构在各种不同的差异性非常大的医学图像分割任务中都取得了优异的语义分割表现。借助于弹性变形的数据增强技术,网络可以在只有少量标注数据的情况下被有效的训练,并且有一个可以令人接受的训练时间,在英伟达 Titan GPU上训练时间仅为10小时。最后,我确信U-Net模型架构可以被很容易地应用于更多不同的医学分割任务中。