麻省理工 | ProxylessNAS自动为目标任务和硬件定制高效CNN结构(文末论文及源码)

时间:2023-02-24 15:59:42


麻省理工 | ProxylessNAS自动为目标任务和硬件定制高效CNN结构(文末论文及源码)

麻省理工 | ProxylessNAS自动为目标任务和硬件定制高效CNN结构(文末论文及源码)

导 读

NAS受限于其过高的计算资源需求,仍无法在大规模任务上直接进行神经网络的学习。

今天分享的这篇文章主要解决NAS代理机制下无法搜索到全局最优的问题,改进搜索策略,一定程度上解决资源消耗的问题。其主要是基于DARTs改的,那就先谈谈DARTs的一些问题:

  • 只搜索cell,然后不断的堆叠起来,实际上只要网络堆叠的够深,性能并不会太差;
  • 一些operation的结构参数非常接近,比如两个op的结构参数分别是20.2%,20.1%,似乎选择20.2%和20.1% 并没有太大的差别;
  • 一开始是在整个Net上进行训练,但是选取的时候只是选择一个子网络,中间可能会出现一些问题。

作者提出了ProxylessNAS,第一个在没有任何Proxy的情况下直接在ImageNet量级的大规模数据集上搜索大设计空间的的NAS算法,并首次专门为硬件定制CNN架构,作者还将模型压缩(减枝、量化)的思想与NAS进行结合,把NAS的计算成本(GPU时间和内存)降低到与常规训练相同的规模,同时保留了丰富的搜索空间,并将神经网络的硬件性能(延时,能耗)也直接纳入到优化目标中。



引言



目前一个普遍的做法是在一个小型的Proxy任务上进行网络学习,然后再迁移到目标任务上。这样的Proxy包括:1) 训练极少量轮数;2)在较小的网络下学习一个结构单元,然后通过重复堆叠同样的block构建一个大的网络;3) 在小数据集(如:CIFAR)上进行搜索。然而,这些在 Proxy 上优化的网络结构在目标任务上并不是最优的。


作者最后提出的ProxylessNAS在CIFAR-10和ImageNet的实验验证了直接搜索」和「为硬件定制」的有效性。在CIFAR-10上,作者新提出的模型仅用5.7M参数就达到了2.08%的测试误差。对比之前的最优模型AmoebaNet-B,ProxylessNAS仅用了六分之一的参数量就达到了更好的结果。在ImageNet上,ProxylessNAS比MobilenetV2高了3.1%的Top-1正确率,并且在GPU上比MobilenetV2快了20%。在同等的top-1准确率下(74.5%以上),ProxylessNAS的手机实测速度是当今业界标准MobileNetV2的1.8倍。在用ProxylessNAS来为不同硬件定制神经网络的同时,作者发现各个平台上搜索到的神经网络在结构上有很大不同。这些发现为之后设计高效CNN结构提供新的思路。



相关背景



近期研究中,神经结构搜索(NAS)已经在各种深度学习任务(如图像识别)的神经网络结构设计自动化方面取得了很大成功。然而,传统NAS算法的计算量巨大,如NASNet需要10^4 GPU hours来运行。以DARTs为代表的Differentiable NAS虽减少了计算需求,可随着搜索空间线性增长的内存成了新的瓶颈。由于这些局限,之前的NAS都利用了Proxy任务,例如仅训练少量Epoch,只学习几个Block,在较小的数据集上搜索(如CIFAR)再迁移。这些Proxy任务上优化的结构,并不一定在目标任务上是最佳的。同时为了实现可转移性,这些方法往往仅搜索少数结构Block,然后重复堆叠。这限制了块的多样性,并导致性能上的损失。依赖Proxy同时也意味着无法在搜索过程中直接权衡延迟等硬件指标。

因此,作者提出了一个简单而有效的方案来解决上述限制,称为ProxylessNAS,其直接在目标任务和硬件上学习结构而不依赖于Proxy。作者还移除了先前NAS工作中的重复块的限制:所有stage都可以*的选择最适合的模块,并允许学习和指定所有块。为此,将体系结构搜索的计算成本(GPU时间和内存)降低到相同水平的常规训练,为了直接在目标硬件上学习专用网络结构,在搜索时我们也考虑了硬件指标。

麻省理工 | ProxylessNAS自动为目标任务和硬件定制高效CNN结构(文末论文及源码)

在CIFAR-10和ImageNet上的实验中,从直接性和专业性的角度出发,提出的新方法可以获得较强的实证结果。在CIFAR-10上,新模型仅需5.7M参数即可达到2.08%的测试误差。在ImageNet上,新模型比MobileNetV 2提高了3.1%,获得了75.1%的top-1精度,同时速度为1.2倍。主要的贡献可归纳如下:

  • ProxylessNAS是第一个在没有任何代理的情况下直接学习大规模数据集上的CNN结构的NAS算法,同时仍然允许大的候选集并消除重复块的限制。它有效地扩大了搜索空间,实现了更好的性能;
  • 为NAS提供了一种新的路径级剪裁视角,显示了NAS与模型压缩之间的紧密联系。通过使用path-level binarization将内存消耗节省一个数量级;
  • 提出了一种新的基于梯度的方法(作为一个正则函数),来处理硬件目标(如:延迟)。针对不同的硬件平台:CPU/GPU/FPGA/TPU/NPU,ProxylessNAS实现了针对目标硬件CNN结构定制。据所知,这是第一个来研究不同硬件结构下的专用神经网络结构的文章;
  • 广泛的实验证明了Directness和Specialization的优势。它在不同硬件平台延迟限制下,在CIFAR-10和ImageNet上实现了最好的性能。作者还分析了专用于不同硬件平台的高效CNN模型的偏好,指出不同硬件平台需要不同的神经网络结构。



