[NAS]ProxylessNAS: Direct Neural Architecture Search on Target Task and Hardware

from ICLR2019

 
MIT HAN Lab

Abstract

神经网络架构搜索可以高效完成神经网络结构的自动搜索。但是搜索过程所需的巨大搜索成本限制了传统NAS算法在大规模数据集上的应用，比如耗费10^4 GPU hours。可微分的NAS通过将离散空间连续化减少了搜索时间但是还存在GPU显存的问题，会随着候选集合的规模线性增长。因此往往需要借助代理任务，即在小规模数据集上训练或者在一部分模块上学习的结果，但这不能保证在目标任务上是最优的。

 
本文提出的Proxyless可以直接在大规模数据集和特定硬件平台上进行搜索学习，在同样规模搜索集合的前提下大大减少了GPU显存和CPU Hour.本文在CIFAR-10和ImageNet数据集上的实验结果显示了这种方法的定制化和有效性。在CIFAR-10数据集上参数量仅为5.7M时测试误差仅有2.08%，参数量比之前的SOTA-AmoebaNet-B压缩了6倍；在ImageNet数据集上比MobileNetV2的top-1精度提升了3.1%,同时GPU延迟加速了1.2倍。本文还将ProxylessNAS用于特定硬件平台时延的搜索，为CNN的NAS搜索提供思路。

Section I Introduction

NAS在自动网络设计取得了优异的成果，如在图像识别和语言建模任务中。但是NAS高昂的搜索成本，需要在一次实验中训练成百上千的模型适得其反难以直接应用在大规模任务中，如ImageNet,也因此使得工业部署难以实现。作为权宜之计，有的研究提出在代理任务上搜索基础模块，如训练更少的epoch,在小规模数据集上先训练或者使用更少的基础模块，然后选择表现最优的模块通过堆叠的方式搭建起目标网络迁移到目标的大规模数据集上。这种设计思路也被后续诸多NAS算法沿用。
但是这种在代理任务上的搜索并不能保证对目标任务是最优的，尤其是将一些硬件约束（如延迟）考虑在内。并且为了保证可移植性，这些方法往往值搜索简单的模块然后堆叠成网络，限制了网络的性能和多样性。
本文提出的ProxylessNAS可以直接在目标任务或硬件上进行搜索，无需使用代理任务。这种算法也避免了之前NAS算法中重复堆叠模块的弊端，使得网络的所有模块都是可以学习定制的。并且本文将搜索成本限制在和常规训练一样的量级。

在GPU Hour层面本文将NAS搜索定义为一个路径级别的剪枝过程，本文直接训练一个过参数化的网路，包含所有候选操作路径。在训练过程中显式引入网络结构参数来学习看哪些路径是冗余的，对冗余路径在训练结束后进行剪枝从而得到精简的优化后的网络结构。通过这种方式只需要训练一个网络而不需超网的参与。
但是如果将所有候选路径都加入考虑弧导致GPU显存爆炸，因为内存会随着候选集合线性增长。因此在GPU memory层面本文采用二值化的方式每次运行只有一条路径被激活，这样可以限制memory。
本文提出一种基于梯度的方法来训练这些二值化连接的网络。更进一步为了处理硬件约束指标如latency，本文将网络延迟建模为一个连续函数，将其作为正则化损失项。除此之外还提出一种REINFORCE-based算法作为处理硬件指标的替代策略。
本文在CIFAR-10和ImageNet数据集上开展实验，得益于算法的直接性和定制化，在CIFAR-10数据集上搜索到的网络在5.7M参数量下测试误差仅为2.08%；ImageNet上搜索的结构达到了75.1%的top-1精度，比MobileV2快了1.2倍。
本文的工作总结如下：

（1）ProxylessNAS可以在目标大规模数据集或硬件上直接进行搜索，无需任何代理任务，也不用通过模块堆叠的方式构建网络。有效保证了原始的搜索空间并且提升了搜索性能。

（2）本文提出一种新的路径级别的NAS剪枝方法，展示了NAS与模型压缩之间的紧密联系，通过路径二值化的方式将搜索现存控制在同一量级。

（3）本文提出一种新的基于梯度将延迟作为正则化项来实现硬件约束目标。这样可以针对不同的硬件平台（GPU,CPU,移动端）进行定制化搜索。据我们所知这是首个研究针对不同硬件平台的NAS算法。

Section II Related Work

机器学习技术，如强化学习、神经进化来代替人类专家设计神经网络这类工作通常被称为神经网络架构搜索。NAS的结构搜索通常被认为是一个元学习过程，通过引入一个元控制器（比如RNN）来在一个给定的搜索空间中训练一个网络，来获得一定的性能指标指导进一步搜索。但是这种方法的计算成本高昂，尤其在大规模数据集上。
有的工作尝试通过简化评价过程来提升这一元学习的过程。比如使用超网再通过采样的方式生成子网进行评估，并在所有子网中共享超网的权重。但是这些算法需要一个超网或者一个RNNcontroller，并且主要聚焦于小规模任务上，而不是大规模数据集。

