U-Net论文笔记

本文是关于U-Net论文《U-Net: Convolutional Networks for Biomedical Image Segmentation》的一些学习笔记，有相当一部分是论文的翻译，也有一些自己的理解。尽管网上同质文章已经不少了，但是有些东西还是要自己亲自总结一下落实到笔头上，印象才能更深刻吧，高手看了还请勿见笑。我也会随着学习的深入长期更新此文，加入一些新的东西，争取写的更加详细。

简介

近年来，深度卷积神经网络已经在很多视觉识别领域表现不俗。然而自卷积神经网络诞生以来，训练集以及网络大小的局限性，一直限制了它的进一步发展。取得突破的是Krizkevsky，在ImageNet上训练了一个具有8层几百万个参数的全监督学习网络，自此之后，人们开始倾向于训练更大层数更深的网络。
卷积神经网络的典型应用就是分类问题，根据一张图片输出一个分类标签。然而在很多视觉任务尤其是医学图像处理领域，我们希望输出的结果包含位置信息，比如某些像素级的分类任务。然而动辄几千张图片的训练集在实际的生物医学工程中往往不容易得到。因此，Ciresan等人训练了一个带有滑动窗口的网络，将每个像素周围的局部区域（patch）作为输入来对像素进行标签分类。首先，这个网络具有定位功能，其次，patch来作为训练数据显然要比图像要多得多。该网络也以较大优势赢得了ISBI 2012的分割挑战赛。

很明显Ciresan等人的方法具有两个缺陷：第一，速度很慢，因为对于每个patch它都要训练一次，并且patch之间会存在一些重叠冗余部分。两个patch之间如果重叠太多，这一部分就会被训练两次，浪费计算资源，甚至可能出现过拟合。第二，定位的准确度和上下文之间二者往往不可兼得。选择大一点的patch，需要更多的最大池化层来对其进行降维，每一层降维都会丢失一些信息，这样就会降低定位准确度。后来的一些方法将多层次的特征信息融合进去，使得二者之间的平衡成为可能。
文中构建了一种更加精确的“全卷积网络（Full Convolution Network, FCN）”方法。首先修改和扩展了网络结构，使得只需很少的训练图像就能产生精确的分割，如图1。

全卷积网络的主要思想就是增加了一些连续层的收缩（contracting）网络，并把池化操作换成上采样。因此这些网络层可以增加输出结果的分辨率。为了定位，在收缩路径（contracting path）上将高分辨率特征和上采样结合起来。基于这些信息，一个连续的卷积层可以学习并得到更精确的输出。

文中给出的网络结构所做的一个重要的修改就是在上采样部分有很多特征通道，使得网络可以将上下文信息传递到高分辨率的层。网络的扩张路径（expansive path）和收缩路径或多或少有对称的部分，形成一个U形的结构。整个网络没有全连接层，并且只使用每个卷积的有效部分。举例来说，分割图仅包含那些具有完整上下文的像素，这种over-tile的策略可以实现对任意大的图片进行无缝分割。对于上下文信息不完整的图像边界区域，可以通过镜像外插值得到，如图2。采用这种策略，可以使得我们可以不受GPU内存的限制处理一些较大的图片。

在实际处理任务中我们往往手里的训练数据十分有限，因此我们要进行数据增强。通过对训练数据应用弹性形变，使得网络能够学习到对这些形变的不变性。这一点其实很重要，因为在生物医学图像分割中，这种弹性形变其实比较常见，并且易于模拟。这种通过数据增强来学习形变不变性的方法可以参考Dosovitskiy的关于无监督特征学习的工作。

细胞分割的另一项挑战就是分割相互接触的同一类物体。关于这一点，我们的建议是使用加权损失函数，通过使得相互连接的细胞间的背景标签取得比较大的权重来达到分割的目的。

通过上面提到的这一系列的改进最后得到的网络可以应对各种生物医学图像分割的问题。文中展示了在ISBI2012和ISBI2015的数据集上取得的优异的实验结果。

网络结构

整个网络的架构示意图如图1所示。它是由收缩路径（左边）和扩张路径（右边）组成。收缩路径是典型的卷积网络结构，由两个3×3的卷积层组成，每一个卷积层后面是一个ReLU单元和2×2的最大池化层作为下采样，池化层的步长是2。在下采样的每一步都把特征通道的数量增加一倍。扩张路径的每一步都是由特征图的2×2“上卷积”来进行上采样操作，每一步上采样操作中特征通道的数量减半，再和收缩路径对应的裁剪后的特征图串联起来，经过两个3×3的卷积，每个卷积后面紧跟一个ReLU单元。由于每次卷积都会丢失一些边界信息，因此裁切是必须的。最后一层用一个1×1的卷积核将每64个特征向量映射到需要分类的类数。整个网络结构共含有23个卷积层。
为了保证输出分割图的无缝拼接（如图2），选择合适的输入图像的尺寸很重要，使得所有的2×2最大池化操作能够应用到x和y为偶数大小的层上去。

训练

训练数据是输入图像以及对应的分割图，训练方法是随机梯度下降法（SGD）。由于对卷积不进行补0，输出图像会比输入图像小一个边界的宽度。为了降低消耗同时最大化利用GPU，我们倾向于对大的batch用大的输入tile，因此降低了单个图像的batch。相对应的我们使用了几个比较高的动量（momentum）0.99，这样大量的训练数据将决定当前优化步骤的更新。
能量函数是由最后的特征图结合交叉熵，用Softmax计算得出。Softmax的定义为：

其中是特征通道k在像素点x的得分数。K代表分类数，最后的交叉熵能量函数为：

其中L(x)是每个像素的标签，w是权重图。

首先我们根据正确的分割结果（ground truth）来预计算一下权重图，然后让网络来学习一下分割的边界部分，也就是之前提到的相互接触的细胞之间的背景部分（见图3c和d）。

分割的边界我们可以用形态学的方法计算得到。权重图的计算公式为：

其中Wc是权重图，d1是到最邻近细胞的边界距离，d2是到第二邻近细胞的边界的距离。w0设为10，σ设为5个像素。
对于一个拥有很多卷积层和多通道的深度网络来说，权重系数的合理初始化是非常重要的。否则的话，网络中某些地方可能会取得过高的权重，而其他部分几乎贡献为0。文中使用高斯分布来初始化权重，标准差为sqrt(N/2)，N是一个神经元输入结点的数量，比如说具有64个特征通道的3×3的卷积层，就等于9×64，即576。

数据增强

数据增强的主要作用在于训练数据有限的情况下，通过在某方面增强数据增大训练样本使得网络来学习我们的图像处理任务所要求的不变性和鲁棒性。比如对于显微镜图像来说，需要大小和旋转不变性，以及对形变的鲁棒性。当训练一个样本标记数据比较少的网络时，在样本中增加任意的弹性形变来增大训练样本使得网络去适应这些形变。文中使用3*3 网格上的随机位移向量来进行平滑弹性变换。这些位移从十像素标准差的随机高斯分布中获得。然后，逐像素的位移用双三线性差值得到。收缩的最后一层用drop out 层也可以达到数据增强的作用。