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

转自：

http://baijiahao.baidu.com/s?id=1600400677130320639&wfr=spider&for=pc

摘要

众所周知，成功地训练一个深度网络需要大量的已标注数据集。本文提出了一种新的网络结构和训练方法，依靠数据增强将有效的标注数据更为有效地使用。该网络由收缩路径（contracting path）和扩张路径（expanding path）组成。其中，收缩路径用于获取上下文信息（context），扩张路径用于精确的定位（localization），且两条路径相互对称。该网络能够从极少图像端对端进行训练，并且在ISBI竞赛中，对于分割电子显微镜中的神经元结构的表现好于以前最好的方法（滑动窗口卷积网络）。除此之外，该网络的运行速度非常快，在最近的一个GPU上，分割一张512*512的图像不到一秒即可完成。

引言

在过去的两年中，深度卷积网络在许多视觉识别任务中表现良好。虽然卷积网络已经存在了很长时间，但其受限于训练集和网络结构的大小。直到Krizhevsky等人在ImageNet数据集上训练了一个具有8层和数百万个参数的大型网络，人们才开始训练更大更深度的网络结构。

卷积网络常用于分类任务，其中图像的输出为单一的类别标签。然而，在许多视觉任务中，尤其是在生物医学图像处理中，所期待的输出结果中应包括如位置和每一像素的类别标签等信息。而且，在生物医学任务中通常难以获取数以千计的训练图像。因此，Ciresan等人在滑动窗口设置中训练网络，以某一像素的邻域（patch）作为输入，用于预测每个像素的类别标签。首先，该网络能够实现定位操作；其次，对像素的邻域而言，训练数据的数量需远大于训练图像的数量。在2012年的ISBI的EM分割竞赛中，该网络成功摘夺桂冠。

上述网络中有两点明显的不足：

该网络运行效率很慢。对于每个邻域，网络都要运行一次，且对于邻域重叠部分，网络会进行重复运算；网络需要在精确的定位和获取上下文信息之间进行权衡。越大的patch需要越多的最大池化层，其会降低定位的精确度，而小的邻域使得网络获取较少的上下文信息。

最近有人提出了一种分类器输出的方法，其考虑了来自多层的特征。这使得同时做到精确的定位和上下文的使用成为可能。

本文将建立一个更为优雅的网络结构，即全卷积网络。基于该网络进行相应的修改和扩展，使得网络能够在较少的训练图像的情况下运行，并做出更为精确的分割操作。如下图所示，其展现了网络的整体结构。在 Fully Convolutional Networks for Semantic Segmentation 一文中，其主要思想为通过连续层补充通常的收缩网络，其中上采样（upsample）操作取代了池化操作。因此，这些层提高了输出结果的分别率。为了进行定位操作，将来自收缩路径的高分辨率特征与上采样相结合。随后连续的卷积层可以基于该信息学习，得出更为精确的输出。

U-net网络结构图

图一：上图为U-net网络结构图（以最低分别率为32*32为例）。每个蓝色框对应一个多通道特征图（map），其中通道数在框顶标，x-y的大小位于框的左下角；白色框表示复制的特征图；箭头表示不同的操作。

在该网络中最重要的修改是在上采样部分有大量的特征通道，其允许网络将上下文信息传播至更高分辨率的层。因而，扩张路径或多或少与收缩路径对称，从而形成一个U型结构。该网络没有全连接层，且仅使用每个卷积的有效部分，例如：分割图仅包含在输入图像中可获得完整上下文信息的像素。该策略允许通过重叠-切片（overlap-tile）策略对任意大的图像进行无缝分割，如下图所示。为了预测图像边缘区域中的像素的类别标签，通过对输入图像进行镜像操作来推断缺失的上下文信息。这种切片策略对于将网络应用于大型图像是重要的，否则图像的分辨率将受限于显存（the GPU memory）。

重叠-切片策略示意图

图二：上图为重叠-切片策略示意图（此处为EM stacks中神经元结构的分割）。预测黄色方框中的分割需要蓝色方框中的图像数据作为输入，缺失的数据由镜像推断。

