chore: import upstream snapshot with attribution

tools/tensorflow-quantization/.gitignore

@@ -0,0 +1,26 @@
+*.pyc
+.idea
+*.pytest_cache
+venv*
+
+tests/logs
+tests/saved_models
+tests/test_*
+temp
+*.json
+build
+tensorflow_quantization.egg-info
+tf-onnx-trt/tests/logs
+tf-onnx-trt/tests/saved_models
+
+qdq_insertion_experiments
+
+tf-trt/tests/logs
+tf-trt/tests/saved_models
+wrappers_test_saved_models
+utils_test_saved_models
+end_to_end_test
+qdq_test_saved_models
+custom_qdq_models
+
+assets/resnet*

tools/tensorflow-quantization/CHANGELOG.md

@@ -0,0 +1,16 @@
+# NVIDIA QAT toolkit for Tensorflow change log
+
+Dates are in YYYY-MM-DD format.
+
+## v0.2.0 (2022-09-09)
+- Fixed bug in `infer_engine.py` (low accuracy due to indentation error).
+- Added Inception-v3 support: code and results.
+- Fixed User Guide links (they were the internal links, not public ones).
+- Fixed bug in `examples/mobilenet/test_qdq_node_placement.py` (we implemented a more general `get_tfkeras_model` for Keras models QAT workflows, but missed that change in the MobileNet test script).
+- Added Pillow requirement in examples.
+- Added note for Conv2DTranspose support in the main README file.
+- Improved generalization of residual branch detection for QDQ node placement.
+
+## v0.1.0 (2022-06-17)
+
+- Initial release of the toolkit.

tools/tensorflow-quantization/LICENSE

@@ -0,0 +1,243 @@
+                                
+                                Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   Copyright 2022 NVIDIA Corporation
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+
+   PORTIONS LICENSED AS FOLLOWS
+
+   > tensorflow_quantization/quantize_wrapper_base.py
+   > tensorflow_quantization/quantizers.py
+   > tensorflow_quantization/quantize.py
+
+    Copyright 2019 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.
+
+   > examples/data_loader.py
+
+    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.
+
+   > examples/utils_finetuning.py
+
+    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.

tools/tensorflow-quantization/README.md

