 在 TensorFlow.org 上查看
|
 在 Google Colab 中运行
|
 在 GitHub 上查看源代码
|
 下载笔记本
|
欢迎使用 TensorFlow 模型优化工具包中的权重聚类综合指南。
本页面记录了各种用例，并展示了如何为每个用例使用 API。 了解所需的 API 后，请在API 文档中查找参数和低级详细信息。
本指南涵盖以下用例
- 定义聚类模型。
- 检查点和反序列化聚类模型。
- 提高聚类模型的准确性。
- 仅用于部署，您必须采取措施才能看到压缩优势。
设置
! pip install -q tensorflow-model-optimization
import tensorflow as tf
import tf_keras as keras
import numpy as np
import tempfile
import os
import tensorflow_model_optimization as tfmot
input_dim = 20
output_dim = 20
x_train = np.random.randn(1, input_dim).astype(np.float32)
y_train = keras.utils.to_categorical(np.random.randn(1), num_classes=output_dim)
def setup_model():
model = keras.Sequential([
keras.layers.Dense(input_dim, input_shape=[input_dim]),
keras.layers.Flatten()
])
return model
def train_model(model):
model.compile(
loss=keras.losses.categorical_crossentropy,
optimizer='adam',
metrics=['accuracy']
)
model.summary()
model.fit(x_train, y_train)
return model
def save_model_weights(model):
_, pretrained_weights = tempfile.mkstemp('.h5')
model.save_weights(pretrained_weights)
return pretrained_weights
def setup_pretrained_weights():
model= setup_model()
model = train_model(model)
pretrained_weights = save_model_weights(model)
return pretrained_weights
def setup_pretrained_model():
model = setup_model()
pretrained_weights = setup_pretrained_weights()
model.load_weights(pretrained_weights)
return model
def save_model_file(model):
_, keras_file = tempfile.mkstemp('.h5')
model.save(keras_file, include_optimizer=False)
return keras_file
def get_gzipped_model_size(model):
# It returns the size of the gzipped model in bytes.
import os
import zipfile
keras_file = save_model_file(model)
_, zipped_file = tempfile.mkstemp('.zip')
with zipfile.ZipFile(zipped_file, 'w', compression=zipfile.ZIP_DEFLATED) as f:
f.write(keras_file)
return os.path.getsize(zipped_file)
setup_model()
pretrained_weights = setup_pretrained_weights()
2024-03-09 12:38:26.610043: E external/local_xla/xla/stream_executor/cuda/cuda_driver.cc:282] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected
定义聚类模型
聚类整个模型（顺序和函数式）
提示，以提高模型准确性
- 您必须将具有可接受准确性的预训练模型传递给此 API。 从头开始训练具有聚类的模型会导致准确性低下。
- 在某些情况下，对某些层的聚类会对模型准确性产生负面影响。 请查看“聚类某些层”以了解如何跳过对影响准确性最大的层的聚类。
要聚类所有层，请将tfmot.clustering.keras.cluster_weights应用于模型。
import tensorflow_model_optimization as tfmot
cluster_weights = tfmot.clustering.keras.cluster_weights
CentroidInitialization = tfmot.clustering.keras.CentroidInitialization
clustering_params = {
'number_of_clusters': 3,
'cluster_centroids_init': CentroidInitialization.KMEANS_PLUS_PLUS
}
model = setup_model()
model.load_weights(pretrained_weights)
clustered_model = cluster_weights(model, **clustering_params)
clustered_model.summary()
Model: "sequential_2"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
cluster_dense_2 (ClusterWe (None, 20) 823
ights)
cluster_flatten_2 (Cluster (None, 20) 0
Weights)
=================================================================
Total params: 823 (4.78 KB)
Trainable params: 423 (1.65 KB)
Non-trainable params: 400 (3.12 KB)
_________________________________________________________________
聚类某些层（顺序和函数式模型）
提示，以提高模型准确性
- 您必须将具有可接受准确性的预训练模型传递给此 API。 从头开始训练具有聚类的模型会导致准确性低下。
- 聚类具有更多冗余参数的后续层（例如
keras.layers.Dense、keras.layers.Conv2D），而不是早期层。 - 在微调期间，冻结聚类层之前的早期层。 将冻结层的数量视为超参数。 从经验上讲，冻结大多数早期层对于当前聚类 API 来说是理想的。
- 避免对关键层（例如注意力机制）进行聚类。
更多：tfmot.clustering.keras.cluster_weights API 文档提供了有关如何按层更改聚类配置的详细信息。
# Create a base model
base_model = setup_model()
base_model.load_weights(pretrained_weights)
# Helper function uses `cluster_weights` to make only
# the Dense layers train with clustering
def apply_clustering_to_dense(layer):
if isinstance(layer, keras.layers.Dense):
return cluster_weights(layer, **clustering_params)
return layer
# Use `keras.models.clone_model` to apply `apply_clustering_to_dense`
# to the layers of the model.
clustered_model = keras.models.clone_model(
base_model,
clone_function=apply_clustering_to_dense,
)
clustered_model.summary()
Model: "sequential_3"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
cluster_dense_3 (ClusterWe (None, 20) 823
ights)
flatten_3 (Flatten) (None, 20) 0
=================================================================
Total params: 823 (4.78 KB)
Trainable params: 423 (1.65 KB)
Non-trainable params: 400 (3.12 KB)
_________________________________________________________________
按通道聚类卷积层
聚类模型可以传递给进一步的优化，例如训练后量化。 如果量化是按通道进行的，则模型也应按通道进行聚类。 这会提高聚类和量化模型的准确性。
要按通道聚类，参数 cluster_per_channel 应设置为 True。 它可以为某些层或整个模型设置。
提示
如果模型要进一步量化，您可以考虑使用聚类保留 QAT 技术。
在按通道应用聚类之前，可以对模型进行剪枝。 使用参数
preserve_sparsity设置为True，在按通道聚类期间会保留稀疏性。 请注意，在这种情况下，应使用稀疏性和聚类保留 QAT 技术。
聚类自定义 Keras 层或指定要聚类的层权重
tfmot.clustering.keras.ClusterableLayer 用于两种用例
- 聚类任何不受原生支持的层，包括自定义 Keras 层。
- 指定要聚类的支持层的哪些权重。
例如，API 默认情况下仅对 Dense 层的内核进行聚类。 下面的示例展示了如何对其进行修改以同时对偏差进行聚类。 请注意，从 keras 层派生时，您需要覆盖函数 get_clusterable_weights，在其中指定要聚类的可训练变量的名称和可训练变量本身。 例如，如果您返回一个空列表 []，则不会聚类任何权重。
常见错误： 对偏差进行聚类通常会过度损害模型精度。
class MyDenseLayer(keras.layers.Dense, tfmot.clustering.keras.ClusterableLayer):
def get_clusterable_weights(self):
# Cluster kernel and bias. This is just an example, clustering
# bias usually hurts model accuracy.
return [('kernel', self.kernel), ('bias', self.bias)]
# Use `cluster_weights` to make the `MyDenseLayer` layer train with clustering as usual.
model_for_clustering = keras.Sequential([
tfmot.clustering.keras.cluster_weights(MyDenseLayer(20, input_shape=[input_dim]), **clustering_params),
keras.layers.Flatten()
])
model_for_clustering.summary()
Model: "sequential_4"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
cluster_my_dense_layer (Cl (None, 20) 846
usterWeights)
flatten_4 (Flatten) (None, 20) 0
=================================================================
Total params: 846 (4.95 KB)
Trainable params: 426 (1.66 KB)
Non-trainable params: 420 (3.28 KB)
_________________________________________________________________
您也可以使用 tfmot.clustering.keras.ClusterableLayer 对 Keras 自定义层进行聚类。为此，您需要像往常一样扩展 keras.Layer 并实现 __init__、call 和 build 函数，但您还需要扩展 clusterable_layer.ClusterableLayer 类并实现
get_clusterable_weights，在其中指定要聚类的权重内核，如上所示。get_clusterable_algorithm，在其中指定权重张量的聚类算法。这是因为您需要指定自定义层权重如何为聚类而整形。返回的聚类算法类应派生自clustering_algorithm.ClusteringAlgorithm类，并且应覆盖get_pulling_indices函数。此函数的示例（支持 1D、2D 和 3D 秩的权重）可以在 此处 找到。
此用例的示例可以在 此处 找到。
检查点和反序列化聚类模型
您的用例： 此代码仅适用于 HDF5 模型格式（不适用于 HDF5 权重或其他格式）。
# Define the model.
base_model = setup_model()
base_model.load_weights(pretrained_weights)
clustered_model = cluster_weights(base_model, **clustering_params)
# Save or checkpoint the model.
_, keras_model_file = tempfile.mkstemp('.h5')
clustered_model.save(keras_model_file, include_optimizer=True)
# `cluster_scope` is needed for deserializing HDF5 models.
with tfmot.clustering.keras.cluster_scope():
loaded_model = keras.models.load_model(keras_model_file)
loaded_model.summary()
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
WARNING:tensorflow:No training configuration found in the save file, so the model was *not* compiled. Compile it manually.
Model: "sequential_5"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
cluster_dense_4 (ClusterWe (None, 20) 823
ights)
cluster_flatten_5 (Cluster (None, 20) 0
Weights)
=================================================================
Total params: 823 (4.78 KB)
Trainable params: 423 (1.65 KB)
Non-trainable params: 400 (3.12 KB)
_________________________________________________________________
/tmpfs/src/tf_docs_env/lib/python3.9/site-packages/tf_keras/src/engine/training.py:3098: UserWarning: You are saving your model as an HDF5 file via `model.save()`. This file format is considered legacy. We recommend using instead the native TF-Keras format, e.g. `model.save('my_model.keras')`.
saving_api.save_model(
提高聚类模型的精度
对于您的特定用例，您可以考虑以下提示
质心初始化在最终优化模型精度中起着关键作用。通常，kmeans++ 初始化优于线性、密度和随机初始化。当不使用 kmeans++ 时，线性初始化往往优于密度和随机初始化，因为它不会倾向于遗漏较大的权重。但是，对于在具有双峰分布的权重上使用很少的聚类的情况，已观察到密度初始化可以提供更好的精度。
设置一个低于训练时使用的学习率，以微调聚类模型。
有关提高模型精度的通用想法，请在“定义聚类模型”下查找适合您的用例的提示。
部署
导出具有大小压缩的模型
常见错误：strip_clustering 和应用标准压缩算法（例如通过 gzip）都是必要的，才能看到聚类的压缩优势。
model = setup_model()
clustered_model = cluster_weights(model, **clustering_params)
clustered_model.compile(
loss=keras.losses.categorical_crossentropy,
optimizer='adam',
metrics=['accuracy']
)
clustered_model.fit(
x_train,
y_train
)
final_model = tfmot.clustering.keras.strip_clustering(clustered_model)
print("final model")
final_model.summary()
print("\n")
print("Size of gzipped clustered model without stripping: %.2f bytes"
% (get_gzipped_model_size(clustered_model)))
print("Size of gzipped clustered model with stripping: %.2f bytes"
% (get_gzipped_model_size(final_model)))
1/1 [==============================] - 1s 990ms/step - loss: 16.1181 - accuracy: 0.0000e+00
final model
Model: "sequential_6"
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
dense_5 (Dense) (None, 20) 420
flatten_6 (Flatten) (None, 20) 0
=================================================================
Total params: 420 (1.64 KB)
Trainable params: 420 (1.64 KB)
Non-trainable params: 0 (0.00 Byte)
_________________________________________________________________
Size of gzipped clustered model without stripping: 3516.00 bytes
WARNING:tensorflow:Compiled the loaded model, but the compiled metrics have yet to be built. `model.compile_metrics` will be empty until you train or evaluate the model.
Size of gzipped clustered model with stripping: 1469.00 bytes