深度学习模型轻量化(下)

深度学习模型轻量化(下)

2.4 蒸馏

2.4.1 蒸馏流程

蒸馏本质是student对teacher的拟合,从teacher中汲取养分,学到知识,不仅仅可以用到模型压缩和加速中。蒸馏常见流程如下图所示

 

 

1.  老师和学生可以是不同的网络结构,比如BERT蒸馏到BiLSTM网络。但一般相似网络结构,蒸馏效果会更好。

2.  总体loss为 soft_label_loss + hard_label_loss。soft_label_loss可以用KL散度或MSE拟合

3.  soft label为teacher模型的要拟合的对象。可以是模型预测输出,也可以是embeddings, 或者hidden layer和attention分布。

针对软标签的定义,蒸馏的方案也是百花齐放,下面分享两篇个人认为非常经典的文章。

2.4.2 distillBERT

DistillBERT: A distilled version of BERT: smaller, faster, cheaper and lighter

DistillBERT由大名鼎鼎的HuggingFace出品。主要创新点为:

1.  Teacher 12层,student 6层,每两层去掉一层。比如student第二层对应teacher第三层

2.  Loss= 5.0 * Lce+2.0 * Lmlm+1.0 * Lcos

  • Lce: soft_label 的KL散度
  • Lmlm: mask LM hard_label 的交叉熵
  • Lcos:hidden state 的余弦相似度

DistilBERT 比 BERT 快 60%,体积比 BERT 小 60%。在glue任务上,保留了 95% 以上的性能。在performance损失很小的情况下,带来了较大的模型压缩和加速效果。

 

 

 2.4.3 TinyBERT

TinyBERT: Distilling BERT for Natural Language Understanding

总体结构

重点来看下 TinyBERT,它是由华为出品,非常值得深入研究。TinyBERT 对 embedding 层,transformer层(包括hidden layer和attention)和 prediction 层均进行了拟合。如下图所示。

 

 

 TinyBERT 蒸馏过程

 其中Embeddings采用MSE, Prediction采用KL散度, Transformer层的hidden layer和attention,均采用MSE。loss如下

 

 

 其中m为层数。

效果分析

 

 

 表2: glue任务上的performance。在glue任务上,可达到bert-base的96%,几乎无损失。表3: tinyBERT模型大小和推理速度。缩小7.5倍,加速9.4倍。压缩和加速效果十分明显。

消融分析

 

 

 表6:分析embedding、prediction、attention、hidden layer软标签作用,其中attention和hidden layer作用最大。这个也很好理解,transformer层本来就是整个BERT中最关键的部分。

 

 

 表7:分析老师学生不同层对应方法的效果,uniform为隔层对应,top为全部对应老师顶部几层,bottom为全部对应老师底部几层。Uniform效果明显好很多。这个也很好理解,浅层可以捕捉低阶特征,深层可以捕捉高阶特征。全是低阶或者高阶显然不合适,我们要尽量荤素搭配。

3 框架层加速

3.1 手机端AI能力

目前移动端AI框架也比较多,包括谷歌的tf-lite,腾讯的NCNN,阿里的MNN,百度的PaddleLite, 小米的MACE等。他们都不同程度的进行了模型压缩和加速的支持。特别是端上推理的加速。这个可以参考“手机端AI性能排名“。

3.2 端侧AI框架加速优化方法

个人总结的主要方法如下,可能有遗漏哈,各位看官请轻拍:

1.  基于基本的C++编译器优化。

a.  打开编译器的优化选项,选择O2等加速选项。

b.  小函数内联,概率大分支优先,避免除法,查表空间换时间,函数参数不超过4个等。

2.  利用C,而不是C++,C++有不少冗余的东西。

3.  缓存优化

a.  小块内存反复使用,提升cache命中率,尽量减少内存申请。比如上一层计算完后,接着用作下一层计算。

b.  连续访问,内存连续访问有利于一次同时取数,相近位置cache命中概率更高。比如纵向访问数组时,可以考虑转置后变为横向访问。

c.  对齐访问,比如224*224的尺寸,补齐为256*224,从而提高缓存命中率。

d.  缓存预取,CPU计算的时候,preload后面的数据到cache中。

4.  多线程。

a.  为循环分配线程。

b.  动态调度,某个子循环过慢的时候,调度一部分循环到其他线程中。

5.  稀疏化

a.  稀疏索引和存储方案,采用eigen的sparseMatrix方案。

6.  内存复用和提前申请

a.  扫描整个网络,计算每层网络内存复用的情况下,最低的内存消耗。推理刚开始的时候就提前申请好。避免推理过程中反复申请和释放内存,避免推理过程中因为内存不足而失败,复用提升内存访问效率和cache命中率。

7.  ARM NEON指令的使用,和ARM的深度融合。NEON可以单指令多取值(SIMD),感兴趣可针对学习,这一块水也很深。

8.  手工汇编,毕竟机器编译出来的代码还是有不少冗余的。可以针对运行频次特别高的代码进行手工汇编优化。当然如果你汇编功底惊天地泣鬼神的强,也可以全方位手工汇编。

9.  算子支持:比如支持GPU加速,支持定点化等。有时候需要重新开发端侧的算子。

4 硬件层加速

硬件层加速比较硬核,小编就连半瓢水都达不到了,为了保证整个方案的全面性,还是硬着头皮东施效颦下。目前AI芯片厂家也是百花齐放,谁都想插一脚,不少互联网公司也来赶集,如下图所示。

 

 

 AI 芯片目前三种方案。GPU目前被英伟达和AMD牢牢把控。ASIC目前最火,TPU、NPU等属于ASIC范畴。

 

 

 5 总结

本文对深度学习模型压缩和加速的几类常用的方法进行了介绍。

参考文献

1.  ALBERT: A Lite BERT for Self-supervised Learning of Language Representations

2.  MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications

3.  Are Sixteen Heads Really Better than One?

4.  DistillBERT: A distilled version of BERT: smaller, faster, cheaper and lighter

5.  TinyBERT: Distilling BERT for Natural Language Understanding

6.  手机端AI性能排名

 

posted @ 2020-05-16 06:52  吴建明wujianming  阅读(1779)  评论(0编辑  收藏  举报