 TensorFlow.orgで表示 |
 Google Colabで実行 |
 GitHub でソースを表示 |
 ノートブックをダウンロード |
概要
TensorFlow Model Optimization ツールキットの一部である重みクラスタリングのエンドツーエンドの例へようこそ。
その他のページ
重みクラスタリングの紹介、およびクラスタリングを使用すべきかどうかの判定（サポート情報も含む）については、概要ページをご覧ください。
ユースケースに合った API を素早く特定するには（16 個のクラスタでモデルを完全クラスタ化するケースを超える内容）、総合ガイドをご覧ください。
内容
チュートリアルでは、次について説明しています。
- MNIST データセットの
tf.kerasモデルを最初からトレーニングする - 重みクラスタリング API を適用してモデルを微調整し、精度を確認する
- クラスタリングによって 6 倍小さな TF および TFLite モデルを作成する
- 重みクラスタリングとポストトレーニング量子化を組み合わせて、8 倍小さな TFLite モデルを作成する
- TF から TFLite への精度の永続性を確認する
セットアップ
この Jupyter ノートブックは、ローカルの virtualenv または Colab で実行できます。依存関係のセットアップに関する詳細は、インストールガイドをご覧ください。
pip install -q tensorflow-model-optimizationimport 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 倍の小さいモデルを作成する、エンドツーエンドの例を確認しました。この新しい機能を試すことをお勧めします。これは、リソースに制約のある環境でのデプロイに特に重要な機能です。