diff --git a/examples/quantization_aware_training/tensorflow/mobilenet_v2/README.md b/examples/quantization_aware_training/tensorflow/mobilenet_v2/README.md new file mode 100644 index 00000000000..8c42efa055e --- /dev/null +++ b/examples/quantization_aware_training/tensorflow/mobilenet_v2/README.md @@ -0,0 +1,31 @@ +# Quantization-Aware Training of MobileNet v2 TensorFlow Model + +This example demonstrates how to use Post-Training Quantization API from Neural Network Compression Framework (NNCF) to quantize and train TensorFlow models on the example of [MobileNet v2](https://huggingface.co/alexsu52/mobilenet_v2_imagenette) quantization aware training, pretrained on [Imagenette](https://github.com/fastai/imagenette) dataset. + +The example includes the following steps: + +- Loading the [Imagenette](https://github.com/fastai/imagenette) dataset (~340 Mb) and the [MobileNet v2 TensorFlow model](https://huggingface.co/alexsu52/mobilenet_v2_imagenette) pretrained on this dataset. +- Quantizing the model using NNCF Post-Training Quantization algorithm. +- Fine tuning quantized model for two epoch to improve quantized model metrics. +- Output of the following characteristics of the quantized model: + - Accuracy drop of the quantized model (INT8) over the pre-trained model (FP32) + - Compression rate of the quantized model file size relative to the pre-trained model file size + - Performance speed up of the quantized model (INT8) + +## Install requirements + +At this point it is assumed that you have already installed NNCF. You can find information on installation NNCF [here](https://github.com/openvinotoolkit/nncf#user-content-installation). + +To work with the example you should install the corresponding Python package dependencies: + +```bash +pip install -r requirements.txt +``` + +## Run Example + +It's pretty simple. The example does not require additional preparation. It will do the preparation itself, such as loading the dataset and model, etc. + +```bash +python main.py +``` diff --git a/examples/quantization_aware_training/tensorflow/mobilenet_v2/main.py b/examples/quantization_aware_training/tensorflow/mobilenet_v2/main.py new file mode 100644 index 00000000000..07e152a7f8d --- /dev/null +++ b/examples/quantization_aware_training/tensorflow/mobilenet_v2/main.py @@ -0,0 +1,199 @@ +# Copyright (c) 2024 Intel Corporation +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# http://www.apache.org/licenses/LICENSE-2.0 +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re +import subprocess +from pathlib import Path +from typing import List + +import openvino as ov +import tensorflow as tf +import tensorflow_datasets as tfds +from rich.progress import track + +import nncf + +ROOT = Path(__file__).parent.resolve() +WEIGHTS_URL = "https://huggingface.co/alexsu52/mobilenet_v2_imagenette/resolve/main/tf_model.h5" +DATASET_CLASSES = 10 + + +def validate(model: ov.Model, val_loader: tf.data.Dataset) -> tf.Tensor: + compiled_model = ov.compile_model(model, device_name="CPU") + output = compiled_model.outputs[0] + + metric = tf.keras.metrics.CategoricalAccuracy(name="acc@1") + for images, labels in track(val_loader): + pred = compiled_model(images.numpy())[output] + metric.update_state(labels, pred) + + return metric.result() + + +def run_benchmark(model_path: Path, shape: List[int]) -> float: + command = [ + "benchmark_app", + "-m", model_path.as_posix(), + "-d", "CPU", + "-api", "async", + "-t", "15", + "-shape", str(shape), + ] # fmt: skip + cmd_output = subprocess.check_output(command, text=True) # nosec + print(*cmd_output.splitlines()[-8:], sep="\n") + match = re.search(r"Throughput\: (.+?) FPS", cmd_output) + return float(match.group(1)) + + +def get_model_size(ir_path: Path, m_type: str = "Mb") -> float: + xml_size = ir_path.stat().st_size + bin_size = ir_path.with_suffix(".bin").stat().st_size + for t in ["bytes", "Kb", "Mb"]: + if m_type == t: + break + xml_size /= 1024 + bin_size /= 1024 + model_size = xml_size + bin_size + print(f"Model graph (xml): {xml_size:.3f} Mb") + print(f"Model weights (bin): {bin_size:.3f} Mb") + print(f"Model size: {model_size:.3f} Mb") + return model_size + + +############################################################################### +# Create a Tensorflow model and dataset + + +def center_crop(image: tf.Tensor, image_size: int, crop_padding: int) -> tf.Tensor: + shape = tf.shape(image) + image_height = shape[0] + image_width = shape[1] + + padded_center_crop_size = tf.cast( + ((image_size / (image_size + crop_padding)) * tf.cast(tf.minimum(image_height, image_width), tf.float32)), + tf.int32, + ) + + offset_height = ((image_height - padded_center_crop_size) + 1) // 2 + offset_width = ((image_width - padded_center_crop_size) + 1) // 2 + + image = tf.image.crop_to_bounding_box( + image, + offset_height=offset_height, + offset_width=offset_width, + target_height=padded_center_crop_size, + target_width=padded_center_crop_size, + ) + + image = tf.compat.v1.image.resize( + image, [image_size, image_size], method=tf.image.ResizeMethod.BILINEAR, align_corners=False + ) + + return image + + +def preprocess_for_eval(image, label): + image = center_crop(image, 224, 32) + image = tf.keras.applications.mobilenet_v2.preprocess_input(image) + image = tf.image.convert_image_dtype(image, tf.float32) + + label = tf.one_hot(label, DATASET_CLASSES) + + return image, label + + +def preprocess_for_train(image, label): + image = tf.image.resize_with_crop_or_pad(image, 256, 256) + image = tf.image.random_crop(image, [224, 224, 3]) + image = tf.image.convert_image_dtype(image, tf.float32) + label = tf.one_hot(label, DATASET_CLASSES) + return image, label + + +train_dataset = tfds.load("imagenette/320px-v2", split="train", shuffle_files=True, as_supervised=True) +train_dataset = train_dataset.map(preprocess_for_train).shuffle(1024).batch(128) + +val_dataset = tfds.load("imagenette/320px-v2", split="validation", shuffle_files=False, as_supervised=True) +val_dataset = val_dataset.map(preprocess_for_eval).batch(128) + +weights_path = tf.keras.utils.get_file("mobilenet_v2_imagenette_weights.h5", WEIGHTS_URL, cache_subdir="models") +tf_model = tf.keras.applications.MobileNetV2(weights=weights_path, classes=DATASET_CLASSES) + +############################################################################### +# Quantize a Tensorflow model +# +# The transformation function transforms a data item into model input data. +# +# To validate the transform function use the following code: +# >> for data_item in val_loader: +# >> model(transform_fn(data_item)) + + +def transform_fn(data_item): + images, _ = data_item + return images + + +# The calibration dataset is a small, no label, representative dataset +# (~100-500 samples) that is used to estimate the range, i.e. (min, max) of all +# floating point activation tensors in the model, to initialize the quantization +# parameters. +# +# The easiest way to define a calibration dataset is to use a training or +# validation dataset and a transformation function to remove labels from the data +# item and prepare model input data. The quantize method uses a small subset +# (default: 300 samples) of the calibration dataset. + +calibration_dataset = nncf.Dataset(val_dataset, transform_fn) +tf_quantized_model = nncf.quantize(tf_model, calibration_dataset) + +tf_quantized_model.compile( + optimizer=tf.keras.optimizers.Adam(learning_rate=1e-5), + loss=tf.keras.losses.CategoricalCrossentropy(), + metrics=[tf.keras.metrics.CategoricalAccuracy()], +) + +tf_quantized_model.fit(train_dataset, epochs=3, verbose=1) + +############################################################################### +# Benchmark performance, calculate compression rate and validate accuracy + +ov_model = ov.convert_model(tf_model, share_weights=False) +ov_quantized_model = ov.convert_model(tf_quantized_model, share_weights=False) + +fp32_ir_path = ROOT / "mobilenet_v2_fp32.xml" +ov.save_model(ov_model, fp32_ir_path, compress_to_fp16=False) +print(f"[1/7] Save FP32 model: {fp32_ir_path}") +fp32_model_size = get_model_size(fp32_ir_path) + +int8_ir_path = ROOT / "mobilenet_v2_int8.xml" +ov.save_model(ov_quantized_model, int8_ir_path) +print(f"[2/7] Save INT8 model: {int8_ir_path}") +int8_model_size = get_model_size(int8_ir_path) + +print("[3/7] Benchmark FP32 model:") +fp32_fps = run_benchmark(fp32_ir_path, shape=[1, 224, 224, 3]) +print("[4/7] Benchmark INT8 model:") +int8_fps = run_benchmark(int8_ir_path, shape=[1, 224, 224, 3]) + +print("[5/7] Validate OpenVINO FP32 model:") +fp32_top1 = validate(ov_model, val_dataset) +print(f"Accuracy @ top1: {fp32_top1:.3f}") + +print("[6/7] Validate OpenVINO INT8 model:") +int8_top1 = validate(ov_quantized_model, val_dataset) +print(f"Accuracy @ top1: {int8_top1:.3f}") + +print("[7/7] Report:") +print(f"Accuracy drop: {fp32_top1 - int8_top1:.3f}") +print(f"Model compression rate: {fp32_model_size / int8_model_size:.3f}") +# https://docs.openvino.ai/latest/openvino_docs_optimization_guide_dldt_optimization_guide.html +print(f"Performance speed up (throughput mode): {int8_fps / fp32_fps:.3f}") diff --git a/examples/quantization_aware_training/tensorflow/mobilenet_v2/requirements.txt b/examples/quantization_aware_training/tensorflow/mobilenet_v2/requirements.txt new file mode 100644 index 00000000000..e1e3c89fe46 --- /dev/null +++ b/examples/quantization_aware_training/tensorflow/mobilenet_v2/requirements.txt @@ -0,0 +1,5 @@ +tensorflow~=2.12.0; python_version < '3.9' +tensorflow~=2.15.1; python_version >= '3.9' +tensorflow-datasets +tqdm +openvino==2024.6