TensorFlow フレームワークを使用した NNCF による量子化対応トレーニング#

%pip install -q "openvino>=2024.0.0" "nncf>=2.9.0" 
%pip install -q "tensorflow-macos>=2.5,<=2.12.0; sys_platform == 'darwin' and platform_machine == 'arm64'" 
%pip install -q "tensorflow>=2.5,<=2.12.0; sys_platform == 'darwin' and platform_machine != 'arm64'" # macOS x86 
%pip install -q "tensorflow>=2.5,<=2.12.0; sys_platform != 'darwin'" 
%pip install -q "tensorflow-datasets>=4.9.0,<4.9.3; platform_system=='Windows'" 
%pip install -q "tensorflow-datasets>=4.9.0"

from pathlib import Path 
import logging 

import tensorflow as tf 
import tensorflow_datasets as tfds 

from nncf import NNCFConfig 
from nncf.tensorflow.helpers.model_creation import create_compressed_model 
from nncf.tensorflow.initialization import register_default_init_args 
from nncf.common.logging.logger import set_log_level 
import openvino as ov 

set_log_level(logging.ERROR) 

MODEL_DIR = Path("model") 
OUTPUT_DIR = Path("output") 
MODEL_DIR.mkdir(exist_ok=True) 
OUTPUT_DIR.mkdir(exist_ok=True) 

BASE_MODEL_NAME = "ResNet-18" 

fp32_h5_path = Path(MODEL_DIR / (BASE_MODEL_NAME + "_fp32")).with_suffix(".h5") 
fp32_ir_path = Path(OUTPUT_DIR / "saved_model").with_suffix(".xml") 
int8_pb_path = Path(OUTPUT_DIR / (BASE_MODEL_NAME + "_int8")).with_suffix(".pb") 
int8_ir_path = int8_pb_path.with_suffix(".xml") 

BATCH_SIZE = 128 
IMG_SIZE = (64, 64) # デフォルトの Imagenet 画像サイズ 
NUM_CLASSES = 10 # Imagenette データセットの場合 

LR = 1e-5 

MEAN_RGB = (0.485 * 255, 0.456 * 255, 0.406 * 255) # Imagenet データセットから 
STDDEV_RGB = (0.229 * 255, 0.224 * 255, 0.225 * 255) # Imagenet データセットから 

fp32_pth_url = 
"https://storage.openvinotoolkit.org/repositories/nncf/openvino_notebook_ckpts/305_resnet18_imagenette_fp32_v1.h5" 
_ = tf.keras.utils.get_file(fp32_h5_path.resolve(), fp32_pth_url) 
print(f"Absolute path where the model weights are saved:\n {fp32_h5_path.resolve()}")

datasets, datasets_info = tfds.load( 
    "imagenette/160px", 
    shuffle_files=True, 
    as_supervised=True, 
    with_info=True, 
    read_config=tfds.ReadConfig(shuffle_seed=0), 
) 
train_dataset, validation_dataset = datasets["train"], datasets["validation"] 
fig = tfds.show_examples(train_dataset, datasets_info)

def preprocessing(image, label): 
    image = tf.image.resize(image, IMG_SIZE) 
    image = image - MEAN_RGB 
    image = image / STDDEV_RGB 
    label = tf.one_hot(label, NUM_CLASSES) 
    return image, label 

train_dataset = train_dataset.map(preprocessing, 
num_parallel_calls=tf.data.experimental.AUTOTUNE).batch(BATCH_SIZE).prefetch(tf.data.experimental.AUTOTUNE) 

validation_dataset = ( 
    validation_dataset.map(preprocessing, 
num_parallel_calls=tf.data.experimental.AUTOTUNE).batch(BATCH_SIZE).prefetch(tf.data.experimental.AUTOTUNE) 
)

def residual_conv_block(filters, stage, block, strides=(1, 1), cut="pre"): 
    def layer(input_tensor): 
        x = tf.keras.layers.BatchNormalization(epsilon=2e-5)(input_tensor) 
        x = tf.keras.layers.Activation("relu")(x) 

        # ショートカット接続を定義 
        if cut == "pre": 
            shortcut = input_tensor 
        elif cut == "post": 
            shortcut = tf.keras.layers.Conv2D( 
                filters, 
                (1, 1), 
                strides=strides, 
                kernel_initializer="he_uniform", 
                use_bias=False, 
            )(x) 

        # 畳み込みレイヤーを続行 
        x = tf.keras.layers.ZeroPadding2D(padding=(1, 1))(x) 
        x = tf.keras.layers.Conv2D( 
            filters, 
            (3, 3), 
            strides=strides, 
            kernel_initializer="he_uniform", 
            use_bias=False, 
        )(x) 

        x = tf.keras.layers.BatchNormalization(epsilon=2e-5)(x) 
        x = tf.keras.layers.Activation("relu")(x) 
        x = tf.keras.layers.ZeroPadding2D(padding=(1, 1))(x) 
        x = tf.keras.layers.Conv2D(filters, (3, 3), kernel_initializer="he_uniform", use_bias=False)(x) 

        # 残余接続を追加 
        x = tf.keras.layers.Add()([x, shortcut]) 
        return x 

    return layer 

