https://blog.tensorflow.org/2021/10/Collaborative-Optimizations.html
2021年10月8日
来自Arm的Mohamed Nour Abouelseoud 和Elena Zhelezina的客座贴
这篇博文介绍了用于边缘设备机器学习模型优化的协作技术，该技术由Arm提出并贡献给TensorFlow模型优化工具包（https://www.tensorflow.org/model_optimization），可从v0.7.0版（https://github.com/tensorflow/model-optimization/releases/tag/v0.7.0）开始使用。
协作优化管道的主要思想是在TensorFlow模型优化工具包中一个接一个地应用不同的优化技术，同时保持部署所需的压缩和准确性之间的平衡。这导致模型尺寸显著减小，并且可以在给定框架和特定于硬件的支持（如Arm Ethos-N和Ethos-U NPU提供的支持）的情况下提高推理速度。
这项工作是工具包路线图的一部分，用于支持开发更小更快的ML模型。您可以查看以前关于培训后量化、量化感知培训、稀疏性和集群的帖子，了解工具包的更多背景信息及其功能。
什么是协作优化？为什么？
协作优化背后的动机与模型优化工具包（TFMOT）背后的动机大体相同，后者旨在实现模型调节和压缩，以改进对边缘设备的部署。边缘计算和面向端点的人工智能的发展对此类工具和技术提出了更高的要求。Collaborative Optimization API将所有可用的模型优化技术堆叠起来，以利用其累积效应，实现最佳的模型压缩，同时保持所需的精度。
考虑到以下优化技术，部署的各种组合是可能的：
- 权重修剪（https://blog.tensorflow.org/2019/05/tf-model-optimization-toolkit-pruning-API.html）
- 权重聚类（https://blog.tensorflow.org/2020/08/tensorflow-model-optimization-toolkit-weight-clustering-api.html）
量化
训练后量化（[https://blog.tensorflow.org/2019/06/tensorflow-integer-quantization.html](https://blog.tensorflow.org/2019/06/tensorflow-integer-quantization.html)） 量化感知训练（QAT）（[https://blog.tensorflow.org/2020/04/quantization-aware-training-with-tensorflow-model-optimization-toolkit.html](https://blog.tensorflow.org/2020/04/quantization-aware-training-with-tensorflow-model-optimization-toolkit.html)）
换句话说，在部署之前，可以应用修剪和聚类中的一个或两个，然后是训练后量化或QAT。
结合这些技术的挑战是API不考虑先前的，每个优化和微调过程不保留前面技术的结果。这破坏了同时应用它们的整体效益；例如，聚类不能保持剪枝过程引入的稀疏性，QAT的微调过程失去了剪枝和聚类的好处。为了克服这些问题，我们引入了以下协作优化技术：
- 稀疏保持聚类：确保零聚类的聚类API，保持模型的稀疏性。https://www.tensorflow.org/model_optimization/guide/combine/sparse_clustering_example
- 稀疏保持量化感知训练（PQAT）：保持模型稀疏性的QAT训练API。https://www.tensorflow.org/model_optimization/guide/combine/pqat_example
- 保留聚类的量化感知训练（CQAT）：进行重新聚类并保留相同数量质心的QAT训练API。https://www.tensorflow.org/model_optimization/guide/combine/cqat_example
- 稀疏和保持簇的量化感知训练（PCQAT）：QAT训练API，用于保持使用稀疏保持簇训练的模型的稀疏性和簇数。https://www.tensorflow.org/model_optimization/guide/combine/pcqat_example
综合考虑，连同训练后量化（而非QAT）选项，它们提供了几个部署路径，如以下部署树中所示，其中叶节点是部署就绪模型，这意味着它们完全量化并采用TFLite格式。绿色填充表示需要重新培训/微调的步骤，红色虚线边框突出显示协作优化步骤。用于在给定节点获取模型的技术在相应的标签中指明。
上图中省略了直接、仅量化（训练后或QAT）部署路径。
其思想是在上述部署树的第三层实现完全优化的模型；然而，任何其他级别的优化都可以证明是令人满意的，并达到所需的推理延迟、压缩和准确性目标，在这种情况下，不需要进一步优化。建议的培训过程是反复检查适用于目标部署场景的部署树级别，查看模型是否满足优化要求，如果不满足，则：如果需要，使用相应的协同优化技术进一步压缩模型并重复，直到模型完全优化（修剪、聚集和量化）。
为了进一步提高与协同优化相关的内存使用率和推理速度，需要专门的运行时或编译器软件以及专用的机器学习硬件。示例包括用于Ethos-N处理器的Arm ML Ethos-N驱动程序堆栈和用于Ethos-U处理器的Ethos-U Vela编译器。目前，这两个示例都需要首先量化并将优化的Keras模型转换为TensorFlow Lite。
下图显示了通过整个协同优化管道的样本权重内核的密度图。
结果是一个量化部署模型，根据训练时指定的目标稀疏度，唯一值的数量减少，稀疏权重的数量显著增加。这将带来显著的模型压缩优势，并显著减少专用硬件上的推断延迟。
压缩和精度结果
下表显示了在流行模型上运行多个实验的结果，展示了应用这些技术带来的压缩优势与精度损失。可以应用更积极的优化，但以牺牲准确性为代价。虽然下表包括TensorFlow Lite模型的测量值，但其他序列化格式也有类似的优点。
稀疏保持量化感知训练（PQAT）
注：DS-CNN-L是为边缘设备设计的关键字定位模型。更多信息可以在Arm的ML示例库中找到。
保持聚类的量化感知训练（CQAT）
稀疏和聚类保持量化感知训练（PCQAT）
应用PCQAT
要应用PCQAT，您需要首先使用剪枝API来剪枝模型，然后使用保持稀疏性的集群API将其与集群连接起来。之后，QATAPI与定制的协同优化量化方案一起使用。下面是一个例子。
import tensorflow_model_optimization as tfmot
model = build_your_model()
# prune model
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(model, ...)
model_for_pruning.fit(...)
# pruning wrappers must be stripped before clustering
stripped_pruned_model = tfmot.sparsity.keras.strip_pruning(pruned_model)修剪模型后，按如下所示进行聚类和拟合
# Sparsity preserving clustering
from tensorflow_model_optimization.python.core.clustering.keras.experimental import (cluster)
# Specify your clustering parameters along
# with the `preserve_sparsity` flag
clustering_params = {
...,
'preserve_sparsity': True
}
# Cluster and fine-tune as usual
cluster_weights = cluster.cluster_weights
sparsity_clustered_model = cluster_weights(stripped_pruned_model_copy, **clustering_params)
sparsity_clustered_model.compile(...)
sparsity_clustered_model.fit(...)
# Strip clustering wrappers before the PCQAT step
stripped_sparsity_clustered_model = tfmot.clustering.keras.strip_clustering(sparsity_clustered_model)然后应用PCQAT。
pcqat_annotate_model = quantize.quantize_annotate_model(stripped_sparsity_clustered_model )
pcqat_model = quantize.quantize_apply(quant_aware_annotate_model,scheme=default_8bit_cluster_preserve_quantize_scheme.Default8BitClusterPreserveQuantizeScheme(preserve_sparsity=True))
pcqat_model.compile(...)
pcqat_model.fit(...)上面的示例显示了实现完全优化的PCQAT模型的训练过程，对于其他技术，请参考CQAT、PQAT和稀疏保持集群示例。请注意，用于PCQAT的API与CQAT的API相同，唯一的区别是使用preserve_sparsity标志以确保在训练期间保留零集群。PQAT API的用法类似，但使用不同的、保持稀疏性的量化方案。
致谢
本文介绍的特性和结果是许多人的工作成果，包括Arm ML工具团队和我们在谷歌TensorFlow模型优化工具包团队中的合作者。
来自Arm- Anton Kachatkou, Aron Virginas-Tar, Ruomei Yan, Saoirse Stewart, Peng Sun, Elena Zhelezina, Gergely Nagy, Les Bell, Matteo Martincigh, Benjamin Klimczak, Thibaut Goetghebuer-Planchon, Tamás Nyíri, Johan Gras。
来自谷歌- David Rim, Frederic Rechtenstein, Alan Chiao, Pulkit Bhuwalka
