论文笔记：用于医学图像分割的卷积网络（U-net）

摘要

• 本文中提出了一种网络结构的训练策略，它依赖于充分利用数据增强技术来更高效的使用带有标签的数据。在U-Net结构中，包括一个捕获上下文信息的收缩路径和一个允许精确定位的对称拓展路径。这种方法可以使用非常少的数据完成端到端的训练，并获得最好的效果。

介绍

• 卷积网络的典型用途是分类任务，其中图像的输出是单个类别标签。 然而，在许多视觉任务中，尤其是在生物医学图像处理中，期望的输出应该包括定位，即，应该将类标签分配给每个像素。（也就是分割）

• Ciresan等人用滑动窗口取像素像素周围的局部区域俩训练网络，训练数据远远大于训练图像的数量。

• 本文提出一种新的完全卷积网络，即U-Net网，主要思想是在收缩网络的后面补充一个与前面类似的网络，其中池化运算符由上采样运算符替换。因此，这些层增加了输出的分辨率。为了定位，从收缩路径的高分辨率特征与上采样输出相结合。然后，连续卷积层可以学习基于该信息组装更精确的输出。

• 作者提出U-Net的本意是将其用于医学图像分割，在以往的CNN中，想将其用于医学图像存在两个困难：

  1. 通常CNN都是应用于分类，生物医学图像更关注的是分割以及定位的任务；
  2. CNN需要获取大量的训练数据，而医学图像很难获得那么大规模的数据。

• 以往解决上面两点困难的方法是使用滑窗的方法，为每一个待分类的像素点取周围的一部分邻域输入。这样的方法有两点好处，首先它完成了定位的工作，其次因为每次取一个像素点周围的邻域，所以大大增加了训练数据的数量。但是这样的方法也有两个缺点，首先通过滑窗所取的块之间具有较大的重叠，所以会导致速度变慢；其次是网络需要在局部准确性和获取上下文之间进行取舍。因为更大的块需要更多的池化层进而降低了定位的准确率，而小的块使网络只看到很小的一部分上下文。

U-Net

• 在本文中，作者修改并扩展了FCN网络结构，使它在使用少量的数据进行训练的情况下获得精确的分割结果，具体结构示意图如下所示：

image

• 在上图中，每一个蓝色块表示一个多通道特征图，特征图的通道数标记在顶部，X-Y尺寸设置在块的左下边缘。不同颜色的箭头代表不同的操作。图的左半部分是收缩路径，右半部分是扩展路径。
• 其中需要注意的是，每经过一次上采样都会将通道数减半，再与收缩路径对应的特征图进行拼接。在拼接之前进行 crop 是必要的(例如在上图中，6464大与5656，为了使这两个特征图能够顺利拼接，取6464中间部分5454的大小，然后拼接)，因为两者的尺寸并不相同（主要是因为 valid conv 造成的）。最后一层使用1 X 1大小的卷积核，将通道数降低至特定的数量（如像素点的类别数量）。
• 网络对于输入的大小也是有要求的。为了使得输出的分割图无缝拼接，重要的是选择输入块的大小，以便所有的2 X 2的池化层都可以应用于偶数的 x 层和 y 层。一个比较好的方法是从最下的分辨率从反向推到，比如说在网络结构中，最小的是32 X 32，沿着收缩路径的反向进行推导可知，输入图像的尺寸应该为572×572。

FCN与U-net的区别

• FCN的主要思想是使用连续的层补充通常的收缩网络，在判别输出的位置添加上采样层，这些层增加了输出层的分别率，为了定位。来自于收缩路径的高分别率与上采样输出相结合，基于这个信息，一个连续的卷积层可以学习组装更精确的输出。（关于FCN详细内容可以参考我的上一篇文章：论文笔记：用于语义分割的全卷积网络 （fully convolutional networks for semantic segmentation））
• 而U-net与FCN的不同在于，U-net的上采样依然有大量的通道，这使得网络将上下文信息向更高层分辨率传播，作为结果，扩展路径与收缩路径对称，形成一个U型的形状（如上图所示）。 网络没有全连接层并且只是用每一个卷积层的有效部分。

