NNCF PTQ API と精度制御を使用した音声認識モデルの量子化

この Jupyter ノートブックは、ローカルへのインストール後にのみ起動できます。

このチュートリアルでは、ポストトレーニング・モード (微調整パイプラインなし) で NNCF (ニューラル・ネットワーク圧縮フレームワーク) の精度制御付き 8 ビット量子化を使用して、Wav2Vec2 として知られる音声認識モデルに精度制御付きの INT8 量子化を適用する方法を示します。このノートブックは、LibriSpeech ASR コーパスでトレーニングされ、微調整された Wav2Vec2-Base-960h PyTorch モデルを使用します。チュートリアルは、カスタムモデルとデータセットに拡張できるように設計されています。これは次の手順で構成されます。

• Wav2Vec2 モデルと LibriSpeech データセットをダウンロードして準備します。

• データの読み込みと精度検証の機能を定義します。

• 精度制御によるモデルの量子化を行います。

• 元の PyTorch モデル、OpenVINO FP16 および INT8 モデルの精度を比較します。

• 元のモデルと量子化されたモデルのパフォーマンスを比較します。

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

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

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

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

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

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

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

目次

• インポート

• モデルの準備

• LibriSpeech データセットの準備

• キャリブレーション・データセットの準備

• 評価関数の準備

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

• モデルの使用例

• 元のモデルと量子化されたモデルの精度を比較

%pip install -q "openvino>=2023.1.0"
%pip install -q "nncf>=2.6.0"
%pip install -q --extra-index-url https://download.pytorch.org/whl/cpu soundfile librosa transformers torch datasets torchmetrics

インポート

import numpy as np
import torch

from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor

2023-10-10 09:32:06.465943: I tensorflow/core/util/port.cc:110] oneDNN custom operations are on. You may see slightly different numerical results due to floating-point round-off errors from different computation orders. To turn them off, set the environment variable TF_ENABLE_ONEDNN_OPTS=0.
2023-10-10 09:32:06.505459: I tensorflow/core/platform/cpu_feature_guard.cc:182] This TensorFlow binary is optimized to use available CPU instructions in performance-critical operations.
To enable the following instructions: AVX2 AVX512F AVX512_VNNI FMA, in other operations, rebuild TensorFlow with the appropriate compiler flags.
2023-10-10 09:32:07.113533: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Could not find TensorRT

モデルの準備

PyTorch モデルクラスをインスタンス化するには、Hugging Face Hub からダウンロードするモデル ID を指定して Wav2Vec2ForCTC.from_pretrained メソッドを使用する必要があります。モデルの重みと構成ファイルは、初回使用時に自動的にダウンロードされます。ファイルのダウンロードには数分かかる場合がありますが、インターネット接続環境によって異なります。

さらに、モデル固有の前処理ステップと後処理ステップを担当するプロセッサー・クラスを作成できます。

BATCH_SIZE = 1
MAX_SEQ_LENGTH = 30480


torch_model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-base-960h", ctc_loss_reduction="mean")
processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-base-960h")

Some weights of Wav2Vec2ForCTC were not initialized from the model checkpoint at facebook/wav2vec2-base-960h and are newly initialized: ['wav2vec2.masked_spec_embed']
You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.

OpenVINO 中間表現 (OpenVINO IR) に変換します。

import openvino as ov


default_input = torch.zeros([1, MAX_SEQ_LENGTH], dtype=torch.float)
ov_model = ov.convert_model(torch_model, example_input=default_input)

WARNING:tensorflow:Please fix your imports. Module tensorflow.python.training.tracking.base has been moved to tensorflow.python.trackable.base. The old module will be deleted in version 2.11.

[ WARNING ]  Please fix your imports. Module %s has been moved to %s. The old module will be deleted in version %s.
/home/ea/work/ov_venv/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py:595: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if attn_weights.size() != (bsz * self.num_heads, tgt_len, src_len):
/home/ea/work/ov_venv/lib/python3.8/site-packages/transformers/models/wav2vec2/modeling_wav2vec2.py:634: TracerWarning: Converting a tensor to a Python boolean might cause the trace to be incorrect. We can't record the data flow of Python values, so this value will be treated as a constant in the future. This means that the trace might not generalize to other inputs!
  if attn_output.size() != (bsz * self.num_heads, tgt_len, self.head_dim):

LibriSpeech データセットの準備

デモでモデルの評価をスピードアップするため、LibriSpeech データセットの短いダミーバージョン (patrickvonplaten/librispeech_asr_dummy) を使用します。モデルの精度は論文の報告と異なる場合があります。元の精度を再現するには、librispeech_asr データセットを使用します。

