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

はじめに

OpenVINO Toolkitの関連ツールであるNeural Network Compression Framework (NNCF)について2本の記事にわたり説明します。
NNCFを使用することで学習済みDNNモデルの圧縮を行うことができます。
これによりOKIのAIエッジコンピュータであるAE2100上でのディープラーニング・モデルの推論パフォーマンスの向上が可能となります。

この記事シリーズはNNCFによりDNNモデルを圧縮し、AE2100での推論スループットを向上させることを目標としています。

1. 準備編(この記事)
2. 圧縮編

今回の記事は準備編としてNNCFの概要について説明します。
またNNCFの環境構築を行い、次回の記事で圧縮の対象とする学習済みモデルを作成します。

なおこの記事の内容はOpenVINO Toolkitのバージョンとして2022.1以降に対応しています。

NNCFとは

Neural Network Compression Framework (NNCF)はディープニューラルネットワーク・モデル(DNNモデル)の圧縮を行うためのフレームワークです。
NNCFを利用することによりOpenVINO Toolkitを利用したモデル推論のパフォーマンス向上を図ることができます。
下図はNNCFとOpenVINO Toolkitの関係を表しています。

NNCFは既存の学習用プログラムにPythonパッケージとしてインポートして使用することができます。

モデルの精度を保持しながら圧縮するため、NNCFではPyTorchやTensorFlowといったディープラーニング・フレームワークを利用したfine-tuningを行いながら圧縮を行います。
モデルの圧縮後、OpenVINO ToolkitのModel Optimizerを利用して圧縮済みモデルをIRへ変換することで、IntelのデバイスとOpenVINO Toolkitを利用した推論が可能となります。

なおNNCFは下記のGitHubリポジトリにて公開されています。
NNCFの詳細については下記リポジトリをご参照ください。

NNCFの特徴

利用方法

NNCFはPythonパッケージとして構成されており、PyTorchやTensorFlowで書かれた既存の学習用プログラムに数行程度のコードを追加することで簡単にモデル圧縮を行うことができます。
またパラメータとして圧縮アルゴリズムの種類、目標圧縮率、圧縮にかけるエポック数を設定するだけで基本的な圧縮を行うことができます。

圧縮アルゴリズム

NNCFで扱うことができる圧縮アルゴリズムとして量子化(quantization)、スパース化(sparsity)、フィルタープルーニング(Filter pruning)があります。
またNNCFでは圧縮アルゴリズムを個別で利用することもできますが、組み合わせて利用することができます。例えば量子化とスパース化を組み合わせた圧縮が可能です。

目標値を設定したモデル圧縮

スパース化やフィルタープルーニングでは圧縮の目標値として目標圧縮率または目標精度のいずれかを設定することができます。
目標圧縮率を設定すると、学習が進むにつれて圧縮率を上げるようモデルの圧縮が行われます。
また目標精度に基づいた圧縮の方法として、精度低下の基準に達したときに圧縮を終了するEarly Exit Trainingや、バリデーションの精度に基づいて圧縮率を調整しながら圧縮を行うAdaptive Compression Level Trainingがあります。

圧縮済みモデルの出力形式

NNCFはOpenVINOに対応したモデル形式で圧縮済みモデルをファイルとして出力することができます。
例えばPyTorchの圧縮済みモデルはONNX形式、TensorFlowの圧縮済みモデルはSavedModel形式またはFrozen Graph形式に対応しています。

圧縮アルゴリズムの概要

量子化、スパース化、フィルタープルーニングについて簡単に説明します。

量子化

量子化はモデルの量子化ビット数を減らすモデル圧縮方法です。
NNCFによる量子化はモデルの重みと活性化関数の演算を8ビット等の桁数へ変換することをサポートしています。

スパース化

スパース化は畳み込み層や全結合層のうち重要度の低い重みをゼロにするモデル圧縮方法です。
スパース化を適用するとモデルが持つ重みテンソル内でゼロが分散した状態になります。
スパース演算に対応したハードウェアでは、重みがゼロの部分の演算をスキップすることができるためモデルの推論スループット向上を図ることができます。
現行のAE2100のデバイスはスパース演算に対応していませんが、次世代のVPUでは対応される見込みです。

フィルタープルーニング

フィルタープルーニングはモデルが持つ畳み込み層のうち重要度の低いフィルターを削減するモデル圧縮方法です。
NNCFのフィルタープルーニングでは重要度としてL1ノルム、L2ノルム、geometoric medianという3種類の基準を扱うことができます。
フィルタープルーニングでは、削減されたフィルターについての演算処理を丸ごとスキップすることができます。
したがってフィルタープルーニングを行ったモデルの推論は、ハードウェアによらず推論スループットの向上が期待できます。

