论文地址：http://www.arxiv.org/pdf/1505.04597.pdf

U-Net：用于生物医学影像分割的卷积网络

摘要

普遍认为深度网络的成功训练需要数千个标记好的训练样本。在本文中，我们提出了一种网络和训练策略，依靠高效的数据增强，以更有效地利用现有的标记样本。该结构由捕获上下文的收缩路径和对称的支持精确定位的展开路径组成。我们证明这样的网络可以从很少的图像中进行端到端的训练，并在ISBI电镜神经元结构分割挑战中优于之前最好的方法(滑动窗口卷积网络)。用该网络在透射光镜图像(相衬和DIC)上进行训练，我们在ISBI2015细胞追踪挑战赛中以巨大优势胜出。此外，网络性能也很高。在最新的GPU上，512x512图像的分割所需时间不到一秒。可在http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net下载基于Caffe的完整实现和训练好的网络。

1.引言

在过去的两年里，深度卷积网络在许多视觉识别任务中表现都优于SOTA，例如[7,3]。虽然卷积网络[8]已经存在很长时间了，但是考虑可用训练集的大小和网络的大小，其成功程度依然有限。Krizhevsky等人[7]的突破在于对拥有100万张训练图像的ImageNet数据集进行8层数百万个参数的大型网络的有监督训练。从那时起，出现了更多更大更深的网络[12]。

卷积网络的典型应用是在分类任务上，其中图像的输出是一个类标签。然而，在许多视觉任务中，特别是在生物医学影像处理中，期望输出应包括定位，也就是说，应该给每个像素指定一个类标签。此外，在生物医学任务中，通常也难以得到数以千计的训练图像。因此，Ciresan等人[1]在滑动窗口中训练网络，通过以像素周围的局部区域(patch)作为输入，来预测每个像素的类标签。首先，这个网络可以本地化。其次，在patch方面的训练数据远多于训练图像的数量。该网络在ISBI 2012上以巨大的优势赢得了EM分割挑战。

不过，显然Ciresan等人的[1]策略有两个缺点。首先，它的速度非常慢，因为每个patch都必须单独运行网络，而由于patch有重叠，存在大量的冗余现象。其次，在定位精度和上下文使用之间需要权衡。较大的patch需要更多的max-pooling层，导致定位精度降低；而较小的patch只允许网络看到少量上下文。最近的方法[11,4]提出了一种利用多层特征的分类器输出，可以同时做到良好的本地化与上下文使用。

在本文中，我们构建了一种更优雅的结构，即所谓的“全卷积网络”[9]。我们修改并扩展了这个结构，使其能在少量训练图像下工作，并产生更精确的分割，见图1；[9]中的主要思想是在连续的层中补充一般的收缩网络，在该层中，池化运算符被上采样运算符代替。因此，这些层提高了输出的分辨率。将收缩路径的高分辨率特征与上采样输出相结合，实现了本地化。连续的卷积层可以根据这些信息组合出更准确的输出。

图1. U-net结构(以最低分辨率的32x32像素为例)。 每个蓝色框对应一个多通道特征图。 通道数标在方框的顶部。 x-y尺寸位于框的左下角。白色框代表复制的特征图。箭头表示不同的操作。

我们结构中的一个重要修改是，在上采样部分我们也有大量的特征通道，这允许网络传播上下文信息到更高分辨率的层。因此，扩展路径与收缩路径或多或少是对称的，并生成一个u型架构。该网络没有全连接层，仅使用每个卷积的有效部分，即，分割图仅包含像素，在输入图像中可以使用其完整上下文。该策略允许使用叠层策略无缝地分割任意大的图像(见图2)。为了预测图像边缘区域的像素点，通过镜像输入图像来推测缺失的上下文。这种平铺策略对于将网络应用至大图像非常重要，因为否则分辨率会受到显存的制约。

图2.任意大图像无缝分割(此处为EM栈中神经元结构的分割)的重叠拼贴策略。对黄色区域中分割的预测需要蓝色区域内的图像数据作为输入。缺失的输入数据则通过镜像来推断。

对于我们的任务，可用的训练数据非常少，通过对现有训练图像进行弹性形变来达到数据增强。这允许网络学习这种变形的不变性，而不需要在标记好的图像语料库中看到这些变换。这在生物医学分割中尤为重要，因为变形是组织中最常见的变异，因此这可以有效地模拟真实的变形。Dosovitskiy等人[2]在无监督特征学习的范畴中展示了数据增强对于学习不变性的价值。

许多细胞分割任务中的另一个挑战是分离同一类别的接触对象，见图3。为此，我们建议使用加权损失，即接触细胞之间的分离背景标签在损失函数中获得较大的权重。

