本期为大家分享一篇经典文章:UNET 3+: A FULL-SCALE CONNECTED UNET FOR MEDICAL IMAGE SEGMENTATION。在论文中作者提出了一种基于UNet和UNet++改进的网络结构,在生物医学图像分割领域有着非常好的应用。
一、研究背景
UNet++虽然名义上通过嵌套和密集跳跃连接进行了多尺度信息的利用,但是从本质上看基本都是短连接,基本上都对解码特征进行了再次处理,再加上各个连接的融合,多尺度信息的原始特征几乎没有得到特别好的利用,信号处理有些矫枉过正或是丢失。尽管取得了良好的性能,但这种方法仍然不能从多尺度中探索足够的信息。
UNet3+在UNet++的基础上改进了多尺度信息的融合方式,并辅之以深度监督和一个分类指导模块,实现了生物医学图像分割进度和分割效率的同步提升。
二、网络结构
1.改进的跳跃路径
图 1.UNET 3+网络架构图
图1为UNet3+网络结构总体的一个概述。对比UNet++的网络结构,我们可以发现它去掉了之前的稠密卷积块,而是提出了一种全尺寸跳跃连接,表示为图中的带箭头虚线部分。无论是连接简单的UNet,还是连接紧密嵌套的UNet++,都缺乏从全尺度探索足够信息的能力,未能明确了解器官的位置和边界。为了弥补UNet和UNet++的缺陷,UNet3+中的每一个解码器层都融合了来自编码器中的小尺度和同尺度的特征图,以及来自解码器的大尺度的特征图,这些特征图捕获了全尺度下的细粒度语义和粗粒度语义。
举例来讲,图2展示了构造X3De特征图的具体过程。在这个过程中为了方便对不同来源的特征图进行融合,我们需要统一这些特征图的尺寸和通道数。论文作者发现64个3×3大小的滤波器进行卷积表现效果较好,卷积后便产生64个通道的特征图,尺寸与同尺度解码器特征图相一致。从图2中我们可以看出,X3De特征图的来源分为三部分。第一部分为编码器中小尺度特征图,也就是图中X1Ee和X2Ee的特征图,作者通过池化操作(使用最大池化法)来缩小特征图的尺寸,通过卷积操作改变特征图的通道数。第二部分为编码器中同尺度特征图,也就是图中X3Ee的特征图,作者通过卷积操作来改变这两个特征图的通道数。第三部分为解码器中大尺度特征图,也就是图中X5De和X4De的特征图,作者通过上采样操作(双线性插值法)来扩大特征图的尺寸,通过卷积操作来改变特征图的尺寸和通道数。最终我们得到五张通道数均为64的特征图,通过通道层面拼接的方式得到一张通道数为320的特征图,然后在经过320个3*3的卷积核进行卷积,最后经过BN+ReLU操作得到X3De的特征图,实现全尺寸特征融合。
图 2.多尺寸特征融合示意图
解码器其他部分的特征图也是按照同样方式得到的。具体的公式表示如下。其中,函数C表示卷积操作,函数H表示特征聚合机制(一个卷积层+一个BN+一个ReLU),函数D和函数U分别表示上采样和下采样操作,[ ]表示通道维度拼接融合。
2.参数量问题
值得一提的是UNet3+的参数量少于UNet和UNet++。其中UNet++的参数明显要多于UNet,我们就直接比较UNet和UNet3+。二者的编码阶段完全一致,参数量相同,因此为我们只需要比较二者解码阶段的参数量。下面两个公式分别为二者解码阶段参数量的计算公式。其中,DF表示卷积核的大小,函数d表示求出节点的深度。
公式看起来较为复杂,本文通过计算解码器的第三阶段的参数多少来说明参数量计算方式。参数主要存在于卷积操作中,池化操作和上采样操作均不涉及参数。其中,一次卷积操作使用的参数量=卷积核大小的平方*输入图像的通道数*输出图像的通道数。
图3为UNet解码器第三阶段的具体操作,其中包含三次卷积(上采样之后特征拼接之前也有一次卷积操作)。因此可以计算第三阶段参数量=3*3*(512*256+256*256+512*256)=2949120。
图2为UNet3+解码器第三阶段的具体操作,其中包含六次卷积。因此可以计算第三阶段参数量=3*3*(64*64+128*64+256*64+512*64+1024*64+320*320)=2064384。
两相对比之下,可以发现,UNet3+第三阶段的参数量是要小于UNet的。根据前面的公式计算二者解码阶段的全部参数量也同样可以得到UNet3+参数要明显少于UNet的结论。其中有一点需要注意的是,在UNet论文中紧跟上采样操作后的卷积操作使用的是2*2的卷积核,而在UNet3+中则使用的是3*3的卷积核,为了互相比较,上述计算中均使用的是3*3的卷积核。
图 3.UNET 卷积示意图
3.深度监督
为了从全尺度的聚合特征图中学习层次表示,UNet3+进一步采用了全尺度深度监督。其与UNet++中的深度监督不同之处在于监督的位置不同,前者监督的是网络解码器每个阶段输出的特征图,后者监督的是网络第一层中的四张特征图(其中三张为跳跃连接中卷积块的输出特征图,一张为解码器最后输出的特征图)。此外,在UNet3+中,为了实现深度监督,每个解码器阶段的最后一层被送入一个普通的3×3卷积层,然后是一个双线性上采样和一个sigmoid函数(这里的上采样是为了放大到全分辨率)。
4.分类引导模块
在大多数医学图像分割中,非器官图像出现假阳性是不可避免的。它很有可能是由于来自背景的噪声信息停留在较浅的层次,导致过度分割的现象。为了实现更精确的分割,我们尝试通过增加一个额外的分类任务来解决这个问题,这个任务是为预测输入图像是否有器官而设计的。
图 4.分类引导模块示意图
图4展示了分类引导模块的具体设计。其中,X5En的输出特征图通过dropout、1*1的卷积、可调节的最大池化和一个sigmoid函数等一系列操作后输出得到一组表示输入图像是否有器官的概率值,并通过交叉熵损失函数来训练这个分类过程。然后这组概率通过Argmax函数转化为0/1的单值函数,表示有/无器官。之后我们将单个分类输出与解码器各个阶段的侧分割输出相乘,实现了对非器官图像过分割的指导。
三、损失函数
结合focal损失函数,MS-SSIM损失函数和iou损失函数,UNet3+作者提出一种混合损失函数用于分割在三个不同层次,即像素级别分割、块分割和图像级别分割,它能捕获大尺度的和精细结构清晰的界限。混合分割损失定义如下。
其中,focal损失函数是对交叉熵损失函数进行改进的一个损失函数,主要关注像素级别分割。iou损失函数则是关注图像级别分割。此外,为了进一步增强器官的边界,我们提出了一个多尺度结构相似度指数损失函数(ms-ssim)来赋予模糊边界更高的权重。受益于此,UNet3+将关注模糊边界,区域分布差异越大,MS-SSIM值越高。
四、实验结果
UNet3+作者使用这个网络架构分别在肝脏和脾脏两个数据集上进行了试验,来验证这个网络架构是否有效。
基于VGG-16和ResNet-101的主干,表1展示了UNet、UNet++和有无深度监督操作的UNet3+架构在肝脏和脾脏两个数据集上的参数数量和分割精度方面的差异。每种方法使用的损失函数都是focal loss。我们可以发现,参数方面,UNet3+远远小于其他两个网络,这也印证了我们之前提出的结论。在分割精度方面,有无深度监督操作的两个UNet3+网络效果都要明显好于其他两个网络,相较于UNet和UNet++,带有深度监督操作的UNet3+网络在两个数据集上执行的两个主干之间获得了3.15和2个点的平均改进。
表 1.分割效果对比表
图5进一步展示了三个网络的分割图像。其中,黄色部分代表假阴性,也就是前景被分成背景了,绿色部分代表假阳性,也就是背景被分成前景了,紫色部分表示真实分割。根据分割结果,我们可以看出UNet3+的结果与真实分割最为接近,而UNet++出现了过度分割的情况,UNet存在分割不足的情况。UNet3+和UNet++分割结果的比较也证明了前者引入分类引导模块有效地改善了非器官图像过度分割的情况。
图 5.分割结果图
表2展示了UNet3+和其他热门分割架构在肝脏和脾脏两个数据集上的分割精度比较。可以很明显地看出,UNet3+的效果要优于其他的网络架构。
表 2.分割效果对比表
五、论文总结
UNet3+在UNet和UNet++的基础上对于跳跃连接和深度监督进行了改进,实现了生物医学图像的分割精度和分割效率的同步提升。
UNet3+的改进主要有以下三部分。第一,引入全尺度的skip connection,该连接结合了来自全尺度特征图的低级语义和高级语义,充分利用了多尺度特征,并且参数更少。第二,进行了深度监督,从全面的聚合特征图中学习层次表示,优化了混合损失函数以增强器官边界。第三,提出分类引导模块,通过图像级分类联合训练,减少非器官图像的过度分割。
最后,通过在肝脏和脾脏数据集上进行的试验,我们发现UNet3+的表现得到了大幅提高并且超过了很多当下最先进最热门的网络架构。