【论文笔记】UNet++:一种用于医学图像分割的嵌套U-Net结构

本文是《UNet++: A Nested U-Net Architecture for Medical Image Segmentation》论文的阅读笔记。强烈建议大家去看下作者对这篇论文的思维过程:研习UNet

文章提出了一个UNet++的网络,它是一个使用了深度监督(deep supervised)的编码器-解码器结构,在解码器和编码器之间有一系列的密集跳跃连接(dense skip connections)。

一、跳跃连接的重要性

有很多效果不错的网络结构都采用了跳跃连接的方式,如UNet网络和FCN(全卷积网络)。不同的是,FCN的跳跃连接是通过元素级的相加来实现的,而UNet的跳跃连接是通过通道的拼接来实现的。此外,ResNet中的shortcut和DenseNet中的密集连接(dense connection)也有异曲同工之妙。跳跃连接可以让解码器中深层的、语义的、粗糙的特征图和编码器中浅层的、低级的、精细的特征图相结合。

由于医学图像是对手术有指导意义,而万一有所偏差则会产生较为严重的后果,所以医学图像的处理要求比自然图像有更高的精度。

二、网络结构

【论文笔记】UNet++:一种用于医学图像分割的嵌套U-Net结构

UNet++网络的结构如上图(a)所示,其中黑色的部分是原来的UNet,绿色和蓝色的部分是在跳跃路径上的密集卷积块(dense convolution blocks),红色的连线表示深度监督。

UNet++与原始的UNet主要有三点不同之处:

  • 在跳跃路径上有卷积层(图中绿色部分),减小了编码器和解码器特征图之间的语义鸿沟(semantic gap)
  • 在跳跃路径上有密集跳跃连接(图中蓝色部分),提升了梯度流
  • 有深度监督(图中红色部分),可以让模型实现剪枝和提升

1. 跳跃连接

xi,jx^{i,j}表示节点Xi,jX^{i,j}的输出,其中ii表示层数,jj表示当前层的第jj个卷积层。xi,jx^{i,j}就可以通过一下公式来计算:
xi,j={H(xi1,j),j=0H([[xi,k]k=0j1,U(xi+1,j1)]),j>0 x^{i, j}=\left\{\begin{array}{l} \mathcal{H}\left(x^{i-1, j}\right), \quad j=0 \\ \mathcal{H}\left(\left[\left[x^{i, k}\right]_{k=0}^{j-1}, \mathcal{U}\left(x^{i+1, j-1}\right)\right]\right), \quad j>0 \end{array}\right.
其中,函数H()\mathcal H(\cdot)表示卷积操作和**函数,U()\mathcal U(\cdot)表示上采样层,[][]表示拼接层(concatenation layer)。UNet++中的密集连接如上图(b)所示。

2. 深度监督

模型一共有四层,也就是有四个分支(语义级),深度监督有两种模式:

  • 精确模式,在该模式下所有分割分支的结果取平均值得到最终的结果;
  • 快速模式,在该模式下只选择一个分割分支,其他的被剪枝。

当某个分支的权重为0时,就相当于是对其进行剪枝了,示意图如上图©所示。

用二值交叉熵和DICE系数的结合来作为四个语义级的损失函数,如下所示:
L(Y,Y^)=1Nb=1N(12YblogY^b+2YbY^bYb+Y^b) \mathcal{L}(Y, \hat{Y})=-\frac{1}{N} \sum_{b=1}^{N}\left(\frac{1}{2} \cdot Y_{b} \cdot \log \hat{Y}_{b}+\frac{2 \cdot Y_{b} \cdot \hat{Y}_{b}}{Y_{b}+\hat{Y}_{b}}\right)
其中,Yb^\hat{Y_b}YbY_b分别表示第bb张图片的预测概率和ground truth,NN表示batch size。

三、实验

【论文笔记】UNet++:一种用于医学图像分割的嵌套U-Net结构

实验过程使用了四个数据集,如表1所示,分别为:。

【论文笔记】UNet++:一种用于医学图像分割的嵌套U-Net结构

用于作比对的网络模型选用的是原始的UNet和自定义的wide UNet模型。评价指标选用的是DICE系数和IoU(交并比)。采用学习率为3e-4的Adam作为训练的优化器,UNet和wide UNet的详细结构如表2所示。第ii层的卷积层使用的是kk个卷积核为3×33\times33×3×33\times3\times3的卷积操作,其中k=32×2ik=32\times2^i。在深度监督时,每个目标节点{x0,jj{1,2,3,4}}\{x^{0,j}|j\in\{1,2,3,4\}\}后跟着的是1×11\times1的卷积操作和sigmoid**函数。

【论文笔记】UNet++:一种用于医学图像分割的嵌套U-Net结构
【论文笔记】UNet++:一种用于医学图像分割的嵌套U-Net结构

实验结果如上图和表3所示,在IoU指标上,没有深度监督(DS)的UNet++比baseline平均提升了2.8到3.3个点。有深度监督的UNet++平均提升了0.6个点。
【论文笔记】UNet++:一种用于医学图像分割的嵌套U-Net结构

上图显示的是不同的剪枝下模型的效果和时间对比图,其中UNet++LiL^i表示的是在第ii个level进行剪枝的UNet++模型。可以发现UNet++L3L^3在损失一点精度的情况下实现了时间的大幅缩减。

四、其他

根据作者《研习UNet》的文章,我下面来捋一下作者的idea产生的过程。

首先作者意识到了跳跃连接会提升模型的效果,但是模型的深度为多少时才最合适呢?当然我们可以创建不同深度的模型然后分别进行实验,找出最优的那个,但是这样太遭罪了。下图是不同深度下的网络模型,所以作者想可不可以将不同深度的模型拼成一个呢?于是就得到了UNet++网络的雏形。

【论文笔记】UNet++:一种用于医学图像分割的嵌套U-Net结构

拼凑的结果如下图所示,但是该模型有个问题,就是红色三角形的部分在反向传播时不能更新,因为X0,3,X1,2,X2,1X^{0,3},X^{1,2},X^{2,1}和右边的操作没有路径连接。

【论文笔记】UNet++:一种用于医学图像分割的嵌套U-Net结构

这时有两种解决办法,一种是上面提到的深度监督,另一种更直接一点,将X0,3,X1,2,X2,1X^{0,3},X^{1,2},X^{2,1}和右边的操作连接起来就得了。

【论文笔记】UNet++:一种用于医学图像分割的嵌套U-Net结构

但是这样又有一个问题,UNet中使用的是长连接,而上图使用的是短连接,而作者认为长连接是有必要的。所以就将其改成了密集连接的形式,如本文最开始的那张图所示。

而深度监督的加入又让模型具有了剪枝的能力,如下图所示,当某个分支的权重为0时,就相当于做了剪枝操作。

【论文笔记】UNet++:一种用于医学图像分割的嵌套U-Net结构

不得不说,妙啊!