本文的工作与One-shot以及DARTS接近，都是移除了controller部分将NAS建模成一个过参数化网络的单次训练过程。具体来说，One-Shot是训练一个过参数化网络然后以一定的概率删除每条路径，使用预训练的估参数化网络来评估，并通过随机置零进行路径的采样。DARTS则是对每条路径使用梯度下降的方式联合训练权重参数和结构参数。但是上述方法都带来了巨大的内存消耗，因此仍然需要使用代理任务。
本文通过二值化路径来解决内存问题。
通过网络剪枝不重要的神经元或通道达到提升网络效率的问题，本文也按照类似的思路。从一个过参数化的网路偶开始减去冗余的部分得到优化后的网络。本文和它们的区别在于之前的剪枝改变的是filter或通道数，而不会改变网络的提频谱结构，本文则是通过路径级别的剪枝来学习有效的网络架构，还允许修剪或者增加网络层数。

Section III Method

本文首先介绍过参数化网络的设计，包含了所有的候选路径；然后介绍如何通过二值化的结构参数来减少内存；最后提出基于梯度的算法来训练网络。本文提出两种方法来处理不可微的约束（如延时）来定制化设计特定硬件上的神经网络。

Part 1 Construction of over-parameterized network

网络表示为N，ei表示特定的边，O包含了所有的候选操作（卷积、池化、identity,zero），为了确保每条边都可以是任意操作搭建一个过参数化的网路哦，本文设定包含N组并行的路径，参见Fig 2，表示为：


对于输入x输出是N条边上的混合操作的结果。比如在DARTS中mo(x)就是所有路径操作的加权和，权重则是通过softmax激活得出的：

那么N条路径输出的特征图会存在memory中，但在训练时只选择其中一条路径。如果使用One-Shot或DARTS那么需要的memory和GPU Hour则是单个网络的N倍，接下来就解决memory消耗的问题。

Part 2 Learning Binarized Path

为了减少memory本文在训练时只保留一条路径。与之前二值化的方法不同本文是对整个路径进行二值化，会引入N个实数值的结构参数，将实数值转化为二进制门序列g：
那么最终所有路径的混合输出表示为：

二值化后在运行时每次只有一条路径会激活，这样memory需求也就下降至和训练一个紧凑模型一样，这不仅下降了一个量级。


Training binarized architecture parameters



Fig 2展示了参数化网路中如何训练权重参数和二值化后的结构参数。在训练权重参数时先冻结结构参数，根据Eq2随机采样二进制序列确定激活的路径。

参见Fig 2左侧；
在训练结构参数时则冻结权重参数，重置而今之门序列然后在验证集上更新结构参数，参见Fig 2右侧。



这两个步骤交替进行，一旦结构参数训练完成就可以剪枝掉冗余的路径获得紧凑的网路哦结构，本文仅选择拥有最高权重的路径。
但是与权重参数不同的是，结构参数不再计算图中因此无法计算梯度进行更新，本文引入一种基于梯度的算法对结构参数进行更新。

本文的结构参数类似BinaryConnect的思路，通过∂L/∂gi代替∂L/∂pi近似估计得到：



后者包含在计算图中，但此时仍然需要存储每条路径的结果oj(x)因此本文继续进行近似。
本文考虑从N个候选路径中选择一条路径，因为如果某条路径是某个位置的最佳选择，那么与其他路径相比它应该也是更好的选择。


遵循这一想法，本文在更新结构参数时会先从路径分布中采样两条路径，其余路径加掩膜。因此这样候选路径的数目从N变为2，对应的路径权重和二进制序列也被reset.
随后根据Eq4来计算路径梯度，经softmax后获得路径权重，此时还需要重新调整这两个更新后的路径权重，主要是为了保障总的路径权值不变，那么就是一条路径增加，另一条路径相应减少，从而保证其他路径的权重不边。这样在更新结构参数时每次只设计两条路径，从而将memory降到和训练紧凑模型相同的水平。

Part 3 Handling Non-differentiable hardware metrics

除了网络精度，延迟（不是FLOPS）也是网络效率的一项重要指标，但是不像精度可以通过计算损失函数的梯度，延迟是不可微的，因此本文提出两种算法来处理这种不可微的约束。

Making latency differentiable 


为了让延迟可微，本文将网络延迟建模为一个连续函数。对于一系列操作的混合输出，每一条路径权重代表选择该条路径的概率，因此期望延迟表示为：



其中F表示延迟预测模型，因为整个网络包含一系列的混合操作输出，因此总的延迟是所有混合操作期望延迟的累加结果：



本文将延迟作为损失函数的一部分，并且权重因子λ2用以调控精度和延迟之间的比重，最终损失函数表示为：


其中分类精度使用的是交叉熵损失函数。
Reinforce-based Approach 
作为BinaryConnect的替代方法本文使用基于强化的方法来训练二值化权重。更新二值化参数的目标是找到能够使得而今之门最大化的reward R。本文假设网络只包含1个混合操作，则二值化参数的更新为：

虽然这种基于增强的方法不是一个序列决策或者RNN控制器的参与，但是它和基于梯度的算法本质上都是对相同的二值化结构参数进行更新的方法，更新规则不同而已，因此可以将二者组合使用。