@@ -0,0 +1,83 @@
+**Note: Tensorflow Quantization development has transitioned to the [TensorRT Model Optimizer](https://github.com/NVIDIA/TensorRT-Model-Optimizer). All developers are encouraged to use the TensorRT Model Optimizer to benefit from the latest advancements on quantization and compression. While the Tensorflow Quantization code will remain available, it will no longer receive further development.**
+
+# <span style="color:green"> **NVIDIA TensorFlow 2.x Quantization** </span>
+
+This TensorFlow 2.x Quantization toolkit quantizes (inserts Q/DQ nodes) TensorFlow 2.x Keras models for Quantization-Aware Training (QAT).
+We follow NVIDIA's QAT recipe, which leads to optimal model acceleration with [TensorRT](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html) on NVIDIA GPUs and hardware accelerators.
+
+### Features
+
+- Implements [NVIDIA Quantization](https://arxiv.org/pdf/2004.09602.pdf) recipe.
+- Supports fully automated or manual insertion of Quantization and DeQuantization (QDQ) nodes in the TensorFlow 2.x model with minimal code.
+- Can easily to add support for new layers.
+- Quantization behavior can be set programmatically.
+- Implements automatic tests for popular architecture blocks such as residual and inception.
+- Offers utilities for TensorFlow 2.x to TensorRT conversion via ONNX.
+- Includes [example workflows](examples).
+
+## Dependencies
+
+**Python** >= 3.8  
+**TensorFlow** >= 2.8  
+**tf2onnx** >= 1.10.1  
+**onnx-graphsurgeon**  
+**pytest**  
+**pytest-html**  
+**TensorRT** (optional) >= 8.4 GA
+
+## Installation
+
+### Docker
+
+Latest TensorFlow 2.x [docker image](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tensorflow/tags) from NGC is recommended.
+
+```bash
+$ cd ~/
+$ git clone https://github.com/NVIDIA/TensorRT.git
+$ docker pull nvcr.io/nvidia/tensorflow:22.03-tf2-py3
+$ docker run -it --runtime=nvidia --gpus all --net host -v ~/TensorRT/tools/tensorflow-quantization:/home/tensorflow-quantization nvcr.io/nvidia/tensorflow:22.03-tf2-py3 /bin/bash
+```
+
+After last command, you will be placed in `/workspace` directory inside the running docker container whereas `tensorflow-quantization` repo is mounted in `/home` directory.
+
+```bash
+$ cd /home/tensorflow-quantization
+$ ./install.sh
+$ cd tests
+$ python3 -m pytest quantize_test.py -rP
+```
+
+If all tests pass, installation is successful.
+
+### Local
+
+```bash
+$ cd ~/
+$ git clone https://github.com/NVIDIA/TensorRT.git
+$ cd TensorRT/tools/tensorflow-quantization
+$ ./install.sh
+$ cd tests
+$ python3 -m pytest quantize_test.py -rP
+```
+
+If all tests pass, installation is successful.
+
+## Documentation
+
+TensorFlow 2.x Quantization toolkit [user guide](https://docs.nvidia.com/deeplearning/tensorrt/tensorflow-quantization-toolkit/docs/index.html).
+
+## Known limitations
+
+1. Only Quantization Aware Training (QAT) is supported as a quantization method.
+2. Only Functional and Sequential Keras models are supported. Original Keras layers are wrapped into quantized layers using TensorFlow's [clone_model](https://www.tensorflow.org/api_docs/python/tf/keras/models/clone_model) method, which doesn't support subclassed models.
+3. Saving the quantized version of a few layers may not be supported in `TensorFlow < 2.8`:
+   - `DepthwiseConv2D` support was added in TF 2.8.
+   - `Conv2DTranspose` is not yet supported by TF (see the open bug [here](https://github.com/tensorflow/model-optimization/issues/964)).
+     However, there's a workaround if you do not need the TF2 SavedModel file and just the ONNX file:
+     1. Implement `Conv2DTransposeQuantizeWrapper`. See our [user guide](https://docs.nvidia.com/deeplearning/tensorrt/tensorflow-quantization-toolkit/docs/docs/add_new_layer_support.html#example) for more information on how to do that.
+     2. Convert the quantized Keras model to ONNX using our provided utility function `convert_keras_model_to_onnx`.
+
+## Resources
+
+- [GTC 2022 talk](https://www.nvidia.com/en-us/on-demand/session/gtcspring22-s41440/)
+- Quantization Basics [whitepaper](https://arxiv.org/abs/2004.09602)

tools/tensorflow-quantization/VERSION

@@ -0,0 +1 @@
+0.2.0

tools/tensorflow-quantization/docs/Makefile

@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = source
+BUILDDIR      = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)

tools/tensorflow-quantization/docs/README.md

@@ -0,0 +1,15 @@
+# Sphinx Docuementation
+
+## Install Sphinx if necessary
+
+Bash script `setup_sphinx.sh` installs Sphinx and all necessary extensions to build this project's documentation.
+
+```bash
+$ ./setup_sphinx.sh
+```
+
+## Build docs
+
+```bash
+$ make clean html
+```

tools/tensorflow-quantization/docs/autobuild.sh

@@ -0,0 +1,5 @@
+#!/bin/sh
+dir1="./source"
+while inotifywait -qqre modify "$dir1"; do
+    make clean html
+done

tools/tensorflow-quantization/docs/make.bat

@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=source
+set BUILDDIR=build
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.https://www.sphinx-doc.org/
+	exit /b 1
+)
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd

tools/tensorflow-quantization/docs/setup_sphinx.sh

@@ -0,0 +1,14 @@
+#!/bin/bash
+
+# Install Sphinx
+python3 -m pip install Sphinx
+
+# Install sphinx-rtd-theme and sphinx-glpi-theme
+python3 -m pip install sphinx-rtd-theme
+python3 -m pip install sphinx-glpi-theme
+
+# Install myst-nb extension to process notebooks and myst markdown files
+python3 -m pip install myst-nb
+
+# Install mermaid for flow diagrams building
+python3 -m pip install sphinxcontrib-mermaid

tools/tensorflow-quantization/docs/source/bqw.rst

@@ -0,0 +1,116 @@
+.. _bqw_api:
+
+**tensorflow_quantization.BaseQuantizeWrapper**
+==================================================
+
+.. autoclass:: tensorflow_quantization.BaseQuantizeWrapper
+   :members:
+
+Example
+
+`Conv2DTranspose` layer is a weighted layer used to perform transformations going in the opposite direction of `Convolution`.
+
+.. note:: `Conv2DTranspose` is a Keras class, thus new wrapper class is `Conv2DTransposeQuantizeWrapper`. This follows toolkit naming conventions.
+
+.. code:: python
+
+    from tensorflow.python.util import tf_inspect
+    from tensorflow_quantization.quantize_wrapper_base import BaseQuantizeWrapper
+
+    class Conv2DTransposeQuantizeWrapper(BaseQuantizeWrapper):
+        def __init__(self, layer, kernel_type="kernel", **kwargs):
+            """
+            Create a quantize emulate wrapper for a keras layer.
+            This wrapper provides options to quantize inputs, outputs amd weights of a quantizable layer.
+            Args:
+            layer: The keras layer to be quantized.
+            kernel_type: Options=['kernel' for Conv2D/Dense, 'depthwise_kernel' for DepthwiseConv2D]
+            **kwargs: Additional keyword arguments to be passed to the keras layer.
+            """
+            self.kernel_type = kernel_type
+            self.channel_axis = kwargs.get("axis", -1)
+            super(Conv2DTransposeQuantizeWrapper, self).__init__(layer, **kwargs)
+
+        def build(self, input_shape):
+            super(Conv2DTransposeQuantizeWrapper, self).build(input_shape)
+
+            self._weight_vars = []
+            self.input_vars = {}
+            self.output_vars = {}
+            self.channel_axis = -1
+            if self.kernel_type == "depthwise_kernel":
+                self.channel_axis = 2
+            # quantize weights only applicable for weighted ops.
+            # By default weights is per channel quantization
+            if self.quantize_weights:
+                # get kernel weights dims.
+                kernel_weights = getattr(self.layer, self.kernel_type)
+                min_weight = self.layer.add_weight(
+                    kernel_weights.name.split(":")[0] + "_min",
+                    shape=(kernel_weights.shape[self.channel_axis]),
+                    initializer=tf.keras.initializers.Constant(-6.0),
+                    trainable=False,
+                )
+                max_weight = self.layer.add_weight(
+                    kernel_weights.name.split(":")[0] + "_max",
+                    shape=(kernel_weights.shape[self.channel_axis]),
+                    initializer=tf.keras.initializers.Constant(6.0),
+                    trainable=False,
+                )
+                quantizer_vars = {"min_var": min_weight, "max_var": max_weight}
+                self._weight_vars.append((kernel_weights, quantizer_vars))
+                # Needed to ensure unquantized weights get trained as part of the wrapper.
+                self._trainable_weights.append(kernel_weights)
+
+            # By default input is per tensor quantization
+            if self.quantize_inputs:
+                input_min_weight = self.layer.add_weight(
+                    self.layer.name + "_ip_min",
+                    shape=None,
+                    initializer=tf.keras.initializers.Constant(-6.0),
+                    trainable=False,
+                )
+                input_max_weight = self.layer.add_weight(
+                    self.layer.name + "_ip_max",
+                    shape=None,
+                    initializer=tf.keras.initializers.Constant(6.0),
+                    trainable=False,
+                )
+                self.input_vars["min_var"] = input_min_weight
+                self.input_vars["max_var"] = input_max_weight
+
+        def call(self, inputs, training=None):
+            if training is None:
+                training = tf.keras.backend.learning_phase()
+
+            # Quantize all weights, and replace them in the underlying layer.
+            if self.quantize_weights:
+                quantized_weights = []
+                quantized_weight = self._last_value_quantizer(
+                    self._weight_vars[0][0],
+                    training,
+                    self._weight_vars[0][1],
+                    per_channel=True,
+                    channel_axis=self.channel_axis
+                )
+                quantized_weights.append(quantized_weight)
+                # Replace the original weights with QDQ weights
+                setattr(self.layer, self.kernel_type, quantized_weights[0])
+
+            # Quantize inputs to the conv layer
+            if self.quantize_inputs:
+                quantized_inputs = self._last_value_quantizer(
+                    inputs, 
+                    training,
+                    self.input_vars,
+                    per_channel=False)
+            else:
+                quantized_inputs = inputs
+
+            args = tf_inspect.getfullargspec(self.layer.call).args
+            if "training" in args:
+                outputs = self.layer.call(quantized_inputs, training=training)
+            else:
+                outputs = self.layer.call(quantized_inputs)
+
+            return outputs

tools/tensorflow-quantization/docs/source/conf.py

@@ -0,0 +1,232 @@
+# -*- coding: utf-8 -*-
+#
+# Configuration file for the Sphinx documentation builder.
+#
+# This file does only contain a selection of the most common options. For a
+# full list see the documentation:
+# http://www.sphinx-doc.org/en/master/config
+
+# -- Path setup --------------------------------------------------------------
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# add these directories to sys.path here. If the directory is relative to the
+# documentation root, use os.path.abspath to make it absolute, like shown here.
+#
+# import os
+# import sys
+# sys.path.insert(0, os.path.abspath('.'))
+
+import sphinx_glpi_theme
+
+# -- Project information -----------------------------------------------------
+
+project = "TensorFlow 2.x Quantization Toolkit"
+copyright = "2022, NVIDIA"
+author = ""
+
+# The short X.Y version
+version = "1.0-beta"
+# The full version, including alpha/beta/rc tags
+release = "1.0.0"
+
+
+# -- General configuration ---------------------------------------------------
+
+# If your documentation needs a minimal Sphinx version, state it here.
+#
+# needs_sphinx = '1.0'
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+extensions = [
+    "sphinx.ext.doctest",
+    "sphinx.ext.intersphinx",
+    "sphinx.ext.todo",
+    "sphinx.ext.coverage",
+    "sphinx.ext.mathjax",
+    "sphinx.ext.ifconfig",
+    "sphinx.ext.viewcode",
+    "sphinx.ext.githubpages",
+    "sphinx.ext.coverage",
+    "sphinx.ext.napoleon",
+    "sphinx.ext.viewcode",
+    "sphinxcontrib.mermaid",
+    "myst_nb",
+]
+
+myst_enable_extensions = ["amsmath", "dollarmath", "smartquotes", "replacements"]
+
+
+napoleon_google_docstring = True
+napoleon_numpy_docstring = True
+napoleon_include_init_with_doc = True
+napoleon_include_private_with_doc = False
+napoleon_include_special_with_doc = False
+napoleon_use_admonition_for_examples = True
+napoleon_use_admonition_for_notes = True
+napoleon_use_admonition_for_references = False
+napoleon_use_ivar = True
+napoleon_use_param = False
+napoleon_use_rtype = False
+
+jupyter_execute_notebooks = "off"
+
+myst_all_links_external = True
+
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ["_templates"]
+
+# The suffix(es) of source filenames.
+# You can specify multiple suffix as a list of string:
+#
+source_suffix = [".rst", ".ipynb", ".md"]
+# source_suffix = '.rst'
+
+# The master toctree document.
+master_doc = "index"
+
+# The language for content autogenerated by Sphinx. Refer to documentation
+# for a list of supported languages.
+#
+# This is also used if you do content translation via gettext catalogs.
+# Usually you set "language" from the command line for these cases.
+language = "Python"
+
+# List of patterns, relative to source directory, that match files and
+# directories to ignore when looking for source files.
+# This pattern also affects html_static_path and html_extra_path.
+exclude_patterns = []
+
+# The name of the Pygments (syntax highlighting) style to use.
+pygments_style = None
+
+# -- Options for HTML output -------------------------------------------------
+
+# The theme to use for HTML and HTML Help pages.  See the documentation for
+# a list of builtin themes.
+#
+html_theme = "glpi"
+html_theme_path = sphinx_glpi_theme.get_html_themes_path()
+
+# Theme options are theme-specific and customize the look and feel of a theme
+# further.  For a list of options available for each theme, see the
+# documentation.
+#
+# html_theme_options = {}
+
+# Add any paths that contain custom static files (such as style sheets) here,
+# relative to this directory. They are copied after the builtin static files,
+# so a file named "default.css" will overwrite the builtin "default.css".
+html_static_path = ["_static"]
+
+# Custom sidebar templates, must be a dictionary that maps document names
+# to template names.
+#
+# The default sidebars (for documents that don't match any pattern) are
+# defined by theme itself.  Builtin themes are using these templates by
+# default: ``['localtoc.html', 'relations.html', 'sourcelink.html',
+# 'searchbox.html']``.
+#
+# html_sidebars = {}
+
+
+# -- Options for HTMLHelp output ---------------------------------------------
+
+# Output file base name for HTML help builder.
+htmlhelp_basename = "tensorflow-quantizationdoc"
+
+
+# -- Options for LaTeX output ------------------------------------------------
+
+latex_elements = {
+    # The paper size ('letterpaper' or 'a4paper').
+    #
+    # 'papersize': 'letterpaper',
+    # The font size ('10pt', '11pt' or '12pt').
+    #
+    # 'pointsize': '10pt',
+    # Additional stuff for the LaTeX preamble.
+    #
+    # 'preamble': '',
+    # Latex figure (float) alignment
+    #
+    # 'figure_align': 'htbp',
+}
+
+# Grouping the document tree into LaTeX files. List of tuples
+# (source start file, target name, title,
+#  author, documentclass [howto, manual, or own class]).
+latex_documents = [
+    (
+        master_doc,
+        "tensorflow-quantization.tex",
+        "tensorflow-quantization Documentation",
+        "NVIDIA",
+        "manual",
+    ),
+]
+
+
+# -- Options for manual page output ------------------------------------------
+
+# One entry per manual page. List of tuples
+# (source start file, name, description, authors, manual section).
+man_pages = [
+    (
+        master_doc,
+        "tensorflow-quantization",
+        "tensorflow-quantization Documentation",
+        [author],
+        1,
+    )
+]
+
+
+# -- Options for Texinfo output ----------------------------------------------
+
+# Grouping the document tree into Texinfo files. List of tuples
+# (source start file, target name, title, author,
+#  dir menu entry, description, category)
+texinfo_documents = [
+    (
+        master_doc,
+        "tensorflow-quantization",
+        "tensorflow-quantization Documentation",
+        author,
+        "tensorflow-quantization",
+        "Qauntization Toolkit for TensorFlow 2.x",
+        "Miscellaneous",
+    ),
+]
+
+
+# -- Options for Epub output -------------------------------------------------
+
+# Bibliographic Dublin Core info.
+epub_title = project
+
+# The unique identifier of the text. This can be a ISBN number
+# or the project homepage.
+#
+# epub_identifier = ''
+
+# A unique identification for the text.
+#
+# epub_uid = ''
+
+# A list of files that should not be packed into the epub file.
+epub_exclude_files = ["search.html"]
+
+
+# -- Extension configuration -------------------------------------------------
+
+# -- Options for intersphinx extension ---------------------------------------
+
+# Example configuration for intersphinx: refer to the Python standard library.
+intersphinx_mapping = {"https://docs.python.org/": None}
+
+# -- Options for todo extension ----------------------------------------------
+
+# If true, `todo` and `todoList` produce output, else they produce nothing.
+todo_include_todos = True

tools/tensorflow-quantization/docs/source/cqdq.rst

@@ -0,0 +1,26 @@
+.. _cqdq_api:
+
+**tensorflow_quantization.CustomQDQInsertionCase**
+==================================================
+
+.. autoclass:: tensorflow_quantization.CustomQDQInsertionCase
+   :members:
+
+Example
+
+.. code:: python
+
+   class EfficientNetQDQCase(CustomQDQInsertionCase):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def info(self):
+        return "In Multiply operation quantize inputs at index 0 and 1."
+
+    def case(self, keras_model: 'tf.keras.Model', qspec: 'QuantizationSpec') -> 'QuantizationSpec':
+        se_block_qspec_object = QuantizationSpec()
+        for layer in keras_model.layers:
+            if isinstance(layer, tf.keras.layers.Multiply):
+                se_block_qspec_object.add(layer.name, quantize_input=True, quantize_weight=False, quantization_index=[0, 1])
+        return se_block_qspec_object
+

tools/tensorflow-quantization/docs/source/docs/add_custom_qdq_cases.md

@@ -0,0 +1,183 @@
+(custom_qdq_case)=
+
+# **Add Custom QDQ Insertion Case**
+
+This toolkit's default quantization behavior for each supported layer is displayed in the [Add New Layer Support](new_layer_support) section.
+
+For the most part, it quantizes (adds Q/DQ nodes to) all inputs and weights (if the layer is weighted) of supported layers. However, the default behavior might not always lead to optimal INT8 fusions in TensorRT(TM). For example, Q/DQ nodes need to be added to residual connections in ResNet models. We provide a more in-depth explanation about this case in the "Custom Q/DQ Insertion Case Quantization" section later in this page.
+
+To tackle those scenarios, we added the `Custom Q/DQ Insertion Case` library feature, which allows users to programmatically decide how a specific layer should be quantized differently in specific situations. Note that providing an object of `QuantizationSpec` class is a hard coded way of achieving the same goal.
+
+Let's discuss the library-provided `ResNetV1QDQCase` to understand how passing custom Q/DQ insertion case objects affect Q/DQ insertion for the `Add` layer.
+
+## **Why is this needed?**
+
+The main goal of the `Custom Q/DQ Insertion` feature is to twick the framework's behavior to meet network-specific quantization requirements. Let's check this through an example.
+
+**Goal**: Perform custom quantization on a ResNet-like model. More specifically, we aim to quantize a model's residual connections.
+
+We show three quantization scenarios: 1) default, 2) custom with `QuantizationSpec` (suboptimal), and 3) custom with `Custom Q/DQ Insertion Case` (optimal).
+
+### **Default Quantization**
+
+```{note}
+Refer to **`Full Default Quantization`** [mode](basic).
+```
+
+The default quantization of the model is done with the following code snippet:
+
+```python
+# Quantize model
+q_nn_model = quantize_model(model=nn_model_original)
+```
+
+Figure 1, below, shows the baseline ResNet residual block and its corresponding quantized block with the default quantization scheme.
+
+<div align="center">
+
+![resnet_base](./assets/special_qdq_base.png)
+![resnet_default](./assets/special_qdq_default.png)
+
+Figure 1. ResNet residual block (left), and default quantized block (right).
+
+</div>
+
+Notice that the default quantization behavior is to not add Q/DQ nodes before `Add` layers. Since `AddQuantizeWrapper`
+is already implemented in the toolkit, and just disabled by default, the simplest way to quantize that layer would be
+to enable quantization of layers of class type `Add`.
+
+### **Custom Quantization with 'QuantizationSpec' (suboptimal)**
+
+```{note}
+Refer to **`Full Custom Quantization`** [mode](basic).
+```
+
+The following code snippet enables quantization of all layers of class type `Add`:
+
+```python
+# 1. Enable `Add` layer quantization
+qspec = QuantizationSpec()
+qspec.add(name='Add', is_keras_class=True)
+
+# 2. Quantize model
+q_nn_model = quantize_model(
+    model=nn_model_original, quantization_spec=qspec
+)
+```
+
+Figure 2, below, shows the standard ResNet residual block and its corresponding quantized block with the suggested custom quantization.
+
+<div align="center">
+
+![resnet_base](./assets/special_qdq_base.png)
+![resnet_default](./assets/special_qdq_qspec.png)
+
+Figure 2. ResNet residual block (left), and Q/DQ node insertion for `Add` layer passed via `QuantizationSpec` (right).
+
+</div>
+
+Notice that all inputs of the `Add` layer were quantized. However, that still does not enable optimal [layer fusions in TensorRT(TM)](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#enable-fusion), where a Convolution layer followed by an ElementWise layer (such as `Add`) can be fused into a single Convolution kernel.
+The [recommendation](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#qdq-placement-recs__xxx), in this case, is to add Q/DQ nodes in the residual connection only (not between `Add` and `Conv`).
+
+### **Custom Quantization with 'Custom Q/DQ Insertion Case' (optimal)**
+
+```{note}
+Refer to **`Full Custom Quantization`** [mode](basic).
+```
+
+The library-provided `ResNetV1QDQCase` class solves this issue by programming `Add` layer class to skip Q/DQ in one path if that path connects to `Conv`.
+This time, we pass an object of `ResNetV1QDQCase` class to the `quantize_model` function:
+
+```python
+# 1. Indicate one or more custom QDQ cases
+custom_qdq_case = ResNetV1QDQCase()
+
+# 3. Quantize model
+q_nn_model = quantize_model(
+    model=nn_model_original, custom_qdq_cases=[custom_qdq_case]
+)
+```
+
+Figure 3, below, shows the standard ResNet residual block and its corresponding quantized block with the suggested custom quantization.
+
+<div align="center">
+
+![resnet_base](./assets/special_qdq_base.png)
+![resnet_special](./assets/special_qdq_customqdqcase.png)
+
+Figure 3. ResNet residual block (left), and Q/DQ node insertion for `Add` layer passed via `ResNetV1QDQCase` (right).
+
+</div>
+
+Notice that Q/DQ nodes are not added to the path coming from `Conv` layer. Additionally, since both outputs of the first `Relu` layer were quantized, it was possible to perform a horizontal fusion with them, resulting in only one pair of Q/DQ nodes at that location.
+This quantization approach leads to an optimal graph for TensorRT INT8 fusions.
+
+## **Library provided custom Q/DQ insertion cases**
+
+We provide custom Q/DQ insertion cases for the models available in the model zoo. The library-provided custom Q/DQ insertion case classes can be imported from `tensorflow_quantization.custom_qdq_cases` module and passed to the `quantize_model` function.
+
+Refer to the [tensorflow_quantization.custom_qdq_cases](../../../tensorflow_quantization/custom_qdq_cases.py) module for more details.
+
+## **How to add a new custom Q/DQ insertion case?**
+
+```{eval-rst}
+
+#. Create a new class by inheriting ``tensorflow_quantization.CustomQDQInsertionCase`` class.
+#. Override two methods:
+
+    1. ``case`` (compulsory)
+
+    This method has fixed signature as shown below. Library automatically calls ``case`` method of all members of ``custom_qdq_cases`` parameter inside ``quantize_model`` function. Logic for changing the default layer behavior should be encoded in this function and an object of ``QuantizationSpec`` class must be returned.
+
+    .. code-block:: python
+
+        (function)CustomQDQInsertionCase.case(
+            self,
+            keras_model : 'tf.keras.Model',
+            qspec : 'QuantizationSpec'
+        ) -> 'QuantizationSpec'
+
+
+    2. ``info`` (optional)
+
+    This is just a helper method explaining the logic inside ``case`` method.
+
+#. Add object of this new class to a list and pass it to the ``custom_qdq_cases`` parameter of the ``quantize_model`` function.
+
+```
+
+```{eval-rst}
+.. ATTENTION::
+
+    If ``CustomQDQInsertionCase`` is written, ``QuantizationSpec`` object MUST be returned.
+
+```
+
+Example,
+
+```python
+class MaxPoolQDQCase(CustomQDQInsertionCase):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def info(self) -> str:
+        return "Enables quantization of MaxPool layers."
+
+    def case(
+        self, keras_model: tf.keras.Model, qspec: QuantizationSpec
+    ) -> QuantizationSpec:
+        mp_qspec = QuantizationSpec()
+        for layer in keras_model.layers:
+            if isinstance(layer, tf.keras.layers.MaxPooling2D):
+                if check_is_quantizable_by_layer_name(qspec, layer.name):
+                    mp_qspec.add(
+                        name=layer.name,
+                        quantize_input=True,
+                        quantize_weight=False
+                    )
+        return mp_qspec
+```
+
+As shown in the above MaxPool custom Q/DQ case class, the `case` method needs to be overridden. The optional `info` method returns a short description string.
+
+The logic written in the `case` method might or might not use the user-provided `QuantizationSpec` object, but it MUST return a new `QuantizationSpec` which holds information on the updated layer behavior. In the `MaxPoolQDQCase` case above, the custom Q/DQ insertion logic is dependent of the user-provided `QuantizationSpec` object (`check_is_quantizable_by_layer_name` checks if the layer name is in the user-provided object and gives priority to that specification).

tools/tensorflow-quantization/docs/source/docs/add_new_layer_support.md

@@ -0,0 +1,137 @@
+(new_layer_support)=
+
+# **Add New Layer Support**
+
+This toolkit uses a TensorFlow Keras [wrapper layer](https://www.tensorflow.org/api_docs/python/tf/keras/layers/Wrapper) to insert QDQ nodes before quantizable layers.
+
+## **Supported Layers**
+
+The following matrix shows the layers supported by this toolkit and their default behavior:
+
+| **Layer**                                 | **Quantize Input** | **Quantize Weight** | **Quantization Indices** |
+|-------------------------------------------|--------------------|---------------------|--------------------------|
+| tf.keras.layers.Conv2D                    | True               | True                | -                        |
+| tf.keras.layers.Dense                     | True               | True                | -                        |
+| tf.keras.layers.DepthwiseConv2D           | True               | True                | -                        |
+| tf.keras.layers.AveragePooling2D          | True               | -                   | -                        |
+| tf.keras.layers.GlobalAveragePooling2D    | True               | -                   | -                        |
+| tf.keras.layers.MaxPooling2D              | False*             | -                   | -                        |
+| tf.keras.layers.BatchNormalization        | False*             | -                   | -                        |
+| tf.keras.layers.Concatenate               | False*             | -                   | None*                    |
+| tf.keras.layers.Add                       | False*             | -                   | None*                    |
+| tf.keras.layers.Multiply                  | False*             | -                   | None*                    |
+
+```{note}
+*Inputs are not quantized by default. However, quantization is possible by passing those layers as `QuantizationSpec` to `quantize_model()`. Alternatively, fine-grained control over the layer's behavior can also be achieved by implementing a [Custom QDQ Insertion Case](custom_qdq_case).
+
+Note that the set of layers to be quantized can be network dependent. For example, `MaxPool` layers need not be quantized in ResNet-v1, but ResNet-v2 requires them to be quantized due to their location in residual connections. This toolkit, thus, offers flexibility to quantize any layer as needed.  
+
+``` 
+
+## **How are wrappers developed?**
+`BaseQuantizeWrapper` is a core quantization class which is inherited from `tf.keras.layers.Wrapper` keras wrapper class as shown in Figure 1 below.  
+
+<div align="center">
+
+![base_wrapper](./assets/inherit_from_keras_base.png)  
+
+Figure 1. BaseQuantizeWrapper inheritance.
+</div>
+
+
+All quantization wrappers are derived from `BaseQuantizeWrapper` class. Each wrapper takes *`layer(tf.keras.layers.Layer)`* as an argument which is handled by the toolkit internally. To simplify the development process, layers are classified as weighted, non-weighted, or other type.
+
+**Weighted Layers**
+
+Weighted layers are inherited from `WeightedBaseQuantizeWrapper` class which itself is inherited from `BaseQuantizeWrapper`, as shown in Figure 2 below. *`layer`* argument to `WeightedBaseQuantizeWrapper` class is handled by the library, however, *`kernel_type`* argument must be selected while developing wrapper. *`kernel_type`* for weighted layer gives access to layer weights.
+
+<div align="center">
+
+![weighted](./assets/inherit_from_weighted.png)  
+
+Figure 2. Inheritance flow for weighted layers.  
+</div>
+
+
+**Non-weighted Layers**
+
+Weighted layers are inherited from `WeightedBaseQuantizeWrapper` class which itself is inherited from `BaseQuantizeWrapper` as shown in Figure 3 below. *`layer`* argument to `WeightedBaseQuantizeWrapper` class is handled by the library.
+
+<div align="center">
+
+![non_weighted](./assets/inherit_from_non_weighted.png)  
+
+Figure 3. Inheritance flow for non-weighted layers. 
+</div>
+
+
+**Other Layers**
+Other layers are inherited from `BaseQuantizeWrapper` directly, as shown in Figure 4 below.
+
+<div align="center">
+
+![other](./assets/inherited_from_base.png)  
+
+Figure 4. Inheritance flow for other layers. 
+</div>
+
+## **How to add a new wrapper?**  
+
+```{eval-rst}
+
+#. Study current wrappers from ``tensorflow_quantization/quantize_wrappers.py`` script.
+#. Create a new class by inheriting one of ``BaseQuantizeWrapper``, ``WeightedBaseQuantizeWrapper`` or ``NonWeightedBaseQuantizeWrapper`` classes based on new layer type.
+#. Update ``build`` and ``call`` methods based upon layer behavior.
+
+.. ATTENTION::
+
+    New class will automatically get registered only if toolkit naming conventions are followed. For a keras layer ``l``, class name must be ``<l.__class__.__name__>QuantizeWrapper``.
+
+
+``` 
+
+## **Example**
+
+Let's see how support for a new Keras layer [GlobalMaxPool2D](https://www.tensorflow.org/api_docs/python/tf/keras/layers/GlobalMaxPool2D) can be added.
+
+This is a non-weighted layer thus we will inherit `NonWeightedBaseQuantizeWrapper`.
+
+Following toolkit naming conventions, this new wrapper should be named `GlobalMaxPool2DQuantizeWrapper`.
+
+```python
+
+from tensorflow_quantization import NonWeightedBaseQuantizeWrapper
+
+class GlobalMaxPool2DQuantizeWrapper(NonWeightedBaseQuantizeWrapper):
+    
+    def __init__(self, layer, **kwargs):
+        """
+        Creates a wrapper to emulate quantization for the GlobalMaxPool2D keras layer.
+        Args:
+        layer: The keras layer to be quantized.
+        **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        super().__init__(layer, **kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+    def call(self, inputs, training=None):
+        return super().call(inputs, training=training)
+
+```
+
+This new wrapper class is the same as the existing `GlobalAveragePooling2D`, `AveragePooling2D`, and `MaxPool2D` wrapper classes found in `tensorflow_quantization/quantize_wrappers.py`.
+
+```{eval-rst}
+
+.. ATTENTION::
+
+    New Class registration is based on tracking child classes of ``BaseQuantizeWrapper`` parent class. Thus, new class won't get registered unless explicitly called (this is current restriction).
+
+    To make sure new wrapper class is registered,
+    
+    #. If new wrapper class is defined in a separate module, import it in the module where ``quantize_model`` function is called.
+
+    #. If new wrapper class if defined in the same module as ``quantize_model`` function, create object of this new class. You don't have to pass that object anywhere.  
+```

tools/tensorflow-quantization/docs/source/docs/basics.md

@@ -0,0 +1,190 @@
+(basic)=
+
+```{eval-rst}
+
+.. admonition:: Attention
+    :class: attention
+
+    #. This toolkit supports only Quantization Aware Training (QAT) as a quantization method.
+    #. Subclassed models are not supported in the current version of this toolkit. Original Keras layers are wrapped into quantized layers using TensorFlow's `clone_model <https://www.tensorflow.org/api_docs/python/tf/keras/models/clone_model>`_ method, which doesn't support subclassed models.
+
+```
+
+# **Basics**
+
+## Quantization Function
+
+`quantize_model` is the only function the user needs to quantize any Keras model. It has the following signature:
+
+```python
+(function) quantize_model:
+(
+    model: tf.keras.Model,
+    quantization_mode: str = "full",
+    quantization_spec: QuantizationSpec = None,
+    custom_qdq_cases : List['CustomQDQInsertionCase'] = None
+) -> tf.keras.Model
+```
+
+````{note}
+Refer to the [Python API](qmodel_api) for more details.
+````
+
+Example
+
+```{eval-rst}
+.. code-block:: python
+    :emphasize-lines: 23
+
+    import tensorflow as tf
+    from tensorflow_quantization import quantize_model, utils
+
+    assets = utils.CreateAssetsFolders("toolkit_basics")
+    assets.add_folder("example")
+
+    # 1. Create a simple model (baseline)
+    input_img = tf.keras.layers.Input(shape=(28, 28, 1))
+    x = tf.keras.layers.Conv2D(filters=2, kernel_size=(3, 3))(input_img)
+    x = tf.keras.layers.ReLU()(x)
+    x = tf.keras.layers.Flatten()(x)
+    x = tf.keras.layers.Dense(10)(x)
+    model = tf.keras.Model(input_img, x)
+
+    # 2. Train model
+    model.fit(train_images, train_labels, batch_size=32, epochs=2, validation_split=0.1)
+
+    # 3. Save model and then convert it to ONNX
+    tf.keras.models.save_model(model, assets.example.fp32_saved_model)
+    utils.convert_saved_model_to_onnx(assets.example.fp32_saved_model, assets.example.fp32_onnx_model)
+
+    # 4. Quantize the model
+    q_model = quantize_model(model)
+
+    # 5. Train quantized model again for a few epochs to recover accuracy (fine-tuning).
+    q_model.fit(train_images, train_labels, batch_size=32, epochs=2, validation_split=0.1)
+    
+    # 6. Save the quantized model with QDQ nodes inserted and then convert it to ONNX
+    tf.keras.models.save_model(q_model, assets.example.int8_saved_model)
+    utils.convert_saved_model_to_onnx(assets.example.int8_saved_model, assets.example.int8_onnx_model)
+
+```
+
+````{note}
+The quantized model `q_model` behaves similar to the original [Keras model](https://www.tensorflow.org/api_docs/python/tf/keras/Model), meaning that the `compile()` and `fit()` functions can also be used to easily fine-tune the model.
+Refer to [Getting Started: End to End](GS_ETE) for more details.
+````
+
+Saved ONNX files can be visualized with [Netron](https://netron.app/). Figure 1, below, shows a snapshot of the original FP32 baseline model.
+
+<div align="center">
+
+![basic_fp32](./assets/basic_example_fp32_onnx_400px.png)
+
+Figure 1. Original FP32 model.
+</div>
+
+The quantization process inserts [Q/DQ](https://www.tensorflow.org/api_docs/python/tf/quantization/quantize_and_dequantize_v2)
+ nodes at the inputs and weights (if layer is weighted) of all supported layers, according to the TensorRT(TM) quantization 
+ [policy](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#work-with-qat-networks). 
+ The presence of a Quantize node (`QuantizeLinear` ONNX op), followed by a Dequantize node (`DequantizeLinear` ONNX op), 
+ for each supported layer, can be verified in the Netron visualization in Figure 2 below.
+
+<div align="center">
+
+![basic_int8](./assets/basic_example_int8_onnx_400px.png)
+
+Figure 2. Quantized INT8 model.
+</div>
+
+TensorRT(TM) converts ONNX models with Q/DQ nodes into an INT8 [engine](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#prog-model), which can take advantage of [Tensor Cores](https://www.nvidia.com/en-us/data-center/tensor-cores/) and other hardware accelerations in the latest NVIDIA(R) GPUs.
+
+## Quantization Modes
+
+There are a few scenarios where one might need to customize the default quantization scheme.
+ We broadly categorize quantization (i.e. the process of adding Q/DQ nodes) into `Full` and `Partial` modes, depending on the set of layers that are quantized.
+ Additionally, Full quantization can be `Default` or `Custom`, while Partial quantization is always `Custom`.
+
+```{eval-rst}
+
+#. Full Default Quantization
+    All supported layers of a given model are quantized as per default toolkit behavior.
+#. Full Custom Quantization
+    Toolkit behavior can be programmed to quantize specific layers differentely by passing an object of ``QuantizationSpec`` class and/or ``CustomQDQInsertionCase`` class. The remaining supported layers are quantized as per `default behavior <new_layer_support>`_. 
+#. Partial Quantization
+    Only layers passed using ``QuantizationSpec`` and/or ``CustomQDQInsertionCase`` class object are quantized.
+```
+
+```{note}
+Refer to the [Tutorials](tut_one) for examples on each mode.
+```
+
+## Terminologies
+
+### Layer Name 
+
+Name of the Keras layer either assigned by the user or Keras. These are unique by default.
+
+```python
+import tensorflow as tf
+l = tf.keras.layers.Dense(units=100, name='my_dense')
+```
+
+Here ‘my_dense’ is a layer name assigned by the user.
+
+```{tip}
+For a given layer `l`, the layer name can be found using `l.name`.
+```
+
+### Layer Class
+
+Name of the Keras layer class.
+
+```python
+import tensorflow as tf
+l = tf.keras.layers.Dense(units=100, name='my_dense')
+```
+
+Here ‘Dense’ is the layer class.
+
+```{tip}
+For a given layer `l`, the layer class can be found using `l.__class__.__name__` or `l.__class__.__module__`.
+```
+
+## NVIDIA(R) vs TensorFlow Toolkit
+
+[TFMOT](https://www.tensorflow.org/model_optimization/api_docs/python/tfmot) is TensorFlow's official quantization toolkit. The quantization recipe used by TFMOT is different to NVIDIA(R)'s in terms of Q/DQ nodes placement, and it is optimized for TFLite inference. The NVIDIA(R) quantization recipe, on the other hand, is optimized for TensorRT(TM), which leads to optimal model acceleration on NVIDIA(R) GPUs and hardware accelerators.
+
+Other differences:
+
+| Feature                     | TensorFlow Model Optimization Toolkit (TFMOT)                  | NVIDIA(R) Toolkit                                                                                                                         |
+|-----------------------------|----------------------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------|
+| QDQ node placements         | Outputs and Weights                                            | Inputs and Weights                                                                                                                        |
+| Quantization support        | Whole model (full) and of some layers (partial by layer class) | Extends TF quantization support: partial quantization by layer name and pattern-base quantization by extending `CustomQDQInsertionCase`   |
+| Quantization scheme         | `tf.quantization.fake_quant_with_min_max_vars`                 | `tf.quantization.quantize_and_dequantize_v2`                                                                                              |
+
+
+## Additional Resources
+**About this toolkit**
+- [GTC 2022 Tech Talk](https://www.nvidia.com/en-us/on-demand/session/gtcspring22-s41440/) (NVIDIA)
+
+**Blogs**
+- [Achieving FP32 Accuracy for INT8 Inference Using Quantization Aware Training with NVIDIA TensorRT](https://developer.nvidia.com/blog/achieving-fp32-accuracy-for-int8-inference-using-quantization-aware-training-with-tensorrt/) (NVIDIA)
+- [Speeding Up Deep Learning Inference Using TensorFlow, ONNX, and NVIDIA TensorRT](https://developer.nvidia.com/blog/speeding-up-deep-learning-inference-using-tensorflow-onnx-and-tensorrt/) (NVIDIA)
+- [Why are Eight Bits Enough for Deep Neural Networks?](https://petewarden.com/2015/05/23/why-are-eight-bits-enough-for-deep-neural-networks/)
+- [What I’ve learned about Neural Network Quantization](https://petewarden.com/2017/06/22/what-ive-learned-about-neural-network-quantization/)
+
+**Videos**
+- [Inference and Quantization](https://www.youtube.com/watch?v=VsGX9kFXjbs)
+- [8-bit Inference with TensorRT Webinar](http://on-demand.gputechconf.com/gtcdc/2017/video/DC7172/)
+
+**Generate per-tensor dynamic range**
+- [Setting Per-Tensor Dynamic Range Using C++](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#set_tensor_mp_c) (NVIDIA)
+- [Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference](https://arxiv.org/pdf/1712.05877.pdf)
+- [Quantizing Deep Convolutional Networks for Efficient Inference: A Whitepaper](https://arxiv.org/pdf/1806.08342.pdf)
+- [8-bit Inference with TensorRT](http://on-demand.gputechconf.com/gtc/2017/presentation/s7310-8-bit-inference-with-tensorrt.pdf) (NVIDIA)
+
+**Documentation** (NVIDIA)
+- [Introduction to NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
+- [Working with TensorRT Using the C++ API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#c_topics)
+- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
+- [Working with INT8 in TensorRT](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#working-with-int8)

tools/tensorflow-quantization/docs/source/docs/example_resnet50v1.ipynb

@@ -0,0 +1,482 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "# **ResNet50 V1**\n",
+    "\n",
+    "This assumes that our toolkits and its base requirements have been met, including access to the ImageNet dataset. Please refer to [\"Requirements\"](https://gitlab-master.nvidia.com/sagshelke/tensorrt_qat/-/tree/main/examples#requirements) in the `examples` folder."
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "## 1. Initial settings"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "outputs": [],
+   "source": [
+    "import os\n",
+    "import tensorflow as tf\n",
+    "from tensorflow_quantization.quantize import quantize_model\n",
+    "from tensorflow_quantization.custom_qdq_cases import ResNetV1QDQCase\n",
+    "from tensorflow_quantization.utils import convert_saved_model_to_onnx"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "HYPERPARAMS = {\n",
+    "    \"tfrecord_data_dir\": \"/media/Data/ImageNet/train-val-tfrecord\",\n",
+    "    \"batch_size\": 64,\n",
+    "    \"epochs\": 2,\n",
+    "    \"steps_per_epoch\": 500,\n",
+    "    \"train_data_size\": None,\n",
+    "    \"val_data_size\": None,\n",
+    "    \"save_root_dir\": \"./weights/resnet_50v1_jupyter\"\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "### Load data"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 12,
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "from examples.data.data_loader import load_data\n",
+    "train_batches, val_batches = load_data(HYPERPARAMS, model_name=\"resnet_v1\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "## 2. Baseline model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "### Instantiate"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 13,
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "model = tf.keras.applications.ResNet50(\n",
+    "    include_top=True,\n",
+    "    weights=\"imagenet\",\n",
+    "    classes=1000,\n",
+    "    classifier_activation=\"softmax\",\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "### Evaluate"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 14,
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "781/781 [==============================] - 41s 51ms/step - loss: 1.0481 - accuracy: 0.7504\n",
+      "Baseline val accuracy: 75.044%\n"
+     ]
+    }
+   ],
+   "source": [
+    "def compile_model(model, lr=0.001):\n",
+    "    model.compile(\n",
+    "        optimizer=tf.keras.optimizers.SGD(learning_rate=lr),\n",
+    "        loss=tf.keras.losses.SparseCategoricalCrossentropy(),\n",
+    "        metrics=[\"accuracy\"],\n",
+    "    )\n",
+    "\n",
+    "compile_model(model)\n",
+    "_, baseline_model_accuracy = model.evaluate(val_batches)\n",
+    "print(\"Baseline val accuracy: {:.3f}%\".format(baseline_model_accuracy*100))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "### Save and convert to ONNX"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 15,
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:tensorflow:Assets written to: ./weights/resnet_50v1_jupyter/saved_model_baseline/assets\n",
+      "ONNX conversion Done!\n"
+     ]
+    }
+   ],
+   "source": [
+    "model_save_path = os.path.join(HYPERPARAMS[\"save_root_dir\"], \"saved_model_baseline\")\n",
+    "model.save(model_save_path)\n",
+    "convert_saved_model_to_onnx(saved_model_dir=model_save_path,\n",
+    "                            onnx_model_path=model_save_path + \".onnx\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "## 3. Quantization-Aware Training model"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "### Quantize"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 16,
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "q_model = quantize_model(model, custom_qdq_cases=[ResNetV1QDQCase()])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "### Fine-tune"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 17,
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/2\n",
+      "500/500 [==============================] - 425s 838ms/step - loss: 0.4075 - accuracy: 0.8898 - val_loss: 1.0451 - val_accuracy: 0.7497\n",
+      "Epoch 2/2\n",
+      "500/500 [==============================] - 420s 840ms/step - loss: 0.3960 - accuracy: 0.8918 - val_loss: 1.0392 - val_accuracy: 0.7511\n"
+     ]
+    },
+    {
+     "data": {
+      "text/plain": "<keras.callbacks.History at 0x7f9cec1e60d0>"
+     },
+     "execution_count": 17,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "compile_model(q_model)\n",
+    "q_model.fit(\n",
+    "    train_batches,\n",
+    "    validation_data=val_batches,\n",
+    "    batch_size=HYPERPARAMS[\"batch_size\"],\n",
+    "    steps_per_epoch=HYPERPARAMS[\"steps_per_epoch\"],\n",
+    "    epochs=HYPERPARAMS[\"epochs\"]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "### Evaluate"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 18,
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "781/781 [==============================] - 179s 229ms/step - loss: 1.0392 - accuracy: 0.7511\n",
+      "QAT val accuracy: 75.114%\n"
+     ]
+    }
+   ],
+   "source": [
+    "_, qat_model_accuracy = q_model.evaluate(val_batches)\n",
+    "print(\"QAT val accuracy: {:.3f}%\".format(qat_model_accuracy*100))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "### Save and convert to ONNX"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "WARNING:absl:Found untraced functions such as conv1_conv_layer_call_fn, conv1_conv_layer_call_and_return_conditional_losses, conv2_block1_1_conv_layer_call_fn, conv2_block1_1_conv_layer_call_and_return_conditional_losses, conv2_block1_2_conv_layer_call_fn while saving (showing 5 of 140). These functions will not be directly callable after loading.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:tensorflow:Assets written to: ./weights/resnet_50v1_jupyter/saved_model_qat/assets\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:tensorflow:Assets written to: ./weights/resnet_50v1_jupyter/saved_model_qat/assets\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "ONNX conversion Done!\n"
+     ]
+    }
+   ],
+   "source": [
+    "q_model_save_path = os.path.join(HYPERPARAMS[\"save_root_dir\"], \"saved_model_qat\")\n",
+    "q_model.save(q_model_save_path)\n",
+    "convert_saved_model_to_onnx(saved_model_dir=q_model_save_path,\n",
+    "                            onnx_model_path=q_model_save_path + \".onnx\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "## 4. QAT vs Baseline comparison"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Baseline vs QAT: 75.044% vs 75.114%\n",
+      "Accuracy difference of +0.070%\n"
+     ]
+    }
+   ],
+   "source": [
+    "print(\"Baseline vs QAT: {:.3f}% vs {:.3f}%\".format(baseline_model_accuracy*100, qat_model_accuracy*100))\n",
+    "\n",
+    "acc_diff = (qat_model_accuracy - baseline_model_accuracy)*100\n",
+    "acc_diff_sign = \"\" if acc_diff == 0 else (\"-\" if acc_diff < 0 else \"+\")\n",
+    "print(\"Accuracy difference of {}{:.3f}%\".format(acc_diff_sign, abs(acc_diff)))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "```{note}\n",
+    "\n",
+    "For full workflow, including TensorRT(TM) deployment, please refer to [examples/resnet](https://github.com/NVIDIA/TensorRT/tree/main/tools/tensorflow-quantization/examples/resnet).\n",
+    "\n",
+    "```"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.7.3"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 1
+}

tools/tensorflow-quantization/docs/source/docs/getting_started.ipynb

@@ -0,0 +1,569 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "(GS_ETE)=\n",
+    "\n",
+    "# **Getting Started: End to End**\n",
+    "\n",
+    "NVIDIA(R) TensorFlow 2.x Quantization toolkit provides a simple API to quantize a given Keras model. At a higher level, Quantization Aware Training (QAT) is a three-step workflow as shown below:\n",
+    "\n",
+    "```{eval-rst}\n",
+    ".. mermaid::\n",
+    "\n",
+    "    flowchart LR\n",
+    "        id1(Pre-trained model) --> id2(Quantize) --> id3(Fine-tune)\n",
+    "\n",
+    "```\n",
+    "Initially, the network is trained on the target dataset until fully converged. The Quantization step consists of inserting Q/DQ nodes in the pre-trained network to simulate quantization during training. Note that simply adding Q/DQ nodes will result in reduced accuracy since the quantization parameters are not yet updated for the given model. The network is then re-trained for a few epochs to recover accuracy in a step called \"fine-tuning\".\n",
+    "\n",
+    "```{eval-rst}\n",
+    "\n",
+    ".. admonition:: Goal\n",
+    "    :class: note\n",
+    "\n",
+    "    #. Train a simple network on the Fashion MNIST dataset and save it as the baseline model.\n",
+    "    #. Quantize the pre-trained baseline network.\n",
+    "    #. Fine-tune the quantized network to recover accuracy and save it as the QAT model.\n",
+    "\n",
+    "```\n",
+    "---\n",
+    "\n",
+    "## 1. Train\n",
+    "Import required libraries and create a simple network with convolution and dense layers."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Model: \"original\"\n",
+      "_________________________________________________________________\n",
+      " Layer (type)                Output Shape              Param #   \n",
+      "=================================================================\n",
+      " nn_input (InputLayer)       [(None, 28, 28)]          0         \n",
+      "                                                                 \n",
+      " reshape_0 (Reshape)         (None, 28, 28, 1)         0         \n",
+      "                                                                 \n",
+      " conv_0 (Conv2D)             (None, 26, 26, 126)       1260      \n",
+      "                                                                 \n",
+      " relu_0 (ReLU)               (None, 26, 26, 126)       0         \n",
+      "                                                                 \n",
+      " conv_1 (Conv2D)             (None, 24, 24, 64)        72640     \n",
+      "                                                                 \n",
+      " relu_1 (ReLU)               (None, 24, 24, 64)        0         \n",
+      "                                                                 \n",
+      " conv_2 (Conv2D)             (None, 22, 22, 32)        18464     \n",
+      "                                                                 \n",
+      " relu_2 (ReLU)               (None, 22, 22, 32)        0         \n",
+      "                                                                 \n",
+      " conv_3 (Conv2D)             (None, 20, 20, 16)        4624      \n",
+      "                                                                 \n",
+      " relu_3 (ReLU)               (None, 20, 20, 16)        0         \n",
+      "                                                                 \n",
+      " conv_4 (Conv2D)             (None, 18, 18, 8)         1160      \n",
+      "                                                                 \n",
+      " relu_4 (ReLU)               (None, 18, 18, 8)         0         \n",
+      "                                                                 \n",
+      " max_pool_0 (MaxPooling2D)   (None, 9, 9, 8)           0         \n",
+      "                                                                 \n",
+      " flatten_0 (Flatten)         (None, 648)               0         \n",
+      "                                                                 \n",
+      " dense_0 (Dense)             (None, 100)               64900     \n",
+      "                                                                 \n",
+      " relu_5 (ReLU)               (None, 100)               0         \n",
+      "                                                                 \n",
+      " dense_1 (Dense)             (None, 10)                1010      \n",
+      "                                                                 \n",
+      "=================================================================\n",
+      "Total params: 164,058\n",
+      "Trainable params: 164,058\n",
+      "Non-trainable params: 0\n",
+      "_________________________________________________________________\n"
+     ]
+    }
+   ],
+   "source": [
+    "import tensorflow as tf\n",
+    "from tensorflow_quantization import quantize_model\n",
+    "from tensorflow_quantization import utils\n",
+    "\n",
+    "assets = utils.CreateAssetsFolders(\"GettingStarted\")\n",
+    "assets.add_folder(\"example\")\n",
+    "\n",
+    "def simple_net():\n",
+    "    \"\"\"\n",
+    "    Return a simple neural network.\n",
+    "    \"\"\"\n",
+    "    input_img = tf.keras.layers.Input(shape=(28, 28), name=\"nn_input\")\n",
+    "    x = tf.keras.layers.Reshape(target_shape=(28, 28, 1), name=\"reshape_0\")(input_img)\n",
+    "    x = tf.keras.layers.Conv2D(filters=126, kernel_size=(3, 3), name=\"conv_0\")(x)\n",
+    "    x = tf.keras.layers.ReLU(name=\"relu_0\")(x)\n",
+    "    x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3), name=\"conv_1\")(x)\n",
+    "    x = tf.keras.layers.ReLU(name=\"relu_1\")(x)\n",
+    "    x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), name=\"conv_2\")(x)\n",
+    "    x = tf.keras.layers.ReLU(name=\"relu_2\")(x)\n",
+    "    x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3, 3), name=\"conv_3\")(x)\n",
+    "    x = tf.keras.layers.ReLU(name=\"relu_3\")(x)\n",
+    "    x = tf.keras.layers.Conv2D(filters=8, kernel_size=(3, 3), name=\"conv_4\")(x)\n",
+    "    x = tf.keras.layers.ReLU(name=\"relu_4\")(x)\n",
+    "    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2), name=\"max_pool_0\")(x)\n",
+    "    x = tf.keras.layers.Flatten(name=\"flatten_0\")(x)\n",
+    "    x = tf.keras.layers.Dense(100, name=\"dense_0\")(x)\n",
+    "    x = tf.keras.layers.ReLU(name=\"relu_5\")(x)\n",
+    "    x = tf.keras.layers.Dense(10, name=\"dense_1\")(x)\n",
+    "    return tf.keras.Model(input_img, x, name=\"original\")\n",
+    "\n",
+    "# create model\n",
+    "model = simple_net()\n",
+    "model.summary()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "Load Fashion MNIST data and split train and test sets."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 38,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# Load Fashion MNIST dataset\n",
+    "mnist = tf.keras.datasets.fashion_mnist\n",
+    "(train_images, train_labels), (test_images, test_labels) = mnist.load_data()\n",
+    "\n",
+    "# Normalize the input image so that each pixel value is between 0 to 1.\n",
+    "train_images = train_images / 255.0\n",
+    "test_images = test_images / 255.0   "
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "Compile the model and train for five epochs."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 39,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/5\n",
+      "422/422 [==============================] - 4s 8ms/step - loss: 0.5639 - accuracy: 0.7920 - val_loss: 0.4174 - val_accuracy: 0.8437\n",
+      "Epoch 2/5\n",
+      "422/422 [==============================] - 3s 8ms/step - loss: 0.3619 - accuracy: 0.8696 - val_loss: 0.4134 - val_accuracy: 0.8433\n",
+      "Epoch 3/5\n",
+      "422/422 [==============================] - 3s 8ms/step - loss: 0.3165 - accuracy: 0.8855 - val_loss: 0.3137 - val_accuracy: 0.8812\n",
+      "Epoch 4/5\n",
+      "422/422 [==============================] - 3s 8ms/step - loss: 0.2787 - accuracy: 0.8964 - val_loss: 0.2943 - val_accuracy: 0.8890\n",
+      "Epoch 5/5\n",
+      "422/422 [==============================] - 3s 8ms/step - loss: 0.2552 - accuracy: 0.9067 - val_loss: 0.2857 - val_accuracy: 0.8952\n",
+      "Baseline test accuracy: 0.888700008392334\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Train original classification model\n",
+    "model.compile(\n",
+    "    optimizer=\"adam\",\n",
+    "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+    "    metrics=[\"accuracy\"],\n",
+    ")\n",
+    "\n",
+    "model.fit(\n",
+    "    train_images, train_labels, batch_size=128, epochs=5, validation_split=0.1\n",
+    ")\n",
+    "\n",
+    "# get baseline model accuracy\n",
+    "_, baseline_model_accuracy = model.evaluate(\n",
+    "    test_images, test_labels, verbose=0\n",
+    ")\n",
+    "print(\"Baseline test accuracy:\", baseline_model_accuracy)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 40,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:tensorflow:Assets written to: GettingStarted/example/fp32/saved_model/assets\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:tensorflow:Assets written to: GettingStarted/example/fp32/saved_model/assets\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "ONNX conversion Done!\n"
+     ]
+    }
+   ],
+   "source": [
+    "# save TF FP32 original model\n",
+    "tf.keras.models.save_model(model, assets.example.fp32_saved_model)\n",
+    "\n",
+    "# Convert FP32 model to ONNX\n",
+    "utils.convert_saved_model_to_onnx(saved_model_dir = assets.example.fp32_saved_model, onnx_model_path = assets.example.fp32_onnx_model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "## 2. Quantize\n",
+    "\n",
+    "Full model quantization is the most basic quantization mode someone can follow. In this mode, Q/DQ nodes are inserted in all supported keras layers, with a single function call:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 41,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [],
+   "source": [
+    "# Quantize model\n",
+    "quantized_model = quantize_model(model)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "Keras model summary shows all supported layers wrapped into QDQ wrapper class."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 42,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Model: \"original\"\n",
+      "_________________________________________________________________\n",
+      " Layer (type)                Output Shape              Param #   \n",
+      "=================================================================\n",
+      " nn_input (InputLayer)       [(None, 28, 28)]          0         \n",
+      "                                                                 \n",
+      " reshape_0 (Reshape)         (None, 28, 28, 1)         0         \n",
+      "                                                                 \n",
+      " quant_conv_0 (Conv2DQuantiz  (None, 26, 26, 126)      1515      \n",
+      " eWrapper)                                                       \n",
+      "                                                                 \n",
+      " relu_0 (ReLU)               (None, 26, 26, 126)       0         \n",
+      "                                                                 \n",
+      " quant_conv_1 (Conv2DQuantiz  (None, 24, 24, 64)       72771     \n",
+      " eWrapper)                                                       \n",
+      "                                                                 \n",
+      " relu_1 (ReLU)               (None, 24, 24, 64)        0         \n",
+      "                                                                 \n",
+      " quant_conv_2 (Conv2DQuantiz  (None, 22, 22, 32)       18531     \n",
+      " eWrapper)                                                       \n",
+      "                                                                 \n",
+      " relu_2 (ReLU)               (None, 22, 22, 32)        0         \n",
+      "                                                                 \n",
+      " quant_conv_3 (Conv2DQuantiz  (None, 20, 20, 16)       4659      \n",
+      " eWrapper)                                                       \n",
+      "                                                                 \n",
+      " relu_3 (ReLU)               (None, 20, 20, 16)        0         \n",
+      "                                                                 \n",
+      " quant_conv_4 (Conv2DQuantiz  (None, 18, 18, 8)        1179      \n",
+      " eWrapper)                                                       \n",
+      "                                                                 \n",
+      " relu_4 (ReLU)               (None, 18, 18, 8)         0         \n",
+      "                                                                 \n",
+      " max_pool_0 (MaxPooling2D)   (None, 9, 9, 8)           0         \n",
+      "                                                                 \n",
+      " flatten_0 (Flatten)         (None, 648)               0         \n",
+      "                                                                 \n",
+      " quant_dense_0 (DenseQuantiz  (None, 100)              65103     \n",
+      " eWrapper)                                                       \n",
+      "                                                                 \n",
+      " relu_5 (ReLU)               (None, 100)               0         \n",
+      "                                                                 \n",
+      " quant_dense_1 (DenseQuantiz  (None, 10)               1033      \n",
+      " eWrapper)                                                       \n",
+      "                                                                 \n",
+      "=================================================================\n",
+      "Total params: 164,791\n",
+      "Trainable params: 164,058\n",
+      "Non-trainable params: 733\n",
+      "_________________________________________________________________\n"
+     ]
+    }
+   ],
+   "source": [
+    "quantized_model.summary()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "Let's check the quantized model's accuracy immediately after Q/DQ nodes are inserted."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 43,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Quantization test accuracy immediately after QDQ insertion: 0.883899986743927\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Compile quantized model\n",
+    "quantized_model.compile(\n",
+    "    optimizer=tf.keras.optimizers.Adam(0.0001),\n",
+    "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+    "    metrics=[\"accuracy\"],\n",
+    ")\n",
+    "# Get accuracy immediately after QDQ nodes are inserted.\n",
+    "_, q_aware_model_accuracy = quantized_model.evaluate(test_images, test_labels, verbose=0)\n",
+    "print(\"Quantization test accuracy immediately after QDQ insertion:\", q_aware_model_accuracy)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "The model's accuracy decreases a bit as soon as Q/DQ nodes are inserted, requiring fine-tuning to recover it.\n",
+    "\n",
+    "```{note}\n",
+    "\n",
+    "Since this is a very small model, accuracy drop is small. For standard models like ResNets, accuracy drop immediately after QDQ insertion can be significant.\n",
+    "\n",
+    "```\n",
+    "\n",
+    "## 3. Fine-tune\n",
+    "Since the quantized model behaves similar to the original keras model, the same training recipe can be used for fine-tuning as well.\n",
+    "\n",
+    "We fine-tune the model for two epochs and evaluate the model on the test dataset."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 46,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/2\n",
+      "1688/1688 [==============================] - 26s 15ms/step - loss: 0.1793 - accuracy: 0.9340 - val_loss: 0.2468 - val_accuracy: 0.9112\n",
+      "Epoch 2/2\n",
+      "1688/1688 [==============================] - 25s 15ms/step - loss: 0.1725 - accuracy: 0.9373 - val_loss: 0.2484 - val_accuracy: 0.9070\n",
+      "Quantization test accuracy after fine-tuning: 0.9075999855995178\n",
+      "Baseline test accuracy (for reference): 0.888700008392334\n"
+     ]
+    }
+   ],
+   "source": [
+    "# fine tune quantized model for 2 epochs.\n",
+    "quantized_model.fit(\n",
+    "    train_images, train_labels, batch_size=32, epochs=2, validation_split=0.1\n",
+    ")\n",
+    "# Get quantized accuracy\n",
+    "_, q_aware_model_accuracy_finetuned = quantized_model.evaluate(test_images, test_labels, verbose=0)\n",
+    "print(\"Quantization test accuracy after fine-tuning:\", q_aware_model_accuracy_finetuned)\n",
+    "print(\"Baseline test accuracy (for reference):\", baseline_model_accuracy)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   },
+   "source": [
+    "```{note}\n",
+    "\n",
+    "If the network is not fully converged, the fine-tuned model's accuracy can surpass the original model's accuracy.\n",
+    "\n",
+    "```"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 45,
+   "metadata": {
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "WARNING:absl:Found untraced functions such as conv_0_layer_call_fn, conv_0_layer_call_and_return_conditional_losses, conv_1_layer_call_fn, conv_1_layer_call_and_return_conditional_losses, conv_2_layer_call_fn while saving (showing 5 of 14). These functions will not be directly callable after loading.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:tensorflow:Assets written to: GettingStarted/example/int8/saved_model/assets\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:tensorflow:Assets written to: GettingStarted/example/int8/saved_model/assets\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "ONNX conversion Done!\n"
+     ]
+    }
+   ],
+   "source": [
+    "# save TF INT8 original model\n",
+    "tf.keras.models.save_model(quantized_model, assets.example.int8_saved_model)\n",
+    "\n",
+    "# Convert INT8 model to ONNX\n",
+    "utils.convert_saved_model_to_onnx(saved_model_dir = assets.example.int8_saved_model, onnx_model_path = assets.example.int8_onnx_model)\n",
+    "\n",
+    "tf.keras.backend.clear_session()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "In this example, accuracy loss due to quantization is recovered in just two epochs.\n",
+    "\n",
+    "This NVIDIA(R) Quantization Toolkit provides an easy interface to create quantized networks, and thus take advantage of INT8 inference on NVIDIA(R) GPUs using TensorRT(TM)."
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "4442e1c252d743d7d1ab28567e302ebe8a15da81acb5d7e7894db75e10bdb29d"
+  },
+  "kernelspec": {
+   "display_name": "Python 3.8.10 64-bit ('base': conda)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.10"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tools/tensorflow-quantization/docs/source/docs/installation.md

@@ -0,0 +1,46 @@
+# **Installation**
+
+## **Docker**
+
+Latest TensorFlow 2.x [docker image](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tensorflow/tags) from NGC is recommended.
+
+Clone the `tensorflow-quantization` repository, pull the docker image, and launch the container.
+
+```{eval-rst}
+    .. code:: console
+
+        $ cd ~/
+        $ git clone https://github.com/NVIDIA/TensorRT.git
+        $ docker pull nvcr.io/nvidia/tensorflow:22.03-tf2-py3
+        $ docker run -it --runtime=nvidia --gpus all --net host -v ~/TensorRT/tools/tensorflow-quantization:/home/tensorflow-quantization nvcr.io/nvidia/tensorflow:22.03-tf2-py3 /bin/bash 
+```
+
+After the last command, you will be placed in the `/workspace` directory inside the running docker container, whereas the `tensorflow-quantization` repository is mounted in the `/home` directory.
+
+```{eval-rst}
+    .. code:: console
+
+        $ cd /home/tensorflow-quantization
+        $ ./install.sh
+        $ cd tests
+        $ python3 -m pytest quantize_test.py -rP
+
+```
+
+If all tests pass, installation is successful.
+
+## **Local**
+
+```{eval-rst}
+
+.. code:: console
+
+    $ cd ~/
+    $ git clone https://github.com/NVIDIA/TensorRT.git
+    $ cd TensorRT/tools/tensorflow-quantization
+    $ ./install.sh
+    $ cd tests
+    $ python3 -m pytest quantize_test.py -rP 
+```
+
+If all tests pass, installation is successful.

tools/tensorflow-quantization/docs/source/docs/intro_to_quantization.md

@@ -0,0 +1,109 @@
+# **Introduction to Quantization**
+
+## What is Quantization?
+[Quantization](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#working-with-int8) is the process of converting continuous values to discrete set of values using linear/non-linear scaling techniques.
+
+## Why Quantization?
+
+* High precision is necessary during training for fine-grained weight updates.
+
+* High precision is not usually necessary during inference and may hinder the deployment of AI models in real-time and/or in resource-limited devices.
+
+* INT8 is computationally less expensive and has lower memory footprint.
+
+* INT8 precision results in faster inference with similar performance.
+
+## Quantization Basics
+See [whitepaper](https://arxiv.org/abs/2004.09602) for more detailed explanations.
+
+Let [β, α] be the range of representable real values chosen for quantization and *`b`* be the bit-width of the signed integer representation.  
+The goal of uniform quantization is to map real values in the range [&beta; , &alpha;] to lie within [-2<sup>b-1</sup>, 2<sup>b-1</sup> - 1]. The real values that lie outside this range are clipped to the nearest bound.  
+
+*Affine Quantization*
+
+Considering 8 bit quantization (*b=8*), a real value within range [&beta;, &alpha;] is quantized to lie within the quantized range `[-128, 127]` (see [source](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Graph/Layers.html#iquantizelayer)):
+
+x<sub>q</sub>=clamp(round(x/scale)+zeroPt)
+
+where,
+scale = (&alpha; - &beta;) / (2<sup>b</sup>-1)  
+
+zeroPt = -round(&beta; * scale) - 2<sup>b-1</sup>  
+
+`round` is a function that rounds a value to the nearest integer. The quantized value is then clamped between -128 to 127.
+
+*Affine DeQuantization*
+
+DeQuantization is the reverse process of quantization (see [source](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Graph/Layers.html#idequantizelayer)):
+
+x=(x<sub>q</sub>−zeroPt)∗scale
+
+## Quantization in TensorRT  
+
+[TensorRT(TM)](https://developer.nvidia.com/tensorrt-getting-started) only supports symmetric uniform quantization, meaning that `zeroPt=0` (i.e. the quantized value of 0.0 is always 0).
+
+Considering 8 bit quantization (*`b=8`*), a real value within range [`min_float`, `max_float`] is quantized to lie within the quantized range `[-127, 127]`, opting not to use `-128` in favor of symmetry. It is important to note that we loose 1 value in symmetric quantization representation, however, loosing 1 out of 256 representable value for 8 bit quantization is insignificant.
+
+*Quantization*
+
+The mathematical representation for symmetric quantization (`zeroPt=0`) is:
+
+x<sub>q</sub>=clamp(round(x/scale))
+
+Since TensorRT supports only symmetric range, the scale is calculated using the max absolute value: `max(abs(min_float), abs(max_float))`. 
+
+Let &alpha; = `max(abs(min_float), abs(max_float))`,
+
+scale = &alpha;/(2<sup>b-1</sup>-1)
+
+Rounding [type](https://en.wikipedia.org/wiki/Rounding#Round_half_to_even) is rounding-to-nearest ties-to-even.
+The quantized value is then clamped between `-127` and `127`.
+
+*DeQuantization*
+Symmetric dequantization is the reverse process of symmetric quantization:
+
+x=(x<sub>q</sub>)∗scale
+
+
+## Intutions
+
+### Quantization Scale
+
+Scaling factor divides a given range of real values into a number of partitions.
+
+Lets understand intution behind scaling factor formula by taking 3 bit quantization as an example. 
+
+*Asymmetric Quantization*
+
+Real values range: [&beta;, &alpha;]
+
+Quantized values range: [-2<sup>3-1</sup>, 2<sup>3-1</sup>-1]  
+i.e. [-4, -3, -2, -1, 0, 1, 2, 3]
+
+As expected there are 8 quantized (2<sup>3</sup>) values for 3 bit quantization.
+
+Scale divides range into partitions. There are 7 (2<sup>3</sup>-1) partitions for 3 bit quantization.
+Thus,  
+scale = (&alpha; - &beta;) / (2<sup>3</sup>-1)
+
+*Symmetric Quantization*
+
+Symmetric quantization brings in two changes
+
+1. Real values are not free now but are restricted. i.e [-&alpha;, &alpha;]  
+where &alpha; = `max(abs(min_float), abs(max_float))`
+2. One value from quantization range is dropped in favor of symmetry leading to a new range [-3, -2, -1, 0, 1, 2, 3].  
+
+There are now 6 (2<sup>3</sup>-2) partitions (unlike 7 for asymmetric quantization).  
+
+Scale divides range into partitions.
+
+scale = 2*&alpha; /(2<sup>3</sup> - 2) = &alpha;/(2<sup>3-1</sup>-1)  
+
+Similar intution holds true for `b` bit quantization.
+
+### Quantization Zero Point
+
+The constant `zeroPt` is of the same type as quantized values x<sub>q</sub>, and is in fact the quantized value x<sub>q</sub> corresponding to the real value 0. This allows us to auto-matically meet the requirement that the real value r = 0 be exactly representable by a quantized value. The motivation for this requirement is that efficient implementation of neural network operators often requires zero-padding of arrays around boundaries.
+
+If we have values with negative data, then the zero point can offset the range. So if our zero point was 128, then unscaled negative values -127 to -1 would be represented by 1 to 127, and positive values 0 to 127 would be represented by 128 to 255.

tools/tensorflow-quantization/docs/source/docs/model_zoo.md

@@ -0,0 +1,97 @@
+# **Model Zoo Results**
+Results obtained on NVIDIA's A100 GPU and TensorRT 8.4.
+
+## [ResNet](https://github.com/NVIDIA/TensorRT/tree/main/tools/tensorflow-quantization/examples/resnet)
+### ResNet50-v1
+
+| Model                 | Accuracy (%) | Latency (ms, bs=1) |
+|-----------------------|--------------|--------------------|
+| Baseline (TensorFlow) | 75.05        | 7.95               |
+| PTQ (TensorRT)        | 74.96        | 0.46               |
+| **QAT** (TensorRT)    | 75.12        | 0.45               |
+
+### ResNet50-v2
+
+| Model                 | Accuracy (%) | Latency (ms, bs=1) |
+|-----------------------|--------------|--------------------|
+| Baseline (TensorFlow) | 75.36        | 6.16               |
+| PTQ (TensorRT)        | 75.48        | 0.57               |
+| **QAT** (TensorRT)    | 75.65        | 0.57               |
+
+### ResNet101-v1
+
+| Model                 | Accuracy (%) | Latency (ms, bs=1) |
+|-----------------------|--------------|--------------------|
+| Baseline (TensorFlow) | 76.47        | 15.92              |
+| PTQ (TensorRT)        | 76.32        | 0.84               |
+| **QAT** (TensorRT)    | 76.26        | 0.84               |
+
+### ResNet101-v2
+
+| Model                 | Accuracy (%) | Latency (ms, bs=1) |
+|-----------------------|--------------|--------------------|
+| Baseline (TensorFlow) | 76.89        | 14.13              |
+| PTQ (TensorRT)        | 76.94        | 1.05               |
+| **QAT** (TensorRT)    | 77.15        | 1.05               |
+
+*QAT fine-tuning hyper-parameters: `bs=32` (`bs=64` was OOM).*
+
+## [MobileNet](https://github.com/NVIDIA/TensorRT/tree/main/tools/tensorflow-quantization/examples/mobilenet)
+### MobileNet-v1
+
+| Model                 | Accuracy (%) | Latency (ms, bs=1)  |
+|-----------------------|--------------|---------------------|
+| Baseline (TensorFlow) | 70.60        | 1.99                |
+| PTQ (TensorRT)        | 69.31        | 0.16                |
+| **QAT** (TensorRT)    | 70.43        | 0.16                |
+
+### MobileNet-v2
+
+| Model                 | Accuracy (%) | Latency (ms, bs=1)    |
+|-----------------------|--------------|-----------------------|
+| Baseline (TensorFlow) | 71.77        | 3.71                  |
+| PTQ (TensorRT)        | 70.87        | 0.30                  |
+| **QAT** (TensorRT)    | 71.62        | 0.30                  |
+
+## [EfficientNet](https://github.com/NVIDIA/TensorRT/tree/main/tools/tensorflow-quantization/examples/efficientnet)
+### 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: `bs=64, ep=10, 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, lr=0.0001, steps_per_epoch=None`*.
+
+## [Inception](https://github.com/NVIDIA/TensorRT/tree/main/tools/tensorflow-quantization/examples/inception)
+### Inception-v3
+
+| Model                 | Accuracy (%) | Latency (ms, bs=1) |
+|-----------------------|--------------|--------------------|
+| Baseline (TensorFlow) | 77.86        | 9.01               |
+| PTQ (TensorRT)        | 77.73        | 0.82               |
+| **QAT** (TensorRT)    | 78.08        | 0.82               |
+
+```{eval-rst}
+
+.. NOTE::
+
+    The results here were obtained with NVIDIA's A100 GPU and TensorRT 8.4.
+    
+    Accuracy metric: Top-1 validation accuracy with the full ImageNet dataset.
+
+    Hyper-parameters
+
+    #. QAT fine-tuning: `bs=64`, `ep=10`, `lr=0.001` (unless otherwise stated).
+    #. PTQ calibration: `bs=64`.
+
+```

tools/tensorflow-quantization/docs/source/docs/qat.md

@@ -0,0 +1,37 @@
+(qat)=
+
+# **Quantization Aware Training (QAT)**
+
+The process of converting continuous to discrete values (Quantization) and vice-versa (Dequantization), requires `scale` and `zeroPt` (zero-point) parameters to be set.
+There are two quantization methods based on how these two parameters are calculated:
+
+```{eval-rst}
+
+#. `Post Training Quantization (PTQ) <https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#enable_int8_c>`_
+    Post Training Quantization computes `scale` after network has been trained. A representative dataset is used to capture the distribution of activations for each activation tensor, then this distribution data is used to compute the `scale` value for each tensor.
+    Each weight's distribution is used to compute weight `scale`.
+
+    TensorRT provides a workflow for PTQ, called `calibration`.
+
+    .. mermaid::
+
+        flowchart LR
+            id1(Calibration data) --> id2(Pre-trained model) --> id3(Capture layer distribution) --> id4(Compute 'scale') --> id5(Quantize model)
+
+#. `Quantization Aware Training (QAT) <https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#work-with-qat-networks>`_
+    Quantization Aware Training aims at computing scale factors during training. Once the network is fully trained, Quantize (Q) and Dequantize (DQ) nodes are inserted into the graph following a specific set of rules. The network is then further trained for few epochs in a process called `Fine-Tuning`. Q/DQ nodes simulate quantization loss and add it to the training loss during fine-tuning, making the network more resilient to quantization. In other words, QAT is able to better preserve accuracy when compared to PTQ.
+
+    .. mermaid::
+
+        flowchart LR
+            id1(Pre-trained model) --> id2(Add Q/DQ nodes) --> id3(Finetune model) --> id4(Store 'scale') --> id5(Quantize model)
+
+```
+
+```{attention}
+This toolkit supports only QAT as a quantization method. Note that we follow the quantization algorithm implemented by TensorRT(TM) when inserting Q/DQ nodes in a model. This leads to a quantized network with optimal layer fusion during the TensorRT(TM) engine building step.
+```
+
+````{note}
+Since TensorRT(TM) only supports symmetric quantization, we assume `zeroPt = 0`.
+````

tools/tensorflow-quantization/docs/source/globals.rst

@@ -0,0 +1,39 @@
+.. _globals_api:
+
+**tensorflow_quantization**
+============================
+
+
+:tensorflow_quantization.G_NUM_BITS:
+
+    8 bit quantization is used by default. However, it can be changed by using ``G_NUM_BITS`` global variable.
+    The following code snippet performs 4 bit quantization.
+
+    .. code:: python
+
+        import tensorflow_quantization
+        # get pretrained model
+        .....
+
+        # perform 4 bit quantization
+        tensorflow_quantization.G_NUM_BITS = 4
+        q_model = quantize_model(nn_model_original)
+
+        # fine-tune model
+        .....
+        
+    Check ``test_end_to_end_workflow_4bit()`` test case from ``quantize_test.py`` test module.
+
+:tensorflow_quantization.G_NARROW_RANGE:
+
+    If True, the absolute value of quantized minimum is the same as the quantized maximum value. For example,
+    minimum of -127 is used for 8 bit quantization instead of -128. TensorRT |tred| only supports G_NARROW_RANGE=True.
+
+:tensorflow_quantization.G_SYMMETRIC:
+
+    If True, 0.0 is always in the center of real min, max i.e. zero point is always 0. 
+    TensorRT |tred| only supports G_SYMMETRIC=True.
+
+.. attention:: When used, set global variables immediately before the ``quantize_model`` function call.
+
+.. |tred|    unicode:: U+2122 .. TRADEMARK SIGN

tools/tensorflow-quantization/docs/source/index.rst

@@ -0,0 +1,55 @@
+.. tensorflow-2.x-quantization documentation master file, created by
+   sphinx-quickstart on Tue Apr 12 14:52:54 2022.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+TensorFlow-2.x-Quantization-Toolkit
+=======================================================
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Toolkit
+
+   docs/installation.md
+   docs/basics.md
+   docs/getting_started.ipynb
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Tutorials
+
+   notebooks/simple_network_quantize_full.ipynb
+   notebooks/simple_network_quantize_partial.ipynb
+   notebooks/simple_network_quantize_specific_class.ipynb
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Examples
+
+   docs/example_resnet50v1.ipynb
+   docs/model_zoo.md
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Advanced Features
+
+   docs/add_new_layer_support.md
+   docs/add_custom_qdq_cases.md
+   
+.. toctree::
+   :maxdepth: 1
+   :caption: API Reference
+
+   globals.rst
+   qmodel.rst
+   qspec.rst
+   cqdq.rst
+   bqw.rst
+   utils.rst
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Quantization Theory
+
+   docs/intro_to_quantization.md
+   docs/qat.md

tools/tensorflow-quantization/docs/source/notebooks/simple_network_quantize_full.ipynb

@@ -0,0 +1,398 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "source": [
+    "(tut_one)=\n",
+    "\n",
+    "# **Full Network Quantization**  \n",
+    "\n",
+    "In this tutorial, we will take a sample network with ResNet-like network and perform ``full`` network quantization.\n",
+    "\n",
+    "\n",
+    "```{eval-rst}\n",
+    "\n",
+    ".. admonition:: Goal\n",
+    "    :class: note\n",
+    "\n",
+    "    #. Take a resnet-like model and train on cifar10 dataset.\n",
+    "    #. Perform full model quantization.\n",
+    "    #. Fine-tune to recover model accuracy.\n",
+    "    #. Save both original and quantized model while performing ONNX conversion.\n",
+    "\n",
+    "```\n",
+    "---"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "outputs": [],
+   "source": [
+    "#\n",
+    "# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n",
+    "# SPDX-License-Identifier: Apache-2.0\n",
+    "#\n",
+    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
+    "# you may not use this file except in compliance with the License.\n",
+    "# You may obtain a copy of the License at\n",
+    "#\n",
+    "# http://www.apache.org/licenses/LICENSE-2.0\n",
+    "#\n",
+    "# Unless required by applicable law or agreed to in writing, software\n",
+    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
+    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
+    "# See the License for the specific language governing permissions and\n",
+    "# limitations under the License.\n",
+    "#\n",
+    "\n",
+    "import tensorflow as tf\n",
+    "from tensorflow_quantization import quantize_model\n",
+    "import tiny_resnet\n",
+    "from tensorflow_quantization import utils\n",
+    "import os\n",
+    "\n",
+    "tf.keras.backend.clear_session()\n",
+    "\n",
+    "# Create folders to save TF and ONNX models\n",
+    "assets = utils.CreateAssetsFolders(os.path.join(os.getcwd(), \"tutorials\"))\n",
+    "assets.add_folder(\"simple_network_quantize_full\")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "outputs": [],
+   "source": [
+    "# Load CIFAR10 dataset\n",
+    "cifar10 = tf.keras.datasets.cifar10\n",
+    "(train_images, train_labels), (test_images, test_labels) = cifar10.load_data()\n",
+    "\n",
+    "# Normalize the input image so that each pixel value is between 0 and 1.\n",
+    "train_images = train_images / 255.0\n",
+    "test_images = test_images / 255.0"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "You must install pydot (`pip install pydot`) and install graphviz (see instructions at https://graphviz.gitlab.io/download/) for plot_model/model_to_dot to work.\n"
+     ]
+    }
+   ],
+   "source": [
+    "nn_model_original = tiny_resnet.model()\n",
+    "tf.keras.utils.plot_model(nn_model_original, to_file = assets.simple_network_quantize_full.fp32 + \"/model.png\")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/10\n",
+      "1407/1407 [==============================] - 16s 9ms/step - loss: 1.7653 - accuracy: 0.3622 - val_loss: 1.5516 - val_accuracy: 0.4552\n",
+      "Epoch 2/10\n",
+      "1407/1407 [==============================] - 13s 9ms/step - loss: 1.4578 - accuracy: 0.4783 - val_loss: 1.3877 - val_accuracy: 0.5042\n",
+      "Epoch 3/10\n",
+      "1407/1407 [==============================] - 13s 9ms/step - loss: 1.3499 - accuracy: 0.5193 - val_loss: 1.3066 - val_accuracy: 0.5342\n",
+      "Epoch 4/10\n",
+      "1407/1407 [==============================] - 13s 9ms/step - loss: 1.2736 - accuracy: 0.5486 - val_loss: 1.2636 - val_accuracy: 0.5550\n",
+      "Epoch 5/10\n",
+      "1407/1407 [==============================] - 13s 9ms/step - loss: 1.2101 - accuracy: 0.5732 - val_loss: 1.2121 - val_accuracy: 0.5670\n",
+      "Epoch 6/10\n",
+      "1407/1407 [==============================] - 12s 9ms/step - loss: 1.1559 - accuracy: 0.5946 - val_loss: 1.1753 - val_accuracy: 0.5844\n",
+      "Epoch 7/10\n",
+      "1407/1407 [==============================] - 13s 9ms/step - loss: 1.1079 - accuracy: 0.6101 - val_loss: 1.1143 - val_accuracy: 0.6076\n",
+      "Epoch 8/10\n",
+      "1407/1407 [==============================] - 13s 9ms/step - loss: 1.0660 - accuracy: 0.6272 - val_loss: 1.0965 - val_accuracy: 0.6158\n",
+      "Epoch 9/10\n",
+      "1407/1407 [==============================] - 13s 9ms/step - loss: 1.0271 - accuracy: 0.6392 - val_loss: 1.1100 - val_accuracy: 0.6122\n",
+      "Epoch 10/10\n",
+      "1407/1407 [==============================] - 13s 9ms/step - loss: 0.9936 - accuracy: 0.6514 - val_loss: 1.0646 - val_accuracy: 0.6304\n",
+      "Baseline FP32 model test accuracy: 61.65\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Train original classification model\n",
+    "nn_model_original.compile(\n",
+    "    optimizer=tiny_resnet.optimizer(lr=1e-4),\n",
+    "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+    "    metrics=[\"accuracy\"],\n",
+    ")\n",
+    "\n",
+    "nn_model_original.fit(\n",
+    "    train_images, train_labels, batch_size=32, epochs=10, validation_split=0.1\n",
+    ")\n",
+    "\n",
+    "# Get baseline model accuracy\n",
+    "_, baseline_model_accuracy = nn_model_original.evaluate(\n",
+    "    test_images, test_labels, verbose=0\n",
+    ")\n",
+    "baseline_model_accuracy = round(100 * baseline_model_accuracy, 2)\n",
+    "print(\"Baseline FP32 model test accuracy: {}\".format(baseline_model_accuracy))"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:tensorflow:Assets written to: /home/nvidia/PycharmProjects/tensorrt_qat/docs/source/notebooks/tutorials/simple_network_quantize_full/fp32/saved_model/assets\n",
+      "WARNING:tensorflow:From /home/nvidia/PycharmProjects/tensorrt_qat/venv38_tf2.8_newPR/lib/python3.8/site-packages/tf2onnx/tf_loader.py:711: extract_sub_graph (from tensorflow.python.framework.graph_util_impl) is deprecated and will be removed in a future version.\n",
+      "Instructions for updating:\n",
+      "Use `tf.compat.v1.graph_util.extract_sub_graph`\n",
+      "ONNX conversion Done!\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Save TF FP32 original model\n",
+    "tf.keras.models.save_model(nn_model_original, assets.simple_network_quantize_full.fp32_saved_model)\n",
+    "\n",
+    "# Convert FP32 model to ONNX\n",
+    "utils.convert_saved_model_to_onnx(saved_model_dir = assets.simple_network_quantize_full.fp32_saved_model, onnx_model_path = assets.simple_network_quantize_full.fp32_onnx_model)\n",
+    "\n",
+    "# Quantize model\n",
+    "q_nn_model = quantize_model(model=nn_model_original)\n",
+    "q_nn_model.compile(\n",
+    "    optimizer=tiny_resnet.optimizer(lr=1e-4),\n",
+    "    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),\n",
+    "    metrics=[\"accuracy\"],\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Test accuracy immediately after quantization:50.45, diff:11.199999999999996\n"
+     ]
+    }
+   ],
+   "source": [
+    "_, q_model_accuracy = q_nn_model.evaluate(test_images, test_labels, verbose=0)\n",
+    "q_model_accuracy = round(100 * q_model_accuracy, 2)\n",
+    "\n",
+    "print(\n",
+    "    \"Test accuracy immediately after quantization: {}, diff: {}\".format(\n",
+    "        q_model_accuracy, (baseline_model_accuracy - q_model_accuracy)\n",
+    "    )\n",
+    ")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "You must install pydot (`pip install pydot`) and install graphviz (see instructions at https://graphviz.gitlab.io/download/) for plot_model/model_to_dot to work.\n"
+     ]
+    }
+   ],
+   "source": [
+    "tf.keras.utils.plot_model(q_nn_model, to_file = assets.simple_network_quantize_full.int8 + \"/model.png\")"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Epoch 1/2\n",
+      "1407/1407 [==============================] - 27s 19ms/step - loss: 0.9625 - accuracy: 0.6630 - val_loss: 1.0430 - val_accuracy: 0.6420\n",
+      "Epoch 2/2\n",
+      "1407/1407 [==============================] - 25s 18ms/step - loss: 0.9315 - accuracy: 0.6758 - val_loss: 1.0717 - val_accuracy: 0.6336\n",
+      "Accuracy after fine-tuning for 2 epochs: 62.27\n",
+      "Baseline accuracy (for reference): 61.65\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Fine-tune quantized model\n",
+    "fine_tune_epochs = 2\n",
+    "\n",
+    "q_nn_model.fit(\n",
+    "    train_images,\n",
+    "    train_labels,\n",
+    "    batch_size=32,\n",
+    "    epochs=fine_tune_epochs,\n",
+    "    validation_split=0.1,\n",
+    ")\n",
+    "\n",
+    "_, q_model_accuracy = q_nn_model.evaluate(test_images, test_labels, verbose=0)\n",
+    "q_model_accuracy = round(100 * q_model_accuracy, 2)\n",
+    "print(\n",
+    "    \"Accuracy after fine-tuning for {} epochs: {}\".format(\n",
+    "        fine_tune_epochs, q_model_accuracy\n",
+    "    )\n",
+    ")\n",
+    "print(\"Baseline accuracy (for reference): {}\".format(baseline_model_accuracy))"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "WARNING:absl:Found untraced functions such as conv2d_layer_call_fn, conv2d_layer_call_and_return_conditional_losses, conv2d_1_layer_call_fn, conv2d_1_layer_call_and_return_conditional_losses, conv2d_2_layer_call_fn while saving (showing 5 of 18). These functions will not be directly callable after loading.\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "INFO:tensorflow:Assets written to: /home/nvidia/PycharmProjects/tensorrt_qat/docs/source/notebooks/tutorials/simple_network_quantize_full/int8/saved_model/assets\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "INFO:tensorflow:Assets written to: /home/nvidia/PycharmProjects/tensorrt_qat/docs/source/notebooks/tutorials/simple_network_quantize_full/int8/saved_model/assets\n"
+     ]
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "ONNX conversion Done!\n"
+     ]
+    }
+   ],
+   "source": [
+    "# Save TF INT8 original model\n",
+    "tf.keras.models.save_model(q_nn_model, assets.simple_network_quantize_full.int8_saved_model)\n",
+    "\n",
+    "# Convert INT8 model to ONNX\n",
+    "utils.convert_saved_model_to_onnx(saved_model_dir = assets.simple_network_quantize_full.int8_saved_model, onnx_model_path = assets.simple_network_quantize_full.int8_onnx_model)\n",
+    "\n",
+    "tf.keras.backend.clear_session()"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%%\n"
+    }
+   }
+  },
+  {
+   "cell_type": "markdown",
+   "source": [
+    "```{note}\n",
+    "ONNX files can be visualized with [Netron](https://netron.app/).\n",
+    "```"
+   ],
+   "metadata": {
+    "collapsed": false,
+    "pycharm": {
+     "name": "#%% md\n"
+    }
+   }
+  }
+ ],
+ "metadata": {
+  "interpreter": {
+   "hash": "4442e1c252d743d7d1ab28567e302ebe8a15da81acb5d7e7894db75e10bdb29d"
+  },
+  "kernelspec": {
+   "display_name": "Python 3.8.10 ('base')",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.8.10"
+  },
+  "orig_nbformat": 4
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tools/tensorflow-quantization/docs/source/notebooks/tiny_resnet.py

@@ -0,0 +1,78 @@
+#
+# 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.
+#
+
+"""
+A small resnet-like network for quick testing.
+"""
+
+import tensorflow as tf
+
+
+def identity_block(input_tensor):
+    """
+    Identity block with no shortcut convolution
+    """
+    y = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), padding="same")(
+        input_tensor
+    )
+    y = tf.keras.layers.ReLU()(y)
+    y = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3), padding="same")(y)
+    out = tf.keras.layers.Add()([y, input_tensor])
+    out = tf.keras.layers.ReLU()(out)
+    return out
+
+
+def identity_block_short_conv(input_tensor):
+    """
+    Identity block with shortcut convolution
+    """
+    y = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), padding="same")(
+        input_tensor
+    )
+    y = tf.keras.layers.ReLU()(y)
+    y = tf.keras.layers.Conv2D(
+        filters=24, kernel_size=(3, 3), strides=(2, 2), padding="same"
+    )(y)
+    ds_input = tf.keras.layers.Conv2D(
+        filters=24, kernel_size=(3, 3), strides=(2, 2), padding="same"
+    )(input_tensor)
+    out = tf.keras.layers.Add()([y, ds_input])
+    out = tf.keras.layers.ReLU()(out)
+    return out
+
+
+def model():
+    """
+    Dummy network with resnet-like architecture.
+    """
+    input_img = tf.keras.layers.Input(shape=(32, 32, 3))
+    x = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3))(input_img)
+    x = tf.keras.layers.ReLU()(x)
+    x = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3))(x)
+    x = tf.keras.layers.ReLU()(x)
+    x = identity_block(x)
+    x = identity_block_short_conv(x)
+    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
+    x = tf.keras.layers.Flatten()(x)
+    x = tf.keras.layers.Dense(100)(x)
+    x = tf.keras.layers.ReLU()(x)
+    x = tf.keras.layers.Dense(10)(x)
+    return tf.keras.Model(input_img, x, name="Dummy_Model")
+
+
+def optimizer(lr=0.001):
+    return tf.keras.optimizers.Adam(learning_rate=lr)

