论文心得U-Net: Convolutional Networks for Biomedical Image Segmentation

学习心得 U-Net: Convolutional Networks for Biomedical Image Segmentation

原文链接

一、翻译

摘要

abstract：
　　人们普遍同意，成功地训练深度网络需要大量已标注的训练样本。在本文中，我们提出了一种网络和训练策略。为了更有效的利用标注数据，我们使用数据扩张的方法(data augmentation)。该体系结构包括两部分：用于捕获上下文的收缩路径（a contracting path）和用于进行精确定位的对称扩展路径（a symmetric expanding path）。我们展示了这样的网络可以在很少的图像上进行端到端训练，并且在ISBI挑战方面优于现有的最佳方法。我们使用这个网络获得了赢得了ISBI cell tracking challenge 2015。不仅如此，这个网络非常的快，对一个512*512的图像，使用一块GPU只需要不到一秒的时间。
　　网络获取地址(基于caffe)： http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net

1. 介绍

在过去的两年中，深度卷积网络在许多视觉识别任务中的表现超越了现有技术。 虽然卷积网络已经存在很长时间了，但是由于可用训练集的大小和网络结构的大小，它们的成功受到限制。 Krizhevsky等人创造了很大的突破，他们对具有8层的大型网络和具有100万个训练图像的ImageNet数据集上的数百万个参数进行了监督训练。从那时起，甚至更大更深的网络也得到了训练。
　　卷积网络的典型用途是用于分类任务，其中图像的输出是单个类别。然而，在许多视觉任务中，特别是在生物医学图像处理中，期望的输出应包括位置，所以应该给每个像素都进行标注。然而在生物医学任务中通常无法获得数千个训练图像。因此，Ciresan等人提出用滑动窗口训练网络，通过提供像素周围的局部区域（patch——每个patch包含很多pixel）作为输入，预测每个像素的类标签。首先，这个网络可以完成定位工作。第二，patch的训练数据远大于训练图像的数量。最终网络赢得了ISBI 2012的胜利。
　　显然，Ciresan等人的策略有两个缺点。首先要分别预测每一个patch的类别，patch之间的重叠会导致大量的冗余。其次这个网络需要在局部准确性和获取整体上下文信息之间平衡，更大的patches需多的最大池化层，但是会降低准确率，小的patches仅允许网络查看很少的上下文信息。最近有方法提出了一种分类器输出，该输出考虑了来自多层的特征，既有准确地定位，又包含了上下文信息。
　　在这篇文章中，我们建立了一个更加优雅的框架，通常被称为“全卷积网络”（fully convolutional network）。我们修改并拓展了这个框架，使其可以仅使用少量训练图片就可以工作，获得更高的分割准确率。网络如下图所示：

　　“全卷积网络”（fully convolutional network）的核心思想是修改一个普通的逐层收缩的网络，用上采样 (up sampling)操作代替网络后部的池化(pooling)操作。因此，这些层增加了输出的分辨率。为了精准定位，在网络收缩过程（路径）中产生的高分辨率特征(high resolution features) ，被连接到了修改后网络的上采样的结果上。在此之后，连续的卷积层基于这些综合信息得到更精确的结果。
　　我们架构的重大改进是，在上采样部分中，我们还拥有大量特征通道，这些特征通道使网络可以将上下文信息传播到更高分辨率的层。结果，扩展路径或多或少地相对于收缩路径对称，并且产生u形结构。 网络不存在任何全连接层(fully connected layers)，并且，只使用每个卷积的有效部分，例如，分割图(segmentation map)只包含这样一些像素点，这些像素点的完整上下文都出现在输入图像中。该策略允许通过重叠-平铺策略对任意大图像进行无缝分割，如下图。（用于任意大图像的无缝分割的重叠拼贴策略。 对黄色区域中的分割的预测需要蓝色区域内的图像数据作为输入。 丢失的输入数据通过镜像推断。）

　　为了预测图像边界区域中的像素，可通过镜像输入图像来推断缺失的上下文。 这种平铺策略对于将网络应用于大图像非常重要，因为否则分辨率会受到GPU内存的限制。
　　至于我们的任务，几乎没有可用的训练数据，我们通过对可用的训练图像应用弹性变形来使用过多的数据增强。这允许网络学习此类变形的不变性，而无需在带注释的图像语料库中查看这些转换。这对医学图像分割是非常重要的，因为组织的形变是非常常见的情况，并且计算机可以很有效的模拟真实的形变。在无监督特征学习的范围内，Dosovitskiy等人已经证明了数据增强对于学习不变性的价值。
　　在细胞分割任务中的另一个挑战是分离同一类别的接触物体。 参见图3。为此，我们使用了加权损失，这些位于touching cells之间的背景在损失函数的权重很高。如图所示：

　　（a）原始图像
　　（b）标注图像实况分割 不同的颜色表示HeLa细胞的不同情况
　　（c）生成分割蒙版（白色：前景，黑色：背景）
　　（d）以像素为单位的权重映射，迫使网络学习边界像素
　　本文提出的网络适用于各种生物医学分割问题。 在本文中，我们展示了有关EM stacks中神经元结构分割的结果（一场持续的竞争始于ISBI 2012），在此方面我们胜过了Ciresan等人的网络。 此外，我们在2015年ISBI cell tracking challenge的光学显微镜图像数据集中中显示了细胞分割的结果。我们在两个最具挑战性的2D数据集上取得了很好的效果。

2. 网络结构

