pcqat_example.ipynb

Kernel: Python 3

Copyright 2021 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.

TensorFlow.org で表示

Google Colab で実行

GitHub でソースを表示

ノートブックをダウンロード

Keras でのスパース性とクラスタリングを保持する量子化認識トレーニング（PCQAT）の例

概要

このチュートリアルでは、TensorFlow モデル最適化ツールキットの協調最適化パイプラインの一部である**スパース性とクラスタリングを保持する量子化認識トレーニング（PCQAT）**API の使用法を実演します。

その他のページ

パイプラインの概要とその他の利用可能な手法については、協調最適化の概要ページを参照してください。

内容

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

1. MNIST データセットの tf.keras モデルを最初からトレーニングします

2. プルーニングを使用してモデルを微調整し、精度を確認して、モデルが正常にプルーニングされたことを確認します。

3. プルーニングされたモデルにスパース性を保持するクラスタリングを適用し、以前に適用されたスパース性が保持されていることを確認します。

4. QAT を適用し、スパース性とクラスタリングの損失を観察します。

5. PCQAT を適用し、前に適用されたスパース性とクラスタリングが保持されていることを確認します。

6. TFLite モデルを生成し、それに PCQAT を適用した場合の影響を観察します。

7. さまざまなモデルのサイズを比較して、スパース性を適用した後の圧縮の利点を観察し、その後、スパース性を保持するクラスタリングと PCQAT の協調最適化手法を観察します。

8. 完全に最適化されたモデルの精度を、最適化されていないベースラインモデルの精度と比較します。

セットアップ

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

! pip install -q tensorflow-model-optimization

import tensorflow as tf

import numpy as np
import tempfile
import zipfile
import os

MNIST の tf.keras モデルをトレーニングし、プルーニングとクラスタリングを使用する

# Load MNIST dataset
mnist = tf.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

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

opt = tf.keras.optimizers.Adam(learning_rate=1e-3)

# Train the digit classification model
model.compile(optimizer=opt,
              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)

モデルをプルーニングし、50％ のスパース性に微調整する

prune_low_magnitude() API を適用して、次のステップでクラスタリングされるプルーニングされたモデルを作成します。プルーニング API の詳細については、包括的なプルーニングガイドを参照してください。

モデルを定義してスパース性 API を適用する

事前にトレーニングされたモデルが使用されることに注意してください。

import tensorflow_model_optimization as tfmot

prune_low_magnitude = tfmot.sparsity.keras.prune_low_magnitude

pruning_params = {
      'pruning_schedule': tfmot.sparsity.keras.ConstantSparsity(0.5, begin_step=0, frequency=100)
  }

callbacks = [
  tfmot.sparsity.keras.UpdatePruningStep()
]

pruned_model = prune_low_magnitude(model, **pruning_params)

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

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

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

3 エポックのプルーニングでモデルを微調整します。

# Fine-tune model
pruned_model.fit(
  train_images,
  train_labels,
  epochs=3,
  validation_split=0.1,
  callbacks=callbacks)

モデルのスパース性とクラスタリングを計算して出力するヘルパー関数を定義します。

def print_model_weights_sparsity(model):
    for layer in model.layers:
        if isinstance(layer, tf.keras.layers.Wrapper):
            weights = layer.trainable_weights
        else:
            weights = layer.weights
        for weight in weights:
            if "kernel" not in weight.name or "centroid" in weight.name:
                continue
            weight_size = weight.numpy().size
            zero_num = np.count_nonzero(weight == 0)
            print(
                f"{weight.name}: {zero_num/weight_size:.2%} sparsity ",
                f"({zero_num}/{weight_size})",
            )

def print_model_weight_clusters(model):
    for layer in model.layers:
        if isinstance(layer, tf.keras.layers.Wrapper):
            weights = layer.trainable_weights
        else:
            weights = layer.weights
        for weight in weights:
            # ignore auxiliary quantization weights
            if "quantize_layer" in weight.name:
                continue
            if "kernel" in weight.name:
                unique_count = len(np.unique(weight))
                print(
                    f"{layer.name}/{weight.name}: {unique_count} clusters "
                )

最初にプルーニングラッパーを削除し、次にモデルカーネルが正しくプルーニングされたことを確認します。

stripped_pruned_model = tfmot.sparsity.keras.strip_pruning(pruned_model)

print_model_weights_sparsity(stripped_pruned_model)

スパース性を保持するクラスタリングを適用し、両方の場合でモデルのスパース性への影響を確認する

次に、プルーニングされたモデルにスパース性を保持するクラスタリングを適用し、クラスタの数を観察して、スパース性が保持されていることを確認します。

import tensorflow_model_optimization as tfmot
from tensorflow_model_optimization.python.core.clustering.keras.experimental import (
    cluster,
)

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

cluster_weights = cluster.cluster_weights

