U-Net: Convolutional Networks for BiomedicalImage Segmentation

Abstract

人们普遍认为，深度网络的成功训练需要成千上万的标注训练样本。在本文中，我们提出了一种网络和训练策略，该策略强烈依赖于数据增强，以更有效地利用现有的标注样本。该架构包括一个用于捕获上下文的收缩路径和一个用于实现精确定位的对称扩展路径。我们展示了这样的网络可以通过非常少量的图像进行端到端的训练，并且在ISBI挑战赛中，对于电子显微镜堆栈中神经元结构分割的任务，它的表现优于之前的最佳方法（一个滑动窗口卷积网络）。使用在透射光显微镜图像（相差和微分干涉对比，DIC）上训练的相同网络，我们在2015年ISBI细胞追踪挑战赛的这些类别中以较大优势获胜。此外，该网络速度很快。在一块近期的GPU上，对一个512*512图像进行分割的处理时间不到一秒。

收缩路径：就是常规的卷积网络，它包含重复的2个3x3卷积（无填充卷积），每个跟随ReLU和步长为2的2x2最大池化操作（池化操作即下采样，在每个下采样步骤将特征通道数x2）

拓展路径：每个步骤包含特征图的上采样，然后是2x2卷积（向上卷积，它将特征通道数目减半），与相对应裁剪的收缩路径特征图相连接。和两个3x3卷积，每个跟随ReLU。

1 Introduction

在过去两年中，深度卷积网络在许多视觉识别任务中超越了当时的最先进技术，例如[7,3]。尽管卷积网络已经存在了很长时间[8]，但由于可用训练集的大小和所考虑网络的规模，它们的成功受到了限制。Krizhevsky等人[7]的突破是由于在拥有100万训练图像的ImageNet数据集上，对一个具有8层和数百万参数的大型网络进行了监督训练。从那时起，人们已经训练了更大和更深的网络[12]。

卷积网络的典型用途是在分类任务上，其中对图像的输出是一个单一的类别标签。然而，在许多视觉任务中，特别是在生物医学图像处理领域，期望的输出应该包括定位，即应该为每个像素分配一个类别标签。此外，在生物医学任务中，通常无法获得成千上万的训练图像。因此，Ciresan等人[1]在滑动窗口设置中训练了一个网络，通过提供该像素周围的局部区域（patch）作为输入来预测每个像素的类别标签。首先，这个网络可以进行定位。其次，就patch而言，训练数据比训练图像的数量要大得多。由此产生的网络在ISBI 2012年的EM分割挑战赛中以较大优势获胜。

Ciresan等人[1]的方法存在两个主要缺点。首先，由于网络必须为每个图像块单独运行，并且由于重叠的图像块存在大量冗余，因此该方法速度相当慢。其次，定位精度和上下文使用之间存在权衡。较大的图像块需要更多的最大池化层，这会降低定位精度，而较小的图像块则使网络只能看到很少的上下文。更近期的方法[11,4]提出了一个分类器输出，该输出考虑了来自多个层的特征。这样可以同时实现良好的定位和上下文的使用。此外，U-Net模型通过结合收缩-扩展路径和跳跃连接，能够逐渐适应长距离的亲和力，从而在收缩阶段提取高低级特征，而扩展阶段则利用收缩阶段学习到的特征（通过跳跃连接）来重建图像到所需的尺寸。

在本文中，我们基于一个更为优雅的架构——所谓的“全卷积网络”（fully convolutional network）[9]。我们修改并扩展了这个架构，使其能够在只有很少训练图像的情况下工作，并产生更精确的分割结果；参见图1。[9]中的主要思想是通过连续的层来补充通常的收缩网络，其中池化操作被上采样操作所取代。因此，这些层提高了输出的分辨率。为了实现定位，将收缩路径中的高分辨率特征与上采样输出结合起来。然后，一个连续的卷积层可以学习如何基于这些信息组装一个更精确的输出。

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

重叠-平铺策略（overlap-tile strategy）是一种用于处理大尺寸图像的技术，它特别适用于无法一次性加载到内存中的图像。这种策略通过将大图像分割成多个小块（tiles），然后对每个小块进行处理，最后将结果拼接起来，以实现对整个图像的分割。

