UNet++:一种用于医学图像分割的嵌套U-Net体系结构

Abstract

在本文中,我们提出了一种新的、更强大的医学图像分割体系结构UNet++。我们的架构本质上是一种deeply-supervised encoder-decoder网络,其中编码器和解码器sub-networks通过一系列nested、dense skip路径连接。重新设计的skip路径连接旨在减少encoder和decoder子网络的feature map的semantic gap语义差距。我们认为,当来自decoder和encoder网络的feature map相似时,优化器会处理更容易的学习任务。我们在多个医学图像分割任务重评估了UNet++和U-Net和wide U-Net架构的比较:胸部低剂量CT扫描中的结节分割、显微镜图像中的细胞核分割、腹部CT扫描中的肝脏分割以及结肠镜检查视频中的息肉分割。我们的实验表明,深度监控的UNet++在U网和宽U网上分别获得了3.9和3.4点的平均IoU增益。

Introduction

最先进的图像分割模型时decoder-encoder结构的变体,如U-Net和全卷积网络(FCN)。这些用于分段的encoder-decoder有一个关键的相似之处:skip connection跳跃连接,它将来自于解码器子网的深度、语义、粗粒度coarse-grained特征图和编码器子网的浅层、低级、细粒度特征图结合起来。事实证明,skip connection在恢复目标对象的细粒度细节方面是有效的;即使在复杂背景下也能生成具有精细细节的分割掩模。skip connection也是实例级分割模型(如mask rcnn)成功的基础,它能够分割被遮挡的对象。可以说自然图像中的图像分割已经达到了令人满意的性能水平,但这些模型是否满足医学图像的严格分割要求?

分割医学图像中的损伤或异常需要比自然图像中期望的更高水平的精度。虽然精确的分割mask在自然图像中可能不重要,但医学图像中的分割错误也会导致临床环境中的用户体验较差。例如,结节周围细微的毛刺图案可能表明结节恶性肿瘤;因此,从临床角度来看,将它们从分割掩模中排除会降低模型的可信度。此外,不准确的分割也可能导致随后的计算机生成的诊断发生重大变化。例如,在纵向研究中对结节生长的错误测量可能导致将不正确的肺放射性疾病分类分配给筛查患者。因此,期望设计更有效的图像分割体系结构,其能够有效地恢复医学图像中目标对象的精细细节。

为了解决医学图像中更精确分割的需要,我们提出了UNet++,一种基于nested and dense skip connections。我们架构背后的基本假设是:当来自于encoder网络的高分辨率特征图在与其对应的来自decoder网络的语义丰富的特征图所融合之前逐渐丰富时候,模型可以更有效的捕捉到前景对象的细粒度细节。我们认为,当来自解码器和编码器网络的特征图在语义上相似时,网络将处理更容易的学习任务。这与通常在U-Net中使用的普通skip连接形成对比,后者直接将高分辨率特征图从编码器快速转发到解码器网络,导致语义不同的特征图的融合。根据我们的实验,所建议的架构是有效的,产生了显著的性能增益。

2. Related Work

龙等人[8]首先引入了完全卷积网络(FCN),而UNet是由龙内伯格等人[9]引入的。他们都有一个关键的想法:skip connection。在FCN,向上采样的特征图与skipped from编码器的特征图相加,而U-Net将它们连接起来,并在每个向上采样步骤之间添加卷积和非线性。skip连接有助于在网络输出端恢复完整的空间分辨率,使完全卷积的方法适用于语义分割。受DenseNet架构[5]的启发,李等人[7]提出了用于肝脏和肝肿瘤分割的DenseNet。本着同样的精神,Drozdzalet等人[2]系统地研究了skip connection的重要性,并在编码器中引入了短跳跃连接。尽管上述体系结构之间存在细微差异,但它们都倾向于融合来自编码器和解码器子网络的语义不同的特征图,根据我们的实验,这可能会降低分割性能。

另外两个最近的相关工作是GridNet [3]和Mask-RCNN [4]。GridNet是一种编码器-解码器架构,其中feature map以网格方式连接,概括了几种经典的分割架构。然而,GridNet在skip connection之间缺少上采样层;因此,它不代表UNet++。MASK-RCNN可能是对象检测、分类和分割最重要的元框架。

我们想指出的是,通过简单的用我们所提出的nested、dense skip connections替换普通的skip connection,可以很容易地将UNet++部署为Mask-RCNN的backbone框架;然而感兴趣的读者可以参考补充资料了解更多的细节。

image.png
图1: (a) UNet++由编码器和解码器组成,它们通过一系列嵌套的密集卷积块连接。UNet++背后的主要思想是在融合之前,弥合编码器和解码器的特征图之间的语义鸿沟。例如,使用具有三个卷积层的密集卷积块来桥接(X0,0,X1,3)之间的语义间隙。在图形摘要中,黑色表示原始的U-Net,绿色和蓝色表示跳过路径上的密集卷积块,红色表示深度监控。红色、绿色和蓝色组件将UNet++与U-Net区分开来。(b)对UNet++的第一skip途径的详细分析。(c)如果在深度监督下训练,可以在推理时删减UNet++代码。

3. Proposed Network Achitecture: UNet++

图1a显示了建议的体系结构的high-level概述。如前所述,UNet++从编码器子网或主干网开始,然后是解码器子网。UNet++与U-Net(图1a中的黑色组件)的区别在于重新设计的连接两个子网络的跳过路径(以绿色和蓝色显示)和深度监控的使用(以红色显示)。

3.1 Re-designed skip pathways