AE2100のデバイスとNNCFの対応

AE2100にはOpenVINO Toolkitに対応したAIアクセラレータであるVPU (Intel Movidius Myriad X)、GPU、CPUが備わっています。
AE2100に搭載されているデバイスと推論パフォーマンスの向上が期待できる圧縮アルゴリズムの対応は下表の通りです。

	VPU	GPU	CPU
量子化			o
スパース化
フィルタープルーニング	o	o	o

AE2100のモデル推論に適したデバイスは主にVPUとGPUであるため、この記事シリーズではフィルタープルーニングを行う方法を説明します。

環境構築

ここからはPythonがインストール済みの環境でNNCFを導入する方法を説明します。
なおNNCFで扱う学習フレームワークとしてTensorFlowとPyTorchがありますが、以下ではNNCFとTensorFlowを導入します。

バージョン

この記事では下記の構成で環境構築を行います。

• Ubuntu 20.04
• Python 3.8.10
• pip 20.0.2
• nncf 2.3.0

手順

1. (推奨)NNCF用に仮想環境を作成します。

   $ cd
   $ python3 -m venv venv_nncf
   

2. 仮想環境を起動します。

   $ source ./venv_nncf/bin/activate
   

3. pipによりNNCFとTensorFlowをインストールします。

   (venv_nncf)$ pip install nncf[tf]==2.3.0
   

   • インストールした際に

     ERROR: tensorboard 2.9.1 has requirement protobuf<3.20,>=3.9.2, but you'll have protobuf 4.21.4 which is incompatible.
     ERROR: tensorflow 2.5.3 has requirement six~=1.15.0, but you'll have six 1.16.0 which is incompatible.
     

     のようなエラーが表示された場合は、pipにより該当のパッケージをバージョンを指定してインストールしてください。

     (venv_nncf)$ pip install protobuf==3.20
     

以上でNNCFの実行環境は構築完了となります。

データセットの用意

次にモデルの学習と評価に用いるデータセットを用意します。
今回は2クラスの画像分類としてcats_and_dogs_filteredデータセットを扱います。
データセットの構成は下記の通りです。なおクラスあたりのデータ数が均等なデータセットです。

• 学習データ: 2,000個
• 検証データ: 1,000個

wgetコマンドによりzip形式のデータセットをダウンロードします。
またunzipコマンドによりダウンロードしたデータセットを解凍します。

$ cd
$ wget https://storage.googleapis.com/mledu-datasets/cats_and_dogs_filtered.zip
$ unzip cats_and_dogs_filtered.zip

上記により~/cats_and_dogs_filteredにデータセットが展開されます。
ディレクトリ構成は下記の通りです。
trainに訓練データが、validationに検証データが格納されています。
またtrainとvalidationのそれぞれにcatsとdogsというクラスごとのディレクトリがあり、クラスごとに画像が格納されています。

cats_and_dogs_filtered/
├ train/
│ ├ cats/
│ │ ├ cat.0.jpg
│ │ └ ...
│ └ dogs/
│   ├ dog.0.jpg
│   └ ... 
└ validation/
  ├ cats/
  │ ├ cat.2000.jpg
  │ └ ...
  └ dogs/
    ├ dog.2000.jpg
    └ ... 

以上でデータセットの用意は完了です。

事前学習

続いてNNCFによる圧縮の対象とするモデルの事前学習を実施します。
今回はImageNet-2012データセットによる学習済みResNet-50を利用して画像分類モデルを作成します。
なお本プログラムではまず最終層のみ学習を行う転移学習を10 epoch行った後、Batch Normalization層を除く層を解凍してfine-tuningを190 epoch行います。

NNCF v2.3.0では圧縮可能なTensorFlowモデルとしてSequentialまたはKeras Functional APIにより作られたモデルのみサポートされています。

ソースコード

下記のソースコードをtrain.pyというファイル名で実行環境に保存してください。

ここをクリックしてソースコードを表示

train.py

from pathlib import Path
import tensorflow as tf

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

# ハイパーパラメータ
batch_size = 32             # バッチサイズ
initial_epochs = 10         # 特徴抽出エポック数
fine_tune_epochs = 190      # fine-tuningエポック数
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)

# 事前学習済みモデルを取得
base_model = tf.keras.applications.ResNet50(
    input_shape=img_shape,
    include_top=False,
    weights="imagenet"
)

