UNet++: A Nested U-Net Architecture for Medical Image Segmentation
UNet ++:用于医学图像分割的嵌套U-Net体系结构
附论文作者对于UNet++的讲解,这才是做科研的方法。https://zhuanlan.zhihu.com/p/44958351
Abstract
在本文中,我们介绍了UNet ++,这是一种用于医学图像分割的功能更强大的新体系结构。我们的架构本质上是一个深度监督的编码器/解码器网络,其中,编码器和解码器子网络通过一系列嵌套的密集跳跃路径连接。重新设计的跳跃路径旨在减小编码器和解码器子网的特征图之间的语义差距。我们认为,当来自解码器和编码器网络的特征映射在语义上相似时,优化器将处理更轻松的学习任务。我们已将UNet ++与跨多个医学图像分割任务的U-Net和广泛的U-Net架构进行了比较:低剂量胸部CT扫描中的结节分割,显微镜图像中的核分割,腹部CT扫描中的肝脏分割以及结肠镜检查视频中的息肉分割。我们的实验表明,经过深层监督的UNet ++分别比U-Net和广泛的U-Net分别获得3.9和3.4点的平均IoU增益。
Introduction
最先进的图像分割模型是编解码器体系结构的变体,如U-Net[9]和全卷积网络(FCN)[8]。这些用于分段的编码器-解码器网络共享一个关键相似性:跳过连接,其将来自解码器子网络的深层、语义、粗粒度特征映射与来自编码子网络的浅、低级别、细粒度特征映射相结合。事实证明,跳过连接在恢复目标对象的细粒度细节方面是有效的;即使在复杂背景下也能生成具有精细细节的分割掩模。跳过连接也是实例级分割模型(如MaskRCNN)成功的基础,它允许分割被遮挡的对象。可以说,自然图像中的图像分割已经达到了令人满意的性能水平,但是这些模型是否满足医学图像的严格分割要求呢?
分割医学图像中的病变或异常需要比自然图像更高的精确度。虽然精确的分割掩模在自然图像中可能不是关键,但即使医学图像中的边缘分割错误也会导致临床环境中的用户体验不佳。例如,结节周围细微的毛刺图案可能表明结节恶性;因此,从临床角度来看,将它们排除在分割掩膜之外将降低模型的可信度。此外,不准确的分割还可能导致随后的计算机生成的诊断中的重大变化。例如,纵向研究中对结节生长的错误测量可能导致将不正确的肺-RADS分类分配给筛查患者。因此,需要设计更有效的图像分割体系结构,能够有效地恢复医学图像中目标对象的精细细节。
为了满足在医学图像中进行更精确分割的需求,我们提出了UNet ++,这是一种基于嵌套和密集跳过连接的新分割架构。我们的体系结构背后的基本假设是,在融合来自解码器网络的相应语义丰富的特征图之前,逐渐丰富来自编码器网络的高分辨率特征图时,该模型可以更有效地捕获前景对象的细粒度细节。我们认为,当来自解码器和编码器网络的特征映射在语义上相似时,网络将处理更轻松的学习任务。这与U-Net中常用的普通跳过连接相反,后者直接从编码器向解码器网络快速转发高分辨率特征图,从而导致语义相异的特征图融合。根据我们的实验,建议的体系结构是有效的,与U-Net和广泛的U-Net相比,可显着提高性能。
Related Work
Long等人。[8]首先介绍了完全卷积网络(FCN),而UNET是由Ronneberger等人提出的。[9]。它们都有一个共同的关键理念:跳过连接。在FCN中,上采样的特征映射与从编码器跳过的特征映射相加,而U-Net将它们连接起来,并在每个上采样步骤之间添加卷积和非线性。跳过连接已被证明有助于恢复网络输出的全空间分辨率,使得完全卷积方法适用于语义分割。受到DenseNet体系结构的启发[5],Li等人。[7]提出了一种用于肝和肝肿瘤分割的H-denseunet算法。本着同样的精神,Drozdzalet al.。[2]系统地研究了跳跃连接的重要性,并在编码器中引入了短跳跃连接。尽管上述体系结构之间存在细微差异,但它们都倾向于融合来自编码器和解码子网络的语义不同的特征映射,根据我们的实验,这可能会降低分割性能。
最近的另外两个相关工作是GridNet[3]和Mask-RCNN[4]。GridNet是一种编码器-解码器体系结构,其中特征地图以网格方式连接,概括了几种经典的分段体系结构。然而,GridNet在跳过连接之间缺少上采样层;因此,它不代表UNET++。MASK-RCNN可能是用于目标检测、分类和分割的最重要的元框架。我们要注意的是,UNET++可以很容易地部署为Mask-RCNN中的主干架构,只需用建议的嵌套密集跳过路径替换普通跳过连接即可。由于篇幅所限,我们无法包含以UNET++为主干架构的Mask RCNN的结果;但是,感兴趣的读者可以参考补充材料了解更多详细信息。
Proposed Network Architecture: UNet++
图1a显示了所建议体系结构的高级概述。如图所示,UNet ++从编码器子网或骨干网开始,然后是解码器子网。 UNet ++与U-Net(图1a中的黑色组件)的不同之处在于重新设计的跳过路径(以绿色和蓝色显示)连接了两个子网,并使用了深度监控(红色)。
图1:(A)UNET++由编码器和解码器组成,编码器和解码器通过一系列嵌套的密集卷积块相连。UNET++的主要思想是在融合之前弥合编码器和解码器的特征映射之间的语义鸿沟。例如,使用具有三个卷积层的密集卷积块来弥合(X0,0,X1,3)之间的语义鸿沟。在图形摘要中,黑色表示原始U网,绿色和蓝色表示跳跃路径上密集的卷积块,红色表示深度监控。红色、绿色和蓝色分量将UNET++与U-net区分开来。(B)详细分析UNET++的第一条跳跃路径。©如果在深度监督下接受培训,UNET++可以在推理时进行剪枝。(UNet++可以在测试时进行剪枝,怎么剪枝,一次减掉一个level,也就是UNet++L2如果可以达到一定的效果,就没必要到UNet++L4,UNet++L4的参数比UNet++L2的参数多很多,但是只能在测试时剪枝,训练时UNet++L4对于前面的部分是有影响的,因为有反向传播。)
Re-designed skip pathways
重新设计的跳跃路径改变了编码器和解码器子网络的连通性。在U-Net中,编码器的特征图直接由解码器接收,而在UNET++中,编码器的特征图经历了一个密集的卷积块,其卷积层数取决于金字塔的层数。例如,节点X0,0和X1,3之间的跳跃路径由具有三个卷积层的密集卷积块组成,其中每个卷积层之前有级联层,该级联层将来自同一密集块的前一卷积层的输出与较低密度块的相应上采样输出进行融合。本质上,密集卷积块使编码器特征映射的语义级更接近于在解码器中等待的特征映射的语义级。假设当接收到的编码器特征图和对应的解码器特征图在语义上相似时,优化器将面临更容易的优化问题。
形式上,我们将跳过路径表示如下:设xi,j表示节点Xi,j的输出,其中i索引沿着编码器的下采样层,j索引沿着跳过路径的密集块的卷积层。由xi,j表示的特征图计算为
其中函数H(·)是卷积运算,然后是激活函数,U(·)表示上采样层,[]表示级联层(cat操作)。基本上,级别j=0的节点只接收来自编码器的前一层的一个输入;级别j=1的节点接收两个输入,这两个输入都来自编码器子网络,但是处于两个连续的级别;并且级别j>1的节点接收j+1个输入,其中j个输入是同一跳过路径中的前j个节点的输出,最后一个输入是来自较低跳过路径的上采样输出。所有先验特征图累积并到达当前节点的原因是我们在每条跳过路径上使用密集卷积块。图1b进一步阐明了公式。公式1通过显示特征地图如何通过UNET++的顶部跳过路径。
Deep supervision
我们建议在UNET++中使用深度监督[6],使模型能够在两种模式下运行:1)精确模式,其中所有分割分支的输出被平均;2)快速模式,其中最终分割图仅从其中一个分割分支中选择,该模式的选择决定了模型剪枝的程度和速度增益。图1c显示了在快速模式下选择分段分支如何导致不同复杂性的架构。
由于跳跃路径的嵌套,UNET++生成了多个语义层次的全分辨率特征图,{x0,j,j∈{1,2,3,4}},可以进行深度监控。我们将二进制交叉熵和Dice系数的组合作为损失函数添加到上述四个语义层次中的每一个层次,其描述为
其中,ˆYb带Yb分别表示图像的展平预测概率和展平ground truth中的第b个,N表示batchsize。
总之,如图1A所示,UNET++与原始U-Net在三个方面不同:1)在跳过路径上具有卷积层(以绿色显示),其桥接编码器和解码器特征地图之间的语义鸿沟;2)在跳过路径上具有密集的跳过连接(以蓝色显示),这改善了梯度流;以及3)具有深度监控(以红色显示),这使得能够进行模型修剪并改进,或者在最坏情况下获得与仅使用一个损失层相当的性能。
Experiments
Datasets:
如表1所示,我们使用四个医学成像数据集进行模型评估,涵盖不同医学成像模式的病变/器官。有关数据集和相应数据预处理的更多详细信息,请读者参阅补充材料。
Baseline models:
为了进行比较,我们使用了原始的U-Net和定制的宽U-Net架构。我们选择U-Net是因为它是图像分割的通用性能基准。我们还设计了一个宽广的U网,其参数数量与我们建议的体系结构相似。这是为了确保我们的体系结构带来的性能提升不是简单地因为增加了参数数量。表2详细说明了U-Net和Wide U-Net架构。
Implementation details:
我们监测Dice系数和IOU,并在验证集上使用提前停止机制。我们还使用了学习速率为3e-4的Adam优化器。表2显示了UNET和Wide U-Net的架构细节。UNET++是从原始的U-Net架构构建的。沿着跳跃路径(xi,j)的所有卷积层使用大小为3×3(或对于3D肺结节分割为3×3×3)的k个核,其中k=32×2^i。为了实现深度监督,向每个目标节点添加了1×1卷积层,其后是S形激活函数:{x0,j | j∈{1,2,3,4}}。因此,UNET++在给定输入图像的情况下生成四个分割图,这些分割图将被进一步平均以生成最终的分割图。更多细节可以在github.com/nesteed-UNET上找到。
Results:
表3比较了U-Net、Wide U-Net和UNET++在肺结节分割、结肠息肉分割、肝脏分割和细胞核分割任务中的数量参数和分割精度。可以看出,宽U-Net的性能一直优于U-Net,除了肝脏分割,这两种架构的性能相当。这一改进归功于宽U网中的更多参数。在没有深度监督的情况下,UNET++在性能上比UNET和Wide U-Net都有显著提高,IOU平均提高了2.8和3.3个点。有深度监督的UNET++比没有深度监督的UNET++平均提高0.6分。具体地说,深度监督的使用导致了肝和肺结节分割的显著改善,但对于细胞核和结肠息肉分割,这种改善消失了。这是因为息肉和肝脏在视频帧和CT切片上以不同的尺度出现;因此,使用所有分割分支的多尺度方法(深度监督)对于准确分割是必不可少的。图2显示了U-Net、Wide U-Net和UNET++的结果之间的定性比较。
分别对(a)细胞核,(b)结肠息肉,(c)肝和(d)肺结节分割任务进行剪枝后,UNet ++的复杂性,速度和准确性。推理时间是使用一个具有12 GB内存的NVIDIA TITAN X(Pascal)处理10k测试图像所花费的时间。
图2.U-Net,宽U-Net和UNet ++之间的定性比较,显示了息肉,肝脏和细胞核数据集的分割结果(仅2D用于独特的可视化)。
Conclusion
为了满足对更精确的医学图像分割的需求,我们提出了UNet ++。建议的体系结构利用了重新设计的跳过路径和深度监控的优势。重新设计的跳过路径旨在减小编码器和解码器子网络的特征图之间的语义差距,从而为优化器解决可能更简单的优化问题。深度监督还可以实现更精确的分割,尤其是对于出现在多个尺度的病变(例如结肠镜检查视频中的息肉)。我们使用四个医学影像数据集(包括肺结节分割,结肠息肉分割,细胞核分割和肝分割)评估了UNet ++。我们的实验表明,经过深层监督的UNet ++分别比U-Net和宽U-Net分别获得了3.9和3.4点的平均IoU增益。