Section IV Experimental and Results

本文在CIFAR-10和ImageNet上进行分类任务，但和之前NAS算法评估方法不同，之前都是在小规模数据集上训练胡迁移到大规模数据集或者训练小的block之后堆叠成大的block，本文直接在目标任何和目标硬件上进行全任务的训练。

Part 1 Experiments on CIFAR-10

Architecture space



以PyramidNet为backbone但是将所有3x3卷积替换为三个cell，每个cell深度为3,2个分支并且以内两个超参数B,F表征网络每一个stage的宽度和深度，一共包含3层stage。


Training details

首先从训练集采样5000张作为验证集，优化器Adam 
学习率分别是0.006和0.01
超网训练 200 子网训练300 epochs

Results

超参数设置：B=18 F=400
因为本文不是通过cell堆叠，每一个cell包含12个可学习的边，因此整个网络结构有12x18x3=648种可能的选择
实验结果展示在Table 1中，可以看到本文的ProxylessNAS不仅拥有最低的测试误差同时还是参数最优的，在参数量只有5.7M时测试误差为2.08%,比AmoebaNet-B参数量压缩了6倍。并且这种无代理NAS方法与有代理的取得了相似的精度甚至更低的误差，表明本文直接探索一个较大的架构空间的有效性。

Part 2 Experiments on ImageNet

ImageNet实验则是探究设计高效的CNN结构因此不仅要高精度还要低时延。这是一个多目标优化任务，其中一个目标还是不可微的，本文测试的三种硬件平台分别是：

GPU:V100 batch_size = 8


CPU: 2 2.40GHz Intel® Xeon® CPU E5-2640 v4 batch_size 1



Mobile: Google Pixel 1




优化目标是：

此外对于移动端本文还是用一个延迟预测器（latency predictor)来代替实际测量的延迟，并且发现预测模型与实际延迟有强相关性，这样就不用实际测量延迟。



Architecture Space




使用MobileNetV2作为backbone，具体来说不需要重复堆叠同样的MBConv，而是使用不同卷积核以及不同扩张率的MBConv。
为了更好的权衡网路哦宽度和深度本文也搭建一个深层次的过参数化网络，并且包含零操作进行skip，这样可以通过skip更多的模块以及使用更大的MBConv来搭建更宽、更浅的网路，或者使用更小的MBConv搭建窄且深的网络。


ImageNet Classification results

Table 2展示了在ImageNet做分类的结果，可以看到比MobileNetV2提升了2.6%的分类精度，但延迟几乎没变；如果精度相同的情况下，MobileNetV2的延迟是143ms而本文只需要 78ms，加速比为1.83，更重要的是本文的搜索成本更低，比如比MnasNet的搜索时间加速了200倍。
除此之外本文还发现如果不加入延迟正则化损失，那么ProxylessNAS倾向于选择计算简单的操作，在Pixel 1的延迟是183ms，加入延迟约束后搜素的网络延迟减少到83ms，并且比不使用延迟约束的网络精度提升2.4%。因此本文人为在网络搜索过程中将延迟作为优化目标是十分有必要的。


Table 3展示的则是在GPU上搜索的结果，可以看到ProxylessNAS搜索到的结果比人工设计或者其他NAS搜索到的网络性能都要优异，比MobileNetV2和MnasNet分别提升3.1%和1.1%。

Table 4展示的则是搜索到的网络在不同硬件平台上的性能和延迟，可以看到在GPU优化的模型在CPU和移动端的表现并不好，反之亦然。说明针对不同硬件设计定制化设计神经网络是十分有必要的，才能在不同硬件平台达到最佳效率。


Fig 6展示了在三个平台上搜素的CNN模型详情，可以看到不同硬件平台的神经网络有不同的偏好：



（1）GPU模型倾向于更浅、更宽，尤其在特征图早期，此时分辨率较高；



（2）GPU模型倾向于使用更大的MBConv比如7x7的MBConv6，而CPU倾向于使用更小的MBConv，可能是因为GPU并行性更高，因此可以利用大型MBConv
；


（3）三种硬件平台上搜索的模型读倾向于在每个stage的第一个block中使用更大的MBConv这样可能使得网络得以在降采样时保留更多的信息，这一特点并没有在之前NAS算法中观察到，因为之前的算法都倾向于堆叠重复的block.




Section V Conclusion

本文提出的ProxylessNAS可以直接在目标任务和硬件平台进行网络搜索而无需使用代理，本文还进一步降低了搜索时间和内存需求，通过引入二值化路径操作。这样可以针对不同硬件平台定制化搜索不同神经网络，设计高性能高效率的神经网络。

Appendix A 
展示了在CIFAR-10任务中的候选操作集合

Appendix B 
介绍了使用的latency predictor

直接在平台上测试延迟又慢成本也高昂，比如测试一次需要20s，并且延迟的测量需要设计一整套流程。本文只是用一部手机，平均延迟是0.75ms.本文使用延迟模型进行搜索使用测量的延迟作为模型最终延迟。从候选空间中采样了5k个网络结构，其中4k用于构建延迟模型其余用来测试，测试使用Tensorflow-Lite