对于我们的任务来说，由于可用的训练数据非常有限，我们通过在现有的训练图像上应用弹性变形（elastic deformations）来进行过度的数据增强。这使得网络能够学习对这种变形的不变性，而无需在标注的图像集合中看到这些变换。这在生物医学分割中尤其重要，因为变形通常是组织中最常见的变化，并且可以有效地模拟真实变形。数据增强在学习不变性方面的价值已经在Dosovitskiy等人[2]的研究中，在无监督特征学习的范围中得到了证明。

在许多细胞分割任务中，另一个挑战是分离同类接触物体；见图3。为此，我们提出了使用加权损失函数的方法，其中在接触细胞之间的分离背景标签在损失函数中获得较大的权重。这种方法有助于模型更好地学习区分相邻的同类物体，提高分割精度。通过给予这些分离背景标签更高的权重，模型在训练过程中会更加关注这些难以分割的区域，从而在分割时能够更准确地区分接触的细胞。

所得到的网络可以应用于各种生物医学分割问题。在本文中，我们展示了在电子显微镜（EM）堆栈中分割神经元结构的结果（这是ISBI 2012年开始的一项正在进行的比赛），我们超越了Ciresan等人[1]的网络。此外，我们还展示了2015年ISBI细胞追踪挑战赛中，来自光学显微镜图像的细胞分割结果。在这项挑战赛中，我们在两个最具挑战性的二维透射光数据集上以较大优势获胜。

2 Network Architecture

网络架构如图1所示。它由一个收缩路径（左侧）和一个扩展路径（右侧）组成。收缩路径遵循卷积网络的典型架构。它由两个3x3卷积（无填充卷积）的重复应用组成，每个卷积后面跟着一个修正线性单元（ReLU）和一个2x2最大池化操作，步幅为2，用于下采样。在每个下采样步骤中，我们将特征通道的数量翻倍。扩展路径的每一步都包括特征图的上采样，随后是一个2x2卷积（“上卷积”），该卷积将特征通道的数量减半，与收缩路径中相应裁剪的特征图进行拼接，然后是两个3x3卷积，每个卷积后面跟着一个ReLU。由于每次卷积都会丢失边界像素，因此需要进行裁剪。在最后一层，使用1x1卷积将每个64维特征向量映射到所需的类别数。整个网络总共有23个卷积层。

为了实现输出分割图的无缝拼接（参见图2），选择输入平铺尺寸非常重要，以确保所有的2x2最大池化操作都应用于x和y尺寸均为偶数的层。这是因为池化操作会减少特征图的尺寸，如果特征图的尺寸是奇数，那么在经过池化后，特征图的尺寸将不再是对称的，这可能导致在拼接时出现不连续或错位的问题。

左半边：

第一块内容：572-3+1，570-3+1

# 由572*572*1变成了570*570*64
self.conv1_1 = nn.Conv2d(in_channels=1, out_channels=64, kernel_size=3, stride=1, padding=0)
self.relu1_1 = nn.ReLU(inplace=True)
# 由570*570*64变成了568*568*64
self.conv1_2 = nn.Conv2d(64, 64, kernel_size=3, stride=1, padding=0)  
self.relu1_2 = nn.ReLU(inplace=True)

最大池化层：卷积核大小是2*2，步长是2，这样图片大小就可以减半

# 采用最大池化进行下采样，图片大小减半，通道数不变，由568*568*64变成284*284*64
self.maxpool_1 = nn.MaxPool2d(kernel_size=2, stride=2)  

以下的内容与第一块儿内容类似。

右半边：右边部分的卷积层也有四个大层，每个大层经过两个卷积层。

# 接下来实现上采样中的up-conv2*2
self.up_conv_1 = nn.ConvTranspose2d(in_channels=1024, out_channels=512, kernel_size=2, stride=2, padding=0) # 28*28*1024->56*56*512

每个大层经过两个卷积层：

self.conv6_1 = nn.Conv2d(in_channels=1024, out_channels=512, kernel_size=3, stride=1, padding=0)  # 56*56*1024->54*54*512
self.relu6_1 = nn.ReLU(inplace=True)
self.conv6_2 = nn.Conv2d(512, 512, kernel_size=3, stride=1, padding=0)  # 54*54*512->52*52*512
self.relu6_2 = nn.ReLU(inplace=True)

3 Training

