神经网络架构搜索——可微分搜索（Latency-DARTS）

本文为华为诺亚的NAS工作，针对DARTS搜索出网络结构往往对硬件不友好的问题，本文在优化中加入可微分的时延损失项，使搜索过程可以在精度和时延之间进行平衡系数的权衡。

论文题目：Latency-Aware Differentiable Neural Architecture Search

论文地址：https://arxiv.org/abs/2001.06392

摘要

可微分的神经架构搜索方法在自动机器学习中盛行，主要是由于其搜索成本低，设计搜索空间灵活。然而，这些方法在优化网络方面存在困难，因此搜索到的网络往往对硬件不友好。本文针对这一问题，在优化中加入可微分的时延损失项，使搜索过程可以在精度和时延之间进行平衡系数的权衡。延迟预测模块（LPM）是对每个网络架构进行编码，并将其输入到一个多层回归器中，通过随机抽样收集训练数据，并在硬件上对其进行评估。本文在NVIDIA Tesla-P100 GPU上评估了该方法。在100K采样架构（需要几个小时）的情况下，延迟预测模块的相对误差低于10%。嵌入延迟预测模块，搜索方法可以减少20%的延迟，同时保留了精度。本文的方法还能简洁的移植到广泛的硬件平台上，或用于优化其他不可微的因素，如功耗。

image

动机

本文从神经网络架构延迟的角度出发，对DARTS的搜索进行了优化。DARTS工作对网络延时的对等指标定义为FLOPs，但是FLOPs仅仅是对模型的计算复杂度实现了量化，并不能完全对等网络的推理延迟。从硬件部署角度考虑，网络模型的推理速度受硬件加速设备，部署框架的优化效果，输入输出数据的延迟等多方面的影响。如下图所示：

image

根据上图的实测结果，我们可以清晰的发现，在相同的FLOPs下，不同的网络结构对应的实际网络延迟也有明显的差异。比如，在400MFLOPs的网络架构中，高延迟能达到25ms，而低延迟仅需要17.1ms。那么，如何在保证模型精度，FLOPs的前提下尽量的做到低延迟呢？

方法

Latency-DARTS架构示意图

延迟预测模块（LPM）

延迟预测模块（LPM）首先是对每个网络架构进行编码，将其输入到一个多层线性回归器中，通过随机抽样收集训练数据，并在硬件上对其进行评估。

数据收集

首先硬件平台是在在NVIDIA Tesla-P100 GPU上，然后以（架构，延迟）对的方式整理数据集。本文从 DARTS 空间中随机抽取了 100K 个架构，并用随机化的网络权重来评估每个架构的延迟。为了使搜索到的架构具有更好的传输性，本文选择在 ImageNet 下测量延迟，其中输入图像大小为224×224，延迟结果是20次测量的平均值。整个过程大约需要9小时。虽然 100K 与整个搜索空间相比是一个小数字（搜索空间有 $1.0 \times 10^9$ 个不同的 Normal Cell），但对于学习任务来说已经足够了。然后，我们将延迟数据分为两部分：80K 对用于训练，剩余 20K 用于验证。

训练细节

在 80K 的训练集上，网络从头开始训练 1000 个 Epochs，使用 200 个批次大小。我们使用固定学习率为 0.01 、动量为 0.9 、权重衰减为 $1 \times 10^{-5}$ 、损失函数为均方误差（MSE）、优化器选择基于动量的 SGD 。

延时预测结果

这里本文使用测试集上预测与 ground-truth 之间的绝对误差和相对误差来评估 LPM 。如下表所示：

延时预测结果

随着训练数据量的增加，测试误差也相应下降。另一方面，当训练数据量大于40K时，精度的提高变得微乎其微。在训练数据量为80K的情况下，延迟预测结果令人满意，绝对误差小于2ms，相对误差小于10%。在根据搜索实验，这样的精度足以找到高效的架构。

可微分时延损失

定义时延损失

$\operatorname{LAT}(\boldsymbol{\alpha})=\mathbb{E}[\operatorname{LPM}(\gamma \mid \gamma \sim \mathcal{S}(\boldsymbol{\alpha}))]$
其中， $\boldsymbol{\alpha}$ 是搜索的架构参数； $\mathcal{S}(\boldsymbol{\alpha})$ 代表 $\boldsymbol{\alpha}$ 的分布，在实际操作中，从 $\boldsymbol{\alpha}$ 中均匀地抽取 14 条边中的 8 条，然后根据当前操作的权重随机选择每条边上的操作;为了保证当前搜索架构延迟预测的稳定性和准确性，使用批量大小为 M=20 ，即对 M 个子网络进行采样， $\left\{\gamma_{m}\right\}_{m=1}^{M}$ ，并取均值作为预测的结果 $\operatorname{LAT}(\boldsymbol{\alpha}) \approx \frac{1}{M} \sum_{m=1}^{M} \operatorname{LPM}\left(\boldsymbol{\gamma}_{m}\right)$ ； $\operatorname{LAT}(\boldsymbol{\alpha})$ 相对于 $\boldsymbol{\alpha}$ 的梯度是通过对 $\gamma_{m}$ 的梯度进行平均来估计的。

嵌入到搜索的Loss中

$\mathcal{L}_{\text {total }}(\boldsymbol{\alpha})=\mathcal{L}_{\text {val }}(\boldsymbol{\alpha})+\lambda \cdot \operatorname{LAT}(\boldsymbol{\alpha})$
这里，平衡系数 $\lambda$ 控制了精度和性能之间的权衡：较小的 $\lambda$ 更倾向于精度而不是延迟，反之亦然。请注意， $\lambda$ 的单位为 $\mathbf{s}^{-1}$ 。本文在实验中表明，选择一个合适的 $\lambda$ 并不困难，调整 $\lambda$ 可以导致架构的不同特性。

实验

CIFAR

CIFAR10实验结果

ImageNet

ImageNet实验结果

总结

本文提出了一种在复杂搜索空间中预测架构延迟的可微分方法，并将该模块融入到可微分架构搜索中。这使得我们可以控制识别精度和推理速度的平衡。本文将延迟预测模块设计成一个多层回归网络，并通过从预设的搜索空间中抽取一些架构进行训练。该流水线可以很容易地移植到广泛的硬件/软件配置中，并有助于设计机器友好的架构。

但是，对于不同的硬件架构和软件部署平台，需要构造不同的数据集才能得到准确的预测结果，因此，本文提出的推理延迟预测模块严重依赖数据集的构建。