介绍

相对ImageNet等通用数据集,医学图像数据集较小。如何在小数据集情况下训练出一个好的模型,是深度学习在医学图像方面的一个难点。本论文提出的U-Net在ISBI 2015 Cell Tracking Challenge取得了第一。在该任务中,训练集中只有30幅512x512果蝇第一龄幼虫腹神经(VNC)的电镜图像。为了充分利用数据,文章中还采用了Data Augmentation的方法扩充数据集。此外,文章中有许多技巧可以应用于数据集很小的情况中。

模型

设计思想

U-Net是基于FCN网络的思想设计的,整个网络只有卷积层,而没有全连接层。
U-Net由一条收缩路径(contracting path)和一条扩张路径(expanding path)组成。在收缩路径中,图像分辨率逐渐降低,上下文信息会逐渐增强;在扩张路径中,通过上采样的方式,让特征图的分辨率逐渐增大。同时,为了结合低层feature map的强位置信息,将收缩路径中的相应部分以concatenate的方式(区别于FCN中的直接相加)结合到扩张路径中。这种架构可以较好地进行位置定位。
1.png

U-Net的创新点包括:

  1. 在上采样部分,feature map的通道数非常大,作者认为这样可以将上下文信息传递到分辨率更高的层当中。这样做的一个结果就是它基本上和浓缩路径对称了,因此看上去像一个U形的结构。
  2. 为了预测图像边界区域的像素点,采用overlap-tile策略补全缺失的context。
  3. 由于训练数据太少,采用大量弹性形变的方式增强数据。这可以让模型更好学习形变不变性。这种增强方式对于医学图像来说很重要。
  4. 在细胞分割任务中的另一个挑战是,如何将同类别的相互接触的目标分开。本文提出了使用一种带权重的损失(weighted loss)。在损失函数中,分割相互接触的细胞像素获得了更大的权重。

Overlap-tile Strategy

2.png
由于图片边缘(黄色框)的预测值需要蓝色框内的图像来计算。为了达到更好的识别效果,文章还采用了Overlap-tile的技巧,将图像边缘进行一定的镜像复制生成边缘图像,这样边缘的识别效果会更好。

Weights initialization

权值初始化对于含有许多卷积层和不同路径的深度网络来说十分重要。如果权值初始化不当,可能导致网络的某些部分贡献很重的的activation,而另一些部分没有贡献。因此,权值的初始化好坏程度会影响最终模型效果以及训练时间。理想的权值初始化是使得网络中的每一个特征图方差都接近1。在U-Net中,可以通过高斯分布随机生成权值,高斯分布的均方差为$\sqrt{2 / N}$,其中N为上一层的参数数量。

网络结构

收缩路径

  1. 注意图像输入是经过tile的。因此网络的输出是和图像的输入大小是不一样的。
  2. 每经过两个3x3 conv(没有padding)之后,会跟上一个2x2的max-pooling进行下采样。

扩张路径

  1. 使用2x2的deconv来进行上采样,上采样的过程中,通道数减半。同时,将上采样的feature map和收缩路径中相应的feature map进行concatenate(注意:由于收缩路径中的conv都是没有做padding动作的,这会导致扩张路径中的feature map和收缩路径中相应位置的feature map大小不一致,这个时候就要将收缩路径中的feature map进行crop再和扩张路径中的feature map进行concatenate)。
  2. 上采样的结果和收缩路径中的相应feature map的裁剪进行concatenate后,再跟上3x3的卷积。由于收缩路径中卷积操作是没有padding的,因此feature map在卷积之后会越来越小,之后在扩张路径进行上采样无法恢复相应的分辨率,因此需要裁剪。
  3. 最后一层使用1x1的卷积,将通道数map到类别数。

训练

损失函数计算

网络输出的是pixel-wise的softmax。表达式如下:

$$ \exp \left(a_{k}(\mathbf{x})\right) /\left(\sum_{k^{\prime}=1}^{K} \exp \left(a_{k^{\prime}}(\mathbf{x})\right)\right) $$

其中, $x$ 为二维平面 $\Omega$ 上的像素位置,$a_{k}(x)$ 表示网络最后输出层中pixel $x$ 对应的第 $k$ 个通道的值,$K$ 是类别总数。$p_{k}(x)$表示像素 $x$ 属于 $k$ 类的概率。

损失函数使用negative cross entropy。cross entropy的数学表达式如下:

$$ E=\sum_{\mathbf{x} \in \Omega} w(\mathbf{x}) \log \left(p_{\ell(\mathbf{x})}(\mathbf{x})\right) $$

其中 $p_{l}(x)$ 表示 $x$ 在真实label所在通道上的输出概率。需要特别注意的是cross entropy中还添加一个权重项 $w(x)$。这是因为考虑到物体间的边界需要更多的关注,所对应的损失权重需要更大。

像素损失权重计算

在损失函数计算中我们讲到对于边界像素我们给的损失权重要更大,但怎么获取这个权重?我们得到一张图片的ground truth是一个二值的mask,本文首先采用形态学方法去计算出物体的边界。然后通过以下的表达式去计算权重图。

$$ w(\mathbf{x})=w_{c}(\mathbf{x})+w_{0} \cdot \exp \left(-\frac{\left(d_{1}(\mathbf{x})+d_{2}(\mathbf{x})\right)^{2}}{2 \sigma^{2}}\right) $$

其中 $w_{c}(x)$ 是类别权重,需要根据训练数据集中的各类别出现的频率来进行统计,类别出现的频率越高,应该给的权重越低,频率越低则给的权重越高(文章没有详细说是怎么计算的)。$d_{1}(x)$ 表示物体像素到最近cell的边界的距离,$d_{2}(x)$ 表示物体像素到第二近的cell的边界的距离。在本文中,设置 $w_{0}$=10,$\sigma$=5。

数据增强

由于本次任务是关于电镜图像下的语义分割,因此主要的挑战之一就是细胞或者组织的形状变化。考虑到这个因素,文章主要将训练集图像变形以产生更多训练图像,来达到扩充数据集的效果。文章将训练图片分格子,然后采用符合高斯分布的变形向量来变形。然后,使用双三次插值来填充变形造成的分辨率降低的问题。

实验

3.png
4.png

参考文献

[1] U-Net: Convolutional Networks for Biomedical Image Segmentation
[2] 知乎专栏:语义分割之U-net
[3] 简书:生物医学图像语义分割(二)U-Net

Last modification:August 11th, 2019 at 11:14 pm
如果觉得我的文章对你有用,请随意赞赏