from datasets import load_dataset


dataset = load_dataset("patrickvonplaten/librispeech_asr_dummy", "clean", split="validation")
test_sample = dataset[0]["audio"]


# define preprocessing function for converting audio to input values for model
def map_to_input(batch):
    preprocessed_signal = processor(batch["audio"]["array"], return_tensors="pt", padding="longest", sampling_rate=batch['audio']['sampling_rate'])
    input_values = preprocessed_signal.input_values
    batch['input_values'] = input_values
    return batch


# apply preprocessing function to dataset and remove audio column, to save memory as we do not need it anymore
dataset = dataset.map(map_to_input, batched=False, remove_columns=["audio"])

Found cached dataset librispeech_asr_dummy (/home/ea/.cache/huggingface/datasets/patrickvonplaten___librispeech_asr_dummy/clean/2.1.0/f2c70a4d03ab4410954901bde48c54b85ca1b7f9bf7d616e7e2a72b5ee6ddbfc)
Loading cached processed dataset at /home/ea/.cache/huggingface/datasets/patrickvonplaten___librispeech_asr_dummy/clean/2.1.0/f2c70a4d03ab4410954901bde48c54b85ca1b7f9bf7d616e7e2a72b5ee6ddbfc/cache-dcb48242e67b91b1.arrow

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

import nncf


def transform_fn(data_item):
    """
    Extract the model's input from the data item.
    The data item here is the data item that is returned from the data source per iteration.
    This function should be passed when the data item cannot be used as model's input.
    """
    return np.array(data_item["input_values"])


calibration_dataset = nncf.Dataset(dataset, transform_fn)

INFO:nncf:NNCF initialized successfully. Supported frameworks detected: torch, tensorflow, onnx, openvino

評価関数の準備

評価関数を定義します。

from torchmetrics import WordErrorRate
from tqdm.notebook import tqdm


def validation_fn(model, dataset):
    """
    Calculate and returns a metric for the model.
    """
    wer = WordErrorRate()
    for sample in dataset:
        # run infer function on sample
        output = model.output(0)
        logits = model(np.array(sample['input_values']))[output]
        predicted_ids = np.argmax(logits, axis=-1)
        transcription = processor.batch_decode(torch.from_numpy(predicted_ids))

        # update metric on sample result
        wer.update(transcription, [sample['text']])

    result = wer.compute()

    return 1 - result

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

キャリブレーション・データセットと評価データセットを提供する必要があります。同じデータセットでもかまいません。

• パラメーター max_drop は精度低下のしきい値を定義します。量子化プロセスは、評価データセットの精度メトリックの低下が max_drop 未満になると停止します。デフォルト値は 0.01 です。max_drop 値に到達しない場合、NNCF は量子化を停止し、エラーを報告します。
• drop_type は、精度の低下の計算方法を定義します: ABSOLUTE (デフォルト) または RELATIVE。
• ranking_subset_size は、精度低下への影響度に基づいてレイヤーをランク付けするために使用されるサブセットのサイズです。デフォルト値は 300 で、サンプル数が多いほどランクが向上する可能性があります。ここでは、実行を高速化するため値 25 を使用します。

注: 実行には数十分かかる場合があり、最大 10 GB の空きメモリーが必要です。

from nncf.quantization.advanced_parameters import AdvancedAccuracyRestorerParameters
from nncf.parameters import ModelType

quantized_model = nncf.quantize_with_accuracy_control(
    ov_model,
    calibration_dataset=calibration_dataset,
    validation_dataset=calibration_dataset,
    validation_fn=validation_fn,
    max_drop=0.01,
    drop_type=nncf.DropType.ABSOLUTE,
    model_type=ModelType.TRANSFORMER,
    advanced_accuracy_restorer_parameters=AdvancedAccuracyRestorerParameters(
        ranking_subset_size=25
    ),
)

Statistics collection:  24%|███████████████████████████████████▎                                                                                                             | 73/300 [00:12<00:37,  5.98it/s]
Applying Smooth Quant: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 50/50 [00:01<00:00, 41.01it/s]

INFO:nncf:36 ignored nodes was found by name in the NNCFGraph

Statistics collection:  24%|███████████████████████████████████▎                                                                                                             | 73/300 [00:22<01:08,  3.31it/s]
Applying Fast Bias correction: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 74/74 [00:23<00:00,  3.09it/s]

INFO:nncf:Validation of initial model was started

