Keras 範例中的權重叢集

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視原始碼 下載筆記本

總覽

歡迎使用權重叢集的端對端範例,這是 TensorFlow Model Optimization Toolkit 的一部分。

其他頁面

如需權重叢集簡介,以及判斷是否應使用權重叢集 (包括支援的項目),請參閱總覽頁面

若要快速找到您的使用案例所需的 API (除了完全叢集化包含 16 個叢集的模型之外),請參閱綜合指南

目錄

在本教學課程中,您將會:

  1. 從頭開始訓練用於 MNIST 資料集的 keras 模型。
  2. 微調模型,方法是套用權重叢集 API 並查看準確度。
  3. 從叢集建立小 6 倍的 TF 和 TFLite 模型。
  4. 從結合權重叢集和訓練後量化建立小 8 倍的 TFLite 模型。
  5. 查看從 TF 到 TFLite 的準確度持久性。

設定

您可以在本機 virtualenvcolab 中執行此 Jupyter Notebook。如需設定依附元件的詳細資訊,請參閱安裝指南

 pip install -q tensorflow-model-optimization
import tensorflow as tf
from tensorflow_model_optimization.python.core.keras.compat import keras

import numpy as np
import tempfile
import zipfile
import os

訓練未經叢集化的 MNIST keras 模型

# Load MNIST dataset
(train_images, train_labels), (test_images, test_labels) = tf.keras.datasets.mnist.load_data()

# Normalize the input image so that each pixel value is between 0 to 1.
train_images = train_images / 255.0
test_images  = test_images / 255.0

# Define the model architecture.
model = keras.Sequential([
    keras.layers.InputLayer(input_shape=(28, 28)),
    keras.layers.Reshape(target_shape=(28, 28, 1)),
    keras.layers.Conv2D(filters=12, kernel_size=(3, 3), activation=tf.nn.relu),
    keras.layers.MaxPooling2D(pool_size=(2, 2)),
    keras.layers.Flatten(),
    keras.layers.Dense(10)
])

