U-Net,NAS-Unet

https://arxiv.org/abs/1505.04597

https://hub.fastgit.org/milesial/Pytorch-UNet

人们普遍认为,成功地训练深层网络需要数千个带标注的训练样本。在本文中,我们提出了一种网络和训练策略,该策略依赖于强大的数据增强功能,以更高效地使用可用的带标注样本。该体系结构包括用于捕获上下文的收缩路径和用于实现精确定位的对称扩展路径。我们表明,这种网络可以从很少的图像进行端到端的训练,并且在电子显微镜堆栈中神经元结构分割的ISBI挑战上优于先前的最佳方法(滑动窗口卷积网络)。通过使用在透射光学显微镜图像(相衬度和DIC)上训练的相同网络,我们在这些类别中以巨大优势赢得了2015年ISBI细胞追踪挑战赛。而且,网络速度很快。在最近的GPU上,512x512图像的分割只需不到一秒钟的时间。完整实现(基于Caffe)和经过培训的网络可在http://lmb.informatik.uni-freiburg.de/people/ronneber/u-net.

1 引言

在过去两年中,深度卷积网络在许多视觉识别任务中的表现超过了最新水平,例如[7,3]。虽然卷积网络已经存在很长时间[8],但由于可用训练集的大小和考虑的网络的大小,它们的成功受到限制。Krizhevsky等人[7]的突破是由于在ImageNet数据集上对一个包含8层和数百万参数的大型网络进行了监督训练,该网络包含100万张训练图像。从那时起,训练了更大更深入的网络[12]。

卷积网络的典型用途是分类任务,其中图像的输出是单个类别标签。然而,在许多视觉任务中,特别是在生物医学图像处理中,期望的输出应该包括定位,即,应该为每个像素分配一个类别标签。此外,在生物医学任务中,成千上万的训练图像通常是遥不可及的。因此,Ciresan等人[1]在滑动窗口设置中训练网络,通过提供像素周围的局部区域(图像块)作为输入来预测每个像素的类别标签。首先,这个网络可以定位。其次,以图像块为单位的训练数据远大于训练图像的数量。由此产生的网络以巨大的优势赢得了2012年ISBI的EM细分挑战。

显然,Ciresan等人[1]的策略有两个缺点。首先,它相当慢,因为网络必须为每个图像块单独运行,并且由于重叠的图像块而存在大量冗余。其次,在定位准确性和上下文使用之间存在一种权衡。较大的图像块需要更多的最大池化层,这会降低定位精度,而较小的图像块只允许网络看到很少的上下文。最近的方法[11,4]提出了一种分类器输出,该输出考虑了来自多个层的特征。良好的定位和上下文的使用是可能的。

在本文中,我们是基于一个更优雅的结构,即所谓的“全卷积网络”[9]。我们修改和扩展了这个架构,使得它可以处理很少的训练图像,并产生更精确的分割;参见图1,[9]中的主要思想是通过连续的层来补充通常的收缩网络,其中池化运算符被上采样运算符取代。因此,这些层提高了输出的分辨率。为了定位,收缩路径的高分辨率特征与上采样输出相结合。随后,连续的卷积层可以学习根据该信息集成更精确的输出。

图1.U-net架构(最低分辨率下32x32像素的示例)。每个蓝色框对应一个多通道特征图。通道的数量在框的顶部表示。x-y尺寸位于框的左下边。白色框表示复制的特征图。箭头表示不同的操作。

我们架构中的一个重要修改是,在上采样部分,我们还有大量的特征通道,允许网络将上下文信息传播到更高分辨率的层。因此,扩展路径或多或少与收缩路径对称,并产生u形结构。该网络没有任何全连接的层,并且仅使用每个卷积的有效部分(valid part),即,分割映射仅包含像素,对于这些像素,完整上下文在输入图像中可得。该策略允许通过重叠平铺策略对任意大的图像进行无缝分割(见图2)。为了预测图像边界区域中的像素,通过镜像复制输入图像来外插缺失的上下文。这种平铺策略对于将网络应用于大型图像非常重要,否则分辨率将受到GPU内存的限制。

图2:用于无缝分割任意大图像的重叠平铺策略(此处为EM堆栈中神经元结构的分割)。黄色区域的分割预测需要蓝色区域内的图像数据作为输入。缺少的输入数据通过镜像复制进行推断

对于我们的任务,可用的训练数据非常少,我们通过对可用的训练图像应用弹性变形来使用过度的数据增强。这允许网络学习对这种变形的不变性,而不需要在标注图像语料库中看到这些变换。这在生物医学分割中尤其重要,因为变形曾经是组织中最常见的变化,可以有效地模拟真实的变形。数据增强对于学习不变性的价值已在Dosovitskiy等人[2]的无监督特征学习范围内得到证明。

许多细胞分割任务中的另一个挑战是分离同一类别的接触对象;参见图3。为此,我们建议使用加权损失,其中,接触单元之间的分离背景标签在损失函数中获得较大权重。

图3.用DIC(微分干涉对比)显微镜记录的玻璃上的HeLa细胞。(a) 原始图像。(b) 覆盖真值分割。不同的颜色表示HeLa细胞的不同实例。(c) 生成的分割掩膜(白色:前景,黑色:背景)。(d) 使用像素级的权重损失映射,强制网络学习边界像素。