Overlap-tile

• 作者在文中介绍了一种Overlap-tile策略，使得任意大小的输入图像都可以获得一个无缝分割，因为输出的分割图它包含的像素点，它们的周围像素点（上下文）都出现在了输入图像中，因此使用Overlap-tile策略对数据进行预处理是有必要的。
• Overlap-tile策略的过程具体如下所示：

image

• 上图是针对任意大小的输入图像的无缝分割的 Overlap-tile 策略。如果我们要预测黄色框内区域（即对黄色的内的细胞进行分割，获取它们的边缘），需要将蓝色框内部分作为输入，如果换色区域在输入图像的边缘的话，那么缺失的数据使用镜像进行补充。如上图左边图像所示，输入图像周围一圈都进行了镜像补充。

• 因为进行的是valid卷积，即上下文只取有效部分，可以理解为padding为0，卷积之后的图像尺寸会改变，所以需要取比黄色框大的图像来保证上下文的信息是有意义的，缺失的部分用镜像的方法补充是填充上下文信息最好的方法了。这种方法通常需要将图像进行分块的时候才使用。

• 那么为什么要对图像分块，不输入整张图像呢？因为内存限制，有的机器内存比较小，需要分块输入。但比之前的滑窗取块要好很多，一方面不用取那么多块，另一方面块之间也没有那么大的区域重叠。通过Overlap-tile 策略可以将图像分块输入，否则的话就只能对图像进行 resize 了，但是这样会降低输入图像的分辨率。

• 此外，如果数据不够的话可以应用弹性形变，对数据进行增强，增加数据量，这允许网络可以学习到这种形变的不变形，并且并不要求带有原始预料库进行到这样的变化（指弹性形变）。

细胞分割

• 进行细胞的分割，另一种挑战是同一类物体的分类，如下图所示：

image

• 上图是用DIC（二次干涉对比）显微技术记录的玻璃上的 HeLa 细胞。其中图 (a) 是原始图像；图 (b) 是基于 gt 的分割覆盖。其中不同的颜色表示不同的 HeLa 细胞示例。图 (c) 是生成的分割掩膜，其中白色部分是前景，黑色部分是后景；图 (d) 是像素级损失权重图，使得网络强制学习边缘像素。
  + 为了解决这个问题，作者使用加权损失，对于位与细胞接触部分的像素加大权重，如图 (d) 中的红色的部分。

Training（训练）

• 作者使用caffe框架，并采用带动量的SGD方法，其中动量的值设为 momentum=0.99。动量设这么大的原因是这样可以使用大量先前看到的训练样本确定当前最优步骤的更新(因为动量的原理就是用先前很多步的状态确定下一步的方向)。相比于大的batchsize，作者喜欢大的input tiles（指的是over-tile）中的那种图像块，因此我们可以将一个batch缩小为一个单张图片的输入。
• 能量函数通过结合交叉熵损失函数的最后特征图上的像素级 soft-max 值来计算，通常多分类问题用soft-max函数作为输出，二分类问题用sigmoid函数作为输出，其中 soft-max 的计算方法如下：

image

• 其中 αk(x)代表在位置 x 处的像素在特征图中的第 k 层的激活值，其中 x∈Ω，，Ω⊂Z^2 ，即 x 属于空间 Ω，而空间 Ω是整数集合Z^2的子集。K 是像素点的类别总数，p_k(x)是近似最大函数。这里面的定义和我们平时使用的 sift-max 是一样的。损失函数是交叉熵损失函数，或者也可以成为 log-likelihood，具体如下所示

image

• 其中需要注意的是这里使用的是加权的损失函数，对于每一个像素点有着自己的权重，这点可以在上面的细胞图中看到。

• 我们通过预先计算权重图来获得每一个像素在损失函数中的权值，这种方法补偿了训练数据每类像素的不同频率，并且使网络更注重学习相互接触的细胞的边缘。

• 分割边界使用形态学运算，特征图的计算方法如下：

image

• 其中的 w_c是用于平衡类别频率的权重图，d_1是该像素点到最近的细胞边界的距离；d_2是该像素点到第二近的细胞边界的距离。在我们的实验中，将 w0设置为10，将σs设置为大约 5 个像素。

• 权值的初始化对于模型的正确训练起着很大的作用，一个好的初始化应该保证网络中的每一个特征图有近似的单位方差。在这里使用服从标准差为√2/N 的高斯分布来进行初始化（实际上就是 He normal），其中的 N 代表着一个神经元的传入节点的数目，比如说某一个卷积层，他的卷积核的大小为 3×3，通道数是 64，那么 N=9×64=576。

  Data Augmentation（数据增强/扩充）

• 当只有少量训练样本的时候，对于让网络学习到所需的不变形和鲁棒性来说，数据增强是必要的。

• 显微镜图像一般需要旋转平移不变性，弹性形变和灰度值变化鲁棒性。训练样本的随机性变形似乎是训练之后少量标注图像的分割网络的关键。

• 此外在本文中，早收缩路径的最后加入了Drop-out，隐式的加强了数据增强。

总结

• 本文提出了很多数据预处理的技巧，有较强的实用性。

本文的创新点

1. 提出了U-net这种网络结构。它同时具备捕捉上下文信息的收缩路径和允许精确定位的对称扩展路径，并且与FCN相比，U-net的上采样过程依然有大量的通道，这使得网络将上下文信息向更高层分辨率传播。
2. Overlap-tile 策略，这种方法用于补全输入图像的上下信息，可以解决由于现存不足造成的图像输入的问题。
3. 使用随机弹性变形进行数据增强。
4. 使用加权损失。预先计算权重图，一方面补偿了训练数据每类像素的不同频率，另一方面是网络更注重学习相互接触的细胞间边缘。

但是上面的这些创新点并不适合所有任务，比如说对刚体进行分割，那么久很难通过弹性变形对数据进行增强。

本文学到的一些实用的技巧

1. 因为使用了 valid conv ，所以采用 Overlap - tile 策略补充图像，其中空白的部分用镜像的方法进行补充；
2. 因为有池化层，因此要保证输入的图像在经过每一次池化的时候都要是边长偶数。这点与一般的卷积神经网络不同，因为一般的网络会使用 padding ，这样会保证卷积前后的大小不变，但是 valid conv 会使卷积后的尺寸变小，所以要特别注意输入图像的尺寸。一个比较好的方法是从最小分辨率出发沿收缩路径的反方向进行计算，得到输入图像的尺寸。
3. 预先计算权重图，以此计算后面的加权损失函数；
4. 加权损失的权重中有一部分是经验值，因此对于不同的任务可以进行调整（只是理论上可以进行调整，并没有试验过）；
5. 使用标准差为√2/N的高斯分布来进行权值初始化，其中需要注意的是，对于不同的卷积层，N 的大小也是不同的。
6. 在收缩路径的最后部加入了 dropout ，隐式地加强了数据增强。

一些疑惑

• 为什么加入 dropout 后可以对数据进行增强？
• 反卷积（上采样）的具体过程是怎么样的？
• 弹性形变数据增强具体是如何对数据增强的？

参考

• 论文原文：U-Net: Convolutional Networks for Biomedical Image Segmentation
• CSDN-独孤呆博：论文精读及分析：U-Net: Convolutional Networks for Biomedical Image Segmentation

微信搜索公众号：轮子工厂。机器学习 | 数据结构与算法 | 源码分析 |　资源与工具分享
后台回复关键字：
1.回复【图书】：获取15本新手自学编程，零基础入门经典学习教材；
2.回复【我要造轮子】：获取100多本我根据知乎上面关于计算机问题的高赞回答里面的介绍整理出来的书籍；
3.回复【开发工具】：获取几大主流编程语言的开发工具～
4.回复【ps教程】：获取ps视频免费教程；
5.回复【内推】：可帮你内推到大厂工作。