输入图像及其相应的分割图用于使用Caffe[6]的随机梯度下降实现来训练网络。由于未填充的卷积操作，输出图像比输入图像小一个恒定的边界宽度。为了最小化开销并最大限度地利用GPU内存，我们倾向于使用大尺寸的输入平铺而不是大批量大小，因此将批量减少到单张图像。相应地，我们使用高动量（0.99），这样先前看到的大量训练样本就能决定当前优化步骤的更新。

计算能量函数的方法，结合像素级的softmax和交叉损失函数，用于图像分割任务中的像素分类。

我们为每个真实分割图预先计算权重图，以补偿训练数据集中某个类别像素的不同频率，并迫使网络学习我们引入的接触细胞之间的小分离边界（参见图3c和d）。分离边界是使用形态学操作计算的。然后计算权重图如下：

在具有许多卷积层和不同网络路径的深度网络中，权重的良好初始化极其重要。否则，网络的某些部分可能会产生过大的激活值，而其他部分则永远不会有贡献。理想情况下，初始权重应该被调整，以便网络中的每个特征图具有大约单位方差。对于具有我们这种架构（交替卷积和ReLU层）的网络，可以通过从一个高斯分布中抽取初始权重来实现这一点，该高斯分布的标准差为 ，其中 N表示一个神经元的传入节点数[5]。例如，对于一个3x3的卷积和前一层的64个特征通道，N=9 * 64=576。

3.1 Data Augmentation

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

4 Experiments

我们展示了U-Net在三个不同的分割任务中的应用。第一个任务是电子显微镜记录中神经元结构的分割。图2展示了数据集和我们获得的分割结果。我们将完整的结果作为补充材料提供。数据集由ISBI 2012年开始的EM分割挑战赛[14]提供，并且仍然对新的贡献开放。训练数据是一组30张图像（512x512像素），这些图像来自果蝇第一龄幼虫腹神经索（VNC）的连续切片透射电子显微镜。每张图像都配有一个对应的完全标注的真实分割图，用于细胞（白色）和膜（黑色）。测试集是公开可用的，但其分割图是保密的。通过将预测的膜概率图发送给组织者，可以获得评估结果。评估是通过在10个不同的阈值水平上对图进行阈值处理，并计算“变形误差”（warping error）、“兰德误差”（Rand error）和“像素误差”（pixel error）[14]来完成的。

U-Net模型（对输入数据的7个旋转版本取平均）在没有任何进一步的预处理或后处理的情况下，达到了0.0003529的变形误差（这是新的最佳分数，见表1）和0.0382的兰德误差（rand-error）。这表明该模型在处理电子显微镜记录中的神经元结构分割任务时具有很高的准确性和鲁棒性。

这明显优于Ciresan等人[1]的滑动窗口卷积网络结果，他们的最佳提交有0.000420的变形误差和0.0504的Rand误差。就Rand误差而言，在该数据集上唯一表现更好的算法，其使用了应用到Ciresan等人[1]概率图上的针对数据集的非常特定后处理方法。

我们还把U-Net应用于光学显微镜图像中的细胞分割任务。这个分割任务是ISBI细胞追踪挑战赛2014年和2015年[10,13]的一部分。第一个数据集“PhC-U373”包含在聚丙烯酰胺基质上记录的胶质母细胞瘤星形细胞U373细胞，使用的是相差显微镜（参见图4a、b和补充材料）。它包含35张部分标注的训练图像。在这里，我们实现了平均92%的交并比（Intersection over Union，IOU），这显著优于第二好的算法，后者为83%（见表2）。第二个数据集“DIC-HeLa”包含在平坦玻璃上记录的HeLa细胞，使用的是微分干涉对比（DIC）显微镜（参见图3、图4c、d和补充材料）。它包含20张部分标注的训练图像。在这里，我们实现了平均77.5%的IOU，这显著优于第二好的算法，后者为46%。

5 Conclusion

U-Net架构在非常不同的生物医学分割应用中取得了非常好的性能。由于使用了弹性变形的数据增强，它只需要非常少的标注图像，并且在NVidia Titan GPU（6 GB）上只需要非常合理的10小时训练时间。我们提供了完整的基于Caffe[6]的实现和训练好的网络[4]。我们确信U-Net架构可以轻松应用于更多的任务。