经典论文解析——Unet和Vnet——图像分割

  前言说点题外话，最近在实习了，所以总结整理的时间不多，之前的系列也会继续做，只是更新速度会放慢一些。我尽量还是本着以质为主的原则写博客，不弄太水的文章，希望对看过的读者能有帮助，以及对自己能有提升。
  言归正传，最近实习上下班在地铁上看了两篇论文，两者相关性还蛮大，所以放一起介绍。它们就是图像分割领域的 Unet(“U-Net Convolutional Networks for BiomedicalImage Segmentation”)和 Vnet(“V-Net Fully Convolutional Neural Networks forVolumetric Medical Image Segmentation”)。

1.Unet

1.1 网络简介

  在论文前言，作者简要介绍了以下几件事：

• 卷积网络的兴起，对图像领域的进步起到了多么多么大的进步；
• 经典卷积网络大部分都是针对图像分类任务的，但是在一些特定场景，如医疗图像处理领域，通常是需要像素级的分类任务，也就是图像分割任务；
• 为了解决医疗图像即需要像素级分割任务，又因为该领域具备缺少大量数据集的特点，所以Ciresan用滑窗法来解决该问题。但是有两点明显不足：a. 效率太低(太慢了)；b. 滑窗法的感受野(有效视野)大小和分割精度呈负相关关系。
• 作者以FCN全卷积神经网络为基础设计了Unet，其中包含两条串联的路径。数据先经过传统的特征提取路径(论文中称压缩路径，contracting path)，将图像压缩为由特征组成的feature maps。然后再经过特征复原路径(论文中称扩展路径， expansive path)，将提取的特征解码为与原始图像尺寸一样的分割后的预测图像。具体结构见下文的网络示意图。
• 在扩展路径上，和FCN网络不同的是，作者保留了大量的feature map通道，从而使更多的信息能流入最终复原的分割图像中。另外，为了降低在压缩路径上损失的图像信息，作者还将压缩路径的feature map叠加在扩展路径同尺寸的feature map上，再进行卷积和上采样工作，以此整合更多信息进行图像分割。
• 由于数据较少，所以数据增广必不可少。作者采用了弹性变形(elastic deformations)的图像增广，以此让网络学习更稳定的图像特征。因为作者的数据集是细胞组织的图像，细胞组织的边界每时每刻都会发生不规则的畸变，所以这种弹性变形的增广是非常有效的。
• 另外，细胞组织图像的一大特点是，多个同类的细胞会紧紧贴合在一起，其中只有细胞壁或膜组织分割，因此，作者在计算损失的过程中，给两个细胞重合的边缘部分增加了损失的权重，以此让网络更加注重这类重合的边缘信息，具体如下图所示(图取自论文)
  

1.2 网络结构

  论文作者还是很亲民的，给了一张Unet的网络结构图，有了这张图，对Unet的网络结构理解可以说是大有好处。小二，上图：
  这图怎么理解呢，从左上方开始，沿着U型一致到右上方，就是Unet的完整路径。以左上角第一层(stage)为例，输入图像是$572\times 572\times 1$的单通道图像，经过蓝色箭头($3\times 3$卷积层)卷积后，转变为$570\times 570\times 64$。为什么图像尺寸缩小了？因为论文里提到"only uses the valid part of each convolution"，所谓的valid convolution就是不使用padding的卷积(题外话，卷积操作其实有三种类型，和padding与否，padding多少有密切关系，大家可以去查一下英文专业术语和卷积类型的对应关系)。所以，经过$3\times 3$的卷积后，图像尺寸缩小了2。同理，再卷积后就变成了$568\times 568\times 64$。
  在左侧一个stage结束后，会经过一个红色箭头($2\times 2$池化层)，然后，图像尺寸就缩小了一倍，变成了$284\times 284\times 64$，进入左侧第二个stage，每次图像尺寸缩小一倍，对应的feature map数量也增加一倍。同样的过程重复操作4次，到第5个stage，也就是Unet结构的最下方时，就完成了压缩路径的操作，此时的feature maps尺寸为$28\times 28\times 1024$。
  然后开始右侧的扩展路径。首先经过一个绿色箭头($2\times 2$上采样卷积，也称为反卷积或转置卷积)，可以将图像尺寸扩大一倍，那么相对应的，也需要将feature map数量缩小一倍，得到$56\times 56\times 512$的数据(即图中右侧倒数第二层输入的蓝色方块部分)。可以看到，右侧倒数第二层(也称扩展路径第4个stage)的输入是白色方块和蓝色方块叠加形成的1024通道的数据。白色方块哪来？从压缩路径第4个stage池化前的feature maps复制过来的。这就是上文说过的，为了补充图像尺寸压缩损失的信息，需要将压缩路径的feature maps和对应扩展路径的feature maps叠加的过程。然后，以$56\times 56\times 1024$的输入进行两次valid卷积，然后再上采样→叠加→卷积…。重复四次后，就得到$388\times 388\times 64$的输出
  最后再用$1\times 1$的卷积把数据压缩成$388\times 388\times 2$，这里是2通道是因为作者只需要将数据集划分为前景和背景两类即可(当然二分类任务也可以只用单通道)，如果是多分类，可以修改为对应的类别数量。综上所述，可以看到，Unet的结构确实名副其实，是一个很对称的U型网络。
  此外，每个$3\times 3$卷积层后使用ReLU函数激活。当然，如果你想可以在此之前再加上BN层，但是原论文是没有说明的。
  这里可能有细心的读者会问了，输入图像明明是$572\times 572$，为什么输出的预测图像却是$388\times 388$？
  因为这里作者在预处理图像时用了一个小技巧，叫**“Overlap-tile"**，具体什么意思呢？可以参考下图：
  首先，实际要进行图像分割的区域是图中的黄色区域；但是如果只用黄色区域进行分割，那很明显在黄色区域边缘的分割效果就可能很差，因为它们周边能提供的特征信息相比中间图像区域要少很多。所以，为了保证边缘的分割效果，输入网络的实际图像尺寸应该要略大于目标区域，也就是图中的蓝色区域。如果蓝色区域超出了原始图像的边界，那也不能用zero-padding的方式填充，而是要用镜像的方式进行padding，这样至少能保证细胞轮廓的完整性。因此，其实蓝色区域就是网络输入尺寸$572\times 572$，而黄色目标区域就是$388\times 388$。这样，在完成下一个$388\times 388$区域的预测时，前后两次预测的蓝色区域会有所重叠，所以才叫"Overlap-tile strategy"。

