chore: import upstream snapshot with attribution

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