chore: import upstream snapshot with attribution

tools/pytorch-quantization/docs/source/_templates/layout.html

@@ -0,0 +1,9 @@
+{% extends '!layout.html' %}
+
+{% block extrahead %}
+<script src="//assets.adobedtm.com/5d4962a43b79/c1061d2c5e7b/launch-191c2462b890.min.js"></script>
+{% endblock %}
+
+{% block footer %}
+<script type="text/javascript">_satellite.pageBottom();</script>
+{% endblock %}

tools/pytorch-quantization/docs/source/calib.rst

@@ -0,0 +1,22 @@
+.. role:: hidden
+    :class: hidden-section
+
+pytorch_quantization.calib
+===================================
+
+.. automodule:: pytorch_quantization.calib
+.. currentmodule:: pytorch_quantization.calib
+
+:hidden:`MaxCalibrator`
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: MaxCalibrator
+    :members:
+    :inherited-members:
+
+:hidden:`HistogramCalibrator`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: HistogramCalibrator
+    :members:
+    :inherited-members:

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

@@ -0,0 +1,228 @@
+#
+# 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.
+#
+
+
+# -*- 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
+from pytorch_quantization import __version__
+
+# -- Project information -----------------------------------------------------
+
+project = 'pytorch-quantization'
+copyright = '2021, NVIDIA'
+
+# The short X.Y version
+version = __version__
+# The full version, including alpha/beta/rc tags
+release = 'master'
+
+author = 'NVIDIA'
+
+
+# -- General configuration ---------------------------------------------------
+
+# If your documentation needs a minimal Sphinx version, state it here.
+#
+needs_sphinx = '1.8'
+
+# 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.autodoc',
+    'sphinx.ext.autosummary',
+    'sphinx.ext.doctest',
+    'sphinx.ext.intersphinx',
+    'sphinx.ext.todo',
+    'sphinx.ext.coverage',
+    'sphinx.ext.napoleon',
+    'sphinx.ext.viewcode',
+]
+
+# Napoleon settings
+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
+
+
+
+# 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', '.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 = None
+
+# 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 = ['_build', 'Thumbs.db', '.DS_Store']
+
+# The name of the Pygments (syntax highlighting) style to use.
+pygments_style = 'sphinx'
+
+
+# -- 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 = 'sphinx_rtd_theme'
+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.
+#
+
+# 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 = 'pytorch-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, 'pytorch-quantization.tex', 'pytorch-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, 'pytorch-quantization', 'pytorch-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, 'pytorch-quantization', 'pytorch-quantization Documentation',
+     author, 'pytorch-quantization', 'One line description of project.',
+     '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 -------------------------------------------------
+# Disable docstring inheritance
+autodoc_inherit_docstrings = False
+
+# Example configuration for intersphinx: refer to the Python standard library.
+intersphinx_mapping = {
+    'python': ('https://docs.python.org/', None),
+    'numpy': ('https://docs.scipy.org/doc/numpy/', None),
+}

tools/pytorch-quantization/docs/source/functional.rst

@@ -0,0 +1,17 @@
+.. role:: hidden
+    :class: hidden-section
+
+pytorch_quantization.nn.functional
+===================================
+
+.. automodule:: pytorch_quantization.nn.functional
+.. currentmodule:: pytorch_quantization.nn.functional
+
+:hidden:`ClipFunction`
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: ClipFunction
+    :members:
+
+
+``clip`` is alias of ``ClipFunction.apply``

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

@@ -0,0 +1,33 @@
+pytorch-quantization's documentation
+================================================
+
+.. toctree::
+   :maxdepth: 1
+   :caption: User Guide
+
+   userguide
+
+.. toctree::
+   :caption: Tutorials
+   :titlesonly:
+   :maxdepth: 2
+
+   tutorials/quant_resnet50
+   tutorials/creating_custom_quantized_modules
+
+
+.. toctree::
+   :maxdepth: 1
+   :caption: Package Reference
+
+   calib
+   nn
+   functional
+   optim
+   tensor_quant
+   utils
+
+Indices
+==================
+
+* :ref:`genindex`

tools/pytorch-quantization/docs/source/nn.rst

