TensorFlow による量子化対応トレーニング (QAT)#

以下は、NNCF の QAT を TensorFlow で作成されたトレーニング・スクリプトに統合する手順を示します:

注

現在、NNCF for TensorFlow は、Keras Sequential API または Functional API を使用して作成されたモデルの最適化をサポートしています。

1. NNCF API をインポート#

トレーニング・スクリプトの先頭に NNCF 関連のインポート文を追加します:

import tensorflow as tf 

from nncf import NNCFConfig 
from nncf.tensorflow import create_compressed_model, register_default_init_args

2. NNCF 構成の作成#

モデル関連のパラメーター ("input_info" セクション) と最適化メソッドのパラメーター ("compression" セクション) で構成される NNCF 構成を定義する必要があります。収束を高速化するため、DL フレームワークに固有のデータセット・オブジェクトを登録することも推奨されます。データ・オブジェクトは、モデル作成ステップで量子化パラメーターを初期化する際に使用されます。

nncf_config_dict = { 
    "input_info": {"sample_size": [1, 3, 224, 224]}, # モデルのトレースに必要な入力形状 
    "compression": { "algorithm": "quantization", # デフォルト設定での 8 ビット量子化 
    }, 
} 
nncf_config = NNCFConfig.from_dict(nncf_config_dict) 
nncf_config = register_default_init_args(nncf_config, train_dataset, batch_size=1) # train_dataset は tf.data.Dataset のインスタンス

3. 最適化の適用#

前のステップで定義した構成を使用して、create_compressed_model() API で元のモデル・オブジェクトをラップする必要があります。このメソッドは、圧縮コントローラーと、元のモデルと同じように使用できるラップされたモデルを返します。このステップでは最適化が適用され、モデルは対応する一連の変換を受け、最適化に必要な追加の操作が含まれます。QAT の場合、圧縮コントローラー・オブジェクトはモデルのエクスポートに使用され、オプションで以下に示すように分散トレーニングにも使用されます。

model = KerasModel() # tensorflow.keras.Model のインスタンス 
compression_ctrl, model = create_compressed_model(model, nncf_config)

4. モデルの微調整#

このステップでは、ベースライン・モデルに対して適用した方法でモデルに微調整を加えることを前提としています。QAT の場合、10^e-5 などの小さな学習率でモデルを数エポックトレーニングする必要があります。原則として、このステップはスキップできます。これは、トレーニング後の最適化がモデルに適用されるためです。

...# #微調整の準備、例: データセット、損失、オプティマイザーのセットアップなど 

# 最適化パラメーターを制御し、圧縮統計をダンプするための圧縮コールバックを作成 
compression_callbacks = create_compression_callbacks(compression_ctrl, log_dir="./compression_log") 
# ベースラインと同じ方法で 5 エポックの量子化モデルを調整 
model.fit(train_dataset, epochs=5, callbacks=compression_callbacks)

5. マルチ GPU 分散トレーニング#

分散マルチ GPU トレーニング (DataParallel ではない) の場合は、微調整する前に compression_ctrl.distributed() を呼び出します。これにより、分散モードで機能するように調整する最適化が通知されます。

compression_ctrl.distributed() # トレーニングの前に呼び出し

注

重みの精度は、モデルを OpenVINO 中間表現に変換した後にのみ INT8 に移行します。その形式でのみ、モデルのフットプリントの削減が期待できます。

これらの手順は、NNCF の QAT メソッドを適用する基本的な手順です。ただし、状況によっては、トレーニング中にモデルのチェックポイントを保存/復元する必要があります。NNCF は元のモデルを独自のオブジェクトでラップするため、これらのニーズに対応する API を提供します。

6. (オプション) チェックポイントの保存#

モデルのチェックポイントを保存するには、次の API を使用します:

from nncf.tensorflow.utils.state import TFCompressionState 
from nncf.tensorflow.callbacks.checkpoint_callback import CheckpointManagerCallback 

checkpoint = tf.train.Checkpoint(model=model, 
                                compression_state=TFCompressionState(compression_ctrl), 
                                ...# 保存する残りのユーザー定義オブジェクト 
                               ) 
callbacks = [] 
callbacks.append(CheckpointManagerCallback(checkpoint, path_to_checkpoint)) 
... 
model.fit(..., callbacks=callbacks)

7. (オプション) チェックポイントから復元#

チェックポイントからモデルを復元ツールを使用は、次の API を使用します:

from nncf.tensorflow.utils.state import TFCompressionStateLoader 

checkpoint = tf.train.Checkpoint(compression_state=TFCompressionStateLoader()) 
checkpoint.restore(path_to_checkpoint) 
compression_state = checkpoint.compression_state.state 

compression_ctrl, model = create_compressed_model(model, nncf_config, compression_state) 
checkpoint = tf.train.Checkpoint(model=model, ...) 

checkpoint.restore(path_to_checkpoint)

NNCF でのチェックポイントの保存/復元の詳細については、対応する NNCF ドキュメントを参照してください。

量子化モデルのデプロイ#

必要に応じてモデルを OpenVINO 中間表現 (IR) に変換し、OpenVINO でコンパイルして実行できます。追加の手順やオプションは必要ありません。

import openvino as ov 

# TensorFlow モデルを OpenVINO モデルに変換 
ov_quantized_model = ov.convert_model(quantized_model) 

# モデルをコンパイルして量子化演算を int8 に変換 
model_int8 = ov.compile_model(ov_quantized_model) 

input_fp32 = ...# FP32 モデル入力 
res = model_int8(input_fp32) 

# モデルを保存 
ov.save_model(ov_quantized_model, "quantized_model.xml")

詳細は、対応するドキュメントをご覧ください。

例#

NNCF を使用した TensorFlow モデルの量子化