[Reading] MnasNet: Platform-Aware Neural Architecture Search for Mobile

1. 概述

  MnasNet: Platform-Aware Neural Architecture Search for Mobile 一文初次提交于 2018/7,文章提出了一种移动端神经架构搜索(mobile neural architecture search,MNAS)的方法,同时将模型的准确率和在真实设备上的延迟作为搜索的目标,以获得二者之间的平衡。文章还提出了一种分层搜索空间(factorized hierarchical
search space),来提高各层结构的多样性。最终得到的 MnasNet 在图像识别和目标检测任务中,相较其他移动端模型,在保持低延迟的同时,获得了 SOTA 性能,如 Figure 2 所示。

Figure 2

Figure 2

  针对如移动端等计算能力有限的场景,需要对模型准确率和延迟进行合理的取舍,增加了人工进行网络设计的难度。因此文章希望通过自动搜索的方法来找到合适的网络架构,通过将准确率和延迟同时作为搜索的目标,来对准确率和延迟进行取舍。

  前人的工作中,通常使用 FLOPS 来评价模型延迟,然而这是不准确的。例如 MobileNet 和 NASNet 具有相近的 FLOPS,分别为 575M 和 564M,然而二者的延迟则有很大差别,分别为 113ms 和 183ms。因此文章使用真实设备来测试模型的延迟,贴合模型的实际使用情况。

  在之前的神经架构搜索(NAS)方法中,为了简化搜索过程,通常只是搜索几个 cell 的架构,然后通过重复堆叠得到完整网络,网络的层结构比较单一。为了提高计算效率,有必要增加层结构的多样性,例如 CNN 在前后阶段的计算量不同,需要使用不同的结构来有效地利用计算资源。因此文章提出了分层搜索空间,使搜索得到的各层可以有不同的结构,平衡活性和搜索空间大小。

2. 搜索方法

2.1. 搜索框架

  文章使用的搜索框架如 Figure 1 所示,其中 Controller 对模型进行采样,得到的模型由 Trainer 训练后,一方面在验证集上计算准确率,一方面在真实设备上计算延迟,二者共同构成收益,据此更新控制器参数。

Figure 1

Figure 1

2.2. 搜索目标

  给定模型 $m$,记模型在目标任务上的准确率为 $ACC(m)$,在目标移动平台上的推理延迟为 $LAT(m)$,目标延迟为 $T$。为了在搜索时同时考虑模型准确率和延迟,常见做法是将 $T$ 作为硬约束,并最大化准确率,即

\begin{align}
\underset{m}{\mathrm{maximize}} \; & ACC(m) \\
\mathrm{subject \; to} \; & LAT(m) \leq T \tag{1}
\end{align}

文章指出,这种方法只是最大化了准确率这个单一指标,无法得到帕累托最优(Pareto optimal),即在不增加延迟的情况下达到最大准确率,或在不降低准确率的情况下达到最低延迟。

  文章使用了准确率和延迟的加权乘积作为优化目标

\begin{equation}
\underset{m}{\mathrm{maximize}} \; ACC(m) \times \bigg[ \frac{LAT(m)}{T} \bigg]^w \tag{2}
\end{equation}

其中 $w$ 为权重因子,定义为

\begin{equation}
w = \begin{cases}
\alpha & if LAT(m) \leq T \\
\beta & \mathrm{otherwise} \end{cases}
\end{equation}

  如果设置 $\alpha = 0, \beta = -1$,则当 $LAT(m) \leq T$ 时,优化目标将只包含 $ACC(m)$;当 $LAT(m) > T$,将进行惩罚,如 Figure 3 上图所示。

Figure 3

Figure 3

  选择 $\alpha$ 和 $\beta$ 的一个经验准则是,保证不同准确率-延迟取舍的帕累托最优解具有相近的收益。文章观察到加倍延迟会给准确率带来大约 5% 的相对提升,假设模型 M1 延迟为 $l$,准确率为 $a$,收益为 $Reward(M1) = a \cdot (l / T)^\beta$,则对于延迟为 $2l$ 的模型 M2,其准确率为 $a \cdot (1 + 5\%)$,收益为 $Reward(M2) = a \cdot (1 + 5\%) \cdot (2l / T)^\beta$。这两个模型应具有相同的收益,即 $Reward(M2) \approx Reward(M1)$,可以解得 $\beta \approx -0.07$,因此文章使用 $\alpha = \beta = -0.07$,如 Figure 3 下图所示。

2.3. 分层搜索空间

  为了提高层结构的多样性,文章将 CNN 网络分成一系列预定义的 block,逐渐降低分辨率并增加过滤器数量,通过对各个 block 进行搜索,使得不同的 block 可以具有不同的层架构,如 Figure 4 所示。

Figure 4

Figure 4

其中每个 block 的子搜索空间包括

  Figure 4 中的 ConvOp、KernelSize、SERatio、SkipOp 和 FilterSize $F_i$ 决定了层的架构,$N_i$ 控制该层在 block 中重复的次数。例如 Figure 4 中的 Layer 4-1 表示选择了一个 inverted bottleneck 和一个残差连接,该层在 Block 4 中重复了 $N_4$ 次。

  分层搜索空间的一大优势是平衡了层的多样性和搜索空间的大小。假设网络有 $B$ 个 block,每个 block 平均有 $N$ 层,子搜索空间为 $S$,对于分层搜索空间,每个 block 只搜索一种层结构,总搜索空间大小为 $S^B$。如果逐层进行搜索,则搜索空间大小为 $S^{B * N}$。典型地,对于 $S = 432, B=5, N = 3$,分层搜索空间的大小为 $10^{13}$,逐层搜索的搜索空间的大小为 $10^{39}$。

