一劳永逸—MIT韩松团队开源神经网络的高效部署

2023-12-19 12:50

本文主要是介绍一劳永逸—MIT韩松团队开源神经网络的高效部署,希望对大家解决编程问题提供一定的参考价值,需要的开发者们随着小编来一起学习吧!

本文首发于微信公众号 CVHub,不得以任何形式转载到其它平台,仅供学习交流,违者必究!

引言

本文为大家介绍一篇神经网络压缩方面相关的经典论文,由MIT 韩松团队发表于ICLR 2020的论文《Once-for-All: Train One Network and Specialize it for Efficient Deployment》.

本文是通过一种 Once for All 的网络,可高效生成 10^19 独立工作的子网络以便适配到不同的硬件平台,包括服务器端各种不同的GPU模型(Nvidia 1080Ti, 2080Ti, V100)、CPU(Inter Xeon 系列)以及移动端各种不同的边缘设备等,实际工程应用价值很大,非常值得一读并实践。

一、背景

为了能够在资源受限型的设备上进行更加高效的推理,以往的方法主要有两种,手工设计网络诸如SqueezeNet,MoblieNetV1-V2和ShuffleNets等;或结合网络架构搜索技术NAS,如MobilieNetV3和NASNet等。其中,手工设计的特征不仅费时费力,并且不具备良好的通用性。另一方面,NAS技术则是结合相关搜索策略在给定的搜索空间内寻找出能尽量拟合当前数据最佳模型,毫无疑问这是非常消耗计算资源的,因为通常来说它需要为每一种设计组合从头训练一遍网络,直至设定好的终止条件。此外,在给定新的推理硬件平台的情况下,这些方法也需要重复架构搜索过程并重新训练模型,从而导致总的开发成本过高。

本文从实际应用角度出发,提出了一种One-For-All(OFA)的通用神经网络。通过解耦训练和搜索来支持不同的架构设置,可以灵活地支持不同的深度宽度卷积核分辨率四个维度,而无需重新训练,利用所提出的渐进式搜索训练(PST)策略在同等的资源消耗下显著的提高了网络模型的性能。

一般的NAS算法通常包含几个步骤:

  • 定义网络的搜索空间(最大层数,每个节点可选择的操作);
  • 设计搜索策略(如何探索空间,目前主流的有强化学习,进化算法,梯度优化);
  • 对每个子网络进行评估。

One-Shot NAS,简称OSN,与一般NAS最大的不同之处在于其是训练一个超网,然后在超网上做减法(类似于剪枝),而之前的NAS更像是从小网络开始做加法。

二、方法

上图是本文的整体预览效果图,下面首先给出OFA的优化目标:

这里,W为大型网络,arch为特定的子结构,C为搜索策略。其核心思想便是利用方法CW中挑选出一个子网络进行训练,使其精度不断逼近arch

2.1 OFA网络

  • Units:单元数,即模块,共5个;
  • Elastic depth:网络深度,搜索空间为{2, 3, 4};
  • Elastic width:网络宽度,搜索空间为{3, 4, 6};
  • Elastic kernel-size:卷积核大小,搜索空间为{3, 5, 7};
  • Arbitrary resolution:图像分辨率,搜索空间为[128, 224],间隔为4,共25个输入分辨率;

因此,候选子网总数为:

  • 所有子网共享超网的部分权重W,因此仅需7.7M的参数量去存储这些子网;
  • 这里的超网可以是vanilla ResnetMobileNet等,限定其最大搜索空间。

2.2 PST策略

动机 如何更加高效的训练多个子网?

方法

  • 暴力法:更新全部子网的梯度(计算代价高)
  • 部分法:更新部分子网的梯度(收敛较困难)
  • 渐进式衰减法:由大到小逐渐逼近(综合性能好)
# 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()

  1. 首先,先训练整个超网,取设定好的内核数、通道数和网络深度;
  2. 在完成步骤1后,保持depthwidth不变,依次改变K的值进行训练直至收敛;
  3. 步骤2,在保持任意两个变量取最大值的情况,依次改变剩余变量的值微调模型。

注意

对于resolution维度的变化是体现在以上训练过程中每个batch上,即多尺度训练;

这里训练完步骤1之后会先采用知识蒸馏得到一个相对较小的模型作为超网。

2.2.1 Elastic Kernel Size

以搜索空间为{3, 5, 7}的卷积核为例,首先训练完7x7大小的卷积核后,中间区域的权重保持共享。这样会存在一个问题就是不同的卷积核可能代表不同的网络分布,如果强制保持一致的话会导致其性能下降。因此,这里采用一个变换矩阵来实现权重共享,实现原理也很简单,可以利用一个MLP来完成此项功能。此外,对于同个stage内,不同通道之间的变换矩阵是共享的;不同stage的话会重新定义一个变幻矩阵。所以所需存储的参数量变为25x25+3x3=706个额外参数来存储每一层的核变换矩阵。

2.2.2 Elastic Depth

如上图5右手边所示,网络深度也是可以动态调整,作者采用的是跳层来实现深度的动态变化。较深的网络可以直接把后几层跳过,例如,深度为4的网络可以只选取前3层,从而得到深度为3的网络。而取前三层的原因是因为权重是共享的,按经验来说,从前往后取的效果更好。

2.2.3 Elastic Width

如图6所示,每层卷积对应的输入和输出通道数也是可以动态调整而非固定的,不过这里会先进行通道重要性指标的排序后,再通过slices来实现,逐步选取前n个重要的通道进行训练逐步衰减。

需要注意的是这里通道重要性计算方法采用的是L1 norm,即数值越大代表对应的通道重要性成都越高。

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

需要注意的是,在实现的过程中需要对所有相关的op进行适配,以动态适应不同超参的输入。

2.3 评估策略

常规的评估方法为在测试集上对每个子网络的效果进行评估,毫无疑问这种方式是非常耗时的。因此作者设计了一个子网络精度评估器,这个评估器的输入是待评估子网的架构参数Oe-Hot Encoding输出是预测该子网精度数值,通过这种方式,只需要做一次推断便可以预测该子网的大致性能。

# accuracy predictor
accuracy_predictor = AccuracyPredictor(pretrained=True,device='cuda:0' if cuda_available else 'cpu'
)print(accuracy_predictor.model)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)
)

一般做模型压缩,轻量化架构和搜索,是为了能在某些设备上更快的运行起来,拿手机设备举个例子,因为存在很多的不同厂商,因此不同的算子在不同的设备上的计算速度存在差异,为此,文中提到了延迟查找表的概念,也就是说我们把我们要用到的算子在设备上进行一次基准测试,然后我们在预估模型在设备上计算时长的时候,就只需要查表就可以了,不用一次次的去部署然后做基准测试。

2.4 选取策略

获得了子网络的精度,以及子网络的运行速度,接下来就可以根据需求对子网络进行搜索和评估。作者使用了演化算法(Regularized Evolution for Image Classifier Architecture Search)来搜索子网络,流程上来说可以分为以下四步:

  1. 随机选择N个子网络作为初始种群;
  2. 选择初始化种群的一部分子网络进行采样;
  3. 选择初始化种群的另外一部分子网络之间进行交叉采样获得杂交后的子网络;
  4. st2st3迭代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, effdef crossover_sample(self, net1, net2):# 先 copy net1new_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, efficiencydef 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 = parentschild_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

三、实验

作者这里用单张V100共花费了1,200GPU hours大约是50天来完成整个训练过程。不过通常来说业务数据集并不会同ImageNet这么大,所以一般情况下几天应该就能搞定全部。

3.2 实验效果

可以看出,采用OFA网络的GPU hours相比传统的NAS方法少了非常多,极大节省了成本开支。其次,结合PS策略可以获得更具性价比的精度,在不同的架构设置下均能取得3个点左右的性能提升。

四、总结

4.1 OFA核心步骤

  • 训练
    • 正常训练一个最大的超网(动态分辨率、最大卷积核、最大宽度、最大深度)。OFA提供了官方训练的模型权重:ofa_D4_E6_K7,可直接使用。
  • 微调
    • 对超网ofa_D4_E6_K7进行渐进式收缩(PS)训练,尽量使得所有子网也能有很好的表现。
  • 搜索
    • 在超网中,使用进化算法和训练好的predictor在给定约束下(如latency FLOPs)进行搜索,得到满足约束的子网,结束算法。

4.2 PS策略实施步骤

train_ofa_net.py就是用于PS训练的脚本。train_ofa_net.py是分步骤分阶段进行的,每一次执行train_ofa_net.py都会将向采样空间添加新的元素进行训练,得到不同阶段的模型权重。训练过程如下:

将{3,5}的卷积核大小加入采样空间进行训练,得到权重ofa_D4_E6_K357

将{3}的深度大小加入采样空间进行训练,得到权重ofa_D34_E6_K357

将{2}的深度大小加入采样空间进行训练,得到权重ofa_D234_E6_K357

将{4}的宽度大小加入采样空间进行训练,得到权重ofa_D234_E46_K357

将{3}的宽度大小加入采样空间进行训练,得到权重ofa_D234_E346_K357

ofa_D234_E346_K357即为最终PS训练完成的模型权重。

4.3 子网的获取步骤

假设我们有一个子网的配置,试图从超网中获得一个独立的子网,需要经过以下三步:

设置子网
  1. 先对超网设置设置子网:
set_active_subnet(ks=net_config['ks'], d=net_config['d'], e=net_config['e'])
  1. 获得子网结构,并从超网中继承权重get_active_subnet().
设置均值和方差

对子网中的BN层的mean和var进行重新估计:

calib_bn(subnet, path, net_config['r'][0], batch_size)
推理验证
validate(subnet, path, net_config['r'][0], data_loader, batch_size, device)

步骤2中在超网的采样训练中,每一个子网进行训练时都会影响bn层中的均值mean方差var。为了达到更好的性能,子网需要对BN层均值mean方差var根据该子网的结构进行重新统计,calib_bn()函数就是用于重新统计meanvar的函数,其构造了data_loader传入set_running_statistics函数进行统计,代码如下:

def set_running_statistics(model, data_loader, distributed=False):#用于记录每个batch下每个bn的mean和varbn_mean = {}bn_var = {}# copy一个模型forward_model = copy.deepcopy(model)for name, m in forward_model.named_modules():if isinstance(m, nn.BatchNorm2d):if distributed:bn_mean[name] = DistributedTensor(name + '#mean')bn_var[name] = DistributedTensor(name + '#var')else:bn_mean[name] = AverageMeter()bn_var[name] = AverageMeter()def new_forward(bn, mean_est, var_est):def lambda_forward(x):#统计每个batch的mean和varbatch_mean = x.mean(0, keepdim=True).mean(2, keepdim=True).mean(3, keepdim=True)  # 1, C, 1, 1batch_var = (x - batch_mean) * (x - batch_mean)batch_var = batch_var.mean(0, keepdim=True).mean(2, keepdim=True).mean(3, keepdim=True)batch_mean = torch.squeeze(batch_mean)batch_var = torch.squeeze(batch_var)#更新记录mean_est.update(batch_mean.data, x.size(0))var_est.update(batch_var.data, x.size(0))# bn forward using calculated mean & var_feature_dim = batch_mean.size(0)return F.batch_norm(x, batch_mean, batch_var, bn.weight[:_feature_dim],bn.bias[:_feature_dim], False,0.0, bn.eps,)return lambda_forwardm.forward = new_forward(m, bn_mean[name], bn_var[name])if len(bn_mean) == 0:# skip if there is no batch normalization layers in the networkreturnwith torch.no_grad():DynamicBatchNorm2d.SET_RUNNING_STATISTICS = Truefor images, labels in data_loader:images = images.to(get_net_device(forward_model))forward_model(images)#inference一次,记录mean和var到bn_mean和bn_var中DynamicBatchNorm2d.SET_RUNNING_STATISTICS = False#将forward_model inference得到的runing_mean和runing_var复制到model中for name, m in model.named_modules():if name in bn_mean and bn_mean[name].count > 0:feature_dim = bn_mean[name].avg.size(0)assert isinstance(m, nn.BatchNorm2d)m.running_mean.data[:feature_dim].copy_(bn_mean[name].avg)#copy meanm.running_var.data[:feature_dim].copy_(bn_var[name].avg)#copy var

4.2 结束语

  • OFA可认为是一种带剪枝算法的one-shot NAS,与常规的剪枝方法相比,它是从多个维度进行裁剪,而非单一维度;
  • 其采用的PST不仅更新子网的权重,也会更新超网的权重使两者尽量对齐保持同步。

总的来说,本文提出了一种OFA的新方法,它将模型训练从架构搜索中分离出来,从而在大量的硬件平台下实现高效的深度学习部署。与以前为每个部署场景设计和训练神经网络的方法不同,它支持不同的体系结构配置,包括弹性深度、宽度、内核大小和分辨率。与传统方法相比,它大大降低了训练成本(GPU时间、能源消耗和二氧化碳排放)。为了防止不同规模的子网络受到干扰,提出了一种渐进收缩算法,使大量的子网络达到与独立训练相同的精度水平。在各种硬件平台和效率约束上的实验证明了该方法的有效性。

五、代码梳理

- ofa- imagenet_classification # 超网训练相关代码- elastic_nn   # 子网代码- modules  # 弹性模块- dynamic_layers.py # 动态层- dynamic_op.py     # 动态op- networks # 弹性网络定义- ofa_mvb3.py       # MobilenetV3- ofa_proxyless.py  # ProxylessNASNets- ofa_resnets.py    # Resnet- training # 渐进式训练- progress_shrinking.py  # PST核心代码- networks     # 常规网络定义- mobilenet_v3.py- proxyless_nets.py- resnets.py- run_manager # 训练核心代码- nas # 搜索部分- accuracy_predictor          # 自定义的评估代码,详情见2.3- acc_dataset.py- acc_predictior.py- arch_encoder.py- efficiency_predictor- latency_looup_tabel.py  # 性能查找表- search_algorithm- evolution.py            # 进化算法,用于选取合适的子网
- train_ofa_net.py # 开始入口

CVHub专注于为大家提供高质量的精选文章和实用开源工具,后期考虑开源一哥一键算法模型API工具,支持限流、队列、轮询和频控能力,同时提供一站式实用算法仓库,敬请,更多精彩内容请关注官方WeChat公众号CVHub,期待关注,感谢大家一路的支持和陪伴,谢谢!

写在最后

如果您也对人工智能和计算机视觉全栈领域感兴趣,强烈推荐您关注有料、有趣、有爱的公众号『CVHub』,每日为大家带来精品原创、多领域、有深度的前沿科技论文解读及工业成熟解决方案!欢迎添加小编微信号:cv_huber,一起探讨更多有趣的话题!

这篇关于一劳永逸—MIT韩松团队开源神经网络的高效部署的文章就介绍到这儿,希望我们推荐的文章对编程师们有所帮助!



http://www.chinasem.cn/article/512234

相关文章

闲置电脑也能活出第二春?鲁大师AiNAS让你动动手指就能轻松部署

对于大多数人而言,在这个“数据爆炸”的时代或多或少都遇到过存储告急的情况,这使得“存储焦虑”不再是个别现象,而将会是随着软件的不断臃肿而越来越普遍的情况。从不少手机厂商都开始将存储上限提升至1TB可以见得,我们似乎正处在互联网信息飞速增长的阶段,对于存储的需求也将会不断扩大。对于苹果用户而言,这一问题愈发严峻,毕竟512GB和1TB版本的iPhone可不是人人都消费得起的,因此成熟的外置存储方案开

高效+灵活,万博智云全球发布AWS无代理跨云容灾方案!

摘要 近日,万博智云推出了基于AWS的无代理跨云容灾解决方案,并与拉丁美洲,中东,亚洲的合作伙伴面向全球开展了联合发布。这一方案以AWS应用环境为基础,将HyperBDR平台的高效、灵活和成本效益优势与无代理功能相结合,为全球企业带来实现了更便捷、经济的数据保护。 一、全球联合发布 9月2日,万博智云CEO Michael Wong在线上平台发布AWS无代理跨云容灾解决方案的阐述视频,介绍了

