clustering_example.ipynb

Kernel: Python 3

Copyright 2020 The TensorFlow Authors.

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

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
)

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

_, 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)

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

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()

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

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

# Fine-tune model
clustered_model.fit(
  train_images,
  train_labels,
  batch_size=500,
  epochs=1,
  validation_split=0.1)

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

_, 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)

クラスタリングによって 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)

次に、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)

実際に 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)))

重みクラスタリングとポストトレーニング量子化を組み合わせて、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)))

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)

まとめ

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