chore: import upstream snapshot with attribution

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: