用于语义分割的U-Net为什么这么强？

背景

对于医学图像，有很多地方都会用到U-Net，包括2018 Data Science Bowl 用于分割细胞核的冠军模型，使用的都是U-Net结构，其用了大量的数据扩增，以及最后用了分水岭算法。由此可见，这个2015年的网络，在如今的医学图像语义分割效果依旧一骑绝尘。为什么针对一般的医学图像，U-Net的表现会这么好呢？为什么只需要几十张图片进行训练，其效果就能接近人类的分割效果呢？这个问题我们留到最后再进行解答。

传统的图像分类网络是针对一整张图像进行多分类，也就是一整张图片丢进我们的黑箱模型，直接出来一个分类结果，但是到语义分割，我们就需要对每个像素进行一个精确的分类。对于此，我们如何既考虑到图像全局的信息，又考虑到图片细节的信息呢？

U-Net为我们提供了一种非常棒的解决思路，下面我们来详细地聊一聊它。附上：U-Net原文。

---

网络结构

首先我们来看其网络结构。左边一侧就是我们传统的卷积层与池化层（也就是所说的pooling层），而到了右边就是利用转置卷积（deconvolution）进行上采样（upsampling）。另外，每次准备池化时进行降采样时，都会将其融合（concatenate）到后面的转置卷积后的feature map上。（如上图的灰色的直连线所示）

只是需要注意的是，这里我们每次卷积处理一次，图像的尺寸就会缩小一部分。因为我们使用卷积时，不会在图像外层补0，使卷积后与卷积前的图像保持相同的大小。这样在每次卷积操作后，图像就会越来越小。

同样在进行转置卷积将图像一层一层放大时，也不会将图像尺寸放大为原先的两倍，而是会比两倍的边长少2（因为卷积核为3 * 3）。

这样就会带来两个问题：

1. 与后面的feature map融合时，如何统一大小？
2. 如何保证real image与对应的mask保持一样的大小？

针对第一个问题，其实很好解决，直接进行剪裁，保留feature map中间部分的信息，与后面的feature map进行拼接融合。而针对第二个问题，论文提出了一个比较巧妙的解决方案：

具体如上图所示。由于网络的设置，最后图像大小会变成388 * 388，而输入是572 * 572。这时注意，我们真实的图片是388 * 388的，保证其和输出的大小一致（从图中的网络结构来看真实图像的大小是388 * 388，但按照论文跑的真实数据对应的label与原图都是512 * 512，此时为保证label大小一致，为 512 * 512，输入图像的尺寸应该长宽均为：512 + (572 - 388) = 696）。对于真实图片，我们将其扩增到572 * 572的大小。论文中不是用补0，亦或是粗暴地拉伸操作，而是在边界处，巧妙地进行翻折，镜像地将图片本身的内容“折”到外侧，将图像的尺寸进行扩增。

当然，后面许多仿U-Net的网络，都会使图像保持着相同的尺寸，在图像外层进行补一圈0的操作，其结果也相当不错，所以所说的操作技巧不是U-Net表现出色的核心所在。其关键还是cancat这个操作，接了额外的四个通路，使其能充分利用图像本身的信息，并且也能很好地缓解BP时的梯度弥散问题。

下面我们对网络的具体细节操作进行说明：

---

卷积与降采样

卷积与降采样都是非常传统的方式，推荐学习网址：An Intuitive Explanation of Convolutional Neural Networks，我们这里搬运几张图片，就不进行过多详细地讲解。

卷积

我们的原本图片（或者feature map）：

3 * 3卷积核：

卷积操作（本质：加权求和）：

上面的卷积操作就非常贴合U-Net中的操作（卷积核为3 * 3，步长为1，同时没有对我们的原图在外边缘补上0）。卷积的作用就是将原图抽象化，起到提取特征的目的。对于同样的输入图像，不同值的卷积核将会生成不同的特征图。所以我们会对一个图，通过多个不同的卷积核，生成多个不同的feature map（或者称channel），丰富特征，提升最后我们的预测准确率。

1 * 1 卷积核