@@ -0,0 +1,140 @@
+.. role:: hidden
+    :class: hidden-section
+
+pytorch_quantization.nn
+===================================
+
+.. automodule:: pytorch_quantization.nn
+.. currentmodule:: pytorch_quantization.nn
+
+:hidden:`TensorQuantizer`
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: TensorQuantizer
+    :members:
+
+Quantized Modules
+----------------------------------
+.. currentmodule:: pytorch_quantization.nn.modules.quant_conv
+
+_QuantConvNd
+~~~~~~~~~~~~
+.. autoclass:: _QuantConvNd
+    :members:
+
+.. currentmodule:: pytorch_quantization.nn
+
+:hidden:`QuantConv1d`
+~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantConv1d
+    :members:
+
+:hidden:`QuantConv2d`
+~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantConv2d
+    :members:
+
+:hidden:`QuantConv3d`
+~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantConv3d
+    :members:
+
+:hidden:`QuantConvTranspose1d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantConvTranspose1d
+    :members:
+
+:hidden:`QuantConvTranspose2d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+
+.. autoclass:: QuantConvTranspose2d
+    :members:
+
+:hidden:`QuantConvTranspose3d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantConvTranspose3d
+    :members:
+
+:hidden:`QuantLinear`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantLinear
+    :members:
+
+:hidden:`QuantMaxPool1d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantMaxPool1d
+    :members:
+
+:hidden:`QuantMaxPool2d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantMaxPool2d
+    :members:
+
+:hidden:`QuantMaxPool3d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantMaxPool3d
+    :members:
+
+:hidden:`QuantAvgPool1d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantAvgPool1d
+    :members:
+
+:hidden:`QuantAvgPool2d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantAvgPool2d
+    :members:
+
+:hidden:`QuantAvgPool3d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantAvgPool3d
+    :members:
+
+:hidden:`QuantAdaptiveAvgPool1d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantAdaptiveAvgPool1d
+    :members:
+
+:hidden:`QuantAdaptiveAvgPool2d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantAdaptiveAvgPool2d
+    :members:
+
+:hidden:`QuantAdaptiveAvgPool3d`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantAdaptiveAvgPool3d
+    :members:
+
+:hidden:`Clip`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: Clip
+    :members:
+
+:hidden:`QuantLSTM`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantLSTM
+    :members:
+
+:hidden:`QuantLSTMCell`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantLSTMCell
+    :members:

tools/pytorch-quantization/docs/source/optim.rst

@@ -0,0 +1,10 @@
+.. role:: hidden
+    :class: hidden-section
+
+pytorch_quantization.optim.helper
+===================================
+
+.. automodule:: pytorch_quantization.optim.helper
+    :members:
+    :undoc-members:
+    :show-inheritance:

tools/pytorch-quantization/docs/source/tensor_quant.rst

@@ -0,0 +1,31 @@
+.. role:: hidden
+    :class: hidden-section
+
+pytorch_quantization.tensor_quant
+===================================
+
+.. automodule:: pytorch_quantization.tensor_quant
+.. currentmodule:: pytorch_quantization.tensor_quant
+
+
+:hidden:`QuantDescriptor`
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: QuantDescriptor
+    :members:
+
+:hidden:`ScaledQuantDescriptor`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: ScaledQuantDescriptor
+    :members:
+
+:hidden:`TensorQuantFunction`
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: TensorQuantFunction
+    :members:
+
+``tensor_quant`` is alias of ``TensorQuantFunction.apply``
+
+``fake_tensor_quant`` is alias of ``FakeTensorQuantFunction.apply``

tools/pytorch-quantization/docs/source/tutorials/creating_custom_quantized_modules.rst

@@ -0,0 +1,169 @@
+Creating Custom Quantized Modules
+=================================
+
+There are several quantized modules provided by the quantization tool as follows:
+
+- QuantConv1d, QuantConv2d, QuantConv3d, QuantConvTranspose1d, QuantConvTranspose2d, QuantConvTranspose3d
+- QuantLinear
+- QuantAvgPool1d, QuantAvgPool2d, QuantAvgPool3d, QuantMaxPool1d, QuantMaxPool2d, QuantMaxPool3d
+
+To quantize a module, we need to quantize the input and weights if present. Following are 3 major use-cases:
+
+#. Create quantized wrapper for modules that have only inputs
+#. Create quantized wrapper for modules that have inputs as well as weights.
+#. Directly add the ``TensorQuantizer`` module to the inputs of an operation in the model graph.
+
+The first two methods are very useful if it's needed to automatically replace the original modules (nodes in the graph) with their quantized versions. The third method could be useful when it's required to manually add the quantization to the model graph at very specific places (more manual, more control).
+
+Let's see each use-case with examples below.
+
+Quantizing Modules With Only Inputs
+-----------------------------------
+A suitable example would be quantizing the ``pooling`` module variants.
+
+Essentially, we need to provide a wrapper function that takes the original module and adds the ``TensorQuantizer`` module around it so that the input is first quantized and then fed into the original module. 
+
+- Create the wrapper by subclassing the original module (``pooling.MaxPool2d``) along with the utilities module (``_utils.QuantInputMixin``).
+
+.. code:: python
+
+    class QuantMaxPool2d(pooling.MaxPool2d, _utils.QuantInputMixin):
+
+- The ``__init__.py`` function would call the original module's init function and provide it with the corresponding arguments. There would be just one additional argument using ``**kwargs`` which contains the quantization configuration information. The ``QuantInputMixin`` utility contains the method ``pop_quant_desc_in_kwargs`` which extracts this configuration information from the input or returns a default if that input is ``None``. Finally the ``init_quantizer`` method is called that initializes the ``TensorQuantizer`` module which would quantize the inputs.
+
+.. code:: python
+
+    def __init__(self, kernel_size, stride=None, padding=0, dilation=1,
+                 return_indices=False, ceil_mode=False, **kwargs):
+        super(QuantMaxPool2d, self).__init__(kernel_size, stride, padding, dilation,
+                                             return_indices, ceil_mode)
+        quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+        self.init_quantizer(quant_desc_input)
+
+- After the initialization, the ``forward`` function needs to be defined in our wrapper module that would actually quantize the inputs using the ``_input_quantizer`` that was initialized in the ``__init__`` function forwarding the inputs to the base module using ``super`` call.
+
+.. code:: python
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        return super(QuantMaxPool2d, self).forward(quant_input)
+
+- Finally, we need to define a getter method for the ``_input_quantizer``. This could, for example, be used  to disable the quantization for a particular module using ``module.input_quantizer.disable()`` which is helpful while experimenting with different layer quantization configuration.
+
+.. code:: python
+
+    @property
+    def input_quantizer(self):
+        return self._input_quantizer
+
+A complete quantized pooling module would look like following:
+
+.. code:: python
+
+    class QuantMaxPool2d(pooling.MaxPool2d, _utils.QuantInputMixin):
+        """Quantized 2D maxpool"""
+        def __init__(self, kernel_size, stride=None, padding=0, dilation=1,
+                    return_indices=False, ceil_mode=False, **kwargs):
+            super(QuantMaxPool2d, self).__init__(kernel_size, stride, padding, dilation,
+                                                return_indices, ceil_mode)
+            quant_desc_input = _utils.pop_quant_desc_in_kwargs(self.__class__, input_only=True, **kwargs)
+            self.init_quantizer(quant_desc_input)
+
+        def forward(self, input):
+            quant_input = self._input_quantizer(input)
+            return super(QuantMaxPool2d, self).forward(quant_input)
+
+        @property
+        def input_quantizer(self):
+            return self._input_quantizer
+
+Quantizing Modules With Weights and Inputs
+------------------------------------------
+We give an example of quantizing the ``torch.nn.Linear`` module. It follows that the only additional change from the previous example of quantizing pooling modules is that we'd need to accomodate the quantization of weights in the Linear module. 
+
+- We create the quantized linear module as follows:
+
+.. code:: python
+
+    class QuantLinear(nn.Linear, _utils.QuantMixin):
+
+- In the ``__init__`` function, we first use the ``pop_quant_desc_in_kwargs`` function to extract the quantization descriptors for both inputs and weights. Second, we initialize the ``TensorQuantizer`` modules for both inputs and weights using these quantization descriptors.
+
+.. code:: python
+
+    def __init__(self, in_features, out_features, bias=True, **kwargs):
+            super(QuantLinear, self).__init__(in_features, out_features, bias)
+            quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+
+            self.init_quantizer(quant_desc_input, quant_desc_weight)
+
+- Also, override the ``forward`` function call and pass the inputs and weights through ``_input_quantizer`` and ``_weight_quantizer`` respectively before passing the quantized arguments to the actual ``F.Linear`` call. This step adds the actual input/weight ``TensorQuantizer`` to the module and eventually the model.
+
+.. code:: python
+
+    def forward(self, input):
+        quant_input = self._input_quantizer(input)
+        quant_weight = self._weight_quantizer(self.weight)
+
+        output = F.linear(quant_input, quant_weight, bias=self.bias)
+
+        return output
+
+- Also similar to the ``Linear`` module, we add the getter methods for the ``TensorQuantizer`` modules associated with inputs/weights. This could be used to, for example, disable the quantization mechanism by calling ``module_obj.weight_quantizer.disable()``
+
+.. code:: python
+
+    @property
+    def input_quantizer(self):
+        return self._input_quantizer
+
+    @property
+    def weight_quantizer(self):
+        return self._weight_quantizer
+
+- With all of the above changes, the quantized Linear module would look like following:
+
+.. code:: python
+
+    class QuantLinear(nn.Linear, _utils.QuantMixin):
+
+        def __init__(self, in_features, out_features, bias=True, **kwargs):
+            super(QuantLinear, self).__init__(in_features, out_features, bias)
+            quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+
+            self.init_quantizer(quant_desc_input, quant_desc_weight)
+
+        def forward(self, input):
+            quant_input = self._input_quantizer(input)
+            quant_weight = self._weight_quantizer(self.weight)
+
+            output = F.linear(quant_input, quant_weight, bias=self.bias)
+
+            return output
+
+        @property
+        def input_quantizer(self):
+            return self._input_quantizer
+
+        @property
+        def weight_quantizer(self):
+            return self._weight_quantizer
+
+
+Directly Quantizing Inputs In Graph
+-----------------------------------
+It is also possible to directly quantize graph inputs without creating wrappers as explained above.
+
+Here's an example:
+
+.. code:: python
+
+    test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+    quantizer = TensorQuantizer(quant_nn.QuantLinear.default_quant_desc_input)
+
+    quant_input = quantizer(test_input)
+
+    out = F.adaptive_avg_pool2d(quant_input, 3)
+
+Assume that there is a ``F.adaptive_avg_pool2d`` operation in the graph and we'd like to quantize this operation. In the example above, we use ``TensorQuantizer(quant_nn.QuantLinear.default_quant_desc_input)`` to define a quantizer that we then use to actually quantize the ``test_input`` and then feed this quantized input to the ``F.adaptive_avg_pool2d`` operation. Note that this quantizer is the same as the ones we used earlier while created quantized versions of torch's modules.

tools/pytorch-quantization/docs/source/tutorials/quant_resnet50.rst

@@ -0,0 +1,317 @@
+Quantizing Resnet50
+===================
+
+Create a quantized model
+------------------------
+
+Import the necessary python modules:
+
+.. code:: python
+
+   import torch
+   import torch.utils.data
+   from torch import nn
+
+   from pytorch_quantization import nn as quant_nn
+   from pytorch_quantization import calib
+   from pytorch_quantization.tensor_quant import QuantDescriptor
+
+   from torchvision import models
+
+   sys.path.append("path to torchvision/references/classification/")
+   from train import evaluate, train_one_epoch, load_data
+
+Adding quantized modules
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+The first step is to add quantizer modules to the neural network graph.
+This package provides a number of quantized layer modules, which contain quantizers for inputs and weights.
+e.g. ``quant_nn.QuantLinear``, which can be used in place of ``nn.Linear``.
+These quantized layers can be substituted automatically, via monkey-patching, or by manually modifying the model definition.
+
+Automatic layer substitution is done with ``quant_modules``. This should be called before model creation.
+
+.. code:: python
+
+    from pytorch_quantization import quant_modules
+    quant_modules.initialize()
+
+This will apply to all instances of each module.
+If you do not want all modules to be quantized you should instead substitute the quantized modules manually.
+Stand-alone quantizers can also be added to the model with ``quant_nn.TensorQuantizer``.
+
+Post training quantization
+--------------------------
+
+For efficient inference, we want to select a fixed range for each quantizer.
+Starting with a pre-trained model, the simplest way to do this is by calibration.
+
+
+Calibration
+~~~~~~~~~~~
+
+We will use histogram based calibration for activations and the default max calibration for weights.
+
+.. code:: python
+
+   quant_desc_input = QuantDescriptor(calib_method='histogram')
+   quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
+   quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)
+
+   model = models.resnet50(pretrained=True)
+   model.cuda()
+
+To collect activation histograms we must feed sample data in to the model.
+First, create ImageNet dataloaders as done in the training script.
+Then, enable calibration in each quantizer and feed training data in to the model.
+1024 samples (2 batches of 512) should be sufficient to estimate the distribution of activations.
+Use training data for calibration so that validation also measures generalization of the selected ranges.
+
+.. code:: python
+
+    data_path = "PATH to imagenet"
+    batch_size = 512
+
+    traindir = os.path.join(data_path, 'train')
+    valdir = os.path.join(data_path, 'val')
+    dataset, dataset_test, train_sampler, test_sampler = load_data(traindir, valdir, False, False)
+
+    data_loader = torch.utils.data.DataLoader(
+        dataset, batch_size=batch_size,
+        sampler=train_sampler, num_workers=4, pin_memory=True)
+
+    data_loader_test = torch.utils.data.DataLoader(
+        dataset_test, batch_size=batch_size,
+        sampler=test_sampler, num_workers=4, pin_memory=True)
+
+.. code:: python
+
+    def collect_stats(model, data_loader, num_batches):
+        """Feed data to the network and collect statistic"""
+
+        # Enable calibrators
+        for name, module in model.named_modules():
+            if isinstance(module, quant_nn.TensorQuantizer):
+                if module._calibrator is not None:
+                    module.disable_quant()
+                    module.enable_calib()
+                else:
+                    module.disable()
+
+        for i, (image, _) in tqdm(enumerate(data_loader), total=num_batches):
+            model(image.cuda())
+            if i >= num_batches:
+                break
+
+        # Disable calibrators
+        for name, module in model.named_modules():
+            if isinstance(module, quant_nn.TensorQuantizer):
+                if module._calibrator is not None:
+                    module.enable_quant()
+                    module.disable_calib()
+                else:
+                    module.enable()
+
+    def compute_amax(model, **kwargs):
+        # Load calib result
+        for name, module in model.named_modules():
+            if isinstance(module, quant_nn.TensorQuantizer):
+                if module._calibrator is not None:
+                    if isinstance(module._calibrator, calib.MaxCalibrator):
+                        module.load_calib_amax()
+                    else:
+                        module.load_calib_amax(**kwargs)
+                print(F"{name:40}: {module}")
+        model.cuda()
+
+   # It is a bit slow since we collect histograms on CPU
+    with torch.no_grad():
+        collect_stats(model, data_loader, num_batches=2)
+        compute_amax(model, method="percentile", percentile=99.99)
+
+After calibration is done, quantizers will have ``amax`` set, which represents the absolute maximum input value representable in the quantized space.
+By default, weight ranges are per channel while activation ranges are per tensor.
+We can see the condensed amaxes by printing each ``TensorQuantizer`` module.
+
+::
+
+   conv1._input_quantizer                  : TensorQuantizer(8bit fake per-tensor amax=2.6400 calibrator=MaxCalibrator(track_amax=False) quant)
+   conv1._weight_quantizer                 : TensorQuantizer(8bit fake axis=(0) amax=[0.0000, 0.7817](64) calibrator=MaxCalibrator(track_amax=False) quant)
+   layer1.0.conv1._input_quantizer         : TensorQuantizer(8bit fake per-tensor amax=6.8645 calibrator=MaxCalibrator(track_amax=False) quant)
+   layer1.0.conv1._weight_quantizer        : TensorQuantizer(8bit fake axis=(0) amax=[0.0000, 0.7266](64) calibrator=MaxCalibrator(track_amax=False) quant)
+   ...
+
+Evaluate the calibrated model
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Next we will evaluate the classification accuracy of our post training quantized model on the ImageNet validation set.
+
+.. code:: python
+
+    criterion = nn.CrossEntropyLoss()
+    with torch.no_grad():
+        evaluate(model, criterion, data_loader_test, device="cuda", print_freq=20)
+
+    # Save the model
+    torch.save(model.state_dict(), "/tmp/quant_resnet50-calibrated.pth")
+
+This should yield 76.1% top-1 accuracy, which is close to the pre-trained model accuracy of 76.2%.
+
+Use different calibration
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+We can try different calibrations without recollecting the histograms, and see which one gets the best accuracy.
+
+.. code:: python
+
+    with torch.no_grad():
+        compute_amax(model, method="percentile", percentile=99.9)
+        evaluate(model, criterion, data_loader_test, device="cuda", print_freq=20)
+
+    with torch.no_grad():
+        for method in ["mse", "entropy"]:
+            print(F"{method} calibration")
+            compute_amax(model, method=method)
+            evaluate(model, criterion, data_loader_test, device="cuda", print_freq=20)
+
+
+MSE and entropy should both get over 76%. 99.9% clips too many values for resnet50 and will get slightly lower accuracy.
+
+Quantization Aware Training
+---------------------
+
+Optionally, we can fine-tune the calibrated model to improve accuracy further.
+
+.. code:: python
+
+   criterion = nn.CrossEntropyLoss()
+   optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)
+   lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)
+
+   # Training takes about one and half hour per epoch on a single V100
+   train_one_epoch(model, criterion, optimizer, data_loader, "cuda", 0, 100)
+
+   # Save the model
+   torch.save(model.state_dict(), "/tmp/quant_resnet50-finetuned.pth")
+
+After one epoch of fine-tuning, we can achieve over 76.4% top-1 accuracy.
+Fine-tuning for more epochs with learning rate annealing can improve accuracy further.
+For example, fine-tuning for 15 epochs with cosine annealing starting with a learning rate of 0.001 can get over 76.7%.
+It should be noted that the same fine-tuning schedule will improve the accuracy of the unquantized model as well.
+
+Further optimization
+~~~~~~~~~~~~~~~~~~~~
+
+For efficient inference on TensorRT, we need know more details about the runtime optimization.
+TensorRT supports fusion of quantizing convolution and residual add.
+The new fused operator has two inputs. Let us call them conv-input and residual-input.
+Here the fused operator’s output precision must match the residual input precision.
+When there is another quantizing node after the fused operator,
+we can insert a pair of quantizing/dequantizing nodes between the residual-input and the Elementwise-Addition node,
+so that quantizing node after the Convolution node is fused with the Convolution node, and the Convolution node is completely quantized with INT8 input and output.
+We cannot use automatic monkey-patching to apply this optimization and we need to manually insert the quantizing/dequantizing nodes.
+
+First create a copy of resnet.py from https://github.com/pytorch/vision/blob/master/torchvision/models/resnet.py,
+modify the constructor, add explicit bool flag ‘quantize’
+
+.. code:: python
+
+    def resnet50(pretrained: bool = False, progress: bool = True, quantize: bool = False, **kwargs: Any) -> ResNet:
+        return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, quantize, **kwargs)
+    def _resnet(arch: str, block: Type[Union[BasicBlock, Bottleneck]], layers: List[int], pretrained: bool, progress: bool,
+                quantize: bool, **kwargs: Any) -> ResNet:
+        model = ResNet(block, layers, quantize, **kwargs)
+    class ResNet(nn.Module):
+        def __init__(self,
+                     block: Type[Union[BasicBlock, Bottleneck]],
+                     layers: List[int],
+                     quantize: bool = False,
+                     num_classes: int = 1000,
+                     zero_init_residual: bool = False,
+                     groups: int = 1,
+                     width_per_group: int = 64,
+                     replace_stride_with_dilation: Optional[List[bool]] = None,
+                     norm_layer: Optional[Callable[..., nn.Module]] = None) -> None:
+            super(ResNet, self).__init__()
+            self._quantize = quantize
+
+When this ``self._quantize`` flag is set to ``True``, we need replace all the ``nn.Conv2d`` with ``quant_nn.QuantConv2d``.
+
+
+.. code:: python
+
+    def conv3x3(in_planes: int,
+                out_planes: int,
+                stride: int = 1,
+                groups: int = 1,
+                dilation: int = 1,
+                quantize: bool = False) -> nn.Conv2d:
+        """3x3 convolution with padding"""
+        if quantize:
+            return quant_nn.QuantConv2d(in_planes,
+                                        out_planes,
+                                        kernel_size=3,
+                                        stride=stride,
+                                        padding=dilation,
+                                        groups=groups,
+                                        bias=False,
+                                        dilation=dilation)
+        else:
+            return nn.Conv2d(in_planes,
+                             out_planes,
+                             kernel_size=3,
+                             stride=stride,
+                             padding=dilation,
+                             groups=groups,
+                             bias=False,
+                             dilation=dilation)
+      def conv1x1(in_planes: int, out_planes: int, stride: int = 1, quantize: bool = False) -> nn.Conv2d:
+          """1x1 convolution"""
+          if quantize:
+              return quant_nn.QuantConv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+          else:
+              return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+
+The residual conv add can be find both in both ``BasicBlock`` and ``Bottleneck``.
+We need first declare quantization node in the ``__init__`` function.
+
+
+.. code:: python
+
+      def __init__(self,
+                   inplanes: int,
+                   planes: int,
+                   stride: int = 1,
+                   downsample: Optional[nn.Module] = None,
+                   groups: int = 1,
+                   base_width: int = 64,
+                   dilation: int = 1,
+                   norm_layer: Optional[Callable[..., nn.Module]] = None,
+                   quantize: bool = False) -> None:
+          # other code...
+          self._quantize = quantize
+          if self._quantize:
+              self.residual_quantizer = quant_nn.TensorQuantizer(quant_nn.QuantConv2d.default_quant_desc_input)
+
+
+Finally we need patch the ``forward`` function in both ``BasicBlock`` and ``Bottleneck``,
+inserting extra quantization/dequantization nodes here.
+
+
+.. code:: python
+
+      def forward(self, x: Tensor) -> Tensor:
+          # other code...
+          if self._quantize:
+              out += self.residual_quantizer(identity)
+          else:
+              out += identity
+          out = self.relu(out)
+
+          return out
+
+The final resnet code with residual quantized can be found in https://github.com/NVIDIA/TensorRT/blob/master/tools/pytorch-quantization/examples/torchvision/models/classification/resnet.py
+
+
+

tools/pytorch-quantization/docs/source/userguide.rst

@@ -0,0 +1,234 @@
+Basic Functionalities
+---------------------
+
+Quantization function
+~~~~~~~~~~~~~~~~~~~~~
+
+``tensor_quant`` and ``fake_tensor_quant`` are 2 basic functions to
+quantize a tensor. ``fake_tensor_quant`` returns fake quantized tensor
+(float value). ``tensor_quant`` returns quantized tensor (integer value)
+and scale.
+
+.. code:: python
+
+    tensor_quant(inputs, amax, num_bits=8, output_dtype=torch.float, unsigned=False)
+    fake_tensor_quant(inputs, amax, num_bits=8, output_dtype=torch.float, unsigned=False)
+
+Example:
+
+.. code:: python
+
+    from pytorch_quantization import tensor_quant
+
+    # Generate random input. With fixed seed 12345, x should be 
+    # tensor([0.9817, 0.8796, 0.9921, 0.4611, 0.0832, 0.1784, 0.3674, 0.5676, 0.3376, 0.2119])
+    torch.manual_seed(12345)
+    x = torch.rand(10)
+
+    # fake quantize tensor x. fake_quant_x will be 
+    # tensor([0.9843, 0.8828, 0.9921, 0.4609, 0.0859, 0.1797, 0.3672, 0.5703, 0.3359, 0.2109])
+    fake_quant_x = tensor_quant.fake_tensor_quant(x, x.abs().max())
+
+    # quantize tensor x. quant_x will be
+    # tensor([126., 113., 127.,  59.,  11.,  23.,  47.,  73.,  43.,  27.])
+    # with scale=128.0057
+    quant_x, scale = tensor_quant.tensor_quant(x, x.abs().max())
+
+Backward of both functions are defined as `Straight-Through Estimator (STE) <https://arxiv.org/abs/1308.3432>`_.
+
+Descriptor and quantizer
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+``QuantDescriptor`` defines how a tensor should be quantized. There are
+also some predefined ``QuantDescriptor``, e.g.
+``QUANT_DESC_8BIT_PER_TENSOR`` and
+``QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL``.
+
+``TensorQuantizer`` is the module for quantizing tensors and defined by
+``QuantDescriptor``.
+
+.. code:: python
+
+    from pytorch_quantization.tensor_quant import QuantDescriptor
+    from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer
+
+    quant_desc = QuantDescriptor(num_bits=4, fake_quant=False, axis=(0), unsigned=True)
+    quantizer = TensorQuantizer(quant_desc)
+
+    torch.manual_seed(12345)
+    x = torch.rand(10, 9, 8, 7)
+
+    quant_x = quantizer(x)
+
+If ``amax`` is given in the :func:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`, :func:`TensorQuantizer <pytorch_quantization.nn.TensorQuantizer>` will use it to quantize. Otherwise, :func:`TensorQuantizer <pytorch_quantization.nn.TensorQuantizer>`  will compute amax then quantize. amax will be computed w.r.t ``axis`` specified. Note that ``axis`` of QuantDescriptor specify remaining axis as oppsed to axis of `max() <https://docs.scipy.org/doc/numpy/reference/generated/numpy.amax.html>`_.
+
+Quantized module
+~~~~~~~~~~~~~~~~
+
+There are 2 major types of module, ``Conv`` and ``Linear``. Both can
+replace ``torch.nn`` version and apply quantization on both weight and
+activation.
+
+Both take ``quant_desc_input`` and ``quant_desc_weight`` in addition to
+arguments of the original module.
+
+.. code:: python
+
+    from torch import nn
+
+    from pytorch_quantization import tensor_quant
+    import pytorch_quantization.nn as quant_nn
+
+    # pytorch's module
+    fc1 = nn.Linear(in_features, out_features, bias=True)
+    conv1 = nn.Conv2d(in_channels, out_channels, kernel_size)
+
+    # quantized version
+    quant_fc1 = quant_nn.Linear(
+        in_features, out_features, bias=True,
+        quant_desc_input=tensor_quant.QUANT_DESC_8BIT_PER_TENSOR,
+        quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW)
+    quant_conv1 = quant_nn.Conv2d(
+        in_channels, out_channels, kernel_size,
+        quant_desc_input=tensor_quant.QUANT_DESC_8BIT_PER_TENSOR,
+        quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL)
+
+Post training quantization
+--------------------------
+
+A model can be post training quantized by simply by calling ``quant_modules.initialize()``
+
+.. code:: python
+
+    from pytorch_quantization import quant_modules
+    model = torchvision.models.resnet50()
+
+If a model is not entirely defined by module, than TensorQuantizer should be 
+manually created and added to the right place in the model.
+
+Calibration
+~~~~~~~~~~~
+
+Calibration is the TensorRT terminology of passing data samples to the
+quantizer and deciding the best amax for activations. 
+We support 3 calibration methods:
+
+-  ``max``: Simply use global maximum absolute value
+-  ``entropy``: TensorRT's entropy calibration
+-  ``percentile``: Get rid of outlier based on given percentile.
+-  ``mse``: MSE(Mean Squared Error) based calibration
+
+In above ResNet50 example, calibration method is set to ``mse``, it can
+be used as the following example:
+
+.. code:: python
+
+    # Find the TensorQuantizer and enable calibration
+    for name, module in model.named_modules():
+        if name.endswith('_quantizer'):
+            module.enable_calib()
+            module.disable_quant()  # Use full precision data to calibrate
+            
+    # Feeding data samples
+    model(x)
+    # ...
+
+    # Finalize calibration
+    for name, module in model.named_modules():
+        if name.endswith('_quantizer'):
+            module.load_calib_amax()
+            module.disable_calib()
+            module.enable_quant()
+            
+    # If running on GPU, it needs to call .cuda() again because new tensors will be created by calibration process
+    model.cuda()
+
+    # Keep running the quantized model
+    # ...
+
+.. note::
+
+    Calibration needs to be performed before exporting the model to ONNX.
+
+Quantization Aware Training
+---------------------
+
+Quantization Aware Training is based on Straight Through Estimator (STE)
+derivative approximation. It is some time known as “quantization aware
+training”. We don’t use the name because it doesn’t reflect the
+underneath assumption. If anything, it makes training being “unaware” of
+quantization because of the STE approximation.
+
+After calibration is done, Quantization Aware Training is simply select a
+training schedule and continue training the calibrated model. Usually,
+it doesn’t need to fine tune very long. We usually use around 10% of the
+original training schedule, starting at 1% of the initial training
+learning rate, and a cosine annealing learning rate schedule that
+follows the decreasing half of a cosine period, down to 1% of the
+initial fine tuning learning rate (0.01% of the initial training
+learning rate).
+
+Some recommendations
+~~~~~~~~~~~~~~~~~~~~
+
+Quantization Aware Training (Essentially a discrete numerical optimization
+problem) is not a solved problem mathematically. Based on our
+experience, here are some recommendations:
+
+-  For STE approximation to work well, it is better to use small
+   learning rate. Large learning rate is more likely to enlarge the
+   variance introduced by STE approximation and destroy the trained
+   network.
+-  Do not change quantization representation (scale) during training, at
+   least not too frequently. Changing scale every step, it is
+   effectively like changing data format (e8m7, e5m10, e3m4, et.al)
+   every step, which will easily affect convergence.
+
+Export to ONNX
+--------------
+
+The goal of exporting to ONNX is to deploy to TensorRT, not to
+ONNX runtime. So we only export fake quantized model into a form TensorRT will take. Fake
+quantization will be broken into a pair of
+QuantizeLinear/DequantizeLinear ONNX ops. TensorRT will take
+the generated ONNX graph, and execute it in int8 in the most optimized way to its
+capability.
+
+.. note::
+
+    Currently, we only support exporting int8 and fp8 fake quantized modules. 
+    Additionally, quantized modules need to be calibrated before exporting to ONNX. 
+
+Fake quantized model can be exported to ONNX as any other Pytorch model. 
+Please learn more about exporting a Pytorch model to ONNX at
+`torch.onnx <https://pytorch.org/docs/stable/onnx.html?highlight=onnx#module-torch.onnx>`__.
+For example:
+
+.. code:: python
+
+   import pytorch_quantization
+   from pytorch_quantization import nn as quant_nn
+   from pytorch_quantization import quant_modules
+
+   quant_modules.initialize()
+   model = torchvision.models.resnet50()
+   
+   # load the calibrated model
+   state_dict = torch.load("quant_resnet50-entropy-1024.pth", map_location="cpu")
+   model.load_state_dict(state_dict)
+   model.cuda()
+
+   dummy_input = torch.randn(128, 3, 224, 224, device='cuda')
+
+   input_names = [ "actual_input_1" ]
+   output_names = [ "output1" ]
+
+   with pytorch_quantization.enable_onnx_export():
+        # enable_onnx_checker needs to be disabled. See notes below.
+        torch.onnx.export(
+            model, dummy_input, "quant_resnet50.onnx", verbose=True, opset_version=10, enable_onnx_checker=False
+            )
+
+.. Note::
+
+    Note that ``axis`` is added to ``QuantizeLinear`` and ``DequantizeLinear`` in opset13.

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

@@ -0,0 +1,21 @@
+.. role:: hidden
+    :class: hidden-section
+
+pytorch_quantization.utils
+===================================
+
+pytorch_quantization.utils.quant\_logging
+--------------------------------------------
+
+.. automodule:: pytorch_quantization.utils.quant_logging
+    :members:
+    :undoc-members:
+    :show-inheritance:
+
+pytorch_quantization.utils.reduce\_amax
+------------------------------------------
+
+.. automodule:: pytorch_quantization.utils.reduce_amax
+    :members:
+    :undoc-members:
+    :show-inheritance: