精度制御を備えた量子化

はじめに

これは、精度メトリックを制御してモデルに 8 ビット量子化を適用する高度な量子化フローです。モデル内で最も影響力のある操作を元の精度に維持することで実現されます。このフローは基本的な 8 ビット量子化に基づいており、次のような違いがあります。

  • 精度メトリックを計算するには、キャリブレーション・データセットと評価データセットが必要です。最も単純なケースでは、両方のデータセットが同じデータを参照できます。

  • 評価関数は、精度メトリックの計算に必要です。ソース・フレームワーク使用可能な関数またはカスタム関数を使用できます。

  • 精度の検証は量子化プロセス中に数回実行されるため、精度制御による量子化は基本 8 ビット量子化フローよりも時間がかかる場合があります。

  • 一部の演算が元の精度で維持されるため、結果のモデルでは、基本 8 ビット量子化フローよりもわずかですがパフォーマンスが向上します。

現在、精度制御を備えた 8 ビット量子化は、OpenVINO および onnx.ModelProto 表現のモデルでのみ使用できます。

精度制御を備えた量子化の手順を以下に説明します。

モデルの準備

FP32 精度の元のモデルを使用する場合は、重みを圧縮せずにモデルをそのまま、精度制御による量子化の入力として使用することを推奨します。これにより、特定の精度の低下に対して最適なパフォーマンスが保証されます。元のモデルの重みを FP16 に圧縮するなど圧縮技術を利用すると、復元されるレイヤー数が大幅に増加し、量子化モデルのパフォーマンスの低下につながる可能性があります。元のモデルが OpenVINO に変換され、精度制御を備えた量子化メソッドで使用する前に openvino.save_model() で保存される場合、compress_to_fp16=False を設定して FP16 への重みの圧縮を無効にします。デフォルトでは、openvino.save_model() はモデルを FP16 に保存するため、この無効化が必要です。

キャリブレーションと評価データセットの準備

このステップは、基本的な 8 ビット量子化フローに似ています。唯一の違いは、キャリブレーションと評価という 2 つのデータセットが必要であることです。

import nncf
import torch

calibration_loader = torch.utils.data.DataLoader(...)

def transform_fn(data_item):
    images, _ = data_item
    return images

calibration_dataset = nncf.Dataset(calibration_loader, transform_fn)
validation_dataset = nncf.Dataset(calibration_loader, transform_fn)
import nncf
import torch

calibration_loader = torch.utils.data.DataLoader(...)

def transform_fn(data_item):
    images, _ = data_item
    return {input_name: images.numpy()} # input_name should be taken from the model,
                                        # e.g. model.graph.input[0].name

calibration_dataset = nncf.Dataset(calibration_loader, transform_fn)
validation_dataset = nncf.Dataset(calibration_loader, transform_fn)

評価関数の準備

評価関数は、モデル・オブジェクトと評価データセットを受け取り、精度メトリック値を返します。モデル・オブジェクトのタイプはフレームワークによって異なります。OpenVINO では、openvino.CompiledModel です。ONNX では、onnx.ModelProto です。次のコードは、OpenVINO と ONNX フレームワーク向けの評価関数の例を示しています。

import numpy as np
import torch
from sklearn.metrics import accuracy_score

import openvino


def validate(model: openvino.CompiledModel, 
             validation_loader: torch.utils.data.DataLoader) -> float:
    predictions = []
    references = []

    output = model.outputs[0]

    for images, target in validation_loader:
        pred = model(images)[output]
        predictions.append(np.argmax(pred, axis=1))
        references.append(target)

    predictions = np.concatenate(predictions, axis=0)
    references = np.concatenate(references, axis=0)
    return accuracy_score(predictions, references)
import numpy as np
import torch
from sklearn.metrics import accuracy_score

import onnx
import onnxruntime


def validate(model: onnx.ModelProto,
             validation_loader: torch.utils.data.DataLoader) -> float:
    predictions = []
    references = []

    input_name = model.graph.input[0].name
    serialized_model = model.SerializeToString()
    session = onnxruntime.InferenceSession(serialized_model, providers=["CPUExecutionProvider"])
    output_names = [output.name for output in session.get_outputs()]

    for images, target in validation_loader:
        pred = session.run(output_names, input_feed={input_name: images.numpy()})[0]
        predictions.append(np.argmax(pred, axis=1))
        references.append(target)

    predictions = np.concatenate(predictions, axis=0)
    references = np.concatenate(references, axis=0)
    return accuracy_score(predictions, references)

精度を制御して量子化を実行

nncf.quantize_with_accuracy_control() 関数は、精度制御を使用して量子化を実行します。次のコードは、OpenVINO と ONNX フレームワークの精度制御による量子化の例を示しています。

model = ... # openvino.Model object

quantized_model = nncf.quantize_with_accuracy_control(
    model,
    calibration_dataset=calibration_dataset,
    validation_dataset=validation_dataset,
    validation_fn=validate,
    max_drop=0.01,
    drop_type=nncf.DropType.ABSOLUTE,
)
import onnx

model = onnx.load("model_path")

quantized_model = nncf.quantize_with_accuracy_control(
    model,
    calibration_dataset=calibration_dataset,
    validation_dataset=validation_dataset,
    validation_fn=validate,
    max_drop=0.01,
    drop_type=nncf.DropType.ABSOLUTE,
)
  • max_drop は精度低下のしきい値を定義します。量子化プロセスは、評価データセットの精度メトリックの低下が max_drop 未満になると停止します。デフォルト値は 0.01 です。max_drop 値に到達しない場合、NNCF は量子化を停止し、エラーを報告します。

  • drop_type は、精度の低下の計算方法を定義します: ABSOLUTE (デフォルト) または RELATIVE

その後、モデルをコンパイルして OpenVINO で実行できます。

import openvino as ov

# compile the model to transform quantized operations to int8
model_int8 = ov.compile_model(quantized_model)

input_fp32 = ... # FP32 model input
res = model_int8(input_fp32)
import openvino as ov

# convert ONNX model to OpenVINO model
ov_quantized_model = ov.convert_model(quantized_model)

# compile the model to transform quantized operations to int8
model_int8 = ov.compile_model(ov_quantized_model)

input_fp32 = ... # FP32 model input
res = model_int8(input_fp32)

モデルを OpenVINO 中間表現 (IR) で保存するには、openvino.save_model() を使用します。FP32 精度で元のモデルを扱う場合、INT8 から FP32 に戻された最も影響力のあるモデル操作で FP32 精度を維持することを推奨します。これには、保存プロセス中に compress_to_fp16=False を使用することを検討してください。この推奨事項は、モデルを FP16 に保存する openvino.save_model() のデフォルト動作に基づいており、この変換によって精度に影響を与える可能性があります。

# save the model with compress_to_fp16=False to avoid an accuracy drop from compression
# of unquantized weights to FP16. This is necessary because
# nncf.quantize_with_accuracy_control(...) keeps the most impactful operations within
# the model in the original precision to achieve the specified model accuracy
ov.save_model(quantized_model, "quantized_model.xml", compress_to_fp16=False)

nncf.quantize_with_accuracy_control() API は、精度制御とカスタム構成を使用してモデルを量子化するため、基本 8 ビット量子化 API のすべてのパラメーターをサポートします。

量子化モデルの精度またはパフォーマンスに満足できない場合は、次のステップとしてトレーニング時の最適化があります。

精度メトリックを制御した NNCF トレーニング後の量子化の例