tools/tensorflow-quantization/docs/source/qmodel.rst

@@ -0,0 +1,33 @@
+.. _qmodel_api:
+
+**tensorflow_quantization.quantize_model**
+============================================
+
+.. automodule:: tensorflow_quantization.quantize
+   :members: quantize_model
+
+.. note:: Currently only Functional and Sequential models are supported.
+
+Examples
+
+.. code:: python
+
+   import tensorflow as tf
+   from tensorflow_quantization.quantize import quantize_model
+
+   # Simple full model quantization.
+   # 1. Create a simple network
+   input_img = tf.keras.layers.Input(shape=(28, 28))
+   r = tf.keras.layers.Reshape(target_shape=(28, 28, 1))(input_img)
+   x = tf.keras.layers.Conv2D(filters=2, kernel_size=(3, 3))(r)
+   x = tf.keras.layers.ReLU()(x)
+   x = tf.keras.layers.Conv2D(filters=2, kernel_size=(3, 3))(x)
+   x = tf.keras.layers.ReLU()(x)
+   x = tf.keras.layers.Flatten()(x)
+   model = tf.keras.Model(input_img, x)
+
+   print(model.summary())
+
+   # 2. Quantize the network
+   q_model = quantize_model(model)
+   print(q_model.summary())

tools/tensorflow-quantization/docs/source/qspec.rst