所得到的网络适用于各种生物医学分割问题。在这篇论文中,我们展示了EM堆栈中神经元结构分割的结果(ISBI 2012年开始的一项持续竞争),我们在这方面的表现超过了Ciresan等人的网络[1]。此外,我们还展示了2015年ISBI细胞追踪挑战赛的光学显微镜图像中的细胞分割结果。在这里,我们在两个最具挑战性的2D透射光数据集上以巨大优势获胜。

2    网络架构

网络架构如图1所示。它由一条收缩路径(左侧)和一条扩展路径(右侧)组成。收缩路径遵循卷积网络的典型结构。它包括重复应用两个3x3卷积(无填充的卷积),每个卷积后面都有一个校正线性单元(ReLU)和一个2x2最大池化运算,步长2用于下采样。在每个下采样步骤中,我们将特征通道的数量增加一倍。扩展路径中的每一步包括特征图的上采样,然后是2x2卷积(“上卷积”),将特征通道的数量减半,与收缩路径中相应的裁剪的特征图进行拼接,以及两个3x3卷积,每个卷积后是一个ReLU。由于每次卷积都会丢失边界像素,因此需要进行裁剪。在最后一层,使用1x1卷积将每个64分量特征向量映射到所需数量的类。该网络总共有23个卷积层。

为了实现对输出的分割映射的无缝平铺(见图2),选择输入平铺大小非常重要,这样所有2x2 最大池化操作都可以应用于即使x-和y-是奇数的层。

3 训练

输入图像及其相应的分割图用于使用Caffe的随机梯度下降实现对网络进行训练[6]。由于使用无填充的卷积,输出图像比输入图像小一个恒定的边界宽度。为了最大限度地减少开销并最大限度地利用GPU内存,我们倾向于使用较大的输入块而不是较大的批处理大小,从而将批处理减少到单个图像。因此,我们使用高动量(0.99),以便大量以前看到的训练样本确定当前优化步骤中的更新。

能量函数由最终特征映射上的像素级软最大值与交叉熵损失函数相结合来计算。软最大值(softmax)定义为,其中表示在位置,处的特征映射的激活输出。是类别个数,是近似最大函数。也就是,就是,第个特征映射的有最大激活输出,而其它的,和近似为0。然后,交叉熵在每个位置惩罚与1的偏差:

    

其中 是每个像素的真值标签,而是一个加权映射,我们是为了在训练中赋予某些像素更多的重要性。

我们预先计算每个真值分割的加权映射,以补偿训练数据集中某类像素的不平衡,并迫使网络学习我们在接触单元之间引入的小分离边界(见图3c和d)。

使用形态学运算计算分离边界。然后将加权映射计算为:

其中,wc:Ω → R是平衡等级频率的权重图,d1:Ω → R表示到最近单元格边界的距离,d2:Ω → R到第二个最近单元格边界的距离。在我们的实验中,我们设置w0=10和σ≈ 5像素。

在具有许多卷积层和通过网络的不同路径的深度网络中,权重的良好初始化非常重要。否则,网络的某些部分可能会提供过多的激活,而其他部分则不会提供。理想情况下,应调整初始权重,以使网络中的每个特征映射具有近似单位方差。对于具有我们架构的网络(交替卷积和ReLU层),这可以通过从标准偏差为p2/N的高斯分布中提取初始权重来实现,其中N表示一个神经元的传入节点数[5]。例如,对于前一层中的3x3卷积和64个特征通道,N=9·64=576。

3.1    数据增强

当只有很少的训练样本可用时,数据增强对于教导网络所需的不变性和鲁棒性是至关重要的。万一

显微图像我们主要需要平移和旋转不变性以及对变形和灰度变化的鲁棒性。特别是训练样本的随机弹性变形似乎是训练具有很少标注图像的分割网络的关键概念。我们使用3×3的粗网格上的随机位移向量生成平滑变形。位移从标准偏差为10像素的高斯分布中采样。然后使用双三次插值计算每像素位移。收缩路径末端的退出层执行进一步的隐式数据增强。

4 实验

我们演示了u-net在三种不同的分割任务中的应用。第一项任务是在电子显微镜记录中分割神经元结构。图2显示了数据集和我们获得的分割的示例。我们提供了完整的结果作为补充资料。该数据集由2012年ISBI启动的EM细分挑战[14]提供,目前仍在接受新的贡献。训练数据是一组30幅图像(512x512像素),来自果蝇一龄幼虫腹神经索(VNC)的连续切片透射电子显微镜。每幅图像都带有相应的完整标注的细胞(白色)和膜(黑色)的基本真实分割图。测试集是公开的,但是它的分割图是保密的。通过将预测的膜概率图发送给组织者,可以获得评估结果。通过在10个不同级别对地图进行阈值化,并计算“扭曲误差”、“兰德误差”和“像素误差”来进行评估[14]。

u-net(输入数据的7个旋转版本的平均值)在不进行任何进一步的预处理或后处理的情况下实现了0.0003529的翘曲误差(新的最佳分数,见表1)和0.0382的兰德误差。

这明显优于Ciresan等人[1]的滑动窗口卷积网络结果,其最佳提交值的翘曲误差为0.000420,兰德误差为0.0504。就随机误差而言,该数据集上唯一性能更好的算法使用了适用于Ciresan等人概率图的高度特定于数据集的后处理方法1。[1]。

我们还将u-net应用于光镜图像中的细胞分割任务。该分离任务是2014年和2015年ISBI细胞追踪挑战赛的一部分[10,13]。第一个数据集“PhC-U373”2包含相衬显微镜记录的在聚丙烯酰亚胺基质上的胶质母细胞瘤-星形细胞瘤U373细胞(见图4a、b和补充材料)。它包含35个部分标注的训练图像。在这里,我们实现了92%的平均IOU(“联合上的交集”),这明显优于83%的次优算法(见表2)。第二个数据集“DIC HeLa”3是通过差分干涉对比(DIC)显微镜记录的平板玻璃上的HeLa细胞(见图3、图4c、d和补充材料)。它包含20个部分标注的训练图像。在这里,我们实现了77.5%的平均IOU,这明显优于46%的次优算法。

5     结论

u-net体系结构在非常不同的生物医学分割应用程序上实现了非常好的性能。由于使用弹性变形增强数据,它只需要很少的带标注的图像,并且在NVidia Titan GPU(6 GB)上的训练时间非常合理,仅为10小时。我们提供完整的基于Caffe[6]的实现和经过培训的网络4。我们确信u-net体系结构可以轻松地应用于更多的任务。



NAS-Unet: Neural Architecture Search for Medical Image Segmentation

https://hub.fastgit.org/tianbaochou/NasUnet

神经结构搜索(NAS)在提高图像分类精度方面有着重要的进展。最近,一些工作尝试将NAS扩展到图像分割,这显示了初步的可行性。然而,它们都专注于自然场景中语义分割的搜索架构。在本文中,我们在搜索空间上设计了三种类型的基元操作集(primitive operation set),以自动找到用于语义图像分割特别是医学图像分割的两个单元结构DownSC和UpSC。受成功应用于各种医学图像分割的U-net网络结构及其变体的启发,我们提出了NAS-Unet,它由U形主干网络上相同数量的DownSC和UpSC堆叠而成。在搜索阶段,DownSC和UpSC的网络结构通过可微网络结构策略同时更新。我们在Promise12、Chaos和超声神经数据集(分别通过磁共振成像、计算机断层扫描和超声采集)上展示了所提方法的良好分割结果。在没有任何预训练的情况下,我们在PASCAL VOC2012上搜索的网络结构在上述三种类型的医学图像数据集上进行评估时,获得了比U-net及其变体更好的性能和更少的参数(约0.8M)。

1. 引言

随着医学影像分析设备的发展和普及,包括磁共振成像(MRI)、计算机断层扫描(CT)和超声,已成为医疗机构进行疾病诊断、手术计划和预后评估不可缺少的设备。磁共振成像是放射成像领域应用最广泛的技术。MRI成像的一个突出特点是成像序列的多样性。在MRI中,图像的对比度取决于相位对比度脉冲序列。最常见的脉冲序列是T1(自旋晶格;也就是说,在与静态磁场相同的方向上磁化)加权和T2加权自旋(自旋自旋;横向于静态磁场)。与CT相比,MRI可能产生不同的信息。MRI扫描可能会带来风险和不适。与CT扫描相比,MRI扫描通常需要更长的时间和更大的声音,并且通常需要受试者进入一个狭窄的限制性管道。超声成像(超声)利用高频声波观察身体内部。与CT和MRI不同,超声图像的分辨率相对较低。

医学图像分析是医学图像分析的第一步,有助于使图像更直观,提高诊断效率。医学图像分割是医学图像分析领域的关键步骤。为了为临床诊断和病理研究提供可靠的依据,帮助医生做出更准确的诊断,需要对医学图像中的聚焦部分进行分割并提取相关特征。最初,医学图像分析是通过顺序应用低级像素处理(例如,基于区域的方法[1]或基于阈值的方法[2])和数学建模来完成的,以构建解决特定任务的基于规则的复合系统[3]。这一时期的分割结果通常没有语义标记。在深度学习时代,图像分割通常是指语义分割,即在像素级识别图像(图像中每个像素所属的对象类别已标记)[4],[5]。例如,在图1中,左侧图像中的医学图像由膀胱壁和其他组织组成,右侧图像是其语义分割的结果,该语义分割将像素语义对象(即标记为黄色的膀胱壁)分割,类似地,其他组织被视为背景并标记为紫色。迄今为止,用于图像分析的最成功的深度学习模型是卷积神经网络(CNN)。

深度学习与大数据、云计算的相互促进,给计算机视觉带来了巨大的发展[6]。CNN是计算机视觉领域中最常用的神经网络,被提出用于解决图像分类问题。图像分割是自然图像和医学图像分析中的一项常见任务。为了解决这个问题,可以简单地使用CNN对图像中的每个像素进行单独分类,方法是在特定像素周围提取面片,并生成与输入图像大小相同的多通道似然图。然而,如果一直保持特征图的尺寸不变,则会消耗大量内存。通常,在几个卷积层之后添加下采样层(如最大池和平均池),以降低特征映射的维数并细化高层上下文。不幸的是,这可能导致输出的分辨率远低于输入。FCNs(完全卷积网络)[7]是为防止分辨率下降而提出的几种方法之一。这是第一个通过将完全连接的层替换为卷积层之后的一系列上采样层来训练端到端像素预测的工作。经典的CNN通常使用一个完全连接的层,在最后一个卷积层之后获得固定长度的特征向量,并将其放入分类器(例如softmax层)。相比之下,FCNs可以接受任意大小的输入图像-最后一个卷积层之后的上采样层可以将其输入的维数恢复为与输入图像相同的维数,从而可以为每个像素生成预测,同时保留原始输入图像中的空间信息,最后通过像素分类将上采样特征映射到期望的像素图像分割。与FCNs类似,U-net[8]由卷积层、下采样层和上采样层组成。与FCN不同,U网络中的下采样层和上采样层以及它们之间的卷积层的数量是相同的。此外,U-net使用跳连操作来连接每对下采样层和上采样层,这使得空间信息直接应用到更深的层和更准确的分割结果。

从最早的LeNet[9]到AlexNet[10]、VggNet[11]、GoogleNet[12]、ResNet[13]和最近的DenseNet[14],CNN模型的性能越来越强大和成熟。许多作品都为特定任务设计了网络结构[15]、[16]。这些流行的网络结构目前由行业专家和学者设计了数月甚至数年。这是因为设计性能优异的网络网络结构通常需要大量该领域的知识。一般的研究人员不具备这种能力,设计过程耗时且劳动密集。基于此,当前卷积神经网络的焦点已转移到神经结构搜索(NAS)[17]。NAS可以看作是AutoML(自动机器学习)的子领域,与超参数优化和元学习有很大的重叠。目前对NAS的研究主要集中在三个方面:搜索空间、搜索策略和性能评估策略。搜索空间定义了理论上可以表示哪些网络结构。结合关于适合任务的属性的先验知识可以减少搜索空间的大小并简化搜索任务。例如,在图像分类中,搜索空间包括在每个搜索步骤选择基本操作和用于定义宏观网络的先验主干结构。搜索策略详细说明了如何探索搜索空间。NAS的目标通常是找到对未见数据具有高评估性能的网络结构(例如,将训练数据集拆分为训练和验证,并搜索训练但通过验证进行评估的网络结构)[17]。在NAS上已经做了很多工作,其中大部分工作集中在图像分类任务[18]–[23]。

尽管NAS在计算机视觉领域具有巨大潜力,但真正的前景取决于它可以扩展到处理图像分类以外的视觉任务,特别是计算机视觉核心问题,如图像语义分割,实例分割和目标检测依赖于高分辨率图像输入和多尺度图像表示。将NAS直接从图像分类引入到图像语义分割中是不可行的:首先,分类任务的搜索空间与分割任务的搜索空间明显不同;其次,从低分辨率图像直接迁移到高分辨率图像是不好的[24]。解决上述两个问题的一个合理想法是为图像分割建立特定的搜索空间,并尝试解决上述两个问题并取得了一些成功——这正是最近工作所遵循的想法[24],[25]( “Searching for efficient multi-scale architectures for dense image prediction,” in NeurIPS, 2018;Auto-deeplab)。

在本文中,我们提出了新的用于医学图像分割的基本操作集。受U-net及其变体在医学图像分割方面的成功启发,我们使用一个U型架构作为主干网络(即宏观网络),并在PASCAL VOC 2012[26]上并行搜索两个基于单元的架构(下采样单元和上采样单元),分别表示为DownSC和UpSC。最后,我们得到了表示为NASUnet的网络结构,它由相同数量的DownSC和UpSC堆叠而成。我们的工作表明,在我们前面提到的所有类型的医学图像数据集中,NasUnets的参数使用效率更高,性能远远优于U-Net和FC DenseNet[27](一种变体U-Net)。总之,我们的贡献如下:

1) 本文是首次尝试将NAS应用于医学图像分割。

2) 在类U骨干网络上,我们分别为DownSC和UpSC提出了不同的原始操作集来进行搜索。在完成搜索后,我们根据经验发现,在我们的UpSC网络结构中,标准的跳连被cweight操作(参见V-A)所取代。

3) 我们表明,在我们评估的所有类型的医学图像分割数据集中,NAS Unet的性能优于U-Net及其变体(FC Densenet),而不使用任何预训练的主干。NAS Unet的训练时间接近U-Net,但参数量仅为6%。FC Densenet的内存成本是我们的两倍。

我们已经在发布了我们的代码https://github.com/tianbaochou/NasUnet.

2. 相关工作

A    基于现代CNNS的医学图像分割

据我们所知,Ciresan等人[28]。首先利用深度神经网络对医学图像进行分割。利用卷积神经网络对电子显微镜图像进行分割“面片”是完成分割的关键思想-为了分割整个堆栈,通过在像素周围提取面片,将分类器应用于滑动窗口中每个切片的每个像素。这种天真的滑动窗口方法的一个缺点是,来自相邻像素的输入面片有大量重叠和冗余计算。[28]中还指出,使用该方法分割整个堆栈会导致时间效率低下,四个GPU上的堆栈至少需要10分钟。Ronneberger等人[8]将完全连接的层重写为卷积,并尝试执行相同的任务以获得更好的结果。如图2(a)所示,作者将FCNs[7]的思想进一步推进,并提出了U-Net网络结构,他们的设计基于编码器-解码器网络框架:将输入图像放入编码器架构中提取高层上下文,然后将该上下文流到解码器架构中,以恢复空间信息和像素分类结果。虽然这不是第一次在卷积神经网络中使用编码器-解码器(例如Long等人[7]使用经过预训练的现代CNN网络作为编码器,使用“上”卷积层作为解码器(FCNs-32)),但是,作者将其与水平跳跃连接相结合,直接连接反向收缩和扩张卷积层。

