RONASMIS:资源优化的NAS用于3D医学图像分割
Abstract
NAS可以自动搜索神经网络,在深度学习领域备受关注。但是目前只有少数工作探索NAS在3D医学图像分割上的应用,因为医学图像一般非常大,受限于GPU算力很难将NAS算法用于如此大的图像。
本文提出资源优化的NAS方法,将训练时间缩短到1.39天/1GB 数据集,单块2080Ti即可满足算力需求。很多NAS算法不用重训练或微调即可实现较好的性能,有的借助权参数共享的一个强化学习学习到的控制器,有的聚焦在宏观结构的搜索而不是微结构的搜索。
实验结果表明本文的RONAS算法比人工设计的网络在3D医学图像上性能更佳。
Section I Introduction
深度学习网络用于医学图像分割取得了优异的性能,展现了如UNet,deep supervision等一系列优异性能的框架,但是这些网络的性能高度依赖于人工处理,如图像后处理、超参微调以及对网络结构的调整和优化。尤其是超参数的调整需要耗费大量的时间成本和计算成本,为了减少人工调参的时间,自动机器学习开始进入人们的视线,尤其是NAS相关算法在自然图像处理得到了大量的研究。
但是对3D医学图像的研究还比较少,因为3D图像一般维度很高,需要大量的算力,目前在3D图像上的NAS探索还没有超过人工设计网络的性能。
本文提出的ROMASMIS最大的特点是训练时间短、算力需求小,每1GB数据集只需要1.39天,在10.8GB的2080Ti上即可完成训练。提出的网络框架关注于宏观架构的搜索,不像之前的NAS算法聚焦的是微结构的搜索,同时会利用3D医学图像的特点。
此外本文还避免重训练网络,而是在结构搜索过程中不断优化子网络。此外还用基于加法的skip connection,用常规卷积替换深度卷积,这些都减小了对memory的需求。而搜索空间只包含了显著影响分割心梗的关键操作,进一步减少了架构的复杂性和对内存的需求。
据我们所知本文是首个将NAS用于3D医学图像分割并超过人工网络性能的研究。
Section II Method
Part 1 Resource-Optimized Search Space for Anisotropic 3D Medical Image and the Basic Architecture
Fig 1展示了本文定义的搜索空间,与传统NAS算法有以下4方面的不同:
1.考虑到大多数3D医学图像各向异性的形状,本文将输入patch的大小和下采样的次数考虑在内。因此用于训练的图像规格包含:4×155×240×240,1×[40]×90∼130]×320×320,2×[11∼24]×24]×[256∼384]×[256∼384]。
在通道数、深度、高度、宽度等方面都不同。
2.为了在编码过程中尽可能扩大感受野同时保留特征的空间信息,本文在搜索空间中你包含了3D空洞卷积和池化操作。
3.本文的目标之一是高效利用GPU memory,而聚焦微结构的搜索需要大量的内存,3D分割使得GPU难以负担。因此本文关注于宏观结构的都多空间,这与其他工作不一样,前人的工作有的同时关注微观和宏观。
4.有的研究显示全局特征之间的连接是至关重要的,而大多数NAS搜索只考虑一两个临近cell的输入;而本文在搜索空间中包含了skip connection,这样可以最大限度的提高跨网络之间的连接。
Table 1展示了搜索空间包含的操作。
此外,其他NAS研究一般在skip connection后使用1x1卷积,这样增加了参数量,也会增加GPU的内存负担。因此本文使用elementwise级别的sum完成skip connection进行宏观结构的搜索。
考虑到可用的GPU大小的受限的,一般会通过级联的方式进行学习,这样就会增加训练时间,并且仍然需要超参数调优。而本文则专注于构建一个资源优化的搜索空间,从而减少超参数调整的时间。
搜索空间还包含了3中激活函数,2种池化,他们不会使用额外的内存;为了方式子网络的过拟合本文还增加了一个零操作-drop_path正则化,会禁用某些操作,从而搜索搭建可靠的网络结构
。
此外本文还将大多数NAS中使用的深度卷积替换为常规卷积,这样可以更好的利用内存。虽然人们普遍认为深度卷积更节省,算数操作更少;但是在Pytorch等框架实现过程中深度卷积无法更好的利用内存,在执行反向传播时需要更多的内存。
下面两个式子展示了如何根据输入图像的维度决定patc size的大小,H,W,D就是设定好的,一般设置H=W,深度是单独设置的。而对于给定的输入patch限制搜索空间基数为5.
当中间阶段的图像维度 小于 最大的池化大小时会调整池化的步幅。因为在训练之前图像大小的范围是未知的。
The proposed Base Architecture
本文是用的基础结构基于UNet,并且结合了DeepLabV3+中的decoder 1x1x1 卷积+skip connection,和deep supervision.详情参见Fig 2.
比如紫色的就是1x1x1 convolution skip connection,主要用来控制从encoder到decoder传播的信息量。一般会将通道数减半然后与decoder的输出进行级联;
未填充的紫色和黑色箭头是stage1,2,3输出结果做element-wise-sum的结果,这是一种改进的deep supervision方法;框图中加粗的部分指的是搜索空间中需要考虑如何设计的部分;
问号则表示需要下采样的位置,由池化stride决定;
绿色虚线箭头表示是否应该使用zero operation;
S表示skip connections.
当skip connenction要连接的特征通道数不一致时大多数NAS框架会使用1x1x1卷积调整通道数后再级联所有特征;而本文则使用一种匹配操作,只对响应通道数的情况才使用1x1x1卷积做通道调整,这样可以减少GPU内存的使用和参数量,然后进行element-wise-sum。
本文希望这种凡是可以使得传输特征包含所有信息,良好的结合不同分辨率的特征。
所有的3D卷积的卷积核为1x1x1或3x3x3,由控制器决定卷积的空洞率。为了训练稳定,stage 1的膨胀率固定为1,stage1和stage2的池化步长固定为2.最深层的特征可以选择分辨率为1/16-1/4不等。
Part 2 Training the Controller’s Parameter for Architecrure Selection
本文使用ENAS提出的参数共享的强化学习来训练Controller,这是一种很好的训练方法,controller只需要计算新建网络结构的reward而无需从头训练子网络,这样减少了训练时间,虽然可能会影响性能;但随着训练次数增加可以达到全局最佳状态。
而本文没有使用最近提出的基于梯度的DARTS相关算法是因为这需要将所有操作加载到GPU,对GPU需求太大。
每次迭代Controller会创建20个子网络,然后通过验证网络的dice socre作为reward.将阈值在0-0.5的才计算得分,Controller根据得分安排下一代的排列顺序。子网的训练使用dice loss。
Section III Implementation Details and Experimental Results
Part 1 Data abnd Implementation Details
数据集:MSD 包含脑、心脏、前列腺的3D医学图像 这种形状各异的数据集非常适合评估NAS的性能。
实施细节:
预处理:每个通道计算z-score进行归一化;
优化器Adam Optimizer。
大脑、心脏、前列腺任务的training epoches分别是150,500,500
Part 2 Experimental Results
与nnUNet和SCNAS进行了对比,对比结果参见Table 2.
nnUNet是目前MSD排名第一的网络框架,SCNAS则是基于梯度的NAS算法的一种,聚焦于微结构的搜索。
本文搜索的结果超过了SCNAS和nnUNet,并且不需要消耗太多内存,只在1个2080Ti上进行训练,三个任务的训练时间分别用了3.1,1,39和0.35天。并且在MSD挑战中许多排名靠前的方法都使用补充的方法,如集成多个神经网络、数据增强、后处理等。
本文使用的数据增强只有水平翻转,也没有使用任何预训练。
以及许多team使用50%重叠的patch级别做inference,这会增加推理时间和计算量;而本文只需要一次推理。
Fig 3可视化了搜索到的用于heart分割的一种最优架构,右侧则是不同任务的reward变化和熵的变化,可以看到熵是逐渐减少的,reward是稳定增加的,说明controller得到了良好的训练;换言之controller会反复推荐相近的子网络才能获得一致的dice分数。实验结果显示无论验证集如何变化,超过80%的时间都会推荐相同的网络。
Section IV Conclusion
本文设计了一个聚焦于宏观结构搜索的高效搜索空间,通过参数共享训练controller选择候选网络,从而用于3D医学图像分割任务。搜索到的网络结构性能优于人工设计的网络以及NAS中的SOTA结果。 此外,本文提出的方法是在不使用任何额外手段下就取得了优异的性能。