网络架构如图所示。它由一个contracting path（收缩路径） 和 expansive path（扩展路径）组成。 收缩路径遵循卷积网络的典型架构。它包括了重复单元：2个3 * 3卷积层（unpadding）、ReLU**函数和一个2 * 2的步长为2的max pooling层。每一次下采样后我们都把特征通道的数量加倍。扩展路径中的每一步都首先使用反卷积(up-convolution)，每次使用反卷积都将特征通道数量减半，特征图大小加倍。反卷积过后，将反卷积的结果与收缩路径中对应步骤的特征图拼接起来，跟随2个3 * 3卷积层（unpadding）、ReLU**函数。由于每次卷积中都会丢失边界像素，因此有必要进行裁剪。 在最后一层，使用1x1卷积将每个64分量特征向量映射到所需的类数。 该网络总共有23个卷积层（18个顺序卷积+4个收缩到扩展的卷积+1个1*1的卷积）。

为了无缝拼接输出分割图，选择输入图块大小非常重要，以保证所有的Max Pooling操作作用于长宽为偶数的feature map。

3. 训练

我们采用随机梯度下降法训练，基于caffe框架。为了最大限度的使用GPU显存，比起输入一个大的batch size，我们更倾向于大量输入tiles，因此实验batch size为1。我们使用了很高的momentum（0.99），大量先前的训练样本确定了当前优化步骤中的更新。损失函数就是pixel-wise softmax + cross_entropy
　　softmax函数，$\ a_k$ ak​代表$\ k$ k通道，$\ x$ x像素位置，$\ K$ K代表类别的数量
$\ p_k(x) = \exp(a_k(x))/(\sum_{k'=1}^K \exp(a_k'(x))$ pk​(x)=exp(ak​(x))/(∑k′=1K​exp(ak′​(x))
　　 交叉熵损失函数，$\ \iota：\Omega\rightarrow({1,2,3...K})$ ι：Ω→(1,2,3...K)代表true label，$\ w$ w是权重
$\ E= \sum\limits_{i\in\Omega}w(x)\log(p_{\iota (x)}(x)$ E=i∈Ω∑​w(x)log(pι(x)​(x)
　　为了使某些像素点更加重要，我们在公式中引入了$\ w(x)$ w(x)。我们对每一张标注图像预计算了一个权重图，来补偿训练集中每类像素的不同频率，使网络更注重学习相互接触的细胞之间的小的分割边界。我们使用形态学操作计算分割边界。权重图计算公式如下：
$\ w(x)=w_c(x)+w_0*\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$ wc​是用于平衡类别频率的权重图，$\ d_1$ d1​代表到最近细胞的边界的距离，$\ d_2$ d2​代表到第二近的细胞的边界的距离。基于经验我们设定$\ w_0$ w0​=10，$\sigma=5$σ=5像素。
　　在具有许多卷积层和通过网络的不同路径的深度网络中，权重的良好初始化非常重要。 否则，网络的某些部分可能会进行过多的**，而其他部分则永远不会起作用。 理想情况下，应调整初始权重，以使网络中的每个特征图都具有大约单位方差。 对于具有我们（交替卷积和ReLU层）的网络结构，这可以通过从具有标准偏差$\ \sqrt {N/2}$ N/2​的高斯分布中提取初始权重来实现，其中N表示一个神经元的传入节点数。 例如。 对于上一层的3x3卷积和64个特征通道，N = 9·64 = 576。

3.1 数据扩充

在只有少量样本的情况况下，要想尽可能的让网络获得不变性和鲁棒性，数据增加是很重要的。因为本论文需要处理显微镜图片，我们需要平移与旋转不变性，并且对形变和灰度变化鲁棒。将训练样本进行随机弹性变形，是训练带有很少标注图像的分割网络的关键。 我们使用随机位移矢量在粗糙的3*3网格上(random displacement vectors on a coarse 3 by 3 grid)产生平滑形变(smooth deformations)。 位移是从10像素标准偏差的高斯分布中采样的。然后使用双三次插值计算每个像素的位移。在contracting path的末尾采用drop-out 层更进一步增加数据。（不懂）

4. 实验

我们演示了unet在三个不同分割任务上的应用，第一项任务是在电子显微镜记录中分割神经元结构，在下图中演示了数据集中的一个例子和我们的分割结果。我们提供了全部结果作为补充材料。数据集是EM分割挑战提供的，这个挑战是从 ISBI 2012开始的，现在依旧开放。训练数据是一组来自果蝇幼虫腹侧腹侧神经索（VNC）的连续切片透射电镜的30张图像（512x512像素）。每个图像都带有一个对应的标注分割图，细胞（白色）和膜（黑色）。测试集是公开可用的，但对应的标注图是保密的，可以通过将预测的膜概率图发送给组织者来获得评估。通过在10个不同级别对结果进行阈值化和计算the “warping error”, the “Rand error” and the “pixel error”（预测的label和实际的label）
　　u-net（输入数据的7个旋转版本的平均值）无需进行任何进一步的预处理或后处理即可获得0.0003529的warping error和Rand error为0.0382。

5. 结论

u-net体系结构在截然不同的生物医学细分应用中实现了非常好的性能。 由于具有弹性变形的数据增强，它仅需要很少的带注释的图像，并且在NVidia Titan GPU（6 GB）上只有10小时的非常合理的训练时间。 我们提供完整的基于Caffe的训练网络。 我们确信u-net架构可以轻松地应用于更多任务。