U-Net网络提出后,在医学图像分割领域表现良好。许多研究人员对U-Net网络网络结构进行了各种改进。çiçek等人[29]首先提出了一种3D U-Net网络架构,该架构通过输入3D图像的连续2D切片序列来实现3D图像分割。Milleri等人[30]提出了一种基于U-net的3D变形架构,称为V-net。V-Net网络结构使用Dice coeffificient损失函数而不是交叉熵损失函数,直接最小化了这种常用的分段误差度量。作者进一步将剩余块引入原始U形设计。这两种方法都扩展了具有三维卷积核的U形结构。Drozdzal等人[31]对长跳跃连接(这意味着两个相距较远的特征图之间的跳跃连接)和短跳跃连接(通常称为单个剩余块)进行了区分,并发现它们都有助于创建医学图像分割的深层结构。Simon等人[27]将密集连接卷积网络(DenseNets)与U形结构相结合,将卷积层替换为主干中的密集块,并扩展到自然图像分割,从而实现良好的性能。

B    神经结构搜索

设计一个好的神经网络网络结构既费时又费力,为了减少人工设计网络网络结构的工作量和资源成本,一些学者将注意力放在了神经网络网络结构搜索(NAS)上。近年来,大多数研究工作集中在搜索CNN结构的图像分类上,很少有研究RNN结构的语言任务。如前所述,NAS包括三个组件:搜索空间、搜索策略和性能评估。搜索算法主要包括启发式算法[19]、[21]、[32]–[34]、强化学习[35]–[38]、[38]、[39]、贝叶斯优化方法[40]、[41]和基于梯度的方法[20]、[42]、[43]。性能评估可以从两个方面来理解。首先,评估候选网络结构的性能,以确定是否为下一次更新保留(或扩展)。其次,我们需要一个由单元(当使用基于单元的搜索空间时)或当前候选网络结构堆叠的更深层次的网络,并将训练数据集放入其中进行训练和评估最终性能。

网络搜索空间包括节点的拓扑结构和每个连接节点之间的操作。前者试图直接构建整个网络架构[36],[44]。然而,由于NASNet[37]在ImageNet上成功地将单元堆叠在一起,最近的工作[20]、[22]、[23]、[45]、[46]更倾向于搜索可重复的单元结构,而保持主干宏观网络固定。后者可以提高搜索效率。近年来,许多关于NAS的研究已经提出了许多有效的算法来生成节点的拓扑结构,而这些算法是基于功能强大但易于处理的网络结构搜索空间的。事实上,如果我们有一个丰富且不过度膨胀(rich and not overly expansive)的搜索空间,即使使用随机搜索也可能获得很好的结果[20],[40]。因此,在本文中,我们将重点放在构建基于细胞的医学图像分割搜索空间上。此外,我们使用当前的差分网络结构搜索方法[20],[22](DARTS;Proxylessnas)作为我们的搜索算法,以加速我们的搜索过程。

C    NAS在图像分割中的应用

NAS自提出以来主要解决图像分类任务。最近有一些工作将NAS应用于图像分割。Chen等人[24]首先引入NAS来解决图像分割问题。作者指出,即使在构建递归搜索空间时使用随机搜索,架构搜索也优于人类发明的架构,并在许多分割数据集上实现了更好的性能。然而,这项工作并没有使用一次性搜索,其重点是搜索称为DPC(类似于解码器)的小型Atrus Spatial Pyramid Pooling(ASPP)模块,并将预先训练的主干(修改的Xception))固定为编码器。Liu等人[25]。提出Auto DeepLab:一个通用的网络级搜索空间,并跨两级结构(网络级和单元级)进行联合搜索。作者指出,搜索空间包括各种现有设计,如DeepLabv3、Conv Deconv和Stacked Hourglass。然而,Auto DeepLab的搜索空间不包括医学图像分割领域最著名的U-Net网络结构。

与我们的工作最相似的是[27],它在收缩阶段和扩展阶段都使用密集块 dense blocks 替换卷积层。然而,我们用NAS方法搜索的单元替换所有预先设计的块。

3. 基于单元的架构搜索空间

在本节中,我们首先描述我们使用的CNN架构的通用表示。我们将展示如何将单元网络结构表示为DAG。之后,我们将介绍我们的医学图像分割搜索空间。最后,我们将详细介绍两种类型的单元网络结构。

A    CNN网络结构表示

有向无环图(DAG)用于表示网络拓扑结构,其中每个节点表示输入图像或特征图,并且每个连边关联的是节点和节点之间的操作(例如卷积操作、池化操作和跳连)。当DAG的生成方法不受限制时,其网络结构空间将非常大,这将给现有的搜索算法带来很大的挑战。因此,我们使用基于单元的网络结构。确定了最佳单元网络结构后,我们可以将单元堆叠到主干网络上形成更深的网络。换言之,单元的架构是由整个网络共享的。