# Train the digit classification model
model.compile(optimizer='adam',
              loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model.fit(
    train_images,
    train_labels,
    validation_split=0.1,
    epochs=10
)
2024-03-09 12:34:47.914475: 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
Epoch 1/10
1688/1688 [==============================] - 21s 4ms/step - loss: 0.2914 - accuracy: 0.9183 - val_loss: 0.1165 - val_accuracy: 0.9680
Epoch 2/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.1108 - accuracy: 0.9686 - val_loss: 0.0783 - val_accuracy: 0.9788
Epoch 3/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0807 - accuracy: 0.9769 - val_loss: 0.0671 - val_accuracy: 0.9833
Epoch 4/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0661 - accuracy: 0.9806 - val_loss: 0.0646 - val_accuracy: 0.9840
Epoch 5/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0575 - accuracy: 0.9826 - val_loss: 0.0596 - val_accuracy: 0.9850
Epoch 6/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0514 - accuracy: 0.9843 - val_loss: 0.0594 - val_accuracy: 0.9852
Epoch 7/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0465 - accuracy: 0.9859 - val_loss: 0.0685 - val_accuracy: 0.9820
Epoch 8/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0422 - accuracy: 0.9873 - val_loss: 0.0622 - val_accuracy: 0.9845
Epoch 9/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0389 - accuracy: 0.9880 - val_loss: 0.0594 - val_accuracy: 0.9837
Epoch 10/10
1688/1688 [==============================] - 7s 4ms/step - loss: 0.0354 - accuracy: 0.9894 - val_loss: 0.0644 - val_accuracy: 0.9843
<tf_keras.src.callbacks.History at 0x7f1be5b37430>

評估基準模型並儲存以供日後使用

_, baseline_model_accuracy = model.evaluate(
    test_images, test_labels, verbose=0)

print('Baseline test accuracy:', baseline_model_accuracy)

_, keras_file = tempfile.mkstemp('.h5')
print('Saving model to: ', keras_file)
keras.models.save_model(model, keras_file, include_optimizer=False)
Baseline test accuracy: 0.98089998960495
Saving model to:  /tmpfs/tmp/tmpigrs28_d.h5
/tmpfs/tmp/ipykernel_29244/3680774635.py:8: 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')`.
  keras.models.save_model(model, keras_file, include_optimizer=False)

使用叢集微調預先訓練的模型

cluster_weights() API 套用至整個預先訓練的模型,以示範其在套用 zip 後減少模型大小,同時保持良好準確度的有效性。如需如何為您的使用案例最佳平衡準確度和壓縮率,請參閱綜合指南中的逐層範例。

定義模型並套用叢集 API

在將模型傳遞至叢集 API 之前,請確定模型已訓練完成,並顯示一些可接受的準確度。

import tensorflow_model_optimization as tfmot

cluster_weights = tfmot.clustering.keras.cluster_weights
CentroidInitialization = tfmot.clustering.keras.CentroidInitialization

clustering_params = {
  'number_of_clusters': 16,
  'cluster_centroids_init': CentroidInitialization.LINEAR
}

# Cluster a whole model
clustered_model = cluster_weights(model, **clustering_params)

# Use smaller learning rate for fine-tuning clustered model
opt = keras.optimizers.Adam(learning_rate=1e-5)

clustered_model.compile(
  loss=keras.losses.SparseCategoricalCrossentropy(from_logits=True),
  optimizer=opt,
  metrics=['accuracy'])

clustered_model.summary()
Model: "sequential"
_________________________________________________________________
 Layer (type)                Output Shape              Param #   
=================================================================
 cluster_reshape (ClusterWe  (None, 28, 28, 1)         0         
 ights)                                                          
                                                                 
 cluster_conv2d (ClusterWei  (None, 26, 26, 12)        244       
 ghts)                                                           
                                                                 
 cluster_max_pooling2d (Clu  (None, 13, 13, 12)        0         
 sterWeights)                                                    
                                                                 
 cluster_flatten (ClusterWe  (None, 2028)              0         
 ights)                                                          
                                                                 
 cluster_dense (ClusterWeig  (None, 10)                40586     
 hts)                                                            
                                                                 
=================================================================
Total params: 40830 (239.13 KB)
Trainable params: 20442 (79.85 KB)
Non-trainable params: 20388 (159.28 KB)
_________________________________________________________________

微調模型並評估與基準的準確度

使用叢集微調模型 1 個 epoch。

# Fine-tune model
clustered_model.fit(
  train_images,
  train_labels,
  batch_size=500,
  epochs=1,
  validation_split=0.1)
108/108 [==============================] - 3s 17ms/step - loss: 0.0397 - accuracy: 0.9875 - val_loss: 0.0671 - val_accuracy: 0.9833
<tf_keras.src.callbacks.History at 0x7f1be5b9a790>

在此範例中,與基準相比,叢集化後測試準確度的損失極小。

_, clustered_model_accuracy = clustered_model.evaluate(
  test_images, test_labels, verbose=0)

print('Baseline test accuracy:', baseline_model_accuracy)
print('Clustered test accuracy:', clustered_model_accuracy)
Baseline test accuracy: 0.98089998960495
Clustered test accuracy: 0.9786999821662903

從叢集建立小 6 倍的模型

需要 strip_clustering 和套用標準壓縮演算法 (例如透過 gzip),才能看到叢集的壓縮優勢。

首先,為 TensorFlow 建立可壓縮模型。在此,strip_clustering 移除叢集僅在訓練期間需要的所有變數 (例如,用於儲存叢集中心點和索引的 tf.Variable),否則這些變數會在推論期間增加模型大小。

final_model = tfmot.clustering.keras.strip_clustering(clustered_model)

_, clustered_keras_file = tempfile.mkstemp('.h5')
print('Saving clustered model to: ', clustered_keras_file)
keras.models.save_model(final_model, clustered_keras_file, 
                           include_optimizer=False)
Saving clustered model to:  /tmpfs/tmp/tmpu8mqv83j.h5
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.
/tmpfs/tmp/ipykernel_29244/2668672504.py:5: 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')`.
  keras.models.save_model(final_model, clustered_keras_file,

然後,為 TFLite 建立可壓縮模型。您可以將叢集化模型轉換為可在目標後端執行的格式。TensorFlow Lite 是您可以部署到行動裝置的範例。

clustered_tflite_file = '/tmp/clustered_mnist.tflite'
converter = tf.lite.TFLiteConverter.from_keras_model(final_model)
tflite_clustered_model = converter.convert()
with open(clustered_tflite_file, 'wb') as f:
  f.write(tflite_clustered_model)
print('Saved clustered TFLite model to:', clustered_tflite_file)
INFO:tensorflow:Assets written to: /tmpfs/tmp/tmpdw8boe6k/assets
INFO:tensorflow:Assets written to: /tmpfs/tmp/tmpdw8boe6k/assets
Saved clustered TFLite model to: /tmp/clustered_mnist.tflite
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
W0000 00:00:1709987777.341582   29244 tf_tfl_flatbuffer_helpers.cc:390] Ignored output_format.
W0000 00:00:1709987777.341632   29244 tf_tfl_flatbuffer_helpers.cc:393] Ignored drop_control_dependency.

定義協助程式函式,以實際透過 gzip 壓縮模型並測量壓縮後的大小。

def get_gzipped_model_size(file):
  # It returns the size of the gzipped model in bytes.
  import os
  import zipfile

  _, zipped_file = tempfile.mkstemp('.zip')
  with zipfile.ZipFile(zipped_file, 'w', compression=zipfile.ZIP_DEFLATED) as f:
    f.write(file)

  return os.path.getsize(zipped_file)