对于只要少量训练数据可用的任务，对可用的训练图像广泛地应用弹性形变的图像增强方法。这使得网络能够学习这种形变的不变性，无需提供带注释的图像语料库（the annotated image corpus）供网络学习。这对于生物医学图像的分割是特别重要的，因为形变是组织中最常见的变化，且实际的形变能有效地模拟。 Dosovitskiy等人在无监督学习中表明了学习不变性的重要性。

在许多细胞分割任务中的另外一个挑战为对同一类别的接触目标的分离，如下图所示。为此，采用加权损失，即在细胞接触之间的分离背景标签中的损失函数中设置较大大的权重。

DIC显微镜记录的HeLa细胞

图三：上图为DIC（微分干涉，differential interference contrast）显微镜记录的HeLa细胞。图a为原始图像；图b为与真实的带标记的分割数据相叠加，其中不同的颜色表示HeLa细胞的不同实例（instance）；图c为生成的分割掩膜，其中白色区域表示前景，黑色区域表示背景；图d为用于网络学习边缘像素的带加权损失的像素图。

U-net网络模型适用于各种生物医学图像分割问题。本文将展示在EM stacks竞赛中对神经元结构分割的结果，以及2015年ISBI的光学显微镜图像中的细胞分割结果。

网络结构

网络结构如图一所示。U-net网络由一个收缩路径（左边）和一个扩张路径（右边）组成。其中，收缩路径遵循典型的卷积网络结构，其由两个重复的3*3卷积核（无填充卷积，unpadded convolution）组成，且均使用修正线性单元（rectified linear unit，ReLU）激活函数和一个用于下采样（downsample）的步长为2的2*2最大池化操作，以及在每一个下采样的步骤中，特征通道数量都加倍。在扩张路径中，每一步都包含对特征图进行上采样（upsample）；然后用2*2的卷积核进行卷积运算（上卷积，up-convolution），用于减少一半的特征通道数量；接着级联收缩路径中相应的裁剪后的特征图；再用两个3*3的卷积核进行卷积运算，且均使用ReLU激活函数。由于在每次卷积操作中，边界像素存在缺失问题，因此有必要对特征图进行裁剪。在最后一层，利用1*1的卷积核进行卷积运算，将每个64维的特征向量映射网络的输出层。总而言之，该网络有23个卷积层。

为了允许输出图的无缝拼接，如图二所示，以及所有2*2最大池化操作能应用于x、y大小相等（even）的层，输入切片（tile）大小的选择是很重要的。

训练

输入图像和其相应的分割图用于训练由基于Caffe的随机梯度下降算法实现的网络。由于未填充的卷积其输出图像的大小小于恒定边界宽度的输入。为了最小化开销和最大限度地利用显存，比起输入一个较大的批次，更倾向于较大的输入切片，因而将批次大小缩减为单张图像。相应地，通过使用高动量（high momentum，0.99）使得大量先前的训练样本在当前的优化步骤中更新。

能量函数（the energy function）是通过与交叉熵损失函数相结合的最终特征图，并利用像素级的soft-max函数来计算的。soft-max函数定义为：

其中，a k (x)表示在第x像素点上第k个特征通道的激活函数。K为类别数，p k (x)为近视最大函数（the approximated maximum-function）。

加权交叉熵定义为：

其中，l={1,...,K}表示每个像素正确的标签，权重函数w，且将权重函数定义为：

其中，w c 表示平衡类别频率的权重图，d 1 表示到最近细胞边界的距离，d 2 表示到第二近细胞边界的距离。基于经验，将w 0 设为10，σ约为5像素。

对于分割边界操作，使用了形态学方法。

权重初始化基于标准差为√(2/N)的高斯分布，其中N表示每个神经元的输入结点数。

数据增强

在3*3的网格上使用随机位移矢量产生平滑形变，其中位移来自于10像素标准差的高斯分布，且通过双三次插值法计算得出。在收缩路径的末尾的drop-out层进一步暗示了数据增强。