B    医学图像分割的搜索空间

在本节中,我们将介绍我们为DownSC和UpSC网络结构选择的基本操作集。之后,我们将描述如何构造它们。

1) 基本作的选择

如何选择合适的基元操作?我们调查了流行的CNN网络结构和在图像分类方面取得巨大成功的前NAS,并总结了我们工作中选择基本操作的重要标准:

(1) 无冗余:这意味着每个基本操作都应该有一些不能被其他操作替换的唯一属性。虽然一些文献[25],[34]表明在搜索过程中可以考虑5×5卷积。大的感受野,如5×5大小的卷积和7×7大小的卷积,可以通过叠加足够多的3×3大小的卷积来代替。因此,所有卷积运算将限制为3×3大小,池化运算为2×2。

(2) 更少的参数:这意味着在搜索过程中消耗更少的内存资源;最初的UNet需要大约31M个参数,这对于移动设备来说是巨大的。在我们的工作中,将引入深度可分离卷积运算,因为它将在不牺牲网络性能的情况下显著减少网络参数。

当滑动步长(步长值)设置为2时,卷积运算可以将特征映射的维度减半或加倍,后者称为上卷积(When set the sliding step (stride value) as 2, the convolution operation can halve the dimension of feature map or double the dimension, the later called Up-convolution.)。这表明向下操作和向上操作可以从相同的基本操作派生。相比之下,与图像分类中的原始操作不同,某些操作的升级版本没有意义(例如identity operation),而池化操作的升级版本(例如平均池化和最大池化)不存在。为了方便起见,我们构建了3种不同类型的基本操作集。

如表1所示,depth conv表示深度可分离操作,而其它的操作,除了dilation conv[47]和cweight[48],是当前NAS方法中普遍使用的。CWweight操作表示挤压和激励[48]操作。在早期的CNN架构中,我们为所有通道生成的特征直接均匀地组合在一起。下一步自然是自动学习每个通道的权重。这正是挤压和激发操作所做的。挤压和激励操作通过为特征通道分配权重来抑制某些冗余特征并增强有用特征。向下cweight操作和向上cweight操作将在通道重新加权之前将特征映射的维数减半或加倍。值得注意的是,当以前的NAS文章展示了它们良好的网络结构时,几乎没有出现膨胀(atrus)卷积运算。然而,此操作的初衷是解决图像分割问题。如前所述,与图像分类任务不同,搜索具有图像分割的架构需要高分辨率输入。巨大的内存消耗无疑是显而易见的。例如,一个512×512的图像大约需要3GB的GPU内存来预测结果,使用原始的U-Net架构,在12GB的Titan pascal GPU上加载模型时,批量大小不超过4。

我们对所有卷积运算使用Conv-ReLU-GN顺序。GN表示组标准化[49],正如Wu等人所表明的那样,这种标准化优于批量标准化,尤其是在批量规模小得多的情况下。由于分割任务的批量比图像分类小得多,因此我们使用组标准化代替批量标准化。

C    两种类型的单元结构

如图2(b)所示,我们设计了两种类型的单元架构,称为基于类U主干的DownSC和UpSC。在两个单元内部,输入节点被定义为前两层中的单元输出[20]、[37]。如图3所示,与输入节点相邻的所有操作都是Down POs或Up POs。设为个中间节点(或称为特征映射层)的集合。与DARTS[20]相同,所有中间节点和输入节点之间的边总数为。

图2.(a)U-Net体系结构(b) Nas Unet体系结构的U形主干,矩形表示需要搜索的单元体系结构。请注意,DownSC和UpSC在NasUnet中都包含它们的down或up操作,而绿色箭头仅表示特征映射(输入图像)的流动。灰色箭头是属于UpSC的变换操作,也是自动搜索的。
图3.单元架构示例。红色箭头表示向下操作(如最大池化),蓝色箭头表示常规操作(例如,identity operation、卷积操作,不会降低特征映射的维数),绿色箭头表示拼接操作。

在收缩步骤中,我们连接L1个单元来学习不同级别的语义上下文信息,并生成一个更小的概率映射,表示为。类似地,在扩展步骤中,我们使用相同数量的单元来恢复中每个概率值的空间信息,并将其扩展到与输入图像一致。最终网络的单元个数为表示为Nas Unet。

与FC densenet架构[27]不同,我们不仅将卷积层替换为这些单元,还将上采样操作和下采样操作移动到单元中。换句话说,正常操作(例如identity operation)和上/下采样操作都被考虑到单元中。如图2(b)所示,变换(transform)也是UpSC中的Norm POs的操作。我们的搜索空间涵盖了许多流行的U型架构,如U-Net[8]和FC DenseNet[27]。值得注意的是,最初的U-Net架构在网络的中间有一个额外的卷积层。然而,在我们的文章中,我们不遵循这种经验,因为我们有一个严格对称的网络结构,由两对几个单元堆叠而成。

4    搜索策略

我们首先描述了如何使用最近的著作[20]、[22]、[50]构造一个具有所有候选路径的过参数化网络。然后,我们引入了一种更有效的架构参数更新策略,以节省GPU(因为CPU的搜索过程要慢得多,我们必须使用GPU)内存[50]。