INFO:nncf:Elapsed Time: 00:00:00
INFO:nncf:Elapsed Time: 00:00:11
INFO:nncf:Metric of initial model: 0.9469565153121948
INFO:nncf:Collecting values for each data item using the initial model
INFO:nncf:Elapsed Time: 00:00:09
INFO:nncf:Validation of quantized model was started
INFO:nncf:Elapsed Time: 00:00:22
INFO:nncf:Elapsed Time: 00:00:11
INFO:nncf:Metric of quantized model: 0.5
INFO:nncf:Collecting values for each data item using the quantized model
INFO:nncf:Elapsed Time: 00:00:06
INFO:nncf:Accuracy drop: 0.4469565153121948 (DropType.ABSOLUTE)
INFO:nncf:Accuracy drop: 0.4469565153121948 (DropType.ABSOLUTE)
INFO:nncf:Total number of quantized operations in the model: 94
INFO:nncf:Number of parallel processes to rank quantized operations: 14
INFO:nncf:ORIGINAL metric is used to rank quantizers
INFO:nncf:Calculating ranking score for groups of quantizers
INFO:nncf:Elapsed Time: 00:04:36
INFO:nncf:Changing the scope of quantizer nodes was started
INFO:nncf:Reverted 1 operations to the floating-point precision:
    __module.wav2vec2.feature_extractor.conv_layers.2.conv/aten::_convolution/Convolution_11
INFO:nncf:Accuracy drop with the new quantization scope is 0.06173914670944214 (DropType.ABSOLUTE)
INFO:nncf:Reverted 1 operations to the floating-point precision:
    __module.wav2vec2.feature_extractor.conv_layers.1.conv/aten::_convolution/Convolution_10
INFO:nncf:Accuracy drop with the new quantization scope is 0.010434746742248535 (DropType.ABSOLUTE)
INFO:nncf:Reverted 1 operations to the floating-point precision:
    __module.wav2vec2.feature_extractor.conv_layers.3.conv/aten::_convolution/Convolution_12
INFO:nncf:Algorithm completed: achieved required accuracy drop 0.006956517696380615 (DropType.ABSOLUTE)
INFO:nncf:3 out of 94 were reverted back to the floating-point precision:
    __module.wav2vec2.feature_extractor.conv_layers.2.conv/aten::_convolution/Convolution_11
    __module.wav2vec2.feature_extractor.conv_layers.1.conv/aten::_convolution/Convolution_10
    __module.wav2vec2.feature_extractor.conv_layers.3.conv/aten::_convolution/Convolution_12

モデルの使用例

import IPython.display as ipd


ipd.Audio(test_sample["array"], rate=16000)

core = ov.Core()

compiled_quantized_model = core.compile_model(model=quantized_model, device_name='CPU')

input_data = np.expand_dims(test_sample["array"], axis=0)

次に、推測を行います。

predictions = compiled_quantized_model([input_data])[0]
predicted_ids = np.argmax(predictions, axis=-1)
transcription = processor.batch_decode(torch.from_numpy(predicted_ids))
transcription

['I E O WE WORD I O O FAGGI  FARE E BO']

元のモデルと量子化されたモデルの精度を比較

• テスト・データセットのデータローダーを定義します。

• PyTorch および OpenVINO モデルの推論を取得する関数を定義します。

• 単語誤り率を計算する関数を定義します。

# inference function for pytorch
def torch_infer(model, sample):
    logits = model(torch.Tensor(sample['input_values'])).logits
    # take argmax and decode
    predicted_ids = torch.argmax(logits, dim=-1)
    transcription = processor.batch_decode(predicted_ids)
    return transcription


# inference function for openvino
def ov_infer(model, sample):
    output = model.output(0)
    logits = model(np.array(sample['input_values']))[output]
    predicted_ids = np.argmax(logits, axis=-1)
    transcription = processor.batch_decode(torch.from_numpy(predicted_ids))
    return transcription


def compute_wer(dataset, model, infer_fn):
    wer = WordErrorRate()
    for sample in tqdm(dataset):
        # run infer function on sample
        transcription = infer_fn(model, sample)
        # update metric on sample result
        wer.update(transcription, [sample['text']])
    # finalize metric calculation
    result = wer.compute()
    return result

ここで、元の PyTorch モデルと量子化モデルの WER を計算します。

pt_result = compute_wer(dataset, torch_model, torch_infer)
quantized_result = compute_wer(dataset, compiled_quantized_model, ov_infer)

print(f'[PyTorch]   Word Error Rate: {pt_result:.4f}')
print(f'[Quantized OpenVINO]  Word Error Rate: {quantized_result:.4f}')

0%|          | 0/73 [00:00<?, ?it/s]

0%|          | 0/73 [00:00<?, ?it/s]

[PyTorch]   Word Error Rate: 0.0530
[Quantized OpenVINO]  Word Error Rate: 0.0600