UNet++: A Nested U-Net Architecture for Medical Image Segmentation
原文链接
作者解读
一.UNet++网络架构
UNet++与UNet相比re-designed skip pathways(图中绿色和蓝色部分)且运用了deep supervision(图中红色部分)。
Re-designed skip pathways
UNet网络并不是越深越好,不同数据集的最优深度是不一样的。给定一个新的数据集,并不知道多深的网络是最优的,因此,需要设计一个网络使得它能够学习不同深度特征。如下图所示,可以在这个网络里面找到1~4层的U-Net,这样做的好处是不管哪个深度的特征有效,全部包含在网络中,让网络自己去学习不同深度的特征的重要性。同时,该网络共享了特征提取器,不需要训练大量的UNet,而是只训练一个encoder,它的不同层次的特征由不同的decoder路径来还原。然而上述网络是不能训练的,因为不会有任何梯度经过红色区域。
为解决网络不能训练的问题,将网络结构进行redesign,如下图所示。将网络中的长连接替换成了一系列短连接,然而长连接是非常必要的,它联系了输入图像的很多信息,有助于还原降采样所带来的信息损失。
因此,提出了既包含长连接又包含短连接的网络结构。将原来空心的UNet填满,抓取不同层次的特征,将它们通过特征叠加的方式整合。这种Redesign改变了编码和解码子网络的连接性。密集的卷积块使编码器的feature map和解码器的feature map在语义级别相似。假设当接收到的编码器的feature map和相应的解码器的feature map在语义上相似时,优化器将更容易进行优化。
用公式表示skip pathways:函数H(·)表示卷积操作,且紧跟着一个激活函数;U(·)表示一个上采样层,[ ]表示连接层。j=0的节点从编码器的前一层只接收一个输入,j>1的节点共接收j+1个输入。
Deep supervision
深度监督使模型可在两种模式下运行:(1)精确模式:最终输出是所有分割分支的输出平均。(2)快速模式:最终的分割图只能从分割分支中选择一个。
(1)整合不同层次的特征,精度提升。
(2)灵活的网络结构配合深监督,在可接受的精度范围内可利用剪枝大幅度缩减参数量。
实现深监督的具体操作是在每个目标结点X0,1,X0,2,X0,3,X0,4后加一个1×1的卷积核followed by a sigmoid activation function,将二元交叉熵和Dice系数的组合作为损失函数添加到四个语义层次中的每一个,监督每一个分支的输出。
这种深监督带来了一个优势: 剪枝。
为什么剪枝???
在深监督过程中,每个子网络的输出均是图像的分割结果,所以如果小的子网络的输出结果已经足够好了,便可以剪掉网络多余的部分。
在测试阶段,由于输入的图像只会前向传播,剪枝之后对前面的输出无任何影响,而在训练阶段,因为既有前向传播又有反向传播,被剪掉的部分可以帮助其他部分更新权重,如果直接拿L1,L2,L3训练,就相当于只训练了不同深度的U-Net,最后的结果会很差。即在测试阶段剪掉的部分对剩余结构不做影响,在训练阶段,剪掉的部分对剩余部分有影响。因此,可以在测试的时候剪枝,根据验证集的结果觉得剪枝多少。
在(a)(/c)(d)三个数据集中,L2的结果和L4非常接近,而L2的参数量为0.5M,L4的参数量为9M,模型的内存可以省下18倍。
二.UNet++优势(对比UNet)
(1)跳跃路径上的卷积层,使编码器的feature map和解码器的feature map在语义级别相似。
(2)跳跃路径上有密集的skip-connection,改善了梯度流。
(3)使用deep-supervision,可以进行模型修剪和改进。
三.Experiments
使用了四种医学图像数据集对UNet++进行评估,并与UNet和wide UNet(强行增加U-Net里面的参数量,增加其每层卷积核个数得到)的表现对比。
**为什么用wide UNet??因为unet++相比unet参数多,为证明并不是无脑增加参数量,模型效果就会好,在设计时使wide unet的参数比unet++稍微多一点。
四种医学图像数据集:小剂量的CT胸透扫描瘤分割、显微镜里的细胞核分割、腹部CT的肝脏分割以及结肠息肉分割
实现细节
监控Dice系数和Intersectionover Union (IoU),并在验证集上使用了提前终止机制。
使用了Adam optimizer,学习率为3e-4。
沿着skip pathway的卷积层使用大小为3×3(3D肺结节分割使用的是3×3×3)的k个kernel,其中k=32×2i。
为实现Deep supervision,在每个目标结点上附加一个1×1的卷积层和一个sigmoid激活函数。
最终,对于给定的一个输入图像,UNet++生成四个分割图,将四个分割图进一步平均,得到最终的分割图。