# 前処理
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"

# 学習ログ保存先
log_dir_path = Path("./model")
log_dir_path.mkdir(parents=True, exist_ok=True)
tensorboard_dir_path = log_dir_path / "tensorboard"
saved_model_path = log_dir_path / "saved_model"
saved_model_path.mkdir(parents=True, exist_ok=True)


def get_model():
    """モデル定義"""
    # Fanctional APIでモデルを定義
    x = base_model.output
    x = tf.keras.layers.Flatten()(x)
    x = tf.keras.layers.Dropout(dropout_rate)(x)
    outputs = tf.keras.layers.Dense(class_num, activation="softmax")(x)
    model = tf.keras.Model(base_model.input, outputs)

    return 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 train(train_ds, val_ds):
    """モデル学習"""
    # 学習済みモデル確認
    base_model.summary()

    # モデル取得
    model = get_model()

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

    # 損失関数
    loss_func = tf.keras.losses.CategoricalCrossentropy

    # コールバック
    callbacks = []
    callbacks.append(tf.keras.callbacks.TensorBoard(
        log_dir=tensorboard_dir_path))

    def feature_extract(model):
        """特徴抽出"""
        print("# Feature extraction")
        # モデル凍結
        base_model.trainable = False
        # モデルコンパイル
        model.compile(
            optimizer=optimizer(learning_rate=learning_rate),
            loss=loss_func(),
            metrics=["accuracy"]
        )
        # モデル確認
        model.summary()
        # 学習前評価
        model.evaluate(val_ds)
        # 学習実行
        history = model.fit(
            train_ds,
            validation_data=val_ds,
            epochs=initial_epochs,
            callbacks=callbacks
        )

        return model, history

    def fine_tune(model, history):
        """Fine-tuning"""
        print("# Fine-tuning")
        # モデル解凍
        base_model.trainable = True

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

        # モデルコンパイル
        model.compile(
            optimizer=optimizer(learning_rate=learning_rate/10),
            loss=loss_func(),
            metrics=["accuracy"]
        )

        # モデル確認
        model.summary()

        # エポック下図計算
        total_epochs = initial_epochs + fine_tune_epochs

        # 学習実行
        history = model.fit(
            train_ds,
            validation_data=val_ds,
            epochs=total_epochs,
            initial_epoch=history.epoch[-1],
            callbacks=callbacks
        )

        return model, history

    model, history = feature_extract(model)
    fine_tune(model, history)

    return model


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

    # 学習
    model = train(train_ds, val_ds)

    # モデルを保存
    model.save(saved_model_path, save_format='tf')


if __name__ == "__main__":
    main()

本プログラムはTensorFlowのチュートリアルを参考に作成しています。

本プログラムではtensorflow.keras.applications.ResNet50によりモデルを取得します。
ただしinput_shapeに指定する入力画像サイズを(224, 224, 3)としています。
また重みとしてweights=imagenetを指定し、2クラス出力とするためinclude_top=Falseを指定しています。

入力画像の前処理はtensorflow.keras.applications.resnet50.preprocess_inputにて実施します。
またデータ拡張はランダム左右反転とランダム回転(20%)としています。

実行

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

$ cd
$ source ./venv_nncf/bin/activate
(venv_nncf)$ python train.py
...
Epoch 198/200
63/63 [==============================] - 12s 134ms/step - loss: 1.2543e-04 - accuracy: 1.0000 - val_loss: 0.0244 - val_accuracy: 0.9950
Epoch 199/200
63/63 [==============================] - 13s 150ms/step - loss: 1.5112e-04 - accuracy: 1.0000 - val_loss: 0.0278 - val_accuracy: 0.9940
Epoch 200/200
63/63 [==============================] - 13s 152ms/step - loss: 0.0021 - accuracy: 0.9990 - val_loss: 0.0267 - val_accuracy: 0.9940

実行する学習の経過が上記のように表示されます。
最終epochのval_accuracyが検証データに対するモデルの精度となります。上記の例では最終の200 epoch目で精度は99.4%です。

また./modelディレクトリが作成され、学習の終了後./model/saved_modelディレクトリに最終epochのモデルがSaved Model形式で保存されます。

次回の記事では圧縮用のプログラムを実行して今回作成したモデルを読み込み、NNCFによるフィルタープルーニングを行います。

おわりに

今回の記事ではNNCFの概要と環境構築手順について説明し、圧縮の対象とする画像分類モデルを作成しました。
次回の記事では今回作成したモデルに対してNNCFを使用したフィルタープルーニングを行います。