嵌入式QT开发:构建高效智能的嵌入式系统

摘要: 本文深入探讨了嵌入式 QT 相关的各个方面。从 QT 框架的基础架构和核心概念出发,详细阐述了其在嵌入式环境中的优势与特点。文中分析了嵌入式 QT 的开发环境搭建过程,包括交叉编译工具链的配置等关键步骤。进一步探讨了嵌入式 QT 的界面设计与开发,涵盖了从基本控件的使用到复杂界面布局的构建。同时也深入研究了信号与槽机制在嵌入式系统中的应用,以及嵌入式 QT 与硬件设备的交互,包括输入输出设

阿里开源语音识别SenseVoiceWindows环境部署

SenseVoice介绍 SenseVoice 专注于高精度多语言语音识别、情感辨识和音频事件检测多语言识别: 采用超过 40 万小时数据训练,支持超过 50 种语言,识别效果上优于 Whisper 模型。富文本识别:具备优秀的情感识别,能够在测试数据上达到和超过目前最佳情感识别模型的效果。支持声音事件检测能力,支持音乐、掌声、笑声、哭声、咳嗽、喷嚏等多种常见人机交互事件进行检测。高效推

高效录音转文字:2024年四大工具精选!

在快节奏的工作生活中,能够快速将录音转换成文字是一项非常实用的能力。特别是在需要记录会议纪要、讲座内容或者是采访素材的时候,一款优秀的在线录音转文字工具能派上大用场。以下推荐几个好用的录音转文字工具! 365在线转文字 直达链接:https://www.pdf365.cn/ 365在线转文字是一款提供在线录音转文字服务的工具,它以其高效、便捷的特点受到用户的青睐。用户无需下载安装任何软件,只

图神经网络模型介绍(1)

我们将图神经网络分为基于谱域的模型和基于空域的模型,并按照发展顺序详解每个类别中的重要模型。 1.1基于谱域的图神经网络         谱域上的图卷积在图学习迈向深度学习的发展历程中起到了关键的作用。本节主要介绍三个具有代表性的谱域图神经网络:谱图卷积网络、切比雪夫网络和图卷积网络。 (1)谱图卷积网络 卷积定理:函数卷积的傅里叶变换是函数傅里叶变换的乘积,即F{f*g}

金融业开源技术 术语

金融业开源技术  术语 1  范围 本文件界定了金融业开源技术的常用术语。 本文件适用于金融业中涉及开源技术的相关标准及规范性文件制定和信息沟通等活动。

安全管理体系化的智慧油站开源了。

AI视频监控平台简介 AI视频监控平台是一款功能强大且简单易用的实时算法视频监控系统。它的愿景是最底层打通各大芯片厂商相互间的壁垒,省去繁琐重复的适配流程,实现芯片、算法、应用的全流程组合,从而大大减少企业级应用约95%的开发成本。用户只需在界面上进行简单的操作,就可以实现全视频的接入及布控。摄像头管理模块用于多种终端设备、智能设备的接入及管理。平台支持包括摄像头等终端感知设备接入,为整个平台提

【C++高阶】C++类型转换全攻略:深入理解并高效应用

📝个人主页🌹:Eternity._ ⏩收录专栏⏪:C++ “ 登神长阶 ” 🤡往期回顾🤡:C++ 智能指针 🌹🌹期待您的关注 🌹🌹 ❀C++的类型转换 📒1. C语言中的类型转换📚2. C++强制类型转换⛰️static_cast🌞reinterpret_cast⭐const_cast🍁dynamic_cast 📜3. C++强制类型转换的原因📝

K8S(Kubernetes)开源的容器编排平台安装步骤详解

K8S(Kubernetes)是一个开源的容器编排平台,用于自动化部署、扩展和管理容器化应用程序。以下是K8S容器编排平台的安装步骤、使用方式及特点的概述: 安装步骤: 安装Docker:K8S需要基于Docker来运行容器化应用程序。首先要在所有节点上安装Docker引擎。 安装Kubernetes Master:在集群中选择一台主机作为Master节点,安装K8S的控制平面组件,如AP