@@ -0,0 +1,166 @@
+.. _qspec_api:
+
+**tensorflow_quantization.QuantizationSpec**
+=============================================
+.. autoclass:: tensorflow_quantization.QuantizationSpec
+   :members:
+
+Examples
+
+Let's write a simple network to use in all examples.
+
+.. code-block:: python
+
+   import tensorflow as tf
+   # Import necessary methods from the Quantization Toolkit
+   from tensorflow_quantization.quantize import quantize_model, QuantizationSpec
+
+   # 1. Create a small network
+   input_img = tf.keras.layers.Input(shape=(28, 28))
+   x = tf.keras.layers.Reshape(target_shape=(28, 28, 1))(input_img)
+   x = tf.keras.layers.Conv2D(filters=126, kernel_size=(3, 3))(x)
+   x = tf.keras.layers.ReLU()(x)
+   x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3))(x)
+   x = tf.keras.layers.ReLU()(x)
+   x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3))(x)
+   x = tf.keras.layers.ReLU()(x)
+   x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3, 3))(x)
+   x = tf.keras.layers.ReLU()(x)
+   x = tf.keras.layers.Conv2D(filters=8, kernel_size=(3, 3))(x)
+   x = tf.keras.layers.ReLU()(x)
+   x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
+   x = tf.keras.layers.Flatten()(x)
+   x = tf.keras.layers.Dense(100)(x)
+   x = tf.keras.layers.ReLU()(x)
+   x = tf.keras.layers.Dense(10)(x)
+   model = tf.keras.Model(input_img, x)
+
+
+#. **Select layers based on layer names**  
+
+   **Goal**: Quantize the 2nd Conv2D, 4th Conv2D and 1st Dense layer in the following network.
+
+   .. code-block:: python
+
+      # 1. Find out layer names
+      print(model.summary())
+
+      # 2. Create quantization spec and add layer names
+      q_spec = QuantizationSpec()
+      layer_name = ['conv2d_1', 'conv2d_3', 'dense']
+
+      """
+      # Alternatively, each layer configuration can be added one at a time:
+      q_spec.add('conv2d_1')
+      q_spec.add('conv2d_3')
+      q_spec.add('dense')
+      """
+
+      q_spec.add(name=layer_name)
+
+      # 3. Quantize model
+      q_model = quantize_model(model, quantization_mode='partial', quantization_spec=q_spec)
+      print(q_model.summary())
+
+      tf.keras.backend.clear_session()
+
+
+#. **Select layers based on layer class**
+   
+   **Goal**: Quantize all `Conv2D` layers.
+
+   .. code-block:: python
+
+      # 1. Create QuantizationSpec object and add layer class
+      q_spec = QuantizationSpec()
+      q_spec.add(name='Conv2D', is_keras_class=True)
+
+      # 2. Quantize model
+      q_model = quantize_model(model, quantization_mode='partial', quantization_spec=q_spec)
+      q_model.summary()
+
+      tf.keras.backend.clear_session()
+
+#. **Select layers based both layer name and layer class**
+
+   **Goal**: Quantize all `Dense` layers and the 3rd `Conv2D` layer.
+
+   .. code-block:: python
+
+      # 1. Create QuantizationSpec object and add layer information
+      q_spec = QuantizationSpec()
+
+      layer_name = ['Dense', 'conv2d_2']
+      layer_is_keras_class = [True, False]
+
+      """
+      # Alternatively, each layer configuration can be added one at a time:
+      q_spec.add(name='Dense', is_keras_class=True)
+      q_spec.add(name='conv2d_2')
+      """
+
+      q_spec.add(name=layer_name, is_keras_class=layer_is_keras_class)
+
+      # 2. Quantize model
+      q_model = quantize_model(model, quantization_mode='partial', quantization_spec=q_spec)
+      q_model.summary()
+
+      tf.keras.backend.clear_session()
+
+#. **Select inputs at specific index for multi-input layers**  
+
+   For layers with multiple inputs, the user can choose which ones need to be quantized. Assume a network that has two layers of class `Add`.
+
+   **Goal**: Quantize index 1 of `add` layer, index 0 of `add_1` layer and the 3rd `Conv2D` layer.
+
+   .. code-block:: python
+
+      # 1. Create QuantizationSpec object and add layer information
+      q_spec = QuantizationSpec()
+
+      layer_name = ['add', 'add_1', 'conv2d_2']
+      layer_q_indices = [[1], [0], None]
+
+      """
+      # Alternatively, each layer configuration can be added one at a time:
+      q_spec.add(name='add', quantization_index=[1])
+      q_spec.add(name='add', quantization_index=[0])
+      q_spec.add(name='conv2d_2')
+      """
+
+      q_spec.add(name=layer_name, quantization_index=layer_q_indices)
+
+      # 2. Quantize model
+      q_model = quantize_model(model, quantization_mode='partial', quantization_spec=q_spec)
+      q_model.summary()
+
+      tf.keras.backend.clear_session()
+
+#. **Quantize only weight and NOT input**  
+
+   **Goal**: Quantize the 2nd Conv2D, 4th Conv2D and 1st Dense layer in the following network. In addition to that, quantize only the weights of the 2nd Conv2D.
+
+   .. code-block:: python
+
+      # 1. Find out layer names
+      print(model.summary())
+
+      # 2. Create quantization spec and add layer names
+      q_spec = QuantizationSpec()
+      layer_name = ['conv2d_1', 'conv2d_3', 'dense']
+      layer_q_input = [False, True, True]
+
+      """
+      # Alternatively, each layer configuration can be added one at a time:
+      q_spec.add('conv2d_1', quantize_input=False)
+      q_spec.add('conv2d_3')
+      q_spec.add('dense')
+      """
+
+      q_spec.add(name=layer_name, quantize_input=layer_q_input)
+
+      # 3. Quantize model
+      q_model = quantize_model(model, quantization_mode='partial', quantization_spec=q_spec)
+      print(q_model.summary())
+
+      tf.keras.backend.clear_session()

tools/tensorflow-quantization/docs/source/utils.rst

@@ -0,0 +1,8 @@
+.. _utils_api:
+
+**tensorflow_quantization.utils**
+============================================
+
+.. automodule:: tensorflow_quantization.utils
+   :members:
+   :undoc-members:

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,
+        }

tools/tensorflow-quantization/install.sh

@@ -0,0 +1,8 @@
+#!/bin/bash
+
+pip install nvidia-pyindex
+pip install onnx-graphsurgeon
+pip install git+https://github.com/onnx/tensorflow-onnx.git
+pip install tensorflow-gpu==2.8.0 tensorflow-datasets
+pip install numpy pytest pytest-html graphviz
+pip install .

tools/tensorflow-quantization/rebuild.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+python -m pip uninstall --yes tensorflow_quantization
+rm -rf  tensorflow_quantization.egg-info
+rm -rf build
+python -m pip install .

tools/tensorflow-quantization/setup.py

@@ -0,0 +1,65 @@
+#
+# 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.
+#
+
+from setuptools import setup, find_packages
+from pathlib import Path
+import os
+
+abspath = os.path.dirname(os.path.realpath(__file__))
+
+license_header = """#
+# 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.
+#
+"""
+
+# Generate version file
+with open(os.path.join(abspath, "VERSION")) as f:
+    version = f.read().strip()
+with open(os.path.join(abspath, "tensorflow_quantization/version.py"), "w") as f:
+    f.write(license_header)
+    f.write(F"__version__ = \"{version}\"")
+
+project_dir = Path(__file__).parent
+
+# Setting up
+setup(
+    name="tensorflow_quantization",
+    version=version,
+    description="NVIDIA TensorFlow 2.x quantization toolkit",
+    long_description=Path("README.md").read_text(),
+    long_description_content_type="text/markdown",
+    packages=["tensorflow_quantization"],
+    python_requires=">=3.6",
+    include_package_data=True,
+    install_requires=["tensorflow-gpu==2.8.0", "tf2onnx==1.10.1"],
+    author="NVIDIA",
+    author_email="nvidia@nvidia.com",
+    license="Apache 2.0",
+)

tools/tensorflow-quantization/tensorflow_quantization/__init__.py

@@ -0,0 +1,37 @@
+#
+# 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.
+#
+
+
+G_NUM_BITS: int = 8
+G_NARROW_RANGE: bool = True
+G_SYMMETRIC: bool = True
+
+from tensorflow_quantization.quantize import QuantizationSpec
+from tensorflow_quantization.quantize import quantize_model
+
+from tensorflow_quantization.quantize_wrapper_base import BaseQuantizeWrapper
+
+from tensorflow_quantization.quantize_wrappers import WeightedBaseQuantizeWrapper
+from tensorflow_quantization.quantize_wrappers import NonWeightedBaseQuantizeWrapper
+
+from tensorflow_quantization.custom_qdq_case_base import CustomQDQInsertionCase
+
+from tensorflow_quantization.utils import CreateAssetsFolders
+from tensorflow_quantization.utils import convert_saved_model_to_onnx
+from tensorflow_quantization.utils import convert_keras_model_to_onnx
+
+from .version import __version__

tools/tensorflow-quantization/tensorflow_quantization/custom_qdq_case_base.py

@@ -0,0 +1,48 @@
+#
+# 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.
+#
+
+
+from abc import ABC
+
+
+class CustomQDQInsertionCase(ABC):
+    """
+    This class helps user to programatically decide toolkit behavior to quantize specific layers.
+    Based on the output of this class 'case' function, toolkit deviates from its standard behavior.
+    """
+
+    def info(self) -> str:
+        return ""
+
+    def case(
+        self, keras_model: "tf.keras.Model", qspec: "QuantizationSpec"
+    ) -> "QuantizationSpec":
+        """
+        This function is called internally by the framework.
+        Given keras model is passed as an argument and object of QuantizationSpec class
+        is expcted in return.
+        Returned QuantzaionSpec class object should contain information about the layers that needs
+        to be treated specially/differently from default framework behavior.
+
+        Args:
+            keras_model (tf.keras.Model): Keras functional or sequentail model
+            qspec (QuantizationSpec): User passed QuantizationSpec object. It is important to note that
+            new special qdq might or might not use quantizations specs user has provided.
+        Returns:
+            A new QuantizationSpec object.
+        """
+        raise NotImplementedError("case method must be overridden by user")

tools/tensorflow-quantization/tensorflow_quantization/custom_qdq_cases.py

@@ -0,0 +1,363 @@
+#
+# 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.
+#
+
+from tensorflow_quantization import CustomQDQInsertionCase
+from tensorflow_quantization import QuantizationSpec
+from tensorflow_quantization import utils
+import tensorflow as tf
+from typing import List
+
+
+def is_parent_type(parent_class: str, class_type="Conv") -> bool:
+    """
+    Checks if 'parent_class' is of type 'type'.
+    Examples of types: Conv, BatchNorm, Dropout, Activation.
+    """
+    return class_type in parent_class
+
+
+def is_parent_pattern(parent_info: dict, pattern: List = ["BatchNorm", "Conv"]) -> bool:
+    """ Checks if parent heritage follows a specific 'pattern'.
+    Args:
+        parent_info (dict): dictionary with parent's information.
+        pattern (List): list containing a layer's parental heritage ([parent, grandparent, great-grandparent, ...]).
+    Returns:
+        bool: indicating whether a layer's parent heritage follows the given pattern.
+    """
+    grandparent_info = parent_info
+    for i, p in enumerate(pattern):
+        if i > 0:
+            grandparent_info = utils._get_previous_layers_class_and_module_and_name(
+                grandparent_info["layer"]
+            )[0]
+        if not is_parent_type(grandparent_info["class"], class_type=p):
+            return False
+    return True
+
+
+def check_is_quantizable_by_layer_name(
+    qspec: QuantizationSpec, current_layer_name: str
+) -> bool:
+    """
+    Checks if 'current_layer_name' is in 'qspec'. It returns True if 'current_layer_name' is NOT in 'qspec' and
+      False if it is. This means that the user's request will get prioritized over our automatic methods.
+
+    Args:
+         qspec (QuantizationSpec): quantization specification.
+
+    Returns:
+        is_quantizable_by_layer_name (bool): boolean indicating whether 'current_layer_name' is quantizable by our
+          method (is NOT in 'qspec'), or not (is in 'qspec', so that configuration should be followed).
+    """
+    def _is_layer_in_user_passed_qspec(layer_name):
+        for l in qspec.layers:
+            if l.name == layer_name:
+                return True
+        return False
+
+    is_quantizable_by_layer_name = qspec is None or (
+        qspec is not None and not _is_layer_in_user_passed_qspec(current_layer_name)
+    )
+    return is_quantizable_by_layer_name
+
+
+###################################################################
+################# General Custom QDQ Cases ########################
+###################################################################
+
+
+class BNQDQCase(CustomQDQInsertionCase):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def info(self) -> str:
+        return "Avoids inserting QDQ before BatchNorm in cases where BN is connected to a Conv layer (since that BN " \
+               "will be fused with previous Conv layer). This case happens in ResNet-v2, where the following pattern " \
+               "exists: BN-ReLU-Conv blocks (pre-activation function). In that scenario, BN is sometimes connected to " \
+               "`Add` layer, which doesn't fuse with BN."
+
+    def case(
+        self, keras_model: tf.keras.Model, qspec: QuantizationSpec
+    ) -> QuantizationSpec:
+        def _check_if_quantizable_bn(layer):
+            layer_parent = layer.input._keras_history.layer
+            parent_class_name = layer_parent.__class__.__name__
+            if not is_parent_type(parent_class_name, class_type="Conv"):
+                if check_is_quantizable_by_layer_name(qspec, layer.name):
+                    return True
+            return False
+
+        bn_qspec = QuantizationSpec()
+        for layer in keras_model.layers:
+            if isinstance(layer, tf.keras.layers.BatchNormalization):
+                """
+                Returns quantizable BatchNorm layers: All BN layers that are not connected to a Conv layer.
+                In other words, don't add QDQ to BN layers in a Conv-BN sequence (and of course, if it shouldn't be
+                ignored due to the user's preference)."
+                """
+                if _check_if_quantizable_bn(layer):
+                    bn_qspec.add(
+                        name=layer.name, quantize_input=True, quantize_weight=False
+                    )
+        return bn_qspec
+
+
+class ResidualConnectionQDQCase(CustomQDQInsertionCase):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def info(self) -> str:
+        info_str = "Goal: To return all quantizable residual inputs. " \
+                   "Rules: Residual connection is represented by the Add layer. The recommendation from the TRT team " \
+                   "        is to add QDQ to all of its inputs except when: " \
+                   "     - the input is Bias. Note that TF sees MatMul+BiasAdd as a Dense layer, so no need to check " \
+                   "       if the input is Bias. " \
+                   "     - in the case of one of the inputs being a simple residual branch and the other Conv or " \
+                   "       Conv+BN, add QDQ nodes to just the residual branch. This is needed to trigger an INT8 " \
+                   "       kernel fusion with Add. " \
+                   "     - in the case of more than one input being Conv or Conv+BN, add QDQ to all inputs except 1. " \
+                   " The last 2 cases are needed to trigger an INT8 kernel fusion with Add. " \
+                   " [ResNet-v1]:  Note that the connection between Conv2D and Add layer is not direct: " \
+                   "                   Conv2D -> BatchNormalization -> Add " \
+                   "               To get to the layer, we need to access `input._keras_history.layer` " \
+                   "               This is the same for EfficientNet-B0. " \
+                   " [ResNet-v2]: Connection is direct ReLU -> Conv2D -> Add " \
+                   " [EfficientNet-B0]: Contains two special patterns: " \
+                   "                      1. Conv -> BatchNorm -> Activation -> Add " \
+                   "                      2. Conv -> BatchNorm -> Activation -> Dropout -> Add"
+        return info_str
+
+    def case(
+        self, keras_model: tf.keras.Model, qspec: QuantizationSpec
+    ) -> QuantizationSpec:
+
+        res_qspec = QuantizationSpec()
+        for layer in keras_model.layers:
+            if isinstance(layer, tf.keras.layers.Add) and check_is_quantizable_by_layer_name(qspec, layer.name):
+                """
+                Returns quantizable inputs to Add layers: all inputs except 1 with 'pattern'.
+                Patterns checked for: Conv, Conv-BN, Conv-BN-Activation, Conv-BN-Activation-Dropout.
+                """
+                layer_parents = utils.find_my_predecessors(keras_model, layer.name)
+
+                # Collect the non-quantizable input (1 branch with Conv pattern)
+                input_indices_convs = []
+                for i, l_parent_info in enumerate(layer_parents):
+                    l_parent_class = l_parent_info["class"]
+                    l_parent_layer = l_parent_info["layer"]
+                    # Check that the input is a Conv pattern
+                    if (
+                            is_parent_type(l_parent_class, class_type="Conv")
+                            or is_parent_pattern(l_parent_info, pattern=["BatchNorm", "Conv"])
+                            or is_parent_pattern(l_parent_info, pattern=["Activation", "BatchNorm", "Conv"])
+                            or is_parent_pattern(l_parent_info, pattern=["Dropout", "Activation", "BatchNorm", "Conv"])
+                    ):
+                        # Check that it's not a residual branch (input does not have more than 1 outbound node)
+                        if hasattr(l_parent_layer, 'outbound_nodes'):
+                            num_outbound_nodes = len(getattr(l_parent_layer, 'outbound_nodes'))
+                            if num_outbound_nodes == 1:
+                                # Branch without QDQ branch is chosen
+                                input_indices_convs.append(i)
+                                break
+
+                # Default behavior: add QDQ in all inputs except 1 with Conv/BN
+                input_indices = list(range(0, len(layer_parents)))
+                if len(input_indices_convs) > 0:
+                    # Don't quantize one of the Conv pattern branches.
+                    index_to_delete = input_indices_convs[-1]
+                    del input_indices[index_to_delete]
+                if len(input_indices) > 0:
+                    res_qspec.add(
+                        layer.name,
+                        quantize_input=True,
+                        quantize_weight=False,
+                        quantization_index=input_indices,
+                    )
+        return res_qspec
+
+
+class MaxPoolQDQCase(CustomQDQInsertionCase):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def info(self) -> str:
+        return "Enables quantization of MaxPool layers. This is needed in cases where MaxPool is added to a residual " \
+               "connection and where the other branches are already quantized (needed to trigger a horizontal fusion " \
+               "in the residual connection. This case happens in ResNet-v2."
+
+    def case(
+        self, keras_model: tf.keras.Model, qspec: QuantizationSpec
+    ) -> QuantizationSpec:
+        mp_qspec = QuantizationSpec()
+        for layer in keras_model.layers:
+            if isinstance(layer, tf.keras.layers.MaxPooling2D):
+                """
+                Returns quantizable MaxPooling2D layers.
+                """
+                if check_is_quantizable_by_layer_name(qspec, layer.name):
+                    mp_qspec.add(
+                        name=layer.name,
+                        quantize_input=True,
+                        quantize_weight=False
+                    )
+        return mp_qspec
+
+
+###################################################################
+############ Network Specific QDQ Cases ###########################
+###################################################################
+
+
+class ResNetV1QDQCase(CustomQDQInsertionCase):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def info(self) -> str:
+        return (
+            "Returns all quantizable nodes in ResNet-v1: "
+            "  1. Residual connections."
+        )
+
+    def case(
+        self, keras_model: tf.keras.Model, qspec: QuantizationSpec
+    ) -> QuantizationSpec:
+        special_qspec = QuantizationSpec()
+
+        # Use Residual connection QDQ
+        residual_cqdq = ResidualConnectionQDQCase()
+        residual_cqdq_qspec = residual_cqdq.case(keras_model, qspec)
+        special_qspec.layers.extend(residual_cqdq_qspec.layers)
+
+        return special_qspec
+
+
+class ResNetV2QDQCase(CustomQDQInsertionCase):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def info(self) -> str:
+        return (
+            "Returns all quantizable nodes in ResNet-v2: "
+            "  1. Residual connections, "
+            "  2. BatchNorm not connected to Conv, "
+            "  3. MaxPool layers."
+        )
+
+    def case(
+        self, keras_model: tf.keras.Model, qspec: QuantizationSpec
+    ) -> QuantizationSpec:
+        special_qspec = QuantizationSpec()
+
+        # Use Residual connection QDQ
+        residual_cqdq = ResidualConnectionQDQCase()
+        residual_cqdq_qspec = residual_cqdq.case(keras_model, qspec)
+        special_qspec.layers.extend(residual_cqdq_qspec.layers)
+
+        # Use BN QDQ Case
+        bn_cqdq = BNQDQCase()
+        bn_cqdq_qspec = bn_cqdq.case(keras_model, qspec)
+        special_qspec.layers.extend(bn_cqdq_qspec.layers)
+
+        # Use MaxPool QDQ Case (necessary for ResNet-v2)
+        mp_cqdq = MaxPoolQDQCase()
+        mp_cqdq_qspec = mp_cqdq.case(keras_model, qspec)
+        special_qspec.layers.extend(mp_cqdq_qspec.layers)
+
+        return special_qspec
+
+
+class EfficientNetQDQCase(CustomQDQInsertionCase):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def info(self) -> str:
+        return (
+            "Returns all quantizable nodes in EfficientNet:"
+            "  1. Residual connections,"
+            "  2. Quantize inputs (0, 1) of Multiply layers in SE (Squeeze-Excite) block."
+        )
+
+    def case(
+        self, keras_model: tf.keras.Model, qspec: QuantizationSpec
+    ) -> QuantizationSpec:
+        special_qspec = QuantizationSpec()
+
+        # Use Residual connection QDQ
+        residual_cqdq = ResidualConnectionQDQCase()
+        residual_cqdq_qspec = residual_cqdq.case(keras_model, qspec)
+        special_qspec.layers.extend(residual_cqdq_qspec.layers)
+
+        # Implement EfficientNet specific case to trigger horizontal fusion in Mul residual branch.
+        #   Gives preference to the user-specified `qspec`.
+        for layer in keras_model.layers:
+            if (
+                    isinstance(layer, tf.keras.layers.Multiply)
+                    and check_is_quantizable_by_layer_name(qspec, layer.name)
+            ):
+                special_qspec.add(
+                    layer.name,
+                    quantize_input=True,
+                    quantize_weight=False,
+                    quantization_index=[0, 1],
+                )
+        return special_qspec
+
+
+class MobileNetQDQCase(CustomQDQInsertionCase):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def info(self) -> str:
+        return (
+            "Returns all quantizable nodes in MobileNet: "
+            "  1. Residual connections."
+        )
+
+    def case(
+        self, keras_model: tf.keras.Model, qspec: QuantizationSpec
+    ) -> QuantizationSpec:
+        special_qspec = QuantizationSpec()
+
+        # Use Residual connection QDQ
+        residual_cqdq = ResidualConnectionQDQCase()
+        residual_cqdq_qspec = residual_cqdq.case(keras_model, qspec)
+        special_qspec.layers.extend(residual_cqdq_qspec.layers)
+
+        return special_qspec
+
+
+class InceptionQDQCase(CustomQDQInsertionCase):
+    def __init__(self) -> None:
+        super().__init__()
+
+    def info(self) -> str:
+        return (
+            "Returns all quantizable nodes in Inception-v3: "
+            "  1. MaxPool layers to trigger horizontal fusion in the output of Concat."
+        )
+
+    def case(
+        self, keras_model: tf.keras.Model, qspec: QuantizationSpec
+    ) -> QuantizationSpec:
+        special_qspec = QuantizationSpec()
+
+        # Use MaxPool QDQ Case
+        mp_cqdq = MaxPoolQDQCase()
+        mp_cqdq_qspec = mp_cqdq.case(keras_model, qspec)
+        special_qspec.layers.extend(mp_cqdq_qspec.layers)
+
+        return special_qspec

tools/tensorflow-quantization/tensorflow_quantization/global_config.py

