chore: import upstream snapshot with attribution

tools/tensorflow-quantization/examples/.gitignore

@@ -0,0 +1,3 @@
+*/weights
+efficientnet/models/*
+*/test_qdq_node_placement

tools/tensorflow-quantization/examples/README.md

@@ -0,0 +1,103 @@
+# About
+This folder contains the Quantization-Aware Training (QAT) workflow for [standard networks](#step-1-model-quantization-and-fine-tuning).
+
+The QAT end-to-end workflow (TF2-to-ONNX) consists of the following steps:
+- Model quantization using the `quantize_model` function with `NVIDIA` quantization scheme.
+- QAT model fine-tuning (saves checkpoints).
+- Baseline vs QAT models accuracy comparison.
+- QAT model conversion to SavedModel format.
+- Conversion of SavedModel to ONNX.
+- TensorRT engine building via ONNX file and inference.
+
+# Requirements
+## 1. Base requirements
+1. Install `tensorflow-quantization` toolkit.
+2. Install additional requirements: `pip install -r requirements.txt`.
+3. (Optional) Install TensorRT for full workflow support (needed for `infer_engine.py`).
+
+**Note**: For CLI run, please go to the cloned repository's root directory and run `export PYTHONPATH=$PWD`, so that the `examples` folder is available for import.
+
+## 2. Data preparation
+### A. Raw data download
+We are using the ImageNet 2012 dataset (task 1 - image classification), which requires manual downloads due to terms of access agreements.
+Please login/sign-up on [the ImageNet website](https://image-net.org/challenges/LSVRC/2012/2012-downloads.php) and download the "train/validation data".
+This is needed for the QAT model fine-tuning, and it is also used to evaluate the Baseline and QAT models.
+
+### B. Conversion to tfrecord
+Our workflow supports `tfrecord` format, so please follow the following instructions (modified from [TensorFlow's instructions](https://github.com/tensorflow/tpu/tree/master/tools/datasets#imagenet_to_gcspy)) to convert the downloaded `.tar` ImageNet files to the required format:
+
+1. Set `IMAGENET_HOME=/path/to/imagenet/tar/files` in [`data/imagenet_data_setup.sh`](data/imagenet_data_setup.sh).
+2. Download [`imagenet_to_gcs.py`](https://github.com/tensorflow/tpu/blob/master/tools/datasets/imagenet_to_gcs.py) to `$IMAGENET_HOME`.
+3. Run `./data/imagenet_data_setup.sh`.
+
+# Workflow
+
+## Step 1: Model quantization and fine-tuning
+Model quantization, fine-tuning, and conversion to ONNX.
+
+Example models:
+
+| Model         | Task             | Script - QAT Workflow        |
+|---------------|------------------|------------------------------|
+| ResNet        | Classification   | [resnet](resnet)             |
+| EfficientNet  | Classification   | [efficientnet](efficientnet) |
+| MobileNet     | Classification   | [mobilenet](mobilenet)       |
+| Inception     | Classification   | [inception](inception)       |
+> For each model's performance results, please refer to the toolkit's User Guide ("Model Zoo").
+
+## Step 2: TensorRT deployment
+Build the TensorRT engine and evaluate its latency and accuracy performances.
+
+#### 2.1. Build TensorRT engine from ONNX
+Convert the ONNX model into a TensorRT engine (also obtains latency measurements):
+
+```sh
+trtexec --onnx=model_qat.onnx --int8 --saveEngine=model_qat.engine --verbose
+```
+
+Arguments:
+* `--onnx`: Path to QAT onnx graph.
+* `--saveEngine`: Output filename of TensorRT engine.
+* `--verbose`: Flag to enable verbose logging.
+
+#### 2.2. TensorRT Inference
+Obtain accuracy results on the validation dataset:
+
+```sh
+python infer_engine.py --engine=<path_to_trt_engine> --data_dir=<path_to_tfrecord_val_data> -b=<batch_size>
+```
+
+Arguments:
+- `-e, --engine`: TensorRT engine filename (to load).
+- `-m, --model_name`: Name of the model, needed to choose the appropriate input pre-processing. Options={`resnet_v1` (default), `resnet_v2`, `efficientnet_b0`, `efficientnet_b3`, `mobilenet_v1`, `mobilenet_v2`}.
+- `-d, --data_dir`: Path to directory of input images in **tfrecord format** (`data["validation"]`).
+- `-k, --top_k_value` (default=1): Value of `K` for the top-K predictions used in the accuracy calculation.
+- `-b, --batch_size` (default=1): Number of inputs to send in parallel (up to max batch size of engine).
+- `--log_file`: Filename to save logs.
+
+Outputs:
+- `.log` file: contains the engine's performance accuracy. 
+
+# Additional resources
+
+The following resources provide a deeper understanding about Quantization aware training, TF2ONNX and importing a model into TensorRT using Python.
+
+**Quantization Aware Training**
+
+* <a href="https://developer.nvidia.com/blog/achieving-fp32-accuracy-for-int8-inference-using-quantization-aware-training-with-tensorrt/">Achieving FP32 Accuracy for INT8 Inference Using Quantization Aware Training with NVIDIA TensorRT</a>
+
+- [Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference](https://arxiv.org/pdf/1712.05877.pdf)
+- [Quantization Aware Training guide](https://www.tensorflow.org/model_optimization/guide/quantization/training)
+- [Deep Residual Learning for Image Recognition](https://arxiv.org/pdf/1512.03385.pdf)
+
+**Parsers**
+
+- [TF2ONNX Converter](https://github.com/onnx/tensorflow-onnx)
+- [ONNX Parser](https://docs.nvidia.com/deeplearning/sdk/tensorrt-api/python_api/parsers/Onnx/pyOnnx.html)
+
+**Documentation**
+
+- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working With TensorRT Using The Python API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#python_topics)
+- [Importing A Model Using A Parser In Python](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#import_model_python)
+- [NVIDIA’s TensorRT Documentation](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html)

tools/tensorflow-quantization/examples/data/data_loader.py

@@ -0,0 +1,350 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+#
+# Copyright 2018 & 2016 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+
+"""
+Changes made by NVIDIA (2022):
+    - Added: load_data_tfrecord_tf() and load_data() functions
+    - Modified preprocess_image_record(): preprocess_image() + tfrecord data deserialization + decode jpeg
+    - Updated global constants with supported models: _DEFAULT_IMAGE_SIZE and _RESIZE_MIN
+
+About this file: Standalone script for ImageNet TFRecord data loading and input image pre-processing for supported
+    models. Follows TensorFlow's codebase data_loading + pre-processing workflow.
+
+Important links:
+    - TF's codebase:
+    https://github.com/tensorflow/models/blob/master/official/legacy/image_classification/resnet/imagenet_preprocessing.py
+    - Deserialize tfrecord:
+    https://github.com/tensorflow/models/blob/master/official/vision/dataloaders/tf_example_decoder.py
+"""
+
+import os
+import tensorflow as tf
+import PIL.Image
+import numpy as np
+from typing import Dict, Union
+
+_SUPPORTED_MODEL_NAMES = [
+    "resnet_v1",
+    "resnet_v2",
+    "efficientnet_b0",
+    "efficientnet_b3",
+    "mobilenet_v1",
+    "mobilenet_v2",
+    "inception_v3",
+]
+
+_NUM_CLASSES = 1000
+_NUM_IMAGES = {
+    "train": 1281167,
+    "validation": 50000,
+}
+
+_DEFAULT_IMAGE_SIZE = {
+    "resnet_v1": 224,
+    "resnet_v2": 299,
+    "efficientnet_b0": 224,
+    "efficientnet_b3": 300,
+    "mobilenet_v1": 224,
+    "mobilenet_v2": 224,
+    "inception_v3": 299,
+}
+_NUM_CHANNELS = 3
+_RESIZE_MIN = {
+    "resnet_v1": 256,
+    "resnet_v2": 342,
+    "efficientnet_b0": 256,
+    "efficientnet_b3": 342,
+    "mobilenet_v1": 256,
+    "mobilenet_v2": 256,
+    "inception_v3": 342,
+}
+
+
+def load_image_np(test_image, model_name: str = "resnet_v1"):
+    # Image is loaded in NHWC format
+    image_np = np.asarray(PIL.Image.open(test_image).convert('RGB'))
+    image = tf.constant(image_np)
+    image = _aspect_preserving_resize(image, _RESIZE_MIN[model_name])
+    image = _central_crop(image, _DEFAULT_IMAGE_SIZE[model_name], _DEFAULT_IMAGE_SIZE[model_name])
+    image = preprocess_model_func(image, model_name)
+    return image
+
+
+def get_filenames(
+    data_dir: str,
+    is_training: bool = False,
+    num_train_files: int = 1024,
+    num_val_files: int = 128,
+):
+    """
+    Returns filenames for dataset.
+
+    Args:
+        data_dir (str): directory where data is stored.
+        is_training (bool): indicates whether to return the 'train' (True) or 'validation' (False) data filenames.
+        num_train_files (int): number of tfrecord shards available for training.
+        num_val_files (int): number of tfrecord shards available for validation.
+
+    Returns:
+        List: list of shards filenames to compose the dataset.
+    """
+    if is_training:
+        return [
+            # Example: train-00000-of-01024
+            os.path.join(data_dir, "train-{:05d}-of-{:05d}".format(i, num_train_files))
+            for i in range(num_train_files)
+        ]
+    else:
+        return [
+            os.path.join(
+                data_dir, "validation-{:05d}-of-{:05d}".format(i, num_val_files)
+            )
+            for i in range(num_val_files)
+        ]
+
+
+def _deserialize_image_record(record):
+    feature_map = {
+        "image/encoded": tf.io.FixedLenFeature([], tf.string, ""),
+        "image/class/label": tf.io.FixedLenFeature([], tf.int64, -1),
+        "image/class/text": tf.io.FixedLenFeature([], tf.string, ""),
+        "image/object/bbox/xmin": tf.io.VarLenFeature(dtype=tf.float32),
+        "image/object/bbox/ymin": tf.io.VarLenFeature(dtype=tf.float32),
+        "image/object/bbox/xmax": tf.io.VarLenFeature(dtype=tf.float32),
+        "image/object/bbox/ymax": tf.io.VarLenFeature(dtype=tf.float32),
+    }
+    with tf.name_scope("deserialize_image_record"):
+        obj = tf.io.parse_single_example(record, feature_map)
+        imgdata = obj["image/encoded"]
+        label = tf.cast(obj["image/class/label"], tf.int32)
+        bbox = tf.stack(
+            [
+                obj["image/object/bbox/%s" % x].values
+                for x in ["ymin", "xmin", "ymax", "xmax"]
+            ]
+        )
+        bbox = tf.transpose(tf.expand_dims(bbox, 0), [0, 2, 1])
+        text = obj["image/class/text"]
+
+        return imgdata, label, bbox, text
+
+
+def _aspect_preserving_resize(image: tf.Tensor, resize_min: Union[int, tf.Tensor]):
+    """Resize images preserving the original aspect ratio.
+
+    Args:
+        image (tf.Tensor): A 3-D image `Tensor`.
+        resize_min (int): A python integer or scalar `Tensor` indicating the size of the smallest side after resize.
+
+    Returns:
+        resized_image (tf.Tensor): A 3-D `Tensor` containing the resized image.
+    """
+    shape = tf.shape(image)
+    height, width = shape[0], shape[1]
+
+    new_height, new_width = _smallest_size_at_least(height, width, resize_min)
+
+    resized_image = tf.image.resize(
+        image, [new_height, new_width], method=tf.image.ResizeMethod.BILINEAR
+    )
+
+    return resized_image
+
+
+def _smallest_size_at_least(height, width, resize_min):
+    resize_min = tf.cast(resize_min, tf.float32)
+
+    # Convert to floats to make subsequent calculations go smoothly.
+    height, width = tf.cast(height, tf.float32), tf.cast(width, tf.float32)
+
+    smaller_dim = tf.minimum(height, width)
+    scale_ratio = resize_min / smaller_dim
+
+    # Convert back to ints to make heights and widths that TF ops will accept.
+    new_height = tf.cast(height * scale_ratio, tf.int32)
+    new_width = tf.cast(width * scale_ratio, tf.int32)
+
+    return new_height, new_width
+
+
+def _central_crop(image, crop_height, crop_width):
+    shape = tf.shape(image)
+    height, width = shape[0], shape[1]
+
+    amount_to_be_cropped_h = height - crop_height
+    crop_top = amount_to_be_cropped_h // 2
+    amount_to_be_cropped_w = width - crop_width
+    crop_left = amount_to_be_cropped_w // 2
+    return tf.slice(image, [crop_top, crop_left, 0], [crop_height, crop_width, -1])
+
+
+def preprocess_image_record(record, min_size=256, image_height=224, image_width=224):
+    """
+    This function performs image cropping so all images in the dataset have the same height and width dimensions.
+        No value pre-processing is done here.
+    """
+
+    imgdata, label, _, _ = _deserialize_image_record(record)
+    # Subtract one so that ImageNet labels are in [0, 1000). This assumes your dataset contains 'background' as 0.
+    label -= 1
+
+    try:
+        image = tf.image.decode_jpeg(
+            imgdata,
+            channels=_NUM_CHANNELS,
+            fancy_upscaling=False,
+            dct_method="INTEGER_FAST",
+        )
+    except:
+        image = tf.image.decode_image(imgdata, channels=_NUM_CHANNELS)
+
+    image = tf.cast(image, tf.float32)
+    image = _aspect_preserving_resize(image, min_size)
+    image = _central_crop(image, image_height, image_width)
+
+    return image, label
+
+
+def preprocess_model_func(image: tf.Tensor, model_name: str = "resnet_v1"):
+
+    if model_name == "resnet_v1":
+        return tf.keras.applications.resnet.preprocess_input(image)
+    elif model_name == "resnet_v2":
+        return tf.keras.applications.resnet_v2.preprocess_input(image)
+    elif model_name == "mobilenet_v1":
+        return tf.keras.applications.mobilenet.preprocess_input(image)
+    elif model_name == "mobilenet_v2":
+        return tf.keras.applications.mobilenet_v2.preprocess_input(image)
+    elif model_name == "inception_v3":
+        return tf.keras.applications.inception_v3.preprocess_input(image)
+    else:
+        # efficientnet doesn't need specific pre-processing (included in the model itself).
+        print("No further pre-processing found for {}".format(model_name))
+
+    return image
+
+
+def load_data_tfrecord_tf(
+    data_dir: str = "./data/imagenet",
+    batch_size: int = 8,
+    num_train_files: int = 1024,
+    num_val_files: int = 128,
+    model_name: str = "resnet_v1",
+) -> Dict[str, tf.data.Dataset]:
+    """
+    Load ImageNet with TensorFlow Datasets (TFDS).
+
+    Args:
+        data_dir (str): directory where data is stored.
+        batch_size (int): batch_size for dataloader.
+        num_train_files (int): number of tfrecord shards available for training.
+        num_val_files (int): number of tfrecord shards available for validation.
+        model_name (str): Model name, used to decide which input pre-processing is needed.
+            Options={supported_model_names}.
+
+    Returns:
+        dataset_dict (Dict[str, tf.data.Dataset]): dictionary with 'train' and 'validation' datasets.
+
+    Raises:
+        ValueError: raised if 'model_name' is not supported.
+    """.format(
+        supported_model_names=_SUPPORTED_MODEL_NAMES
+    )
+
+    # 1. Load ImageNet2012 train dataset - needs to manually download the full ImageNet2012 dataset first.
+    assert os.path.exists(data_dir)
+    if model_name not in _SUPPORTED_MODEL_NAMES:
+        raise ValueError(
+            "Invalid model name ",
+            model_name,
+            " provided. Please select among {}".format(_SUPPORTED_MODEL_NAMES),
+        )
+
+    # 2. Make train/validation datasets
+    dataset_dict = {}
+    for key, is_training in zip(["train", "validation"], [True, False]):
+        filenames = get_filenames(
+            data_dir,
+            is_training=is_training,
+            num_train_files=num_train_files,
+            num_val_files=num_val_files,
+        )
+        dataset = tf.data.TFRecordDataset(filenames)
+
+        # Image cropping and resizing
+        if model_name in _DEFAULT_IMAGE_SIZE and model_name in _RESIZE_MIN:
+            dataset = dataset.map(
+                lambda record: preprocess_image_record(
+                    record,
+                    min_size=_RESIZE_MIN[model_name],
+                    image_height=_DEFAULT_IMAGE_SIZE[model_name],
+                    image_width=_DEFAULT_IMAGE_SIZE[model_name],
+                )
+            )
+        else:
+            dataset = dataset.map(preprocess_image_record)
+
+        dataset = dataset.map(lambda image, label: (preprocess_model_func(image, model_name), label))
+        # Divide dataset into batches
+        dataset = dataset.batch(batch_size, drop_remainder=True)
+        dataset_dict[key] = dataset
+
+    return dataset_dict
+
+
+def load_data(
+    hyperparams: Dict, model_name: str = "resnet_v1"
+) -> [tf.data.Dataset, tf.data.Dataset]:
+    """ Loads ImageNet data in `tfrecord` format (requires manual data download).
+
+    Args:
+        hyperparams (Dict): dictionary with necessary hyper-parameters for data loading.
+        model_name (str): Model name, used to decide which input pre-processing is needed.
+            Options={supported_model_names}.
+
+    Returns:
+        train_batches (tf.data.Dataset): 'train' dataset.
+        val_batches (tf.data.Dataset): 'validation' dataset.
+    """.format(
+        supported_model_names=_SUPPORTED_MODEL_NAMES
+    )
+
+    data_batches = load_data_tfrecord_tf(
+        data_dir=hyperparams["tfrecord_data_dir"],
+        batch_size=hyperparams["batch_size"],
+        model_name=model_name,
+    )
+    train_batches, val_batches = (data_batches["train"], data_batches["validation"])
+    if hyperparams["train_data_size"] is not None:
+        train_batches = train_batches.take(hyperparams["train_data_size"])
+    if hyperparams["val_data_size"] is not None:
+        val_batches = val_batches.take(hyperparams["val_data_size"])
+
+    return train_batches, val_batches

tools/tensorflow-quantization/examples/data/imagenet_data_setup.sh

@@ -0,0 +1,44 @@
+export IMAGENET_HOME=/media/Data/imagenet_data
+# Setup folders
+mkdir -p $IMAGENET_HOME/validation
+mkdir -p $IMAGENET_HOME/train
+
+# ###### Modification 1: set .tar files path to $IMAGENET_HOME #############
+# Extract validation and training
+tar xf $IMAGENET_HOME/ILSVRC2012_img_val.tar -C $IMAGENET_HOME/validation
+tar xf $IMAGENET_HOME/ILSVRC2012_img_train.tar -C $IMAGENET_HOME/train
+# ##########################################################################
+
+# Extract and then delete individual training tar files This can be pasted
+# directly into a bash command-line or create a file and execute.
+cd $IMAGENET_HOME/train
+
+for f in *.tar; do
+  d=`basename $f .tar`
+  mkdir $d
+  tar xf $f -C $d
+done
+
+cd $IMAGENET_HOME # Move back to the base folder
+
+# [Optional] Delete tar files if desired as they are not needed
+rm $IMAGENET_HOME/train/*.tar
+
+# ###### Modification 2: Updated deprecated link #############
+# Download labels file.
+wget -O $IMAGENET_HOME/synset_labels.txt \
+https://raw.githubusercontent.com/tensorflow/models/master/research/slim/datasets/imagenet_2012_validation_synset_labels.txt
+# ############################################################
+
+# Process the files. Remember to get the script from github first. The TFRecords
+# will end up in the --local_scratch_dir. To upload to gcs with this method
+# leave off `nogcs_upload` and provide gcs flags for project and output_path.
+python imagenet_to_gcs.py \
+  --raw_data_dir=$IMAGENET_HOME \
+  --local_scratch_dir=$IMAGENET_HOME/tf_records \
+  --nogcs_upload
+
+# ######## Modification 3: move train and validation files to root dir #######################
+mv $IMAGENET_HOME/tf_records/train* $IMAGENET_HOME/tf_records
+mv $IMAGENET_HOME/tf_records/validation* $IMAGENET_HOME/tf_records
+# ############################################################################################

tools/tensorflow-quantization/examples/data/test_data_loader.py

@@ -0,0 +1,168 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+#
+
+
+"""
+This module contains test cases for our data loader, which contains data loading and pre-processing functions
+    for the ImageNet2012 dataset in 'tfrecord' format.
+NOTE: the user needs to manually download the full ImageNet2012 dataset first.
+"""
+
+import tensorflow as tf
+from examples.data.data_loader import load_data
+import numpy as np
+from collections import defaultdict
+from typing import Dict
+import pytest
+
+DATA_HYPERPARAMS = {
+    "tfrecord_data_dir": "/media/Data/ImageNet/train-val-tfrecord",
+    "batch_size": 64,
+    "train_data_size": 100,  # If 'None', consider all data, otherwise, consider subset.
+    "val_data_size": 100,  # If 'None', consider all data, otherwise, consider subset.
+}
+
+
+def load_tfrecord_mean_min_max(model_name: str) -> [Dict, Dict, Dict]:
+    """
+    Loads `tfrecord` dataset and calculates the data's mean, min, and max values.
+
+    Args:
+        model_name (str): model name for data pre-processing.
+
+    Returns:
+        total_mean_dict: dictionary with MEAN values in 'R', 'G', 'B'.
+        total_min_dict: dictionary with MIN values in 'R', 'G', 'B'.
+        total_max_dict: dictionary with MAX values in 'R', 'G', 'B'.
+    """
+    # 1. Data loading
+    train_batches, val_batches = load_data(
+        hyperparams=DATA_HYPERPARAMS, model_name=model_name
+    )
+    assert isinstance(train_batches, tf.data.Dataset) and isinstance(
+        val_batches, tf.data.Dataset
+    )
+
+    # 2. Test input preprocessing
+    mean = defaultdict(list)
+    min = defaultdict(list)
+    max = defaultdict(list)
+    for batch in [train_batches]:
+        for examples in batch:
+            image, label = examples
+            image_dict = defaultdict()
+            for i, c in zip([0, 1, 2], ["R", "G", "B"]):
+                image_dict[c] = image[:, :, :, i]
+                mean[c].append(tf.math.reduce_mean(image_dict[c]))
+                min[c].append(tf.math.reduce_min(image_dict[c]))
+                max[c].append(tf.math.reduce_max(image_dict[c]))
+
+    total_mean_dict = defaultdict(list)
+    total_min_dict = defaultdict(list)
+    total_max_dict = defaultdict(list)
+    for c in ["R", "G", "B"]:
+        total_mean_dict[c] = tf.math.reduce_mean(mean[c])
+        total_min_dict[c] = tf.math.reduce_min(min[c])
+        total_max_dict[c] = tf.math.reduce_max(max[c])
+    return total_mean_dict, total_min_dict, total_max_dict
+
+
+def test_imagenet_tfrecord_efficientnetb0():
+    """
+    Tests data loading and pre-processing for EfficientNet-B0.
+    Note that EfficientNet doesn't have any input pre-processing methods besides image resizing and cropping.
+      See `data_loader.preprocess_image_record()`.
+    """
+    print("------------ EfficientNet-B0 pre-processing test -------------")
+
+    # 1. Data loading and get mean, max, min of data without input preprocessing
+    total_mean_dict, total_min_dict, total_max_dict = load_tfrecord_mean_min_max(
+        model_name="efficientnet_b0"
+    )
+
+    # 2. Check if mean is as expected and max/min values
+    mean_RGB = [123.68, 116.779, 103.939]
+    mean_RGB_obtained = list(
+        total_mean_dict.values()
+    )  # [total_mean_dict['R'], total_mean_dict['G'], total_mean_dict['B']]
+    mean_diff = abs(np.array(mean_RGB_obtained) - np.array(mean_RGB))
+    print("  Expected mean (RGB): {}".format(mean_RGB))
+    print("  Calculated mean (RGB): {}".format(np.array(mean_RGB_obtained)))
+    print("  Difference: {}".format(np.array(mean_diff)))
+    assert (mean_diff <= 3.0).all()
+
+    # 3. Values expected to be between 0 and 255
+    total_min = min(total_min_dict["R"], min(total_min_dict["G"], total_min_dict["B"]))
+    total_max = max(total_max_dict["R"], max(total_max_dict["G"], total_max_dict["B"]))
+    assert total_min >= 0.0 and total_max <= 255.0
+
+
+def test_imagenet_tfrecord_resnetv1():
+    """Tests data loading and pre-processing for ResNetv1.
+
+    ResNetv1 input pre-processing:
+        - Resizing + cropping
+        - "The images are converted from RGB to BGR, then each color channel is zero-centered with respect to the
+        ImageNet dataset, without scaling."
+        - Zero-center: (data - mean(data) / std(data)) -> In this case, std(data) = None.
+    """
+    print("------------ ResNetv1 pre-processing test -------------")
+
+    # 1. Data loading and get mean, max, min of data without input preprocessing
+    total_mean_dict, total_min_dict, total_max_dict = load_tfrecord_mean_min_max(
+        model_name="resnet_v1"
+    )
+
+    # 2.1 Data should be zero-centered (mean=0)
+    mean_RGB_obtained = list(total_mean_dict.values())
+    print("  Expected mean (RGB): [0, 0, 0]")
+    print("  Calculated mean (RGB): {}".format(np.array(mean_RGB_obtained)))
+    print("  Min values: {}".format(np.array(list(total_min_dict.values()))))
+    print("  Max values: {}".format(np.array(list(total_max_dict.values()))))
+    assert (abs(np.array(mean_RGB_obtained)) <= 3.0).all()
+
+    # 2.2 No scaling, meaning values are between -255 and 255 (after zero-centering)
+    assert (np.array(list(total_min_dict.values())) >= -255.0).all()
+    assert (np.array(list(total_max_dict.values())) <= 255.0).all()
+
+
+def test_imagenet_tfrecord_resnetv2():
+    """Tests data loading and pre-processing for ResNetv2.
+
+    ResNetv2 input pre-processing:
+        - Resizing + cropping
+        - "The inputs pixel values are scaled between -1 and 1, sample-wise."
+        - Sample-wise normalization: https://stackoverflow.com/questions/37625272/keras-batchnormalization-what-exactly-is-sample-wise-normalization
+
+    MobileNet-v1/v2 input pre-processing:
+        - "The inputs pixel values are scaled between -1 and 1, sample-wise."
+        - This is the same as ResNet-v2, with the difference that the MobileNet model takes input shape 224x224x3
+          (same as ResNet-v1).
+    """
+    print("------------ ResNetv2 pre-processing test -------------")
+
+    # 1. Data loading and get mean, max, min of data without input preprocessing
+    total_mean_dict, total_min_dict, total_max_dict = load_tfrecord_mean_min_max(
+        model_name="resnet_v2"
+    )
+
+    # 2. Check that values are between -1 and 1
+    print("  Min values: {}".format(np.array(list(total_min_dict.values()))))
+    print("  Max values: {}".format(np.array(list(total_max_dict.values()))))
+    print("  Mean values: {}".format(np.array(list(total_mean_dict.values()))))
+    assert (np.array(list(total_min_dict.values())) >= -1.0).all()
+    assert (np.array(list(total_max_dict.values())) <= 1.0).all()

tools/tensorflow-quantization/examples/efficientnet/README.md

@@ -0,0 +1,102 @@
+## About
+This script presents a QAT end-to-end workflow (TF2-to-ONNX) for [EfficientNet](https://github.com/tensorflow/models/tree/master/official/legacy/image_classification/efficientnet).
+
+### Contents
+[Requirements](#requirements) • [Workflow](#workflow) • [Results](#results)  
+
+## Requirements
+1. Install base requirements and prepare data. Please refer to [examples' README](../README.md).
+
+2. Clone the models from Tensorflow model garden:
+
+```
+git clone https://github.com/tensorflow/models.git
+pushd models && git checkout tags/v2.8.0 && popd
+export PYTHONPATH=$PWD/models:$PYTHONPATH
+pip install -r models/official/requirements.txt
+```
+
+> cd models && git submodule init && git submodule update
+
+3. Download pretrained checkpoints:
+   1. B0: https://tfhub.dev/tensorflow/efficientnet/b0/classification/1
+   2. B3: https://tfhub.dev/tensorflow/efficientnet/b3/classification/1
+
+
+## Workflow
+
+### Step 1: Model Quantization and Fine-tuning
+
+* In `run_qat_workflow.py`, please set the `pretrained_ckpt_path` field to the directory of the downloaded checkpoint to start fine-tuning with QAT. All the required hyper-parameters can be set in the `HYPERPARAMS` dictionary.
+
+  Please run the following to quantize, fine-tune, and save the final graph in SavedModel format.
+
+  ```sh
+  python run_qat_workflow.py
+  ```
+  > Update `MODEL_VERSION` to the EfficientNet version you wish to quantize.
+
+### Step 2: Exporting a QAT SavedModel
+
+Once you've fine-tuned the QAT model, export it by running
+
+```sh
+python export.py --ckpt <path_to_pretrained_ckpt> --output <saved_model_output_name> --model_version b0
+```
+
+This script applies quantization to the model, restores the checkpoint, and exports it in a SavedModel format. This script will generate `eff`  which is a directory containing saved model. We set the overall graph data format to `NCHW` by using `tf.keras.backend.set_image_data_format('channels_first')`. TensorRT expects `NCHW` format for graphs trained with QAT for better optimizations.
+
+Arguments:
+
+* `--ckpt` : Path to fine-tuned QAT checkpoint to be loaded.
+* `--output` : Name of output TF saved model.
+* `--model_version` : EfficientNet model version, currently supports {`b0`, `b3`}.
+
+### Step 3: Conversion to ONNX
+
+Convert the saved model into ONNX by running
+
+```sh
+python -m tf2onnx.convert --saved-model <path_to_saved_model> --output model_qat.onnx  --opset 13
+```
+
+By default, tf2onnx uses TF's graph optimizers to performs constant folding after a saved model is loaded.
+
+Arguments:
+
+* `--saved-model` : Name of TF SavedModel
+* `--output` : Name of ONNX output graph
+* `--opset` : ONNX opset version (opset 13 or higher must be used)
+
+### Step 4: TensorRT Deployment
+Please refer to the [examples' README](../README.md).
+
+## Results
+
+This section presents the validation accuracy for the full ImageNet dataset on NVIDIA's A100 GPU and TensorRT 8.4 GA.
+
+### EfficientNet-B0
+| Model                 | Accuracy (%) | Latency (ms, bs=1) |
+|-----------------------|--------------|--------------------|
+| Baseline (TensorFlow) | 76.97        | 6.77               |
+| PTQ (TensorRT)        | 71.71        | 0.67               |
+| **QAT** (TensorRT)    | 75.82        | 0.68               |
+> QAT fine-tuning hyper-parameters: `bs64, ep10, lr=0.001, steps_per_epoch=None`
+
+### EfficientNet-B3
+| Model                 | Accuracy (%) | Latency (ms, bs=1) |
+|-----------------------|--------------|--------------------|
+| Baseline (TensorFlow) | 81.36        | 10.33              |
+| PTQ (TensorRT)        | 78.88        | 1.24               |
+| **QAT** (TensorRT)    | 79.48        | 1.23               |
+> QAT fine-tuning hyper-parameters: `bs=32, ep20, steps_per_epoch=None, lr=0.0001`
+
+### Notes
+- QAT fine-tuning hyper-parameters:
+  - Optimizer: `piecewise_sgd`, `lr_schedule=[(1.0, 1), (0.1, 3), (0.01, 6), (0.001, 9), (0.001, 15)]`.
+  - Other hyper-parameters are under each model's results table.
+- PTQ calibration: `bs=64`
+- EfficientNet model quantization:
+  - QDQ nodes added in Residual connection (fix added to ResidualQDQCustomCase for `Conv-BN-Activation-Dropout` pattern),
+  - Global Average Pooling,
+  - Multiply layer in SE block.

tools/tensorflow-quantization/examples/efficientnet/export.py

@@ -0,0 +1,59 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 tensorflow as tf
+import argparse
+from utils import create_efficientnet_model
+from tensorflow_quantization.quantize import quantize_model
+from tensorflow_quantization.custom_qdq_cases import EfficientNetQDQCase
+
+
+def export_saved_model(model_version="b0"):
+    model = create_efficientnet_model(model_version=model_version)
+    q_model = quantize_model(model, custom_qdq_cases=[EfficientNetQDQCase()])
+    if args.ckpt:
+        q_model.load_weights(args.ckpt).expect_partial()
+
+    tf.keras.models.save_model(q_model, args.output)
+    print("Exported the model to {}".format(args.output))
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Export saved model for efficientnet_b0"
+    )
+    parser.add_argument(
+        "--ckpt",
+        type=str,
+        default="qat/checkpoints_best",
+        help="Path to pretrained QAT efficientnet checkpoint.",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="qat/saved_model",
+        help="Path to pretrained QAT saved model.",
+    )
+    parser.add_argument(
+        "--model_version",
+        type=str,
+        default="b0",
+        help="EfficientNet model version, currently supports {'b0', 'b3'}.",
+    )
+    args = parser.parse_args()
+    export_saved_model(args.model_version)

tools/tensorflow-quantization/examples/efficientnet/run_qat_workflow.py

@@ -0,0 +1,116 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 os
+import tensorflow as tf
+
+from tensorflow_quantization.quantize import quantize_model
+from tensorflow_quantization.custom_qdq_cases import EfficientNetQDQCase
+
+from examples.data.data_loader import load_data
+from examples.utils_finetuning import fine_tune
+
+import numpy as np
+import random
+from utils import create_efficientnet_model
+
+MODEL_VERSION = "b0"  # Options={b0, b3}
+
+HYPERPARAMS = {
+    # ################ Data loading ################
+    "tfrecord_data_dir": "/media/Data/imagenet_data/tf_records",
+    "batch_size": 64,
+    "train_data_size": None,  # If 'None', consider all data, otherwise, consider subset.
+    "val_data_size": None,  # If 'None', consider all data, otherwise, consider subset.
+    # ############## Fine-tuning ##################
+    "pretrained_ckpt_path": "./weights/efficientnet_{}/baseline".format(MODEL_VERSION),
+    "epochs": 10,
+    "steps_per_epoch": None,  # 'None' if you want to use the default number of steps. If you use this, make sure the number of steps is <= the number of shards (total number of samples / batch_size). Otherwise, an error will occur.
+    "base_lr": 0.001,
+    "optimizer": "piecewise_sgd",  # Options={sgd, piecewise_sgd, adam}
+    "save_ckpt_dir": "./weights/efficientnet_{}/qat".format(MODEL_VERSION),
+    # ############## Enable/disable tasks ##################
+    "evaluate_baseline_model": True,
+    "evaluate_qat_model": True,
+    "seed": 42,
+}
+
+# Set seed for reproducible results
+os.environ["PYTHONHASHSEED"] = str(HYPERPARAMS["seed"])
+random.seed(HYPERPARAMS["seed"])
+np.random.seed(HYPERPARAMS["seed"])
+tf.random.set_seed(HYPERPARAMS["seed"])
+
+
+def evaluate_acc(model, validation_data):
+    # Compile model (needed to evaluate model)
+    model.compile(
+        optimizer="sgd",
+        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        metrics=["sparse_categorical_accuracy"],
+    )
+
+    _, val_accuracy = model.evaluate(validation_data)
+
+    return val_accuracy
+
+
+def main():
+    # ------------- Initial settings -------------
+    # Load data
+    train_batches, val_batches = load_data(HYPERPARAMS, model_name="efficientnet_"+MODEL_VERSION)
+
+    # ------------- Baseline model -------------
+    model = create_efficientnet_model(model_version=MODEL_VERSION)
+
+    # Load pre-trained weights
+    if HYPERPARAMS["pretrained_ckpt_path"]:
+        model.load_weights(HYPERPARAMS["pretrained_ckpt_path"]).expect_partial()
+
+    if HYPERPARAMS["evaluate_baseline_model"]:
+        baseline_model_accuracy = evaluate_acc(model, val_batches)
+        print("Baseline model accuracy:", baseline_model_accuracy)
+
+    # ------------- QAT model -------------
+    # Quantize model
+    q_model = quantize_model(model, custom_qdq_cases=[EfficientNetQDQCase()])
+
+    # Fine-tuning + saving new checkpoints
+    print("Fine-tuning and saving QAT model checkpoint...")
+    lr_schedule_array = [(1.0, 1), (0.1, 3), (0.01, 6), (0.001, 9), (0.001, 15)]
+    fine_tune(
+        q_model, train_batches, val_batches,
+        qat_save_finetuned_weights=HYPERPARAMS["save_ckpt_dir"],
+        hyperparams=HYPERPARAMS,
+        lr_schedule_array=lr_schedule_array,
+        enable_tensorboard_callback=False
+    )
+    print("Fine-tuning done!")
+
+    # Loads best weights if they exist
+    best_checkpoint_path = os.path.join(HYPERPARAMS["save_ckpt_dir"], "checkpoints_best")
+    if os.path.exists(best_checkpoint_path + ".index"):
+        q_model.load_weights(best_checkpoint_path).expect_partial()
+
+    if HYPERPARAMS["evaluate_qat_model"]:
+        print("\nEvaluating QAT model...")
+        qat_model_accuracy = evaluate_acc(q_model, val_batches)
+        print("QAT val accuracy:", qat_model_accuracy)
+
+
+if __name__ == "__main__":
+    main()

tools/tensorflow-quantization/examples/efficientnet/test_qdq_node_placement.py

@@ -0,0 +1,74 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 sys
+import tensorflow as tf
+from tensorflow_quantization.custom_qdq_cases import EfficientNetQDQCase
+from utils import create_efficientnet_model
+from tests.onnx_graph_qdq_validator import validate_quantized_model
+from tensorflow_quantization.utils import CreateAssetsFolders
+import pytest
+
+# Create a directory to save test models
+test_assets = CreateAssetsFolders("test_qdq_node_placement")
+
+
+def test_efficientnet_b0_quantize_full():
+    """
+    EfficientNet-B0: Full model quantization.
+
+    Contains special patterns connected to Add layer:
+        1. (Conv->BatchNorm->Activation)->Add
+        2. (Conv->BatchNorm->Activation->Dropout)->Add
+    Previously, only (Conv)->Add and (Conv->BatchNorm)->Add checks were done when checking for quantizable Residual
+        connections (branches to `Add` layer). The new EfficientNet patterns have now also been added to the
+        ResidualCustomQDQCases check.
+    """
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Instantiate Baseline model
+    nn_model_original = create_efficientnet_model()
+
+    custom_qdq_cases = [EfficientNetQDQCase()]
+    q_model, validated = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name,
+        custom_qdq_cases=custom_qdq_cases
+    )
+    assert validated, "ONNX QDQ validation for full network quantization failed!"
+    # necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()
+
+
+def test_efficientnet_b3_quantize_full():
+    """
+    EfficientNet-B3: Full model quantization.
+    Contains the same special patterns as EfficientNet-B0.
+    """
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Instantiate Baseline model
+    nn_model_original = create_efficientnet_model(model_version="b3")
+
+    custom_qdq_cases = [EfficientNetQDQCase()]
+    q_model, validated = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name,
+        custom_qdq_cases=custom_qdq_cases
+    )
+    assert validated, "ONNX QDQ validation for full network quantization failed!"
+    # necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()

tools/tensorflow-quantization/examples/efficientnet/utils.py

@@ -0,0 +1,44 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 os
+import sys
+import tensorflow as tf
+tf_models_path = os.path.realpath("./models")
+sys.path.insert(1, tf_models_path)
+try:
+    from official.legacy.image_classification.efficientnet import efficientnet_model
+except Exception:
+    print("Error importing TF official models codebase.")
+
+
+def create_efficientnet_model(model_version="b0"):
+    model_name = "efficientnet-" + model_version
+    model_configs = dict(efficientnet_model.MODEL_CONFIGS)
+    assert model_name in model_configs, "Model name is not valid!"
+    config = model_configs[model_name]
+
+    # Set the dataformat of the model to NCHW for training and inference
+    tf.keras.backend.set_image_data_format("channels_first")
+    # B0=(224, 224, 3); B3=(300, 300, 3)
+    image_input = tf.keras.layers.Input(
+        shape=(config.resolution, config.resolution, config.input_channels), name="image_input", dtype=tf.float32
+    )
+    outputs = efficientnet_model.efficientnet(image_input, config)
+    model = tf.keras.Model(inputs=image_input, outputs=outputs)
+
+    return model

tools/tensorflow-quantization/examples/inception/README.md

@@ -0,0 +1,43 @@
+## About
+This script presents a QAT end-to-end workflow (TF2-to-ONNX) for [Inception models](https://keras.io/api/applications/inceptionv3/) in `tf.keras.applications`.
+
+### Contents
+[Requirements](#requirements) • [Workflow](#workflow) • [Results](#results)  
+
+## Requirements
+Install base requirements and prepare data. Please refer to [examples' README](../README.md).
+
+## Workflow
+
+### Step 1: Model Quantization and Fine-tuning
+> Similar to [ResNet](../resnet): different model and different input pre-processing.
+
+Please run the following to quantize, fine-tune, and save the final graph in SavedModel format (checkpoints are also saved).
+
+```sh
+python run_qat_workflow.py
+```
+
+### Step 2: Conversion to ONNX
+Step 1 already does the conversion from SavedModel to ONNX automatically. For manual steps, please see step 3 in [EfficientNet's README](../efficientnet_b0/README.md).
+
+### Step 3: TensorRT Deployment
+Please refer to the [examples' README](../README.md).
+
+## Results
+Results obtained on NVIDIA's A100 GPU and TensorRT 8.4.2.4 (GA Update 1).
+
+### Inception-v3
+
+| Model    | TF (%) | TF latency (ms, bs=1) | TRT(%) | TRT latency (ms, bs=1) |
+|----------|--------|-----------------------|--------|------------------------|
+| Baseline | 77.86  | 9.01                  | 77.86  | 1.39                   |
+| PTQ      | -      | -                     | 77.73  | 0.82                   |
+| **QAT**  | 78.11  | 101.97                | 78.08  | 0.82                   |
+
+### Notes
+- Optimization: MaxPool needs to be quantized to trigger horizontal fusion in Concat layer.
+- QAT fine-tuning hyper-params:
+  - Optimizer: `piecewise_sgd`, `lr_schedule=[(1.0, 1), (0.1, 2), (0.01, 7)]` (default)
+  - Hyper-parameters: `bs=64, ep=10, lr=0.001, steps_per_epoch=500`
+- PTQ calibration: `bs=64`.

tools/tensorflow-quantization/examples/inception/run_qat_workflow.py

@@ -0,0 +1,178 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 os
+import tensorflow as tf
+
+from tensorflow_quantization.quantize import quantize_model
+from tensorflow_quantization.utils import convert_saved_model_to_onnx
+from tensorflow_quantization.custom_qdq_cases import InceptionQDQCase
+
+from examples.utils import ensure_dir, get_tfkeras_model
+from examples.data.data_loader import load_data
+from examples.utils_finetuning import (
+    get_finetuned_weights_dirname,
+    fine_tune,
+    compile_model,
+)
+import gc
+import numpy as np
+import random
+import sys
+import logging
+
+
+MODEL_NAME = "inception_v3"  # Options=[inception_v3]
+
+HYPERPARAMS = {
+    # ################ Data loading ################
+    "tfrecord_data_dir": "/media/Data/imagenet_data/tf_records",
+    "batch_size": 64,
+    "train_data_size": None,  # Only for `tfrecord`. If None, consider all data, otherwise, consider subset.
+    "val_data_size": None,  # Only for `tfrecord`. If None, consider all data, otherwise, consider subset.
+    # ############## Fine-tuning ##################
+    "epochs": 10,
+    "steps_per_epoch": 500,  # 500,  # 'None' if you want to use the default number of steps. If you use this, make sure the number of steps is <= the number of shards (total number of samples / batch_size). Otherwise, an error will occur.
+    "base_lr": 0.001,  # 0.0001
+    "optimizer": "piecewise_sgd",  # Options={sgd, piecewise_sgd, adam}
+    "save_root_dir": "./weights/{}".format(
+        MODEL_NAME
+    ),  # DIR is updated to reflect hyperparams
+    # ############## Enable/disable tasks ##################
+    "finetune_qat_model": True,  # If True, finetune QAT model. Otherwise, just quantize and load weights if existent.
+    "rewrite_weights_qat_finetuning": True,  # If True, rewrites existing fine-tuned weights. Otherwise, just load weights if they exist.
+    "evaluate_baseline_model": True,
+    "evaluate_qat_model": True,
+    "save_baseline_model": True,
+    "seed": 42,
+}
+
+# Set seed for reproducible results
+os.environ["PYTHONHASHSEED"] = str(HYPERPARAMS["seed"])
+random.seed(HYPERPARAMS["seed"])
+np.random.seed(HYPERPARAMS["seed"])
+tf.random.set_seed(HYPERPARAMS["seed"])
+
+# Create logger and save to out.log
+LOGGER = logging.getLogger()
+LOGGER.setLevel(logging.INFO)
+
+
+def main():
+    # ------------- Initial settings -------------
+    # Create directory to save the fine-tuned weights + add relevant hyperparameters in the name
+    qat_save_finetuned_weights = get_finetuned_weights_dirname(HYPERPARAMS)
+    ensure_dir(qat_save_finetuned_weights)
+
+    # Add terminal and file handlers to logger
+    output_file_handler = logging.FileHandler(
+        os.path.join(qat_save_finetuned_weights, "out.log"), mode="w"
+    )
+    stdout_handler = logging.StreamHandler(sys.stdout)
+    LOGGER.addHandler(output_file_handler)
+    LOGGER.addHandler(stdout_handler)
+
+    # Load data
+    train_batches, val_batches = load_data(HYPERPARAMS, model_name=MODEL_NAME)
+
+    # ------------- Baseline model -------------
+    LOGGER.info("------------- Baseline model -------------")
+
+    # Instantiate Baseline model
+    model = get_tfkeras_model(model_name=MODEL_NAME)
+
+    if HYPERPARAMS["evaluate_baseline_model"]:
+        # Compile model (needed to evaluate model)
+        compile_model(model)
+        _, baseline_model_accuracy = model.evaluate(val_batches)
+        LOGGER.info("Baseline val accuracy: {}".format(baseline_model_accuracy))
+
+    if HYPERPARAMS["save_baseline_model"]:
+        tf.keras.models.save_model(
+            model, os.path.join(HYPERPARAMS["save_root_dir"], "saved_model_baseline")
+        )
+        convert_saved_model_to_onnx(
+            saved_model_dir=os.path.join(
+                HYPERPARAMS["save_root_dir"], "saved_model_baseline"
+            ),
+            onnx_model_path=os.path.join(
+                HYPERPARAMS["save_root_dir"], "model_baseline.onnx"
+            ),
+        )
+
+    # ------------- QAT model -------------
+    # Quantize model
+    LOGGER.info("\n------------- QAT model -------------")
+    q_model = quantize_model(model, custom_qdq_cases=[InceptionQDQCase()])
+    # q_model = quantize_model(model)
+
+    finetuned_qat_weights_path = os.path.join(
+        qat_save_finetuned_weights, "checkpoints_best"
+    )
+    # Performs fine-tuning if `rewrite` is enabled or if fine-tuned weights don't exist yet
+    #   (1st time fine-tuning model).
+    if HYPERPARAMS["finetune_qat_model"] and (
+        HYPERPARAMS["rewrite_weights_qat_finetuning"]
+        or not os.path.exists(finetuned_qat_weights_path)
+    ):
+        # Fine-tuning + saving new checkpoints
+        LOGGER.info("\nFine-tuning model...")
+        fine_tune(
+            q_model,
+            train_batches,
+            val_batches,
+            qat_save_finetuned_weights,
+            HYPERPARAMS,
+            LOGGER,
+        )
+        LOGGER.info("Fine-tuning done!")
+
+    # Loads best weights if they exist
+    if os.path.exists(finetuned_qat_weights_path + ".index"):
+        LOGGER.info("Loading fine-tuned weights...")
+        q_model.load_weights(finetuned_qat_weights_path).expect_partial()
+        LOGGER.info("Loaded complete!")
+    compile_model(q_model)
+
+    if HYPERPARAMS["evaluate_qat_model"]:
+        LOGGER.info("\nEvaluating QAT model...")
+        _, qat_model_accuracy = q_model.evaluate(val_batches)
+        LOGGER.info("QAT val accuracy: {}".format(qat_model_accuracy))
+
+    # Save quantized model
+    LOGGER.info("\nSaving QAT model")
+    tf.keras.models.save_model(
+        q_model, os.path.join(qat_save_finetuned_weights, "saved_model")
+    )
+
+    # Clear GPU and invoke Garbage Collector to avoid script ending during ONNX conversion
+    tf.keras.backend.clear_session()
+    gc.collect()
+    del model
+    del q_model
+
+    # Convert SavedModel to ONNX
+    LOGGER.info("\nONNX conversion...")
+    convert_saved_model_to_onnx(
+        saved_model_dir=os.path.join(qat_save_finetuned_weights, "saved_model"),
+        onnx_model_path=os.path.join(qat_save_finetuned_weights, "model.onnx"),
+    )
+
+
+if __name__ == "__main__":
+    main()

tools/tensorflow-quantization/examples/inception/test_qdq_node_placement.py

@@ -0,0 +1,61 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 sys
+import tensorflow as tf
+from tensorflow_quantization.custom_qdq_cases import InceptionQDQCase
+from examples.utils import get_tfkeras_model
+from tests.onnx_graph_qdq_validator import validate_quantized_model
+from tensorflow_quantization.utils import CreateAssetsFolders
+from tensorflow_quantization.quantize import LayerConfig
+import pytest
+
+# Create a directory to save test models
+test_assets = CreateAssetsFolders("test_qdq_node_placement")
+
+EXPECTED_QDQ_INSERTION = [
+    LayerConfig(name="Conv2D", is_keras_class=True),
+    LayerConfig(name="Dense", is_keras_class=True),
+    LayerConfig(name="DepthwiseConv2D", is_keras_class=True),
+    LayerConfig(
+        name="AveragePooling2D", is_keras_class=True, quantize_weight=False
+    ),
+    LayerConfig(
+        name="GlobalAveragePooling2D", is_keras_class=True, quantize_weight=False
+    )
+]
+
+
+def test_inceptionv3_quantize_full():
+    """
+    Inception-v3: Full model quantization
+    """
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Instantiate Baseline model
+    nn_model_original = get_tfkeras_model(model_name="inception_v3")
+
+    custom_qdq_cases = [InceptionQDQCase()]
+    q_model, validated = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name,
+        custom_qdq_cases=custom_qdq_cases,
+        expected_qdq_insertion=EXPECTED_QDQ_INSERTION
+    )
+    assert validated, "ONNX QDQ validation for full network quantization failed!"
+    # necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()

tools/tensorflow-quantization/examples/infer_engine.py

@@ -0,0 +1,276 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 os
+import argparse
+import numpy as np
+import tensorrt as trt
+import pycuda.driver as cuda
+
+# If you face the following issue:
+#  "pycuda._driver.LogicError: explicit_context_dependent failed: invalid device context - no currently active context?"
+#  Add "import pycuda.autoinit", this is needed to initialize cuda!
+import pycuda.autoinit
+import tensorflow as tf
+from examples.data.data_loader import load_data_tfrecord_tf, load_image_np, _SUPPORTED_MODEL_NAMES
+
+TRT_DYNAMIC_DIM = -1
+
+
+class HostDeviceMem(object):
+    """Simple helper data class to store Host and Device memory."""
+
+    def __init__(self, host_mem, device_mem):
+        self.host = host_mem
+        self.device = device_mem
+
+    def __str__(self):
+        return "Host:\n" + str(self.host) + "\nDevice:\n" + str(self.device)
+
+    def __repr__(self):
+        return self.__str__()
+
+
+def allocate_buffers(engine: trt.ICudaEngine, batch_size: int) -> [list, list, list]:
+    """
+    Function to allocate buffers and bindings for TensorRT inference.
+
+    Args:
+        engine (trt.ICudaEngine):
+        batch_size (int): batch size to be used during inference.
+
+    Returns:
+        inputs (List): list of input buffers.
+        outputs (List): list of output buffers.
+        dbindings (List): list of device bindings.
+    """
+    inputs = []
+    outputs = []
+    dbindings = []
+
+    for binding in engine:
+        binding_shape = engine.get_binding_shape(binding)
+        if binding_shape[0] == TRT_DYNAMIC_DIM:  # dynamic shape
+            size = batch_size * abs(trt.volume(binding_shape))
+        else:
+            size = abs(trt.volume(binding_shape))
+        dtype = trt.nptype(engine.get_binding_dtype(binding))
+        # Allocate host and device buffers
+        host_mem = cuda.pagelocked_empty(size, dtype)
+        device_mem = cuda.mem_alloc(host_mem.nbytes)
+        # Append the device buffer to device bindings
+        dbindings.append(int(device_mem))
+
+        # Append to the appropriate list (input/output)
+        if engine.binding_is_input(binding):
+            inputs.append(HostDeviceMem(host_mem, device_mem))
+        else:
+            outputs.append(HostDeviceMem(host_mem, device_mem))
+
+    return inputs, outputs, dbindings
+
+
+def infer(
+    engine_path: str,
+    val_batches,
+    batch_size: int = 8,
+    top_k_value: int = 1,
+) -> None:
+    """
+    Performs inference in TensorRT engine.
+
+    Args:
+        engine_path (str): path to the TensorRT engine.
+        val_batches (tf.data.Dataset): validation dataset (batches).
+        batch_size (int): batch size used for inference and dataset batch splitting.
+        top_k_value (int): value of `K` for the top K predictions used in the accuracy calculation.
+
+    Raises:
+        RuntimeError: raised when loading images in the host fails.
+    """
+
+    def override_shape(shape: tuple) -> tuple:
+        """Overrides batch dimension if dynamic."""
+        if TRT_DYNAMIC_DIM in shape:
+            shape = tuple(
+                [batch_size if dim == TRT_DYNAMIC_DIM else dim for dim in shape]
+            )
+        return shape
+
+    # Open engine as runtime
+    with open(engine_path, "rb") as f, trt.Runtime(
+        trt.Logger(trt.Logger.ERROR)
+    ) as runtime:
+        engine = runtime.deserialize_cuda_engine(f.read())
+
+        # Allocate buffers and create a CUDA stream.
+        inputs, outputs, dbindings = allocate_buffers(engine, batch_size)
+
+        # Initiate test_accuracy
+        test_accuracy = tf.keras.metrics.SparseTopKCategoricalAccuracy(
+            k=top_k_value, name="top_k_accuracy", dtype=tf.float32
+        )
+        test_accuracy.reset_states()
+
+        # Contexts are used to perform inference.
+        with engine.create_execution_context() as context:
+
+            # Resolves dynamic shapes in the context
+            for binding in engine:
+                binding_idx = engine.get_binding_index(binding)
+                binding_shape = engine.get_binding_shape(binding_idx)
+                if engine.binding_is_input(binding_idx):
+                    binding_shape = override_shape(binding_shape)
+                    context.set_binding_shape(binding_idx, binding_shape)
+
+            if isinstance(val_batches, tf.Tensor):
+                # Load images in Host (flatten and copy to page-locked buffer in Host)
+                data = val_batches.numpy().astype(np.float32).ravel()
+                pagelocked_buffer = inputs[0].host
+                np.copyto(pagelocked_buffer, data)
+                inp = inputs[0]
+                # Transfer input data from Host to Device (GPU)
+                cuda.memcpy_htod(inp.device, inp.host)
+                # Run inference
+                context.execute_v2(dbindings)
+                # Transfer predictions back to Host from GPU
+                out = outputs[0]
+                cuda.memcpy_dtoh(out.host, out.device)
+
+                softmax_output = np.array(out.host)
+                top1_idx = np.argmax(softmax_output)
+                output_confidence = softmax_output[top1_idx]
+                print("Top-1 Index of the image : {} Confidence: {}".format(top1_idx, output_confidence))
+
+            elif isinstance(val_batches, tf.data.Dataset):
+                # Loop over number of steps to evaluate entire validation dataset
+                for step, example in enumerate(val_batches):
+                    images, labels = example
+                    if step % 100 == 0 and step != 0:
+                        print(
+                            "Evaluating batch {}: {:.4f}".format(
+                                step, test_accuracy.result()
+                            )
+                        )
+                    try:
+                        # Load images in Host (flatten and copy to page-locked buffer in Host)
+                        data = images.numpy().astype(np.float32).ravel()
+                        pagelocked_buffer = inputs[0].host
+                        np.copyto(pagelocked_buffer, data)
+                    except RuntimeError:
+                        raise RuntimeError(
+                            "Failed to load images in Host at step {}".format(step)
+                        )
+
+                    inp = inputs[0]
+                    # Transfer input data from Host to Device (GPU)
+                    cuda.memcpy_htod(inp.device, inp.host)
+                    # Run inference
+                    context.execute_v2(dbindings)
+                    # Transfer predictions back to Host from GPU
+                    out = outputs[0]
+                    cuda.memcpy_dtoh(out.host, out.device)
+
+                    # Split 1-D output of length N*labels into 2-D array of (N, labels)
+                    batch_outs = np.array(np.split(np.array(out.host), batch_size))
+                    # Update test accuracy
+                    test_accuracy.update_state(labels, batch_outs)
+
+                # Print final accuracy and save to log file
+                print("\n======================================\n")
+                result_str = "Top-{} accuracy: {:.4f}\n".format(
+                    top_k_value, test_accuracy.result()
+                )
+                print(result_str)
+                # Save logs to file
+                results_dir = "/".join(args.engine.split("/")[:-1])
+                with open(os.path.join(results_dir, args.log_file), "w") as log_file:
+                    log_file.write(result_str)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Run inference on TensorRT engines for Imagenet-based Classification models."
+    )
+    parser.add_argument(
+        "-e", "--engine", type=str, required=True, help="Path to TensorRT engine"
+    )
+    parser.add_argument(
+        "--image", type=str, help="Path to an image to perform single image inference"
+    )
+    parser.add_argument(
+        "-m",
+        "--model_name",
+        type=str,
+        default="resnet_v1",
+        help="Name of the model, needed to choose the appropriate input pre-processing."
+        "Options include {}".format(_SUPPORTED_MODEL_NAMES),
+    )
+    parser.add_argument(
+        "-d",
+        "--data_dir",
+        default="/media/Data/ImageNet/train-val-tfrecord",
+        type=str,
+        help="Path to directory of input images in tfrecord format (val data).",
+    )
+    parser.add_argument(
+        "-k",
+        "--top_k_value",
+        default=1,
+        type=int,
+        help="Value of `K` for the top-K predictions used in the accuracy calculation.",
+    )
+    parser.add_argument(
+        "-b",
+        "--batch_size",
+        default=1,
+        type=int,
+        help="Number of inputs to send in parallel (up to max batch size of engine).",
+    )
+    parser.add_argument(
+        "--log_file",
+        type=str,
+        default="engine_accuracy.log",
+        help="Filename to save logs.",
+    )
+    args = parser.parse_args()
+
+    if args.model_name not in _SUPPORTED_MODEL_NAMES:
+        raise ValueError(
+            "Invalid model name ",
+            args.model_name,
+            " provided. Please select among {}".format(_SUPPORTED_MODEL_NAMES),
+        )
+
+    # Load the test data and pre-process input
+    val_batches = None
+    if args.image:
+        val_batches = load_image_np(args.image, args.model_name)
+    else:
+        data_batches = load_data_tfrecord_tf(
+            data_dir=args.data_dir, batch_size=args.batch_size, model_name=args.model_name
+        )
+        val_batches = data_batches["validation"]
+
+    # Perform inference
+    infer(
+        args.engine,
+        val_batches,
+        batch_size=args.batch_size,
+        top_k_value=args.top_k_value,
+    )

tools/tensorflow-quantization/examples/mobilenet/README.md

@@ -0,0 +1,55 @@
+## About
+This script presents a QAT end-to-end workflow (TF2-to-ONNX) for [MobileNet models](https://keras.io/api/applications/mobilenet/) in `tf.keras.applications`.
+
+### Contents
+[Requirements](#requirements) • [Workflow](#workflow) • [Results](#results)  
+
+## Requirements
+Install base requirements and prepare data. Please refer to [examples' README](../README.md).
+
+## Workflow
+
+### Step 1: Model Quantization and Fine-tuning
+> Similar to [ResNet](../resnet): different model and different input pre-processing (`mobilenet`).
+
+Please run the following to quantize, fine-tune, and save the final graph in SavedModel format (checkpoints are also saved).
+
+```sh
+python run_qat_workflow.py
+```
+
+### Step 2: Conversion to ONNX
+Step 1 already does the conversion from SavedModel to ONNX automatically. For manual steps, please see step 3 in [EfficientNet's README](../efficientnet_b0/README.md).
+
+### Step 3: TensorRT Deployment
+Please refer to the [examples' README](../README.md).
+
+## Results
+Results obtained on NVIDIA's A100 GPU and TensorRT 8.4.10.1.
+
+### MobileNet-v1
+
+| Model    | TF (%)      | TF latency (ms, bs=1) | TRT(%) | TRT latency (ms, bs=1) |
+|----------|-------------|-----------------------|--------|------------------------|
+| Baseline | 70.60       | 1.99                  | 70.60  | 0.32                   |
+| PTQ      | -           | -                     | 69.31  | 0.16                   |
+| **QAT**  | 70.51 (ep2) | 50.49                 | 70.43  | 0.16                   |
+
+**Note**: no residual connections exist in MobileNet-v1.
+
+### MobileNet-v2
+
+| Model    | TF (%)      | TF latency (ms, bs=1) | TRT(%)   | TRT latency (ms, bs=1) |
+|----------|-------------|-----------------------|----------|------------------------|
+| Baseline | 71.77       | 3.71                  | 71.77    | 0.55                   |
+| PTQ      | -           | -                     | 70.87    | 0.30                   |
+| **QAT**  | 71.68 (ep1) | 74.27                 | 71.62    | 0.30                   |
+
+**Note**: residual connections exist in MobileNet-v2.
+
+### Notes
+- QAT fine-tuning hyper-parameters:
+  - Optimizer: `piecewise_sgd`, `lr_schedule=[(1.0, 1), (0.1, 2), (0.01, 7)]` (default)
+  - Hyper-parameters: `bs=64, ep=10, lr=0.001`
+- PTQ calibration: `bs=64`.
+- MobileNet-v3 might not show good acceleration in TensorRT due to its architecture (`Conv->BN->((Add->Clip->Mul), ())->Mul`), which is not a kernel fusion in TRT. 

tools/tensorflow-quantization/examples/mobilenet/run_qat_workflow.py

@@ -0,0 +1,179 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 os
+import tensorflow as tf
+
+from tensorflow_quantization.quantize import quantize_model
+from tensorflow_quantization.utils import convert_saved_model_to_onnx
+from tensorflow_quantization.custom_qdq_cases import MobileNetQDQCase
+
+from examples.utils import ensure_dir
+from examples.data.data_loader import load_data
+from examples.utils import get_tfkeras_model
+from examples.utils_finetuning import (
+    get_finetuned_weights_dirname,
+    fine_tune,
+    compile_model,
+)
+import gc
+import numpy as np
+import random
+import sys
+import logging
+
+
+MODEL_NAME = "mobilenet_v2"  # Options=[mobilenet_v1, mobilenet_v2]
+
+HYPERPARAMS = {
+    # ################ Data loading ################
+    "tfrecord_data_dir": "/media/Data/imagenet_data/tf_records",
+    "batch_size": 64,
+    "train_data_size": None,  # Only for `tfrecord`. If None, consider all data, otherwise, consider subset.
+    "val_data_size": None,  # Only for `tfrecord`. If None, consider all data, otherwise, consider subset.
+    # ############## Fine-tuning ##################
+    "epochs": 10,
+    "steps_per_epoch": 500,  # 500,  # 'None' if you want to use the default number of steps. If you use this, make sure the number of steps is <= the number of shards (total number of samples / batch_size). Otherwise, an error will occur.
+    "base_lr": 0.001,  # 0.0001
+    "optimizer": "piecewise_sgd",  # Options={sgd, piecewise_sgd, adam}
+    "save_root_dir": "./weights/{}".format(
+        MODEL_NAME
+    ),  # DIR is updated to reflect hyperparams
+    # ############## Enable/disable tasks ##################
+    "finetune_qat_model": True,  # If True, finetune QAT model. Otherwise, just quantize and load weights if existent.
+    "rewrite_weights_qat_finetuning": True,  # If True, rewrites existing fine-tuned weights. Otherwise, just load weights if they exist.
+    "evaluate_baseline_model": True,
+    "evaluate_qat_model": True,
+    "save_baseline_model": False,
+    "seed": 42,
+}
+
+# Set seed for reproducible results
+os.environ["PYTHONHASHSEED"] = str(HYPERPARAMS["seed"])
+random.seed(HYPERPARAMS["seed"])
+np.random.seed(HYPERPARAMS["seed"])
+tf.random.set_seed(HYPERPARAMS["seed"])
+
+# Create logger and save to out.log
+LOGGER = logging.getLogger()
+LOGGER.setLevel(logging.INFO)
+
+
+def main():
+    # ------------- Initial settings -------------
+    # Create directory to save the fine-tuned weights + add relevant hyperparameters in the name
+    qat_save_finetuned_weights = get_finetuned_weights_dirname(HYPERPARAMS)
+    ensure_dir(qat_save_finetuned_weights)
+
+    # Add terminal and file handlers to logger
+    output_file_handler = logging.FileHandler(
+        os.path.join(qat_save_finetuned_weights, "out.log"), mode="w"
+    )
+    stdout_handler = logging.StreamHandler(sys.stdout)
+    LOGGER.addHandler(output_file_handler)
+    LOGGER.addHandler(stdout_handler)
+
+    # Load data
+    # MobileNet requires the input to be pre-processed like ResNet-v2 but with input 224x224x3 (as in ResNet-v1)
+    train_batches, val_batches = load_data(HYPERPARAMS, model_name=MODEL_NAME)
+
+    # ------------- Baseline model -------------
+    LOGGER.info("------------- Baseline model -------------")
+
+    # Instantiate Baseline model
+    model = get_tfkeras_model(model_name=MODEL_NAME)
+
+    if HYPERPARAMS["evaluate_baseline_model"]:
+        # Compile model (needed to evaluate model)
+        compile_model(model)
+        _, baseline_model_accuracy = model.evaluate(val_batches)
+        LOGGER.info("Baseline val accuracy: {}".format(baseline_model_accuracy))
+
+    if HYPERPARAMS["save_baseline_model"]:
+        tf.keras.models.save_model(
+            model, os.path.join(HYPERPARAMS["save_root_dir"], "saved_model_baseline")
+        )
+        convert_saved_model_to_onnx(
+            saved_model_dir=os.path.join(
+                HYPERPARAMS["save_root_dir"], "saved_model_baseline"
+            ),
+            onnx_model_path=os.path.join(
+                HYPERPARAMS["save_root_dir"], "model_baseline.onnx"
+            ),
+        )
+
+    # ------------- QAT model -------------
+    # Quantize model
+    LOGGER.info("\n------------- QAT model -------------")
+    q_model = quantize_model(model, custom_qdq_cases=[MobileNetQDQCase()])
+
+    finetuned_qat_weights_path = os.path.join(
+        qat_save_finetuned_weights, "checkpoints_best"
+    )
+    # Performs fine-tuning if `rewrite` is enabled or if fine-tuned weights don't exist yet
+    #   (1st time fine-tuning model).
+    if HYPERPARAMS["finetune_qat_model"] and (
+        HYPERPARAMS["rewrite_weights_qat_finetuning"]
+        or not os.path.exists(finetuned_qat_weights_path)
+    ):
+        # Fine-tuning + saving new checkpoints
+        LOGGER.info("\nFine-tuning model...")
+        fine_tune(
+            q_model,
+            train_batches,
+            val_batches,
+            qat_save_finetuned_weights,
+            HYPERPARAMS,
+            LOGGER,
+        )
+        LOGGER.info("Fine-tuning done!")
+
+    # Loads best weights if they exist
+    if os.path.exists(finetuned_qat_weights_path + ".index"):
+        LOGGER.info("Loading fine-tuned weights...")
+        q_model.load_weights(finetuned_qat_weights_path).expect_partial()
+        LOGGER.info("Loaded complete!")
+    compile_model(q_model)
+
+    if HYPERPARAMS["evaluate_qat_model"]:
+        LOGGER.info("\nEvaluating QAT model...")
+        _, qat_model_accuracy = q_model.evaluate(val_batches)
+        LOGGER.info("QAT val accuracy: {}".format(qat_model_accuracy))
+
+    # Save quantized model
+    LOGGER.info("\nSaving QAT model")
+    tf.keras.models.save_model(
+        q_model, os.path.join(qat_save_finetuned_weights, "saved_model")
+    )
+
+    # Clear GPU and invoke Garbage Collector to avoid script ending during ONNX conversion
+    tf.keras.backend.clear_session()
+    gc.collect()
+    del model
+    del q_model
+
+    # Convert SavedModel to ONNX
+    LOGGER.info("\nONNX conversion...")
+    convert_saved_model_to_onnx(
+        saved_model_dir=os.path.join(qat_save_finetuned_weights, "saved_model"),
+        onnx_model_path=os.path.join(qat_save_finetuned_weights, "model.onnx"),
+    )
+
+
+if __name__ == "__main__":
+    main()

tools/tensorflow-quantization/examples/mobilenet/test_qdq_node_placement.py

@@ -0,0 +1,66 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 sys
+import tensorflow as tf
+from tensorflow_quantization.custom_qdq_cases import MobileNetQDQCase
+from examples.utils import get_tfkeras_model
+from tests.onnx_graph_qdq_validator import validate_quantized_model
+from tensorflow_quantization.utils import CreateAssetsFolders
+import pytest
+
+# Create a directory to save test models
+test_assets = CreateAssetsFolders("test_qdq_node_placement")
+
+
+def test_mobilenetv1_quantize_full():
+    """
+    MobileNet-v1: Full model quantization
+    """
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Instantiate Baseline model
+    nn_model_original = get_tfkeras_model(model_name="mobilenet_v1")
+
+    custom_qdq_cases = [MobileNetQDQCase()]
+    q_model, validated = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name,
+        custom_qdq_cases=custom_qdq_cases
+    )
+    assert validated, "ONNX QDQ validation for full network quantization failed!"
+    # necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()
+
+
+def test_mobilenetv2_quantize_full():
+    """
+    MobileNet-v2: Full model quantization
+    """
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Instantiate Baseline model
+    nn_model_original = get_tfkeras_model(model_name="mobilenet_v2")
+
+    custom_qdq_cases = [MobileNetQDQCase()]
+    q_model, validated = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name,
+        custom_qdq_cases=custom_qdq_cases
+    )
+    assert validated, "ONNX QDQ validation for full network quantization failed!"
+    # necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()

tools/tensorflow-quantization/examples/requirements.txt

@@ -0,0 +1,9 @@
+# ImageNet: conversion to tfrecord
+gcloud
+google-cloud-storage
+
+# Data loader
+Pillow
+
+# TensorRT: build/infer engine
+pycuda

tools/tensorflow-quantization/examples/resnet/README.md

@@ -0,0 +1,67 @@
+## About
+This script presents a QAT end-to-end workflow (TF2-to-ONNX) for [ResNet models](https://keras.io/api/applications/resnet/) in `tf.keras.applications`.
+
+### Contents
+[Requirements](#requirements) • [Workflow](#workflow) • [Results](#results)  
+
+## Requirements
+Install base requirements and prepare data. Please refer to [examples' README](../README.md).
+
+## Workflow
+
+### Step 1: Model Quantization and Fine-tuning
+
+Please run the following to quantize, fine-tune, and save the final graph in SavedModel format (checkpoints are also saved).
+
+```sh
+python run_qat_workflow.py
+```
+
+### Step 2: Conversion to ONNX
+Step 1 already does the conversion from SavedModel to ONNX automatically. For manual steps, please see step 3 in [EfficientNet's README](../efficientnet_b0/README.md).
+
+### Step 3: TensorRT Deployment
+Please refer to the [examples' README](../README.md).
+
+## Results
+Results obtained on NVIDIA's A100 GPU and TensorRT 8.4 EA.
+
+### ResNet50-v1
+
+| Model    | TF (%)      | TF latency (ms, bs=1) | TRT(%) | TRT latency (ms, bs=1) |
+|----------|-------------|-----------------------|--------|------------------------|
+| Baseline | 75.05       | 7.95                  | 75.05  | 1.96                   |
+| PTQ      | -           | -                     | 74.96  | 0.46                   |
+| **QAT**  | 75.11 (ep5) | -                     | 75.12  | 0.45                   |
+
+### ResNet50-v2
+
+| Model    | TF (%)       | TF latency (ms, bs=1) | TRT(%)  | TRT latency (ms, bs=1) |
+|----------|--------------|-----------------------|---------|------------------------|
+| Baseline | 75.36        | 6.16                  | 75.37   | 2.35                   |
+| PTQ      | -            | -                     | 75.48   | 0.57                   |
+| **QAT**  | 75.59 (ep5)  | -                     | 75.65   | 0.57                   |
+
+### ResNet101-v1
+
+| Model    | TF (%)       | TF latency (ms, bs=1) | TRT(%) | TRT latency (ms, bs=1) |
+|----------|--------------|-----------------------|--------|------------------------|
+| Baseline | 76.47        | 15.92                 | 76.48  | 3.84                   |
+| PTQ      | -            | -                     | 76.32  | 0.84                   |
+| **QAT**  | 76.33 (ep30) | -                     | 76.26  | 0.84                   |
+
+### ResNet101-v2
+
+| Model    | TF (%) | TF latency (ms, bs=1) | TRT(%) | TRT latency (ms, bs=1) |
+|----------|--------|-----------------------|--------|------------------------|
+| Baseline | 76.89  | 14.13                 | 76.88  | 4.55                   |
+| PTQ      | -      | -                     | 76.94  | 1.05                   |
+| **QAT**  | 77.20  | -                     | 77.15  | 1.05                   |
+> QAT fine-tuning hyper-parameters for ResNet101-v2: `bs=32` (`bs=64` was OOM).
+
+### Notes
+- QAT fine-tuning hyper-parameters:
+  - Optimizer: `piecewise_sgd`, `lr_schedule=[(1.0,1),(0.1,2),(0.01,7)]` (default)
+  - Hyper-parameters: `bs=64, ep=10, lr=0.001`.
+  - Added QDQ nodes in Residual connection.
+- PTQ calibration: `bs=64`.

tools/tensorflow-quantization/examples/resnet/run_qat_workflow.py

@@ -0,0 +1,188 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 os
+import tensorflow as tf
+
+from tensorflow_quantization.quantize import quantize_model
+from tensorflow_quantization.utils import convert_saved_model_to_onnx
+from tensorflow_quantization.custom_qdq_cases import ResNetV1QDQCase, ResNetV2QDQCase
+
+from examples.utils import ensure_dir
+from examples.data.data_loader import load_data
+from examples.utils_finetuning import (
+    get_finetuned_weights_dirname,
+    fine_tune,
+    compile_model,
+)
+from utils import get_resnet_model
+import gc
+import numpy as np
+import random
+import sys
+import logging
+
+
+RESNET_DEPTH = "50"  # Options=[50, 101]
+RESNET_VERSION = "v1"  # Options=[v1, v2]
+
+# For reference: the baseline model used bs=256, train_epochs=90, train_steps=10.
+HYPERPARAMS = {
+    # ################ Data loading ################
+    "tfrecord_data_dir": "/media/Data/imagenet_data/tf_records",
+    "batch_size": 64,
+    "train_data_size": None,  # If None, consider all data, otherwise, consider subset.
+    "val_data_size": None,  # If None, consider all data, otherwise, consider subset.
+    # ############## Fine-tuning ##################
+    "epochs": 2,
+    "steps_per_epoch": 500,  # Set as 'None' if you want to use the default number of steps. If you use this, make sure the number of steps is <= the number of shards (total number of samples / batch_size). Otherwise, an error will occur.
+    "base_lr": 0.001,
+    "optimizer": "piecewise_sgd",  # Options={sgd, piecewise_sgd, adam}
+    "save_root_dir": "./weights/resnet{}{}_test".format(
+        RESNET_DEPTH, RESNET_VERSION
+    ),  # DIR is updated to reflect hyperparams
+    # ############## Enable/disable tasks ##################
+    "finetune_qat_model": True,  # If True, finetune QAT model. Otherwise, just quantize and load weights if existent.
+    "rewrite_weights_qat_finetuning": True,  # If True, rewrites existing fine-tuned weights. Otherwise, just load weights if they exist.
+    "evaluate_baseline_model": True,
+    "evaluate_qat_model": True,
+    "save_baseline_model": True,
+    "seed": 42,
+}
+
+# Set seed for reproducible results
+os.environ["PYTHONHASHSEED"] = str(HYPERPARAMS["seed"])
+random.seed(HYPERPARAMS["seed"])
+np.random.seed(HYPERPARAMS["seed"])
+tf.random.set_seed(HYPERPARAMS["seed"])
+
+# Create logger and save to out.log
+LOGGER = logging.getLogger()
+LOGGER.setLevel(logging.INFO)
+
+
+def main():
+    # ------------- Initial settings -------------
+    # Create directory to save the fine-tuned weights + add relevant hyperparameters in the name
+    qat_save_finetuned_weights = get_finetuned_weights_dirname(HYPERPARAMS)
+    ensure_dir(qat_save_finetuned_weights)
+
+    # Add terminal and file handlers to logger
+    output_file_handler = logging.FileHandler(
+        os.path.join(qat_save_finetuned_weights, "out.log"), mode="w"
+    )
+    stdout_handler = logging.StreamHandler(sys.stdout)
+    LOGGER.addHandler(output_file_handler)
+    LOGGER.addHandler(stdout_handler)
+
+    # Load data
+    train_batches, val_batches = load_data(
+        HYPERPARAMS, model_name="resnet_{}".format(RESNET_VERSION)
+    )
+
+    # ------------- Baseline model -------------
+    LOGGER.info("------------- Baseline model -------------")
+
+    # Instantiate Baseline model
+    model = get_resnet_model(
+        resnet_depth=RESNET_DEPTH,
+        resnet_version=RESNET_VERSION,
+    )
+
+    if HYPERPARAMS["evaluate_baseline_model"]:
+        # Compile model (needed to evaluate model)
+        compile_model(model)
+        _, baseline_model_accuracy = model.evaluate(val_batches)
+        LOGGER.info("Baseline val accuracy: {}".format(baseline_model_accuracy))
+
+    if HYPERPARAMS["save_baseline_model"]:
+        tf.keras.models.save_model(
+            model, os.path.join(HYPERPARAMS["save_root_dir"], "saved_model_baseline")
+        )
+        convert_saved_model_to_onnx(
+            saved_model_dir=os.path.join(
+                HYPERPARAMS["save_root_dir"], "saved_model_baseline"
+            ),
+            onnx_model_path=os.path.join(
+                HYPERPARAMS["save_root_dir"], "model_baseline.onnx"
+            ),
+        )
+
+    # ------------- QAT model -------------
+    # Quantize model
+    LOGGER.info("\n------------- QAT model -------------")
+    if "v1" == RESNET_VERSION:
+        q_model = quantize_model(model, custom_qdq_cases=[ResNetV1QDQCase()])
+    else:
+        q_model = quantize_model(model, custom_qdq_cases=[ResNetV2QDQCase()])
+
+    finetuned_qat_weights_path = os.path.join(
+        qat_save_finetuned_weights, "checkpoints_best"
+    )
+    # Performs fine-tuning if `rewrite` is enabled or if fine-tuned weights don't exist yet
+    #   (1st time fine-tuning model).
+    if HYPERPARAMS["finetune_qat_model"] and (
+        HYPERPARAMS["rewrite_weights_qat_finetuning"]
+        or not os.path.exists(finetuned_qat_weights_path)
+    ):
+        # Fine-tuning + saving new checkpoints
+        LOGGER.info("\nFine-tuning model...")
+        fine_tune(
+            q_model,
+            train_batches,
+            val_batches,
+            qat_save_finetuned_weights,
+            HYPERPARAMS,
+            LOGGER,
+        )
+        LOGGER.info("Fine-tuning done!")
+
+    # Loads best weights if they exist
+    if os.path.exists(finetuned_qat_weights_path + ".index"):
+        LOGGER.info("Loading fine-tuned weights...")
+        q_model.load_weights(finetuned_qat_weights_path).expect_partial()
+        LOGGER.info("Loaded complete!")
+    compile_model(q_model)
+
+    if HYPERPARAMS["evaluate_qat_model"]:
+        LOGGER.info("\nEvaluating QAT model...")
+        _, qat_model_accuracy = q_model.evaluate(val_batches)
+        LOGGER.info("QAT val accuracy: {}".format(qat_model_accuracy))
+
+    # Save quantized model
+    LOGGER.info("\nSaving QAT model")
+    tf.keras.models.save_model(
+        q_model, os.path.join(qat_save_finetuned_weights, "saved_model")
+    )
+
+    # Clear GPU and invoke Garbage Collector to avoid script ending during ONNX conversion
+    tf.keras.backend.clear_session()
+    gc.collect()
+    del model
+    del q_model
+
+    # Convert SavedModel to ONNX
+    LOGGER.info("\nONNX conversion...")
+    convert_saved_model_to_onnx(
+        saved_model_dir=os.path.join(qat_save_finetuned_weights, "saved_model"),
+        onnx_model_path=os.path.join(qat_save_finetuned_weights, "model.onnx"),
+    )
+
+
+if __name__ == "__main__":
+    main()

tools/tensorflow-quantization/examples/resnet/scripts/README.md

@@ -0,0 +1,22 @@
+## Scripts for TRT Deployment
+For both baseline and QAT, change:
+- `RESNET_DEPTH` for 50 or 101,
+- `RESNET_VERSION` for v1 or v2,
+- `BS` for which batch sizes you wish to evaluate the engine on.
+
+#### Baseline
+```
+./scripts/deploy_engine_baseline.sh
+```
+> Change `ROOT_DIR` to where your ONNX file is.
+
+#### QAT
+```
+./scripts/deploy_engine_qat.sh
+```
+> Change `QAT_SUBDIR` and `ROOT_DIR` to where your ONNX file is.
+
+### Only accuracy
+```
+./scripts/infer_engine.sh
+```

tools/tensorflow-quantization/examples/resnet/scripts/deploy_engine_baseline.sh

@@ -0,0 +1,47 @@
+#!/usr/bin/env bash
+
+# Single run:
+# ../../engine_builder/build_engine_single.py --root_dir=/home/nvidia/PycharmProjects/tensorflow-quantization/examples/resnet/weights/resnet50v1 --onnx=model_baseline_dynamic.onnx --engine=model_baseline_dynamic.engine --input=224,224,3 --min_bs=1 --max_bs=1 --opt_bs=1 --precision=fp32
+#
+
+RESNET_DEPTH=50
+RESNET_VERSION=v1
+ROOT_DIR=../weights/resnet${RESNET_DEPTH}${RESNET_VERSION}
+LOGS_SUBDIR=baseline_engines_trtSource
+LOGS_DIR=${ROOT_DIR}/${LOGS_SUBDIR}
+mkdir $LOGS_DIR
+
+echo "1/3. Building engine"
+# bs=32 OOM in workstation
+ONNX=model_baseline_dynamic.onnx
+ENGINE=${LOGS_SUBDIR}/model_baseline_dynamic_bs{min1,opt8,max16}.engine
+python ../../../engine_builder/build_engine_single.py --root_dir=$ROOT_DIR \
+                                                      --onnx=$ONNX \
+                                                      --engine=$ENGINE \
+                                                      --input=224,224,3 \
+                                                      --min_bs=1 --opt_bs=8 --max_bs=16 \
+                                                      --precision=fp32
+wait
+
+for BS in 8 16;  do # 8 32 128; do
+  echo "Model evaluation..."
+  echo "############### bs=${BS} ###############"
+
+  # Latency calculation from built engine
+  echo "2/3. Latency evaluation"
+  trtexec --device=0 \
+          --loadEngine=${ROOT_DIR}/${ENGINE} \
+          --shapes=input_1:0:${BS}x224x224x3 \
+          --workspace=2048 \
+          --separateProfileRun \
+          --dumpProfile \
+          --explicitBatch &> ${LOGS_DIR}/trtexec_latency_bs${BS}.log
+  wait
+
+  echo "3/3. Accuracy evaluation"
+  python ../infer_engine.py --engine=${ROOT_DIR}/${ENGINE} \
+                            --log_file=engine_accuracy_bs${BS}.log \
+                            --model_name=resnet_$RESNET_VERSION \
+                            -b=$BS
+  wait
+done

tools/tensorflow-quantization/examples/resnet/scripts/deploy_engine_qat.sh

@@ -0,0 +1,49 @@
+#!/usr/bin/env bash
+
+# Single run:
+# ../../engine_builder/build_engine_single.py --root_dir=/home/nvidia/PycharmProjects/tensorflow-quantization/examples/resnet/weights/resnet50v1 --onnx=model_baseline_dynamic.onnx --engine=model_baseline_dynamic.engine --input=224,224,3 --min_bs=1 --max_bs=1 --opt_bs=1 --precision=fp32
+#
+
+RESNET_DEPTH=50
+RESNET_VERSION=v1
+QAT_SUBDIR=qat_tfrecord_ep10_steps500_l2False_baselr0.0001_piecewise_sgd_bs128
+ROOT_DIR=../weights/resnet${RESNET_DEPTH}${RESNET_VERSION}/${QAT_SUBDIR}
+LOGS_SUBDIR=engines_trtSource
+LOGS_DIR=${ROOT_DIR}/${LOGS_SUBDIR}
+mkdir $LOGS_DIR
+
+echo "1/3. Building engine"
+# bs=32 OOM in workstation
+ONNX=model_dynamic.onnx
+ENGINE=${LOGS_SUBDIR}/model_baseline_bs{min1,opt8,max128}.engine
+python ../../../engine_builder/build_engine_single.py --root_dir=$ROOT_DIR \
+                                                      --onnx=$ONNX \
+                                                      --engine=$ENGINE \
+                                                      --input=224,224,3 \
+                                                      --min_bs=1 --opt_bs=8 --max_bs=128 \
+                                                      --precision=int8
+wait
+
+for BS in 1 8 128;  do # 8 32 128; do
+  echo "Model evaluation..."
+  echo "############### bs=${BS} ###############"
+
+  # Latency calculation from built engine
+  echo "2/3. Latency evaluation"
+  trtexec --device=0 \
+          --loadEngine=${ROOT_DIR}/${ENGINE} \
+          --shapes=input_1:0:${BS}x224x224x3 \
+          --workspace=1024 \
+          --separateProfileRun \
+          --dumpProfile \
+          --explicitBatch \
+          --int8 &> ${LOGS_DIR}/trtexec_latency_bs${BS}.log
+  wait
+
+  echo "3/3. Accuracy evaluation"
+  python ../infer_engine.py --engine=${ROOT_DIR}/${ENGINE} \
+                            --log_file=engine_accuracy_bs${BS}.log \
+                            --model_name=resnet_$RESNET_VERSION \
+                            -b=$BS
+  wait
+done

tools/tensorflow-quantization/examples/resnet/scripts/infer_engine.sh

@@ -0,0 +1,17 @@
+ROOT_DIR="/home/nvidia/PycharmProjects/tensorrt_qat/examples/resnet/"
+
+RESNET_DEPTH="50"
+RESNET_VERSION="v1"
+MODEL_TYPE="baseline"  # "qat"
+PRECISION="fp32"  # "int8"
+
+ENGINES_DIR="engines_gtc_trt8.4_gittrt/${MODEL_TYPE}"
+LOGS_DIR="logs_gtc_trt8.4_gittrt/${MODEL_TYPE}"
+
+for BS in 1; do
+  SUBDIR="resnet${RESNET_DEPTH}${RESNET_VERSION}_${PRECISION}_${BS}_sparsity_disable_DLA_disabled"
+
+  python ../infer_engine.py --engine=${ROOT_DIR}/${ENGINES_DIR}/${SUBDIR}.plan \
+                            --log_file=${ROOT_DIR}/${LOGS_DIR}/${SUBDIR}_accuracy.log \
+                            --model_name=resnet_$RESNET_VERSION -b=1
+done

tools/tensorflow-quantization/examples/resnet/test_qdq_node_placement.py

@@ -0,0 +1,113 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 sys
+import tensorflow as tf
+from tensorflow_quantization.quantize import QuantizationSpec
+from tensorflow_quantization.custom_qdq_cases import ResNetV1QDQCase, ResNetV2QDQCase
+from examples.resnet.utils import get_resnet_model
+from tests.onnx_graph_qdq_validator import validate_quantized_model
+from tensorflow_quantization.utils import CreateAssetsFolders
+import pytest
+
+# Create a directory to save test models
+test_assets = CreateAssetsFolders("test_qdq_node_placement")
+
+
+def test_resnet50v1_quantize_full():
+    """
+    ResNet-v1: Full model quantization
+    """
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Instantiate Baseline model
+    nn_model_original = get_resnet_model(resnet_depth="50", resnet_version="v1")
+
+    custom_qdq_cases = [ResNetV1QDQCase()]
+    q_model, validated = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name,
+        custom_qdq_cases=custom_qdq_cases
+    )
+    assert validated, "ONNX QDQ validation for full network quantization failed!"
+    # necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()
+
+
+def test_resnet50v1_quantize_full_special():
+    """
+    ResNet-v1: Full model quantization with the first Conv layer (conv2_block1_1_conv) not being quantized.
+    """
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Instantiate Baseline model
+    nn_model_original = get_resnet_model(resnet_depth="50", resnet_version="v1")
+
+    # Full quantization
+    custom_qdq_cases = [ResNetV1QDQCase()]
+    # Special case
+    qspec = QuantizationSpec()
+    qspec.add(name="conv2_block1_1_conv", quantize_input=False, quantize_weight=False)
+
+    q_model, validated = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name,
+        qspec=qspec, custom_qdq_cases=custom_qdq_cases
+    )
+
+    assert validated, "ONNX QDQ validation for full network quantization failed!"
+    # necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()
+
+
+def test_resnet50v2_quantize_full():
+    """
+    ResNet-v2: Full model quantization
+    """
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Instantiate Baseline model
+    nn_model_original = get_resnet_model(resnet_depth="50", resnet_version="v2")
+
+    custom_qdq_cases = [ResNetV2QDQCase()]
+    q_model, validated = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name,
+        custom_qdq_cases=custom_qdq_cases
+    )
+    assert validated, "ONNX QDQ validation for full network quantization failed!"
+    # necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()
+
+
+def test_resnet50v2_quantize_full_special():
+    """
+    ResNet-v2: Full model quantization with the first MaxPool layer (pool1_pool) not being quantized.
+    """
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Instantiate Baseline model
+    nn_model_original = get_resnet_model(resnet_depth="50", resnet_version="v2")
+
+    custom_qdq_cases = [ResNetV2QDQCase()]
+    qspec = QuantizationSpec()
+    qspec.add(name="pool1_pool", quantize_input=False, quantize_weight=False)
+    q_model, validated = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name,
+        qspec=qspec, custom_qdq_cases=custom_qdq_cases
+    )
+    assert validated, "ONNX QDQ validation for full network quantization failed!"
+    # necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()

tools/tensorflow-quantization/examples/resnet/utils.py

@@ -0,0 +1,45 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 tensorflow as tf
+from examples.data.data_loader import _DEFAULT_IMAGE_SIZE, _NUM_CHANNELS
+from examples.utils import get_tfkeras_model
+
+
+def get_resnet_model(resnet_depth: str = "50", resnet_version: str = "v1") -> tf.keras.Model:
+    """
+    Creates a native tf.keras ResNet model.
+
+    Args:
+        resnet_depth (str): ResNet depth. Options=[50 (default), 101, 152].
+        resnet_version (str): ResNet version. Options=[v1 (default), v2].
+
+    Returns:
+        model (tf.keras.Model): model corresponding to 'resnet_depth' and 'resnet_version'.
+    """
+
+    shape = (
+        _DEFAULT_IMAGE_SIZE["resnet_{}".format(resnet_version)],
+        _DEFAULT_IMAGE_SIZE["resnet_{}".format(resnet_version)],
+        _NUM_CHANNELS,
+    )
+
+    model_name = "resnet_" + resnet_depth + resnet_version
+    model = get_tfkeras_model(model_name=model_name, shape=shape)
+
+    return model

tools/tensorflow-quantization/examples/utils.py

@@ -0,0 +1,95 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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 os
+import tensorflow as tf
+from examples.data.data_loader import _NUM_CLASSES, _DEFAULT_IMAGE_SIZE, _NUM_CHANNELS
+from typing import Tuple
+
+MODELS_CLASSES_DICT = {
+    "resnet_50v1": tf.keras.applications.ResNet50,
+    "resnet_101v1": tf.keras.applications.ResNet101,
+    "resnet_152v1": tf.keras.applications.ResNet152,
+    "resnet_50v2": tf.keras.applications.ResNet50V2,
+    "resnet_101v2": tf.keras.applications.ResNet101V2,
+    "resnet_152v2": tf.keras.applications.ResNet152V2,
+    "mobilenet_v1": tf.keras.applications.MobileNet,
+    "mobilenet_v2": tf.keras.applications.MobileNetV2,
+    "inception_v3": tf.keras.applications.InceptionV3,
+}
+
+
+def get_tfkeras_model(model_name: str = "mobilenet_v1", shape: Tuple = None) -> tf.keras.Model:
+    """
+    Creates a native tf.keras.applications model.
+
+    Args:
+        model_name (str): Options={model_name_options}.
+
+    Returns:
+        model (tf.keras.Model): model corresponding to 'model_name'.
+
+    Raises:
+        ValueError: raised when 'model_name' is not supported.
+    """.format(
+        model_name_options=list(MODELS_CLASSES_DICT.keys())
+    )
+    try:
+        model_class = MODELS_CLASSES_DICT[model_name]
+    except ValueError:
+        raise ValueError("Model {} was not found!".format(model_name))
+    print("Loading model as {}".format(model_class))
+
+    if shape is None:
+        shape = (
+            _DEFAULT_IMAGE_SIZE[model_name],
+            _DEFAULT_IMAGE_SIZE[model_name],
+            _NUM_CHANNELS,
+        )
+
+    input_img = tf.keras.layers.Input(shape=shape, name="input_1")
+    model = model_class(
+        include_top=True,
+        weights="imagenet",
+        input_tensor=input_img,
+        input_shape=None,
+        pooling=None,
+        classes=_NUM_CLASSES,
+        classifier_activation="softmax",
+    )
+
+    return model
+
+
+def print_model_weights_shapes(model):
+    """
+    Print shape of each layer weight.
+    Args:
+        model: Keras model
+    """
+    print([model.get_weights()[i].shape for i in range(len(model.get_weights()))])
+
+
+def ensure_dir(dirname):
+    """
+    Create directory is doesn't exist already.
+    Args:
+        dirname: Name of the directory to create.
+    """
+    if not os.path.exists(dirname):
+        os.makedirs(dirname)

tools/tensorflow-quantization/examples/utils_finetuning.py

@@ -0,0 +1,279 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+#
+# Copyright 2022 The TensorFlow Authors. All Rights Reserved.
+#
+# 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.
+
+"""
+The only code snippet inherited from TensorFlow is the 'PiecewiseConstantDecayWithWarmup' class.
+    See that class description for the exact modifications.
+"""
+
+import os
+import tensorflow as tf
+import numpy as np
+from examples.data.data_loader import _NUM_IMAGES
+from datetime import datetime
+from examples.utils import ensure_dir
+from typing import Dict, List
+import logging
+
+
+def get_finetuned_weights_dirname(hyperparams: Dict) -> str:
+    """
+    Generates the directory name to save all files relevant to the model's quantization.
+
+    Args:
+        hyperparams (Dict): dictionary with necessary fine-tuning hyper-parameters.
+
+    Returns:
+        full_dirpath (str): path to directory where the fine-tuned model, log files, ... will be saved.
+    """
+    dirname = (
+        "qat_"
+        + "ep"
+        + str(hyperparams["epochs"])
+        + "_steps"
+        + str(hyperparams["steps_per_epoch"])
+        + "_baselr"
+        + str(hyperparams["base_lr"])
+        + "_"
+        + str(hyperparams["optimizer"])
+        + "_bs"
+        + str(hyperparams["batch_size"])
+    )
+    full_dirpath = os.path.join(hyperparams["save_root_dir"], dirname)
+    return full_dirpath
+
+
+def compile_model(model):
+    model.compile(
+        optimizer="sgd",
+        loss=tf.keras.losses.SparseCategoricalCrossentropy(),
+        metrics=["accuracy"],
+    )
+
+
+def fine_tune(
+    q_model: tf.keras.Model,
+    train_batches: tf.data.Dataset,
+    val_batches: tf.data.Dataset,
+    qat_save_finetuned_weights: str,
+    hyperparams: Dict,
+    logger: logging.RootLogger = None,
+    lr_schedule_array: List[tuple] = [(1.0, 1), (0.1, 2), (0.01, 7)],
+    enable_tensorboard_callback: bool = True
+) -> None:
+    """
+    Helper function to fine-tune QAT model.
+
+    Args:
+        q_model (tf.keras.Model): Keras model.
+        train_batches (tf.data.Dataset): train dataset split in batches.
+        val_batches (tf.data.Dataset): validation dataset split in batches.
+        qat_save_finetuned_weights (str): path to directory where the fine-tuned model, log files, ... will be saved.
+        hyperparams (Dict): dictionary with necessary fine-tuning hyper-parameters.
+        logger (logging.RootLogger): used to save logs.
+        lr_schedule_array (List[tuple]): list of tuples in the format '(multiplier, epoch to start)'.
+        enable_tensorboard_callback (bool): enables tensorboard callback if True.
+
+    Returns:
+        None
+
+    Raises:
+        ValueError: raised when the given optimizer is not supported.
+    """
+
+    if hyperparams["optimizer"] == "piecewise_sgd":
+        lr_schedule = PiecewiseConstantDecayWithWarmup(
+            batch_size=hyperparams["batch_size"],
+            epoch_size=_NUM_IMAGES["train"],  # for tfrecord
+            warmup_epochs=lr_schedule_array[0][1],
+            boundaries=list(p[1] for p in lr_schedule_array[1:]),
+            multipliers=list(p[0] for p in lr_schedule_array),
+            compute_lr_on_cpu=True,
+            base_lr=hyperparams["base_lr"],
+        )
+        optimizer = tf.keras.optimizers.SGD(learning_rate=lr_schedule, momentum=0.9)
+    elif hyperparams["optimizer"] == "sgd":
+        optimizer = tf.keras.optimizers.SGD(
+            learning_rate=hyperparams["base_lr"], momentum=0.0
+        )
+    elif hyperparams["optimizer"] == "adam":
+        optimizer = tf.keras.optimizers.Adam(hyperparams["base_lr"])
+    else:
+        raise ValueError("Optimizer `{}` is not supported. Please add support.".format(hyperparams["optimizer"]))
+
+    q_model.compile(
+        optimizer=optimizer,
+        loss=tf.keras.losses.SparseCategoricalCrossentropy(),
+        metrics=["accuracy"],
+    )
+
+    # Initialize TensorBoard visualization
+    callbacks = []
+    if enable_tensorboard_callback:
+        logdir_root = os.path.join(qat_save_finetuned_weights, "logs")
+        ensure_dir(logdir_root)
+        logdir = os.path.join(logdir_root, datetime.now().strftime("%Y%m%d-%H%M%S"))
+        tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=logdir)
+        callbacks.append(tensorboard_callback)
+    # Initialize ModelCheckpoint callback
+    ckpt_callback = tf.keras.callbacks.ModelCheckpoint(
+        filepath=os.path.join(qat_save_finetuned_weights, "checkpoints_best"),
+        save_weights_only=True,
+        monitor="val_accuracy",
+        mode="max",  # Save ckpt with max 'val_accuracy' (best)
+        save_best_only=True,
+    )
+    callbacks.append(ckpt_callback)
+
+    history = q_model.fit(
+        train_batches,
+        validation_data=val_batches,
+        batch_size=hyperparams["batch_size"],
+        epochs=hyperparams["epochs"],
+        steps_per_epoch=hyperparams["steps_per_epoch"],
+        callbacks=callbacks,
+        # verbose=2 if save_log is True else 1  # 0 = silent, 1 = progress bar, 2 = one line per epoch.
+    )
+
+    # Save fine-tuning history to logfile
+    if logger:
+        logger.info("------ Per epoch -------")
+        for ep in history.epoch:
+            log_str = "Epoch {ep}/{total_ep}".format(
+                ep=ep + 1, total_ep=history.params["epochs"]
+            )
+            for metric_name, metric_value in history.history.items():
+                log_str += " - " + metric_name + ": {}".format(metric_value[ep])
+            logger.info(log_str)
+        logger.info("------------------------")
+
+        # Save fine-tuned checkpoints
+        logger.info("Saving fine-tuned checkpoints")
+    q_model.save_weights(os.path.join(qat_save_finetuned_weights, "checkpoints_last"))
+
+
+class PiecewiseConstantDecayWithWarmup(
+    tf.keras.optimizers.schedules.LearningRateSchedule
+):
+    """
+    Piecewise constant decay with warmup schedule.
+
+    Original codebase: TensorFlow's "official.vision.image_classification.resnet.common"
+    Original URL: https://github.com/tensorflow/models/blob/master/official/legacy/image_classification/resnet/common.py
+
+    Modification: base learning rate `base_lr` given as parameter instead of a global constant.
+        PREVIOUS: self.rescaled_lr = BASE_LEARNING_RATE * batch_size / base_lr_batch_size
+        CURRENT:  self.rescaled_lr = base_lr
+    """
+
+    def __init__(
+        self,
+        batch_size,
+        epoch_size,
+        warmup_epochs,
+        boundaries,
+        multipliers,
+        compute_lr_on_cpu=True,
+        name=None,
+        base_lr=0.1,
+    ):
+        super(PiecewiseConstantDecayWithWarmup, self).__init__()
+        if len(boundaries) != len(multipliers) - 1:
+            raise ValueError(
+                "The length of boundaries must be 1 less than the "
+                "length of multipliers"
+            )
+
+        steps_per_epoch = epoch_size // batch_size
+
+        self.rescaled_lr = base_lr
+        self.step_boundaries = [np.int64(steps_per_epoch * x) for x in boundaries]
+        self.lr_values = [self.rescaled_lr * m for m in multipliers]
+        self.warmup_steps = warmup_epochs * steps_per_epoch
+        self.compute_lr_on_cpu = compute_lr_on_cpu
+        self.name = name
+
+        self.learning_rate_ops_cache = {}
+
+    def __call__(self, step):
+        if tf.executing_eagerly():
+            return self._get_learning_rate(step)
+
+        # In an eager function or graph, the current implementation of optimizer
+        # repeatedly call and thus create ops for the learning rate schedule. To
+        # avoid this, we cache the ops if not executing eagerly.
+        graph = tf.compat.v1.get_default_graph()
+        if graph not in self.learning_rate_ops_cache:
+            if self.compute_lr_on_cpu:
+                with tf.device("/device:CPU:0"):
+                    self.learning_rate_ops_cache[graph] = self._get_learning_rate(step)
+            else:
+                self.learning_rate_ops_cache[graph] = self._get_learning_rate(step)
+        return self.learning_rate_ops_cache[graph]
+
+    def _get_learning_rate(self, step):
+        """Compute learning rate at given step."""
+        with tf.name_scope("PiecewiseConstantDecayWithWarmup"):
+
+            def warmup_lr(step):
+                return self.rescaled_lr * (
+                    tf.cast(step, tf.float32) / tf.cast(self.warmup_steps, tf.float32)
+                )
+
+            def piecewise_lr(step):
+                if step.dtype == tf.float32:
+                    self.step_boundaries = [
+                        np.float32(bound) for bound in self.step_boundaries
+                    ]
+                elif step.dtype == tf.int64:
+                    self.step_boundaries = [
+                        np.int64(bound) for bound in self.step_boundaries
+                    ]
+
+                step_lr = tf.compat.v1.train.piecewise_constant(
+                    step, self.step_boundaries, self.lr_values
+                )
+                return step_lr
+
+            return tf.cond(
+                step < self.warmup_steps,
+                lambda: warmup_lr(step),
+                lambda: piecewise_lr(step),
+            )
+
+    def get_config(self):
+        return {
+            "rescaled_lr": self.rescaled_lr,
+            "step_boundaries": self.step_boundaries,
+            "lr_values": self.lr_values,
+            "warmup_steps": self.warmup_steps,
+            "compute_lr_on_cpu": self.compute_lr_on_cpu,
+            "name": self.name,
+        }