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

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

1. モデルにトレーニング後の量子化を適用#

トレーニング後の量子化メソッドを使用してモデルを量子化します。

model = TorchModel() # torch.nn.Module のインスタンス 
model = nncf.quantize(model, ...)

2. モデルの微調整#

このステップでは、ベースライン・モデルに対して適用した方法でモデルに微調整を加えることを前提としています。QAT の場合、10e-5 などの小さな学習率でモデルを数エポックトレーニングする必要があります。量子化モデルは、前方パスと後方パスの両方で量子化エラーをモデル化することにより、微調整中にすべての計算を浮動小数点精度で実行します。

...# データセット、損失、オプティマイザーの設定などの準備の微調整 
# ベースラインとして 5 エポックの量子化モデルを調整 
for epoch in range(0, 5): 
     for i, data in enumerate(train_loader):          ...# トレーニング・ループ本体

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

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

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

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

checkpoint = { 
     'state_dict': model.state_dict(), 
     'nncf_config': model.nncf.get_config(), ...# 保存する残りのユーザー定義オブジェクト 
} 
torch.save(checkpoint, path_to_checkpoint)

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

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

resuming_checkpoint = torch.load(path_to_checkpoint) 
nncf_config = resuming_checkpoint['nncf_config'] 
quantized_model = nncf.torch.load_from_config(model, nncf_config, example_input) 
state_dict = resuming_checkpoint['state_dict'] 
model.load_state_dict(state_dict)

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

必要に応じてモデルを OpenVINO 中間表現 (IR) に変換し、コンパイルして追加の手順なしで OpenVINO で実行できます。

import openvino as ov 

input_fp32 = ...# FP32 モデル入力 

# Torch モデルを OpenVINO モデルに変換 
ov_quantized_model = ov.convert_model(quantized_model, example_input=input_fp32) 

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

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

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

#