A    过参数化单元结构

给定一个单元结构,其中表示DAG中的某个边。设为上述三种基本操作类型之一,其中包含N个候选操作。我们将每条边设置为具有N条并行连接(如图4所示)的混合操作,表示为,而不是将每条边与确定操作关联。因此,过度参数化的单元架构可以表示为。混合操作MixO的输出是基于其N个路径的输出定义的:

    

如公式1所示,表示的权重,在OneShot中[50]是常数1,但在DARTS中[20]是对N个架构参数的实值取softmax计算的:。的初始值设为。

图4.过参数化的单元架构示例。每条边与来自不同基元操作集的N个候选操作相关联。

B    一种GPU内存节省更新策略

在上面,每个边的输出是N个候选基元操作的混合操作,这意味着只有当所有操作都加载到GPU内存中时,才能计算所有N条路径的输出特征映射。然而,训练一个紧凑的模型只使用一条路径。因此,相比训练紧凑型,[20]和[50]大约需要N倍的GPU内存。在本文中,我们使用Cai[22]等人提出的二值门(binary gate)来学习二值化路径(binarized path),而不是N条路径。DARTS和二值门方法(表示ProxylessNAS)的区别在于前者在每一步都通过梯度下降更新所有架构参数,而后者只更新其中一个。如图5所示,在训练网络权重参数时,我们首先冻结架构参数,对于输入数据的每个批量,我们随机采样二值门。然后通过训练数据集上的标准梯度下降更新活动路径的权重参数。当训练架构参数时,权重参数被冻结,然后我们重置二值门并更新验证集上的架构参数(图5(a))。这两个更新步骤以交替方式执行。一旦完成架构参数的训练,我们就可以通过修剪冗余路径来导出紧凑的架构。在这项工作中,我们只需选择个(在我们的工作中,k=2)最高权重的路径作为输入的路径。总之,通过这种方式,无论N的值是多少,在架构参数的每个更新步骤中只涉及两条路径,从而将内存需求降低到训练紧凑模型的相同水平。值得注意的是,ProxylessNAS方法在每个更新步骤中只考虑两条路径进行更新,这将导致不在这两条路径上的训练的操作的度要远低于在这两条路径上的操作的度(the trained degree of operation not on the two paths being much lower than the operation on)。因此,我们需要更多的迭代来更新。此外,将不在GPU内存中的特征映射移动到GPU时会花费额外的时间。

图5.DART和ProxylessNAS之间的不同更新策略(a)更新权重参数时仅更新一条路径,但(b)更新所有路径。

V. 实验结果

在此,我们报告了如何实现NAS Unet的详细信息。之后,我们将在基准数据集上报告医学图像分割结果,我们的网络由最佳发现的单元堆叠。

A.NASUNET实施细节

我们考虑了DownSC和UpSC中的中间的数目,并且单元总数,DownSC搜索空间近似为,UpSC搜索空间约为,因此搜索空间的总大小为的数量级,远小于[25 ](Auto-deeplab)。与DARTS不同,在将特征映射的边长减半时,我们不遵循将卷积核数量加倍的做法。

我们在PASCAL VOC 2012数据集[26]上进行架构搜索,用于医学图像分割。更具体地说,我们使用480×480随机图像裁剪。我们随机选择训练集中的一半图像作为验证集。当我们使用DARTS搜索策略时,批量大小为2,架构搜索优化总共进行了120次。当我们使用二值门更新策略时,批处理大小可以是8,但需要200个epoch(参见第IV-B节)。

由于本文致力于为医学图像分割构建一个高效的单元搜索空间,任何可微搜索策略的方法都是可行的[20]、[22]、[23]、[50]。我们希望在非常复杂的图像数据集(PASCAL VOC 2012数据集[26])上搜索代理单元架构,并将其迁移到医学图像数据集。所以在我们的实验中,我们使用了DARTS更新策略。使用ProxylessNAS也是可行的,但需要分别搜索这些数据集。

在学习网络权重时,我们使用动量为0.95的SGD优化器,余弦学习率从0.025衰减到0.01,权重衰减0.0003[20]。在学习架构时,我们使用Adam优化器[51],学习率为0.0003,权重衰减为0.0001。我们经验性地发现,我们从一开始就优化或者在一定的epochs(例如50)后优化,mIoU和pixAcc都是缓慢增长(如图6所示)。因此,我们从一开始就优化。整个架构搜索优化在一个Titan Pascal GPU上大约需要1.5天。我们搜索的DownSC和UpSC如图7(a)和(b)所示。

图6.PASCAL VOC2012数据集上DownSC和UpSC堆叠的候选最佳体系结构网络的验证损失和mIoU
图7.中间数(intermediate number)为4时搜索的单元架构,在PASCAL VOC2012数据集上搜索。 (a)是DownSC体系结构,(b)是UpSC体系结构。 在DownSC中, 和 表示之前的两个单元;而在UpSC中,它们表示之前的一个单元以及水平对应的单元。(The ck−1 and ck−2 indicate the previous two cells in DownSC, but the previous cell and horizontal pre-correspondence cell in UpSC.)

