特邀博文 / Mohamed Nour Abouelseoud 和 Anton Kachatkou，来自 Arm
我们很高兴能将 Arm 提供的 Weight Clustering API 引入 TensorFlow 模型优化工具包。
权重聚类 (Weight Clustering) 技术用较少 唯一值 (Unique Value) 替换大量不同的参数值，缩减模型的存储和传输数据的大小。这一优势适用于所有部署。除了特定框架和硬件的支持，例如在 Arm Ethos-N 和 Ethos-U 机器学习处理器中的支持，权重聚类还可以改善内存占用量并提高推理速度。
这是工具包路线图的一部分，支持开发更小更快的 ML 模型。在有关 训练后整形量化、训练时量化 和 稀疏性 的往期文章中，有工具包背景及功能的更多相关信息。
Arm 和 TensorFlow 团队在这一领域长期合作，着力改善移动和 IoT 设备的部署。
什么是权重聚类？
从智能手机到农业传感器和医疗器械，深度学习应用在资源受限的环境中也越发普及。这种向资源受限环境的转移需要更小更高效的模型架构，也引起了对 剪枝 和 量化 和 稀疏性 等模型优化技术的日益重视。
权重聚类优化算法可以减少模型的存储量和网络传输量。下图简要说明了这一概念。
如上图所示。若模型的一个层包含一个 4x4 的权重矩阵（即上图中的“权重矩阵 (Weight Matrix)”表示）。每个权重使用一个 float32 值存储。保存模型时，将在磁盘上存储 16 个唯一的 float32 值。
权重聚类用相同的值替换层中的相似权重来缩减模型的大小。在模型的训练权重中通过聚类算法来找到这些值。用户可以指定聚类的数量（在本例中为 4）。上图“计算质心”中展示了此步骤，并将 4 个质心记录在“质心表（即图中的 Centroid）”中，以为每个质心提供对应的索引值 (0-3)。
接下来，将权重矩阵中的每个权重替换为对应的质心索引。此步骤见“分配索引 (Assign Indices)”。接下来便可通过由权重聚类算法进行修改后的索引矩阵（即图中的 “Pull indices”）来替代用于存储权重的原始矩阵。
上图所示案中将矩阵规模从 16 个不同的浮点数减少到 4 个浮点数质心和 16 个 2 位宽索引。矩阵规模越大，节省的成本也就越多。
注意，即使仍然存储了 16 个浮点数，它们现在也只有 4 个不同的值。常用压缩工具（如 Zip）现在可以利用数据中的冗余来实现更高的压缩率。
聚类的技术实现源于 Deep Compression: Compressing Deep Neural Networks with Pruning, Trained Quantization and Huffman Coding。论文阐述了关于梯度更新和权重检索的更多细节。
聚类可以通过简单的 Keras API 实现，其中任何 Keras 模型（或层）都可以被包装和微调。参见以下用法示例。
权重聚类的优势
权重聚类在缩减跨序列化格式的模型存储和传输大小方面具有明显优势，因为共享参数模型的压缩率比无共享参数的模型高得多。这与稀疏（剪枝）模型类似，不同之处在于压缩的好处由减少唯一权重的数量实现，而剪枝则通过将低于一定阈值的权重设置为零来实现。Keras 模型聚类后，将其传递给任一常用压缩工具即可获得减小占用空间大小的好处。
为了进一步改进聚类相关推理时的内存使用量和速度，需要专用的运行时或编译器软件以及专用的机器学习硬件。以 Ethos-N 处理器的 Arm ML Ethos-N 驱动栈和 Ethos-U 处理器的 Ethos-U Vela 编译器为例。目前，这两个示例都需要先将优化的 Keras 模型量化并转换到 TensorFlow Lite。
聚类可以单独完成，也可以作为级联 Deep Compression 优化流水线的一部分，进一步缩减尺寸和提升推理速度。
压缩和准确率结果
通过在几个主流模型上进行实验，证明了权重聚类的压缩优势。还可以采用更激进的优化，但会牺牲准确率。虽然下表仅为 TensorFlow Lite 模型的测量结果，但对于 SavedModel 等其他序列化格式的益处也显而易见。
下表展示了实现结果的聚类配置。有些模型更容易因激进的聚类而导致准确率下降，在这种情况下，对优化稳健性更高的层使用选择性聚类。
模型聚类
TensorFlow 模型优化工具包从版本 v0.4.0 开始提供聚类 API。模型聚类前，需要先进行全面训练，然后再传递至聚类 API。完整模型聚类的代码片如下所示。
import tensorflow_model_optimization as tfmot
cluster_weights = tfmot.clustering.keras.cluster_weights
pretrained_model = pretrained_model()
clustering_params = {
'number_of_clusters': 32,
'cluster_centroids_init': tfmot.clustering.keras.CentroidInitialization.LINEAR
}
clustered_model = cluster_weights(pretrained_model, **clustering_params)
# Fine-tune
clustered_model.fit(...)
# Prepare model for serving by removing training-only variables.
model_for_serving = tfmot.clustering.keras.strip_clustering(clustered_model)
...
要对模型中的选定层进行聚类，可以在构造模型时将相同的聚类方法应用于这些层。
clustered_model = tf.keras.Sequential([
Dense(...),
cluster_weights(Dense(...,
kernel_initializer=pretrained_weights,
bias_initializer=pretrained_bias),
**clustering_params),
Dense(...)
])
当一个层进行选择性聚类前，仍然需要经过完全训练。因此，可以使用该层的 kernel_initializer 参数初始化权重。也可以使用 tf.keras.models.clone_model。
文档
要详细了解如何使用 API，可以从这个简单的 端到端聚类示例 Colab 开始。此处 包含更全面的指南和其他提示。
致谢
本文所介绍的功能和结果是多人合作的成果，包括 Arm ML Tooling 团队和我们在 Google TensorFlow 模型优化工具包团队的协作者。
来自 Arm - Anton Kachatkou、Aron Virginas-Tar、Ruomei Yan、Konstantin Sofeikov、Saoirse Stewart、Peng Sun、Elena Zhelezina、Gergely Nagy、Les Bell、Matteo Martincigh、Grant Watson、Diego Russo、Benjamin Klimczak、Thibaut Goetghebuer-Planchon。
来自 Google - Alan Chiao、Raziel Alvarez
原文：TensorFlow Model Optimization Toolkit — Weight Clustering API
中文：TensorFlow 公众号