技术方法



Over-parameterized网络的构建

首先描述了具有所有候选路径的over-parameterized网络的构造,然后介绍了如何利用二值化的体系结构参数将训练over-parameterized网络的内存消耗降低到与常规训练相同的水平。提出了一种基于梯度的二值化结构参数训练算法。最后,提出了两种处理不可微目标(如延迟)的技术,用于在目标硬件上专门化神经网络。

麻省理工 | ProxylessNAS自动为目标任务和硬件定制高效CNN结构(文末论文及源码)

学习BINARIZED PATH

为了减少内存占用,在训练over-parameterized网络时,只保留一条路径。与Courbariaux等人不同。对单个权重进行二值化后,对整个路径进行二值化。引入n个实值结构参数{αI},然后将实值路径权重转换为二进制门:

麻省理工 | ProxylessNAS自动为目标任务和硬件定制高效CNN结构(文末论文及源码)


然而直接去训练这样一个over-parameterized网络是有问题的,因为其GPU显存会随着候选操作的数量线性增长。这里,作者就利用到路径级二值化的思想来解决这个问题:即将路径上的架构参数二值化,并使得在训练过程中只有一个路径处于激活状态。这样一来GPU显存的需求就降到和正常训练一个水平。在训练这些二值化的架构参数的时候,作者采用类似BinaryConnect的思想,使用对应的Binary Gate的梯度来更新架构参数:

麻省理工 | ProxylessNAS自动为目标任务和硬件定制高效CNN结构(文末论文及源码)

为了解决这个问题,考虑将从n个候选对象中选择一条路径的任务分解为多个二进制选择任务。直觉是,如果一条道路是在某一特定位置上的最佳选择,那么与任何其他路径相比,它应该是更好的选择。

HANDLING NON-DIFFERENTIABLE HARDWARE METRICS

除了准确率之外,在设计高效神经网络结构的时侯,延迟是另一个非常重要的目标。与可以使用损失函数的梯度优化的准确率不同,延迟这一指标是不可微的。在本节中,我们提出了两种算法来处理这种不可微分的目标。

麻省理工 | ProxylessNAS自动为目标任务和硬件定制高效CNN结构(文末论文及源码)

通过将缩放因子λ2(>0)乘以控制精度和延迟之间的权衡,将网络的预期延迟合并到正常的损失函数中。最后的损失函数如上图所示(右)。

对于二进制参数,我们有以下更新:

麻省理工 | ProxylessNAS自动为目标任务和硬件定制高效CNN结构(文末论文及源码)



实验结果



麻省理工 | ProxylessNAS自动为目标任务和硬件定制高效CNN结构(文末论文及源码)

在CIFAR-10和ImageNet上进行了实验。不同于之前的NAS工作,作者直接在目标数据集上进行神经网络结构学习,为目标硬件进行优化,同时允许每一个block*地选择操作。

在ImageNet上的结果

麻省理工 | ProxylessNAS自动为目标任务和硬件定制高效CNN结构(文末论文及源码)

在移动端,与MobilenetV2相比,提出的ProxylessNAS在维持同等的延迟的前提下,TOP-1准确率提升了2.6%。此外,在各个不同的延迟设定下,新模型大幅优于MobilenetV2:为了达到74.6%的精度,MobilenetV2需要143ms的推理时间,而新模型仅需要78ms(1.83x)。

与MnasNet相比,新模型在提升0.6%Top-1的同时保持略低的推理时间。值得一提的是,新模型所消耗的搜索资源要比Mnas少得多(1/200)。

麻省理工 | ProxylessNAS自动为目标任务和硬件定制高效CNN结构(文末论文及源码)

下图展示了在三个硬件平台上搜索到的CNN模型的详细结构:GPU/CPU/Mobile。

我们注意到,当针对不同平台时,网络结构呈现出不同的偏好:

1)GPU模型短而宽,尤其是在feature map较大时;

2)GPU模型更喜欢大MBConv操作(如:7x7 MBConv6),而CPU模型则倾向于小操作。这是因为GPU比CPU有更高的并行度,因此它可以更好地利用大MBConv。

另一个有趣的观察是,当特征图被下采样时,所有的网络结构都倾向于选择一个更大的MBConv。这可能是因为MBConv操作有利于网络在下采样时保留更多信息。值得注意的是,这是之前强制block之间共享结构的NAS方法无法发现的。

麻省理工 | ProxylessNAS自动为目标任务和硬件定制高效CNN结构(文末论文及源码)

END