Olaf Ronneberger, Philipp Fischer, and Thomas Brox
Computer Science Department and BIOSS Centre for Biological Signalling Studies,
University of Freiburg, Germany
摘要
人们普遍认为,对深层网络的成功训练需要数千份已标注的训练样本。在本文中,我们提出了一个网络和训练策略,为了更有效的利用标注数据,我们是用数据扩张的方法(data augmentation )的方法。该体系结构由两部分组成:包括一个收缩路径(contracting path )来获取上下文信息以及一个支持精确定位的对称扩展路径(symmetric expanding path)。我们证明,这样的网络可以从非常少的图像中进行端到端的训练,并且在ISBI(生物医学成像国际研讨会)挑战中优于先前的最佳方法(一个滑动窗口的卷积网络),用于在电子显微镜的堆栈中分割神经元结构。使用相同的网络,在透射光显微镜图像上(阶段对比DIC(鉴别干涉对比显微镜)),我们在这些类别中以很大的优势赢得了ISBI细胞跟踪挑战。此外,网络速度很快。在最近的GPU上,一个512x512图像的分割需要不到一秒的时间。完整的实现(基于Caffe)和经过训练的网络可以在:
一、介绍
在过去的两年里,在许多视觉识别任务中,深度卷积神经网络的表现超过了艺术状态。虽然卷积神经网络已经存在很长时间了,但由于可用的训练集的大小和被考虑的网络的大小,它们的成功是有限的。Krizhevsky等人的突破是由于对一个大型网络的监督训练,该网络有8层,数百万个参数在ImageNet数据集上,包括了100万个训练图像。从那以后,甚至更大、更深入的网络被训练。
[图片上传中...(image.png-71d741-1523449829707-0)]
Fig. 1. U-net architecture (example for 32x32 pixels in the lowest resolution).
Each blue box corresponds to a multi-channel feature map. The number of channels is denoted on top of the box. The x-y-size is provided at the lower left edge of the box. White boxes represent copied feature maps. The arrows denote the different operations.
卷积神经网络的典型应用是分类任务上,对图像的输出是一个单独的类标签。然而,在许多视觉任务中,特别是在生物医学图像处理中,所期望的输出应该包括局部定位。一个类标签应该被分配给每个像素。此外,在生物医学任务中,数以千计的训练图像通常是无法触及的。因此,Ciresan等人建立了一个滑动窗口去训练了一个网络,通过预测每个像素点的类别对医学图像进行分割。首先,这个网络可以完成定位工作。其次,由于预测时要切patch,所以训练数据量比训练图片的数量大得多。由此产生的网络在2012年ISBI的EM分割挑战中获得了很大的优势。
很明显的是,这个网络有两个很明显的缺点:首先,它相当慢,因为每个patch都必须单独运行,由于重叠的patch会存在大量的冗余。(要分别预测每一个patch的类别,patch之间的重叠导致每次预测都要重复计算同一个点。)然后这个网络需要在局部准确性和获取整体上下文信息之间取舍。更大的patches需要更多最大池化层来降低定位精度,而小的patches则允许网络只看到很少的上下文。最近的方法提出了一个分类器输出,它考虑了来自多个层的特性。良好的定位和上下文的使用是可能的。
在这篇论文中,我们建立了一个更优雅的架构,即所谓的“全卷积网络”。我们修改并扩展了这个架构,使它可以使用非常少的训练图像就可以工作,并获得更高的分割准确率。fcn主要的想法是:修改一个普通的逐层收缩网络,用上采样(up sampling)操作代替网络后部的池化(pooling)操作。因此,这些层增加了输出的分辨率。为了定位,在网络收缩过程(路径)中产生的高分辨率特征(high resolution features) ,被连接到了修改后网络的上采样的结果上。在此之后,一个卷积层基于这些信息综合得到更精确的结果。
与fcn不同的是,我们的网络在上采样部分依然有大量的特性通道(feature channels)这使得网络可以将环境信息向更高的分辨率层(higher resolution layers)传播。结果是,扩张路径基本对称于收缩路径,并产生一个u型的体系结构。网络不存在任何全连接层(fully connected layers),并且,只使用每个卷积的有效部分,例如,分割图(segmentation map)只包含这样一些像素点,这些像素点的完整上下文都出现在输入图像中。overlap-tile strategy 允许无缝地分割任意大的图像(参见图2),为了预测图像边界区域的像素点,我们采用镜像图像的方式补全缺失的环境像素。这个tiling方法在使用网络分割大图像时是非常有用的,因为如果不这么做,GPU显存会限制图像分辨率。
Fig. 2. Overlap-tile strategy for seamless segmentation of arbitrary large images (here
segmentation of neuronal structures in EM stacks). Prediction of the segmentation in
the yellow area, requires image data within the blue area as input. Missing input data
is extrapolated by mirroring.
对于我们的任务,我们的训练数据非常少,我们通过对现有的训练图像应用弹性形变的方式来增加数据。这使得模型学习得到形变不变性,不需要在带标注的图像语料库中看到这些转换。这在生物医学的分割中尤其重要,因为组织的形变是非常常见的情况,并且计算机可以很有效的模拟真实的形变,在无监督特征学习的范围内,Dosovitskiy等人的学习增加数据去获得不变性的的价值已经显现出来。
在许多细胞分割任务中,另一个挑战是如何将同类别的相互接触的目标分开,为了达到这个目的,我们建议使用加权损失,我们提出了使用一种带权重的损失(weighted loss)。在损失函数中,分割相互接触的细胞获得了更大的权重。
由此产生的网络适用于各种生物医学分割问题。在本文中,我们展示了EM栈中神经元结构的分割(在ISBI 2012年开始的一场持续的竞争),我们的表现超过了Ciresan等人的网络。此外,我们还展示了2015年ISBI细胞跟踪挑战的光学显微镜图像的细胞分割结果。在这两个最具挑战性的2D传输光数据集上,我们获得了很大的优势。
二、网络结构
图1展示了网络结构,它由contracting path 和 expansive path组成。contracting path是典型的卷积网络架构。它的架构是一种重复结构,每次重复中都有2个卷积层和一个pooling层,卷积层中卷积核大小均为33,激活函数使用ReLU,两个卷积层之后是一个22的步长为2的max pooling层。每一次下采样后我们都把特征通道的数量加倍。expanding path中的每一步都首先使用反卷积(up-convolution),每次使用反卷积都将特征通道数量减半,特征图大小加倍。反卷积过后,将反卷积的结果与contracting path中对应步骤的特征图拼接起来。contracting path中的特征图尺寸稍大,将其修剪过后进行拼接。对拼接后的map进行2次33的卷积。最后一层的卷积核大小为11,将64通道的特征图转化为特定深度(分类数量,二分类为2)的结果。网络总共23层。
为了允许无缝地平铺输出分割图(参见图2),选择输入tile大小是很重要的,这样所有2x2的最大池采样操作都被应用到一个具有偶数x和y大小的层上。
三、训练
利用输入图像及其相应的分割图,利用Caffe的随机梯度下降法对网络进行了训练。由于没有填充的卷积,输出图像要比输入的小,因为它是一个恒定的边界宽度。为了最小化开销并最大限度地使用GPU显存,比起输入一个大的batch size,我们更倾向于大量输入tiles,因此我们使用了一个高动量(high momentum)(0.99),这样大量以前看到的训练样本决定了当前优化步骤中的更新。
能量函数是由一个像素级的softmax在最终的特征图和交叉熵损失函数之间计算出来的。Soft-max被定义为
。
表示在像素位置x的特征通道k的数目,,,K表示为类别的数量,是近似的最大值函数,即是对于最大的激活量来说的,对于所有其他的k来说。交叉熵在每个位置都受到惩罚,,表示每个像素的真实标签
(交叉熵函数)
是我们引入的一个权重图,在训练中凸显某些像素的重要性。我们对每一张标注图像预计算了一个权重图,来补偿训练集中每类像素的不同频率,使网络更注重学习相互接触的细胞之间的小的分割边界。我们使用形态学操作计算分割边界。权重图计算公式如下:
wc是用于平衡类别频率的权重图,d1代表到最近细胞的边界的距离,d2代表到第二近的细胞的边界的距离。基于经验我们设定w0=10,σ≈5像素。
Fig. 3. HeLa cells on glass recorded with DIC (differential interference contrast) microscopy.
(a)raw image. (b) overlay with ground truth segmentation. Different colors indicate different instances of the HeLa cells. (c) generated segmentation mask (white:foreground, black: background). (d) map with a pixel-wise loss weight to force the network to learn the border pixels.
网络中权重的初始化:我们的网络的权重由高斯分布初始化,分布的标准差为(N/2)^0.5,N为每个神经元的输入节点数量。例如,对于一个上一层是64通道的33卷积核来说,N=964。
3.1 数据增加
在只有少量样本的情况下,要想尽可能的让网络获得不变性和鲁棒性,数据增加是必不可少的。因为本论文需要处理显微镜图片,我们需要平移与旋转不变性,并且对形变和灰度变化鲁棒。将训练样本进行随机弹性形变是训练分割网络的关键。我们使用随机位移矢量在粗糙的3*3网格上(random displacement vectors on a coarse 3 by 3 grid)产生平滑形变(smooth deformations)。 位移是从10像素标准偏差的高斯分布中采样的。然后使用双三次插值计算每个像素的位移。在contracting path的末尾采用drop-out 层更进一步增加数据。
四、实验
我们演示了u-net对三个不同的分割任务的应用。第一个任务是在电子显微镜记录中分割神经元结构。图2显示了数据集和我们获得的分割的一个例子。我们提供完整的结果作为补充材料。该数据集是由2012年ISBI挑战提供的,目前仍对新贡献开放。训练数据是由果蝇第一个幼虫腹侧神经索(VNC)的连续段透射电子显微镜的30张图像(512x512像素)组成的。每个图像都有一个对应的完全带标注的分割图像(白色)和膜(黑色)。测试集是公开的,但是它的分割图是保密的。通过将预测的膜概率图发送给组织者,可以获得评估。该评估是通过在10个不同的层次上对图进行阈值计算,并计算“warping error”、“Rand error”和“pixel error”。
u-net(平均超过7个旋转版本的输入数据)在没有任何进一步的预处理或后处理错误的情况下实现了0.0003529的“warping error”(新的最好的分数,见表1)和一个0.0382的“rand error”。这比Ciresan等人的滑动窗口卷积网络的结果要好得多,后者的最佳“warping error”是0.000420和“rand error”为0.0504。在“rand error”方面,在这个数据集上唯一更好的执行算法使用高度数据集特定的后处理方法1应用于Ciresan等人的概率图。
Fig. 4. Result on the ISBI cell tracking challenge. (a) part of an input image of the“PhC-U373” data set. (b) Segmentation result (cyan mask) with manual ground truth(yellow border) (c) input image of the “DIC-HeLa” data set. (d) Segmentation result (random colored masks) with manual ground truth (yellow border).
我们还将u-net应用到一个细胞分割任务中,在微观图像中。这一分段任务是ISBI细胞跟踪挑战的一部分,2014年和2015年10月13日。第一个数据集“PhC-U373”包含由相衬显微术记录的聚丙烯酰胺-星形胶质细胞瘤细胞瘤(见图4a、b和Supp.材料)。它包含35个部分带注释的训练图像。这里我们实现平均IOU(“十字路口在联盟”)的92%,这是明显比第二个最好的算法为83%(见表2)。第二个数据集“DIC-HeLa”海拉细胞在一个玻璃(DIC)微分干涉对比显微镜记录。它包含20个部分带标注的训练图像。在这里,我们得到的平均IOU是77.5%,这明显优于第二种最好的算法46%。
五、结论
u-net体系结构在不同的生物医学分割应用上取得了很好的性能。由于有弹性的数据处理功能,它只需要很少的标注图像,而且在NVidia Titan GPU(6 GB)上只有10个小时的训练时间。我们提供完整的基于Caffe的实现和训练有素的网络。我们确信,u-net体系结构可以很容易地应用到更多的任务中。