权重聚类综合指南

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看源代码

下载笔记本

欢迎阅读 TensorFlow Model Optimization Toolkit 中权重聚类的综合指南。

本页面记录了各种用例，并展示了如何将 API 用于每种用例。了解需要哪些 API 后，可在 API 文档中找到参数和底层详细信息：

如果要查看权重聚类的好处以及支持的功能，请查看概述。
有关单个端到端示例，请参阅权重聚类示例。

本指南涵盖了以下用例：

定义聚类模型。
为聚类模型设置检查点和进行反序列化。
提高聚类模型的准确率。
仅对于部署而言，您必须采取措施才能看到压缩的好处。

设置

! pip install -q tensorflow-model-optimization

import tensorflow as tf
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 = tf.keras.utils.to_categorical(np.random.randn(1), num_classes=output_dim)

def setup_model():
  model = tf.keras.Sequential([
      tf.keras.layers.Dense(input_dim, input_shape=[input_dim]),
      tf.keras.layers.Flatten()
  ])
  return model

def train_model(model):
  model.compile(
      loss=tf.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()

定义聚类模型

聚类整个模型（序贯模型和函数式模型）

提高模型准确率的提示：

您必须将具有可接受准确率的预训练模型传递给此 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.DENSITY_BASED
}

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 (ClusterWeig (None, 20)                423       
_________________________________________________________________
cluster_flatten_2 (ClusterWe (None, 20)                0         
=================================================================
Total params: 423
Trainable params: 23
Non-trainable params: 400
_________________________________________________________________

聚类某些层（序贯模型和函数式模型）

提高模型准确率的提示：

您必须将具有可接受准确率的预训练模型传递给此 API。使用聚类从头开始训练模型会导致准确率不佳。
与前面的层相反，使用更多冗余参数（例如 tf.keras.layers.Dense 和 tf.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, tf.keras.layers.Dense):
    return cluster_weights(layer, **clustering_params)
  return layer

# Use `tf.keras.models.clone_model` to apply `apply_clustering_to_dense` 
# to the layers of the model.
clustered_model = tf.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 (ClusterWeig (None, 20)                423       
_________________________________________________________________
flatten_3 (Flatten)          (None, 20)                0         
=================================================================
Total params: 423
Trainable params: 23
Non-trainable params: 400
_________________________________________________________________

为聚类模型设置检查点和进行反序列化

您的用例：仅 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 = tf.keras.models.load_model(keras_model_file)

loaded_model.summary()

WARNING:tensorflow:No training configuration found in the save file, so the model was *not* compiled. Compile it manually.
Model: "sequential_4"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
cluster_dense_4 (ClusterWeig (None, 20)                423       
_________________________________________________________________
cluster_flatten_4 (ClusterWe (None, 20)                0         
=================================================================
Total params: 423
Trainable params: 23
Non-trainable params: 400
_________________________________________________________________

提高聚类模型的准确率

对于您的特定用例，您可以考虑以下提示：

形心初始化在最终优化的模型准确率中起到关键作用。通常，线性初始化优于密度和随机初始化，因为它不会丢失较大的权重。但是，对于在具有双峰分布的权重上使用极少簇的情况，已经观察到密度初始化可以提供更出色的准确率。
微调聚类模型时，将学习率设置为低于训练中使用的学习率。
有关提高模型准确率的总体思路，请在“定义聚类模型”下查找您的用例对应的提示。

部署

导出大小经过压缩的模型

常见误区：strip_clustering 和应用标准压缩算法（例如通过 gzip）对于看到聚类压缩的好处必不可少。

model = setup_model()
clustered_model = cluster_weights(model, **clustering_params)

clustered_model.compile(
    loss=tf.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 [==============================] - 0s 2ms/step - loss: 16.1181 - accuracy: 0.0000e+00
final model
Model: "sequential_5"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
dense_5 (Dense)              (None, 20)                420       
_________________________________________________________________
flatten_5 (Flatten)          (None, 20)                0         
=================================================================
Total params: 420
Trainable params: 420
Non-trainable params: 0
_________________________________________________________________


Size of gzipped clustered model without stripping: 1865.00 bytes
Size of gzipped clustered model with stripping: 1462.00 bytes