图3.用DIC(微分干涉对比)显微镜记录的玻璃上的HeLa细胞。( a )原始图像( b )覆盖与ground truth分割。不同的颜色表示不同的HeLa细胞。( c )生成的分割蒙版(白色：前景，黑色：背景)。( d )使用逐像素损失权重进行映射，以强制网络学习边缘像素。

所得到的网络适用于各种生物医学分割问题。在本文中，我们展示了EM栈中神经元结构分割的结果(一项从ISBI2012开始至今的比赛)，在那里我们超过了Ciresan等人[1]的网络。此外，我们展示了ISBI2015细胞追踪比赛的光镜图像细胞分割结果。在这两个最具挑战性的2D透射光数据集上，我们以较大的优势胜出。

2.网络结构

网络结构如图1所示。它由收缩路径(左边)和扩展路径(右边)组成。收缩路径遵循卷积网络的典型结构。其包含重复的两个3x3卷积(无padding)，每个卷积后都有一个ReLU单元和一个2x2的最大池化操作，用于下采样。在每个下采样步骤中，我们将特征通道的数量翻倍。扩展路径中的每个步骤都包括对特征图进行上采样，然后是将特征通道数量减半的2x2卷积（“向上卷积”），与来自收缩路径的相应裁剪的特征图的串联以及两个3x3 卷积，每个后跟一个ReLU。由于每次卷积都会丢失边缘像素，因此有必要进行裁剪。在最后一层，使用1x1卷积映射每个64组件特征向量到所需的类数。该网络总共有23个卷积层。

为了无缝拼接输出分割图(见图2)，输入图块大小的选择很重要，以便将所有2x2最大池化操作应用于x、y大小为偶数的层。

3.训练

利用输入图像及其对应的分割图，采用Caffe[6]的随机梯度下降来对网络进行训练。由于没有进行padding，输出图像比输入图像小一个恒定的边界宽度。为了最小化开销并最大限度地利用显存，我们倾向于使用大的输入块而不是大的batch size，以将batch减少至单个图像。相应的，我们使用高动量(0.99)，以用大量先前所见的训练样本来确定当前优化步骤中的更新。

能量函数由最终特征图上的像素级soft-max与交叉熵损失函数组合而成。定义soft-max为 p k ( x ) = e x p ( a k ( x ) ) / ( ∑ k ′ = 1 K exp ⁡ ( a k ′ ( x ) ) ) {p_k}(x) = exp({a_k}(x))/(\sum\nolimits_{k' = 1}^K {\exp ({a_{k'}}(x))} ) pk​(x)=exp(ak​(x))/(∑k′=1K​exp(ak′​(x)))，其中 a k ( x ) {a_k}(x) ak​(x)表示特征通道 k k k在像素位置 x ∈ Ω x \in \Omega x∈Ω和
Ω ⊂ Z 2 \Omega \subset {{\Bbb Z}^2} Ω⊂Z2
的激活情况。 K K K为类的个数， p k ( x ) {p_k}(x) pk​(x)为近似最大函数。即对于具有最大激活 a k ( x ) {a_k}(x) ak​(x)的 k k k， p k ( x ) ≈ 1 {p_k}(x) \approx 1 pk​(x)≈1，而对于其他的 k k k， p k ( x ) ≈ 0 {p_k}(x) \approx 0 pk​(x)≈0。 然后，交叉熵则利用
E = ∑ x ∈ Ω w ( x ) l o g ( p l ( x ) ( x ) ) E = \sum\limits_{x \in \Omega } {w(x)log(} {p_{l(x)}}(x)) E=x∈Ω∑​w(x)log(pl(x)​(x))
惩罚每个位置上 p l ( x ) ( x ) {p_{l(x)}}(x) pl(x)​(x)与1的偏差。

