FBNet系列论文解析

FBNet系列论文是由Facebook推出的NAS算法搜索到的网络，V1、V2采用的是和DARTS一样的方法，通过构建Supernet和微分梯度方法计算出最佳的网络；V3采用的是自己独特的方法——JointNAS，将超参数和训练策略都作为搜索空间，通过粗粒度和细粒度两个stage搜索出来最好的网络和训练参数。

FBNetV1

网络结构和搜索空间设计

FBNet的训练方法和构建网络的方法基本上沿用了DARTS的方法，不同的是DARTS的Supernet主要体现在Cell结构里，搜索空间包括cell内部的连接和路径上的操作；而FBNet体现在整个网络的主体里，连接是确定的，只搜索路径上的操作。

图1. 可微NAS的过程

图1是经典的可微NAS的搜索过程。搜索空间是一系列的可选操作，在super net里，每个搜索路径上都并列放置着搜索空间里的操作，有不同的架构参数用于表示这些操作的权重；FBNet在损失函数里加入端侧Latency的项，和准确率一起作为优化的目标；在搜索结束后，根据架构参数生成每条路径上的概率分布，从中采样出最好的网络作为输出（一般选取每条路径上概率值最高的操作）。

表1. FBNet的整体网络架构

FBNet的整体网络架构如表1所示。左起第一栏表示每个block的输入shape；第2栏表示block的模块设计，除了第一个阶段是$3\times 3$的卷积和最后三层conv+avgpool+fc外，中间全部都是TBS模块，TBS模块的结构如图2所示；第3栏表示每个阶段的输出channel大小；第4栏表示每个stage模块的重复个数；最右边一栏表示每个stage的缩放的倍数，也就是stride的大小。

从表1可以看出，FBNet的整体结构类似ResNet，从头到尾是一个直筒的按顺序操作的架构，并没有旁路的分叉，不像DARTS的cell里面还有不同block之间的连接。

图2. TBS内部结构

表2. TBS的可选搜索空间

图2显示的是TBS block的内部结构，这个结构的设计借鉴于MobileNetV2和ShiftNet的思想。TBS由两条路组成，左边先通过$1\times 1$的分组卷积将channel放大（后面跟一个shuffle的操作，混合不同group的特征），然后再用核为$K\times K$大小的可分离卷积进行计算（如果某个stage的stride为2，在该stage的第一个TBS block的可分离卷积stride设置成2），最后再用$1\times 1$的分组卷积将channel调整成stage的输出通道数大小；右边的分支就是一个跳跃连接，如果stride为1，就是直接连接，如果stride为2，通过stride=2的卷积将特征图调整成和左边一致的shape大小。

表2表示的TBS的9个可选操作，名称也很容易理解，k表示可分离卷积核大小，e表示中间channels的放大倍数，g表示分组的个数。

损失函数设计

FBNet的损失函数不只有标准的分类交叉熵损失项，还加上了网络的时延Latency，具体公式如下所示：

$$L(a,w_a)=CE(a,w_a)\cdot \alpha log(LAT(a){)}^{\beta }$$

其中，$a$表示网络架构，$w_a$表示网络的权重参数，$CE$表示的是交叉熵函数，$\alpha$是时延项的系数，$LAT$表示网络架构的时延Latency。从这个公式可以看出，该损失函数同时考虑了分类任务的交叉熵函数和网络时延的平衡，目的是找到一个准确率高且推理时延低的子网络。

在时延的计算中，作者采用的是组合的方式去计算总的网络时延。首先对各个stage中可能出现的算子去计算单算子的时延，汇总成表，总的网络时延是这些算子时延在所有网络层的加权和。公式为：

$$LAT(a)=\sum _l\sum _i{m}_{l,i}\cdot LAT(b_{l,i})$$

其中，$l$表示第$l$个网络层，$i$表示该网络层的第$i$个可选操作。$m_{l,i}$表示第$l$个网络中第$i$个操作的权重参数，具体计算公式为：

$$m_{l,i}=\frac{exp[({\theta }_{l,i}+g_{l,i})/\tau ]}{{\sum }_{i}exp[({\theta }_{l,i}+g_{l,i})/\tau ]}$$

这个公式叫做Gumbel Softmax Function，经常用来转换那些原本不可微的损失函数。其中$\theta$就是网络的架构参数，$g$是一个（0，1）区间的随机噪声，$\tau$是温度超参，$\tau$接近0时，这个公式就接近于离散采样分布，$\tau$越大时，就变成了一个连续的随机变量。

