U-Net网络:Convolutional Networks for Biomedical Image Segmentation

U-Net:卷积网络用于生物医学图像分割

摘要

  人们普遍认为，对深度网络作品的成功训练需要成千上万的注释训练样本。在本文中，我们提出了一种网络和培训策略，该策略依靠大量使用数据增强功能来更有效地使用可用的带注释的样本,该体系结构由捕获上下文的收缩路径和支持精确定位的对称展开路径组成,我们证明，这样的网络可以从很少的图像中得到端到端的训练，并且在电子显微镜堆栈中分割神经结构的ISBI挑战中优于之前最好的方法(滑动窗口卷积网络),使用同样的网络对透射光学显微镜图像进行训练(相位对比和DIC)，我们在2015年的ISBI细胞跟踪挑战赛中以巨大优势胜出.

1.介绍

  在过去的两年中，深度卷积网络在许多视觉识别任务上的表现超过了目前的先进水平,虽然卷积网络已经存在很长时间了，但是由于可用训练集的大小和考虑的网络的大小，其成功程度受到了限制.Krizhevsky等人的突破在于对拥有100万张训练图像的ImageNet数据集进行8层数百万个参数的大型网络的监督训练,从那时起，更大更深入的网络得到了训练.
  卷积网络的典型应用是在分类任务上，图像的输出是一个类标签,然而，在许多视觉任务中，特别是在生物医学图像处理中，期望的输出应该包括定位，也就是说，应该给每个像素分配一个类标签.此外，在生物医学任务中，数以千计的训练图像通常是无法触及的.因此，Ciresan等人在滑动窗口设置中训练网络，通过提供像素周围的局部区域(patch)作为输入，来预测每个像素的类标签.首先，这个网络可以本地化。其次，在patch方面的训练数据远远大于训练图像的数量.
  显然，Ciresan等人的策略有两个缺点。首先，它的速度非常慢，因为每个patch都必须单独运行网络，而且由于重叠的patch存在大量冗余,其次，在定位精度和上下文使用之间需要权衡,较大的patch需要更多的max-pooling层，这会降低定位精度，而较小的patch只允许网络看到很少的上下文。
  在本文中，我们构建了一个更优雅的架构，即所谓的“全卷积网络”。我们修改和扩展了这个架构，这样它就可以在很少的训练图像上工作，并产生更精确的分割,如图1所示，主要思想是通过连续的层来补充一个通常的收缩网络，其中池操作符被上采样操作符所取代，因此，这些层增加了输出的分辨率。为了实现局部化，将收缩路径中的高分辨率特征与上采样的输出结合在一起，接二连三的卷积层可以学习根据这些信息组装出更精确的输出。

图1所示:U-net架构(例如最低分辨率为32x32像素),每个蓝框对应一个多通道特征图.通道的数量显示在盒子的顶部,x-y大小提供在盒子的左下边缘.白盒代表复制的特征图。箭头表示不同的操作

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

  图2:无缝分割任意大图像的重叠策略(这里是EM堆栈中的神经结构分割).预测分割在黄色区域，需要蓝色区域内的图像数据作为输入,缺少的输入数据是通过镜像来推断的
  对于我们的任务，可用的训练数据非常少，我们通过对可用的训练图像应用弹性变形来使用过量的数据增强.这允许网络学习这种变形的不变性，而不需要在带注释的图像语料库中看到这些变换.这在生物医学分割中尤为重要，因为变形是组织中最常见的变异，可以有效地模拟真实的变形.
  在许多细胞分割任务的另一个挑战是分离触摸对象的同一类;参见图3。为此，我们建议使用加权损失，即在接触单元之间分离的背景标签在损失函数中获得很大的权重.
  所得到的网络适用于各种生物医学分割问题.在本文中，我们展示了EM栈中神经元结构分割的结果(在2012年ISBI开始的一项正在进行的竞争)，在那里我们超过了Ciresan的网络.

2.网络架构

  网络架构如图1所示。它由收缩路径(左侧)和扩展路径(右侧)组成。收缩路径遵循卷积网络的典型结构.它由重复应用两个3x3卷积（未填充的卷积）组成，每个卷积后跟一个整流线性单位（ReLU）和一个2x2最大池运算，步长为2用于下采样。在每一个下采样步骤中，我们将特征通道的数量翻倍。扩展路径中的每个步骤都包括对特征图进行上采样，然后进行2x2卷积（向上卷积），将特征通道的数量减半，并与从收缩路径中相应裁剪的特征图进行级联，再进行两次3x3卷积， 每个后跟一个ReLU。由于每次卷积中都会丢失边界像素，因此有必要进行裁剪。 在最后一层，使用1x1卷积将每个64分量特征向量映射到所需的类数。 该网络总共有23个卷积层。为了无缝拼接输出分割图（请参见图2），重要的是选择输入图块大小，以便将所有2x2 max-pooling操作应用于具有x和y大小均等的图层。

3.训练

  利用输入图像及其对应的分割映射，采用Caffe的随机梯度下降实现对网络进行训练.由于没有填充卷积，输出图像比输入图像小一个恒定的边界宽度.为了最大程度地减少开销并最大程度地利用GPU内存，我们建议在大批量时使用较大的输入图块，从而将批量减少为单个图像.因此，我们使用一个高动量(0.99)，以便之前看到的大量训练样本决定当前优化步骤中的更新。
  结合交叉熵损失函数，在最终的特征图上通过像素级的soft-max来计算能量函数。softmax被定义为$p_k(x)=exp(a_x(x))/{\sum }_{k^{\prime }=1}^{K}exp(a_{k^{\prime }}(x))$，其中$a_k(x)$表示在像素位置$x\in \Omega$处的特征通道k的激活，K为类的个数，$p_k(x)$为最大近似函数.当k具有最大激活$a_k(x)$时$p_k(x)\approx 1$，而对于其他k$p_k(x)\approx 0$

其中$l:\Omega \to 1,...,K$表示每个像素的真实标签，且$w:\Omega \to \mathbb{R}$是一个权重图，我们引入它是为了让一些像素在训练中更重要.
  我们预先计算每个ground truth分割的权值图，以补偿训练数据集中某一类像素的不同频率，并迫使网络学习我们在接触细胞之间引入的小的分离边界.
  使用形态学操作计算分离边界。然后将权重映射计算为:

其中$w_c:\Omega \to \mathbb{R}$是一个权衡图来平衡类频率,$d_1:\Omega \to \mathbb{R}$表示到最近的细胞边界的距离且$d_2:\Omega \to \mathbb{R}$表示到第二个最近细胞边界的距离.
  在具有许多卷积层和通过网络的不同路径的深度网络中，良好的权值初始化是极其重要的,否则，网络的某些部分可能会过度激活，而其他部分则永远不会激活.理想情况下，初始权值应该调整为网络中的每个特征图都具有近似单位方差.对于具有我们架构的网络(交替卷积和ReLU层)，这可以通过从一个标准差为$\sqrt{2/N}$的高斯分布中绘制初始权值来实现，其中N表示一个神经元的传入节点数。例如，对于3x3卷积，前一层64个特征通道$N=9\cdot 64=576$