重新设计的skip路径改变了编码器和解码器子网的连通性。在U-Net中,编码器的特征图直接在解码器中接收;然而,在U-Net++中,他们经历密集的卷积块,其卷积层数取决于pyramid level。例如,节点X0,0和X1,3之间的跳跃路径由具有三个卷积层的密集卷积块组成,其中每个卷积层之前是级联层,该级联层将来自相同密集块的前一卷积层的输出与较低密集块的相应上采样输出融合。本质上,密集卷积块使编码器特征图的语义级别更接近于解码器中等待的特征映射的语义级别。假设当接收的编码器特征图和相应的解码器特征图在语义上相似时,优化器将面临更容易的优化问题。

形式上,我们将skip路径公式化如下: 让xi,j 表示节点Xi,j的输出,i指沿着编码器的下采样层,j指沿着skip路径的密集块的卷积层。由xi、j表示的要素图堆栈计算如下:
image.png
其中函数H()是卷积运算,后面是激活函数,U()表示上采样层,[ ]表示级联层。基本上,j = 0级的节点只接收来自编码器前一层的一个输入;j = 1级的节点接收两个输入,两个输入都来自编码器子网络,但处于两个连续的级;并且j > 1级的节点接收j + 1个输入,其中j个输入是同一skip路径中的前j个节点的输出,最后一个输入是来自较低跳过路径的上采样输出。所有先前的特征映射累积并到达当前节点的原因是因为我们利用了沿着每个跳过路径的密集卷积块。图1b进一步阐明了等式。Eq 1通过显示功能图如何通过UNet++的顶部跳转路径。

3.2 Deep supervision

我们建议在UNet++中使用深度监控[6],使模型能够以两种模式运行:1)精确模式,其中对来自所有分割分支的输出进行平均;2)快速模式,其中最终分割图仅从一个分割分支中选择,其选择决定了模型修剪和速度增益的程度。图1c示出了快速模式中分段分支的选择如何导致不同复杂性的体系结构。
image.png

由于嵌套的skip路径,U-Net++在多个语义级别{x0,j , j ∈{1,2,3,4}}生成全分辨率的特征图,这些是可以接受深度监管的。我们将二元交叉熵和骰子系数的组合作为损失函数添加到上述四个语义级别中的每一个,描述如下:
image.png
其中,Yb和Yb分别表示一个批次的flatten预测概率和平坦地面真实值,N表示批次大小。

总之,如图1a所示,UNet++在三个方面不同于原始的U-Net:1)在跳过路径上具有卷积层(以绿色示出),这弥合了编码器和解码器特征映射之间的语义鸿沟;2)在skip路径上具有密集的skip连接(以蓝色显示),这改善了梯度流动;以及3)具有深度监控(以红色显示),如将在第4节中显示的,这使得模型修剪成为可能,并且改善了或者在最坏的情况下实现了与仅使用一个损失层相当的性能。

4 Experiment

数据集:如表1所示,我们使用四个医学成像数据集进行模型评估,涵盖来自不同医学成像模式的病变/器官。关于数据集和相应的数据预处理的更多细节,我们请读者参考补充材料。

基线模型:为了比较,我们使用了原始的U-Net和定制的wide U-Net架构。我们选择U-Net是因为它是图像分割的通用性能基线。我们还设计了一个wide U-Net,其参数数量与我们建议的架构相似。这是为了确保我们的架构产生的性能增益不仅仅是由于参数数量的增加。表2详细说明了U-Net和wide U-Net系结构。

实现细节:我们监控dice系数和联合交集(IoU),并在验证集上使用提前停止机制。我们也使用adam优化器,学习率为3e-4。UNet和wide U-Net的架构细节如表2所示。UNet++是从最初的U-Net架构构建的。沿跳跃路径(Xi,j)的所有卷积层使用大小为3×3(或3×3×3用于3D肺结节分割)的k个核,其中k = 32 × 2^i。为了实现深度监控,1×1卷积层后接sigmoid激活函数被附加到每个目标节点:{x0,j| j ∈ {1,2,3,4}}。因此,在给定输入图像的情况下,UNet++生成四个分割图,这些分割图将被进一步平均以生成最终的分割图。更多细节可以在github.com/Nested-UNet.找到

image.png
image.png
图3:分别在(a)细胞核、(b)结肠息肉、(c)肝脏和(d)肺结节分割任务上修剪后的UNet++的复杂性、速度和准确性。推断时间是使用一个NVIDIA TITAN X(帕斯卡)和12 GB内存处理10k测试图像所花费的时间。

结果:表3在肺结节分割、结肠息肉分割、肝脏分割和细胞核分割任务的数量参数和分割精度方面比较了U-Net、宽U-Net和UNet++。正如所看到的,除了肝分割之外,宽U网的性能一直优于U网,肝分割是两种架构的性能相当。这种改进归因于宽U网的参数数量较大。没有深度监管的UNet++在UNet和宽U-Net上实现了显著的性能提升,IoU平均提高了2.8和3.3个点。深度监管的UNet++比没有深度监管的UNet++平均提高了0.6个百分点。具体来说,深度监控的使用导致肝脏和肺结节分割的显著改善,但是这种改善对于细胞核和结肠息肉分割而言是消失的。这是因为息肉和肝脏在视频帧计算机断层扫描中以不同的比例出现;因此,使用所有分割分支的多尺度方法(深度监督)对于精确分割是必不可少的。图2显示了U-Net、宽U-Net和UNet++的结果之间的定性比较。

模型修剪:图3显示了在应用不同级别的修剪后,UNet++的分割性能。我们使用UNet++ Lito表示在第一级修剪的UNet++(更多细节见图1c)。正如所看到的,Unet++ L3平均减少了32.2%的推理时间,而IoU仅降低了0.6个百分点。更积极的修剪进一步减少了推理时间,但代价是显著的准确性下降。