其中 l : Ω → { 1, ⋯ ,K}  l:\Omega \to {\text{\{ 1,}} \cdots {\text{,K\} }} l:Ω→{ 1,⋯,K} 是每个像素的真实标签， w : Ω → R w:\Omega \to R w:Ω→R是我们引入的一个权重图，其能在训练阶段给予一些像素更高权重。
我们为每个ground truth分割预先计算权重图，以补偿训练数据集中某些像素频率不同的类别，并迫使网络学习我们在相连细胞之间引入的小的分隔边界(见图3c和图d）。
使用形态学操作计算分离边界。映射权重等于
w ( x ) = w c ( x ) + w 0 ⋅ exp ⁡ ( − ( d 1 ( x ) + d 2 ( x ) ) 2 2 σ 2 ) w(x) = {w_c}(x) + {w_0} \cdot \exp ( - \frac{{{{({d_1}(x) + {d_2}(x))}^2}}}{{2{\sigma ^2}}}) w(x)=wc​(x)+w0​⋅exp(−2σ2(d1​(x)+d2​(x))2​)
其中 w c ( x ) : Ω → R {w_c}(x):\Omega \to R wc​(x):Ω→R为权重图，用于平衡类频率； d 1 : Ω → R {d_1}:\Omega \to R d1​:Ω→R表示距最近细胞边缘的距离， d 2 : Ω → R {d_2}:\Omega \to R d2​:Ω→R表示矩第二近细胞边缘的距离。在我们的实验中，令 w 0 = 10 {w_0} = 10 w0​=10， σ ≈ 5   pixels \sigma \approx 5\,{\text{pixels}} σ≈5pixels。

在具有多重卷积与多重路径的深度网络中，良好的权值初始化非常重要。否则，网络的某些部分可能会过度激活，而其他部分则永远不会激活。理想情况下，应调整初始权重，使得网络中的每个特征图具有相近单位方差。对于我们的网络结构(交替卷积和ReLU层)，这可以通过从高斯分布(标准差为 2 / N \sqrt {2/N} 2/N ​)中提取初始权重来实现，其中 N N N表示一个神经元[5]的传入节点数。例如，对于包含3x3卷积与64个特征通道的上一层， N = 9 ⋅ 64 = 576 N = 9 \cdot 64 = 576 N=9⋅64=576。

3.1.数据增强

当只有少量训练样本可用时，数据增强对于向网络传授所需的不变性和鲁棒性至关重要。对于显微镜图像，我们首先需要平移和旋转不变性以及对形变和灰度变化的鲁棒性。特别是训练样本的随机弹性形变似乎是训练只有少量标记图像的分割网络的关键概念。我们在一个粗糙的3×3网格上使用随机位移向量生成光滑形变。位移采样自一个服从10像素标准差的高斯分布。然后使用双三次插值计算每个像素的位移。收缩路径末端的Drop-out层进一步执行隐式的数据增强。

4.实验

我们演示了u-net在三种不同分割任务中的应用。第一个任务是在电子显微记录中分割神经元结构。图2展示了数据集与我们得到的一个分割的例子。我们提供完整结果作为补充材料。数据集由EM分割挑战赛提供，该挑战始于ISBI 2012，目前仍接受新的提交。训练数据为30幅果蝇一龄幼虫腹侧神经索(VNC)系列切片透射电镜图像(512x512像素)。每幅图像都带有一个标注好的ground truth分割图，用于表示细胞(白色)和膜(黑色)。测试集是公开的，但分割图保密。通过将预测得到的膜概率图发送给主办方，可以得到一个评价。通过对图进行10级的阈值化，并计算“warping error”，“rand error”和“pixel error”来进行评估[14]。

u-net(7个旋转版本输入得到的平均值)在没有任何进一步的预处理或后处理的情况下实现warping error为0.0003529(新的最高分，见表1)，rand error为0.0382。

表1. EM分割挑战排名[14] (2015.3.6)，按warping error排序

图4. ISBI细胞追踪挑战的结果。( a )“PhC-U373”数据集输入图像的一部分。( b )人工得到的ground truth(黄色边框)分割结果(青色蒙版)。( c )“ DIC-HeLa”数据集的输入图像。( d )人工得到的ground truth(黄色边框)分割结果(随机彩色蒙版)。

表2.ISBI2015细胞追踪挑战的细分结果(IOU)

这明显优于Ciresan等[1]滑动窗口卷积网络得到的结果，其最佳提交的warping error为0.000420，rand error为0.0504。在rand error方面，在此数据集上唯一性能更好的算法使用了高度数据集相关的后处理方法，这些方法应用于Ciresan等人的概率图。

我们还将u-net用于光镜图像下的细胞分割任务。该分割任务是ISBI2014、2015细胞追踪挑战的一部分[10,13]。第一个数据集“PhC-U373”，包含在相衬显微镜下记录的于聚丙烯酰胺底物上的胶质母细胞瘤-星形细胞瘤U373细胞(见图4a,b和补充材料)。它包含35个部分标注的训练图像。这里我们得到的平均IOU(intersection over union)为92%，明显优于次优算法的83%(见表2)。第二个数据集“DIC-HeLa”是平板玻璃上的HeLa细胞，通过差分干涉对比(DIC)显微镜记录(见图3、图4c、d和支撑材料)。其包含20个部分标注的训练图像。在这里，我们实现了平均IOU为77.5%，明显优于次优算法(46%)。

5.结论

在不同的生物医学分割应用中，u-net架构都取得了很好的性能。由于使用了带弹性形变的数据增强，它只需要少量标注好的图像，就能在NVidia Titan(6gb)上达到10小时这一合理的训练时间。我们提供了完整的基于Caffe的实现与训练好的网络。我们相信u-net架构可以很容易地应用到更多的任务中。