clustering_params = {
  'number_of_clusters': 8,
  'cluster_centroids_init': CentroidInitialization.KMEANS_PLUS_PLUS,
  'preserve_sparsity': True
}

sparsity_clustered_model = cluster_weights(stripped_pruned_model, **clustering_params)

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

print('Train sparsity preserving clustering model:')
sparsity_clustered_model.fit(train_images, train_labels,epochs=3, validation_split=0.1)

最初にクラスタリングラッパーを削除し、モデルが正しくプルーニングおよびクラスタリングされていることを確認します。

stripped_clustered_model = tfmot.clustering.keras.strip_clustering(sparsity_clustered_model)

print("Model sparsity:\n")
print_model_weights_sparsity(stripped_clustered_model)

print("\nModel clusters:\n")
print_model_weight_clusters(stripped_clustered_model)

QAT と PCQAT を適用し、モデルのクラスタリングとスパース性への影響を確認する

次に、スパースなクラスタリングされたモデルに QAT と PCQAT の両方を適用し、PCQAT がモデル内の重みのスパース性とクラスタリングを保持することを確認します。クラスタリングラッパーが削除されたモデルが QAT および PCQAT API に渡されることに注意してください。

# QAT
qat_model = tfmot.quantization.keras.quantize_model(stripped_clustered_model)

qat_model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
print('Train qat model:')
qat_model.fit(train_images, train_labels, batch_size=128, epochs=1, validation_split=0.1)

# PCQAT
quant_aware_annotate_model = tfmot.quantization.keras.quantize_annotate_model(
              stripped_clustered_model)
pcqat_model = tfmot.quantization.keras.quantize_apply(
              quant_aware_annotate_model,
              tfmot.experimental.combine.Default8BitClusterPreserveQuantizeScheme(preserve_sparsity=True))

pcqat_model.compile(optimizer='adam',
              loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
              metrics=['accuracy'])
print('Train pcqat model:')
pcqat_model.fit(train_images, train_labels, batch_size=128, epochs=1, validation_split=0.1)

print("QAT Model clusters:")
print_model_weight_clusters(qat_model)
print("\nQAT Model sparsity:")
print_model_weights_sparsity(qat_model)
print("\nPCQAT Model clusters:")
print_model_weight_clusters(pcqat_model)
print("\nPCQAT Model sparsity:")
print_model_weights_sparsity(pcqat_model)

PCQAT モデルの圧縮のメリットを確認する

zip 形式のモデルファイルを取得するためのヘルパー関数を定義します。

def get_gzipped_model_size(file):
  # It returns the size of the gzipped model in kilobytes.

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

モデルにスパース性、クラスタリング、および PCQAT を適用すると、圧縮に大きなメリットがもたらされることに注目してください。

# QAT model
converter = tf.lite.TFLiteConverter.from_keras_model(qat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
qat_tflite_model = converter.convert()
qat_model_file = 'qat_model.tflite'
# Save the model.
with open(qat_model_file, 'wb') as f:
    f.write(qat_tflite_model)

# PCQAT model
converter = tf.lite.TFLiteConverter.from_keras_model(pcqat_model)
converter.optimizations = [tf.lite.Optimize.DEFAULT]
pcqat_tflite_model = converter.convert()
pcqat_model_file = 'pcqat_model.tflite'
# Save the model.
with open(pcqat_model_file, 'wb') as f:
    f.write(pcqat_tflite_model)

print("QAT model size: ", get_gzipped_model_size(qat_model_file), ' KB')
print("PCQAT model size: ", get_gzipped_model_size(pcqat_model_file), ' KB')

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(f"Evaluated on {i} results so far.")
    # 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(pcqat_model_file)
interpreter.allocate_tensors()

pcqat_test_accuracy = eval_model(interpreter)

print('Pruned, clustered and quantized TFLite test_accuracy:', pcqat_test_accuracy)
print('Baseline TF test accuracy:', baseline_model_accuracy)

結論

このチュートリアルでは、モデルを作成し、prune_low_magnitude() API を使用してモデルをプルーニングし、cluster_weights() API を使用してスパース性を保持するクラスタリングを適用して重みのクラスタリングを実行する際にスパース性を保持する方法を学習しました。

次に、スパース性とクラスタリングを保持する量子化認識トレーニング (PCQAT) を適用して、QAT を使用する際にモデルのスパース性とクラスタリングを保持しました。最終的な PCQAT モデルを QAT モデルと比較して、スパース性とクラスタリングが前者で保持され、後者で失われることを示しました。

次に、モデルを TFLite に変換して、スパース性の連鎖、クラスタリング、および PCQAT モデルの最適化手法の圧縮のメリットを示し、TFLite モデルを評価して、TFLite バックエンドで精度が維持されることを確認しました。

最後に、PCQAT TFLite モデルの精度を最適化前のベースラインモデルの精度と比較して、元のモデルと同様の精度を維持しながら、協調最適化手法で圧縮のメリットを実現できることを示しました。