1.3 网络训练技巧

• 基本配置：使用momentum=0.99的SGD进行训练，momentum取0.99是作者希望梯度能结合尽量多的之前的动量。为了尽可能利用GPU，作者选用的网络输入图像尺寸是尽可能大。
• 损失函数：网络输出$388\times 388\times 2$的数据后，将$388\times 388$的每个像素经过softmax操作，然后对每个像素求取交叉熵损失。其中，每个像素的交叉熵损失是带有权重的，正如上文所述，重合的边界权重较大。权重的计算和设计可以参见论文，这里不详细展开了。
• 网络参数的初始化：对于深度较深的神经网络，参数初始化是非常重要的。作者采用标准差为$\sqrt{2/N}$的高斯分布作为参数初始化方式，其中N表示一个神经元的输入参数数量，如$3\times 3\times 64$的卷积操作，$N=3\times 3\times 64=576$。
• 数据增广：
  • 图像平移；
  • 图像旋转；
  • 图像弹性变形，使用双三次插值进行变形后的填充；
  • 压缩路径最后使用了dropout层;

2.Vnet

2.1 网络简介

  相比于Unet，Vnet其实就是Unet的一个变型。所以论文就不展开介绍了，了解了Unet的同学再去看Vnet的文章是比较轻松的。Vnet论文主要是针对医学图像提出的，作者使用的数据集是三维图像，而不是二维，这个和常见数据是有点区别，但其实做法类似。最后输出是单通道的三维数据，代表每个像素的前景/背景概率，如果概率>0.5则是前景。因此，网络的输出最终要经过sigmoid函数激活。

2.2 网络结构

  从图像上可以看出，Vnet和Unet确实是很像的。但是由于论文是针对三维图像提出的，所以图中用方格代表feature map。另外，还可以看出，Vnet也借用了Unet从压缩路径叠加feature map，从而补充损失信息的方法(橙色线路)。这里需要特别说明的，也是Vnet和Unet最大的不同，就是在每个stage中，Vnet采用了ResNet的短路连接方式(灰色路线)。相当于在Unet中引入ResBlock。这是Vnet最大的改进之处。
  另外，Vnet第一个stage在短路部分只进行了一次卷积操作，在第二个stage进行了两次…这和Unet每个stage卷积操作次数相同的结构是略有不同的。
  最后，Vnet由于是医学图像领域的网络，且其数据集是二分类任务，所以他们采用了医学图像常用的Dice loss。简单来说，Dice其实就是分割领域的类似IOU一样的指标。假设A是预测出的所有前景像素集合，B是ground truth的前景像素集合，则:
$$Dice=\frac{2|A\bigcap B|}{|A|+|B|},DiceLoss=1-\frac{2|A\bigcap B|}{|A|+|B|}$$
  具体的Dice系数细节，可以参阅这篇博客。