来自 Arm 的 Mohamed Nour Abouelseoud 和 Elena Zhelezina 的客座文章。
本博文介绍了针对边缘设备的机器学习模型优化协同技术，由 Arm 提出并贡献给 TensorFlow 模型优化工具包，从版本 v0.7.0 开始提供。
协同优化管道的核心思想是在保持部署所需压缩和精度平衡的情况下，将 TensorFlow 模型优化工具包 中的不同优化技术逐一应用。这将显著减少模型大小，并能提高推理速度，前提是框架和硬件支持，例如 Arm 提供的 Ethos-N 和 Ethos-U NPU。
这项工作是 工具包路线图 的一部分，旨在支持开发更小更快 ML 模型。您可以在关于 训练后量化、量化感知训练、稀疏性 和 聚类 的以前文章中了解有关工具包及其功能的更多背景信息。
协同优化背后的动机与模型优化工具包 (TFMOT) 的总体动机相同，即为改进边缘设备的部署而对模型进行条件化和压缩。边缘计算和 面向端点的 AI 的发展趋势，对这类工具和技术产生了巨大的需求。协同优化 API 将所有可用的模型优化技术堆叠在一起，以利用它们的累积效应，在保持所需精度的情况下实现最佳模型压缩。
考虑到以下优化技术，各种可能的部署组合:
换句话说，可以在部署之前应用剪枝和聚类中的一种或两种，然后进行训练后量化或 QAT。
组合这些技术的挑战在于，API 不考虑之前的技术，每个优化和微调过程都不会保留之前技术的成果。这会破坏同时应用它们带来的整体优势；例如，聚类不会保留剪枝过程引入的稀疏性，而 QAT 的微调过程会丢失剪枝和聚类的优势。为了克服这些问题，我们引入了以下协同优化技术:
将这些技术与使用训练后量化而不是 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 示例 存储库。
要应用 PCQAT，您需要首先使用修剪 API 来修剪模型，然后将其与使用稀疏性保持集群 API 的集群链接起来。之后，QAT API 与自定义协作优化量化方案一起使用。下面是一个示例。
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 工具团队和我们在 Google 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。
来自 Google - David Rim、Frederic Rechtenstein、Alan Chiao、Pulkit Bhuwalka
