Mohamed Nour Abouelseoud 的客座文章，以及Arm 的 Anton Kachatkou

我们很高兴为 TensorFlow 模型优化工具包 引入权重聚类 API，该 API 由 Arm 提出并贡献。

权重聚类是一种通过用更少数量的唯一值替换许多唯一参数值来减少模型存储和传输大小的技术。这种优势适用于所有部署。除了框架和硬件特定的支持（例如在 Arm Ethos-N 和 Ethos-U 机器学习处理器中），权重聚类还可以进一步改善内存占用和推理速度。

这项工作是 工具包路线图 的一部分，旨在支持开发更小、更快的 ML 模型。您可以查看有关 训练后量化、量化感知训练 和 稀疏性 的先前文章，以了解更多有关工具包及其功能的背景信息。

Arm 和 TensorFlow 团队一直在合作 改善对移动设备和物联网设备的部署。

什么是权重聚类？

越来越多的深度学习应用程序正在进入资源受限的环境，从智能手机到农业传感器和医疗仪器。这种向资源受限环境的转变导致了人们对更小、更高效的模型架构的努力，以及对模型优化技术的重视，例如 剪枝 和 量化。

权重聚类是一种优化算法，用于减少模型的存储和网络传输大小。简而言之，该想法在下面的图表中解释。

以下是图表说明。例如，假设模型中的一个层包含一个 4x4 的权重矩阵（如上图中的“权重矩阵”所示）。每个权重使用 float32 值存储。保存模型时，您将 16 个唯一的 float32 值存储到磁盘。

权重聚类通过将层中相似的权重替换为相同的值来减小模型的大小。这些值是通过对模型的训练权重运行聚类算法找到的。用户可以指定聚类的数量（在本例中为 4）。此步骤在上面的图表中显示为“获取质心”，并且 4 个质心值显示在“质心”表中。每个质心值都有一个索引（0-3）。

接下来，权重矩阵中的每个权重都将替换为其质心的索引。此步骤在“分配索引”中显示。现在，权重聚类算法可以存储在“拉取索引”中显示的修改后的矩阵（包含质心值的索引）以及质心值本身，而不是存储原始权重矩阵。

在本例中，我们将大小从 16 个唯一的浮点数减少到 4 个浮点数和 16 个 2 位索引。随着矩阵大小的增加，节省量也会增加。

请注意，即使我们仍然存储了 16 个浮点数，但它们现在只有 4 个不同的值。常见的压缩工具（如 zip）现在可以利用数据中的冗余来实现更高的压缩率。

聚类的技术实现来自 深度压缩：使用剪枝、训练量化和霍夫曼编码压缩深度神经网络。有关梯度更新和权重检索的更多详细信息，请参阅本文。

聚类可以通过简单的 Keras API 使用，其中任何 Keras 模型（或层）都可以被包装和微调。请参阅以下用法示例。

权重聚类的优势

权重聚类在跨序列化格式减少模型存储和传输大小方面具有直接优势，因为具有共享参数的模型比没有共享参数的模型具有更高的压缩率。这类似于稀疏（剪枝）模型，只是压缩优势是通过减少唯一权重的数量来实现的，而剪枝是通过将低于一定阈值的权重设置为零来实现的。一旦 Keras 模型被聚类，就可以通过任何常见的压缩工具传递它来获得减小大小的好处。

为了进一步释放与聚类相关的内存使用和推理速度方面的改进，需要专用运行时或编译器软件和专用机器学习硬件。例如，Arm ML Ethos-N 驱动程序堆栈 用于 Ethos-N 处理器，而 Ethos-U Vela 编译器 用于 Ethos-U 处理器。目前，这两个示例都需要先将优化的 Keras 模型量化并转换为 TensorFlow Lite。

聚类可以单独进行，也可以作为级联 深度压缩 优化管道的一部分来进行，以实现进一步的尺寸减小和推理速度。

压缩和准确性结果

在几个流行的模型上进行了实验，证明了权重聚类的压缩优势。可以应用更积极的优化，但会以准确性为代价。虽然下表包含了 TensorFlow Lite 模型的测量值，但对于其他序列化格式（如 SavedModel）也观察到了类似的优势。

下表演示了如何配置聚类以实现结果。有些模型更容易受到积极聚类的准确性下降的影响，在这种情况下，对更能抵抗优化的层使用了选择性聚类。

聚类模型

聚类 API 在 TensorFlow 模型优化工具包的 v0.4.0 版本中可用。要聚类模型，需要先完全训练模型，然后再将其传递给聚类 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 工具团队以及我们在 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