[NAS]RONASMIS：Resource Optimized Neural Architecture Search for 3D Medical Image Segmentation

Paper

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结果。
此外，本文提出的方法是在不使用任何额外手段下就取得了优异的性能。