U-Net: Convolutional Networks for Biomedical Image Segmentation

摘要

人们普遍认为，深度网络的成功训练需要数千个带注释的训练样本。在本文中，我们提出了一种网络和训练策略，它依赖于数据扩充的强大使用，以更有效地使用可用的带注释的样本。该体系结构由捕获上下文的收缩路径和支持精确定位的对称扩展路径组成。我们证明这样的网络可以从非常少的图像端到端的训练，并且在ISBI挑战中在电子显微镜栈中神经结构的分割上胜过先前的最佳方法(滑动窗口卷积网络)。我们使用相同的网络训练透射光学显微镜图像(相位对比和DIC)，在2015年ISBI细胞跟踪挑战赛中，我们在这些类别中获得了巨大的优势。此外，网络速度很快。在最近的GPU上，512x512图像的分割需要不到一秒的时间。

1、简介

在过去的两年中，深度卷积网络在许多视觉识别任务中的表现都超过了当前的技术水平。虽然卷积网络已经存在很长时间了，但是由于可用训练集的大小和考虑的网络的大小，它们的成功是有限的。Krizhevsky等人对[7]的突破是由于对一个8层的大网络进行了监督训练，在100万张训练图像的ImageNet数据集上有数百万个参数。从那时起，甚至更大更深的网络被训练成[12]。

卷积网络的典型用法是在分类任务中，图像的输出是单个类标签。然而，在许多视觉任务中，特别是在生物医学图像处理中，需要的输出应该包括定位，即定位。类标签应该被分配给每个像素。此外，成千上万的训练图像通常超出了生物医学任务的范围。因此，Ciresan等人在滑动窗口设置中训练网络，通过提供像素周围的局部区域(patch)作为输入来预测每个像素的类标签。首先，这个网络可以本地化。其次，以patch表示的训练数据远远大于训练图像的数量。由此产生的网络在2012年ISBI会议上以较大优势赢得了EM细分挑战。

显然，Ciresan等人的[1]策略有两个缺点。首先，它非常慢，因为每个补丁都必须单独运行网络，而且由于重叠的补丁存在大量冗余。其次，本地化的准确性和上下文的使用之间存在权衡。较大的补丁需要更多的最大池层，这会降低定位精度，而较小的补丁只允许网络看到很少的上下文。最近的一些方法[11,4]提出了一种考虑多层特征的分类器输出。好的本地化和上下文的使用是可以同时进行的。

在本文中，我们构建了一个更优雅的架构，即所谓的“全卷积网络”。我们修改和扩展了这个架构，使它工作在非常少的训练图像和产生更精确的分割;参见图1。因此，这些层增加了输出的分辨率。为了进行定位，收缩路径的高分辨率特征与上采样输出相结合。然后，一个连续的卷积层可以学习根据这些信息组装更精确的输出。

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

对于我们的任务，可用的训练数据很少，我们通过对可用的训练图像应用弹性变形来使用过量的数据扩充。这使得网络可以学习这些变形的不变性，而不需要在带注释的图像语料库中查看这些转换。这在生物医学分割中尤为重要，因为变形曾经是组织中最常见的变化，而真实的变形可以有效地模拟。Dosovitskiy等人的[2]研究表明，在无监督特征学习的范围内，数据增强对学习不变性的价值。在许多细胞分割任务中的另一个挑战是分离相同类的触摸对象;参见图3。为此，我们建议使用加权损失，即接触细胞之间的分离背景标签在损失函数中获得较大的权重。

得到的网络适用于各种生物医学分割问题。在这篇论文中，我们展示了在EM栈(一个正在进行的竞争始于ISBI 2012年)中神经元结构的分割结果，我们在这方面的表现优于Ciresan等人的[1]网络。此外，我们在2015年ISBI细胞追踪挑战赛的光镜图像中展示了细胞分割的结果。在这里，我们赢得了两个最具挑战性的2D透射光数据集的巨大优势。

2、网络结构

网络架构如图1所示。它由收缩路径(左侧)和扩展路径(右侧)组成。契约路径遵循典型的卷积网络架构。它包括重复应用两个3x3的卷积(无填充卷积)，每个卷积后面都有一个校正的线性单元(ReLU)和一个2x2最大池操作，其中stride 2用于下行采样。在每个向下采样的步骤中，我们将特征通道的数量增加一倍。每一步的扩张路径由一个upsampling特性映射后跟2 x2卷积(“up-convolution”)部分频道的数量特性,与相应的裁剪功能连接图从缔约路径,和两个3 x3的隆起,每一个ReLU紧随其后。裁剪是必要的，因为在每个卷积中边界像素的损失。在最后一层，使用1x1卷积将每个64分量的特征向量映射到所需的类数。网络总共有23个卷积层。为了允许输出分割映射的无缝平铺(参见图2)，必须选择输入平铺的大小，以便将所有2x2最大池操作应用到具有均匀x和y大小的层。

3、训练

利用Caffe的随机梯度下降实现，利用输入图像及其对应的分割映射训练网络。由于未经填充的卷积，输出图像比输入图像的边界宽度要小一个常数。为了最小化开销并最大限度地利用GPU内存，我们倾向于使用大的输入块而不是大的批处理大小，从而将批处理减少到单个图像。因此，我们使用一个高动量(0.99)，以便大量以前看到的训练样本确定当前优化步骤中的更新。

