UNet3+阅读笔记

Paper:https://arxiv.org/pdf/2004.08790
Code:https://github.com/ZJUGiveLab/UNet-Version

以前和大家分享过UNet和UNet++，UNet是一种编码器解码器结构的卷积神经网络，常用于医学图像分割；而UNet++是在UNet的基础上添加嵌套和密集跳过连接的卷积神经网络。这次给大家分享UNet系列的一篇新作UNet3+，它使用全尺度的跳跃连接把不同尺度的特征图相融合，并且通过深度监督从多尺度聚合的特征图中学习特征表达，还提出了一种混合损失函数，将分类任务和分割任务相结合可以增强器官边界和减少非器官图像的过度分割，从而获得更准确的分割结果。
UNet3+网络结构如下图1所示，通过重新设计跳跃连接来融合多尺度特征，并利用全面的深度监督，提供更少的参数，从而产生更准确的分割结果。
图1 UNet3+网络模型

1. 全尺度连接
   为了弥补UNet和UNet++不能精确分割图像中器官的位置和边界，UNet3+中每一个解码器都结合了全部编码器的特征，这些不同尺度的特征能够获取细粒度的细节和粗粒度的语义。比如，下图2展示了第三个解码器$X_{De}^3$的特征图构造，它不仅直接接收来自相同尺度编码器层的特征图$X_{En}^3$，还对上面两层进行不同尺度的最大池化操作，将$X_{En}^1$和$X_{En}^2$进行下采样，保证下采样后特征图的尺寸和$X_{En}^3$的大小一样；同时还对下面两层进行不同尺度的上采样操作，将$X_{En}^4$和$X_{En}^5$进行上采样，确保上采样后的特征图的尺寸和$X_{En}^3$的大小一样。最后，作者发现采用64个3x3的卷积核进行卷积效果较好，那么在编码器和解码器之间构建的全尺度连接采用的就是64个3x3的卷积核进行卷积，在对应的解码器模块进行连接时，先采用320个(64x5=320)3x3的卷积核进行卷积，然后在经过BN和ReLU层就得到$X_{De}^3$。
   
   图2 $X_{En}^3$的全尺度连接特征图

2. 全尺度监督
   为了从全尺度的聚合特征图中学习层次表达，UNet3+进一步采用了全尺度的深度监督。与UNet++对每个嵌套的子网络进行监督不同的是，在UNet3+中每一个解码器模块都有一个输出，与GT(ground truth)进行比较计算loss，从而实现全尺度的监督。具体的流程是，每一个解码器模块的最后一层被送入一个3x3的卷积层，然后是一个双线性上采样得到和GT一样大小的特征图，再使用sigmoid函数得到分割结果。
   为了实现器官边界的精确分割，作者提出了一个多尺度结构相似指数损失函数(MS-SSIM，如公式1所示)来给予模糊边界更高的权重。因此，UNet3+将更加关注模糊边界，区域分布差异越大，损失函数值(MS-SSIM)越大。
   
   其中，M表示尺度的总数量(为了和UNet、UNet++保持一致，作者将M设置为5)，${\mu }_p$,${\mu }_g$和${\sigma }_p$,${\sigma }_g$分别是p、g的均值和方差，${\sigma }_{pg}$表示它们的协方差。${\beta }_m$,${\gamma }_m$定义每个尺度中这两个部分的相对重要性，$C_1=0.01^2$和$C_2=0.03^2$是避免被0除。
   结合Focal loss, MS-SSIM loss和IOU loss，作者提出一种混合损失函数用于三个不同层次的分割任务(像素分割、块分割和图像分割)中，它能获取不同尺度目标的清晰的边界。混合损失函数的公式如公式2所示。
   

3. 分类引导模块
   为了防止非器官图像的过度分割，和提高模型的分割精度，作者通过添加一个额外的分类任务来预测输入图像是否有器官，从而实现更精准的分割。如图3所示，对不同尺度的解码器模块的结果进行一系列的操作，包括Dropout，1x1的卷积，max pooling和 sigmoid函数，就可以得到一个分类结果(0表示没有器官，1表示有器官)。然后，将分类结果和分割结果相乘，这样就可以抑制对非器官图像的过度分割，从而提高模型的分割精度，实现精准分割。
   
   图3 分类引导模块示意图

4. 实验结果
   图4展示了UNet、UNet++和UNet3+在肝脏和脾脏数据集上面的分割结果，我们可以看出UNet3+表现最好；图5展示了UNet3+和5个SOTA分割模型在肝脏和脾脏数据集上面的分割结果，我们可以看出UNet3+也表现的很好。

图4 UNet、UNet++和UNet3+在肝脏和脾脏数据集上面的分割结果

图5 UNet3+和5个SOTA在肝脏和脾脏数据集上面的分割结果