@@ -0,0 +1,64 @@
+#
+# 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 holds the quantization config class object that is accessed globally by library modules.
+DIRECT USE OF THIS MODULE BY USER IS PROHIBITED.
+"""
+
+# List that holds quantization config class object, Length is always one!
+# Object is added automatically on class creation
+G_CONFIG_OBJECT = []
+
+
+def add_config_object(config_object: "BaseConfig") -> None:
+    """
+    Add instance of quantize config class to the global list.
+    Args:
+        config_object : Instance of one of four quantize config class
+    """
+    assert (
+            len(G_CONFIG_OBJECT) == 0
+    ), "Looks like previous quatize object is alive. Did you call clear() on the object?"
+    G_CONFIG_OBJECT.append(config_object)
+
+
+def remove_config_object() -> None:
+    """
+    Remove instance of quantize config class from the global list.
+    """
+    if G_CONFIG_OBJECT:
+        G_CONFIG_OBJECT.clear()
+
+
+def get_config_object() -> "BaseConfig":
+    """
+    Return quantize config class object
+    """
+    assert (
+            len(G_CONFIG_OBJECT) == 1
+    ), "Have you created quantize config object before calling `quantize_model`?"
+    if G_CONFIG_OBJECT:
+        return G_CONFIG_OBJECT[0]
+
+
+def is_config_object_created() -> bool:
+    """
+    Sanity check function for whether quantize config class object is created.
+    """
+    return len(G_CONFIG_OBJECT) == 1

tools/tensorflow-quantization/tensorflow_quantization/quantize.py

@@ -0,0 +1,445 @@
+#
+# 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 2019 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.
+
+from typing import List, Union
+import tensorflow as tf
+import tensorflow_quantization.quantize_wrappers as quantize_wrappers
+import tensorflow_quantization.global_config as cfg
+import tensorflow_quantization.quantize_config as quantize_config
+from dataclasses import dataclass
+from tensorflow_quantization.quantize_wrappers import DISABLED_LAYER_QUANTIZATION_DEFAULT
+
+
+@dataclass
+class LayerConfig:
+    """
+    Internal dataclass for a single layer config.
+    Args:
+        name (str): Name of the layer. As seen from utilities such as `model.summary()`
+        is_keras_class (bool) : Set this to True if layer_name passed represents a layer class from Keras.
+            Default is False.
+        quantize_input (bool): Set this to True if input to the layers should be quantized. Default is True
+            since default behavior is following Nvidia quantization recipe.
+        quantize_weight (bool): Set this to True if weights to the layers should be quantized. Default is True
+            since default behavior is following Nvidia quantization recipe. For weightless layers, value is
+            ignored.
+        quantization_index (List): Indices on inputs to which quantization is applied for the layers with
+            multiple inputs. E.g Add, Concatenate
+    Returns:
+        None
+    """
+
+    name: str = None
+    is_keras_class: bool = False
+    quantize_input: bool = True
+    quantize_weight: bool = True
+    quantization_index: list = None
+
+
+class QuantizationSpec:
+    """
+    Helper class holding config objects for all layers to quantize.
+    """
+
+    def __init__(self) -> None:
+        self.layers = []
+
+    def __str__(self) -> str:
+        for l in self.layers:
+            print(l)
+        return ""
+
+    def add(
+        self,
+        name: Union[str, List],
+        is_keras_class: Union[bool, List] = False,
+        quantize_input: Union[bool, List] = True,
+        quantize_weight: Union[bool, List] = True,
+        quantization_index: Union[List, List[List]] = None,
+    ) -> None:
+        """
+        Takes user parameters and adds LayerConfig object to a list for each add call.
+
+        Args:
+            name (Union[str, List]): Name of the layer. As seen from utilities such as `model.summary()`
+            is_keras_class (Union[bool, List]): List or a single value. Set this to True if layer_name passed represents a layer class from Keras.
+                Default is False.
+            quantize_input (Union[bool, List]): List or a single value. Set this to True if input to the layers should be quantized. Default is True
+                since default behavior is following Nvidia quantization recipe.
+            quantize_weight (Union[bool, List]): List or a single value. Set this to True if weights to the layers should be quantized. Default is True
+                since default behavior is following Nvidia quantization recipe. For weightless layers, value is
+                ignored.
+            quantization_index (Union[List, List[List]]): List or List of List. List with indices on inputs to which quantization is applied for the layers with
+                multiple inputs. E.g Add, Concatenate
+        Returns:
+            None
+        """
+        if not isinstance(name, list):
+            self.layers.append(
+                LayerConfig(
+                    name=name,
+                    is_keras_class=is_keras_class,
+                    quantize_input=quantize_input,
+                    quantize_weight=quantize_weight,
+                    quantization_index=quantization_index,
+                )
+            )
+        else:
+            # layer names is passed as a list
+            if isinstance(is_keras_class, list):
+                assert len(name) == len(
+                    is_keras_class
+                ), "[E] `is_keras_class` is a list but length is not same as layer `name` list"
+            if isinstance(quantize_input, list):
+                assert len(name) == len(
+                    quantize_input
+                ), "[E] `quantize_input` is a list but length is not same as layer `name` list"
+            if isinstance(quantize_weight, list):
+                assert len(name) == len(
+                    quantize_weight
+                ), "[E] `quantize_weight` is a list but length is not same as layer `name` list"
+            if isinstance(quantization_index, list):
+                assert len(name) == len(
+                    quantization_index
+                ), "[E] `quantization_index` is list but length is not same as layer `name` list"
+            for i, e in enumerate(name):
+                cl_name = e
+                cl_is_keras_class = (
+                    is_keras_class[i]
+                    if isinstance(is_keras_class, list)
+                    else is_keras_class
+                )
+                cl_quantize_input = (
+                    quantize_input[i]
+                    if isinstance(quantize_input, list)
+                    else quantize_input
+                )
+                cl_quantize_weight = (
+                    quantize_weight[i]
+                    if isinstance(quantize_weight, list)
+                    else quantize_weight
+                )
+                cl_quantization_index = (
+                    quantization_index[i]
+                    if isinstance(quantization_index, list)
+                    else quantization_index
+                )
+                self.layers.append(
+                    LayerConfig(
+                        name=cl_name,
+                        is_keras_class=cl_is_keras_class,
+                        quantize_input=cl_quantize_input,
+                        quantize_weight=cl_quantize_weight,
+                        quantization_index=cl_quantization_index,
+                    )
+                )
+
+
+def _skip_layer(layer: tf.keras.layers.Layer) -> bool:
+    """
+    Decide whether quantization wrapping should be skipped for the given layer.
+    The decision is made based on an internal quantize config object parameters.
+
+    Args:
+        layer (tf.keras.layers.Layer): Keras model layer
+    Returns:
+        bool: True if given layer should not be quantized else False
+    """
+    config_object = cfg.get_config_object()
+
+    # Check if any layer with Disabled Quantization by default are in the 'config_object.layer_classes_to_quantize'.
+    #   If so, that layer will be enabled for quantization. Otherwise, skip (return True).
+    layer_class_name = layer.__class__.__name__
+    if layer_class_name in DISABLED_LAYER_QUANTIZATION_DEFAULT:
+        if layer_class_name not in config_object.layer_classes_to_quantize:
+            if layer.name in config_object.get_layer_config():
+                # User can enable a single layer even if the default behavior of a Class is to not quantize.
+                # The decision of whether to quantize this layer or not will be left for later checks, such as when
+                #   quantize_input and quantize_weight = False.
+                pass
+            else:
+                # Default behavior: skip layer
+                return True
+
+    # 1. When quantize_input = False, quantize_weight = False and quantization_index=None, don't even wrap the layer.
+    if layer.name in config_object.get_layer_config():
+        current_layer_config = config_object.get_layer_config()[layer.name]
+        if (
+            current_layer_config["qbool_list"][0] == False  # quantize_input
+            and current_layer_config["qbool_list"][1] == False  # quantize_weight
+            and "qindex_list" not in current_layer_config
+        ):
+            print(
+                "[I] Layer `{layer_name}` is not quantized. There is nothing to quantize since "
+                "quantize_input = False, quantize_weight = False and quantization_index=None".format(
+                    layer_name=layer.name
+                )
+            )
+
+            return True
+
+    # 2. Called when quantization_mode is `partial`
+    if config_object.config_class_id == 2:
+        # A. Skip current `layer class` if current layer class is not in user provided QuantizationSpec class
+        #      object. However, when current layer name is passed by user to quantize, don't skip the layer.
+        if (
+                len(config_object.layer_classes_to_quantize) != 0
+                and layer.__class__.__name__ not in config_object.layer_classes_to_quantize
+        ):
+            if layer.name in config_object.get_layer_config():
+                return False
+            else:
+                print(
+                    "[I] Layer class `{layer_class_name}` is not quantized. Partial quantization is enabled "
+                    "and layer class is not in user provided QuantizationSpec class object".format(
+                        layer_class_name=layer.__class__.__name__
+                    )
+                )
+                return True
+        # B. Skip current layer if `layer.name` is not in user provided QuantizationSpec class object.
+        #      However, if current layer class is passed by user to quantize, don't skip the layer.
+        elif layer.name not in config_object.get_layer_config():
+            if layer.__class__.__name__ in config_object.layer_classes_to_quantize:
+                return False
+            else:
+                print(
+                    "[I] Layer `{layer_name}` is not quantized. Partial quantization is enabled and layer name is not "
+                    "in user provided QuantizationSpec class object".format(
+                        layer_name=layer.name
+                    )
+                )
+                return True
+
+    return False
+
+
+def _quantize_model_layer_clone_function(
+    layer: tf.keras.layers.Layer,
+) -> "BaseQuantizeWrapper":
+    """
+    Wrap or leave given layer based on quantize config object parameters.
+    Args:
+        layer (tf.keras.layers.Layer): Keras model layer
+    Returns:
+        BaseQuantizeWrapper: layer wrapped in BaseQuantizeWrapper class.
+    """
+    layer_wrapper = layer
+    if _skip_layer(layer):
+        # Skip the layers not specified by the user.
+        pass
+    else:
+        child_wrappers_dict = quantize_wrappers.BaseQuantizeWrapper.CHILD_WRAPPERS
+        possible_wrapper_name_for_this_layer = (
+            layer.__class__.__name__ + "QuantizeWrapper"
+        )
+        if possible_wrapper_name_for_this_layer in child_wrappers_dict:
+            wrapper_function = child_wrappers_dict[possible_wrapper_name_for_this_layer]
+            layer_wrapper = wrapper_function(layer)
+    return layer_wrapper
+
+
+def _execute_quantize_model(
+    model: tf.keras.Model, class_id: int, qspec: QuantizationSpec = None
+) -> tf.keras.Model:
+    """
+    clone the model and apply quantization to specific layers based on quantize config object parameters.
+    Args:
+        model (tf.keras.Model): Keras functional or sequential model.
+            * Currently Subclassed models are not supported
+        class_id (int): internal quantization class ID
+        qspec (QuantizationSpec): object of QuantizationSpec class. If few layers or layer classes are to be treated
+            differently, LayerConfig class objects for that layer/layer class are created internally and
+            added to QuantizationSpec class.
+    Returns:
+        tf.keras.Model: Quantized model with QDQ nodes added.
+    """
+    config_id_class_name_map = {
+        0: "FullNetworkQuantization",
+        1: "FullNetworkSpecialQuantization",
+        2: "PartialNetworkQuantization",
+    }
+
+    # 1. Create quantize config object
+    q_config_object = getattr(quantize_config, config_id_class_name_map[class_id])()
+
+    # 2. Update object attributes
+    if qspec:
+        q_config_object.add_quantization_spec_object(qspec, model.layers)
+
+    assert (
+        cfg.is_config_object_created()
+    ), "[E] Have you created the quantization config object before calling `quantize_model`?"
+
+    # 3. Ensure that the original model is kept untouched.
+    #    This step is needed as `clone_model` with our custom `clone_function` wraps layers in a destructive manner.
+    #    TODO: delete later if a better solution is found, most likely inside our custom `clone_function`.
+    cloned_model = tf.keras.models.clone_model(model)
+    cloned_model.set_weights(model.get_weights())
+
+    # 4. Wrap quantizable layers
+    quant_model = tf.keras.models.clone_model(
+        cloned_model, input_tensors=None, clone_function=_quantize_model_layer_clone_function
+    )
+
+    # 5. Clean global space afterwards
+    q_config_object.clean()
+
+    return quant_model
+
+
+def _recognize_config_class_id(
+    quantization_mode: str = "full", qspec: QuantizationSpec = None
+) -> int:
+    """
+    Interpret internal quantize config class based on parameters passed by user to
+    `quantize_model` function.
+    Args:
+        quantization_mode (str): Either 'full' or 'partial' quantization mode
+        qspec (QuantizationSpec): object of QuantizationSpec class. If few layers or layer classes are to be treated
+            differently, LayerConfig class objects for that layer/layer class are created internally and
+            added to QuantizationSpec class.
+    Returns:
+        int: ID for quantization category class used internally.
+    Raises:
+        Exception: if no class can be interpreted for given parameter combination
+    """
+    if quantization_mode == "full" and qspec is None:
+        return 0
+    elif quantization_mode == "full" and qspec is not None:
+        return 1
+    elif quantization_mode == "partial" and qspec is not None:
+        return 2
+    else:
+        raise Exception(
+            "Could not recognize config class ID."
+            " Are parameters passed to `quantize_model` function correct?"
+        )
+
+
+def _validate_config(
+    quantization_mode: str = "full", qspec: QuantizationSpec = None
+) -> None:
+    """
+    Validate if parameters passed to `quantize_model` makes sense.
+    Args:
+        quantization_mode (str): quantization mode can be either 'full' or 'partial'
+        qspec (QuantizationSpec): object of QuantizationSpec class. If few layers or layer classes are to be treated
+            differently, LayerConfig class objects for that layer/layer class are created internally and
+            added to QuantizationSpec class.
+    Returns:
+        None
+    Raises:
+        AssertionError: when configuration is not valid.
+    """
+
+    def _verify_support_for_all_layer_classes(qspec: QuantizationSpec):
+        for layer in qspec.layers:
+            if layer.is_keras_class:
+                # Layer class name is provided.
+                child_wrappers_dict = (
+                    quantize_wrappers.BaseQuantizeWrapper.CHILD_WRAPPERS
+                )
+                possible_wrapper_name_for_this_layer = layer.name + "QuantizeWrapper"
+                assert possible_wrapper_name_for_this_layer in child_wrappers_dict, (
+                    "[E] layer class `{layer_name}` is not supported yet! Either there is no native wrapper or user "
+                    "provided wrapper registration failed.".format(
+                        layer_name=layer.name
+                    )
+                )
+
+    if qspec:
+        _verify_support_for_all_layer_classes(qspec)
+
+    if quantization_mode == "partial":
+        assert (
+            qspec is not None
+        ), "[E] `QuantizationSpec` class object must be passed when `quantization_mode=partial`."
+
+
+def quantize_model(
+    model,
+    quantization_mode: str = "full",
+    quantization_spec: QuantizationSpec = None,
+    custom_qdq_cases: List["CustomQDQInsertionCase"] = None,
+) -> tf.keras.Model:
+    """
+    Insert Q/DQ nodes in Keras model and return a copy. Weights are preserved unlike native keras clone.
+
+    Args:
+        model(tf.keras.Model): Keras Functional or Sequential model.subclassed models are not yet supported.
+        quantization_mode(str): quantization mode can be either 'full' or 'partial'
+        quantization_spec(QuantizationSpec) : object of QuantizationSpec class. If few layers or layer classes are to
+            be treated differently, LayerConfig class objects for that layer/layer class are created internally and
+            added to QuantizationSpec class.
+        custom_qdq_cases(List[CustomQDQInsertionCase]) : `Case` method on every object in this list is called by passing
+            model and user passed quantization_spec as arguments. Each member of this list is an object of a class
+            inherited from CustomQDQInsertionCase class.
+
+    Raises:
+        AssertionError: When passed model is subclassed.
+        AssertionError: When CustomQDQInsertionCase does not return QuantizationSpec object.
+        AssertionError: When quantization mode is `partial` but QuantizationSpec object is not passed.
+        AssertionError: When quantization wrapper is not found for desired layer class.
+        ExceptionError: When internal quantization class ID can't be detected. This happens when passed parameters
+            do not make sense.
+    Returns:
+        tf.keras.Model: Quantized model with QDQ nodes inserted according to NVIDIA quantization recipe.
+    """
+    supported_model_classes = {"Functional", "Sequential"}
+    assert (
+        model.__class__.__name__ in supported_model_classes
+    ), "[E] Currently only `Functional` or `Sequential` model quantization is supported."
+
+    # Update quantization_spec object based on output of special QDQ cases.
+    custom_quantization_spec = QuantizationSpec()
+    if custom_qdq_cases:
+        for custom_qdq_case in custom_qdq_cases:
+            qspec_case_object = custom_qdq_case.case(model, quantization_spec)
+            if qspec_case_object:
+                assert isinstance(
+                    qspec_case_object, QuantizationSpec
+                ), "[E] {} \
+                does not return an object of QuantizationSpec.".format(
+                    qspec_case_object.__class__.__name__
+                )
+                custom_quantization_spec.layers.extend(qspec_case_object.layers)
+
+    # if user has passed quantization_spec then extend it with custom_quantization_spec
+    # else use just custom_quantization_spec
+    if quantization_spec:
+        quantization_spec.layers.extend(custom_quantization_spec.layers)
+    else:
+        if len(custom_quantization_spec.layers) != 0:
+            quantization_spec = custom_quantization_spec
+
+    # Check if config is valid and quantize model
+    _validate_config(quantization_mode, quantization_spec)
+    cid = _recognize_config_class_id(quantization_mode, quantization_spec)
+    return _execute_quantize_model(model, cid, quantization_spec)

tools/tensorflow-quantization/tensorflow_quantization/quantize_config.py

@@ -0,0 +1,273 @@
+#
+# 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 implements classes to configure three supported quantization modes:
+1. Full: quantize all layers with standard protocol based NVIDIA quantization scheme.
+2. Full special: quantize few layers in a specific way and remaining with standard protocol based on NVIDIA
+                 quantization scheme.
+3. Partial: quantize ONLY few layers.
+
+Each quantization mode can quantize all supported Keras layer classes or only subset of it.
+"""
+
+from abc import ABC
+import tensorflow_quantization.global_config as global_config
+import warnings
+from typing import List, Dict
+
+
+class BaseConfig(ABC):
+    """
+    Base class from which four quantize config classes are derived.
+    Default quantization recipe is Nvidia's recommendation.
+    """
+
+    def __new__(cls):
+        instance = super().__new__(cls)
+        # Add instance to global list
+        global_config.add_config_object(instance)
+        return instance
+
+    def __init__(self) -> None:
+        self.quantization_mode: str = "full"
+        self.layerwise_config: dict = {}  # holds special layers information.
+        self.layer_classes_to_quantize: set = set()
+        self.config_class_id: int = 0
+
+    def __str__(self) -> str:
+        return (
+            " quantization_mode: {quant_mode} \n "
+            "layerwise_config: {layerwise_config} \n "
+            "specific_layer_class: {specific_layer_class} \n "
+            "config_class_id: {config_class_id} \n".format(
+                quant_mode=self.quantization_mode,
+                layerwise_config=self.layerwise_config,
+                specific_layer_class=self.specific_layer_class,
+                config_class_id=self.config_class_id,
+            )
+        )
+
+    @staticmethod
+    def _validate_layer_names(
+        user_passed_layer_names: List, model_layers: List
+    ) -> None:
+        """
+        Check whether user passed layer names exists in  Keras model being quantized.
+        Args:
+            user_passed_layer_names (List): Layer names passed by user to treat specially.
+            model_layers (List): Keras model layers passed as a list.
+        Returns:
+            None
+        Raise:
+            Warning : when specific layer name is not found. Such layers are simply ignored.
+        """
+        model_layer_name_set = set()
+        for l in model_layers:
+            model_layer_name_set.add(l.name)
+
+        for ul in user_passed_layer_names:
+            if ul not in model_layer_name_set:
+                warnings.warn(
+                    "layer name {} is passed by user but could not find layer with this name in model.".format(
+                        ul
+                    )
+                )
+
+    def add_quantization_spec_object(
+        self, qspec: "QuantizationSpec", original_model_layers: List
+    ) -> None:
+        """
+        This method parses object of QuantizationSpec class and fill in `layerwise_config` dictionary
+        holding information about layers that need to be treated specially.
+        Specific layer classes that need to be treated specially are also here.
+        Args:
+            qspec (QuantizationSpec): object of QuantizationSpec class. If few layers or layer classes are to be treated
+                differently, LayerConfig class objects for that layer/layer class are created internally and
+                added to QuantizationSpec class.
+            original_model_layers (List): Keras model layers passed as a list.
+        Returns:
+            None
+        """
+        for layer in qspec.layers:
+            if layer.is_keras_class:
+                self.add_special_layer_class(layer.name)
+            else:
+                layer_config_dict = {"qbool_list": [False, False]}
+                layer_config_dict["qbool_list"][0] = layer.quantize_input
+                layer_config_dict["qbool_list"][1] = layer.quantize_weight
+                if layer.quantization_index:
+                    layer_config_dict["qindex_list"] = layer.quantization_index
+                self.add_special_layer(
+                    layer_name=layer.name, config_dict=layer_config_dict
+                )
+        # Validate whether added layers exist in the model
+        self._validate_layer_names(
+            list(self.layerwise_config.keys()), original_model_layers
+        )
+
+    def add_special_layer(self, layer_name: str, config_dict: Dict) -> None:
+        """
+        Add layer specific quantization information to quantize config object.
+        Args:
+            layer_name (str): layer name
+            config_dict (Dict): Layer specific quantization parameter dictionary in the
+                following format.
+                There are only two accepted keys `qbool_list` and `qindex_list`.
+                `qbool_list` is list of length two where each value is
+                [<True/False quantize inputs>, <True/False quantize weights>]
+                e.g.
+                To quantize inputs and weights, `qbool_list`=[True, True]
+
+                `qindex_list` is a list of specific indices to quaintize for layers such as Add, Concatenate
+                where more than two inputs are present.
+
+                Based on above information,
+                1. config_dict for weighted layer with name `dense_2`, to quantize inputs and weights will be
+                {'qbool_list':[True, True]} with laye_name=`dense_2`
+                2. config_dict for non weighted layer with name `add_3` to quantize input at index 1 will be
+                {'qbool_list':[True, False], 'qindex_list':[1]} with layer_name=`add_3`
+        Returns:
+            None
+        Raises:
+            Exception: When invalid keys are detected.
+        """
+        self.layerwise_config[layer_name] = config_dict
+
+    def remove_layer(self, layer_name: str) -> None:
+        """
+        Remove specific layer based on name from quantize config object.
+        Args:
+            layer_name (str): layer name
+        Returns:
+            None
+        """
+        if layer_name in self.layerwise_config:
+            del self.layerwise_config[layer_name]
+
+    def remove_layers(self, layers_name: List) -> None:
+        """
+        Bulk remove specific layers based on names from quantize config object.
+        Args:
+            layers_name (List): layers names, list of strings
+        Returns:
+            None
+        """
+        for layer_name in layers_name:
+            self.remove_layer(layer_name=layer_name)
+
+    def get_layer_config(self) -> Dict:
+        """
+        Return dictionary with information about layers to quantize for quantize
+        config object.
+        Args:
+            None
+        Returns:
+            Dict: a dictionary with layerwise configuration parameters.
+        """
+        return self.layerwise_config
+
+    def is_empty(self) -> bool:
+        """
+        Return True if no layer specific quantization information is available in quantize
+        config object.
+        Args:
+            None
+        Returns:
+            bool: True if no special layers are passed else return False
+        """
+        return not self.layerwise_config
+
+    def clear_layer_config(self) -> None:
+        """
+        Clear layer config information from quanize config object
+        Args:
+            None
+        Returns:
+            None
+        """
+        self.layerwise_config.clear()
+
+    def add_special_layer_class(self, layer_class_name: str) -> None:
+        """
+        Add class name to quantize config object so that only layers with specific class are quantized.
+        Args:
+            layer_class_name : String that represents keras class
+        Returns:
+            None
+        """
+        self.layer_classes_to_quantize.add(layer_class_name)
+
+    def clean(self):
+        """
+        Clean quantize config object from global space. Calling this is important to use `quantize_model` multiple times
+        within a single module.
+        Args:
+            None
+        Returns:
+            None
+        """
+        global_config.remove_config_object()
+
+
+class FullNetworkQuantization(BaseConfig):
+    """
+    Quantize all layers based on NV scheme.
+
+    Nvidia recommended recipe for quantization is using Q/DQ only wth inputs/weights.
+    Q/DQ output support is just to compare engine performance/accuracy when other quantization
+    scheme is used.
+
+    NV: Add Q/DQ at input and weights
+    TF: Add Q/DQ at output and weights
+
+    This is config class with index `0` which is default.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.config_class_id = 0
+
+
+class FullNetworkSpecialQuantization(BaseConfig):
+    """
+    Quantize few layers in specific way and remaining network in standard way based on NV scheme.
+    Layers are selected based on 'names' which can be via 'model.summary()' for functional
+    and sequential models.
+    Subclassed model layer information can be found using `KerasModelTraveller` class from utils.
+
+    This is config class with index 1.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.config_class_id = 1
+
+
+class PartialNetworkQuantization(BaseConfig):
+    """
+    Quantize only specific layers and not the entire network.
+    Layers are selected based on name.
+
+    This is config class with index 2.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.quantization_mode = "partial"
+        self.config_class_id = 2

tools/tensorflow-quantization/tensorflow_quantization/quantize_wrapper_base.py

@@ -0,0 +1,236 @@
+#
+# 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 2019 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.
+
+import tensorflow as tf
+import tensorflow_quantization.quantizers as quantizers
+import tensorflow_quantization.global_config as cfg
+from abc import abstractmethod
+
+deserialize_keras_object = tf.keras.utils.deserialize_keras_object
+serialize_keras_object = tf.keras.utils.serialize_keras_object
+
+NO_WEIGHT_LAYERS = {
+    "Concatenate",
+    "Add",
+    "AveragePooling2D",
+    "GlobalAveragePooling2D",
+    "MaxPooling2D",
+    "BatchNormalization",
+}
+
+
+class BaseQuantizeWrapper(tf.keras.layers.Wrapper):
+    """Base wrapper class which all layer wrappers inherit"""
+
+    CHILD_WRAPPERS = {}
+
+    def __init_subclass__(cls, **kwargs) -> None:
+        super().__init_subclass__(**kwargs)
+        cls.CHILD_WRAPPERS[cls.__name__] = cls
+
+    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
+        """Create a quantize emulate wrapper for a keras layer.
+        This wrapper provides options to quantize inputs and weights of the layer.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        if layer is None:
+            raise ValueError("`layer` cannot be None.")
+
+        # Check against keras.Model since it is an instance of keras.layers.Layer.
+        if not isinstance(layer, tf.keras.layers.Layer) or isinstance(
+            layer, tf.keras.Model
+        ):
+            raise ValueError(
+                "`layer` can only be a `tf.keras.layers.Layer` instance. "
+                "You passed an instance of type: {input}.".format(
+                    input=layer.__class__.__name__
+                )
+            )
+        if "name" not in kwargs:
+            kwargs["name"] = self._make_layer_name(layer)
+        super(BaseQuantizeWrapper, self).__init__(layer, **kwargs)
+
+        # get quantize config object that holds all the information about how quantization should be performed.
+        quantize_config_object = cfg.get_config_object()
+
+        # set all initial quantization parameters to False/None
+        self.quantize_inputs = False
+        self.quantize_weights = False
+        self.quantize_specific_input_indices = None
+
+        layer_class_name_t = layer.__class__.__name__  # Layer class name
+        layer_name_t = layer.name  # Actual layer name
+
+        def _configure_singular_quantize():
+            self.quantize_inputs = True
+            if layer_class_name_t in NO_WEIGHT_LAYERS:
+                self.quantize_weights = False
+            else:
+                self.quantize_weights = True
+
+        def _configure_special_quantize(
+            quantize_bool_list: list, layer_name_t: str, index_list_if_any: list = None
+        ):
+            assert (len(quantize_bool_list)) == 2, (
+                "Three boolean values (representing whether to quantize [inputs, weights]) must be provided in "
+                "quantize_config for layer: {layer_name_t}. If quantization does not apply for specific part, "
+                "pass None. e.g. For layer ( e.g. Concatenate, Add) with no weights, `qbool_list` to quantize "
+                "input can be [True, False]".format(layer_name_t=layer_name_t)
+            )
+
+            self.quantize_inputs = quantize_bool_list[0]
+            if layer_class_name_t in NO_WEIGHT_LAYERS:
+                self.quantize_weights = False
+            else:
+                self.quantize_weights = quantize_bool_list[1]
+
+            if index_list_if_any:
+                self.quantize_specific_input_indices = index_list_if_any
+
+        if quantize_config_object.config_class_id == 0:
+            # This is straight forward full network quantization
+            _configure_singular_quantize()
+        else:
+            # Config class id 1 or 2.
+            # User has provided layer (name) specific quantization information
+            quantize_config_dict = quantize_config_object.get_layer_config()
+            if layer_name_t in quantize_config_dict:
+                # This layer needs to be quantized in specific way
+                if "qindex_list" in quantize_config_dict[layer_name_t]:
+                    _configure_special_quantize(
+                        quantize_config_dict[layer_name_t]["qbool_list"],
+                        layer_name_t,
+                        quantize_config_dict[layer_name_t]["qindex_list"],
+                    )
+                else:
+                    _configure_special_quantize(
+                        quantize_config_dict[layer_name_t]["qbool_list"], layer_name_t
+                    )
+            else:
+                _configure_singular_quantize()
+
+        self._track_trackable(layer, name="layer")
+
+    @staticmethod
+    def _make_layer_name(layer):
+        return "{}_{}".format("quant", layer.name)
+
+    @staticmethod
+    def _weight_name(name):
+        """Extracts the weight name from the full TensorFlow variable name.
+        For example, returns 'kernel' for 'dense_2/kernel:0'.
+        Args:
+          name: TensorFlow variable name.
+        Returns:
+          Extracted weight name.
+        """
+        return name.split(":")[0].split("/")[-1]
+
+    def build(self, input_shape):
+        super(BaseQuantizeWrapper, self).build(input_shape)
+
+        self.optimizer_step = self.add_weight(
+            "optimizer_step",
+            initializer=tf.keras.initializers.Constant(-1),
+            dtype=tf.dtypes.int32,
+            trainable=False,
+        )
+
+    def compute_output_shape(self, input_shape):
+        return self.layer.compute_output_shape(self.layer.input_shape)
+
+    def _last_value_quantizer(
+        self, x, training, quantizer_vars, per_channel=False, channel_axis=-1
+    ):
+        """Use currying to return True/False specialized fns to the cond."""
+        from tensorflow_quantization import G_NUM_BITS, G_SYMMETRIC, G_NARROW_RANGE
+
+        return quantizers.LastValueQuantize(
+            x,
+            quantizer_vars["min_var"],
+            quantizer_vars["max_var"],
+            per_channel=per_channel,
+            channel_axis=channel_axis,
+            is_training=training,
+            num_bits=G_NUM_BITS,
+            narrow_range=G_NARROW_RANGE,
+            symmetric=G_SYMMETRIC,
+        )
+
+    @abstractmethod
+    def call(self, inputs, training=None):
+        raise NotImplementedError
+
+    def get_config(self):
+        base_config = super(BaseQuantizeWrapper, self).get_config()
+        config = {"quantize_config": None}
+        return dict(list(base_config.items()) + list(config.items()))
+
+    @classmethod
+    def from_config(cls, config):
+        config = config.copy()
+
+        # BaseQuantizeWrapper may be constructed with any QuantizeConfig and the
+        # wrapper itself cannot know all the possible config classes.
+        # The deserialization code should ensure the QuantizeConfig is in keras
+        # serialization scope.
+        quantize_config = deserialize_keras_object(
+            config.pop("quantize_config"), module_objects=globals(), custom_objects=None
+        )
+
+        layer = tf.keras.layers.deserialize(config.pop("layer"))
+
+        return cls(layer=layer, quantize_config=quantize_config, **config)
+
+    @property
+    def trainable(self):
+        return self.layer.trainable
+
+    @trainable.setter
+    def trainable(self, value):
+        self.layer.trainable = value
+
+    @property
+    def trainable_weights(self):
+        return self.layer.trainable_weights + self._trainable_weights
+
+    @property
+    def non_trainable_weights(self):
+        return self.layer.non_trainable_weights + self._non_trainable_weights
+
+    @property
+    def updates(self):
+        return self.layer.updates + self._updates
+
+    @property
+    def losses(self):
+        return self.layer.losses + self._losses

tools/tensorflow-quantization/tensorflow_quantization/quantize_wrappers.py