def ResNet18(input_shape=None): 
    """Instantiates the ResNet18 architecture.""" 
    img_input = tf.keras.layers.Input(shape=input_shape, name="data") 

    # ResNet18 ボトム 
    x = tf.keras.layers.BatchNormalization(epsilon=2e-5, scale=False)(img_input) 
    x = tf.keras.layers.ZeroPadding2D(padding=(3, 3))(x) 
    x = tf.keras.layers.Conv2D(64, (7, 7), strides=(2, 2), kernel_initializer="he_uniform", use_bias=False)(x) 
    x = tf.keras.layers.BatchNormalization(epsilon=2e-5)(x) 
    x = tf.keras.layers.Activation("relu")(x) 
    x = tf.keras.layers.ZeroPadding2D(padding=(1, 1))(x) 
    x = tf.keras.layers.MaxPooling2D((3, 3), strides=(2, 2), padding="valid")(x) 

    # ResNet18 ボディー 
    repetitions = (2, 2, 2, 2) 
    for stage, rep in enumerate(repetitions): 
        for block in range(rep): 
            filters = 64 * (2**stage) 
            if block == 0 and stage == 0: 
                x = residual_conv_block(filters, stage, block, strides=(1, 1), cut="post")(x) 
            elif block == 0: 
                x = residual_conv_block(filters, stage, block, strides=(2, 2), cut="post")(x) 
            else: 
                x = residual_conv_block(filters, stage, block, strides=(1, 1), cut="pre")(x) 
    x = tf.keras.layers.BatchNormalization(epsilon=2e-5)(x) 
    x = tf.keras.layers.Activation("relu")(x) 

    # ResNet18 トップ 
    x = tf.keras.layers.GlobalAveragePooling2D()(x) 
    x = tf.keras.layers.Dense(NUM_CLASSES)(x) 
    x = tf.keras.layers.Activation("softmax")(x) 

    # モデルを作成 
    model = tf.keras.models.Model(img_input, x) 

    return model

IMG_SHAPE = IMG_SIZE + (3,) 
fp32_model = ResNet18(input_shape=IMG_SHAPE)

# 浮動小数点の重みをロード
 fp32_model.load_weights(fp32_h5_path) 

# 浮動小数点モデルをコンパイル 
fp32_model.compile(.compile( 
    loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.1), 
    metrics=[tf.keras.metrics.CategoricalAccuracy(name="acc@1")], 
) 

# 浮動小数点モデルを検証 
test_loss, acc_fp32 = fp32_model.evaluate( 
    validation_dataset, 
    callbacks=tf.keras.callbacks.ProgbarLogger(stateful_metrics=["acc@1"]), 
) 
print(f"\nAccuracy of FP32 model: {acc_fp32:.3f}")

nncf_config_dict = { 
    "input_info": {"sample_size": [1, 3] + list(IMG_SIZE)}, 
    "log_dir": str(OUTPUT_DIR), # NNCF 固有のログ出力のログ・ディレクトリー
    "compression": { 
        "algorithm": "quantization", # ここでアルゴリズムを指定     }, 
} 
nncf_config = NNCFConfig.from_dict(nncf_config_dict)

nncf_config = register_default_init_args(nncf_config=nncf_config, 
data_loader=train_dataset, batch_size=BATCH_SIZE)

compression_ctrl, int8_model = create_compressed_model(fp32_model, nncf_config)

# INT8 モデルをコンパイル 
Int8_model.compile(.compile( 
    optimizer=tf.keras.optimizers.Adam(learning_rate=LR), 
    loss=tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.1), 
    metrics=[tf.keras.metrics.CategoricalAccuracy(name="acc@1")], 
) 

# INT8 モデルを検証 
test_loss, test_acc = int8_model.evaluate( 
    validation_dataset, 
    callbacks=tf.keras.callbacks.ProgbarLogger(stateful_metrics=["acc@1"]), 
)

print(f"\nAccuracy of INT8 model after initialization: {test_acc:.3f}") 

# INT8 モデルをトレーニング 
int8_model.fit(train_dataset, epochs=2) 

# NT8 モデルを検証 
test_loss, acc_int8 = int8_model.evaluate( 
    validation_dataset, 
    callbacks=tf.keras.callbacks.ProgbarLogger(stateful_metrics=["acc@1"]), 
) 
print(f"\nAccuracy of INT8 model after fine-tuning: {acc_int8:.3f}") 
print(f"\nAccuracy drop of tuned INT8 model over pre-trained FP32 model: {acc_fp32 - acc_int8:.3f}")

model_ir_fp32 = ov.convert_model(fp32_model)

model_ir_int8 = ov.convert_model(int8_model)

ov.save_model(model_ir_fp32, fp32_ir_path, compress_to_fp16=False) 
ov.save_model(model_ir_int8, int8_ir_path, compress_to_fp16=False)

import ipywidgets as widgets 

# OpenVINO ランタイムを初期化 
core = ov.Core() 
device = widgets.Dropdown( 
    options=core.available_devices, 
    value="CPU", 
    description="Device:", disabled=False, 
) 

device

def parse_benchmark_output(benchmark_output): 
    parsed_output = [line for line in benchmark_output if "FPS" in line] 
    print(*parsed_output, sep="\n") 

print("Benchmark FP32 model (IR)") 
benchmark_output = ! benchmark_app -m $fp32_ir_path -d $device.value -api async -t 15 -shape [1,64,64,3]
parse_benchmark_output(benchmark_output) 

print("\nBenchmark INT8 model (IR)") 
benchmark_output = ! benchmark_app -m $int8_ir_path -d $device.value -api async -t 15 -shape [1,64,64,3] 
parse_benchmark_output(benchmark_output)

core = ov.Core() 
core.get_property(device.value, "FULL_DEVICE_NAME")