能量函数的计算采用像素级软极大值法，并结合交叉熵损失函数。softmax的定义为

p_{k}(\mathbf{x})=\exp \left(a_{k}(\mathbf{x})\right) /\left(\sum_{k^{\prime}=1}^{K} \exp \left(a_{k^{\prime}}(\mathbf{x})\right)\right)

，其中a_{k}(x) 代表第k个通道，第x \subset \Omega 个位置的激活，其中\Omega \subset Z^{2} 。K是类的数量，p_{k}(x) 是近似的最大函数。例如，对k，p_{k}(x) \approx 1 有最大的激活a_{k}(x) ，对其他k，p_{k}(x) \approx 1 。交叉熵在每个位置惩罚p_{\ell(\mathbf{x})}(\mathbf{x}) 与1的偏差。

E=\sum_{\mathbf{x} \in \Omega} w(\mathbf{x}) \log \left(p_{\ell(\mathbf{x})}(\mathbf{x})\right)

其中，\iota: \Omega \rightarrow\{1, \ldots, K\} 是每个像素的真实标签，\Omega \rightarrow R 是一个权值图，我们介绍给一些像素在训练中更重要。我们pre-compute每个地面实况的重量地图分割来弥补像素从某个类的不同频率的训练数据集,并迫使网络学习小分离之间的边界,我们引入接触细胞(参见图3 c和d)。分离边界计算使用形态学操作。然后计算权重图

w(\mathbf{x})=w_{c}(\mathbf{x})+w_{0} \cdot \exp \left(-\frac{\left(d_{1}(\mathbf{x})+d_{2}(\mathbf{x})\right)^{2}}{2 \sigma^{2}}\right)

其中w_{c}: \Omega \rightarrow R 是平衡类频率的权重映射，d_{1}: \Omega \rightarrow R 代表最近细胞到边界的距离，d_{2}: \Omega \rightarrow R 代表第二近细胞到边界的距离，在实验中设置为\omega_{0}=10 和\sigma=5 像素。

否则，网络的某些部分可能会提供过多的激活，而其他部分则永远不会提供。理想情况下，初始权值应该调整，使网络中的每个特征图具有近似单位方差。对于我们这种结构的网络(交替的卷积层和ReLU层)，这可以通过从标准偏差为\sqrt{2 / N} 的高斯分布中提取初始权值来实现，其中N表示一个神经元[5]的传入节点数。例如，对于前一层的3x3卷积和64个特征通道，N = 9·64 = 576。

3.1、数据增强

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

4、实验

我们演示了U-Net在三个不同的细分任务中的应用。第一个任务是电子显微记录中神经元结构的分割。图2显示了一个数据集示例和我们获得的分段。我们提供完整的结果作为补充资料。数据集是由EM分割挑战[14]提供的，该挑战始于2012年ISBI会议，目前仍对新的贡献开放。训练数据是一组30张果蝇一龄幼虫腹神经索(VNC)连续切片透射电镜图像(512x512像素)。每个图像都配有相应的完全注释的细胞(白色)和细胞膜(黑色)的地面真相分割地图。测试集是公开的，但其分割地图是保密的。通过将预测的膜概率图发送给组织者，可以得到评价。评估是通过在10个不同的水平上对地图进行阈值化，并计算“扭曲误差”、“Rand误差”和“像素误差”[14]。u-net(平均超过7个旋转版本的输入数据)在没有进一步预处理或后处理的情况下获得了0.0003529的翘曲误差(新的最佳值，见表1)和0.0382的随机误差。

这明显优于Ciresan等人的滑动窗口卷积网络结果，[1]的最佳提交的翘翘误差为0.000420,rand误差为0.0504。就rand误差而言，在这个数据集上唯一性能更好的算法是使用Ciresan等人的概率图上应用的高度数据集特定的后处理方法。我们也将U-Net应用到光镜下的细胞分割任务中。这项分段任务是2014和2015年ISBI细胞跟踪挑战的一部分[10,13]。第一组数据集“PhC-U373”2包含胶质母细胞瘤-星形细胞瘤U373细胞，这些细胞位于聚丙烯酰亚胺基质上，通过相衬显微镜记录(见图4a、b和Supp.材料)。它包含35个部分注释的训练图像。这里我们实现平均借据(“十字路口在联盟”)的92%,这是明显比第二个最好的算法为83%(见表2)。第二个数据集“DIC-HeLa”3海拉细胞在平板玻璃(DIC)微分干涉对比显微镜记录的(参见图3,图4 c, d和增刊。材料)。它包含20个部分注释的训练图像。在这里，我们实现了77.5%的平均IOU，明显好于第二好的算法(46%)。

5、结论

U-Net架构在非常不同的生物医学分割应用上取得了非常好的性能。由于数据增加与弹性变形，它只需要非常少的注释图像，并有一个非常合理的训练时间只有10小时，在NVidia Titan GPU (6GB)。我们提供完整的基于[6]的实现和训练有素的网络4。我们确信u-net体系结构可以很容易地应用于更多的任务。