从图7可以看出,我们搜索空间中的搜索处理更倾向于选择cweight版本操作(包括down-cweight操作、up-cweight操作和cweight操作),因为cweight版本操作在DownSC和UpSC网络结构中都占很大比例。值得注意的是,cweight操作取代了标准的跳过连接,以在向下采样和向上采样路径之间传递高分辨率信息(包括更精确的空间信息和高级语义信息)(图2(b)中的灰色箭头表示变换)。这意味着传递高分辨率信息不是简单的拼接,而是加权拼接。

B    医学图像分割结果

为了评估NAS Unet的性能,我们使用了三种类型的医学图像数据集(磁共振成像(MRI)、计算机断层扫描(CT)和超声):Promise12[52]、Chaos[53]和NERVE[54]数据集。通过最小化负 Dice Similarity Coeffificient (DSC) 损失来更新所有模型的权重。我们使用DSC和联合平均交集(mIOU)来评估模型性能。基线方法为U-Net[8],FC Densenet[27]。为了公平起见,我们通过Pytorch[55]重新实现了它们,并使用了相同的数据增强(我们还试图改善一些噪声图像的质量[56],[57])。此外,我们使用相同的Adam优化器,初始学习率为3.0e-4,权重衰减为5.0e-5,训练200个epoch。

1) Promise12

Promise12[52]包含50例训练病例,这些病例包括前列腺的横向T2加权MR图像。训练集有大约1250个带有相应标签的图像(仅体素值为0和1)。每个2d MRI切片的尺寸调整为256×256,并使用对比度受限自适应直方图均衡化(CLAHE)对直方图进行均衡。训练数据集分为40个培训案例和10个验证案例。如表2所示,我们的模型在没有任何预训练的情况下优于所有基线方法。Train Time和GM分别代表批次大小为2时的训练时间成本(总天数和小时数),以及批次大小为2时的GPU内存成本(如下表所示)。

2) Chaos

Chaos[53]挑战将于2014年4月8日至11日在意大利威尼斯举行的IEEE生物医学成像国际研讨会(ISBI)上举行。挑战将从ISBI会议开始。五场比赛将使用两个数据库(腹部CT和MRI),我们选择两个:肝脏分割(仅CT)和腹部器官分割(仅MRI)。第一个挑战是从计算机断层扫描(CT)数据集中分割肝脏,第二个挑战是从磁共振成像(MRI)数据集中分割四个腹部器官(即肝脏、脾脏、右肾和左肾)。这两个数据库中的每个数据集对应于属于单个患者的一系列DICOM图像。第一个数据库包含40名不同患者的CT图像。共提供2874个切片(每个切片为512×512)用于培训,1408个切片用于测试。第二个数据库包括来自两个不同MRI序列的120个DICOM数据集(T1-DUAL同相(40个数据集)、异相(40个数据集)和T2-SPIR(40个数据集))。数据集的分辨率为256×256,切片数在26到50之间(平均36)。总共将提供1594个切片(每个序列532个切片)用于培训,1537个切片将用于测试。

目前,我们使用2874层的CT图像和940层的MR图像来评估我们的模型。如表3所示,在没有任何预训练的情况下,我们的模型在CT图像和MR图像上都比所有基线方法实现了更好的性能。值得一提的是,在MR图像数据集中,我们通过将五个类别重新加权到骰子损失函数中,改进了im平衡类别问题(这五个类别的频率比为1066:40:3.7:4.1)。批量大小也设置为400。

3) Ultrasound Nerve

超声神经分割是2016年的一项挑战。本次比赛的任务是在超声图像中分割一组称为臂丛神经(BP)的神经。一些图像(约60%的训练集)不包含臂丛神经区域。图像的大小为580×420像素。共有5635张培训图片和5508张测试图片(其中20%用于公开排名,80%用于最终排名)。训练数据集包含许多相互矛盾的图像,这意味着两个图像非常相似,但一个图像具有非空掩码,另一个具有空掩码(如图8所示)。因此,我们遵循Julieans方法[58],通过计算每个图像的签名来删除相互矛盾的图像。然后,二次相似性是两幅图像的两个特征向量之间的余弦距离。这将生成所有训练集图像的距离矩阵,然后对其进行阈值化,以决定应删除哪些图像。最后,我们保存了4456张训练图像(在5635张图像中)。最后,我们随机分割0.2张训练图像作为验证集。如表4所示,我们的模型在没有任何预培训的情况下比所有基线方法都能获得更好的性能。

图8.两个矛盾图像的示例,(a)与(b)相似,但没有空遮罩。与所有人类标记的数据一样,真值中的潜在错误很容易发生。

六、结论

在本文中,我们尝试将神经结构搜索扩展到医学图像分割。我们为我们的搜索空间和基于搜索单元的架构设计了三种类型的基本操作集,它们由DownSC和UpSC堆叠而成。我们选择U型主干(我们的搜索空间包括U型网络及其许多变体)进行搜索,并引入了一种记忆保存搜索算法(二进制门)[22]来加速搜索过程。通过从头开始对医学图像分割数据集进行培训,评估搜索结果NAS Unet。在Promise2上,NAS Unet显著优于基线方法。在Chaos和Ultrasound Nerve方面,NAS Unet也优于这些基线方法。

你可能感兴趣的:(U-Net,NAS-Unet)