实验结果

图3.FBNet在imagenet上的实验结果

论文在实验部分提到，在搜索过程中，只选取了ImageNet其中的100类作为训练数据，训练90个epoch，每个epoch中，80%的数据用来训练网络参数，20%的数据用来训练架构参数，统计latency的设备采用三星Galaxy S8。搜索完成后，从中挑取了三个不同大小的代表网络进行训练，结果如图3所示。

FBNet-A和FBNet-B、FBNet-C的区别在于最后一个卷积的输出channel不一样。从实验中可以看出，在相同水平的时延条件下，FBNet的网络具有相对更好地分类精度。

FBNetV2

在FBNetV1中，搜索空间主要是每一层网络的一些可选操作，输入输出的channel都是固定的，每个stage里面的特征分辨率也是固定的。FBNetV2在V1的基础上，增加了对这两项的搜索，既能搜索V1里面的操作，也能搜索每个stage的输入输出channel和stage内部的分辨率缩放值。

网络结构和搜索空间的设计

FBNetV2的整体网络结构基本上沿用了V1的设计（表1），在TBS的可选操作上做了调整（表2）。另外最重要的是，在网络架构参数上，加入了通道和分辨率两个维度。作者在论文里面讨论了很多设计的过程，从最初不太合理的设计到如何调整成相对合理的设计，最后的结论其实都汇总在图1里。

图1. 通道调整设计和分辨率调整设计

图1是作者最后给出他们的设计方式，左图是通道搜索的设计，右图是分辨率降采样搜索的设计。

在通道搜索中，作者是假设不同的每一层的通道路径都共享一个卷积核参数，不同channel的输出取决于一个叫做mask的向量。比如左图中$g_1$、$g_2$和$g_3$是不同channel选择的架构参数，对应的白蓝条状是mask向量，其中白色表示1，蓝色表示0，分别和卷积相乘，那么白色对应的部分表示该路径上的输出channel个数，蓝色部分表示该路径上没有这部分channel，但是在整体计算上维度是保持一致的，不同mask向量分别和卷积核参数相乘，再用channel的架构参数加权和。这样在搜索的过程中，通过架构参数就可以知道每个stage中的卷积核该选择多少个输出channel。

在分辨率搜索中，共享的是特征向量，不同的分辨率会从共享特征向量里面去抽取。抽取的方法是先确保角落块被选中，然后根据缩放比例去调整抽取的步长。比如，缩放比为2的话，选中的就是四个角落1/4大小的特征（右图中间），缩放比为3的话，选中的四个角落1/9大小的特征，剩下再按平均分布的原则去抽取特征（右图上侧）。抽取完后，把这些特征拼在一起，用卷积核去计算，然后再反变换回原先的位置。其他未被选中位置的特征补零。

每个stage的channel和分辨率搜索空间如表1所示。

表1. 整体网络架构设计和channel、分辨率的搜索空间

表2. TBS模块的搜索空间

损失函数设计

在损失函数上，权重系数沿用了V1中的Gumbel Softmax的方法，将Latency的损失项换成了FLOPS/Params，FLOPS/Params的计算方法与标准方法无异，唯一的区别在于它们也是通过权重系数来做加权和计算。比如在输出通道上，计算方法为：

$${\overline{C}}_{out}^l=\sum _i{g}_i^l\cdot C_{i,out}^l$$

其中$C_{i,out}^l$表示第$l$层第$i$条路径的通道数，和图1里对应的mask向量中1的个数一致。宽和高的计算方法和通道的一致。

实验结果

论文同样是在ImageNet数据集上做实验，搜索的过程中，搜索参数的设置和V1的一致。随机选取10%的类别数据，训练90个epoch，每个epoch中80%数据用于训练网络参数，20%的数据用于训练架构参数。搜索结束后，选取几个代表网络进行Fully Train，实验结果如表3所示。

从该表格可以看出，在相同级别FLOPS的网络中，FBNet-V2同时具有更小的计算量和更高的分类精度。

表3. FBNetV2在ImageNet数据集上的实验结果

FBNetV3

FBNetV3使用的搜索方法和V1、V2都不同，不再采用对SuperNet梯度下降的方法，而是使用一种全新的方法，叫做JointNAS，分粗粒度和细粒度两个阶段，对网络架构和训练超参都进行搜索。