2.4. 搜索算法

  文章延续了 NASNet 的思路,将学习生成神经网络架构看成是一个强化学习问题,将控制器基于参数 $\theta$ 输出的网络结构参数看成是动作 $a_{1:T}$,目标是最大化期望收益

\begin{equation}
J = E_{P(a_{1:T};\theta)}[R(m)] \tag{5}
\end{equation}

其中 $m$ 是由动作 $a_{1:T}$ 决定的模型,$R(m)$ 是式 $(2)$ 定义的目标值。

  文章也使用了 RNN 作为控制器,控制器首先根据其参数 $\theta$ 采样一批模型,对每个模型 $m$,训练后获得准确率 $ACC(m)$,然后在真实设备上执行推断获得延迟 $LAT(m)$,计算得到收益 $R(m)$,据此更新控制器参数 $\theta$ 以最大化式 $(5)$ 定义期望收益。

3. 实验结果

  由于直接在如 ImageNet 和 COCO 这些大数据集上进行搜索的计算量太大,之前的 NAS 通常使用如 CIFAR-10 的小数据集。但文章指出,在引入延迟之后,小数据集不再适用,因此选择通过在 ImageNet 上训练少量(5个) epoch 来进行搜索,在 64 个 TPU 上,每次搜索耗时 4.5 天。控制器总共采样了 8K 个模型,选择了最好的 15 个迁移到 ImageNet,选择了最好的 1 个迁移到 COCO。

  网络在 ImageNet 上的性能如 Table 1 所示,搜索的目标延迟与 MobileNetV2 类似,设置 $T = 75 ms$,$\alpha = \beta = -0.07$,最终选择了具有不同延迟-准确率取舍的三个最佳模型参与对比。可见参数最少的 MnasNet-A1 就已经达到了 SOTA 性能,相较 MobileNetV2 (1.4x),准确率高 0.5%,速度快 1.8 倍。相较同样是搜索得到的 NASNet-A,MnasNet-A1 准确率高 1.2%,速度快 2.3 倍。

Table 1

Table 1

  MNasNet 中使用了 squeeze-and-excitation(SE),为了排除 SE 的影响,文章比较了不带 SE 的 MnasNet 的效果如 Table 2 所示。可见即便排除了 SE,MNasNet 的性能依旧优于 MobileNetV2。

Table 2

Table 2

  实际场景中,经常要通过增加深度或通道数的方法,对模型进行缩放,来获得不同准确率和延迟的模型。文章比较了 MNasNet 的缩放性能,如 Table 5 所示,可见 MNasNet 在不同缩放尺度上都稳定优于 MobileNetV2。

Figure 5

Figure 5

此外文章还专门为不同延迟单独进行了搜索,并比较了搜索和缩放两种方法的性能,如 Table 4 所示。可见与缩放模型相比,单独搜索的模型具有更好的性能。

Table 4

Table 4

  文章将 Table 2 中的 MNasNet 用于 SSDLite,测试了在 COCO 目标检测任务上的性能,如 Table 3 所示。可见 MnasNet-A1 + SSDLite 的性能显著优于 MobileNetV2 + SSDLite,与 SSD300 相当,而其参数数量和乘加数量远少于 SSD300。

Table 3

Table 3

4. 讨论

4.1. 软/硬约束

  文章测试了在式 $(2)$ 中使用不同 $\alpha$ 和 $\beta$ 的效果,如 Figure 6 所示。如果令 $\alpha = 0, \beta = -1$,则延迟的要求变成了硬约束,控制器会更偏向生成低延迟的模型;如果令 $\alpha = \beta = -0.07$,则延迟的要求是软约束,控制器搜索延迟的范围会更大,可以得到更多的帕累托最优解。

Figure 6

Figure 6

4.2. 搜索空间和收益

  文章中使用的多目标收益和分层搜索空间都有助于获得更好的模型。为了区分二者的贡献,文章分别测试了单/多目标收益和 NASNet(Cell-based) / MnasNet 搜索空间的组合,如 Table 5 所示。可见对于相同的搜索空间,相比于单目标收益,多目标收益可以在损失准确率的情况下,获得更快的模型;对于相同的收益,MnasNet 搜索空间可以找到更快和更准确率更高的模型。

Table 5

Table 5

4.3. 层的多样性

  分层搜索空间有助于提高层结构的多样性。Figure 7(a) 展示了 MnasNet 架构,可见它确实具有更丰富的层结构。

Figure 7

Figure 7

  Table 6 比较了 MnasNet 和只重复一种层结构的效果,可见 MnasNet 可以更好地平衡准确率和延迟。

Table 6

Table 6

5. 总结

  文章提出了一种可以对准确率和延迟进行平衡的神经架构搜索方法,通过使用多目标搜索和分层搜索空间,来生成适用于如移动端的资源受限场景下使用的 CNN 网络。通过搜索得到的 MNASNet 在图像识别和目标检测任务中,都达到了移动端模型的 SOTA 性能。