-
Notifications
You must be signed in to change notification settings - Fork 0
Home
模型自动压缩器(Auto Model Compressor, AMC)是基于Pytorch开发的,用于压缩深度神经网络计算量(浮点型乘加计算次数,Floating-Point-Fused-Multiply-Adds)的框架(暂内部开源)。
AMC主要由四个模块组成:数学压缩器(Mathematical Compressors),训练器(Trainer),超参数调整器(Hyper-Parameters Tuners)及实用工具集(Utils)。如下图所示:

-
数学压缩器(Mathematical Compressors):是指使用数学方法对某个神经网络的卷积层进行低损压缩(减少乘加计算次数)的方法。现框架主要支持两大类型方法——通道剪枝(Channel Pruning)和低秩分解(Low-rank Decomposition)。通道剪枝是利用统计学方法分析卷积层的稀疏性,选择性地保留对卷积层计算结果影响最大的特征通道,从而减少卷积层的乘加计算次数。而低秩分解则是利用卷积层本质是矩阵运算的特点,利用矩阵分解,把一个卷积层分解为两个乘加计算次数更小的卷积层相连,从而达到减少计算量的目的。两种方法各有千秋,低秩分解能较高精度地压缩稀疏性大的卷积层,但增加了神经网络的深度,这对部署时的GPU设备并不友好。而通道剪枝能不改变原有神经网络结构的情况下,压缩神经网络的计算量,但精度却不如低秩分解高。另外,试验表明,不同应用类型的神经网络,对两种方法的敏感度也不同,需要针对性选择方法。
-
训练器(Trainers):是指对压缩后的神经网络进行再次训练(或微调,Finetune)的一般实现。训练器还实现了非常多有利于训练效果的损失函数和量化训练工具。
-
超参数调整器(Hyper-Parameters Tuners):是指用于调整网络压缩的压速率,再训练的学习率等参数的“代理者”,可通过强化学习方法(Deep Deterministic Policy Gradient, DDPG)或遗传算法(Genetic Algorithms, GA)自动化调整上述参数,解放人频繁调参的双手。
-
实用工具集(Utils):包含大量方便获取Pytorch神经网络图信息、计算量等有用信息的实现。包含解析Pytorch模型并生成部署框架Venus所需二进制文件的实现。包括一些基础类定义。
给定有一个已训练好的原始模型,使用Utils提供的工具把模型解析成特定特定类:class ReconstructedNetwork(nn.Module)。
-
Step 1:使用数学压缩器,对原始网络进行压缩,生成“压缩网络(Compressed Network)”。
-
Step 2:压缩网络提供了一个非常好的初值,使用训练器对压缩网络进行再训练,得到一个“再训练网络(Retrained Network)”
把再训练网络作为原始网络,可以循环Step 1和Step 2进行交替迭代。最终得到一个在计算量和效果平衡的最优网络。利用Utils提供的工具,直接转换该最优到.vnmodel二进制文件,用于三端的部署。
具体流程如下图:

Code:
如上图所示,初始化compressor的输入,origin_net(torch.nn.Module)是原始网络,trainloader(torch.utils.data.DataLoader)装载训练net的train数据,valloader(torch.utils.data.DataLoader)是测试net的test数据,trainset_ratio是压缩使用的训练数据的比例,compress_ratio是每一个卷积层压缩的计算量的比例,sampled_pixels_per_img是每层卷积层输入数据采样的像素数量,checkpointfolder是输出数据的目录,device指定原始网络和压缩网络所在的设备,accuracy_first是压缩过程中是否关注精度,accuracy_threshold代表关注精度的阈值。
对于低秩分解,nonlinear_case参数代表使用线性模式还是非线性模式。
对于通道剪枝,lars_alpha_init代表LARS解LASSO问题时的惩罚因子的初值。
调用compress()方法,直接压缩,压缩后的网络(torch.nn.Module)放到class.compressed_net中。
AMC的训练器(Trainer),本质上是固化了Pytorch神经网络的训练任务,通过自定义该固定训练任务中的某些步骤,来自定义某个网络的训练。一般来说,用户设计好Pytorch的神经网络(nn.Module)后,再自定义以下几个步骤,几乎可以实现所有情况的训练任务:
- 1)封装训练/测试用的数据与标签为torch.utils.data.DataLoader;
- 2)封装Loss-Function为criterion(torch.nn.Module);
- 3)如何从torch.utils.data.DataLoader的遍历枚举中获得criterion的predict输入。
- 4)定义训练优化器Optimizer(nn.optim)和学习率调度器scheduler(nn.optim.lr_scheduler)。
- [1] He Y, Zhang X, Sun J. Channel pruning for accelerating very deep neural networks[C]//Proceedings of the IEEE International Conference on Computer Vision. 2017: 1389-1397.
- [2] Zhang X, Zou J, He K, et al. Accelerating very deep convolutional networks for classification and detection[J]. IEEE transactions on pattern analysis and machine intelligence, 2015, 38(10): 1943-1955.
- [3] Lillicrap T P, Hunt J J, Pritzel A, et al. Continuous control with deep reinforcement learning (2015)[J]. arXiv preprint arXiv:1509.02971, 2016.
- [4] Man K F, Tang K S, Kwong S. Genetic algorithms: concepts and applications [in engineering design][J]. IEEE transactions on Industrial Electronics, 1996, 43(5): 519-534.
- [5] Hinton G, Vinyals O, Dean J. Distilling the knowledge in a neural network[J]. arXiv preprint arXiv:1503.02531, 2015.
- [6] Chen G, Choi W, Yu X, et al. Learning efficient object detection models with knowledge distillation[C]//Advances in Neural Information Processing Systems. 2017: 742-751.
- [7] Krizhevsky A, Hinton G. Learning multiple layers of features from tiny images[R]. Technical report, University of Toronto, 2009.