FBNet系列论文是由Facebook推出的NAS算法搜索到的网络,V1、V2采用的是和DARTS一样的方法,通过构建Supernet和微分梯度方法计算出最佳的网络;V3采用的是自己独特的方法——JointNAS,将超参数和训练策略都作为搜索空间,通过粗粒度和细粒度两个stage搜索出来最好的网络和训练参数。
FBNet的训练方法和构建网络的方法基本上沿用了DARTS的方法,不同的是DARTS的Supernet主要体现在Cell结构里,搜索空间包括cell内部的连接和路径上的操作;而FBNet体现在整个网络的主体里,连接是确定的,只搜索路径上的操作。
图1是经典的可微NAS的搜索过程。搜索空间是一系列的可选操作,在super net里,每个搜索路径上都并列放置着搜索空间里的操作,有不同的架构参数用于表示这些操作的权重;FBNet在损失函数里加入端侧Latency的项,和准确率一起作为优化的目标;在搜索结束后,根据架构参数生成每条路径上的概率分布,从中采样出最好的网络作为输出(一般选取每条路径上概率值最高的操作)。
表1. FBNet的整体网络架构
FBNet的整体网络架构如表1所示。左起第一栏表示每个block的输入shape;第2栏表示block的模块设计,除了第一个阶段是 3 × 3 3\times3 3×3的卷积和最后三层conv+avgpool+fc外,中间全部都是TBS模块,TBS模块的结构如图2所示;第3栏表示每个阶段的输出channel大小;第4栏表示每个stage模块的重复个数;最右边一栏表示每个stage的缩放的倍数,也就是stride的大小。
从表1可以看出,FBNet的整体结构类似ResNet,从头到尾是一个直筒的按顺序操作的架构,并没有旁路的分叉,不像DARTS的cell里面还有不同block之间的连接。
图2显示的是TBS block的内部结构,这个结构的设计借鉴于MobileNetV2和ShiftNet的思想。TBS由两条路组成,左边先通过 1 × 1 1\times1 1×1的分组卷积将channel放大(后面跟一个shuffle的操作,混合不同group的特征),然后再用核为 K × K K\times K K×K大小的可分离卷积进行计算(如果某个stage的stride为2,在该stage的第一个TBS block的可分离卷积stride设置成2),最后再用 1 × 1 1\times1 1×1的分组卷积将channel调整成stage的输出通道数大小;右边的分支就是一个跳跃连接,如果stride为1,就是直接连接,如果stride为2,通过stride=2的卷积将特征图调整成和左边一致的shape大小。
表2表示的TBS的9个可选操作,名称也很容易理解,k表示可分离卷积核大小,e表示中间channels的放大倍数,g表示分组的个数。
FBNet的损失函数不只有标准的分类交叉熵损失项,还加上了网络的时延Latency,具体公式如下所示:
L ( a , w a ) = C E ( a , w a ) ⋅ α l o g ( L A T ( a ) ) β L(a, w_a)=CE(a, w_a)\cdot \alpha log(LAT(a))^{\beta} L(a,wa)=CE(a,wa)⋅αlog(LAT(a))β
其中, a a a表示网络架构, w a w_a wa表示网络的权重参数, C E CE CE表示的是交叉熵函数, α \alpha α是时延项的系数, L A T LAT LAT表示网络架构的时延Latency。从这个公式可以看出,该损失函数同时考虑了分类任务的交叉熵函数和网络时延的平衡,目的是找到一个准确率高且推理时延低的子网络。
在时延的计算中,作者采用的是组合的方式去计算总的网络时延。首先对各个stage中可能出现的算子去计算单算子的时延,汇总成表,总的网络时延是这些算子时延在所有网络层的加权和。公式为:
L A T ( a ) = ∑ l ∑ i m l , i ⋅ L A T ( b l , i ) LAT(a)=\sum_{l}\sum_{i}m_{l,i}\cdot LAT(b_{l,i}) LAT(a)=l∑i∑ml,i⋅LAT(bl,i)
其中, l l l表示第 l l l个网络层, i i i表示该网络层的第 i i i个可选操作。 m l , i m_{l,i} ml,i表示第 l l l个网络中第 i i i个操作的权重参数,具体计算公式为:
m l , i = e x p [ ( θ l , i + g l , i ) / τ ] ∑ i e x p [ ( θ l , i + g l , i ) / τ ] m_{l,i}=\frac{exp[(\theta_{l,i}+g_{l,i})/\tau]}{\sum_iexp[(\theta_{l,i}+g_{l,i})/\tau]} ml,i=∑iexp[(θl,i+gl,i)/τ]exp[(θl,i+gl,i)/τ]
这个公式叫做Gumbel Softmax Function,经常用来转换那些原本不可微的损失函数。其中 θ \theta θ就是网络的架构参数, g g g是一个(0,1)区间的随机噪声, τ \tau τ是温度超参, τ \tau τ接近0时,这个公式就接近于离散采样分布, τ \tau τ越大时,就变成了一个连续的随机变量。
论文在实验部分提到,在搜索过程中,只选取了ImageNet其中的100类作为训练数据,训练90个epoch,每个epoch中,80%的数据用来训练网络参数,20%的数据用来训练架构参数,统计latency的设备采用三星Galaxy S8。搜索完成后,从中挑取了三个不同大小的代表网络进行训练,结果如图3所示。
FBNet-A和FBNet-B、FBNet-C的区别在于最后一个卷积的输出channel不一样。从实验中可以看出,在相同水平的时延条件下,FBNet的网络具有相对更好地分类精度。
在FBNetV1中,搜索空间主要是每一层网络的一些可选操作,输入输出的channel都是固定的,每个stage里面的特征分辨率也是固定的。FBNetV2在V1的基础上,增加了对这两项的搜索,既能搜索V1里面的操作,也能搜索每个stage的输入输出channel和stage内部的分辨率缩放值。
FBNetV2的整体网络结构基本上沿用了V1的设计(表1),在TBS的可选操作上做了调整(表2)。另外最重要的是,在网络架构参数上,加入了通道和分辨率两个维度。作者在论文里面讨论了很多设计的过程,从最初不太合理的设计到如何调整成相对合理的设计,最后的结论其实都汇总在图1里。
图1是作者最后给出他们的设计方式,左图是通道搜索的设计,右图是分辨率降采样搜索的设计。
在通道搜索中,作者是假设不同的每一层的通道路径都共享一个卷积核参数,不同channel的输出取决于一个叫做mask的向量。比如左图中 g 1 g_1 g1、 g 2 g_2 g2和 g 3 g_3 g3是不同channel选择的架构参数,对应的白蓝条状是mask向量,其中白色表示1,蓝色表示0,分别和卷积相乘,那么白色对应的部分表示该路径上的输出channel个数,蓝色部分表示该路径上没有这部分channel,但是在整体计算上维度是保持一致的,不同mask向量分别和卷积核参数相乘,再用channel的架构参数加权和。这样在搜索的过程中,通过架构参数就可以知道每个stage中的卷积核该选择多少个输出channel。
在分辨率搜索中,共享的是特征向量,不同的分辨率会从共享特征向量里面去抽取。抽取的方法是先确保角落块被选中,然后根据缩放比例去调整抽取的步长。比如,缩放比为2的话,选中的就是四个角落1/4大小的特征(右图中间),缩放比为3的话,选中的四个角落1/9大小的特征,剩下再按平均分布的原则去抽取特征(右图上侧)。抽取完后,把这些特征拼在一起,用卷积核去计算,然后再反变换回原先的位置。其他未被选中位置的特征补零。
每个stage的channel和分辨率搜索空间如表1所示。
在损失函数上,权重系数沿用了V1中的Gumbel Softmax的方法,将Latency的损失项换成了FLOPS/Params,FLOPS/Params的计算方法与标准方法无异,唯一的区别在于它们也是通过权重系数来做加权和计算。比如在输出通道上,计算方法为:
C ‾ o u t l = ∑ i g i l ⋅ C i , o u t l \overline{C}_{out}^{l}=\sum_ig_i^l \cdot C_{i,out}^l Coutl=i∑gil⋅Ci,outl
其中 C i , o u t l C_{i,out}^l Ci,outl表示第 l l l层第 i i i条路径的通道数,和图1里对应的mask向量中1的个数一致。宽和高的计算方法和通道的一致。
论文同样是在ImageNet数据集上做实验,搜索的过程中,搜索参数的设置和V1的一致。随机选取10%的类别数据,训练90个epoch,每个epoch中80%数据用于训练网络参数,20%的数据用于训练架构参数。搜索结束后,选取几个代表网络进行Fully Train,实验结果如表3所示。
从该表格可以看出,在相同级别FLOPS的网络中,FBNet-V2同时具有更小的计算量和更高的分类精度。
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部分,具体迭代步骤为:
在第一步迭代的时候,还要确定早停方案,早停主要是为了找到样本训练(网络本身的训练,不是对Predictor的训练)的epoch参数,确定方法为:
使用早停策略得到的网络准确率,就可以用来训练更新预测器了,在更新预测器时也有几点tricks。首先,使用Huber Loss减少不正常样本的影响;其次,开始时冻结embedding层,只训练accuracy Predictor,50个epoch;最后,再训练整个Predictor,逐步减少学习率,50个epoch。
细粒度搜索的空间是网络架构+训练超参,搜索的方法是自适应遗传算法,如图2中的Stage所示,搜索步骤为:
如一开头所说的,FBNetV3的搜索空间包括了训练超参和网络架构。训练超参的搜索空间包括了优化器类型、初始学习率、参数正则化比例、mixup比例、dropout比例、随机深度drop比例和是否使用EMA等。网络架构的搜索空间是逆残差模型的参数,包括输入分辨率、卷积核大小、中间通道放大比例、每一stage网络的通道数和深度等。FBNetV3的搜索空间如表1所示。
作者在论文中提到,如果同时搜索训练策略和网络架构,会导致搜索空间的可能性急剧膨胀,基本上很难让搜索算法收敛到比较好的效果。为了加快搜索的效率,作者采用先搜索训练超参再搜索网络架构的方法。
首先固定网络(FBNetV2-L3)作为基础网络去搜索训练策略,得到了一组较好的训练超参。在实验的过程中普遍发现使用EMA的模型在训练中间过程上效果都会好于原始模型,所以在后面搜索网络架构时都会使用EMA方法。
作者在搜索训练策略和网络架构上均使用了Two-Stage的方法,根据不同的FLOPS约束,搜索到了不同的V3网络结果。具体实验结果如表2所示,从表中可以看出,在同等精度的条件下,相比其他手工设计和NAS的网络,FBNetV3大幅减小了计算量FLOPS。