U-Net阅读笔记

U-Net: Convolutional Networks for Biomedical Image Segmentation

提出时间：2015年

引用次数：18470次

摘要

这篇文章提出了一种简单高效的网络结构和一种通过数据增强来提升训练数据利用效率的训练策略。网络结构可以视为由两条路径构成，一条为压缩路径（指的是图像的大小不但变小），用来提取特征，捕获上下文信息；还有一条为对称的扩张路径（图像大小不断变大）。使用这种网络结构和训练策略，能够在很少的图片数据上实现端对端训练。本文的方法超过了ISBI电镜神经元结构分割挑战赛上的最好结果。并且使用相同的网络结构，在透射光显微镜图像和以相当大的优势赢得了2015年ISBI细胞追踪挑战赛。并且，这种网络结构非常高效，在一块GPU上处理一张$512\times 512$大小的图片耗费的时间少于1秒。基于Caffe实现的代码和预训练模型可以在这个网址上找到：http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net。

简介

卷积神经网络的经典应用场景是分类任务，但是对于很多视觉问题，特别是生物医学图像处理，希望输出的内容中包含位置信息，通过分割的方法，能够给每个空间像素分配一个类别信息，是一种比较好的输出。同时，在生物医学任务中，获取成千上万的训练图像也是不现实的。因此这篇文章核心要达到的目标是：使用较少的训练数据，获得精确的图像分割结果。上图为网络的整体架构。

U-Net使用的架构是FCN中提出的全卷积结构，并且对这种结构进行了修改，使其能够在更少的训练数据上得到更好的分割效果。FCN中的主要思想是在特征提取后添加一些层来增强定位信息，提高输出特征图的分辨率，在这些层中使用了上采样操作替代池化操作。来自压缩路径（特征提取路径）中的高分辨率信息与上采样的输出进行结合，然后使用一系列卷积层来学习如何将这些信息结合起来得到更精确的输出。

U-Net与FCN的一个重要区别是其在上采样通道中依然有很大量的特征通道（U-Net中使用了拼接操作concatenation，而FCN中使用的是加和），这样就使得网络能够将上下文信息传递到高分辨率的层。因此，这个扩张路径就看上去与压缩路径相比显得很对称，就像是一种U形结构。

网络不使用任何全连接层，并且只使用每个卷积中有效的部分，比如说分割图中只包含那些完整的上下文都在输入图像中的像素，而对于边界上的像素，使用镜像操作，即用输入图像对应镜像区域的内容进行填补。通过一种层叠铺砖的方法，能够允许对任意大小的输入图像进行无缝分割。这种铺砖tiling（个人认为类似于裁剪）技术对于在很大的图像数据上进行训练非常重要，因为图像的像素是受GPU显存限制的。

在这篇文章的工作中，由于训练数据很少，广泛使用了弹性形变这种数据增强技术，能够使得网络学习到对这种形变的不变性。这种增强方式对分割问题很有帮助。

在细胞分割问题中的另外一个挑战是相同类别的细胞之间相互接触，很难区分开来。因此在这篇文章中提出了使用一种权重损失函数，用来分割相互接触的细胞。在相互接触的细胞中间分隔开的背景标签会在损失函数中得到一个较大的权重。

网络结构

图一展示了网络的整体结构。由压缩路径（左侧）和扩张路径（右侧）构成，压缩路径是一个经典的卷积结构，重复使用$3\times 3$卷积（无padding）之后加上一个ReLU和一个stride=2的$2\times 2$最大池化操作（没有BN层），在每个下采样步骤中，都将特征图的通道数乘2。在扩张路径中，包含了许多$2\times 2$大小的up-convolution操作，这个操作将通道数减半，并且与来自压缩路径中裁剪过来的对应的特征图进行拼接。之后是两个$3\times 3$的卷积，每个都跟着一个ReLU激活函数。裁剪操作是很有必要的，因为在每次卷积中都损失了边界的像素。最后一层是一个$1\times 1$的卷积用来将64个特征向量映射到对应的类别数量（class_number）。网络整体有23个卷积层。

要能够产生无缝的输出分割图，选择合适的tile size非常重要，这样可以使得所有的$2\times 2$最大池化操作被应用到x和y大小都是偶数的特征图上。

训练

训练使用了随机梯度下降SGD。为了充分利用GPU显存资源，因此设置了更大的tile size而不是更大的batch size，因此这里的batch size设置为1。动量0.99。损失函数使用逐像素的SoftMax+CrossEntropy。SoftMax定义如下：

$$p_k=\exp (a_k(\mathrm{x}))/(\sum _{k^{\prime }}^K\exp (a_{k^{\prime }}(\mathrm{x})))$$

交叉熵定义如下：

$$E=\sum _{\mathrm{x}\in \Omega }w(x)\log (p_{\ell (\mathrm{x})}(\mathrm{x}))$$

$w:\Omega \to \mathbb{R}$是一个预先计算好的权重图，用来分隔细胞间边界。计算公式如下：

$$w(\mathrm{x})=w_c(\mathrm{x})+w_0\cdot \exp (-\frac{(d_1(\mathrm{x})+d_2(\mathrm{x}{)}^2}{2{\sigma }^2})$$

实验中设置$w_0=10,\sigma \approx 5$。

神经网络参数的初始化使用的是高斯分布，标准差为$\sqrt{2/N}$，$N$代表的是该层中神经元的数量，比如输入通道为64，$3\times 3$卷积，则$N=3\times 3\times 64=576$。

实验

本文使用U-Net在三种不同的分割数据集上进行了实验：

电子显微镜神经元结构分割，训练集为30张图片（$512\times 512$像素），来自ISBI 2012 EM分割挑战。

光镜细胞追踪任务，2015年ISBI细胞追踪挑战的一部分。第一个数据集为PhC-U373，包含了35张标注图像，实现了平均IoU=92%，第二名为83%。第二个数据集为DIC-HeLa，20张标注图像，IoU=77.5%，第二名为46%。