越来越多的模型部署在手机端,而端侧可用的计算能力有限、电池容量有限,为了高效的部署推理AI模型,产业界已经有很多算法框架,高通也推出了针对端侧的优化工具库AIMET。
很多小伙伴对于AIMET表示兴趣,我也打算学习和使用这个工具,将在这里做一些学习笔记,欢迎留言。
- 简介Blog: AI Model Efficiency Toolkit
- Github开源项目: quic/aimet: AIMET
- AIMET文档:AI Model Efficiency Toolkit User Guide
- Model Zoo: GitHub - quic/aimet-model-zooh ttps://github.com/quic/aimet-model-zoo
本文包括
- AIMET简介
- 模型量化工具
- 模型压缩工具
- 可视化工具
- Model Zoo
一、AIMET 简介
AIMET全称AI Model Efficiency Toolkit,顾名思义,它是用来提升模型效率的工具。
(1)AIMET 架构
AIMET是一个软件工具包,能够压缩和量化已经训练好的模型,它是基于TensorFlow和Pytorch的,也就是说需要首先安装TF1.15(似乎还不支持TF2.0) 和Pytorch。
(2)AIMET使用流程:
AIMET 使用流程
- 一个训练好的TF或者Torch模型加载到AIMET里面;
- AIMET做模型的量化、压缩,生成一个新的模型;
- 如果新的模型满足速度和精度的预期,那么可以导出模型;
- 如果生成的模型精度损失严重,AIMET也支持fine-tune,来恢复量化压缩损失;
- 最后,TF可以保持成PB,Torch可以导出成ONNX,通过SNPE进行部署。
可以看出,AIMET主要专注于模型量化(Quantization)和模型压缩(Compression)。
二、模型量化
AIMET 量化步骤
- (1)(2) 在TF或者Torch上创建和训练模型;
- (3) 创建一个AIMET量化模型,如果量化模型精度满足要求,跳到第(5)步,不然进行第(4)步;
- (4) 用户可以进一步训练这个AIMET量化模型;
- (5) 保存模型weights和量化encoding;
- (6) 部署模型到定点加速硬件。
(1)AIMET 量化模拟器
Doc:AIMET Quantization Simulation
让模型跑在量化计算的硬件加速单元上需要将模型从浮点量化成定点,这一过程会引入计算精度的损失。AIMET提供了模拟这种量化损失的功能,让用户可以进一步重新训练模型(微调),从而恢复精度上的损失。
AIMET Simulation向前传播时对于量化损失的模拟方式和TF的量化训练方案基本一直。下图是向前推理时的量化损失引入方式,选定层的weights和activations先量化然后反量化。
AIMET 前传引入量化噪音
(2)后量化
Doc: AIMET Post-Training Quantization Techniques
AIMET 提供了三种后量化算法。
AdaRound
- CEL和BC的简单原理可以参考:SNPE中的后量化 ,详细的论文可以参考 https://arxiv.org/abs/1906.04721
- AdaRound 详细原理:https://arxiv.org/abs/2004.10568 , 有机会进一步研究。
(3)量化感知训练AIMET Quantization Aware Training
和其他框架下的QAT或者微调一样,向前传播的时候模拟量化带来的精度损失,然后反向传播采用浮点计算更新权重,这样可以让模型可以更好的部署到定点硬件上。
AIMET 量化simulator 提供了多种量化方案,可以在QAT的时候使用。
AIMET TensorFlow Quantization SIM API
AIMET PyTorch Quantization SIM API
三、模型压缩
Doc:
AIMET Visualization for Quantization
AIMET提供了模型压缩工具,可以用来减少模型的MAC和内存,同时保证一定的精度,AIMET支持各种压缩方案,如权重SVD、空间SVD和通道修剪。
使用流程
- AIMET允许用户将训练好的模型压缩到所需的压缩率,然后进一步微调。
- AIMET中的所有压缩方案都使用两步流程--1.选择压缩率;2.模型压缩。
具体的方案和步骤后续进一步学习研究。
四、可视化
Doc: AIMET Visualization
AIMET可视化工具可以帮助大家分析模型的量化和压缩,从而让问题和结果更加直观。
可视化工具是基于Bokeh开发。
AIMET 可视化工具
(1)权重范围
可视化模型的权重。例如。可以查看每一层的中各通道的数值范围,从而直观的看出是否可以用CLE优化。
权重可视化示例
(2)压缩灵敏度
直观地看出模型中任何层对压缩的反馈。
每层压缩率示例
五、Model Zoo
同时,他们提供了流行的神经网络在使用AIMET之后的模型,并给出了它们的浮点和量化的性能,结果表明,量化模型可以提供良好的精度,与浮点模型相媲美。
同时,每一个模型都给出了对应的评估脚本。量化模型如果是采用后量化(PTQ)评估脚本会在评估前执行PTQ,如果是使用QAT,则提供模型的checkpoint。
总结
到此为止,我们知道了AIMET是什么,可以做什么,至于应该怎么用,后面将进一步尝试安装和尝试AIMET的功能。
期待和大家进一步交流。
参考:
- Open sourcing the AI Model Efficiency Toolkit
- AI Model Efficiency Toolkit User Guide
- Here’s why quantization matters for AI
- GitHub - quic/aimet-model-zoo
- https://github.com/quic/aimet
作者:Wenhao