粗粒度搜索

粗粒度阶段的搜索主要是为了训练网络预测器（Predictor），这个预测器是一个多层感知器构成的小型网络，包含了两个部分，一个代理预测器（Proxy Predictor），一个是准确率预测器（Accuracy Predictor）。具体结构如图1所示。

图1. 网络架构预测器

图1中的第一排部分就是用来预测Proxy任务的，使用的标签数据是FLOPS和Params。左边的architecture presentation指的是网络架构的编码描述，以one-hot或者int的形成一组矩阵向量，用来表示网络的结构（参考DARTS的架构参数），architecture encoder就是多层感知器，用来提取架构描述的特征，生成architecture embedding，也就是低维特征向量。这个低维特征向量首先用于Proxy Predictor的训练过程，由于任何一个网络架构描述都可以对应到一个实际的网络模型，同时也能对应它的计算量和参数量数值，所以在这个阶段的训练，并不需要额外的数据。

在预训练好第一排的网络后，就要来迭代训练第二排的准确率预测器了，这个预测器的输入数据为训练超参表示加上低维特征向量，而输出数据则是这个网络架构+训练超参的结果准确率。迭代优化的算法参考图2中Stage1部分，具体迭代步骤为：

1. 基于已经预测的准确率，选择一组候选集，选择的方法为quosi蒙特卡洛（QMC）；
2. 训练和评估候选集的网络模型，得到评估准确率；
3. 使用所有历史候选集的准确率和表示输入（架构+超参）去更新这个预测器。

在第一步迭代的时候，还要确定早停方案，早停主要是为了找到样本训练（网络本身的训练，不是对Predictor的训练）的epoch参数，确定方法为：

1. 选取n个网络，对完整训练的结果和早停训练的结果分别作排序；
2. 计算两种排序的相关性；
3. 如果早停epoch达到某个值时，两种排序的相关性达到阈值（文中是0.92），则认为确定好了早停的epoch，否则重新加大epoch，进入步骤1重新开始。

使用早停策略得到的网络准确率，就可以用来训练更新预测器了，在更新预测器时也有几点tricks。首先，使用Huber Loss减少不正常样本的影响；其次，开始时冻结embedding层，只训练accuracy Predictor，50个epoch；最后，再训练整个Predictor，逐步减少学习率，50个epoch。

图2. 两个Stage的架构搜索（粗粒度+细粒度）

细粒度搜索

细粒度搜索的空间是网络架构+训练超参，搜索的方法是自适应遗传算法，如图2中的Stage所示，搜索步骤为：

1. 选择最好的几个样本作为种群（第一代从粗粒度中选）；
2. 在给定约束的条件下，对这些样本使用变异的方法产生一个子代的种群；
3. 使用准确率预测器生成子代种群的得分，如果最好的个体不再更新了，就停止迭代，否则进入步骤1。

如一开头所说的，FBNetV3的搜索空间包括了训练超参和网络架构。训练超参的搜索空间包括了优化器类型、初始学习率、参数正则化比例、mixup比例、dropout比例、随机深度drop比例和是否使用EMA等。网络架构的搜索空间是逆残差模型的参数，包括输入分辨率、卷积核大小、中间通道放大比例、每一stage网络的通道数和深度等。FBNetV3的搜索空间如表1所示。

表1. FBNetV3的搜索空间

实验结果

作者在论文中提到，如果同时搜索训练策略和网络架构，会导致搜索空间的可能性急剧膨胀，基本上很难让搜索算法收敛到比较好的效果。为了加快搜索的效率，作者采用先搜索训练超参再搜索网络架构的方法。

首先固定网络（FBNetV2-L3）作为基础网络去搜索训练策略，得到了一组较好的训练超参。在实验的过程中普遍发现使用EMA的模型在训练中间过程上效果都会好于原始模型，所以在后面搜索网络架构时都会使用EMA方法。

作者在搜索训练策略和网络架构上均使用了Two-Stage的方法，根据不同的FLOPS约束，搜索到了不同的V3网络结果。具体实验结果如表2所示，从表中可以看出，在同等精度的条件下，相比其他手工设计和NAS的网络，FBNetV3大幅减小了计算量FLOPS。

表2. FBNetV3不同计算量级别的实验结果以及与其他网络的对比