在卷积操作中，有个特殊的存在：1 * 1 卷积核。这个操作我们在U-Net的最后一层也能够看到。其目的不是为了对单个channel进行卷积操作，而是为了沟通多个channel，使其能够在channel这一维度上做线性组合，这就是通道间的信息交互。

举个栗子：原本的338 * 338 * 64的特征图，进行了1 * 1卷积核的卷积操作后，变为了338 * 338 * 2的输出，其主要是将64层channel变为了2层channel。这个例子相当于训练了两个1 * 1 * 64的卷积核权重。

池化

上图展现了最大池化，这里也不再进行过多的介绍。池化的主要作用是：

1. 减少参数和计算复杂度，控制过拟合，提高模型的泛化能力；
2. 使网络对于输入图像中一些小的变化，变得不那么敏感（输入的微小冗余将不会改变池化的输出），并且具有平移、旋转、尺度不变性。

---

上采样

这部分在用深度学习做语义分割中是比较精华的部分。关键问题是我们将图片一层一层缩小之后，如何还原回原图的大小？

这就需要用到我们的转置卷积。关于转置卷积，这两篇教程写得非常详细：1. A guide to convolution arithmetic for deep learning，2. Convolution arithmetic tutorial。下面这两篇教程，说下我对转置卷积的理解：

对于一般的卷积操作，本质上我们可以将输入图片拉直，然后进行进行矩阵的乘法，最后输出一个向量（也就是卷积之后的feature map拉直的结果）。

我们以一个3 * 3的卷积核为例，输入是一个4 * 4的图片，拉直后也就是一个16维的列向量，记为$x$。

卷积核我们用下面的矩阵$C$表示，

则卷积运算可以变为：
$$y=Cx$$
这时我们就可以得到一个4维的列向量$y$，也就是我们卷积后的输出矩阵拉直的结果。到此，我们就进行了一个降维的操作。

那如果我们想将一个4维的列向量$y$，映射为16维的列向量$x$需要怎么操作呢？首先的想法是前面的卷积运算，左右两边均左乘$C^{-1}$，但是很明显，$C$连一个方阵都不是，是没有逆的。这时，我们考虑$C$的转置：

进行如下的运算：
$$x_{(new)}=C^{T}y$$
这样我们同样也能够将一个低维的向量空间映射到高维（同样，矩阵里面的权重也是需要通过训练得到）。只是需要注意：转置不是逆，我们通过卷积得到$y$后，再通过转置卷积变换回$x_{(new)}$，$x_{(new)}$与原先的$x$只是维数相同，值是不一样的，这点需要注意。

接下来我们通过可视化的方式来进一步理解转置卷积。这是卷积的常规操作（蓝->绿）

而下面就是转置卷积的操作，其实也是卷积，只是在外面补0。如果我们的步长为1，就是下面的情况：

如果步长为2，则需要将原本的矩阵打散，间隔变为1（以此类推），再一格一格的滑动。注意，这里的步长变化不是真的移动卷积核时跳着移，而是将矩阵每个像素之间打散，中间用0填补（通常是0，也可以用差值的方法）。

---

预处理

文中的预处理方式是：

1. 平移
2. 旋转
3. 随机弹性形变（文中说，对于细胞的分割，这个非常重要）
4. 随机灰度值变化

这里不做过多的探讨。

---

损失函数

文中的损失函数使用的是cross entropy loss function，准确的说，里面还需要加一个softmax函数。这里是针对每个像素，每个像素来进行计算，最后将一张图上所有像素点的求和。

$a_k(x)$为第$k$个feature channel（最后一层有几个类别就有几个feature channel，一共$K$个）激活函数（文中是relu）作用后的值，然后再作用一个softmax函数计算$p_k(x)$，也就是属于第$k$类的概率：

最后我们计算cross entropy loss function：

上面这个式子其实简化了对类别的求和。因为针对一个像素点$x$而言，其只属于一个类别，也就是对label向量，只有一个类别是1，其余都是0。所以针对一个像素点，我们可以只将真实是第$l(x)$类的那一项进行保留，然后再将所有的像素点的loss相加即可。

