本文介绍了One-Shot NAS中的OFA网络架构搜索方法，它通过训练超网并进行子网采样来实现模型压缩。OFA采用Progressive Shrinking策略训练超网，并利用子网精度预测器和延迟查找表加速搜索过程。文章展示了在实际应用中，通过对GAN网络的模型压缩，找到性能与计算量之间的平衡点。
之前的文章中讲到了轻量化网络架构(点这里查看)的设计，也提到了模型压缩除了轻量化架构外，还有
- 模型剪枝
- 模型架构搜索
模型架构搜索（NAS）也是一个非常有效的模型压缩方法，相比人工设计架构和剪枝，机器搜索架构更高效且效果更好，但是之前基于遗传和强化算法的模型架构搜索方法需要大量的 GPU 资源支持，且训练时间长，限制了 NAS 的应用，随着研究进展，出现了一些不那么耗时耗资源的方法。
Once for All (简称 OFA)就是这样一种网络架构搜索方法，这里会结合论文+源码的方式来进行讲解
首先，为了让大家更好的理解，我们需要先了解一下什么是 One-Shot NAS，因为 OFA 是属于这个流派的算法之一
什么是 One-Shot NAS
一般的 NAS 算法都会包含以下几个步骤，首先是定义网络的搜索空间（最大层数，每个节点可选择的操作），然后则是设计搜索策略（如何探索空间，目前主流的有强化学习，进化算法，梯度优化），最后则是对每个子网络进行评估。
One-Shot NAS (简称 OSN)与一般 NAS 最大的不同之处在于他是训练一个超网，然后在超网上做减法（其实和剪枝有点像），而之前的 NAS 更像是从小网络开始做加法。
在搜索和训练上，OSN 主要包含两个阶段，分别是 training stage 和 searching stage，在训练阶段，他会训练一次 super-net(超网)，然后便可以针对不同的 Constraint，从超网中抽取得到各种不同的 Sub-network。
OFA 中超网的定义
把原版的 Resnet 和 MobileNet V3 当成了超网，限定了最大的搜索空间。
因为我们后续需要从超网中抽取子网，因此还需要对超网架构进行参数化处理，这里限定了卷积核大小取值，卷积层数（深度）取值，通道数（宽度）取值，根据调节不同的参数，我们便可以在超网中获得子网。
弹性分辨率：训练的时候，分辨率的大小是动态的，而不是固定的，这个通过修改 dataloader 便可以做到。
弹性卷积核：每层卷积所对应的卷积核大小也是可以动态调整，而不是固定的，但是因为这里涉及到对应的权重参数，文中认为如果大卷积核如果训练的好的话，其中心的小卷积核通过简单的非线性映射（F.linear）也可适用
弹性宽度：每层卷积对应的输入和输出通道数也是可以动态调整，而不是固定的，不过这里会先进行通道重要性指标的排序后，再通过 slices 来实现
上述的参数便可以组成构建超网的基础组件（DynamicConv），源码的核心实现代码如下：
class DynamicConv2D(nn.Module):
def __init__(self,
kernel_list = [3, 5, 7],
in_channel_list = [3],
out_channel_list = [32, 64, 128]):
# 这里的 active 都取最大，但在训练中
# 会使用 random 来选择动态的参数
self.max_kernel = max(kernel_list)
self.active_kernel = max(kernel_list)
self.active_in_channel = max(in_channel_list)
self.active_out_channel = max(out_channel_list)
self.conv = nn.Conv2d(self.active_in_channel, self.active_out_channel, self.active_kernel)
def forward(self, x):
# 取中心的卷积核对应的下标
start, end = sub_filter_start_end(self.max_kernel, self.active_kernel)
# Filters 的 Size = [out_channel, in_channel, kernel_size, kernel_size]
filters = self.conv.weight[:self.active_out_channel, :self.active_in_channel, start:end, start:end]
# 需要先 flatten，然后与 transform matrix 相乘
transfer_filter = F.linear(filters, transform_matrix)
y = F.conv2d(x, transfer_filter, self.active_in_channel, self.active_out_channel)
return y
弹性深度：架构所对应的深度也是可以动态调整，这里则利用跳层来实现的深度动态变化
这里只描述了卷积过程中常见的几个参数，但这些参数的变动是会引起其他算子的连锁变化的，在源码里针对常见的算子都做了动态调整的代码对应的实现，出于篇幅，这里就不展开讲，感兴趣的可以详细的去看看源码实现 <
训练超网
OSN 方法中，超网只需要训练一次，这个流程和正常我们训练模型的最大差别在于 forward pass 的时候，他会对超网的每一层对应的参数进行均匀采样，我这样直接讲或许有些抽象，结合下面伪代码则有更好的理解
# CAUTION: 伪代码
# 定义动态网络
net = dynamic_net()
for step in len(samples) / batch_size:
# 获得动态 Block
blobs = net.blocks
# 对每个动态 block 均匀采样
for b in blobs:
b.active_kernel = random.uniform([3, 5, 7])
b.active_out_channel = random.uniform([32, 64, 128, 256])
b.active_depth = random.uniform([3, 5, 7])
net.train()
net.update_parameters()
这样训练的方式，类似于之前的 dropout 层，不同的是它现在是在卷积的深度宽度和卷积上都做了，这样子在训练的时候每次更新的网络架构都是不一样的，参数也是不一样的。
OFA 论文中并没有均匀采样方式训练超网，而是通过 Progressive Shrinking 的方法来训练超网。这种方法文中认为可以更好的防止在训练过程中子网之间会互相影响。
简单来说，PS 训练方法分以下几个步骤
固定网络参数为最大值，文中 kernel_size =7, depth=4, width=6，训练直至收敛
然后动态调整 kernel_size = {3, 5, 7}，固定 depth = 4, width = 6, 训练直至收敛
同步骤 2 分别动态调整 depth 和 width，训练直至收敛
从在论文的实验结果来看，PS 方法对于最后出来的模型在精度上是有提高的
另外，文中提到他们训练超网是在一张 V100 的 GPU 上训练了大概50天，但其实我们在实际的应用中，数据量其实是远远没有 ImageNet 大的，一般情况下几天便可完成训练和搜索。
其实超网训练方式在业界都还存在有很多疑问点，其中一个疑问点便是所有子网参数共享且一起训练，他们之间就一定存在耦合，真的能够搜索到最佳的子网络架构吗？PS 虽然说能够缓解，但也还是在权重共享的模式下训练，但目前这块并没有进一步更好的理论解释，因此大家也就这么用着了。
子网络精度预测器
我们可以在测试集上对每个子网络的效果进行评估，但这种方式速度还是过慢了，因此文中提出了一种新的子网络评估方法，简单来说就是新增了一个三层神经网络模型，这个模型的输入是待评估子网的架构参数（One-Hot Encoding），输出是预测该子网精度数值，通过这种方式，我们只需要做一次推断便可以预测该子网的大致性能。
# accuracy predictor
accuracy_predictor = AccuracyPredictor(
pretrained=True,
device='cuda:0' if cuda_available else 'cpu'
)
print(accuracy_predictor.model)
# 评估网络的精度的网络架构如下图
# 其实就是三层 FC 构成
Sequential(
(0): Linear(in_features=128, out_features=400, bias=True)
(1): ReLU()
(2): Linear(in_features=400, out_features=400, bias=True)
(3): ReLU()
(4): Linear(in_features=400, out_features=400, bias=True)
(5): ReLU()
(6): Linear(in_features=400, out_features=1, bias=True)
)
延迟查找表
一般做模型压缩，轻量化架构和搜索，是为了能在某些设备上更快的运行起来，拿手机设备举个例子，因为存在很多的不同厂商，因此不同的算子在不同的设备上的计算速度存在差异，为此，文中提到了延迟查找表的概念，也就是说我们把我们要用到的算子在设备上进行一次基准测试，然后我们在预估模型在设备上计算时长的时候，就只需要查表就可以了，不用一次次的去部署然后做基准测试。
从查找表中可以看到，针对每个算子不同参数下所对应的结构都在设备上进行基准测试，我们搜索的架构也都是由这些算子所构成的，从而在评估耗时时，只需要查表便可以了。
搜索子网络
上述我们已经获得了子网络的精度的预测器了，以及如何评估子网络的运行速度，接下来我们便可以根据需求对子网络进行搜索和评估了。
对于子网络评估如果时间上允许的话，我们可以为每个子网络都单独计算指标，但目前就算是限定在 MobileNet V3 作为超网，可选的子网络也是有 10**5 的量级的，因此文中使用了进化算法（envolution_finder.py）来搜索子网络，流程上来说可以分为以下四步
首先随机选择 N 个子网络作为初始种群
选择初始化种群的一部分子网络进行再进行采样
选择初始化种群的另外一部分子网络之间进行交叉采样获得杂交后的子网络
循环一次 2~3 的过程，为种群的一次迭代，这里循环了 500 个迭代，然后返回精度最高且满足耗时要求的网络
对应的实现伪代码则如下所示：
# CAUTION: 伪代码，删减了一些边界判断和变量声明，完整的请查阅源码
def random_sample_net(self, net):
# 随机采样超网，获得子网
subnet = copy(net)
for b in subnet:
b.active_kernel = random.uniform([3, 5, 7])
b.active_out_channel = random.uniform([32, 64, 128, 256])
b.active_depth = random.uniform([3, 5, 7])
eff = self.efficiency_predictor.predict_efficiency(subnet)
return subnet, eff
def crossover_sample(self, net1, net2):
# 先 copy net1
new_sample = copy(net1)
for key in new_sample.keys():
if not isinstance(new_sample[key], list):
continue
# 然后对应的网络参数，在 net1 和 net2 中选择一个，模拟交叉的过程
for i in range(len(new_sample[key])):
new_sample[key][i] = random.choice([sample1[key][i], net2[key][i]])
efficiency = self.efficiency_predictor.predict_efficiency(new_sample)
return new_sample, efficiency
def run_evolution_search(self):
# 这里先获取初始种群和对应的耗时
for _ in range(population_size):
sample, efficiency = self.random_sample(super_net)
child_pool.append(sample)
efficiency_pool.append(efficiency)
# 获得种群个体的精度预测值
accs = self.accuracy_predictor.predict_accuracy(child_pool)
for i in range(population_size):
population.append((accs[i].item(), child_pool[i], efficiency_pool[i]))
# 完成种群初始化后，开始执行进化迭代
for iter in tqdm(range(max_time_budget)):
# 对所有子网络的精度进行倒序排序
# 获取 top-k 精度的子网络作为进化迭代的种群
parents = sorted(population, key=lambda x: x[0])[::-1][:parents_size]
acc = parents[0][0]
print('Iter: {} Acc: {}'.format(iter - 1, parents[0][0]))
# 看是否大于历史最大精度，如大于，记录下来
if acc > best_valids[-1]:
best_valids.append(acc)
best_info = parents[0]
else:
best_valids.append(best_valids[-1])
population = parents
child_pool = []
efficiency_pool = []
# 抽取一部分进行变异，也就是再采样该网络
for i in range(mutation_numbers):
par_sample = population[np.random.randint(parents_size)][1]
new_sample, efficiency = self.mutate_sample(par_sample)
child_pool.append(new_sample)
efficiency_pool.append(efficiency)
# 抽取另外一部分进行杂交，也就是让两个网络的参数进行交叉
for i in range(population_size - mutation_numbers):
par_sample1 = population[np.random.randint(parents_size)][1]
par_sample2 = population[np.random.randint(parents_size)][1]
new_sample, efficiency = self.crossover_sample(par_sample1, par_sample2)
child_pool.append(new_sample)
efficiency_pool.append(efficiency)
accs = self.accuracy_predictor.predict_accuracy(child_pool)
for i in range(population_size):
population.append((accs[i].item(), child_pool[i], efficiency_pool[i]))
return best_valids, best_info
实际应用
下图是我在实际应用中，对 GAN 网络做模型压缩，对全部子网进行评估后，再根据精度排序后的结果，横轴为子网络的计算量（GMACs），纵轴为对应模型的评估指标（FID, 越小越好）
如图红框所示的取值范围，在模型压缩到一个临界点的时候，指标会出现骤然的下跌，说明，在当前这个任务和数据集下，该模型的计算量的临界值我们可以知道个大概，然后再根据临界值所对应的模型在进行蒸馏来 fine-tune，实际实验下来，结果也还是可以的。
总结
个人认为 OFA 算法其实本质上是一个剪枝算法，只不过与剪枝相比，在评估子网络性能这个环节做到了自动化，从而大大的减少算法工程师的工作量
从 OFA 最终实验的指标上来看，搜索的方法也达到了 One Shot NAS 方法里的 SOTA，对于想要尝试 NAS 的朋友来说 OFA 是不错的选择，算法不算复杂，效果也不错
参考
Once-for-All: Train One Network and Specialize it for Efficient Deployment
GAN Compression: Efficient Architectures for Interactive Conditional GANs
1007
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
