DNNモデル圧縮ツール「NNCF」で圧縮したモデルをAE2100で動かしてみよう(2) ―圧縮編―

## はじめに

OpenVINO Toolkitの関連ツールであるNeural Network Compression Framework (NNCF)について2本の記事にわたり説明します。
この記事シリーズはNNCFによりDNNモデルを圧縮し、AE2100での推論スループットを向上させることを目標としています。

1. [準備編](https://qiita.com/sa9moto/items/0bf908cef741636d504a)
2. 圧縮編(この記事)

今回の記事は圧縮編として、NNCFを利用して学習済みモデルの圧縮を行います。
またAE2100を使用して圧縮済みモデルの動作確認を行います。

圧縮対象のモデルは[前回の記事](https://qiita.com/sa9moto/items/0bf908cef741636d504a)にて作成していますので、この記事のプログラムを実行する前に実施してください。

なお本記事ではOpenVINO Toolkitのバージョンとして2022.1を利用します。
またAE2100での動作確認にあたり「AE2100 標準コンテナイメージ（OpenVINO有（Ubuntu20.04版）） 2022年10月版」を使用しています。

## モデル圧縮

NNCFを利用してモデルの圧縮を行います。
以下ではNNCF APIを利用した圧縮プログラムと圧縮に関する設定を記述したNNCF設定ファイルを作成し、モデルを圧縮します。
また、圧縮したモデルに対して後ほどOpenVINO ToolkitのModel Optimizerを用いるために、圧縮したモデルをSavedModel形式として出力します。

### NNCF APIの利用方法

NNCFを利用した圧縮プログラムを作成するには、NNCFのPythonパッケージとしてインポートします。
[公式ドキュメント](https://github.com/openvinotoolkit/nncf/tree/release_v230#usage)に倣うと、TensorFlowとNNCFを利用した圧縮プログラムは下記の流れで作成します。

``` python
import tensorflow as tf
from nncf import NNCFConfig
from nncf.tensorflow import register_default_init_args
from nncf.tensorflow import create_compressed_model

# データセットを読み込む
dataset = ...
# 圧縮対象のモデルを読み込む
model = ...

# 圧縮のパラメータを記述した設定ファイル(JSON形式)を読み込む
nncf_config = NNCFConfig.from_json(...)
# データセットを引数として与えて圧縮アルゴリズムを初期化する
nncf_config = register_default_init_args(nncf_config, dataset, batch_size=1)

# 圧縮アルゴリズムをモデルに適用する
compression_ctrl, compressed_model = create_compressed_model(model, nncf_config)

# compressed_modelの圧縮を開始する
compressed_model.fit(...)

# 圧縮済みモデルをTensorFlowのモデル形式(Frozen Graph, TensorFlow SavedModel, .h5)にて出力する
compression_ctrl.export_model(model_path, save_format=...)
```

圧縮プログラムは学習プログラムに対して主に下記の処理を追加することで作成します。

- `nncf.NNCFConfig.from_json()`
- `nncf.tensorflow.register_default_init_args()`
- `nncf.tensorflow.create_compressed_model()`

`nncf.NNCFConfig.from_json()`ではJSON形式で圧縮に関するハイパーパラメータを記述したNNCF設定ファイルを読み込みます。
NNCF設定ファイルの作成方法は次節にて説明します。
次に`nncf.tensorflow.register_default_init_args()`にデータセットを引数として与え、圧縮アルゴリズムの初期化を行います。
その後`nncf.tensorflow.create_compressed_model()`に圧縮対象のモデルとNNCFの設定を引数として与えることで、圧縮アルゴリズムをモデルに適用します。
以上により得たモデルの`.fit(...)`を用いることで学習を行いながらモデルの圧縮を行います。

### NNCF設定ファイル

モデルの圧縮に関するパラメータをJSON形式で記述したNNCF設定ファイルを作成します。
NNCF設定ファイルを作成するには[公式ドキュメント](https://github.com/openvinotoolkit/nncf/blob/release_v230/docs/compression_algorithms/Pruning.md#pruning)や[NNCF設定ファイルのサンプル](https://github.com/openvinotoolkit/nncf/blob/release_v230/examples/tensorflow/classification/configs/pruning/resnet50_imagenet_pruning_geometric_median.json)が参考になります。
今回はフィルタープルーニングを行うため、下記のNNCF設定ファイルを`~/config.json`に保存してください。

```json:config.json
{
    "input_info": { // 入力情報(必須)
        "sample_size": [ // 入力データのサイズ
            1,
            224,
            224,
            3
        ]
    },
    "compression": { // 圧縮設定
        "algorithm": "filter_pruning", // 圧縮アルゴリズムを指定(フィルタープルーニング)
        "pruning_init": 0.0, // 初期プルーニングレベル。デフォルトは0.0。
        "params": {
            "schedule": "exponential", // プルーニングレベルのスケジュール。デフォルトは`exponential`。
            "pruning_target": 0.5, // プルーニングレベルの目標値。デフォルトは0.5。
            "num_init_steps": 0, // 事前学習を行うエポック数。デフォルトは0。
            "pruning_steps": 1000, // プルーニングを行うエポック数。プルーニングレベルは`pruning_init`の値から`pruning_target`の値まで`pruning_steps`分のepochにわたり増加する。
            "filter_importance": "geometric_median" // フィルタの重要度基準。デフォルトは`L2`。
        }
    }
}
```

フィルタープルーニングを行う場合、NNCF設定ファイルには主に下記のパラメータを設定します。

- `input_info`
- `compression`
    - `algorithm`
    - `params`
        - `pruning_target`
        - `pruning_steps`

まずモデルの入力に関する情報を`input_info`に指定する必要があります。
今回は1入力の画像分類モデル(ResNet-50)を扱うため、その入力サイズ`[1,224,224,3]`を`sample_size`に指定しています。

また圧縮アルゴリズムに関する設定は`compression`に記述します。
適用する圧縮アルゴリズムを指定するには`algorithm`にアルゴリズム名を指定します。
今回はフィルタープルーニングを行うため`algorithm`に`filter_pruning`を指定しています。

また圧縮アルゴリズムに関する詳細なパラメータを`params`に記述します。
どの程度のフィルタープルーニングを行うかの目標値(目標プルーニングレベル)を設定するには`pruning_target`を指定します。
上記の設定では目標プルーニングレベルを50%とするため`pruning_target`を`0.5`としています。
フィルタープルーニングにかける学習epoch数を設定するには`pruning_steps`を指定します。
上記の設定では1000 epochかけてフィルタープルーニングを行うため`pruning_steps`を`1000`としています。

設定ファイルには他にも様々なパラメータが用意されています。
詳細は[公式ドキュメント](https://github.com/openvinotoolkit/nncf/blob/release_v230/docs/compression_algorithms/Pruning.md#pruning)をご覧ください。

以上でNNCF設定ファイルの作成は完了です。

### ソースコード

圧縮プログラムとして下記のソースコードを`~/compress.py`に保存してください。

<details>
<summary>ここをクリックしてソースコードを表示</summary>

``` python:compress.py
from pathlib import Path
import glob
import numpy as np
import tensorflow as tf
from nncf import NNCFConfig
from nncf.tensorflow import create_compressed_model
from nncf.tensorflow import register_default_init_args
from nncf.tensorflow import create_compression_callbacks
from nncf.tensorflow.utils.state import TFCompressionState
from nncf.tensorflow.utils.state import TFCompressionStateLoader

# シード固定
tf.random.set_seed(0)

# ハイパーパラメータ
batch_size = 32             # バッチサイズ
epochs = 1200               # エポック数
img_size = 224              # 画像サイズ
class_num = 2               # クラス数
dropout_rate = 0.2          # ドロップアウト率
learning_rate = 0.00001     # 学習率
shuffle_buffer_size = 1000

# 画像データの形状
img_shape = (img_size, img_size, 3)

# 前処理
preprocess = tf.keras.applications.resnet50.preprocess_input

# データ拡張
augment_li = [
    tf.keras.layers.experimental.preprocessing.RandomFlip("horizontal"),
    tf.keras.layers.experimental.preprocessing.RandomRotation(0.2),
]
augment = tf.keras.Sequential(augment_li)

# データセットの読込元
data_dir_path = Path("./cats_and_dogs_filtered")
train_dir_path = data_dir_path / "train"
val_dir_path = data_dir_path / "validation"

# 学習済みモデルの読込元
trained_saved_model_path = "./model/saved_model"

# 設定ファイル
config_path = Path("./config.json")

# 圧縮ログ保存先
log_dir_path = Path("./compressed_model")
log_dir_path.mkdir(parents=True, exist_ok=True)
tensorboard_dir_path = log_dir_path / "tensorboard"
ckpt_dir_path = log_dir_path / "checkpoint"
ckpt_path = ckpt_dir_path / "ckpt"
saved_model_path = log_dir_path / "saved_model"


def get_datasets(dir_path, shuffle_flg=False, augment_flg=False):
    """データセット取得"""
    # ディレクトリからデータを読込み
    ds = tf.keras.preprocessing.image_dataset_from_directory(
        directory=dir_path,
        image_size=(img_size, img_size),
        label_mode="categorical",
        batch_size=batch_size
    )
    # シャッフル
    if shuffle_flg:
        ds = ds.shuffle(shuffle_buffer_size, reshuffle_each_iteration=True)
    # データ拡張
    if augment_flg:
        ds = ds\
            .map(
                lambda x, y: (augment(x, training=True), y),
                num_parallel_calls=tf.data.AUTOTUNE
            )
    # 前処理
    ds = ds.map(lambda x, y: (preprocess(x), y))
    ds = ds.prefetch(buffer_size=tf.data.AUTOTUNE)

    return ds


def load_model(saved_model_path):
    """Saved Model読込"""
    return tf.keras.models.load_model(saved_model_path)


def update_ckpt_name():
    """チェックポイントの名前を変更"""
    tmp_ckpt_paths = glob.glob(str(ckpt_dir_path / "ckpt-*"))
    for tmp_ckpt_path in tmp_ckpt_paths:
        tmp_ckpt_path = Path(tmp_ckpt_path)
        tmp_ckpt_path.rename(
            ckpt_path.with_suffix(tmp_ckpt_path.suffix))


def compress(model, train_ds, val_ds):
    """モデル圧縮"""
    print("# Compression")

    # 圧縮の設定ファイルを読込
    nncf_config = NNCFConfig.from_json(config_path)

    # 圧縮アルゴリズムの初期化処理にデータセットを使用
    nncf_config = register_default_init_args(
        nncf_config,
        train_ds,
        batch_size=batch_size
    )

    # 重みをすべて解凍
    model.trainable = True

    # Batch Normalization層のみ凍結
    for layer in model.layers:
        if layer.name.endswith('bn') \
                or layer.name.startswith('batch_normalization') \
                or layer.name.endswith('BatchNorm'):
            layer.trainable = False

    # モデル圧縮のための変形
    compression_ctrl, compressed_model = create_compressed_model(
        model,
        nncf_config
    )

    # 最適化アルゴリズム
    optimizer = tf.keras.optimizers.Adam(learning_rate=learning_rate)

    # モデルの構成を確認
    compressed_model.summary()

    # チェックポイント用意
    checkpoint = tf.train.Checkpoint(
        model=compressed_model,
        compression_state=TFCompressionState(compression_ctrl)
    )
    checkpoint_maneger = tf.train.CheckpointManager(
        checkpoint, str(ckpt_dir_path), max_to_keep=1)

    # 圧縮用のコールバックを作成
    callbacks = []

    class BestCheckpointAfterPruning(tf.keras.callbacks.Callback):
        """Pruning後のモデル保存用コールバック"""

        def __init__(self, checkpoint_maneger, pruning_steps):
            self.checkpoint_maneger = checkpoint_maneger
            self.best_acc = -np.Inf
            self.best_epoch = -np.Inf
            self.pruning_steps = pruning_steps

        def on_epoch_end(self, epoch, logs):
            # val_accが最良のモデルを保存
            acc = logs.get("val_acc")
            if epoch > self.pruning_steps and acc > self.best_acc:
                self.checkpoint_maneger.save()
                self.best_acc = acc
                self.best_epoch = epoch
                print("Save the current model")

    pruning_steps = nncf_config.get("compression", {})\
        .get("params", {}).get("pruning_steps", epochs-1)
    ckpt_callback = BestCheckpointAfterPruning(
        checkpoint_maneger, pruning_steps)
    callbacks.append(ckpt_callback)
    callbacks.append(
        tf.keras.callbacks.TensorBoard(log_dir=tensorboard_dir_path))
    compression_callbacks = create_compression_callbacks(
        compression_ctrl,
        log_dir=str(tensorboard_dir_path)
    )
    callbacks.append(compression_callbacks)

    # 精度
    acc = tf.keras.metrics.CategoricalAccuracy(name='acc'),
    # 損失関数
    loss_obj = tf.keras.losses.CategoricalCrossentropy(name='loss')

    metrics = [acc, loss_obj]

    # モデルのコンパイル
    compressed_model.compile(
        optimizer=optimizer,
        loss=loss_obj,
        metrics=metrics
    )

    # 学習
    compressed_model.fit(
        train_ds,
        validation_data=val_ds,
        epochs=epochs,
        callbacks=callbacks
    )

    # 圧縮済みモデルの精度を表示
    print("[Best] epoch: {} val_acc: {:.4f}".format(
        ckpt_callback .best_epoch+1, ckpt_callback.best_acc))

    # Checkpointのファイル名を変更
    update_ckpt_name()


def load_compressed_model(model):
    """圧縮済みモデル読込"""

    # 圧縮状態の復元
    checkpoint = tf.train.Checkpoint(
        compression_state=TFCompressionStateLoader())
    checkpoint.restore(str(ckpt_path))

    # 圧縮の設定ファイルを読込
    nncf_config = NNCFConfig.from_json(config_path)

    compression_state = checkpoint.compression_state.state
    compression_ctrl, compress_model = create_compressed_model(
        model, nncf_config, compression_state)

    # モデルの復元
    checkpoint = tf.train.Checkpoint(model=compress_model)
    checkpoint.restore(str(ckpt_path))

    return compression_ctrl, compress_model, compression_state


def export_saved_model(model):
    """Saved Model出力"""

    print("# Export the best model as SavedModel")

    # Checkpoint読込
    compression_ctrl, compress_model, compression_state\
        = load_compressed_model(model)

    # SavedModel出力
    compression_ctrl.export_model(str(saved_model_path), save_format="tf")


def main():
    # データセット準備
    train_ds = get_datasets(
        train_dir_path, shuffle_flg=True, augment_flg=True)
    val_ds = get_datasets(val_dir_path)

    # 学習済みモデルロード
    model = load_model(trained_saved_model_path)

    # 圧縮
    compress(model, train_ds, val_ds)

    # SavedModel変換
    export_saved_model(model)


if __name__ == "__main__":
    main()
```

</details>

圧縮プログラムについて説明します。

まずハイパーパラメータとして`epoch`を1200としています。
前節にて作成したNNCF設定ファイルで`pruning_steps`を`1000`としたため、フィルタープルーニングを1000 epochかけて行った後、200 epoch学習を継続します。
その他、前回の記事の学習プログラムと共通なハイパーパラメータや前処理、データ拡張は、学習プログラムと同様のものとしています。

次にデータセットと圧縮対象のモデルを読み込みます。
データセットとして、学習プログラムで使用したものと同じcats_and_dogs_filteredデータセットを読み込みます。
また圧縮対象のモデルとして、学習プログラムで作成したSavedModel形式の学習済みモデルを`model`として読み込みます。

続いてNNCF APIを利用している部分について説明します。
前節にて作成したNNCF設定ファイルのパスを`NNCFConfig.from_json(...)`に指定し、圧縮の設定を`nncf_config`として読み込みます。
また圧縮対象となる`model`を`create_compressed_model(model, nncf_config)`にて`compressed_model`へ変換します。
その後`compressed_model.fit(...)`により圧縮を開始します。

圧縮後、学習を継続(1000-1200 epoch)し精度が最良となったモデルのCheckpointを1個保存するため、`compressed_model.fit(...)`のコールバックとして`BestCheckpointAfterPruning(...)`を定義しています。
最終的に精度最良モデルのCheckpointを再度読み込み、SavedModel形式へ変換します。

### 実行

仮想環境を起動し、プログラムを実行します。

``` console
$ cd
$ source ./venv_nncf/bin/activate
(venv_nncf)$ python compress.py
...
Epoch 1200/1200
63/63 [==============================] - 10s 128ms/step - loss: 0.0255 - acc: 0.9955 - val_loss: 0.1743 - val_acc: 0.9600
INFO:nncf:Statistics by pruned layers:
+--------------------+-------------------+----------------+--------------------+
|    Layer's name    |  Weight's shape   |  Mask's shape  |   Filter pruning   |
|                    |                   |                |       level        |
+====================+===================+================+====================+
| conv2_block1_1_con | [1, 1, 64, 64]    | [1, 1, 1, 64]  | 0.500              |
| v/kernel_mask:0    |                   |                |                    |
+--------------------+-------------------+----------------+--------------------+
...
+--------------------+-------------------+----------------+--------------------+
| conv5_block3_2_bn/ | [512]             | [512]          | 0.500              |
| beta_mask:0        |                   |                |                    |
+--------------------+-------------------+----------------+--------------------+
Statistics of the pruned model:
+---------+--------+---------+---------------+
|    #    |  Full  | Current | Pruning level |
+=========+========+=========+===============+
| GFLOPS  | 7.712  | 3.351   | 0.565         |
+---------+--------+---------+---------------+
| MParams | 23.656 | 11.496  | 0.514         |
+---------+--------+---------+---------------+
| Filters | 26562  | 22722   | 0.145         |
+---------+--------+---------+---------------+
Prompt: statistic pruning level = 1 - statistic current / statistic full.
Statistics of the filter pruning algorithm:
+---------------------------------------+-------+
|           Statistic's name            | Value |
+=======================================+=======+
| Filter pruning level in current epoch | 0.500 |
+---------------------------------------+-------+
| Target filter pruning level           | 0.500 |
+---------------------------------------+-------+
[Best] epoch: 1171 val_acc: 0.9880
...
```

実行すると学習の経過が上記のように表示されます。

まず精度や損失関数の値がepochごとに表示されます。
次にフィルタープルーニングの適用状況がepochごとに表示されます。
「Statistics by pruned layers」にはレイヤーごとのプルーニングレベルが表示されます。
「Statistics of the pruned model」にはGFLOPS、パラメータ数、フィルターに対するプルーニングレベルの適用状況が表示されます。
「Statistics of the filter pruning algorithm」には現在のepochにおいてのプルーニングレベルと設定ファイルにて指定した目標プルーニングレベルが表示されます。
Epochが進行するごとにプルーニングレベルが上昇することが確認できます。

全epochが終了すると、フィルタープルーニング後の学習で最良となった精度とそのepochが表示されます。
上記の例では1,171 epoch目で精度98.8%となりました。
前回の記事より圧縮前のモデルは99.4%でしたので、フィルタープルーニング後の精度低下は0.6%となります。

最後に上記の最良精度の圧縮済みモデルが`~/compressed_model/saved_model`にSavedModel形式で出力されます。
以上でモデル圧縮は完了です。

## Model OptimizerによるIRへの変換

Model Optimizerを使用して圧縮済みモデルをIRへ変換します。

``` console
$ cd
$ source /opt/intel/openvino/venv/bin/activate
(venv)$ mo \
--saved_model_dir ./compressed_model/saved_model \
--output_dir ./compressed_model/ir \
--framework tf \
--data_type FP16 \
--source_layout nhwc \
--target_layout nchw \
--input_shape [1,224,224,3] \
--mean_values [123.68,116.779,103.939] \
--transform Pruning
```

実行すると`~/compressed_model/ir`ディレクトリが作成され、その下に圧縮済みモデルのIR(xmlファイルとbinファイル)が出力されます。

[APIリファレンス](https://www.tensorflow.org/api_docs/python/tf/keras/applications/resnet50/preprocess_input)によると、前処理として入力画像をRGBからBGRへ変換し、スケーリングなしかつImageNetデータセットに基づいてゼロ中心化した値へ変換します。
具体的には`[0, 255]`のスケールを持つRGB画像からそれぞれ`(123.68, 116.779, 103.939)`が減算された値へ変換します。
この前処理を踏まえ、引数`--mean_values [123.68,116.779,103.939]`を指定します。
また、フィルタープルーニングを適用するために`--transform Pruning`を指定します。

以上でIRへの変換は完了です。

## AE2100での推論

OpenVINOのサンプルプログラムを実行し、圧縮済みモデルの動作を確認します。

### AE2100へのデプロイ

下記のファイルをAE2100へデプロイします。

- 画像分類サンプルプログラム(`classification_sample_async`)
- 圧縮済みモデルのIR(`~/compressed_model/ir/saved_model{.xml,.bin}`)
- ラベルファイル(`~/compressed_model/ir/saved_model.labels`)
- 猫／犬の画像(`~/cats_and_dogs_filtered/validation`に保存されている画像数枚)

今回は2クラス分類ですので[画像分類サンプルプログラム(classification_sample_async/main.cpp) 33行目](https://github.com/openvinotoolkit/openvino/blob/60099a19bd8df2b48d4f574d9bea0ddfde96e448/samples/cpp/classification_sample_async/main.cpp#L33)の`N_TOP_RESULTS`の値を`10`から`2`に変更してビルドしました。

分類結果を確認するためのラベルファイルとしてIRと同じファイル名+拡張子`.labels`(`saved_model.labels`)のテキストファイルを下記のように作成してください。

``` plaintext:saved_model.labels
cat
dog
```

### 動作確認

圧縮済みモデルによる画像分類の結果を確認するにはAE2100標準コンテナのシェルを起動して`classification_sample_async`を実行します。
例えば圧縮済みモデル`saved_model.xml`で猫の画像`cat.2000.jpg`の分類を行うには下記のようにコマンドを実行します。

``` console
root@92d0d660f309:~# ./classification_sample_async \
-i cat.2000.jpg \
-m saved_model.xml \
-d HDDL
```

実行結果が下記のように表示され、猫の画像を正しく分類できていることがわかります。

``` plaintext
...
Top 2 results:

Image ../img/cat.2000.jpg

classid probability label
------- ----------- -----
0       1.0000000   cat
1       0.0000000   dog
```

また圧縮前のモデルによる画像分類の結果は、下記のように得られました。

``` plaintext
...
Top 2 results:

Image ../img/cat.2000.jpg

classid probability label
------- ----------- -----
0       1.0000000   cat
1       0.0000000   dog
```

圧縮前のモデルと圧縮済みモデルとで分類結果が一致しました。

なおOpenVINO Toolkitのベンチマークプログラム(`benchmark_app`)で圧縮前後のモデルのスループットを測定すると、下記の表の通りとなりました。

| 目標プルーニングレベル<br>(`pruning_target`) | スループット(VPU)  | スループット(GPU) |
| :---: | ---: | ---: |
| 未圧縮 | 59.1 FPS | 25.7 FPS |
| 50%    | 91.2 FPS | 44.1 FPS |

圧縮前のモデルに対してフィルタープルーニングを行うことにより、推論スループットはVPUで約1.5倍、GPUで約1.7倍となりました。

## (参考)ResNet-50(ImageNet-2012)のフィルタープルーニング結果

参考までにImageNet-2012データセットを事前学習したResNet-50モデルのフィルタープルーニングを行った結果を掲載します。
なお下記の測定ではNNCFのサンプルプログラム(Image Classification Sample)を利用してフィルタープルーニングを行いました。

https://github.com/openvinotoolkit/nncf/tree/release_v230/examples/tensorflow/classification

| 目標プルーニングレベル<br>(`pruning_target`) | 精度<br>(Accuracy) | スループット(VPU)  | スループット(GPU) |
| :--:  | :---: | ---: | ---: |
| 未圧縮 | 75.04% | 71.5 FPS | 27.8 FPS |
| 40%    | 74.81% | 87.9 FPS | 40.8 FPS |
| 60%    | 74.03% | 125.2 FPS | 51.0 FPS |
| 80%    | 71.79% | 147.7 FPS | 72.5 FPS |

目標プルーニングレベル60%では精度の低下を約1%に抑えつつ、推論スループットはVPU／GPUともに約1.8倍となりました。
一方で目標プルーニングレベル80%では推論スループットはVPUで約2.1倍、GPUで約2.6倍となっているものの、精度は約3%の低下となりました。

目標プルーニングレベルに対してどの程度精度を保持できるかはデータセットによって異なります。
上記の表を参考にお手持ちのデータセットでモデルの圧縮を試してみてください。

## おわりに

本記事では学習済みの画像分類モデルにNNCFを適用し、モデル圧縮を行いました。
また圧縮済みモデルをAE2100へデプロイし、動作を確認しました。

AE2100での推論パフォーマンス向上を図る際はNNCFの活用をご検討ください。