comprehensive_guide.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.

プルーニングの総合ガイド

TensorFlow.org で表示

Google Colab で実行

GitHubでソースを表示

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

Keras 重みプルーニングの総合ガイドへようこそ。

このページでは、さまざまなユースケースを示し、それぞれで API を使用する方法を説明します。どの API が必要であるかを特定したら、API ドキュメントでパラメータと詳細を確認してください。

• プルーニングのメリットとサポート対象を確認する場合は、概要をご覧ください。

• 単一のエンドツーエンドの例については、プルーニングの例をご覧ください。

次のユースケースについて説明しています。

• プルーニングされたモデルの定義とトレーニング

  • Sequential と Functional

  • Keras model.fit とカスタムトレーニングループ

• プルーニングされたモデルのチェックポイントと逆シリアル化

• プルーニングされたモデルのデプロイと圧縮のメリットの確認

プルーニングアルゴリズムの構成については、tfmot.sparsity.keras.prune_low_magnitude API ドキュメントをご覧ください。

セットアップ

必要な API の特定と目的の理解については、次を実行できますが、このセクションを読まずに進むことができます。

! pip install -q tensorflow-model-optimization

import tensorflow as tf
import numpy as np
import tensorflow_model_optimization as tfmot

%load_ext tensorboard

import tempfile

input_shape = [20]
x_train = np.random.randn(1, 20).astype(np.float32)
y_train = tf.keras.utils.to_categorical(np.random.randn(1), num_classes=20)

def setup_model():
  model = tf.keras.Sequential([
      tf.keras.layers.Dense(20, input_shape=input_shape),
      tf.keras.layers.Flatten()
  ])
  return model

def setup_pretrained_weights():
  model = setup_model()

  model.compile(
      loss=tf.keras.losses.categorical_crossentropy,
      optimizer='adam',
      metrics=['accuracy']
  )

  model.fit(x_train, y_train)

  _, pretrained_weights = tempfile.mkstemp('.tf')

  model.save_weights(pretrained_weights)

  return pretrained_weights

def get_gzipped_model_size(model):
  # Returns size of gzipped model, in bytes.
  import os
  import zipfile

  _, keras_file = tempfile.mkstemp('.h5')
  model.save(keras_file, include_optimizer=False)

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

  return os.path.getsize(zipped_file)

setup_model()
pretrained_weights = setup_pretrained_weights()

モデルを定義する

すべてのモデルをプルーニングする（Sequential と Functional）

モデルの精度を高めるためのヒント:

• 「一部のレイヤーをプルーニングする」を試して、最も精度を低下させるレイヤーのプルーニングを省略します。

• 一般的に、始めからトレーニングするよりも、プルーニングで微調整する方が優れています。

プルーニングでモデル全体をトレーニングするには、モデルに tfmot.sparsity.keras.prune_low_magnitude を適用します。

base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended.

model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model)

model_for_pruning.summary()

一部のレイヤーをプルーニングする（Sequential と Functional）

モデルのプルーニングによって制度に悪影響が及ぶことがあります。そのため、選択的にモデルのレイヤーをプルーニングすることで、精度、速度、およびモデルサイズ間のトレードオフを探ることができます。

モデルの精度を高めるためのヒント:

• 一般的に、始めからトレーニングするよりも、プルーニングで微調整する方が優れています。

• 初期のレイヤーの代わりに後のレイヤーのプルーニングを試します。

• クリティカルレイヤー（注意メカニズムなど）のプルーニングを回避します。

その他:

• tfmot.sparsity.keras.prune_low_magnitude API ドキュメントには、レイヤーごとにプルーニング構成を変える方法が示されています。

次の例では、Dense レイヤーのみをプルーニングしています。

# Create a base model
base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended for model accuracy

# Helper function uses `prune_low_magnitude` to make only the 
# Dense layers train with pruning.
def apply_pruning_to_dense(layer):
  if isinstance(layer, tf.keras.layers.Dense):
    return tfmot.sparsity.keras.prune_low_magnitude(layer)
  return layer

# Use `tf.keras.models.clone_model` to apply `apply_pruning_to_dense` 
# to the layers of the model.
model_for_pruning = tf.keras.models.clone_model(
    base_model,
    clone_function=apply_pruning_to_dense,
)

model_for_pruning.summary()

この例ではプルーニングするものを決定するためにレイヤーの種類が使用されていますが、特定のレイヤーをプルーニングする上で最も簡単な方法は、name プロパティを設定し、clone_function でその名前を探す方法です。

print(base_model.layers[0].name)

可読性が高くても、モデルの精度を潜在的に低下させる

これは、プルーニングによる微調整とは使用できません。そのため、微調整をサポートする上記の例よりも制度に劣る可能性があります。

初期のモデルを定義する際に prune_low_magnitude を適用することも可能ですが、後で重みを読み込む場合、以下の例が機能しません。

Functional の例

# Use `prune_low_magnitude` to make the `Dense` layer train with pruning.
i = tf.keras.Input(shape=(20,))
x = tfmot.sparsity.keras.prune_low_magnitude(tf.keras.layers.Dense(10))(i)
o = tf.keras.layers.Flatten()(x)
model_for_pruning = tf.keras.Model(inputs=i, outputs=o)

model_for_pruning.summary()

Sequential の例

# Use `prune_low_magnitude` to make the `Dense` layer train with pruning.
model_for_pruning = tf.keras.Sequential([
  tfmot.sparsity.keras.prune_low_magnitude(tf.keras.layers.Dense(20, input_shape=input_shape)),
  tf.keras.layers.Flatten()
])

model_for_pruning.summary()

カスタム Keras レイヤーのプルーニングまたはプルーニングするレイヤーの部分の変更

一般的な過ち: バイアスをプルーニングすると通常、モデルの精度を著しく悪化させてしまう。

tfmot.sparsity.keras.PrunableLayer は、2 つのユースケースに役立ちます。

1. カスタム Keras レイヤーのプルーニング

2. プルーニングする組み込み Keras レイヤーの部分の変更

たとえば、API はデフォルトで Dense レイヤーのカーネルのプルーニングのみを行います。以下の例ではバイアスもプルーニングします。

class MyDenseLayer(tf.keras.layers.Dense, tfmot.sparsity.keras.PrunableLayer):

  def get_prunable_weights(self):
    # Prune bias also, though that usually harms model accuracy too much.
    return [self.kernel, self.bias]

# Use `prune_low_magnitude` to make the `MyDenseLayer` layer train with pruning.
model_for_pruning = tf.keras.Sequential([
  tfmot.sparsity.keras.prune_low_magnitude(MyDenseLayer(20, input_shape=input_shape)),
  tf.keras.layers.Flatten()
])

model_for_pruning.summary()

モデルのトレーニング

Model.fit

トレーニング中に tfmot.sparsity.keras.UpdatePruningStep コールバックを呼び出します。

トレーニングをデバッグできるようにするには、tfmot.sparsity.keras.PruningSummaries コールバックを使用します。

# Define the model.
base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended for model accuracy
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model)

log_dir = tempfile.mkdtemp()
callbacks = [
    tfmot.sparsity.keras.UpdatePruningStep(),
    # Log sparsity and other metrics in Tensorboard.
    tfmot.sparsity.keras.PruningSummaries(log_dir=log_dir)
]

model_for_pruning.compile(
      loss=tf.keras.losses.categorical_crossentropy,
      optimizer='adam',
      metrics=['accuracy']
)

model_for_pruning.fit(
    x_train,
    y_train,
    callbacks=callbacks,
    epochs=2,
)

#docs_infra: no_execute
%tensorboard --logdir={log_dir}

Colab を使用していないユーザーは、TensorBoard.dev で、このノートブックの前回の実行結果を閲覧できます。

カスタムトレーニングループ

トレーニング中に tfmot.sparsity.keras.UpdatePruningStep コールバックを呼び出します。

トレーニングをデバッグできるようにするには、tfmot.sparsity.keras.PruningSummaries コールバックを使用します。

# Define the model.
base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended for model accuracy
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model)

# Boilerplate
loss = tf.keras.losses.categorical_crossentropy
optimizer = tf.keras.optimizers.Adam()
log_dir = tempfile.mkdtemp()
unused_arg = -1
epochs = 2
batches = 1 # example is hardcoded so that the number of batches cannot change.

# Non-boilerplate.
model_for_pruning.optimizer = optimizer
step_callback = tfmot.sparsity.keras.UpdatePruningStep()
step_callback.set_model(model_for_pruning)
log_callback = tfmot.sparsity.keras.PruningSummaries(log_dir=log_dir) # Log sparsity and other metrics in Tensorboard.
log_callback.set_model(model_for_pruning)

step_callback.on_train_begin() # run pruning callback
for _ in range(epochs):
  log_callback.on_epoch_begin(epoch=unused_arg) # run pruning callback
  for _ in range(batches):
    step_callback.on_train_batch_begin(batch=unused_arg) # run pruning callback

    with tf.GradientTape() as tape:
      logits = model_for_pruning(x_train, training=True)
      loss_value = loss(y_train, logits)
      grads = tape.gradient(loss_value, model_for_pruning.trainable_variables)
      optimizer.apply_gradients(zip(grads, model_for_pruning.trainable_variables))

  step_callback.on_epoch_end(batch=unused_arg) # run pruning callback

#docs_infra: no_execute
%tensorboard --logdir={log_dir}

Colab を使用していないユーザーは、TensorBoard.dev で、このノートブックの前回の実行結果を閲覧できます。

プルーニングされたモデルの精度を改善する

まず、tfmot.sparsity.keras.prune_low_magnitude API ドキュメントを確認し、プルーニングスケジュールが何か、また各種プルーニングスケジュールの計算について理解してください。

ヒント:

• モデルがプルーニングする際に、高すぎず低すぎない学習率を使用します。プルーニングスケジュールをハイパーパラメータとすることを検討してください。

• 簡易テストとして、トレーニング始めに最終的なスパース性までモデルを実験的にプルーニングします。tfmot.sparsity.keras.ConstantSparsity スケジュールで begin_step を 0 に設定します。良い結果が得る可能性があります。

• モデルに回復する時間を与えられうように、あまり頻繁にプルーニングしないようにします。プルーニングスケジュールには十分なデフォルトの頻度が指定されています。

• モデルの精度を改善するための一般的なアイデアについては、「モデルを定義する」に記載のケース別のヒントをご覧ください。

チェックポイントと逆シリアル化

チェックポイント作成時には、オプティマイザのステップを保持する必要があります。つまり、チェックポイント作成に Keras HDF5 モデルを使用することはできますが、Keras HDF5 の重みは使用できません。

# Define the model.
base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended for model accuracy
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model)

_, keras_model_file = tempfile.mkstemp('.h5')

# Checkpoint: saving the optimizer is necessary (include_optimizer=True is the default).
model_for_pruning.save(keras_model_file, include_optimizer=True)

上記は一般的に適用されます。次のコードは、HDF5 モデル形式のみで必要です（HDF5 重みまたはその他の形式では不要です）。

# Deserialize model.
with tfmot.sparsity.keras.prune_scope():
  loaded_model = tf.keras.models.load_model(keras_model_file)

loaded_model.summary()

プルーニングされたモデルをデプロイする

サイズ圧縮によるモデルのエクスポート

一般的な過ち: strip_pruning と標準圧縮アルゴリズム（gzip など）の適用は、プルーニングの圧縮のメリットを確認する上で必要です。

# Define the model.
base_model = setup_model()
base_model.load_weights(pretrained_weights) # optional but recommended for model accuracy
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model)

# Typically you train the model here.

model_for_export = tfmot.sparsity.keras.strip_pruning(model_for_pruning)

print("final model")
model_for_export.summary()

print("\n")
print("Size of gzipped pruned model without stripping: %.2f bytes" % (get_gzipped_model_size(model_for_pruning)))
print("Size of gzipped pruned model with stripping: %.2f bytes" % (get_gzipped_model_size(model_for_export)))

ハードウェア固有の最適化

異なるバックエンドでプルーニングによるレイテンシの改善が可能になったら、ブロックのスパース性を使用することで、特定のハードウェアのレイテンシを改善することができます。

ブロックサイズを大きくすると、ターゲットモデルの精度を得るために達成可能なピークスパース性が低下します。これにも関わらず、レイテンシは改善されることがあります。

ブロックスパース性のサポート状況に関する詳細は、tfmot.sparsity.keras.prune_low_magnitude API ドキュメントをご覧ください。

base_model = setup_model()

# For using intrinsics on a CPU with 128-bit registers, together with 8-bit
# quantized weights, a 1x16 block size is nice because the block perfectly
# fits into the register.
pruning_params = {'block_size': [1, 16]}
model_for_pruning = tfmot.sparsity.keras.prune_low_magnitude(base_model, **pruning_params)

model_for_pruning.summary()