Keras での重みクラスタリングの例

TensorFlow.orgで表示 Google Colabで実行 GitHub でソースを表示 ノートブックをダウンロード

概要

TensorFlow Model Optimization ツールキットの一部である重みクラスタリングのエンドツーエンドの例へようこそ。

その他のページ

重みクラスタリングの紹介、およびクラスタリングを使用すべきかどうかの判定(サポート情報も含む)については、概要ページをご覧ください。

ユースケースに合った API を素早く特定するには(16 個のクラスタでモデルを完全クラスタ化するケースを超える内容)、総合ガイドをご覧ください。

内容

チュートリアルでは、次について説明しています。

  1. MNIST データセットの tf.keras モデルを最初からトレーニングする
  2. 重みクラスタリング API を適用してモデルを微調整し、精度を確認する
  3. クラスタリングによって 6 倍小さな TF および TFLite モデルを作成する
  4. 重みクラスタリングとポストトレーニング量子化を組み合わせて、8 倍小さな TFLite モデルを作成する
  5. TF から TFLite への精度の永続性を確認する

セットアップ

この Jupyter ノートブックは、ローカルの virtualenv または Colab で実行できます。依存関係のセットアップに関する詳細は、インストールガイドをご覧ください。

 pip install -q tensorflow-model-optimization
import tensorflow as tf
from tensorflow import keras

import numpy as np
import tempfile
import zipfile
import os

クラスタを使用せずに、MNIST の tf.keras モデルをトレーニングする

# Load MNIST dataset
mnist = keras.datasets.mnist
(train_images, train_labels), (test_images, test_labels) = 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=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])

model.fit(
    train_images,
    train_labels,
    validation_split=0.1,
    epochs=10
)
Epoch 1/10
1688/1688 [==============================] - 7s 2ms/step - loss: 0.5279 - accuracy: 0.8552 - val_loss: 0.1275 - val_accuracy: 0.9650
Epoch 2/10
1688/1688 [==============================] - 4s 2ms/step - loss: 0.1340 - accuracy: 0.9622 - val_loss: 0.0869 - val_accuracy: 0.9785
Epoch 3/10
1688/1688 [==============================] - 4s 2ms/step - loss: 0.0928 - accuracy: 0.9731 - val_loss: 0.0747 - val_accuracy: 0.9800
Epoch 4/10
1688/1688 [==============================] - 4s 2ms/step - loss: 0.0755 - accuracy: 0.9771 - val_loss: 0.0706 - val_accuracy: 0.9820
Epoch 5/10
1688/1688 [==============================] - 4s 2ms/step - loss: 0.0625 - accuracy: 0.9819 - val_loss: 0.0634 - val_accuracy: 0.9832
Epoch 6/10
1688/1688 [==============================] - 4s 2ms/step - loss: 0.0586 - accuracy: 0.9819 - val_loss: 0.0606 - val_accuracy: 0.9830
Epoch 7/10
1688/1688 [==============================] - 4s 2ms/step - loss: 0.0530 - accuracy: 0.9842 - val_loss: 0.0616 - val_accuracy: 0.9847
Epoch 8/10
1688/1688 [==============================] - 4s 2ms/step - loss: 0.0444 - accuracy: 0.9864 - val_loss: 0.0622 - val_accuracy: 0.9843
Epoch 9/10
1688/1688 [==============================] - 4s 2ms/step - loss: 0.0407 - accuracy: 0.9874 - val_loss: 0.0620 - val_accuracy: 0.9840
Epoch 10/10
1688/1688 [==============================] - 4s 2ms/step - loss: 0.0372 - accuracy: 0.9888 - val_loss: 0.0684 - val_accuracy: 0.9818
<tensorflow.python.keras.callbacks.History at 0x7fa994f7a588>

ベースラインモデルを評価して後で使用できるように保存する

_, 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)
tf.keras.models.save_model(model, keras_file, include_optimizer=False)
Baseline test accuracy: 0.9789999723434448
Saving model to:  /tmp/tmp8m1tzkhh.h5

クラスタを使ってトレーニング済みのモデルを微調整する

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 = tf.keras.optimizers.Adam(learning_rate=1e-5)

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

clustered_model.summary()
Model: "sequential"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
cluster_reshape (ClusterWeig (None, 28, 28, 1)         0         
_________________________________________________________________
cluster_conv2d (ClusterWeigh (None, 26, 26, 12)        136       
_________________________________________________________________
cluster_max_pooling2d (Clust (None, 13, 13, 12)        0         
_________________________________________________________________
cluster_flatten (ClusterWeig (None, 2028)              0         
_________________________________________________________________
cluster_dense (ClusterWeight (None, 10)                20306     
=================================================================
Total params: 20,442
Trainable params: 54
Non-trainable params: 20,388
_________________________________________________________________

モデルを微調整し、ベースラインに対する精度を評価する

1 エポック、クラスタでモデルを微調整します。

# Fine-tune model
clustered_model.fit(
  train_images,
  train_labels,
  batch_size=500,
  epochs=1,
  validation_split=0.1)
108/108 [==============================] - 1s 6ms/step - loss: 0.0683 - accuracy: 0.9770 - val_loss: 0.0977 - val_accuracy: 0.9747
<tensorflow.python.keras.callbacks.History at 0x7faa032a4d68>

この例では、ベースラインと比較し、クラスタリング後のテスト精度に最小限の損失があります。

_, 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.9789999723434448
Clustered test accuracy: 0.9702000021934509

クラスタリングによって 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)
tf.keras.models.save_model(final_model, clustered_keras_file, 
                           include_optimizer=False)
Saving clustered model to:  /tmp/tmpz9c4ugbj.h5

次に、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: /tmp/tmp41s2kzub/assets
Saved clustered TFLite model to: /tmp/clustered_mnist.tflite

実際に gzip でモデルを圧縮し、zip 圧縮されたサイズを測定するヘルパー関数を定義します。

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: 78047.00 bytes
Size of gzipped clustered Keras model: 12419.00 bytes
Size of gzipped clustered TFlite model: 11920.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: /tmp/tmpg3y6r2ll/assets
INFO:tensorflow:Assets written to: /tmp/tmpg3y6r2ll/assets
Saved quantized and clustered TFLite model to: /tmp/tmp6p9j3n3w.tflite
Size of gzipped baseline Keras model: 78047.00 bytes
Size of gzipped clustered and quantized TFlite model: 9045.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)
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.9698
Clustered TF test accuracy: 0.9702000021934509

まとめ

このチュートリアルでは、TensorFlow Model Optimization Toolkit API を使用してクラスタモデルを作成する方法を確認しました。より具体的には、精度の違いを最小限に抑えて MNIST の 8 倍の小さいモデルを作成する、エンドツーエンドの例を確認しました。この新しい機能を試すことをお勧めします。これは、リソースに制約のある環境でのデプロイに特に重要な機能です。