计算cross entropy loss function时，我们需要考虑不同的类别，其造成的损失是不一样的，因此会赋予不同的权重：$w(x)$。对于罕见病而言，患病个体的权重是要远远高于健康个体的，而对于一个如下图所示的细胞分割问题，我们需要重点分割出来的不是大片黑色的无细胞区域（图c），而是细胞与细胞之间狭小的细胞膜。

因此，作者给出了如下的权重计算方式：

我们来理解一下。$w_c(x)$就是我们常规的为样本进行加权的方法，其是用于平衡不同类别的不平衡性（样本多的权重小，样本小的权重大）。后面一项就是比较精华的地方。$d_1(x)$表示边界到最近细胞的距离，$d_2(x)$表示边界到第二近的细胞的距离。

可以想象对于两个细胞薄薄细胞膜上的像素点，$d_1(x)$与$d_2(x)$都非常小，那么取负号再取指数后，变得很大。而对远离两个细胞的空白位置上的像素点，$d_1(x)$与$d_2(x)$都很大，那么取负号再取指数后，权重则相对较小。

$\sigma$则是用于类似标准化的常数，文中取5，$w_0$取10，这些都是经过多次试验得到较好的值。

---

思考

现在我们再思考一下，对于U-Net这样的一个网络结构，其效果为什么这么突出？

为什么好？

我们可以看到，它能够结合图像的全局与局部细节方面的特征，进行综合的考虑。每经过几次卷积后，便会concatenate到后面的几层，直接传递到同高度decoder上进行信息融合。在刚开始卷积时，图像的feature map保留的是图像高分辨率的细节信息（病灶的具体形态，条状还是点状，边缘平滑亦或是粗糙等等），它能帮助最终我们的图像提供精细分割；而经过一次一次卷积池化后，最后在U-Net的最下层，此时的feature map中包含的是整幅图像的全局信息（病灶的总体位置、分布等），它能提供分割目标在整个图像中上下文语义信息，反应目标和周围环境之间的关系。

最后再经过一层一层上采样，将不同层级的信息一步一步融合，便可融合多尺度的信息，得到最终我们的判别分割结果。

什么情况下好？

用过U-Net的童鞋应该都知道，其很好的应用于少样本的医学图像。

首先医学图像等样本，问题空间是很小的，例如，细胞核分割、细胞膜分割、视网膜血管分割等等。变来变去都是这几种样式，而且我们还能用多种方式进行数据的扩增，由几十张图片扩增到数十万张训练集图片。所以小样本的医学图像就能够表现得非常出色了。而对于一般生活中的图片，则比较难处理，其问题空间非常大（风景、建筑、人物等等），U-Net这么小的一个结构难以完全刻画其中的关系。

另外，针对一般的医学图像，边界模糊、梯度复杂，需要较多的高分辨率信息用于精准分割。同时一般的医学分割目标在人体图像中的分布很具有规律，低分辨率信息能够提供这一信息，用于目标物体的识别。两者结合，正好是U-Net完美适用的场景。

---

延伸

当然，对于U-Net，很多大神们在其基础上进行了各种各样的改进，包括：ENet（部分医学图像表现得比U-Net要好），U-Net++等等。下面贴上这两个网络的结构图，有兴趣的童鞋可以看看论文。

ENet

U-Net++

---

参考

1. U-Net: Convolutional Networks for Biomedical Image Segmentation
2. An Intuitive Explanation of Convolutional Neural Networks
3. Convolution arithmetic tutorial
4. A guide to convolution arithmetic for deep learning
5. A technical report on convolution arithmetic in the context of deep learning
6. U-Net分割网络为什么对样本量小的训练集，效果依然很好？
7. Unet神经网络为什么会在医学图像分割表现好？
8. 图像语义分割入门+FCN/U-Net网络解析
9. ENet: A Deep Neural Network Architecture for Real-Time Semantic Segmentation
10. UNet++: A Nested U-Net Architecture for Medical Image Segmentation
11. 研习U-Net