比較並查看模型因叢集化而小 6 倍

print("Size of gzipped baseline Keras model: %.2f bytes" % (get_gzipped_model_size(keras_file)))
print("Size of gzipped clustered Keras model: %.2f bytes" % (get_gzipped_model_size(clustered_keras_file)))
print("Size of gzipped clustered TFlite model: %.2f bytes" % (get_gzipped_model_size(clustered_tflite_file)))
Size of gzipped baseline Keras model: 78177.00 bytes
Size of gzipped clustered Keras model: 13053.00 bytes
Size of gzipped clustered TFlite model: 12638.00 bytes

從結合權重叢集和訓練後量化建立小 8 倍的 TFLite 模型

您可以將訓練後量化套用至叢集化模型,以獲得額外的好處。

converter = tf.lite.TFLiteConverter.from_keras_model(final_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
tflite_quant_model = converter.convert()

_, quantized_and_clustered_tflite_file = tempfile.mkstemp('.tflite')

with open(quantized_and_clustered_tflite_file, 'wb') as f:
  f.write(tflite_quant_model)

print('Saved quantized and clustered TFLite model to:', quantized_and_clustered_tflite_file)
print("Size of gzipped baseline Keras model: %.2f bytes" % (get_gzipped_model_size(keras_file)))
print("Size of gzipped clustered and quantized TFlite model: %.2f bytes" % (get_gzipped_model_size(quantized_and_clustered_tflite_file)))
INFO:tensorflow:Assets written to: /tmpfs/tmp/tmpzkdsr0o3/assets
INFO:tensorflow:Assets written to: /tmpfs/tmp/tmpzkdsr0o3/assets
W0000 00:00:1709987778.159317   29244 tf_tfl_flatbuffer_helpers.cc:390] Ignored output_format.
W0000 00:00:1709987778.159352   29244 tf_tfl_flatbuffer_helpers.cc:393] Ignored drop_control_dependency.
Saved quantized and clustered TFLite model to: /tmpfs/tmp/tmpylkyahiy.tflite
Size of gzipped baseline Keras model: 78177.00 bytes
Size of gzipped clustered and quantized TFlite model: 9792.00 bytes

查看從 TF 到 TFLite 的準確度持久性

定義協助程式函式,以評估測試資料集上的 TFLite 模型。

def eval_model(interpreter):
  input_index = interpreter.get_input_details()[0]["index"]
  output_index = interpreter.get_output_details()[0]["index"]

  # Run predictions on every image in the "test" dataset.
  prediction_digits = []
  for i, test_image in enumerate(test_images):
    if i % 1000 == 0:
      print('Evaluated on {n} results so far.'.format(n=i))
    # Pre-processing: add batch dimension and convert to float32 to match with
    # the model's input data format.
    test_image = np.expand_dims(test_image, axis=0).astype(np.float32)
    interpreter.set_tensor(input_index, test_image)

    # Run inference.
    interpreter.invoke()

    # Post-processing: remove batch dimension and find the digit with highest
    # probability.
    output = interpreter.tensor(output_index)
    digit = np.argmax(output()[0])
    prediction_digits.append(digit)

  print('\n')
  # Compare prediction results with ground truth labels to calculate accuracy.
  prediction_digits = np.array(prediction_digits)
  accuracy = (prediction_digits == test_labels).mean()
  return accuracy

您評估已叢集化和量化的模型,然後查看從 TensorFlow 持續到 TFLite 後端的準確度。

interpreter = tf.lite.Interpreter(model_content=tflite_quant_model)
interpreter.allocate_tensors()

test_accuracy = eval_model(interpreter)

print('Clustered and quantized TFLite test_accuracy:', test_accuracy)
print('Clustered TF test accuracy:', clustered_model_accuracy)
INFO: Created TensorFlow Lite XNNPACK delegate for CPU.
WARNING: Attempting to use a delegate that only supports static-sized tensors with a graph that has dynamic-sized tensors (tensor#13 is a dynamic-sized tensor).
Evaluated on 0 results so far.
Evaluated on 1000 results so far.
Evaluated on 2000 results so far.
Evaluated on 3000 results so far.
Evaluated on 4000 results so far.
Evaluated on 5000 results so far.
Evaluated on 6000 results so far.
Evaluated on 7000 results so far.
Evaluated on 8000 results so far.
Evaluated on 9000 results so far.


Clustered and quantized TFLite test_accuracy: 0.9791
Clustered TF test accuracy: 0.9786999821662903

結論

在本教學課程中,您已瞭解如何使用 TensorFlow Model Optimization Toolkit API 建立叢集化模型。更具體來說,您已完成建立小 8 倍 MNIST 模型且準確度差異極小的端對端範例。我們鼓勵您試用這項新功能,這對於在資源受限環境中部署可能尤其重要。