@@ -0,0 +1,556 @@
+#
+# 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 tensorflow.python.util import tf_inspect
+from tensorflow_quantization.quantize_wrapper_base import BaseQuantizeWrapper
+import warnings
+
+"""
+Naming convention for keras `layer` quantize wrapper is
+<layer.__class__.__name__>QuantizeWrapper
+"""
+
+DISABLED_LAYER_QUANTIZATION_DEFAULT = [
+    "MaxPooling2D",
+    "BatchNormalization",
+    "Add",
+    "Multiply",
+    "Concatenate"
+]
+
+
+# ##############################################
+# ############# Weighted Layers ################
+# ##############################################
+
+
+class WeightedBaseQuantizeWrapper(BaseQuantizeWrapper):
+    """
+    BaseQuantizeWrapper for weighted layers: Conv2D, DepthwiseConv2D, and Dense layer.
+        These layers share a lot of the same code except for a few modifications. Conv2D and Dense share the same code.
+        Layers that inherit this class support weight and input QDQ nodes.
+
+    TRT Rule:
+        One Q/DQ pair should be attached to the input activation, and another Q/DQ pair should be attached to weights.
+        Weights tensor is per-channel quantized:
+            For the Q/DQ attached to weight tensor, set axis=0 and axis=1 for Conv and ConvTransposed respectively.
+        Input tensor is per-tensor quantized.
+    """
+
+    def __init__(
+        self, layer: tf.keras.layers.Layer, kernel_type: str = "kernel", **kwargs
+    ):
+        """
+        Creates a wrapper to emulate quantization for a keras layer.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          kernel_type (str): Options=['kernel' for Conv2D/Dense, 'depthwise_kernel' for DepthwiseConv2D]
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        self.kernel_type = kernel_type
+        self.channel_axis = kwargs.get("axis", -1)
+        super().__init__(layer, **kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+        self._weight_vars = []
+        self.input_vars = {}
+        self.output_vars = {}
+        self.channel_axis = -1
+        if self.kernel_type == "depthwise_kernel":
+            self.channel_axis = 2
+        # quantize weights only applicable for weighted ops.
+        # By default weights is per channel quantization
+        if self.quantize_weights:
+            # get kernel weights dims.
+            kernel_weights = getattr(self.layer, self.kernel_type)
+            min_weight = self.layer.add_weight(
+                kernel_weights.name.split(":")[0] + "_min",
+                shape=(kernel_weights.shape[self.channel_axis]),
+                initializer=tf.keras.initializers.Constant(-6.0),
+                trainable=False,
+            )
+            max_weight = self.layer.add_weight(
+                kernel_weights.name.split(":")[0] + "_max",
+                shape=(kernel_weights.shape[self.channel_axis]),
+                initializer=tf.keras.initializers.Constant(6.0),
+                trainable=False,
+            )
+            quantizer_vars = {"min_var": min_weight, "max_var": max_weight}
+            self._weight_vars.append((kernel_weights, quantizer_vars))
+            # Needed to ensure unquantized weights get trained as part of the wrapper.
+            self._trainable_weights.append(kernel_weights)
+
+        # By default input is per tensor quantization
+        if self.quantize_inputs:
+            input_min_weight = self.layer.add_weight(
+                self.layer.name + "_ip_min",
+                shape=None,
+                initializer=tf.keras.initializers.Constant(-6.0),
+                trainable=False,
+            )
+            input_max_weight = self.layer.add_weight(
+                self.layer.name + "_ip_max",
+                shape=None,
+                initializer=tf.keras.initializers.Constant(6.0),
+                trainable=False,
+            )
+            self.input_vars["min_var"] = input_min_weight
+            self.input_vars["max_var"] = input_max_weight
+
+    def call(self, inputs, training=None):
+        if training is None:
+            training = tf.keras.backend.learning_phase()
+
+        # Quantize all weights, and replace them in the underlying layer.
+        if self.quantize_weights:
+            quantized_weights = []
+            quantized_weight = self._last_value_quantizer(
+                self._weight_vars[0][0],
+                training,
+                self._weight_vars[0][1],
+                per_channel=True,
+                channel_axis=self.channel_axis,
+            )
+            quantized_weights.append(quantized_weight)
+            # Replace the original weights with QDQ weights
+            setattr(self.layer, self.kernel_type, quantized_weights[0])
+
+        # Quantize inputs to the conv layer
+        if self.quantize_inputs:
+            quantized_inputs = self._last_value_quantizer(
+                inputs, training, self.input_vars, per_channel=False
+            )
+        else:
+            quantized_inputs = inputs
+
+        args = tf_inspect.getfullargspec(self.layer.call).args
+        if "training" in args:
+            outputs = self.layer.call(quantized_inputs, training=training)
+        else:
+            outputs = self.layer.call(quantized_inputs)
+
+        return outputs
+
+
+class Conv2DQuantizeWrapper(WeightedBaseQuantizeWrapper):
+    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
+        """
+        Creates a wrapper to emulate quantization for the Conv2D keras layer.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        self.kernel_type = "kernel"
+        super().__init__(layer, kernel_type=self.kernel_type, **kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+    def call(self, inputs, training=None):
+        return super().call(inputs, training=training)
+
+
+class DenseQuantizeWrapper(WeightedBaseQuantizeWrapper):
+    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
+        """
+        Creates a wrapper to emulate quantization for the Dense keras layer.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        self.kernel_type = "kernel"
+        super().__init__(layer, kernel_type=self.kernel_type, **kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+    def call(self, inputs, training=None):
+        return super().call(inputs, training=training)
+
+
+class DepthwiseConv2DQuantizeWrapper(WeightedBaseQuantizeWrapper):
+    """Requires TF >= 2.8.0"""
+
+    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
+        """
+        Creates a wrapper to emulate quantization for the DepthwiseConv2D keras layer.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        self.kernel_type = "depthwise_kernel"
+        super().__init__(layer, kernel_type=self.kernel_type, **kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+    def call(self, inputs, training=None):
+        return super().call(inputs, training=training)
+
+
+# ##############################################
+# ########### Non-Weighted Layers ##############
+# ######### with Single Input/Output ###########
+# ##############################################
+
+
+class NonWeightedBaseQuantizeWrapper(BaseQuantizeWrapper):
+    """
+    BaseQuantizeWrapper for non-weighted layers with Single Input/Output: AveragePooling2D, GlobalAveragePooling,
+        MaxPooling2D and BatchNormalization.
+
+    Supports 1 input and 1 output QDQ. Similar to Concat, except that Concat supports multiple inputs.
+    NonWeightedBaseQuantizeWrapper can use WeightedBaseQuantizeWrapper by giving quantize_weigths=False.
+    """
+
+    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
+        """
+        Creates a wrapper to emulate quantization for non-weighted keras layers.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        super().__init__(layer, **kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+        self.input_vars = {}
+
+        # By default input is per tensor quantization
+        if self.quantize_inputs:
+            input_min_weight = self.layer.add_weight(
+                self.layer.name + "_ip_min",
+                shape=None,
+                initializer=tf.keras.initializers.Constant(-6.0),
+                trainable=False,
+            )
+            input_max_weight = self.layer.add_weight(
+                self.layer.name + "_ip_max",
+                shape=None,
+                initializer=tf.keras.initializers.Constant(6.0),
+                trainable=False,
+            )
+            self.input_vars["min_var"] = input_min_weight
+            self.input_vars["max_var"] = input_max_weight
+
+    def call(self, inputs, training=None):
+        if training is None:
+            training = tf.keras.backend.learning_phase()
+
+        # Quantize inputs to the conv layer
+        if self.quantize_inputs:
+            quantized_inputs = self._last_value_quantizer(
+                inputs, training, self.input_vars, per_channel=False
+            )
+        else:
+            quantized_inputs = inputs
+
+        args = tf_inspect.getfullargspec(self.layer.call).args
+        if "training" in args:
+            outputs = self.layer.call(quantized_inputs, training=training)
+        else:
+            outputs = self.layer.call(quantized_inputs)
+
+        return outputs
+
+
+class AveragePooling2DQuantizeWrapper(NonWeightedBaseQuantizeWrapper):
+    """
+    TRT Rule:
+        Add Q/DQ to its input if the ops follows is quantized.
+        Quantize average pooling will introduce small variance compared to float because of the rounding change.
+        TensorRT doesn’t have Int8 in and fp32 out average pool support.
+        If the op follows average pooling is not quantized, it is users choice between running average pooling
+        in int8 then convert to fp32 for the following op and run average pooling in fp32.
+
+    Currently, we're adding QDQ to all AveragePooling2D layers.
+    """
+
+    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
+        """
+        Creates a wrapper to emulate quantization for the AveragePooling2D keras layer.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        super().__init__(layer, **kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+    def call(self, inputs, training=None):
+        return super().call(inputs, training=training)
+
+
+class GlobalAveragePooling2DQuantizeWrapper(NonWeightedBaseQuantizeWrapper):
+    """
+    TRT Rule:
+        No explicit rule from the TRT team. Following the same as AveragePooling2D.
+
+    Residual block v2: Add to MaxPool (branch1) and BN (branch2).
+
+    Supports 1 input and 1 output QDQ. Same as AveragePooling2DQuantizeWrapper.
+    """
+
+    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
+        """
+        Creates a wrapper to emulate quantization for the GlobalAveragePooling2D keras layer.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        super().__init__(layer, **kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+    def call(self, inputs, training=None):
+        return super().call(inputs, training=training)
+
+
+class MaxPooling2DQuantizeWrapper(NonWeightedBaseQuantizeWrapper):
+    """
+    TRT Rule:
+        Max pooling is precision-neutral. But unlike ReLU, input and output of max pooling will have different
+            histograms which will lead to different calibration results.
+        The recommendation is to let TensorRT optimize precision neutral ops.
+        There are cases where adding Q/DQ before maxpool can enable additional optimization.
+
+    Residual block v2: Add to MaxPool (branch1) and BN (branch2).
+
+    Supports 1 input and 1 output QDQ. Same as AveragePooling2DQuantizeWrapper.
+    """
+
+    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
+        """
+        Creates a wrapper to emulate quantization for the MaxPooling2D keras layer.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        super().__init__(layer, **kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+    def call(self, inputs, training=None):
+        return super().call(inputs, training=training)
+
+
+class BatchNormalizationQuantizeWrapper(NonWeightedBaseQuantizeWrapper):
+    """
+    TRT Rule:
+        Keep batch normalization untouched, don't add Q/DQ to its input and not necessary to fold it before exporting
+            graph. TensorRT supports Batch normalization folding. It can take a graph with batch normalization, fold it
+            into previous convolution and create a new graph.
+        If batch normalization is folded before exporting the graph, TensorRT can still import and execute the graph as
+            it becomes regular convolutions.
+
+    Exception for Residual block v2:
+        BN-ReLU-Conv2D -> need to add Q/DQ before BN in order to run in INT8.
+        In order to do that, we add a check in 'quantize_model()' to check if BN's parent is a Conv layer. If it is, set
+            quantize_inputs to False. The reason why we don't add this check here is to allow the user to add QDQ nodes
+            before BN if they so wish.
+
+    Supports 1 input and 1 output QDQ. Same as AveragePooling2DQuantizeWrapper.
+    """
+
+    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
+        """
+        Creates a wrapper to emulate quantization for the BatchNormalization keras layer.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        super().__init__(layer, **kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+    def call(self, inputs, training=None):
+        return super().call(inputs, training=training)
+
+
+# ##############################################
+# ########### Non-Weighted Layers ##############
+# #### with Multiple Inputs, Single Output #####
+# ##############################################
+
+
+class NonWeightedBaseQuantizeWrapperForMultipleInputs(BaseQuantizeWrapper):
+    """
+    BaseQuantizeWrapper for non-weighted layers with Multiple Inputs: Concat, Add, and Multiply.
+    Supports multiple inputs and 1 output QDQ. Similar to AveragePooling2D, except pooling supports only a single input.
+
+    TRT Rule:
+        Add Q/DQ to all inputs of the layer.
+    """
+
+    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
+        """
+        Creates a wrapper to emulate quantization for the keras layer.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        super().__init__(layer, **kwargs)
+
+    def _should_quantization_this_index(self, i):
+        if not self.quantize_specific_input_indices:
+            return True
+        else:
+            # This is a small list so iterating makes sense
+            for e in self.quantize_specific_input_indices:
+                if e == i:
+                    return True
+                elif e >= self.num_inputs:
+                    warnings.warn(
+                        "{layer_name} has {num_inputs} inputs but quantization index {e} is passed.".format(
+                            layer_name=self.layer.name, num_inputs=self.num_inputs, e=e
+                        )
+                    )
+            return False
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+        self.input_vars = []  # list of dictionaries
+        self.num_inputs = len(input_shape)
+
+        # By default input is per tensor quantization
+        if self.quantize_inputs:
+            # for concat input is list of Tensors
+            layer_name_key_idx = 0
+            for i in range(self.num_inputs):
+                if self._should_quantization_this_index(i):
+                    input_min_weight = self.layer.add_weight(
+                        self.layer.name + "_ip{}_min".format(layer_name_key_idx),
+                        shape=None,
+                        initializer=tf.keras.initializers.Constant(-6.0),
+                        trainable=False,
+                    )
+                    input_max_weight = self.layer.add_weight(
+                        self.layer.name + "_ip{}_max".format(layer_name_key_idx),
+                        shape=None,
+                        initializer=tf.keras.initializers.Constant(6.0),
+                        trainable=False,
+                    )
+                    self.input_vars.append(
+                        {"min_var": input_min_weight, "max_var": input_max_weight}
+                    )
+                    layer_name_key_idx += 1
+
+    def call(self, inputs, training=None):
+        if training is None:
+            training = tf.keras.backend.learning_phase()
+
+        # Quantize inputs to the conv layer
+        quantized_inputs = inputs[:]
+        if self.quantize_inputs:
+            input_vars_idx = 0
+            for i in range(len(inputs)):
+                if self._should_quantization_this_index(i):
+                    quantized_inputs[i] = self._last_value_quantizer(
+                        inputs[i],
+                        training,
+                        self.input_vars[input_vars_idx],
+                        per_channel=False,
+                    )
+                    input_vars_idx += 1
+
+        args = tf_inspect.getfullargspec(self.layer.call).args
+        if "training" in args:
+            outputs = self.layer.call(quantized_inputs, training=training)
+        else:
+            outputs = self.layer.call(quantized_inputs)
+
+        return outputs
+
+
+class MultiplyQuantizeWrapper(NonWeightedBaseQuantizeWrapperForMultipleInputs):
+    """
+    TRT Rule:
+        Add Q/DQ to all inputs of Multiply layer in SE block.
+    """
+
+    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
+        """
+        Creates a wrapper to emulate quantization for the Multiply keras layer.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        super().__init__(layer, **kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+    def call(self, inputs, training=None):
+        return super().call(inputs, training=training)
+
+
+class ConcatenateQuantizeWrapper(NonWeightedBaseQuantizeWrapperForMultipleInputs):
+    """
+    TRT Rule:
+        Add Q/DQ to all inputs.
+        Alternative: If there is Q/DQ attached to the input of the op after concat, don't add Q/DQ to input to concat,
+            let TensorRT pull from the next op.
+    """
+
+    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
+        """
+        Creates a wrapper to emulate quantization for the Concatenate keras layer.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        super().__init__(layer, **kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+    def call(self, inputs, training=None):
+        return super().call(inputs, training=training)
+
+
+class AddQuantizeWrapper(NonWeightedBaseQuantizeWrapperForMultipleInputs):
+    """
+    TRT Rule:
+        If the add is NOT bias. Attach Q/DQ to all of its input.
+        Exception: add in residual block. To trigger fusion, Attach Q/DQ to the residual being added to output of
+            convolution.
+    """
+
+    def __init__(self, layer: tf.keras.layers.Layer, **kwargs):
+        """
+        Creates a wrapper to emulate quantization for the Add keras layer.
+        Args:
+          layer (tf.keras.layers.Layer): The keras layer to be quantized.
+          **kwargs: Additional keyword arguments to be passed to the keras layer.
+        """
+        super().__init__(layer, **kwargs)
+
+    def build(self, input_shape):
+        super().build(input_shape)
+
+    def call(self, inputs, training=None):
+        return super().call(inputs, training=training)

tools/tensorflow-quantization/tensorflow_quantization/quantizers.py

@@ -0,0 +1,185 @@
+#
+# 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 2019 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.
+
+"""
+This module is borrowed from TFMOT repository and updated.
+It implements QDQ insertion based on "Last Value Quantization".
+"""
+
+import tensorflow as tf
+
+
+def LastValueQuantize(
+    inputs,
+    min_var,
+    max_var,
+    per_channel=False,
+    channel_axis=-1,
+    name_prefix="LastValueQuant",
+    is_training=True,
+    num_bits=8,
+    narrow_range=False,
+    symmetric=False,
+):
+    """Adds a layer that collects quantization ranges as last input ranges.
+    LastValueQuantize creates variables called 'min' and 'max', representing the
+    interval used for quantization and clamping.
+    Args:
+      inputs: a tensor containing values to be quantized.
+      per_channel: (Optional) a boolean specifying whether to use different
+        quantization ranges per output channel.
+      init_min: a float scalar, the initial value for variable min.
+      init_max: a float scalar, the initial value for variable max.
+      name_prefix: name_prefix for created nodes.
+      is_training: Whether the op is applied to a training or eval graph.
+      num_bits: Number of bits to use for quantization, must be between 2 and 8.
+      narrow_range: Whether to use the narrow quantization range
+        [1; 2^num_bits - 1] or wide range [0; 2^num_bits - 1].
+      symmetric: If true, use symmetric quantization limits instead of training
+        the minimum and maximum of each quantization range separately.
+    Returns:
+      a tensor containing quantized values.
+    """
+    with tf.name_scope(name_prefix):
+        input_shape = inputs.get_shape()
+        input_dim = len(input_shape)
+        if channel_axis == -1:
+            channel_axis += input_dim
+
+        if not is_training:
+            return _QuantizeAndDequantize(
+                inputs,
+                min_var,
+                max_var,
+                per_channel=per_channel,
+                channel_axis=channel_axis,
+                num_bits=num_bits,
+                narrow_range=narrow_range,
+            )
+
+        if per_channel:
+            if input_dim == 2:
+                reduce_dims = [0]
+            elif input_dim == 4:
+                reduce_dims = [i for i in range(input_dim) if i != channel_axis]
+
+        if per_channel:
+            if input_dim >= 2:
+                batch_min = tf.math.reduce_min(
+                    inputs, axis=reduce_dims, name="BatchMin"
+                )
+            else:
+                batch_min = inputs
+        else:
+            batch_min = tf.math.reduce_min(inputs, name="BatchMin")
+
+        if per_channel:
+            if input_dim >= 2:
+                batch_max = tf.math.reduce_max(
+                    inputs, axis=reduce_dims, name="BatchMax"
+                )
+            else:
+                batch_max = inputs
+        else:
+            batch_max = tf.math.reduce_max(inputs, name="BatchMax")
+
+        if symmetric:
+            if narrow_range:
+                min_max_ratio = -1
+            else:
+                # In two's complement notation, the negative range is slightly larger
+                # than the positive range.
+                min_max_ratio = -((1 << num_bits) - 2) / (1 << num_bits)
+
+            # TFLite requires that 0.0 if always in the [min; max] range. Because
+            # batch_min <= batch_max, it follows that range_min <= 0 <= range_max.
+            range_min = tf.math.minimum(batch_min, batch_max / min_max_ratio)
+            range_max = tf.math.maximum(batch_max, batch_min * min_max_ratio)
+        else:
+            # TFLite requires that 0.0 if always in the [min; max] range.
+            range_min = tf.math.minimum(batch_min, 0.0)
+            range_max = tf.math.maximum(batch_max, 0.0)
+
+        assign_min = min_var.assign(range_min, name="AssignMinLast")
+        assign_max = max_var.assign(range_max, name="AssignMaxLast")
+
+        return _QuantizeAndDequantize(
+            inputs,
+            assign_min,
+            assign_max,
+            per_channel=per_channel,
+            channel_axis=channel_axis,
+            num_bits=num_bits,
+            narrow_range=narrow_range,
+        )
+
+
+def _QuantizeAndDequantize(
+    inputs, min_var, max_var, per_channel, channel_axis, num_bits, narrow_range
+):
+    """Adds a fake quantization operation.
+    Depending on value of per_channel, this operation may do global quantization
+    or per channel quantization.  min_var and max_var should have corresponding
+    shapes: [1] when per_channel == False and [d] when per_channel == True.
+    Args:
+      inputs: a tensor containing values to be quantized.
+      min_var: a variable containing quantization range lower end(s).
+      max_var: a variable containing quantization range upper end(s).
+      per_channel: a boolean specifying whether to use per-channel quantization.
+      num_bits: Number of bits to use for quantization, must be between 2 and 8.
+      narrow_range: Whether to use the narrow quantization range
+        [1; 2^num_bits - 1] or wide range [0; 2^num_bits - 1].
+    Returns:
+      a tensor containing quantized values.
+    """
+
+    if per_channel:
+
+        return tf.quantization.quantize_and_dequantize_v2(
+            inputs,
+            min_var,
+            max_var,
+            num_bits=num_bits,
+            narrow_range=narrow_range,
+            axis=channel_axis,
+            range_given=True,
+        )
+    else:
+        assert min_var.get_shape() == []  # pylint: disable=g-explicit-bool-comparison
+        assert max_var.get_shape() == []  # pylint: disable=g-explicit-bool-comparison
+
+        return tf.quantization.quantize_and_dequantize_v2(
+            inputs,
+            min_var,
+            max_var,
+            num_bits=num_bits,
+            narrow_range=narrow_range,
+            range_given=True,
+        )

tools/tensorflow-quantization/tensorflow_quantization/utils.py

@@ -0,0 +1,364 @@
+#
+# 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 collections import deque
+from typing import List
+import os
+import shutil
+from tf2onnx import tf_loader, utils, convert
+import copy
+
+
+def ensure_and_clean_dir(dir_path, do_clean_dir=True) -> None:
+    """Create a directory to save test logs
+
+    Args:
+        dir_path (str): directory to create / clean.
+        do_clean_dir (bool): boolean indicating whether to clean the directory if it already exists (remove+create).
+    Returns:
+        None
+    """
+    if not os.path.exists(dir_path):
+        os.makedirs(dir_path)
+    elif do_clean_dir:
+        shutil.rmtree(dir_path)
+        os.makedirs(dir_path)
+
+
+class Folder:
+    """
+    Folder class that tracks all files for a single experiment.
+    """
+
+    def __init__(self, folder_name) -> None:
+        self.base = folder_name
+        ensure_and_clean_dir(self.base)
+        self.fp32 = os.path.join(self.base, "fp32")
+        ensure_and_clean_dir(self.fp32)
+        self.fp32_saved_model = os.path.join(
+            self.fp32, "saved_model"
+        )  # location of fp32 saved keras model
+        self.fp32_onnx_model = os.path.join(
+            self.fp32, "original.onnx"
+        )  # location of fp32 onnx model
+
+        self.int8 = os.path.join(self.base, "int8")
+        ensure_and_clean_dir(self.int8)
+        self.int8_saved_model = os.path.join(
+            self.int8, "saved_model"
+        )  # location of int8 saved keras model
+        self.int8_onnx_model = os.path.join(
+            self.int8, "quantized.onnx"
+        )  # location of int8 onnx model
+
+
+class CreateAssetsFolders:
+    """Create empty folders to save the original and quantized TensorFlow models and their respective ONNX
+    models for each experiment.
+
+    The following directory structure is created: base_directory -> experiment_directory (created by `add_folder` method) -> (fp32 [saved_model, .onnx model]),
+    (int8 [saved_model, .onnx model]).
+    """
+
+    def __init__(self, base_experiment_directory) -> None:
+        self.base = base_experiment_directory
+        if not os.path.exists(self.base):
+            os.mkdir(self.base)
+
+    def add_folder(self, folder_name: str) -> None:
+        """
+        Create the experiment directory (sub-folder in the base directory passed to this class).
+
+        Args:
+            folder_name (str): name of folder
+
+        Returns:
+            None
+        """
+        setattr(self, folder_name, Folder(os.path.join(self.base, folder_name)))
+
+
+def convert_saved_model_to_onnx(
+    saved_model_dir: str, onnx_model_path: str, opset=13
+) -> None:
+    """Convert Keras saved model into ONNX format.
+    Works directly with CreateAssetsFolder object path.
+
+    Args:
+        saved_model_dir (str): Path to keras saved model.
+        onnx_model_path (str): Full path to ONNX model file.
+
+    Returns:
+        None
+    """
+    # 1. Let TensorRT optimize QDQ nodes instead of TF
+    from tf2onnx.optimizer import _optimizers
+
+    updated_optimizers = copy.deepcopy(_optimizers)
+    del updated_optimizers["q_dq_optimizer"]
+    del updated_optimizers["const_dequantize_optimizer"]
+
+    # 2. Extract graph definition from SavedModel
+    graph_def, inputs, outputs = tf_loader.from_saved_model(
+        model_path=saved_model_dir,
+        input_names=None,
+        output_names=None,
+        tag="serve",
+        signatures=["serving_default"],
+    )
+
+    # 3. Convert tf2onnx and save onnx file
+    model_proto, _ = convert._convert_common(
+        graph_def,
+        opset=opset,
+        input_names=inputs,
+        output_names=outputs,
+        output_path=onnx_model_path,
+        optimizers=updated_optimizers,
+    )
+
+    utils.save_protobuf(onnx_model_path, model_proto)
+    print("ONNX conversion Done!")
+
+
+def convert_keras_model_to_onnx(
+    keras_model: tf.keras.Model, onnx_model_path: str, opset=13
+) -> None:
+    """Convert in-memory Keras model into ONNX format.
+    Works directly with CreateAssetsFolder object path.
+
+    Args:
+        keras_model (tf.keras.Model): Keras model.
+        onnx_model_path (str): Full path to ONNX model file.
+
+    Returns:
+        None
+    """
+    # 1. Let TensorRT optimize QDQ nodes instead of TF
+    from tf2onnx.optimizer import _optimizers
+
+    updated_optimizers = copy.deepcopy(_optimizers)
+    del updated_optimizers["q_dq_optimizer"]
+    del updated_optimizers["const_dequantize_optimizer"]
+
+    # 2. Convert keras model directly and save onnx file.
+    onnx_model_proto, _ = convert.from_keras(keras_model, opset=opset, optimizers=updated_optimizers)
+    utils.save_protobuf(onnx_model_path, onnx_model_proto)
+
+
+class KerasModelTraveller:
+    """
+    Utility class to travel Keras model and print out detailed layer information.
+    """
+
+    def __init__(self, print_layer_config=False) -> None:
+        self._pc = print_layer_config
+        self.model_list = deque([])
+        # Used to filter which classes you want printed, by layer.__class__
+        self._filter_by_class = None
+        self._layer_names = []
+        self._print_basic_info = None
+
+    def _print_layer_info(self, layer):
+        assert isinstance(layer, tf.keras.layers.Layer)
+        if self._filter_by_class is None or layer.__class__ in self._filter_by_class:
+            self._layer_names.append(layer.name)
+            if self._print_basic_info:
+                print(
+                    "layer name:{layer_name}, layer class:{layer_class}".format(
+                        layer_name=layer.name, layer_class=layer.__class__
+                    )
+                )
+            if self._pc:
+                print(layer.get_config())
+            if self._print_basic_info:
+                print("-----------------")
+
+    def _dissect(self):
+        if not self.model_list:
+            return
+        number_of_models = len(self.model_list)
+        for _ in range(number_of_models):
+            # Get a subclassed model
+            current_model = self.model_list.pop()
+            print("Keras Subclassed Model: {}".format(current_model.__class__.__name__))
+            assert isinstance(current_model, tf.keras.Model)
+            for l in current_model.layers:
+                if isinstance(l, tf.keras.Model):
+                    # This is another subclassed model inside
+                    # Add this model to model queue for further analysis
+                    self.model_list.appendleft(l)
+                    self._dissect()
+                else:
+                    # This is a layer
+                    self._print_layer_info(l)
+
+    def _travel(
+        self, keras_model: tf.keras.Model, filter_by_class=None, print_basic_info=False
+    ):
+        """Gets layer info by dissecting the model (need for multi-layered models)
+
+        Args:
+            keras_model (tf.keras.Model): Keras model
+            filter_by_class (str): None or array of layer.__class__ to print
+
+        Returns:
+            None
+        """
+        self.filter_by_class = filter_by_class
+        self._print_basic_info = print_basic_info
+        assert isinstance(
+            keras_model, tf.keras.Model
+        ), "Model passed is not Keras model"
+        self.model_list.appendleft(keras_model)
+        self._dissect()
+        self.filter_by_class = None
+
+    def get_layer_names(self, keras_model: tf.keras.Model, filter_by_class=None):
+        """Get name of all layers in the model.
+
+        Args:
+            keras_model (tf.keras.Model): Keras model
+            filter_by_class (str): None or array of layer.__class__ to print
+
+        Returns:
+            None
+        """
+        self._travel(keras_model=keras_model, filter_by_class=filter_by_class)
+        return self._layer_names
+
+    def get_layer_information(self, keras_model: tf.keras.Model, filter_by_class=None):
+        """Print information about all layers.
+
+        Args:
+            keras_model (tf.keras.Model): Keras model
+            filter_by_class (str): None or array of layer.__class__ to print
+
+        Returns:
+            None
+        """
+        self._travel(
+            keras_model=keras_model,
+            filter_by_class=filter_by_class,
+            print_basic_info=True,
+        )
+
+
+def _get_layer_info(layer: tf.keras.layers.Layer) -> dict:
+    """
+    Returns the layer's class, module, and name
+    """
+    return {
+        "class": layer.__class__.__name__,
+        "module": layer.__class__.__module__,
+        "name": layer.name,
+        "layer": layer,
+    }
+
+
+def _get_previous_layers_class_and_module_and_name(
+    layer: tf.keras.layers.Layer,
+) -> List[dict]:
+    """
+    For a given layer return a dictionary with name, module and class information of all previous layers.
+    """
+    r = []
+    if isinstance(layer.input, list):
+        for layer_input_tensor in layer.input:
+            ip_tensor_parent_layer = layer_input_tensor._keras_history.layer
+            r.append(_get_layer_info(ip_tensor_parent_layer))
+    else:
+        ip_tensor_parent_layer = layer.input._keras_history.layer
+        r.append(_get_layer_info(ip_tensor_parent_layer))
+    return r
+
+
+def find_my_predecessors(model: tf.keras.Model, current_layer_name: str) -> List[dict]:
+    """
+    Given a layer name, find all predecessors of that layer.
+
+    Args:
+        model (tf.keras.Model): Keras functional model
+        current_layer_name (str): name of a model layer for which predecessors has to be found.
+
+    Returns:
+        List[dict]: List of predecessors. Each dictionary has three keys as follows,
+        ::
+            {'class':<pred_layer_class>, 'module':<pred_layer_module>, 'name':<pred_layer_name>}
+
+    Raises:
+        AssertionError: If model is subclassed or current_layer_name is not string.
+    """
+    supported_model_classes = {"Functional", "Sequential"}
+    assert isinstance(current_layer_name, str), "current layer name should be passed."
+    assert (
+        model.__class__.__name__ in supported_model_classes
+    ), "model should be Functional or Sequential."
+
+    for layer in model.layers:
+        if layer.name == current_layer_name:
+            return _get_previous_layers_class_and_module_and_name(layer)
+
+
+def find_my_successors(model: tf.keras.Model, current_layer_name: str) -> List[dict]:
+    """
+    Given a layer name, find all successors of that layer.
+
+    Args:
+        model (tf.keras.Model): Keras functional model
+        current_layer_name (str): name of a model layer for which successors has to be found.
+
+    Returns:
+        List[dict]: List of predecessors. Each dictionary has three keys as follows,
+        ::
+            {'class':<pred_layer_class>, 'module':<pred_layer_module>, 'name':<pred_layer_name>}
+
+    Raises:
+        AssertionError: If model is subclassed or current_layer_name is not string.
+    """
+    supported_model_classes = {"Functional", "Sequential"}
+    assert isinstance(current_layer_name, str), "current layer name should be passed."
+    assert (
+        model.__class__.__name__ in supported_model_classes
+    ), "model should be Functional or Sequential."
+
+    def _check_all_next_layers_with_connection_to_current(
+        next_layers: List[tf.keras.layers.Layer],
+        current_layer_name: str,
+        current_layer_class: str,
+    ):
+        successors = []
+        for layer in next_layers:
+            p_layers = _get_previous_layers_class_and_module_and_name(layer)
+            for p_layer in p_layers:
+                if (
+                    p_layer["class"] == current_layer_class
+                    and p_layer["name"] == current_layer_name
+                ):
+                    successors.append(_get_layer_info(layer))
+        return successors
+
+    all_layers = model.layers
+    for i, layer in enumerate(all_layers):
+        if layer.name == current_layer_name:
+            next_layers = all_layers[i + 1 :]
+            layer_info = _get_layer_info(layer)
+            return _check_all_next_layers_with_connection_to_current(
+                next_layers, layer_info["name"], layer_info["class"]
+            )

tools/tensorflow-quantization/tests/custom_qdq_cases_test.py

@@ -0,0 +1,216 @@
+#
+# 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.
+#
+
+
+from network_pool import pippin_28_28
+from tensorflow_quantization import custom_qdq_cases
+import pytest
+from tensorflow_quantization import quantize_model
+from tensorflow_quantization.utils import convert_saved_model_to_onnx
+from tensorflow_quantization.utils import CreateAssetsFolders
+
+import tensorflow as tf
+
+test_assets = CreateAssetsFolders("test_custom_qdq_cases")
+
+
+def test_resnet_residual_qdq_case():
+    model = pippin_28_28()
+    test_assets.add_folder("pipin_28_28")
+    tf.keras.models.save_model(model, test_assets.pipin_28_28.fp32_saved_model)
+    convert_saved_model_to_onnx(
+        saved_model_dir=test_assets.pipin_28_28.fp32_saved_model,
+        onnx_model_path=test_assets.pipin_28_28.fp32_onnx_model,
+    )
+
+    resnet_residual_qdq = custom_qdq_cases.ResNetV1QDQCase()
+    r = resnet_residual_qdq.case(model, None)
+    expected_qdq_insertion = {
+        "add": 1,
+        "add_1": "any",
+    }
+    assert (
+        len(r.layers) == 2
+    ), "There should be 2 custom layers, but found {}".format(len(r.layers))
+    for l in r.layers:
+        if l.name not in expected_qdq_insertion:
+            raise Exception(
+                "Layer {} is not expected to be treated as custom layer".format(l.name)
+            )
+        else:
+            if l.quantization_index != None:
+                if expected_qdq_insertion[l.name] == "any":
+                    continue
+                assert (
+                    l.quantization_index[0] == expected_qdq_insertion[l.name]
+                ), "For layer {l_name}, only {expected_qdq} indices should be quantized".format(
+                    l_name=l.name, expected_qdq=expected_qdq_insertion[l.name]
+                )
+
+
+def assert_add_bn_expected_layers(
+        r, expected_add_layer_behavior, expected_bn_layer_behavior, expected_mp_layer_behavior
+):
+    assert len(r.layers) == (
+            len(expected_add_layer_behavior) + len(expected_bn_layer_behavior) + len(expected_mp_layer_behavior)
+    ), "Not all expected layers are captured for ResNet custom QDQ case."
+    for l in r.layers:
+        assert (
+                l.name in expected_add_layer_behavior
+                or l.name in expected_bn_layer_behavior
+                or l.name in expected_mp_layer_behavior
+        ), "layer {} is not expected to be captured for ResNet custom QDQ case".format(
+            l.name
+        )
+        if "add" in l.name:
+            if expected_add_layer_behavior[l.name] == "any":
+                continue
+            assert l.quantization_index[0] == expected_add_layer_behavior[l.name], (
+                "For layer {l_name}, expected quantization index is {expected_add_behavior} but index {l_quant_id} "
+                "is captured in ResNet custom QDQ case.".format(
+                    l_name=l.name,
+                    expected_add_behavior=expected_add_layer_behavior[l.name],
+                    l_quant_idx=l.quantization_index[0],
+                )
+            )
+
+
+def test_resnet50_residual_qdq_case():
+    resnet50 = tf.keras.applications.resnet50.ResNet50(weights=None)
+    test_assets.add_folder("resnet50_v1")
+    tf.keras.models.save_model(resnet50, test_assets.resnet50_v1.fp32_saved_model)
+    convert_saved_model_to_onnx(
+        saved_model_dir=test_assets.resnet50_v1.fp32_saved_model,
+        onnx_model_path=test_assets.resnet50_v1.fp32_onnx_model,
+    )
+
+    resnet_custom_qdq_case = custom_qdq_cases.ResNetV1QDQCase()
+    r = resnet_custom_qdq_case.case(resnet50, None)
+    for r_1 in r.layers:
+        print("\"{}\",".format(r_1.name))
+
+    expected_add_layer_behavior = {
+        "conv2_block1_add": "any",
+        "conv2_block2_add": 0,
+        "conv2_block3_add": 0,
+        "conv3_block1_add": "any",
+        "conv3_block2_add": 0,
+        "conv3_block3_add": 0,
+        "conv3_block4_add": 0,
+        "conv4_block1_add": "any",
+        "conv4_block2_add": 0,
+        "conv4_block3_add": 0,
+        "conv4_block4_add": 0,
+        "conv4_block5_add": 0,
+        "conv4_block6_add": 0,
+        "conv5_block1_add": "any",
+        "conv5_block2_add": 0,
+        "conv5_block3_add": 0,
+    }
+
+    # Empty, no BatchNorm layers should be quantized in ResNet-v1
+    expected_bn_layer_behavior = {}
+
+    # MaxPool quantization is actually not needed in ResNet-v1
+    expected_mp_layer_behavior = {}
+
+    assert_add_bn_expected_layers(
+      r, expected_add_layer_behavior, expected_bn_layer_behavior, expected_mp_layer_behavior
+    )
+
+    q_resnet50 = quantize_model(resnet50, custom_qdq_cases=[resnet_custom_qdq_case])
+    tf.keras.models.save_model(q_resnet50, test_assets.resnet50_v1.int8_saved_model)
+    convert_saved_model_to_onnx(
+        saved_model_dir=test_assets.resnet50_v1.int8_saved_model,
+        onnx_model_path=test_assets.resnet50_v1.int8_onnx_model,
+    )
+
+
+def test_resnet50v2_bn_qdq_case():
+    resnet50_v2 = tf.keras.applications.resnet_v2.ResNet50V2(weights=None)
+    test_assets.add_folder("resnet50_v2")
+    tf.keras.models.save_model(resnet50_v2, test_assets.resnet50_v2.fp32_saved_model)
+    convert_saved_model_to_onnx(
+        saved_model_dir=test_assets.resnet50_v2.fp32_saved_model,
+        onnx_model_path=test_assets.resnet50_v2.fp32_onnx_model,
+    )
+
+    resnet_custom_qdq_case = custom_qdq_cases.ResNetV2QDQCase()
+    r = resnet_custom_qdq_case.case(resnet50_v2, None)
+    for r_1 in r.layers:
+        print("\"{}\",".format(r_1.name))
+
+    expected_add_layer_behavior = {
+        "conv2_block1_out": 0,
+        "conv2_block2_out": 0,
+        "conv2_block3_out": 0,
+        "conv3_block1_out": 0,
+        "conv3_block2_out": 0,
+        "conv3_block3_out": 0,
+        "conv3_block4_out": 0,
+        "conv4_block1_out": 0,
+        "conv4_block2_out": 0,
+        "conv4_block3_out": 0,
+        "conv4_block4_out": 0,
+        "conv4_block5_out": 0,
+        "conv4_block6_out": 0,
+        "conv5_block1_out": 0,
+        "conv5_block2_out": 0,
+        "conv5_block3_out": 0,
+    }
+
+    # ResNet-v2 quantizes BatchNorms that are not connected to Conv layers
+    expected_bn_layer_behavior = {
+        "conv2_block1_preact_bn",
+        "conv2_block2_preact_bn",
+        "conv2_block3_preact_bn",
+        "conv3_block1_preact_bn",
+        "conv3_block2_preact_bn",
+        "conv3_block3_preact_bn",
+        "conv3_block4_preact_bn",
+        "conv4_block1_preact_bn",
+        "conv4_block2_preact_bn",
+        "conv4_block3_preact_bn",
+        "conv4_block4_preact_bn",
+        "conv4_block5_preact_bn",
+        "conv4_block6_preact_bn",
+        "conv5_block1_preact_bn",
+        "conv5_block2_preact_bn",
+        "conv5_block3_preact_bn",
+        "post_bn",
+    }
+
+    # ResNet-v2 quantizes all MaxPool layers
+    expected_mp_layer_behavior = {
+        "pool1_pool",
+        "max_pooling2d",
+        "max_pooling2d_1",
+        "max_pooling2d_2",
+    }
+
+    assert_add_bn_expected_layers(
+        r, expected_add_layer_behavior, expected_bn_layer_behavior, expected_mp_layer_behavior
+    )
+
+    q_resnet50_v2 = quantize_model(
+        resnet50_v2, custom_qdq_cases=[resnet_custom_qdq_case]
+    )
+    tf.keras.models.save_model(q_resnet50_v2, test_assets.resnet50_v2.int8_saved_model)
+    convert_saved_model_to_onnx(
+        saved_model_dir=test_assets.resnet50_v2.int8_saved_model,
+        onnx_model_path=test_assets.resnet50_v2.int8_onnx_model,
+    )

tools/tensorflow-quantization/tests/network_pool.py

@@ -0,0 +1,277 @@
+#
+# 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 tiny networks used for testing across different modules.
+They are named after famous Hobbits for obvious reasons.
+"""
+
+import tensorflow as tf
+
+##################################################
+###### Tiny, VGG like network ####################
+##################################################
+def bilbo_28_28():
+    """
+    Network with VGG like architecture.
+    """
+    input_img = tf.keras.layers.Input(shape=(28, 28), name="nn_input")
+    x = tf.keras.layers.Reshape(target_shape=(28, 28, 1), name="reshape_0")(input_img)
+    x = tf.keras.layers.Conv2D(filters=516, kernel_size=(3, 3), name="conv_0")(x)
+    x = tf.keras.layers.ReLU(name="relu_0")(x)
+    x = tf.keras.layers.Conv2D(filters=252, kernel_size=(3, 3), name="conv_1")(x)
+    x = tf.keras.layers.ReLU(name="relu_1")(x)
+    x = tf.keras.layers.Conv2D(filters=126, kernel_size=(3, 3), name="conv_2")(x)
+    x = tf.keras.layers.ReLU(name="relu_2")(x)
+    x = tf.keras.layers.Conv2D(filters=64, kernel_size=(3, 3), name="conv_3")(x)
+    x = tf.keras.layers.ReLU(name="relu_3")(x)
+    x = tf.keras.layers.Conv2D(filters=32, kernel_size=(3, 3), name="conv_4")(x)
+    x = tf.keras.layers.ReLU(name="relu_4")(x)
+    x = tf.keras.layers.Conv2D(filters=16, kernel_size=(3, 3), name="conv_5")(x)
+    x = tf.keras.layers.ReLU(name="relu_5")(x)
+    x = tf.keras.layers.Conv2D(filters=8, kernel_size=(3, 3), name="conv_6")(x)
+    x = tf.keras.layers.ReLU(name="relu_6")(x)
+    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2), name="max_pool_0")(x)
+    x = tf.keras.layers.Flatten(name="flatten_0")(x)
+    x = tf.keras.layers.Dense(100, name="dense_0")(x)
+    x = tf.keras.layers.ReLU(name="relu_7")(x)
+    x = tf.keras.layers.Dense(10, name="dense_1")(x)
+    return tf.keras.Model(input_img, x, name="Bilbo")
+
+
+#####################################################
+###### Tiny, ResNet like network ####################
+#####################################################
+def identity_block_plain(input_tensor):
+    """
+    Identity block with no shortcut convolution
+    """
+    y = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), padding="same")(
+        input_tensor
+    )
+    y = tf.keras.layers.ReLU()(y)
+    y = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3), padding="same")(y)
+    out = tf.keras.layers.Add()([y, input_tensor])
+    out = tf.keras.layers.ReLU()(out)
+    return out
+
+
+def identity_block_short_conv_plain(input_tensor):
+    """
+    Identity block with shortcut convolution
+    """
+    y = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), padding="same")(
+        input_tensor
+    )
+    y = tf.keras.layers.ReLU()(y)
+    y = tf.keras.layers.Conv2D(
+        filters=24, kernel_size=(3, 3), strides=(2, 2), padding="same"
+    )(y)
+    ds_input = tf.keras.layers.Conv2D(
+        filters=24, kernel_size=(3, 3), strides=(2, 2), padding="same"
+    )(input_tensor)
+    out = tf.keras.layers.Add()([y, ds_input])
+    out = tf.keras.layers.ReLU()(out)
+    return out
+
+
+def frodo_32_32():
+    """
+    Dummy network with resnet like architecture.
+    """
+    input_img = tf.keras.layers.Input(shape=(32, 32, 3))
+    x = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3))(input_img)
+    x = tf.keras.layers.ReLU()(x)
+    x = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3))(x)
+    x = tf.keras.layers.ReLU()(x)
+    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
+    x = identity_block_plain(x)
+    x = identity_block_short_conv_plain(x)
+    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
+    x = tf.keras.layers.Flatten()(x)
+    x = tf.keras.layers.Dense(100)(x)
+    x = tf.keras.layers.ReLU()(x)
+    x = tf.keras.layers.Dense(10)(x)
+    return tf.keras.Model(input_img, x, name="Frodo")
+
+
+def sam_32_32():
+    """
+    Dummy network with resnet like architecture.
+    """
+    input_img = tf.keras.layers.Input(shape=(32, 32, 3))
+    x = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3))(input_img)
+    x = tf.keras.layers.ReLU()(x)
+    x = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3))(x)
+    x = tf.keras.layers.ReLU()(x)
+    x = identity_block_plain(x)
+    x = identity_block_plain(x)
+    x = identity_block_short_conv_plain(x)
+    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
+    x = tf.keras.layers.Flatten()(x)
+    x = tf.keras.layers.Dense(100)(x)
+    x = tf.keras.layers.ReLU()(x)
+    x = tf.keras.layers.Dense(10)(x)
+    return tf.keras.Model(input_img, x, name="Sam")
+
+
+##############################################
+###### Popular network blocks ################
+##############################################
+def relu_bn(input):
+    """
+    Block with BN+ReLU
+    """
+    bn = tf.keras.layers.BatchNormalization()(input)
+    relu = tf.keras.layers.ReLU()(bn)
+    return relu
+
+
+def bn(input):
+    return tf.keras.layers.BatchNormalization()(input)
+
+
+def relu(input):
+    return tf.keras.layers.ReLU()(input)
+
+
+def inception_block(input_tensor):
+    """
+    Inception block from GoogleNet
+    """
+    b1x1 = tf.keras.layers.Conv2D(filters=12, kernel_size=(1, 1), padding="same")(
+        input_tensor
+    )
+    b1x1 = relu_bn(b1x1)
+
+    b5x5 = tf.keras.layers.Conv2D(filters=12, kernel_size=(1, 1), padding="same")(
+        input_tensor
+    )
+    b5x5 = tf.keras.layers.Conv2D(filters=24, kernel_size=(5, 5), padding="same")(b5x5)
+    b5x5 = relu_bn(b5x5)
+
+    b3x3 = tf.keras.layers.Conv2D(filters=12, kernel_size=(1, 1), padding="same")(
+        input_tensor
+    )
+    b3x3 = tf.keras.layers.Conv2D(filters=20, kernel_size=(3, 3), padding="same")(b3x3)
+    b3x3 = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3), padding="same")(b3x3)
+    b3x3 = relu_bn(b3x3)
+
+    out = tf.keras.layers.Concatenate()([b1x1, b5x5])
+    return out
+
+
+def identity_block_bn(input_tensor):
+    """
+    Identity block with no shortcut convolution
+    """
+    y = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), padding="same")(
+        input_tensor
+    )
+    y = relu_bn(y)
+    y = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3), padding="same")(y)
+    y = bn(y)
+    out = tf.keras.layers.Add()([y, input_tensor])
+    out = relu(out)
+    return out
+
+
+def identity_block_short_conv_bn(input_tensor):
+    """
+    Identity block with shortcut convolution
+    """
+    y = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3), padding="same")(
+        input_tensor
+    )
+    y = relu_bn(y)
+    y = tf.keras.layers.Conv2D(
+        filters=24, kernel_size=(3, 3), strides=(2, 2), padding="same"
+    )(y)
+    y = bn(y)
+    ds_input = tf.keras.layers.Conv2D(
+        filters=24, kernel_size=(3, 3), strides=(2, 2), padding="same"
+    )(input_tensor)
+    ds_input = bn(ds_input)
+    out = tf.keras.layers.Add()([y, ds_input])
+    out = relu(out)
+    return out
+
+
+def otho_28_28():
+    input_img = tf.keras.layers.Input(shape=(28, 28), name="input_0")
+    r = tf.keras.layers.Reshape(target_shape=(28, 28, 1), name="reshape_0")(input_img)
+    x = tf.keras.layers.Conv2D(filters=2, kernel_size=(3, 3), name="conv_0")(r)
+    x = tf.keras.layers.ReLU(name="relu_0")(x)
+    x = tf.keras.layers.Conv2D(filters=2, kernel_size=(3, 3), name="conv_1")(x)
+    x = tf.keras.layers.ReLU(name="relu_1")(x)
+    x = tf.keras.layers.Flatten(name="flatten_0")(x)
+    return tf.keras.Model(input_img, x, name="Otho")
+
+
+def lotho_28_28():
+    input_img = tf.keras.layers.Input(shape=(28, 28), name="input_0")
+    r = tf.keras.layers.Reshape(target_shape=(28, 28, 1), name="reshape_0")(input_img)
+    x = tf.keras.layers.DepthwiseConv2D(kernel_size=(3, 3), name="dconv_0")(r)
+    x = tf.keras.layers.ReLU(name="relu_0")(x)
+    x = tf.keras.layers.DepthwiseConv2D(kernel_size=(3, 3), name="dconv_1")(x)
+    x = tf.keras.layers.ReLU(name="relu_1")(x)
+    x = tf.keras.layers.Flatten(name="flatten_0")(x)
+    return tf.keras.Model(input_img, x, name="Lotho")
+
+
+def lobelia_28_28():
+    input_img = tf.keras.layers.Input(shape=(28, 28), name="input_0")
+    r = tf.keras.layers.Reshape(target_shape=(28, 28, 1), name="reshape_0")(input_img)
+    x = tf.keras.layers.Flatten(name="flatten_0")(r)
+    x = tf.keras.layers.Dense(100, name="dense_0")(x)
+    x = tf.keras.layers.ReLU(name="relu_0")(x)
+    x = tf.keras.layers.Dense(10, name="dense_1")(x)
+    return tf.keras.Model(input_img, x, name="Lobelia")
+
+
+def merry_28_28():
+    input_img = tf.keras.layers.Input(shape=(28, 28))
+    x = tf.keras.layers.Reshape(target_shape=(28, 28, 1))(input_img)
+    x = tf.keras.layers.Conv2D(filters=8, kernel_size=(3, 3))(x)
+    x = relu_bn(x)
+    x = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3))(x)
+    x = relu_bn(x)
+    x = inception_block(x)
+    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
+    x = tf.keras.layers.Flatten()(x)
+    x = tf.keras.layers.Dense(100)(x)
+    x = tf.keras.layers.ReLU()(x)
+    x = tf.keras.layers.Dense(10)(x)
+    return tf.keras.Model(input_img, x, name="Merry")
+
+
+def pippin_28_28():
+    input_img = tf.keras.layers.Input(shape=(28, 28))
+    x = tf.keras.layers.Reshape(target_shape=(28, 28, 1))(input_img)
+    x = tf.keras.layers.Conv2D(filters=12, kernel_size=(3, 3))(x)
+    x = relu_bn(x)
+    x = tf.keras.layers.Conv2D(filters=24, kernel_size=(3, 3))(x)
+    x = relu_bn(x)
+    x = identity_block_bn(x)
+    x = identity_block_short_conv_bn(x)
+    x = tf.keras.layers.MaxPooling2D(pool_size=(2, 2))(x)
+    x = tf.keras.layers.Flatten()(x)
+    x = tf.keras.layers.Dense(100)(x)
+    x = tf.keras.layers.ReLU()(x)
+    x = tf.keras.layers.Dense(10)(x)
+    return tf.keras.Model(input_img, x, name="Pippin")

tools/tensorflow-quantization/tests/onnx_graph_qdq_validator.py

@@ -0,0 +1,551 @@
+#
+# 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 onnx
+import onnx_graphsurgeon as gs
+import tensorflow as tf
+from tensorflow_quantization.quantize import LayerConfig, quantize_model
+from typing import List, Tuple
+from tensorflow_quantization.utils import convert_saved_model_to_onnx
+import copy
+
+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="Concatenate",
+        is_keras_class=True,
+        quantize_weight=False,
+        quantization_index=["all"],
+    ),
+    LayerConfig(
+        name="AveragePooling2D", is_keras_class=True, quantize_weight=False
+    ),
+    LayerConfig(
+        name="GlobalAveragePooling2D", is_keras_class=True, quantize_weight=False
+    )
+]
+
+
+class ONNXQDQValidator:
+    """
+    Validate ONNX file for correct QDQ insertion.
+    All onnx-graphsurgeon terminologies are used in the explanations.
+    """
+
+    def __init__(self) -> None:
+        self.expected_qdq_layer_behavior = {}
+        self.graph = None
+        self.data_format = tf.keras.backend.image_data_format()
+
+    @staticmethod
+    def _extract_layer_names_from_class_type(
+            expected_qdq_behavior, original_keras_model
+    ):
+        """Checks if expected_qdq_behavior has items where is_keras_class=True and extract all layers relevant to it.
+        Also checks if the user didn't specifically name that layer in expected_qdq_behavior.
+        """
+
+        def layer_is_class_type(class_specs, origin_layer):
+            for c in class_specs:
+                if origin_layer.__class__.__name__ == c.name:
+                    return c
+            return None
+
+        def skip_layer_name(layer_name):
+            for layer in expected_qdq_behavior:
+                if layer_name == layer.name:
+                    return True
+            return False
+
+        expected_qdq_behavior_class = [
+            layer for layer in expected_qdq_behavior if layer.is_keras_class
+        ]
+        expected_qdq_behavior_layers = [
+            layer for layer in expected_qdq_behavior
+            # Skip if quantize_input and quantize_weight=False
+            if not layer.is_keras_class and (layer.quantize_input or layer.quantize_weight)
+        ]
+
+        if original_keras_model is not None:
+            for original_layer in original_keras_model.layers:
+                class_type = layer_is_class_type(
+                    expected_qdq_behavior_class, original_layer
+                )
+                if class_type is not None and not skip_layer_name(original_layer.name):
+                    # Skip if quantize_input and quantize_weight=False
+                    if class_type.quantize_input or class_type.quantize_weight:
+                        expected_qdq_behavior_layers.append(
+                            LayerConfig(
+                                name=original_layer.name,
+                                quantize_input=class_type.quantize_input,
+                                quantize_weight=class_type.quantize_weight,
+                                quantization_index=class_type.quantization_index,
+                            )
+                        )
+
+        return expected_qdq_behavior_layers
+
+    def _collect_layer_names(self, expected_qdq_behavior):
+        """
+        Populates the global variable 'self.expected_qdq_layer_behavior', a dictionary in the format:
+            key (string) : Layer name after quantization wrapper is applied.
+            value (list) : List with layer specific parameters.
+                value[0] (bool) = True if this layer is a keras class.
+                value[1] (bool) = True if input to this layer should be quantized.
+                value[2] (bool) = True if layer weight should be quantized.
+                value[3] (list) = List of quantization index, if any.
+                value[4] (bool) = Set to False initially but when quantization of this layer is verified, set to True.
+        """
+        for layer in expected_qdq_behavior:
+            self.expected_qdq_layer_behavior["quant_" + layer.name] = []
+            self.expected_qdq_layer_behavior["quant_" + layer.name].append(
+                layer.is_keras_class
+            )
+            self.expected_qdq_layer_behavior["quant_" + layer.name].append(
+                layer.quantize_input
+            )
+            self.expected_qdq_layer_behavior["quant_" + layer.name].append(
+                layer.quantize_weight
+            )
+            self.expected_qdq_layer_behavior["quant_" + layer.name].append(
+                layer.quantization_index if layer.quantization_index is not None else []
+            )
+            self.expected_qdq_layer_behavior["quant_" + layer.name].append(False)
+
+    def _load_onnx_graph(self, onnx_model_path):
+        self.graph = gs.import_onnx(onnx.load(onnx_model_path))
+
+    def _get_tf_name_of_node(self, onnx_node):
+        splitted_node_name = onnx_node.name.split("/")
+        if len(splitted_node_name) > 1:
+            # This is other node than QuantizeLinear or DequantizeLinear
+            node_op = onnx_node.op
+
+            # Most layers have their name in position -2
+            #   List: Conv, BatchNormalization, Relu, Add, MatMul, Softmax, Pad, MaxPool, GlobalAveragePool
+            # Exceptions: Squeeze, Transpose, Reshape
+            if node_op == "Squeeze" or node_op == "Transpose" or node_op == "Reshape":
+                return splitted_node_name[-1]
+
+            # Quantized layers
+            for exp_qdq_layer_name in self.expected_qdq_layer_behavior.keys():
+                if exp_qdq_layer_name + "/" in onnx_node.name:
+                    return exp_qdq_layer_name
+
+            # Other layers
+            return splitted_node_name[-2]
+        else:
+            return None
+
+    def _get_input_tensor_parent(self, onnx_node, input_idx):
+        """
+        Get input Tensors parent recursively.
+        Here we want to know id DequantizeLinear is tensors parent.
+        Recursively we go up the graph since reshape layers are added while onnx conversion between QDQ and node.
+        """
+        current_node_ip_tensor = onnx_node.inputs[input_idx]
+        try:
+            current_node_ip_tensor_parent = current_node_ip_tensor.inputs[0]
+        except IndexError:
+            # Example for weight Tensor, parent is None
+            return None
+        while (
+                current_node_ip_tensor_parent.op == "Transpose"
+                or current_node_ip_tensor_parent.op == "Reshape"
+        ):  # and self.data_format == "channels_last":
+            # When image data format is 'channels_last' or Conv is of type 'Depthwise', Transpose and/or Reshape
+            #   layers are added between QDQ and target layer. Always select input at index 0 since it's the
+            #   variable coming from the previous node. Other indices, if present, are constant inputs to the node.
+            current_node_ip_tensor = current_node_ip_tensor_parent.inputs[0]
+            try:
+                current_node_ip_tensor_parent = current_node_ip_tensor.inputs[0]
+            except IndexError:
+                # We can't move upwards anymore in the graph
+                break
+
+        return current_node_ip_tensor_parent.op
+
+    def _weighted_qdq_behavior(self, node, tf_node_name):
+        """
+        For weighted layers such as MatMul/Conv2D and DepthwiseConv2D, only quantize_input and quantize_weight options
+          are valid.
+        Input has length of 2 usually.
+        Index 0 is Variable i.e. output of previous op
+        Index 1 is Constant i.e. weight
+        NOTE: In general, for node with more than one inputs, index 0 is variable coming out of previous node.
+              Other indices are constant inputs to the node.
+        """
+        # case 1. Only one of quantize_input or quantize_weight is True
+        if (
+                not self.expected_qdq_layer_behavior[tf_node_name][1]
+                or not self.expected_qdq_layer_behavior[tf_node_name][2]
+        ):
+            # subcase 1. When quantize_input=False
+            if not self.expected_qdq_layer_behavior[tf_node_name][1]:
+                if self._get_input_tensor_parent(node, 0) == "DequantizeLinear":
+                    print(
+                        "[E] quantize_input=False but still input is quantized for weighted layer `{}`".format(
+                            tf_node_name
+                        )
+                    )
+                    return False
+                # send update that correct quantization is found for intended layer.
+                self.expected_qdq_layer_behavior[tf_node_name][-1] = True
+            # subcase 2. When quantize_weight=False
+            if not self.expected_qdq_layer_behavior[tf_node_name][2]:
+                if self._get_input_tensor_parent(node, 1) == "DequantizeLinear":
+                    print(
+                        "[E] quantize_weight=False but still weight is quantized for weighted layer `{}`".format(
+                            tf_node_name
+                        )
+                    )
+                    return False
+                # send update that correct quantization is found for intended layer.
+                self.expected_qdq_layer_behavior[tf_node_name][-1] = True
+        else:
+            # case 2. Both quantize_input=True, quantize_weight=True
+            # Every input should be output of DequantizeLinear op
+            parent_check = []
+            for idx in range(len(node.inputs)):
+                input_tensor_parent = self._get_input_tensor_parent(node, idx)
+                if input_tensor_parent != "DequantizeLinear":
+                    parent_check.append(0)
+                else:
+                    parent_check.append(1)
+            # Check if both input and weight are quantized (2 inputs == 'DequantizeLinear')
+            #   This takes into consideration that Conv sometimes has a BiasAdd input, which is not quantized.
+            if sum(parent_check) < 2:
+                print(
+                    "[E] quantize_weight=True and quantize_input=True but still not all inputs are quantized for "
+                    "weighted layer `{}`".format(tf_node_name)
+                )
+                return False
+            # send update that correct quantization is found for intended layer.
+            self.expected_qdq_layer_behavior[tf_node_name][-1] = True
+        return True
+
+    def _pool_qdq_behavior(self, node, tf_node_name):
+        """
+        Pool layer has just one input which is variable coming from the previous op.
+        """
+        if self._get_input_tensor_parent(node, 0) != "DequantizeLinear":
+            print(
+                "[E] Variable input for MaxPool layer `{}` is not quantized.".format(
+                    tf_node_name
+                )
+            )
+            return False
+        # send update that correct quantization is found for intended layer.
+        self.expected_qdq_layer_behavior[tf_node_name][-1] = True
+        return True
+
+    def _bn_qdq_behavior(self, node, tf_node_name):
+        """
+        BN has one variable input and four (scale, beta, mean, var) constant inputs.
+        Remember variable input is always at index 0
+        For quantization, just check QDQ nodes insertion in variable input.
+        """
+        # Check if the parent node is not DequantizeLinear and input should be quantized.
+        # Reason: in the ResNet CustomQDQCase, BN is only quantized when preceded by Conv. Otherwise, quantize_input
+        #  (and quantize_weight) is set to False.
+        quantize_input = self.expected_qdq_layer_behavior[tf_node_name][1]
+        if (
+                self._get_input_tensor_parent(node, 0) != "DequantizeLinear"
+                and quantize_input
+        ):
+            print(
+                "[E] Variable input for BatchNormalization layer `{}` is not quantized.".format(
+                    tf_node_name
+                )
+            )
+            return False
+        # send update that correct quantization is found for intended layer.
+        self.expected_qdq_layer_behavior[tf_node_name][-1] = True
+        return True
+
+    def _multi_input_qdq_behavior(self, node, tf_node_name):
+        """
+        For layers with multiple inputs, we need to check whether each intended layer is quantized.
+        """
+        # There is quantization index list, check if provided indices are output of DequantizeLinear
+        for _, e in enumerate(self.expected_qdq_layer_behavior[tf_node_name][3]):
+            if e in ["any", "all"]:
+                all_inputs = len(node.inputs)
+                q_inputs = 0
+                for inp_idx in range(all_inputs):
+                    if node.i(inp_idx).op == "DequantizeLinear":
+                        q_inputs += 1
+                if e == "any" and q_inputs != 1:
+                    print(
+                        "[E] quantization_index=['{}'] thus only one input should be quantized, but {} out of {} "
+                        "inputs are quantized for layer `{}`".format(e, q_inputs, all_inputs, tf_node_name)
+                    )
+                    return False
+                elif e == "all" and q_inputs != all_inputs:
+                    print(
+                        "[E] quantization_index=['{}'] thus all inputs should be quantized, but {} out of {} "
+                        "inputs are quantized for layer `{}`".format(e, q_inputs, all_inputs, tf_node_name)
+                    )
+                    return False
+                # send update that correct quantization is found for intended layer.
+                self.expected_qdq_layer_behavior[tf_node_name][-1] = True
+            else:
+                if node.i(e).op != "DequantizeLinear":
+                    print(
+                        "[E] Input at index {e} in layer `{tf_node_name}` should be quantized but it is not.".format(
+                            e=e, tf_node_name=tf_node_name
+                        )
+                    )
+                    return False
+                # send update that correct quantization is found for intended layer.
+                self.expected_qdq_layer_behavior[tf_node_name][-1] = True
+        return True
+
+    def _non_quantized_layer_qdq_behavior(self, node, tf_node_name):
+        """
+        Squeeze layer should not be quantized.
+        """
+        if self._get_input_tensor_parent(node, 0) == "DequantizeLinear":
+            print(
+                "[E] Variable input for {node_op} layer `{tf_node_name}` is quantized.".format(
+                    node_op=node.op, tf_node_name=tf_node_name
+                )
+            )
+            return False
+        # send update that correct quantization is found for intended layer.
+        self.expected_qdq_layer_behavior[tf_node_name][-1] = True
+        return True
+
+    def _qdq_monitor(self, node, tf_node_name):
+        m = {
+            "MatMul": self._weighted_qdq_behavior,
+            "Conv": self._weighted_qdq_behavior,
+            "MaxPool": self._pool_qdq_behavior,
+            "AveragePool": self._pool_qdq_behavior,
+            "GlobalAveragePool": self._pool_qdq_behavior,
+            "BatchNormalization": self._bn_qdq_behavior,
+            "Concat": self._multi_input_qdq_behavior,
+            "Add": self._multi_input_qdq_behavior,
+            "Mul": self._multi_input_qdq_behavior,
+        }
+        if node.op not in m:  # Squeeze, Softmax, ...
+            m[node.op] = self._non_quantized_layer_qdq_behavior
+        return m[node.op](node, tf_node_name)
+
+    def check_onnx_node(self, node_name):
+        for node in self.graph.nodes:
+            if node.name == node_name:
+                print(node)
+
+    def _unintended_layer_quantize_check_pass(self):
+        """
+        Check whether un-intended layer is quantized.
+        If any un-intended layer is quantized, checking fails immediately.
+        """
+        for node in self.graph.nodes:
+            tf_node_name = self._get_tf_name_of_node(node)
+            if tf_node_name and "quant" in tf_node_name:
+                if tf_node_name not in self.expected_qdq_layer_behavior:
+                    print(
+                        "[E] layer `{}` should not be quantized.".format(tf_node_name)
+                    )
+                    return False
+        return True
+
+    def _intended_layer_quantize_check_pass(self):
+        """
+        Checks if the layers exists and whether all expected layers are quantized.
+        If any intended layer is not quantized, checking fails immediately.
+        """
+        for k, v in self.expected_qdq_layer_behavior.items():
+            check_quant_layer_exists = any([k + "/" in node.name for node in self.graph.nodes])
+            check_original_layer_exists = any([k.replace("quant_", "") + "/" in node.name for node in self.graph.nodes])
+            if not check_quant_layer_exists:
+                if check_original_layer_exists:
+                    print("[E] layer `{}` should have been quantized but wasn't.".format(k.replace("quant_", "")))
+                    return False
+                else:
+                    print("[W] layer `{}` does not exist.".format(k))
+                    continue
+            elif not v[-1]:
+                print("[E] layer `{}` should be quantized but it did not.".format(k))
+                return False
+        return True
+
+    def _qdq_insertion_check_pass(self):
+        """
+        Validate QDQ insertion.
+        """
+        check_status = True
+        for node in self.graph.nodes:
+            tf_node_name = self._get_tf_name_of_node(node)
+            if tf_node_name and "quant" in tf_node_name:
+                check_status = check_status and self._qdq_monitor(node, tf_node_name)
+                if not check_status:
+                    return check_status
+        return check_status
+
+    def validate(
+            self, onnx_model_path, expected_qdq_behavior, original_keras_model=None
+    ):
+        self._load_onnx_graph(onnx_model_path)
+        expected_qdq_behavior = self._extract_layer_names_from_class_type(
+            expected_qdq_behavior, original_keras_model
+        )
+        # Populate 'self.expected_qdq_layer_behavior'
+        self._collect_layer_names(expected_qdq_behavior)
+        ulcp = self._unintended_layer_quantize_check_pass()
+        if not ulcp:
+            print("[I] Unintended layer quantization check failed.")
+            return False
+        qicp = self._qdq_insertion_check_pass()
+        if not qicp:
+            print("[I] Quantize insertion check failed.")
+            return False
+        ilcp = self._intended_layer_quantize_check_pass()
+        if not ilcp:
+            print("[I] Intended layer quantization check failed.")
+            return False
+        return True
+
+
+def get_expected_qdq_insertion(
+        nn_model_original: tf.keras.Model,
+        qspec_test: "QuantizationSpec" = None,
+        custom_qdq_cases: List["CustomQDQInsertionCase"] = None,
+        quantization_mode: str = "full",
+        expected_qdq_insertion_user: List[LayerConfig] = None
+) -> List[LayerConfig]:
+    """
+    Gets expected QDQ insertion.
+
+    Args:
+        nn_model_original (tf.keras.Model): baseline model (non-quantized), needed to obtain all layers quantized with
+            Custom QDQ Case.
+        qspec_test (QuantizationSpec): Quantization specification to test the quantized model with.
+        custom_qdq_cases (List[CustomQDQInsertionCase]): indicates layers with custom QDQ placements
+            (i.e., ResidualConnectionQDQCase).
+        quantization_mode (str): quantization mode, can be "full" or "partial".
+        expected_qdq_insertion_user (List[LayerConfig]): List of layer configs specified by the user. If 'None', use
+            the default quantization behavior.
+    Returns:
+        expected_qdq_insertion (List[LayerConfig]): list with expected QDQ node placements.
+    """
+    # 1. Establish QDQ node placement behavior for all relevant classes
+    if expected_qdq_insertion_user is not None:
+        # User-specified QDQ behavior
+        expected_qdq_insertion = expected_qdq_insertion_user
+    else:
+        if quantization_mode == "partial":
+            # No classes are quantized by default
+            expected_qdq_insertion = []
+        else:
+            # Default quantization behavior
+            expected_qdq_insertion = copy.deepcopy(EXPECTED_QDQ_INSERTION)
+
+    # 2. Extend quantization behavior with the user's specifications
+    if qspec_test is not None:
+        # Only add layers that are being quantized (don't add when `quantize_input` or 'quantize_weight`=False)
+        expected_qdq_insertion.extend(qspec_test.layers)
+
+    # 3. Extend quantization behavior with the Custom QDQ Cases
+    if custom_qdq_cases is not None:
+        for custom_qdq_case in custom_qdq_cases:
+            qspec_case_object = custom_qdq_case.case(nn_model_original, qspec=qspec_test)
+            expected_qdq_insertion.extend(qspec_case_object.layers)
+
+    # 4. Check if Multiple Input classes have empty or None 'quantization_index'. If so, update it to
+    #   'quantization_index=["all"]'.
+    for exp_insertion in expected_qdq_insertion:
+        if exp_insertion.is_keras_class and exp_insertion.name in ['Add', 'Multiply', 'Concatenate']:
+            if not exp_insertion.quantization_index:  # None or []
+                exp_insertion.quantization_index = ["all"]
+
+    return expected_qdq_insertion
+
+
+# ###############################################
+# ######### Full QAT workflow test ##############
+# ###############################################
+
+
+def validate_quantized_model(
+        test_assets: "CreateAssetsFolders",
+        nn_model_original: tf.keras.Model,
+        quantization_mode: str = "full",
+        qspec: "QuantizationSpec" = None,
+        custom_qdq_cases: List["CustomQDQInsertionCase"] = None,
+        test_name: str = "test",
+        expected_qdq_insertion: List["LayerConfig"] = None
+) -> Tuple[tf.keras.Model, bool]:
+    """
+    Full test workflow: quantization, obtain expected QDQ node placements, check node placements against expected.
+
+    Args:
+        test_assets (CreateAssetsFolders): Folder organizer.
+        nn_model_original (tf.keras.Model): Keras model.
+        quantization_mode (str): quantization mode, can be "full" or "partial".
+        qspec (QuantizationSpec): QuantizationSpec for model quantization.
+        custom_qdq_cases (List[CustomQDQInsertionCase]): list of custom QDQ cases for model quantization.
+        test_name (str): name for this test workflow.
+        expected_qdq_insertion (List[LayerConfig]): expected QDQ insertion classes and/or layers.
+    Returns:
+        q_model (tf.keras.Model): quantized model.
+        validated (bool): indicates whether the quantized ONNX file is correct or not (according to QDQ node placements).
+    """
+    # Create test folders
+    test_assets.add_folder(test_name)
+    test_assets_attr = getattr(test_assets, test_name)
+
+    # Save baseline model
+    tf.keras.models.save_model(nn_model_original, test_assets_attr.fp32_saved_model)
+    convert_saved_model_to_onnx(
+        saved_model_dir=test_assets_attr.fp32_saved_model,
+        onnx_model_path=test_assets_attr.fp32_onnx_model,
+    )
+
+    # Quantize model
+    q_model = quantize_model(
+        model=nn_model_original,
+        quantization_mode=quantization_mode,
+        quantization_spec=copy.deepcopy(qspec),
+        custom_qdq_cases=custom_qdq_cases
+    )
+    # Save quantized model
+    tf.keras.models.save_model(q_model, test_assets_attr.int8_saved_model)
+    convert_saved_model_to_onnx(
+        saved_model_dir=test_assets_attr.int8_saved_model,
+        onnx_model_path=test_assets_attr.int8_onnx_model,
+    )
+
+    # Validate QDQ node placements in ONNX file
+    expected_qdq_insertion = get_expected_qdq_insertion(
+        tf.keras.models.clone_model(nn_model_original),
+        qspec_test=copy.deepcopy(qspec),
+        quantization_mode=quantization_mode,
+        custom_qdq_cases=custom_qdq_cases,
+        expected_qdq_insertion_user=expected_qdq_insertion
+    )
+
+    v = ONNXQDQValidator()
+    validated = v.validate(
+        test_assets_attr.int8_onnx_model, expected_qdq_insertion, original_keras_model=nn_model_original
+    )
+    return q_model, validated

tools/tensorflow-quantization/tests/quantize_config_test.py

@@ -0,0 +1,69 @@
+#
+# 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 tensorflow_quantization import quantize_config
+import tensorflow_quantization.global_config as global_config
+from tensorflow_quantization import QuantizationSpec
+from network_pool import bilbo_28_28
+
+
+def test_global_object_creation():
+    fnq = quantize_config.FullNetworkQuantization()
+    assert (
+            len(global_config.G_CONFIG_OBJECT) == 1
+    ), "quantization config class object is not added to the global list"
+    assert isinstance(
+        global_config.G_CONFIG_OBJECT[0], quantize_config.FullNetworkQuantization
+    )
+    fnq.clean()
+    tf.keras.backend.clear_session()
+
+
+def test_quantization_config_layer_names_add():
+    model = bilbo_28_28()
+    fnq = quantize_config.FullNetworkQuantization()
+    qspec = QuantizationSpec()
+    qspec.add(name="conv_0")
+    qspec.add(name="conv_2")
+    qspec.add(name="conv_4")
+    fnq.add_quantization_spec_object(qspec, model.layers)
+    assert (
+        "conv_0" in fnq.layerwise_config
+    ), "There seems to be an issue with layer name addition in `add_special_layers` function"
+    assert (
+        "conv_2" in fnq.layerwise_config
+    ), "There seems to be an issue with layer name addition in `add_special_layers` function"
+    assert (
+        "conv_4" in fnq.layerwise_config
+    ), "There seems to be an issue with layer name addition in `add_special_layers` function"
+    fnq.clean()
+    tf.keras.backend.clear_session()
+
+
+def test_quantization_config_layer_class_add():
+    model = bilbo_28_28()
+    fnq = quantize_config.FullNetworkQuantization()
+    qspec = QuantizationSpec()
+    qspec.add(name="Dense", is_keras_class=True)
+    fnq.add_quantization_spec_object(qspec, model.layers)
+    assert (
+        "Dense" in fnq.layer_classes_to_quantize
+    ), "There seems to be an issue with layer class name addition in `add_special_layers` function"
+    fnq.clean()
+    tf.keras.backend.clear_session()

tools/tensorflow-quantization/tests/quantize_qdq_insertion_test.py

@@ -0,0 +1,158 @@
+#
+# 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 import QuantizationSpec
+from tensorflow_quantization.custom_qdq_cases import ResNetV1QDQCase
+from network_pool import frodo_32_32
+from onnx_graph_qdq_validator import validate_quantized_model
+from tensorflow_quantization.utils import CreateAssetsFolders
+import pytest
+
+test_assets = CreateAssetsFolders("test_quantize_qdq_insertion")
+
+# ############################################
+# ######### Full Quantize Test ###############
+# ############################################
+
+
+def test_quantize_full():
+    this_function_name = sys._getframe().f_code.co_name
+
+    nn_model_original = frodo_32_32()
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model_original, test_name=this_function_name
+    )
+
+    assert vr, "ONNX QDQ Validation for full network quantization failed!"
+    # Necessary to clear model layer names from the memory
+    tf.keras.backend.clear_session()
+
+
+def test_quantize_full_residual():
+    this_function_name = sys._getframe().f_code.co_name
+
+    nn_model_original = frodo_32_32()
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model_original, custom_qdq_cases=[ResNetV1QDQCase()], test_name=this_function_name
+    )
+
+    assert vr, "ONNX QDQ Validation for quantizing full network with special residual failed!"
+    tf.keras.backend.clear_session()
+
+
+# ############################################
+# ######### Full Special Quantize Test #######
+# ############################################
+
+
+def test_quantize_full_special_layer():
+    this_function_name = sys._getframe().f_code.co_name
+
+    nn_model_original = frodo_32_32()
+
+    # Create a Quantization Spec (dictionary telling how `add` layer should be treated differently).
+    qspec = QuantizationSpec()
+    qspec.add(name="add", quantization_index=[0])
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model_original, qspec=qspec, test_name=this_function_name
+    )
+
+    assert vr, "QDQ Validation for full network but one special layer quantization failed!"
+    tf.keras.backend.clear_session()
+
+
+# ##########################################
+# ######### Partial Quantize Test ##########
+# ##########################################
+
+
+def test_quantize_partial():
+    this_function_name = sys._getframe().f_code.co_name
+
+    nn_model_original = frodo_32_32()
+
+    # Create a qspec dictionary to quantize only two layers named 'conv2d_2' and 'dense'
+    qspec = QuantizationSpec()
+    qspec.add(name="conv2d_2")
+    qspec.add(name="dense")
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model_original, quantization_mode="partial", qspec=qspec, test_name=this_function_name
+    )
+
+    assert vr, "ONNX QDQ Validation for partial network quantization failed!"
+    tf.keras.backend.clear_session()
+
+
+# ####################################################
+# ######### Subset layers Test - Full quantize #######
+# ####################################################
+
+
+def test_quantize_specific_class_maxpool():
+    this_function_name = sys._getframe().f_code.co_name
+
+    nn_model_original = frodo_32_32()
+
+    # Create a list with keras layer classes to quantize
+    qspec = QuantizationSpec()
+    qspec.add(name="MaxPooling2D", is_keras_class=True)
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model_original, qspec=qspec, test_name=this_function_name
+    )
+
+    assert vr, "ONNX QDQ Validation for specific class `Dense` quantization failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_quantize_specific_class_add():
+    this_function_name = sys._getframe().f_code.co_name
+
+    nn_model_original = frodo_32_32()
+
+    # Create a list with keras layer classes to quantize
+    qspec = QuantizationSpec()
+    qspec.add(name="Add", is_keras_class=True)
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model_original, qspec=qspec, test_name=this_function_name
+    )
+
+    assert vr, "ONNX QDQ Validation for quantizing specific class `Add` failed!"
+    tf.keras.backend.clear_session()
+
+
+# ####################################################
+# ####### Subset layers Test - Partial quantize ######
+# ####################################################
+
+
+def test_quantize_specific_class_conv2d_partial():
+    this_function_name = sys._getframe().f_code.co_name
+
+    nn_model_original = frodo_32_32()
+
+    # Create a list with keras layer classes to quantize
+    qspec = QuantizationSpec()
+    qspec.add(name="Conv2D", is_keras_class=True)
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model_original, quantization_mode="partial", qspec=qspec, test_name=this_function_name
+    )
+
+    assert vr, "ONNX QDQ Validation for quantizing specific class `Conv2D` and `conv2d_1` layer failed!"
+    tf.keras.backend.clear_session()

tools/tensorflow-quantization/tests/quantize_test.py

@@ -0,0 +1,206 @@
+#
+# 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 `quantize_model` feature.
+`quantize_model` feature quantizes all supported layers in the given Keras model with `NVIDIA` quantization scheme.
+
+Tests if weights were copied correctly after quantization and end-to-end training accuracy.
+"""
+
+import tensorflow as tf
+from tensorflow_quantization import quantize
+from tensorflow_quantization import quantize_model
+from network_pool import lobelia_28_28
+from network_pool import bilbo_28_28
+import pytest
+import tensorflow_quantization
+from tensorflow_quantization.utils import (
+    CreateAssetsFolders,
+    convert_saved_model_to_onnx,
+)
+
+
+def _print_model_weights_shapes(model):
+    """
+    Print shapes of all weights
+    Args:
+        model: Keras model
+    """
+    print([model.get_weights()[i].shape for i in range(len(model.get_weights()))])
+
+
+def test_clone_numerics_quantize_whole_model(debug=False):
+    """
+    Checks whether weights are copied correctly when a dummy model is quantized.
+    """
+    model = lobelia_28_28()
+    if debug:
+        _print_model_weights_shapes(model)
+    om_l0_test_weights = model.get_weights()[0][10, :5]
+    om_l1_test_weights = model.get_weights()[2][10, :5]
+
+    # Quantize model
+    q_model = quantize_model(model)
+
+    if debug:
+        _print_model_weights_shapes(q_model)
+    qm_l0_test_weights = q_model.get_weights()[1][10, :5]
+    qm_l1_test_weights = q_model.get_weights()[8][10, :5]
+    assert all([a == b for a, b in zip(om_l0_test_weights, qm_l0_test_weights)])
+    assert all([a == b for a, b in zip(om_l1_test_weights, qm_l1_test_weights)])
+    tf.keras.backend.clear_session()
+
+
+def test_adding_one_layer_at_a_time():
+    qspec = quantize.QuantizationSpec()
+    qspec.add(name="conv2d_1")
+    qspec.add(name="Dense", is_keras_class=True)
+
+    assert isinstance(
+        qspec.layers[0], quantize.LayerConfig
+    ), "LayerConfig object is not created for newly added layer."
+    assert (
+        len(qspec.layers) == 2
+    ), "New layers are not added to layer list of QuantizationSpec."
+
+
+def test_adding_layer_name_list():
+    qspec = quantize.QuantizationSpec()
+    layer_name = ["conv2d", "conv2d_1", "conv2d_7", "dense"]
+    layer_qip = [True, False, True, False]
+    layer_idx = [None, [0], None, None]
+    qspec.add(name=layer_name, quantize_input=layer_qip, quantization_index=layer_idx)
+    assert (
+        len(qspec.layers) == 4
+    ), "Four layers are not added to qspec object as expected."
+
+
+def train_quantize_fine_tune(exp_folder: "Folder", perform_four_bit_quantization: bool = False) -> None:
+    """
+    Train, quantize and fine-tune Keras model using NVIDIA's QAT wrapper library.
+
+    Args:
+        exp_folder (Folder): Base experiment folder object.
+        perform_four_bit_quantization (bool): If True, 4 bit quantization is performed. 8 bit quantization is default.
+
+    Returns:
+        None
+    """
+
+    # Load MNIST dataset
+    mnist = tf.keras.datasets.fashion_mnist
+    (train_images, train_labels), (test_images, test_labels) = mnist.load_data()
+
+    # Normalize the input image so that each pixel value is between 0 to 1.
+    train_images = train_images / 255.0
+    test_images = test_images / 255.0
+
+    nn_model_original = bilbo_28_28()
+    # Train original classification model
+    nn_model_original.compile(
+        optimizer="adam",
+        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        metrics=["accuracy"],
+    )
+
+    nn_model_original.fit(
+        train_images, train_labels, batch_size=128, epochs=5, validation_split=0.1
+    )
+
+    # get baseline model accuracy
+    _, baseline_model_accuracy = nn_model_original.evaluate(
+        test_images, test_labels, verbose=0
+    )
+    print("Baseline test accuracy:", baseline_model_accuracy)
+
+    tf.keras.models.save_model(nn_model_original, exp_folder.fp32_saved_model)
+    convert_saved_model_to_onnx(
+        saved_model_dir=exp_folder.fp32_saved_model,
+        onnx_model_path=exp_folder.fp32_onnx_model,
+    )
+
+    if perform_four_bit_quantization:
+        tensorflow_quantization.G_NUM_BITS = 4
+    
+    # quantize entire model using `quantize_model` feature
+    q_model = quantize_model(nn_model_original)
+
+    # fine tune annotated model
+    q_model.compile(
+        optimizer="adam",
+        loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True),
+        metrics=["accuracy"],
+    )
+    q_model.fit(
+        train_images, train_labels, batch_size=32, epochs=5, validation_split=0.1
+    )
+
+    # Get quantized accuracy
+    _, q_aware_model_accuracy = q_model.evaluate(test_images, test_labels, verbose=0)
+    print("Quant test accuracy:", q_aware_model_accuracy)
+
+    assert (
+        q_aware_model_accuracy >= baseline_model_accuracy or
+        abs(baseline_model_accuracy - q_aware_model_accuracy) * 100 <= 2.0
+    ), "QAT accuracy is not acceptable: {:.2f} vs {:.2f} for baseline".format(
+        q_aware_model_accuracy * 100, baseline_model_accuracy * 100
+    )
+
+    # save quantized model and convert to ONNX
+    tf.keras.models.save_model(q_model, exp_folder.int8_saved_model)
+    convert_saved_model_to_onnx(
+        saved_model_dir=exp_folder.int8_saved_model,
+        onnx_model_path=exp_folder.int8_onnx_model,
+    )
+
+
+def test_end_to_end_workflow():
+    """
+    Test end-to-end QAT workflow using the `quantize_model` function.
+    The following steps are included:
+        1. Create a dummy model (baseline)
+        2. Train model on Fashion MNIST dataset
+        3. Calculate baseline FP32 model accuracy
+        4. Perform 4 bit (default) quantization and fine-tuning
+        5. Convert QAT model to ONNX
+    """
+
+    test_assets = CreateAssetsFolders("test_quantize_end_to_end")
+    test_assets.add_folder("test_end_to_end_workflow")
+
+    train_quantize_fine_tune(test_assets.test_end_to_end_workflow)
+    tf.keras.backend.clear_session()
+
+@pytest.mark.skip(reason="Just used to test 4 bit quantization feature.")
+def test_end_to_end_workflow_4bit():
+    """
+    Test end-to-end QAT workflow using the `quantize_model` function for 4 bit quantization.
+    The following steps are included:
+        1. Create a dummy model (baseline)
+        2. Train model on Fashion MNIST dataset
+        3. Calculate baseline FP32 model accuracy
+        4. Perform 4 bit quantization and fine-tuning
+        5. Convert QAT model to ONNX
+    """
+
+    test_assets = CreateAssetsFolders("test_quantize_end_to_end")
+    test_assets.add_folder("test_end_to_end_workflow_4bit")
+
+    train_quantize_fine_tune(test_assets.test_end_to_end_workflow_4bit, perform_four_bit_quantization=True)
+    tf.keras.backend.clear_session()

tools/tensorflow-quantization/tests/quantize_wrapper_base_test.py

@@ -0,0 +1,55 @@
+#
+# 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.
+#
+
+
+from tensorflow_quantization.quantize_wrapper_base import BaseQuantizeWrapper
+import copy
+import pytest
+
+
+EXPECTED_WRAPPERS = [
+    "WeightedBaseQuantizeWrapper",
+    "Conv2DQuantizeWrapper",
+    "DenseQuantizeWrapper",
+    "DepthwiseConv2DQuantizeWrapper",
+    "NonWeightedBaseQuantizeWrapper",
+    "AveragePooling2DQuantizeWrapper",
+    "GlobalAveragePooling2DQuantizeWrapper",
+    "MaxPooling2DQuantizeWrapper",
+    "BatchNormalizationQuantizeWrapper",
+    "NonWeightedBaseQuantizeWrapperForMultipleInputs",
+    "MultiplyQuantizeWrapper",
+    "ConcatenateQuantizeWrapper",
+    "AddQuantizeWrapper",
+]
+
+
+def test_old_wrappers_registration():
+    all_wrappers = BaseQuantizeWrapper.CHILD_WRAPPERS
+    assert EXPECTED_WRAPPERS == list(all_wrappers.keys())
+
+
+def test_new_wrapper_registration():
+    class TestWrapper(BaseQuantizeWrapper):
+        def __init__(self, layer, **kwargs):
+            super().__init__(layer, **kwargs)
+
+    all_wrappers = BaseQuantizeWrapper.CHILD_WRAPPERS
+    expected = copy.deepcopy(EXPECTED_WRAPPERS)
+    expected.append("TestWrapper")
+
+    assert expected == list(all_wrappers.keys())

tools/tensorflow-quantization/tests/quantize_wrappers_test.py

@@ -0,0 +1,632 @@
+#
+# 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 import QuantizationSpec
+from tensorflow_quantization.quantize import LayerConfig
+from onnx_graph_qdq_validator import validate_quantized_model
+from tensorflow_quantization.utils import CreateAssetsFolders
+from network_pool import (
+    otho_28_28,
+    lotho_28_28,
+    lobelia_28_28,
+    merry_28_28,
+    pippin_28_28,
+)
+import pytest
+
+# Create a directory to save wrapper test data
+test_assets = CreateAssetsFolders("test_quantize_wrappers")
+
+
+# ###################################################
+# ####### Conv2D layer wrapper tests ################
+# ###################################################
+
+
+def test_conv2d_wrapper_quant_full_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = otho_28_28()
+
+    # Quantization
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="conv_0"),
+        LayerConfig(name="conv_1"),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_conv2d_wrapper_quant_partial_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = otho_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(name="conv_1")
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="conv_1"),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_conv2d_wrapper_quant_partial_only_input_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = otho_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(name="conv_0", quantize_weight=False)
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="conv_0", quantize_weight=False),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_conv2d_wrapper_quant_partial_only_weight_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = otho_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(name="conv_0", quantize_input=False)
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="conv_0", quantize_input=False),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+# ###################################################
+# ####### DepthwiseConv2D layer wrapper tests #######
+# ###################################################
+
+
+def test_depthwise_conv2d_wrapper_quant_full_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = lotho_28_28()
+
+    # Quantization
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="dconv_0"),
+        LayerConfig(name="dconv_1"),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_depthwise_conv2d_wrapper_quant_partial_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = lotho_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(name="dconv_1")
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="dconv_1"),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_depthwise_conv2d_wrapper_quant_partial_only_input_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = lotho_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(name="dconv_1", quantize_weight=False)
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="dconv_1", quantize_weight=False),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_depthwise_conv2d_wrapper_quant_partial_only_weight_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = lotho_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(name="dconv_1", quantize_input=False)
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="dconv_1", quantize_input=False),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+# ###################################################
+# ####### Dense layer wrapper tests #################
+# ###################################################
+
+
+def test_dense_wrapper_quant_full_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = lobelia_28_28()
+
+    # Quantization
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="dense_0"),
+        LayerConfig(name="dense_1"),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_dense_wrapper_quant_partial_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = lobelia_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(name="dense_0")
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="dense_0"),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_dense_wrapper_quant_partial_only_input_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = lobelia_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(name="dense_0", quantize_weight=False)
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="dense_0", quantize_weight=False),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_dense_wrapper_quant_partial_only_weight_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = lobelia_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(name="dense_1", quantize_input=False)
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="dense_1", quantize_input=False),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+
+    tf.keras.backend.clear_session()
+
+
+# ###################################################
+# ####### Concatenation layer wrapper tests #########
+# ###################################################
+
+
+def test_concat_wrapper_quant_full_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = merry_28_28()
+
+    # Quantization
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="conv2d"),
+        LayerConfig(name="conv2d_1"),
+        LayerConfig(name="conv2d_3"),
+        LayerConfig(name="conv2d_4"),
+        LayerConfig(name="conv2d_2"),
+        LayerConfig(name="dense"),
+        LayerConfig(name="dense_1"),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_concat_wrapper_quant_full_quant_bn_concat_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = merry_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(name="batch_normalization_3")
+    qspec.add(name="concatenate", quantization_index=[0, 1])
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="conv2d"),
+        LayerConfig(name="conv2d_1"),
+        LayerConfig(name="conv2d_3"),
+        LayerConfig(name="conv2d_4"),
+        LayerConfig(name="batch_normalization_3"),
+        LayerConfig(name="conv2d_2"),
+        LayerConfig(name="concatenate", quantization_index=[0, 1]),
+        LayerConfig(name="dense"),
+        LayerConfig(name="dense_1"),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_concat_wrapper_quant_specific_index_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = merry_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(
+        name="concatenate",
+        quantize_input=True,
+        quantize_weight=False,
+        quantization_index=[0],
+    )
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="concatenate", quantization_index=[0]),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+# ###################################################
+# ####### Add layer wrapper tests ###################
+# ###################################################
+
+
+# Use KerasModelLayersSurgeon() from utils to find layer names.
+def test_add_wrapper_quant_full_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = pippin_28_28()
+
+    # Quantization
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="conv2d"),
+        LayerConfig(name="conv2d_1"),
+        LayerConfig(name="conv2d_2"),
+        LayerConfig(name="conv2d_3"),
+        LayerConfig(name="conv2d_4"),
+        LayerConfig(name="conv2d_6"),
+        LayerConfig(name="conv2d_5"),
+        LayerConfig(name="dense"),
+        LayerConfig(name="dense_1"),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_add_wrapper_quant_partial_specific_index_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = pippin_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(
+        name="add", quantize_input=True, quantize_weight=False, quantization_index=[1]
+    )
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="add", quantization_index=[1])
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode="partial", qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+def test_add_wrapper_quant_full_specific_index_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = pippin_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(
+        name="add", quantize_input=True, quantize_weight=False, quantization_index=[1]
+    )
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="conv2d"),
+        LayerConfig(name="conv2d_1"),
+        LayerConfig(name="conv2d_2"),
+        LayerConfig(name="conv2d_3"),
+        LayerConfig(name="add", quantization_index=[1]),
+        LayerConfig(name="conv2d_4"),
+        LayerConfig(name="conv2d_6"),
+        LayerConfig(name="conv2d_5"),
+        LayerConfig(name="dense"),
+        LayerConfig(name="dense_1"),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+# ########################################################
+# ############ Test subset layer class selection #########
+# ########################################################
+
+
+def test_subset_layer_class_selection_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = pippin_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(name="Conv2D", is_keras_class=True)
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="conv2d"),
+        LayerConfig(name="conv2d_1"),
+        LayerConfig(name="conv2d_2"),
+        LayerConfig(name="conv2d_3"),
+        LayerConfig(name="conv2d_4"),
+        LayerConfig(name="conv2d_6"),
+        LayerConfig(name="conv2d_5"),
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode='partial', qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+# ########################################################
+# ############ Test missing layer name warning ###########
+# ########################################################
+
+
+def test_missing_layer_name_warning_nv():
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = pippin_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(
+        name="add", quantize_input=True, quantize_weight=False, quantization_index=[1]
+    )
+    qspec.add(name="wrong_layer")
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="add", quantization_index=[1])
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode='partial', qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()
+
+
+# ########################################################
+# ##### Test Add,Concat out of range index warning #######
+# ########################################################
+
+
+@pytest.mark.skip(
+    reason="When quantization index out of range does not give error but still wraps \
+add layer without quantizing any input"
+)
+def test_out_of_range_index():
+
+    # Create experiment specific directory
+    this_function_name = sys._getframe().f_code.co_name
+
+    # Baseline model
+    nn_model = pippin_28_28()
+
+    # Quantization
+    qspec = QuantizationSpec()
+    qspec.add(
+        name="add", quantize_input=True, quantize_weight=False, quantization_index=[3]
+    )
+    qspec.add(name="dense")
+    # (Optional) QDQ node placement check
+    #   Here just to explicitly show the user which layers are quantized.
+    expected_qdq_insertion = [
+        LayerConfig(name="add"), LayerConfig(name="dense")
+    ]
+    q_model, vr = validate_quantized_model(
+        test_assets, nn_model, quantization_mode='partial', qspec=qspec, test_name=this_function_name,
+        expected_qdq_insertion=expected_qdq_insertion
+    )
+
+    assert vr, "ONNX QDQ Validation failed!"
+    tf.keras.backend.clear_session()

