目前在深度学习领域,一方面需要追求更高的性能,采用强大、复杂的模型网络和实验方法;另一方面又需要关注如何将算法更稳定、高效地在硬件平台上落地。复杂的模型固然展现更好的性能,但过高的存储空间需求和计算资源消耗,是影响在各硬件平台上的落地的重要原因之一。尤其在NLP领域,以BERT、GPT为代表的预训练模型规模越来越大。
针对部署难题,可以使用剪枝、量化、蒸馏等传统的模型压缩技术,但是这些模型压缩技术整体的流程较长,直接串行结合的效果不佳。而OFA(Once For All)[1]技术巧妙地结合了剪枝、蒸馏和结构搜索策略,不仅提升了压缩效果,还简化了压缩流程。
百度飞桨模型压缩工具PaddleSlim新增支持OFA这一实用功能,并在BERT和GAN模型上做了验证:对于BERT模型实现了近2倍的加速;而对GAN模型则实现了33倍的体积压缩。除此以外,PaddleSlim还为用户提供了简单易用且低侵入的轻量级接口,用户无需修改模型训练代码,即可完成OFA压缩。
什么是OFA?
OFA(Once For All)是一种结合One-Shot NAS和蒸馏的压缩方案。其优势是可以基于已有的预训练模型进行压缩,不需要重新训练新的预训练模型;同时只需要训练一个超网络就可以从中选择满足不同延时要求的子模型。
OFA的搜索空间和One-Shot NAS不同,之前的搜索是在不同OP之间进行选择,比如把降采样中的卷积替换成Pooling,或者把3x3卷积替换成5x5。OFA中的搜索空间不是OP的替换,而是对当前OP进行属性变换,比如卷积核大小的变化和参数通道数的变化。
说明:
卷积核大小变化指的是对于网络中已有的5x5的卷积,在实际训练过程中随机把它修改成为3x3,而它与5x5的卷积是权重共享的。
参数通道数变化指的是在训练过程中随机把参数通道数修改成搜索空间中设置的任意一个通道数。
OFA的搜索空间还可以包含对模型深度的搜索。以BERT-base模型为例,一共有12个transformer encoder block。如果深度相关的搜索空间设置为[8, 12]两个选项,则在训练过程中会随机选择当前训练的子模型是包含8个还是12个transformer encoder block。如果训练的是包含8个transformer encoder block组成的子网络,则根据一些规则去掉原始12个block中的4个block。
通过上面训练,每次都会随机训练超网络中任意一个子网络,该网络仅包含超网络部分参数。
OFA中的蒸馏使用的是自蒸馏的方式,教师网络选择和超网络中最大子模型相同的模型,