tools/tensorflow-quantization/tests/run_all_tests.sh

@@ -0,0 +1,30 @@
+#!/bin/bash
+
+# clean
+rm -rf wrappers_test_saved_models
+rm -rf quantize_model_test_saved_models
+rm -rf utils_test_saved_models
+rm -rf qdq_test_saved_models
+rm -rf __pycache__
+rm -rf custom_qdq_models
+
+# Run quantize_config tests
+python -m pytest quantize_config_test.py -rP
+
+# Run QDQ insertion tests
+python -m pytest quantize_qdq_insertion_test.py -rP
+
+# Run wrappers tests
+python -m pytest quantize_wrappers_test.py -rP
+
+# Run wrappers base tests
+python -m pytest quantize_wrapper_base_test.py -rP
+
+# Run end to end training test
+python -m pytest quantize_test.py -rP
+
+# Run special qdq insertion tests
+python -m pytest custom_qdq_cases_test.py -rP
+
+# Run utils test
+python -m pytest utils_test.py -rP

tools/tensorflow-quantization/tests/utils_test.py

@@ -0,0 +1,101 @@
+#
+# 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_quantization.utils as utils
+import tensorflow as tf
+from tensorflow_quantization import quantize_model
+from tensorflow_quantization.utils import (
+    CreateAssetsFolders,
+    convert_saved_model_to_onnx,
+)
+from network_pool import sam_32_32
+import pytest
+
+test_assets = CreateAssetsFolders("test_utils")
+
+
+def test_keras_traveller():
+    kmt = utils.KerasModelTraveller()
+    model = sam_32_32()
+    layer_names = kmt.get_layer_names(keras_model=model)
+    expected_layer_names = [
+        "input_1",
+        "conv2d",
+        "re_lu",
+        "conv2d_1",
+        "re_lu_1",
+        "conv2d_2",
+        "re_lu_2",
+        "conv2d_3",
+        "add",
+        "re_lu_3",
+        "conv2d_4",
+        "re_lu_4",
+        "conv2d_5",
+        "add_1",
+        "re_lu_5",
+        "conv2d_6",
+        "re_lu_6",
+        "conv2d_7",
+        "conv2d_8",
+        "add_2",
+        "re_lu_7",
+        "max_pooling2d",
+        "flatten",
+        "dense",
+        "re_lu_8",
+        "dense_1",
+    ]
+    assert layer_names == expected_layer_names, "Keras model traveller failed."
+    tf.keras.backend.clear_session()
+
+
+def test_convert_to_onnx():
+    test_assets.add_folder("test_convert_to_onnx")
+
+    model = sam_32_32()
+    q_model = quantize_model(model)
+    # Create experiment specific directory
+    tf.keras.models.save_model(
+        q_model, test_assets.test_convert_to_onnx.int8_saved_model
+    )
+    convert_saved_model_to_onnx(
+        saved_model_dir=test_assets.test_convert_to_onnx.int8_saved_model,
+        onnx_model_path=test_assets.test_convert_to_onnx.int8_onnx_model,
+    )
+    tf.keras.backend.clear_session()
+
+
+def test_find_my_predecessors():
+    resnet50 = tf.keras.applications.resnet.ResNet50(weights=None)
+    r = utils.find_my_predecessors(resnet50, "conv2_block1_add")
+    assert r[0]["class"] == "BatchNormalization"
+    assert r[0]["name"] == "conv2_block1_0_bn"
+    assert r[1]["class"] == "BatchNormalization"
+    assert r[1]["name"] == "conv2_block1_3_bn"
+
+
+def test_find_my_successors():
+    resnet50 = tf.keras.applications.resnet.ResNet50(weights=None)
+    r = utils.find_my_successors(resnet50, "pool1_pool")
+    assert r[0]["class"] == "Conv2D"
+    assert r[0]["name"] == "conv2_block1_1_conv"
+    assert r[1]["class"] == "Conv2D"
+    assert r[1]["name"] == "conv2_block1_0_conv"