chore: import upstream snapshot with attribution

tools/pytorch-quantization/.coveragerc

@@ -0,0 +1,10 @@
+[run]
+omit = pytorch_quantization/utils/amp_wrapper.py, pytorch_quantization/utils/quant_logging.py, *__init__.py
+
+[report]
+exclude_lines =
+    def backward
+    RNNReLUCell
+    RNNTanhCell
+    GRU
+    

tools/pytorch-quantization/.gitignore

@@ -0,0 +1,157 @@
+## copied from PyTorch
+
+.coverage
+.mypy_cache
+*/*.pyc
+*/*.so*
+*/**/__pycache__
+*/**/*.dylib*
+*/**/*.pyc
+*/**/*.pyd
+*/**/*.so*
+*/**/**/*.pyc
+*/**/**/**/*.pyc
+*/**/**/**/**/*.pyc
+aten/build/
+aten/src/ATen/Config.h
+aten/src/ATen/cuda/CUDAConfig.h
+build/
+dist/
+
+
+# IPython notebook checkpoints
+.ipynb_checkpoints
+
+# Editor temporaries
+*.swn
+*.swo
+*.swp
+*.swm
+*~
+
+# macOS dir files
+.DS_Store
+
+# Symbolic files
+tools/shared/cwrap_common.py
+
+# Ninja files
+.ninja_deps
+.ninja_log
+compile_commands.json
+*.egg-info/
+docs/source/scripts/activation_images/
+
+## General
+
+# Compiled Object files
+*.slo
+*.lo
+*.o
+*.cuo
+*.obj
+
+# Compiled Dynamic libraries
+*.so
+*.dylib
+*.dll
+
+# Compiled Static libraries
+*.lai
+*.la
+*.a
+*.lib
+
+# Compiled protocol buffers
+*.pb.h
+*.pb.cc
+*_pb2.py
+
+# Compiled python
+*.pyc
+*.pyd
+
+# Compiled MATLAB
+*.mex*
+
+# IPython notebook checkpoints
+.ipynb_checkpoints
+
+# Editor temporaries
+*.swn
+*.swo
+*.swp
+*~
+
+# Sublime Text settings
+*.sublime-workspace
+*.sublime-project
+
+# Eclipse Project settings
+*.*project
+.settings
+
+# QtCreator files
+*.user
+
+# PyCharm files
+.idea
+
+# OSX dir files
+.DS_Store
+
+## Caffe2
+
+# build, distribute, and bins (+ python proto bindings)
+build
+build_host_protoc
+build_android
+build_ios
+/build_*
+.build_debug/*
+.build_release/*
+distribute/*
+*.testbin
+*.bin
+cmake_build
+.cmake_build
+gen
+.setuptools-cmake-build
+.pytest_cache
+aten/build/*
+
+# Bram
+plsdontbreak
+
+# Generated documentation
+docs/_site
+docs/gathered
+_site
+doxygen
+docs/dev
+
+# LevelDB files
+*.sst
+*.ldb
+LOCK
+LOG*
+CURRENT
+MANIFEST-*
+
+# generated version file
+caffe2/version.py
+
+# setup.py intermediates
+.eggs
+caffe2.egg-info
+
+# Atom/Watchman required file
+.watchmanconfig
+
+# BEGIN NOT-CLEAN-FILES (setup.py handles this marker. Do not change.)
+#
+# Below files are not deleted by "setup.py clean".
+
+# Visual Studio Code files
+.vscode
+.vs

tools/pytorch-quantization/.pylintrc

@@ -0,0 +1,407 @@
+[MASTER]
+
+# Specify a configuration file.
+#rcfile=
+
+# Python code to execute, usually for sys.path manipulation such as
+# pygtk.require().
+#init-hook=
+
+# Add files or directories to the blacklist. They should be base names, not
+# paths.
+ignore=CVS
+
+# Add files or directories matching the regex patterns to the blacklist. The
+# regex matches against base names, not paths.
+ignore-patterns=
+
+# Pickle collected data for later comparisons.
+persistent=yes
+
+# List of plugins (as comma separated values of python modules names) to load,
+# usually to register additional checkers.
+load-plugins=
+
+# Use multiple processes to speed up Pylint.
+jobs=1
+
+# Allow loading of arbitrary C extensions. Extensions are imported into the
+# active Python interpreter and may run arbitrary code.
+unsafe-load-any-extension=no
+
+# A comma-separated list of package or module names from where C extensions may
+# be loaded. Extensions are loading into the active Python interpreter and may
+# run arbitrary code
+extension-pkg-whitelist=
+
+# Allow optimization of some AST trees. This will activate a peephole AST
+# optimizer, which will apply various small optimizations. For instance, it can
+# be used to obtain the result of joining multiple strings with the addition
+# operator. Joining a lot of strings can lead to a maximum recursion error in
+# Pylint and this flag can prevent that. It has one side effect, the resulting
+# AST will be different than the one from reality. This option is deprecated
+# and it will be removed in Pylint 2.0.
+optimize-ast=no
+
+
+[MESSAGES CONTROL]
+
+# Only show warnings with the listed confidence levels. Leave empty to show
+# all. Valid levels: HIGH, INFERENCE, INFERENCE_FAILURE, UNDEFINED
+confidence=
+
+# Enable the message, report, category or checker with the given id(s). You can
+# either give multiple identifier separated by comma (,) or put this option
+# multiple time (only on the command line, not in the configuration file where
+# it should appear only once). See also the "--disable" option for examples.
+#enable=
+
+# Disable the message, report, category or checker with the given id(s). You
+# can either give multiple identifiers separated by comma (,) or put this
+# option multiple times (only on the command line, not in the configuration
+# file where it should appear only once).You can also use "--disable=all" to
+# disable everything first and then reenable specific checks. For example, if
+# you want to run only the similarities checker, you can use "--disable=all
+# --enable=similarities". If you want to run only the classes checker, but have
+# no Warning level messages displayed, use"--disable=all --enable=classes
+# --disable=W"
+disable=old-raise-syntax,unpacking-in-except,using-cmp-argument,filter-builtin-not-iterating,nonzero-method,dict-iter-method,basestring-builtin,unicode-builtin,oct-method,cmp-method,delslice-method,round-builtin,range-builtin-not-iterating,cmp-builtin,old-division,old-octal-literal,raising-string,coerce-method,input-builtin,buffer-builtin,parameter-unpacking,hex-method,map-builtin-not-iterating,getslice-method,execfile-builtin,long-builtin,print-statement,raw_input-builtin,apply-builtin,standarderror-builtin,xrange-builtin,setslice-method,import-star-module-level,intern-builtin,indexing-exception,reduce-builtin,no-absolute-import,long-suffix,coerce-builtin,suppressed-message,dict-view-method,metaclass-assignment,backtick,unichr-builtin,useless-suppression,old-ne-operator,zip-builtin-not-iterating,reload-builtin,file-builtin,next-method-called,import-error
+
+
+[REPORTS]
+
+# Set the output format. Available formats are text, parseable, colorized, msvs
+# (visual studio) and html. You can also give a reporter class, eg
+# mypackage.mymodule.MyReporterClass.
+output-format=text
+
+# Put messages in a separate file for each module / package specified on the
+# command line instead of printing them on stdout. Reports (if any) will be
+# written in a file name "pylint_global.[txt|html]". This option is deprecated
+# and it will be removed in Pylint 2.0.
+files-output=no
+
+# Tells whether to display a full report or only the messages
+reports=yes
+
+# Python expression which should return a note less than 10 (10 is the highest
+# note). You have access to the variables errors warning, statement which
+# respectively contain the number of errors / warnings messages and the total
+# number of statements analyzed. This is used by the global evaluation report
+# (RP0004).
+evaluation=10.0 - ((float(5 * error + warning + refactor + convention) / statement) * 10)
+
+# Template used to display messages. This is a python new-style format string
+# used to format the message information. See doc for all details
+#msg-template=
+
+
+[VARIABLES]
+
+# Tells whether we should check for unused import in __init__ files.
+init-import=no
+
+# A regular expression matching the name of dummy variables (i.e. expectedly
+# not used).
+dummy-variables-rgx=(_+[a-zA-Z0-9]*?$)|dummy
+
+# List of additional names supposed to be defined in builtins. Remember that
+# you should avoid to define new builtins when possible.
+additional-builtins=
+
+# List of strings which can identify a callback function by name. A callback
+# name must start or end with one of those strings.
+callbacks=cb_,_cb
+
+# List of qualified module names which can have objects that can redefine
+# builtins.
+redefining-builtins-modules=six.moves,future.builtins
+
+
+[BASIC]
+
+# Good variable names which should always be accepted, separated by a comma
+good-names=i,j,k,ex,Run,_,x,lr,s
+
+# Bad variable names which should always be refused, separated by a comma
+bad-names=foo,bar,baz,toto,tutu,tata
+
+# Colon-delimited sets of names that determine each other's naming style when
+# the name regexes allow several styles.
+name-group=
+
+# Include a hint for the correct naming format with invalid-name
+include-naming-hint=no
+
+# List of decorators that produce properties, such as abc.abstractproperty. Add
+# to this list to register other decorators that produce valid properties.
+property-classes=abc.abstractproperty
+
+# Regular expression matching correct class names
+class-rgx=[A-Z_][a-zA-Z0-9]+$
+
+# Naming hint for class names
+class-name-hint=[A-Z_][a-zA-Z0-9]+$
+
+# Regular expression matching correct function names
+function-rgx=[a-z_][a-z0-9_]{2,30}$
+
+# Naming hint for function names
+function-name-hint=[a-z_][a-z0-9_]{2,30}$
+
+# Regular expression matching correct attribute names
+attr-rgx=[a-z_][a-z0-9_]{2,30}$
+
+# Naming hint for attribute names
+attr-name-hint=[a-z_][a-z0-9_]{2,30}$
+
+# Regular expression matching correct constant names
+const-rgx=(([A-Z_][A-Z0-9_]*)|(__.*__))$
+
+# Naming hint for constant names
+const-name-hint=(([A-Z_][A-Z0-9_]*)|(__.*__))$
+
+# Regular expression matching correct module names
+module-rgx=(([a-z_][a-z0-9_]*)|([A-Z][a-zA-Z0-9]+))$
+
+# Naming hint for module names
+module-name-hint=(([a-z_][a-z0-9_]*)|([A-Z][a-zA-Z0-9]+))$
+
+# Regular expression matching correct variable names
+variable-rgx=[a-z_][a-z0-9_]{2,30}$
+
+# Naming hint for variable names
+variable-name-hint=[a-z_][a-z0-9_]{2,30}$
+
+# Regular expression matching correct argument names
+argument-rgx=[a-z_][a-z0-9_]{2,30}$
+
+# Naming hint for argument names
+argument-name-hint=[a-z_][a-z0-9_]{2,30}$
+
+# Regular expression matching correct class attribute names
+class-attribute-rgx=([A-Za-z_][A-Za-z0-9_]{2,30}|(__.*__))$
+
+# Naming hint for class attribute names
+class-attribute-name-hint=([A-Za-z_][A-Za-z0-9_]{2,30}|(__.*__))$
+
+# Regular expression matching correct inline iteration names
+inlinevar-rgx=[A-Za-z_][A-Za-z0-9_]*$
+
+# Naming hint for inline iteration names
+inlinevar-name-hint=[A-Za-z_][A-Za-z0-9_]*$
+
+# Regular expression matching correct method names
+method-rgx=[a-z_][a-z0-9_]{2,30}$
+
+# Naming hint for method names
+method-name-hint=[a-z_][a-z0-9_]{2,30}$
+
+# Regular expression which should only match function or class names that do
+# not require a docstring.
+no-docstring-rgx=^_
+
+# Minimum line length for functions/classes that require docstrings, shorter
+# ones are exempt.
+docstring-min-length=-1
+
+
+[ELIF]
+
+# Maximum number of nested blocks for function / method body
+max-nested-blocks=5
+
+
+[SPELLING]
+
+# Spelling dictionary name. Available dictionaries: none. To make it working
+# install python-enchant package.
+spelling-dict=
+
+# List of comma separated words that should not be checked.
+spelling-ignore-words=
+
+# A path to a file that contains private dictionary; one word per line.
+spelling-private-dict-file=
+
+# Tells whether to store unknown words to indicated private dictionary in
+# --spelling-private-dict-file option instead of raising a message.
+spelling-store-unknown-words=no
+
+
+[SIMILARITIES]
+
+# Minimum lines number of a similarity.
+min-similarity-lines=4
+
+# Ignore comments when computing similarities.
+ignore-comments=yes
+
+# Ignore docstrings when computing similarities.
+ignore-docstrings=yes
+
+# Ignore imports when computing similarities.
+ignore-imports=no
+
+
+[FORMAT]
+
+# Maximum number of characters on a single line.
+max-line-length=120
+
+# Regexp for a line that is allowed to be longer than the limit.
+ignore-long-lines=^\s*(# )?<?https?://\S+>?$
+
+# Allow the body of an if to be on the same line as the test if there is no
+# else.
+single-line-if-stmt=no
+
+# List of optional constructs for which whitespace checking is disabled. `dict-
+# separator` is used to allow tabulation in dicts, etc.: {1  : 1,\n222: 2}.
+# `trailing-comma` allows a space between comma and closing bracket: (a, ).
+# `empty-line` allows space-only lines.
+no-space-check=trailing-comma,dict-separator
+
+# Maximum number of lines in a module
+max-module-lines=1000
+
+# String used as indentation unit. This is usually "    " (4 spaces) or "\t" (1
+# tab).
+indent-string='    '
+
+# Number of spaces of indent required inside a hanging  or continued line.
+indent-after-paren=4
+
+# Expected format of line ending, e.g. empty (any line ending), LF or CRLF.
+expected-line-ending-format=
+
+
+[MISCELLANEOUS]
+
+# List of note tags to take in consideration, separated by a comma.
+notes=FIXME,XXX,TODO
+
+
+[LOGGING]
+
+# Logging modules to check that the string format arguments are in logging
+# function parameter format
+logging-modules=logging
+
+
+[TYPECHECK]
+
+# Tells whether missing members accessed in mixin class should be ignored. A
+# mixin class is detected if its name ends with "mixin" (case insensitive).
+ignore-mixin-members=yes
+
+# List of module names for which member attributes should not be checked
+# (useful for modules/projects where namespaces are manipulated during runtime
+# and thus existing member attributes cannot be deduced by static analysis. It
+# supports qualified module names, as well as Unix pattern matching.
+ignored-modules=numpy,torch
+
+# List of class names for which member attributes should not be checked (useful
+# for classes with dynamically set attributes). This supports the use of
+# qualified names.
+ignored-classes=optparse.Values,thread._local,_thread._local
+
+# List of members which are set dynamically and missed by pylint inference
+# system, and so shouldn't trigger E1101 when accessed. Python regular
+# expressions are accepted.
+generated-members=
+
+# List of decorators that produce context managers, such as
+# contextlib.contextmanager. Add to this list to register other decorators that
+# produce valid context managers.
+contextmanager-decorators=contextlib.contextmanager
+
+
+[CLASSES]
+
+# List of method names used to declare (i.e. assign) instance attributes.
+defining-attr-methods=__init__,__new__,setUp
+
+# List of valid names for the first argument in a class method.
+valid-classmethod-first-arg=cls
+
+# List of valid names for the first argument in a metaclass class method.
+valid-metaclass-classmethod-first-arg=mcs
+
+# List of member names, which should be excluded from the protected access
+# warning.
+exclude-protected=_asdict,_fields,_replace,_source,_make
+
+
+[IMPORTS]
+
+# Deprecated modules which should not be used, separated by a comma
+deprecated-modules=optparse
+
+# Create a graph of every (i.e. internal and external) dependencies in the
+# given file (report RP0402 must not be disabled)
+import-graph=
+
+# Create a graph of external dependencies in the given file (report RP0402 must
+# not be disabled)
+ext-import-graph=
+
+# Create a graph of internal dependencies in the given file (report RP0402 must
+# not be disabled)
+int-import-graph=
+
+# Force import order to recognize a module as part of the standard
+# compatibility libraries.
+known-standard-library=
+
+# Force import order to recognize a module as part of a third party library.
+known-third-party=enchant
+
+# Analyse import fallback blocks. This can be used to support both Python 2 and
+# 3 compatible code, which means that the block might have code that exists
+# only in one or another interpreter, leading to false positives when analysed.
+analyse-fallback-blocks=no
+
+
+[DESIGN]
+
+# Maximum number of arguments for function / method
+max-args=10
+
+# Argument names that match this expression will be ignored. Default to name
+# with leading underscore
+ignored-argument-names=_.*
+
+# Maximum number of locals for function / method body
+max-locals=15
+
+# Maximum number of return / yield for function / method body
+max-returns=6
+
+# Maximum number of branch for function / method body
+max-branches=12
+
+# Maximum number of statements in function / method body
+max-statements=50
+
+# Maximum number of parents for a class (see R0901).
+max-parents=7
+
+# Maximum number of attributes for a class (see R0902).
+max-attributes=7
+
+# Minimum number of public methods for a class (see R0903).
+min-public-methods=2
+
+# Maximum number of public methods for a class (see R0904).
+max-public-methods=20
+
+# Maximum number of boolean expressions in a if statement
+max-bool-expr=5
+
+
+[EXCEPTIONS]
+
+# Exceptions that will emit a warning when being caught. Defaults to
+# "Exception"
+overgeneral-exceptions=Exception

tools/pytorch-quantization/.style.yapf

@@ -0,0 +1,7 @@
+[style]
+based_on_style = google
+indent_width = 4
+COLUMN_LIMIT = 120
+SPACES_BEFORE_COMMENT = 2
+EACH_DICT_ENTRY_ON_SEPARATE_LINE = 1
+CONTINUATION_INDENT_WIDTH = 4

tools/pytorch-quantization/CONTRIBUTING.md

@@ -0,0 +1,119 @@
+# Contributing guidelines
+
+## Style
+
+#### C++ coding style
+
+Follow [Google C++ Style Guide](https://google.github.io/styleguide/cppguide.html), with max-line-length extended to 120.
+
+Run `cpplint` before committing code.
+
+#### Python coding style
+
+Follow [Google Python Style Guide,](https://google.github.io/styleguide/pyguide.html) with max-line-length extended to 120. Exceptions are allowed if it feels more natural to follow Pytorch style. For example, Pytorch allows import relative path, also class name.
+
+Run `pylint` before committing code. It doesn't mean every issue has to be corrected nor check has to be manually disabled. Just make sure you are aware of the remaining issues and you are comfort with all of them. But, don't leave lint error there. Disable it explicitly if it is not a real error.
+
+Install `pylint`
+
+```bash
+pip install pylint
+```
+
+To check a file with `pylint`:
+
+```bash
+pylint --rcfile=.pylintrc myfile.py
+```
+
+#### Yapf
+
+[yapf](https://github.com/google/yapf/) is an auto format tool owned by Google (not a Google product). To save the time of arguing code style during code review, use yapf to format the code is a good option. Note that it doesn't reformat comment.
+
+Install `yapf`
+
+```bash
+pip install yapf
+```
+
+Format code with yapf
+
+```bash
+yapf myfile.py --style .style.yapf
+```
+
+There are Sublime and Vim plugins.
+
+## Test
+
+Use [googletest](https://github.com/google/googletest) for c++ code.
+
+Use [pytest](https://docs.pytest.org/en/latest/) for python code.
+
+To run all the tests:
+
+```shell
+pytest --verbose
+```
+
+To run a particular test file
+
+```shell
+pytest --verbose mytestfile.py
+```
+
+
+
+## Conventions
+
+Quantization is a very overloaded word, many things related to it can create a lot of confusions. Let's try to avoid confusions as much as possible by following existing conventions. Generally, if there is a similar Tensorflow or numpy function, follow its convention. Though Tensorflow uses `quantized`, `quantization` and `quant`, let's stick with the shortest one only.
+
+### Naming
+
+##### Function and class name
+
+When developing quantized version of a function or module, add`Quant` to class name, add `quant_` to function name, e.g.
+
+```python
+class Linear(...)
+class QuantLinear(...)
+
+def linear(...)
+def quant_linear(...)
+```
+
+##### Variable name
+
+Add prefix `quant_mode_`, `num_bits_`  etc. to name of tensors will be quantized, e.g.
+
+```python
+def matmul(a, b)
+def quant_matmul(a, b)
+```
+
+Don't use prefix/suffix `weight` or `act` if tensor being quantized doesn't have them explicitly in name. From function's perspective, it takes tensors, not necessarily weight and activation tensors. e.g. `a` and `b` of `matmul` can both be either weight or activation.
+
+##### Quantization mode
+
+There only convention here we can adopt is `per_channel`. Other things, like there is no convention to follow of per row/column scale of matrix multiply. And though we usually absolute max value based scaling factor, there are other ways to decide it, like KL-divergence. 
+
+Our API design is flexible enough to support any granularity of quantization. The main concept is `axis`.
+
+```python
+# axis=None means per tensor
+# For 2d convolution weight with layout KCRS, axis=(1, 2, 3) means perchannel quantization
+# more example below
+QUANT_DESC_8BIT_PER_TENSOR = QuantDescriptor(num_bits=8)
+QUANT_DESC_UNSIGNED_8BIT_PER_TENSOR = QuantDescriptor(num_bits=8, unsigned=True)
+QUANT_DESC_8BIT_CONV1D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW = QuantDescriptor(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONVTRANSPOSE1D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(1))
+QUANT_DESC_8BIT_CONVTRANSPOSE2D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(1))
+QUANT_DESC_8BIT_CONVTRANSPOSE3D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(1))
+
+```
+
+### Misc
+

tools/pytorch-quantization/LICENSE

@@ -0,0 +1,283 @@
+
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   Copyright 2020 NVIDIA Corporation
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+
+
+   PORTIONS LICENSED AS FOLLOWS
+
+   > pytorch_quantization/nn/modules/quant_bert.py
+
+      Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+
+      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.
+
+   > pytorch_quantization/nn/modules/*
+
+      From PyTorch:
+
+      Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
+      Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
+      Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
+      Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
+      Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
+      Copyright (c) 2011-2013 NYU                      (Clement Farabet)
+      Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
+      Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
+      Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
+
+      From Caffe2:
+
+      Copyright (c) 2016-present, Facebook Inc. All rights reserved.
+
+      All contributions by Facebook:
+      Copyright (c) 2016 Facebook Inc.
+      
+      All contributions by Google:
+      Copyright (c) 2015 Google Inc.
+      All rights reserved.
+      
+      All contributions by Yangqing Jia:
+      Copyright (c) 2015 Yangqing Jia
+      All rights reserved.
+      
+      All contributions from Caffe:
+      Copyright(c) 2013, 2014, 2015, the respective contributors
+      All rights reserved.
+      
+      All other contributions:
+      Copyright(c) 2015, 2016 the respective contributors
+      All rights reserved.
+      
+      Caffe2 uses a copyright model similar to Caffe: each contributor holds
+      copyright over their contributions to Caffe2. The project versioning records
+      all such contribution and copyright details. If a contributor wants to further
+      mark their specific copyright on a particular contribution, they should
+      indicate their copyright solely in the commit message of the change when it is
+      committed.
+
+      All rights reserved.
+
+      Redistribution and use in source and binary forms, with or without
+      modification, are permitted provided that the following conditions are met:
+
+      1. Redistributions of source code must retain the above copyright
+         notice, this list of conditions and the following disclaimer.
+
+      2. Redistributions in binary form must reproduce the above copyright
+         notice, this list of conditions and the following disclaimer in the
+         documentation and/or other materials provided with the distribution.
+
+      3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America
+         and IDIAP Research Institute nor the names of its contributors may be
+         used to endorse or promote products derived from this software without
+         specific prior written permission.
+
+      THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+      AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+      IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+      ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+      LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+      CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+      SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+      INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+      CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+      ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+      POSSIBILITY OF SUCH DAMAGE.

tools/pytorch-quantization/MANIFEST.in

@@ -0,0 +1,3 @@
+include LICENSE
+include requirements.txt
+recursive-include src *

tools/pytorch-quantization/README.md

@@ -0,0 +1,45 @@
+**Note: Pytorch Quantization development has transitioned to the [TensorRT Model Optimizer](https://github.com/NVIDIA/TensorRT-Model-Optimizer). All developers are encouraged to use the TensorRT Model Optimizer to benefit from the latest advancements on quantization and compression. While the Pytorch Quantization code will remain available, it will no longer receive further development.**
+
+# Pytorch Quantization
+
+PyTorch-Quantization is a toolkit for training and evaluating PyTorch models with simulated quantization. Quantization can be added to the model automatically, or manually, allowing the model to be tuned for accuracy and performance. Quantization is compatible with NVIDIAs high performance integer kernels which leverage integer Tensor Cores. The quantized model can be exported to ONNX and imported by TensorRT 8.0 and later.
+
+## Install
+
+#### Binaries
+
+```bash
+pip install pytorch-quantization --extra-index-url https://pypi.ngc.nvidia.com
+```
+
+#### From Source
+
+```bash
+git clone https://github.com/NVIDIA/TensorRT.git
+cd tools/pytorch-quantization
+```
+
+Install PyTorch and prerequisites
+```bash
+pip install -r requirements.txt
+# for CUDA 10.2 users
+pip install torch>=1.9.1
+# for CUDA 11.1 users
+pip install torch>=1.9.1+cu111
+```
+
+Build and install pytorch-quantization
+```bash
+# Python version >= 3.7, GCC version >= 5.4 required
+python setup.py install
+```
+
+#### NGC Container
+
+`pytorch-quantization` is preinstalled in NVIDIA NGC PyTorch container, e.g. `nvcr.io/nvidia/pytorch:22.12-py3`
+
+## Resources
+
+* Pytorch Quantization Toolkit [userguide](https://docs.nvidia.com/deeplearning/tensorrt/pytorch-quantization-toolkit/docs/index.html)
+* Quantization Basics [whitepaper](https://arxiv.org/abs/2004.09602)
+

tools/pytorch-quantization/VERSION

@@ -0,0 +1 @@
+2.2.1

tools/pytorch-quantization/docs/Makefile

@@ -0,0 +1,36 @@
+#
+# 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.
+#
+
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line.
+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/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:

tools/pytorch-quantization/examples/calibrate_quant_resnet50.ipynb

@@ -0,0 +1,749 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import datetime\n",
+    "import os\n",
+    "import sys\n",
+    "import time\n",
+    "import collections\n",
+    "\n",
+    "import torch\n",
+    "import torch.utils.data\n",
+    "from torch import nn\n",
+    "\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "import torchvision\n",
+    "from torchvision import transforms\n",
+    "\n",
+    "from pytorch_quantization import nn as quant_nn\n",
+    "from pytorch_quantization import calib\n",
+    "from pytorch_quantization.tensor_quant import QuantDescriptor\n",
+    "\n",
+    "from absl import logging\n",
+    "logging.set_verbosity(logging.FATAL)  # Disable logging as they are too noisy in notebook\n",
+    "\n"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# For simplicity, import train and eval functions from the train script from torchvision instead of copything them here\n",
+    "# Download torchvision from https://github.com/pytorch/vision\n",
+    "sys.path.append(\"/raid/skyw/models/torchvision/references/classification/\")\n",
+    "from train import evaluate, train_one_epoch, load_data"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Set default QuantDescriptor to use histogram based calibration for activation"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "quant_desc_input = QuantDescriptor(calib_method='histogram')\n",
+    "quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)\n",
+    "quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Initialize quantized modules"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pytorch_quantization import quant_modules\n",
+    "quant_modules.initialize()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create model with pretrained weight"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "ResNet(\n",
+       "  (conv1): QuantConv2d(\n",
+       "    3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False\n",
+       "    (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "    (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "  )\n",
+       "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "  (relu): ReLU(inplace=True)\n",
+       "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
+       "  (layer1): Sequential(\n",
+       "    (0): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "      (downsample): Sequential(\n",
+       "        (0): QuantConv2d(\n",
+       "          64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "        )\n",
+       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "    )\n",
+       "    (1): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (2): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "  )\n",
+       "  (layer2): Sequential(\n",
+       "    (0): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "      (downsample): Sequential(\n",
+       "        (0): QuantConv2d(\n",
+       "          256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
+       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "        )\n",
+       "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "    )\n",
+       "    (1): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (2): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (3): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "  )\n",
+       "  (layer3): Sequential(\n",
+       "    (0): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "      (downsample): Sequential(\n",
+       "        (0): QuantConv2d(\n",
+       "          512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
+       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "        )\n",
+       "        (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "    )\n",
+       "    (1): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (2): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (3): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (4): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (5): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "  )\n",
+       "  (layer4): Sequential(\n",
+       "    (0): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "      (downsample): Sequential(\n",
+       "        (0): QuantConv2d(\n",
+       "          1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
+       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "        )\n",
+       "        (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "    )\n",
+       "    (1): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (2): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "  )\n",
+       "  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
+       "  (fc): QuantLinear(\n",
+       "    in_features=2048, out_features=1000, bias=True\n",
+       "    (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=dynamic calibrator=HistogramCalibrator quant)\n",
+       "    (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=dynamic calibrator=MaxCalibrator quant)\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "model = torchvision.models.resnet50(pretrained=True, progress=False)\n",
+    "model.cuda()"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Create data loader"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Loading data\n",
+      "Loading training data\n",
+      "Took 3.580507755279541\n",
+      "Loading validation data\n",
+      "Creating data loaders\n"
+     ]
+    }
+   ],
+   "source": [
+    "data_path = \"/raid/data/imagenet/imagenet_pytorch\"\n",
+    "batch_size = 512\n",
+    "\n",
+    "traindir = os.path.join(data_path, 'train')\n",
+    "valdir = os.path.join(data_path, 'val')\n",
+    "_args = collections.namedtuple('mock_args', ['model', 'distributed', 'cache_dataset'])\n",
+    "dataset, dataset_test, train_sampler, test_sampler = load_data(traindir, valdir, _args(model='resnet50', distributed=False, cache_dataset=False))\n",
+    "\n",
+    "data_loader = torch.utils.data.DataLoader(\n",
+    "    dataset, batch_size=batch_size,\n",
+    "    sampler=train_sampler, num_workers=4, pin_memory=True)\n",
+    "\n",
+    "data_loader_test = torch.utils.data.DataLoader(\n",
+    "    dataset_test, batch_size=batch_size,\n",
+    "    sampler=test_sampler, num_workers=4, pin_memory=True)\n"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Calibrate the model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def collect_stats(model, data_loader, num_batches):\n",
+    "    \"\"\"Feed data to the network and collect statistic\"\"\"\n",
+    "\n",
+    "    # Enable calibrators\n",
+    "    for name, module in model.named_modules():\n",
+    "        if isinstance(module, quant_nn.TensorQuantizer):\n",
+    "            if module._calibrator is not None:\n",
+    "                module.disable_quant()\n",
+    "                module.enable_calib()\n",
+    "            else:\n",
+    "                module.disable()\n",
+    "\n",
+    "    for i, (image, _) in tqdm(enumerate(data_loader), total=num_batches):\n",
+    "        model(image.cuda())\n",
+    "        if i >= num_batches:\n",
+    "            break\n",
+    "\n",
+    "    # Disable calibrators\n",
+    "    for name, module in model.named_modules():\n",
+    "        if isinstance(module, quant_nn.TensorQuantizer):\n",
+    "            if module._calibrator is not None:\n",
+    "                module.enable_quant()\n",
+    "                module.disable_calib()\n",
+    "            else:\n",
+    "                module.enable()\n",
+    "            \n",
+    "def compute_amax(model, **kwargs):\n",
+    "    # Load calib result\n",
+    "    for name, module in model.named_modules():\n",
+    "        if isinstance(module, quant_nn.TensorQuantizer):\n",
+    "            if module._calibrator is not None:\n",
+    "                if isinstance(module._calibrator, calib.MaxCalibrator):\n",
+    "                    module.load_calib_amax()\n",
+    "                else:\n",
+    "                    module.load_calib_amax(**kwargs)\n",
+    "#             print(F\"{name:40}: {module}\")\n",
+    "    model.cuda()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "100%|██████████| 2/2 [04:50<00:00, 111.13s/it]"
+     ]
+    }
+   ],
+   "source": [
+    "# It is a bit slow since we collect histograms on CPU\n",
+    "with torch.no_grad():\n",
+    "    collect_stats(model, data_loader, num_batches=2)\n",
+    "    compute_amax(model, method=\"percentile\", percentile=99.99)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Now evaluate the calibrated model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Test:  [ 0/98]  eta: 0:05:53  loss: 0.5656 (0.5656)  acc1: 85.7422 (85.7422)  acc5: 96.0938 (96.0938)  time: 3.6079  data: 2.8152  max mem: 5880\n",
+      "Test:  [20/98]  eta: 0:01:07  loss: 0.6741 (0.6825)  acc1: 82.8125 (82.4219)  acc5: 95.8984 (95.7682)  time: 0.7343  data: 0.0002  max mem: 5882\n",
+      "Test:  [40/98]  eta: 0:00:46  loss: 0.6995 (0.7157)  acc1: 80.0781 (81.4024)  acc5: 96.0938 (95.7412)  time: 0.7226  data: 0.0002  max mem: 5882\n",
+      "Test:  [60/98]  eta: 0:00:29  loss: 1.1064 (0.8590)  acc1: 71.4844 (78.2627)  acc5: 91.0156 (94.1150)  time: 0.7259  data: 0.0002  max mem: 5882\n",
+      "Test:  [80/98]  eta: 0:00:13  loss: 1.1220 (0.9372)  acc1: 72.4609 (76.7072)  acc5: 89.6484 (93.1375)  time: 0.7220  data: 0.0002  max mem: 5882\n",
+      "Test: Total time: 0:01:13\n",
+      " * Acc@1 76.138 Acc@5 92.916\n"
+     ]
+    }
+   ],
+   "source": [
+    "criterion = nn.CrossEntropyLoss()\n",
+    "with torch.no_grad():\n",
+    "    evaluate(model, criterion, data_loader_test, device=\"cuda\", print_freq=20)\n",
+    "    \n",
+    "# Save the model\n",
+    "torch.save(model.state_dict(), \"/tmp/quant_resnet50-calibrated.pth\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## We can also try different calibrations and see which one works the best"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 10,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Test:  [ 0/98]  eta: 0:05:27  loss: 0.6037 (0.6037)  acc1: 84.9609 (84.9609)  acc5: 95.3125 (95.3125)  time: 3.3411  data: 2.6190  max mem: 5882\n",
+      "Test:  [20/98]  eta: 0:01:06  loss: 0.6760 (0.7041)  acc1: 81.2500 (81.7522)  acc5: 95.7031 (95.4892)  time: 0.7243  data: 0.0002  max mem: 5882\n",
+      "Test:  [40/98]  eta: 0:00:45  loss: 0.7241 (0.7351)  acc1: 79.1016 (80.7784)  acc5: 95.8984 (95.4459)  time: 0.7243  data: 0.0002  max mem: 5882\n",
+      "Test:  [60/98]  eta: 0:00:29  loss: 1.1162 (0.8793)  acc1: 71.4844 (77.6383)  acc5: 90.8203 (93.7948)  time: 0.7204  data: 0.0002  max mem: 5882\n",
+      "Test:  [80/98]  eta: 0:00:13  loss: 1.1498 (0.9603)  acc1: 71.4844 (76.0368)  acc5: 89.4531 (92.7156)  time: 0.7164  data: 0.0002  max mem: 5882\n",
+      "Test: Total time: 0:01:12\n",
+      " * Acc@1 75.438 Acc@5 92.486\n"
+     ]
+    }
+   ],
+   "source": [
+    "with torch.no_grad():\n",
+    "    compute_amax(model, method=\"percentile\", percentile=99.9)\n",
+    "    evaluate(model, criterion, data_loader_test, device=\"cuda\", print_freq=20)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "mse calibration\n",
+      "Test:  [ 0/98]  eta: 0:06:34  loss: 0.5700 (0.5700)  acc1: 85.1562 (85.1562)  acc5: 96.2891 (96.2891)  time: 4.0243  data: 3.3231  max mem: 5882\n",
+      "Test:  [20/98]  eta: 0:01:08  loss: 0.6758 (0.6838)  acc1: 82.8125 (82.5707)  acc5: 96.0938 (95.7868)  time: 0.7204  data: 0.0002  max mem: 5882\n",
+      "Test:  [40/98]  eta: 0:00:46  loss: 0.7047 (0.7163)  acc1: 80.2734 (81.4834)  acc5: 96.2891 (95.7746)  time: 0.7178  data: 0.0002  max mem: 5882\n",
+      "Test:  [60/98]  eta: 0:00:29  loss: 1.1127 (0.8585)  acc1: 71.0938 (78.3395)  acc5: 90.8203 (94.1278)  time: 0.7192  data: 0.0002  max mem: 5882\n",
+      "Test:  [80/98]  eta: 0:00:13  loss: 1.1261 (0.9367)  acc1: 72.6562 (76.7530)  acc5: 89.8438 (93.1785)  time: 0.7176  data: 0.0002  max mem: 5882\n",
+      "Test: Total time: 0:01:13\n",
+      " * Acc@1 76.186 Acc@5 92.926\n",
+      "entropy calibration\n",
+      "Test:  [ 0/98]  eta: 0:05:28  loss: 0.5648 (0.5648)  acc1: 85.3516 (85.3516)  acc5: 96.0938 (96.0938)  time: 3.3558  data: 2.6268  max mem: 5882\n",
+      "Test:  [20/98]  eta: 0:01:05  loss: 0.6724 (0.6815)  acc1: 82.8125 (82.5428)  acc5: 95.8984 (95.7589)  time: 0.7196  data: 0.0002  max mem: 5882\n",
+      "Test:  [40/98]  eta: 0:00:45  loss: 0.7090 (0.7149)  acc1: 80.6641 (81.4929)  acc5: 96.0938 (95.7269)  time: 0.7214  data: 0.0002  max mem: 5882\n",
+      "Test:  [60/98]  eta: 0:00:29  loss: 1.1077 (0.8571)  acc1: 72.0703 (78.3779)  acc5: 90.6250 (94.0798)  time: 0.7198  data: 0.0002  max mem: 5882\n",
+      "Test:  [80/98]  eta: 0:00:13  loss: 1.1253 (0.9356)  acc1: 72.2656 (76.7626)  acc5: 90.0391 (93.1231)  time: 0.7192  data: 0.0002  max mem: 5882\n",
+      "Test: Total time: 0:01:12\n",
+      " * Acc@1 76.206 Acc@5 92.900\n"
+     ]
+    }
+   ],
+   "source": [
+    "with torch.no_grad():\n",
+    "    for method in [\"mse\", \"entropy\"]:\n",
+    "        print(F\"{method} calibration\")\n",
+    "        compute_amax(model, method=method)\n",
+    "        evaluate(model, criterion, data_loader_test, device=\"cuda\", print_freq=20)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "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.6.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tools/pytorch-quantization/examples/finetune_quant_resnet50.ipynb

@@ -0,0 +1,605 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "Following the calibrate_quant_resnet50 example, now we fine tune the model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "WARNING: Logging before flag parsing goes to stderr.\n",
+      "W0608 21:25:39.018203 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
+      "W0608 21:25:39.019082 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
+      "W0608 21:25:39.019555 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
+      "W0608 21:25:39.020030 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
+      "W0608 21:25:39.020492 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
+      "W0608 21:25:39.020947 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n",
+      "W0608 21:25:39.021392 140228493526848 tensor_quant.py:96] Meaning of axis has changed since v2.0. Make sure to update.\n"
+     ]
+    }
+   ],
+   "source": [
+    "import datetime\n",
+    "import os\n",
+    "import sys\n",
+    "import time\n",
+    "\n",
+    "import torch\n",
+    "import torch.utils.data\n",
+    "from torch import nn\n",
+    "\n",
+    "from tqdm import tqdm\n",
+    "\n",
+    "import torchvision\n",
+    "from torchvision import transforms\n",
+    "\n",
+    "from pytorch_quantization import nn as quant_nn\n",
+    "from pytorch_quantization import calib\n",
+    "from pytorch_quantization.tensor_quant import QuantDescriptor\n",
+    "\n",
+    "from absl import logging\n",
+    "logging.set_verbosity(logging.FATAL)  # Disable logging as they are too noisy in notebook"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# For simplicity, import train and eval functions from the train script from torchvision instead of copything them here\n",
+    "sys.path.append(\"/raid/skyw/models/torchvision/references/classification/\")\n",
+    "from train import evaluate, train_one_epoch, load_data"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "QuantResNet(\n",
+       "  (conv1): QuantConv2d(\n",
+       "    3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False\n",
+       "    (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=2.6387 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "    (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.7817](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "  )\n",
+       "  (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "  (relu): ReLU(inplace=True)\n",
+       "  (maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)\n",
+       "  (layer1): Sequential(\n",
+       "    (0): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=2.9730 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.7266](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.0971 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.4679](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3318 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3936](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "      (downsample): Sequential(\n",
+       "        (0): QuantConv2d(\n",
+       "          64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=2.9730 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.9879](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        )\n",
+       "        (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "    )\n",
+       "    (1): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4872 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.2618](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.0466 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0564, 0.5201](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5106 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.2946](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (2): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        256, 64, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5250 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0533, 0.1921](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9980 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0810, 0.2856](64) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.6532 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0197, 0.2752](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "  )\n",
+       "  (layer2): Sequential(\n",
+       "    (0): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        256, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5499 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0670, 0.3532](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1606 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0694, 0.2993](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1425 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3917](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "      (downsample): Sequential(\n",
+       "        (0): QuantConv2d(\n",
+       "          256, 512, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
+       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5499 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.5662](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        )\n",
+       "        (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "    )\n",
+       "    (1): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4626 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0370, 0.2522](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8304 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0142, 0.2998](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1722 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3038](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (2): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4864 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0653, 0.2383](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9450 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0646, 0.2556](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8535 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0162, 0.3522](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (3): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        512, 128, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5229 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0648, 0.2814](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9247 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0595, 0.2210](128) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        128, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9747 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0127, 0.2956](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "  )\n",
+       "  (layer3): Sequential(\n",
+       "    (0): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        512, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5941 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0842, 0.3425](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        256, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3565 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0544, 0.2008](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.0293 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3212](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "      (downsample): Sequential(\n",
+       "        (0): QuantConv2d(\n",
+       "          512, 1024, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
+       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.5941 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3460](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        )\n",
+       "        (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "    )\n",
+       "    (1): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3305 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0407, 0.2942](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9844 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0432, 0.2634](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9333 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.4969](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (2): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3388 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0469, 0.2715](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8617 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0397, 0.2100](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.7507 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3538](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (3): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3554 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0553, 0.2390](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9257 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0455, 0.2792](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8117 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0000, 0.3126](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (4): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4199 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0598, 0.2722](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.9274 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0459, 0.1919](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8702 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0160, 0.3161](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (5): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        1024, 256, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.4258 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0618, 0.3995](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.2256 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0498, 0.2236](256) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        256, 1024, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3560 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0180, 0.3288](1024) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "  )\n",
+       "  (layer4): Sequential(\n",
+       "    (0): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        1024, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3915 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0744, 0.3415](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        512, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1571 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0432, 0.3993](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1295 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0198, 0.3546](2048) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "      (downsample): Sequential(\n",
+       "        (0): QuantConv2d(\n",
+       "          1024, 2048, kernel_size=(1, 1), stride=(2, 2), bias=False\n",
+       "          (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.3915 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "          (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0197, 0.6413](2048) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        )\n",
+       "        (1): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      )\n",
+       "    )\n",
+       "    (1): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=3.9348 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0482, 0.7003](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1277 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0379, 0.2257](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=0.8992 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0211, 0.2427](2048) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "    (2): Bottleneck(\n",
+       "      (conv1): QuantConv2d(\n",
+       "        2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=5.2181 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0601, 0.4541](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv2): QuantConv2d(\n",
+       "        512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.2051 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0338, 0.1416](512) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn2): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (conv3): QuantConv2d(\n",
+       "        512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False\n",
+       "        (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=1.1045 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "        (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.0119, 0.2798](2048) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "      )\n",
+       "      (bn3): BatchNorm2d(2048, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
+       "      (relu): ReLU(inplace=True)\n",
+       "    )\n",
+       "  )\n",
+       "  (avgpool): AdaptiveAvgPool2d(output_size=(1, 1))\n",
+       "  (fc): QuantLinear(\n",
+       "    in_features=2048, out_features=1000, bias=True\n",
+       "    (_input_quantizer): TensorQuantizer(8bit fake per-tensor amax=5.5345 calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "    (_weight_quantizer): TensorQuantizer(8bit fake axis=0 amax=[0.1716, 0.7371](1000) calibrator=MaxCalibrator(track_amax=False) quant)\n",
+       "  )\n",
+       ")"
+      ]
+     },
+     "execution_count": 3,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "from pytorch_quantization import quant_modules\n",
+    "quant_modules.initialize()\n",
+    "\n",
+    "# Create and load the calibrated model\n",
+    "model = torchvision.models.resnet50()\n",
+    "model.load_state_dict(torch.load(\"/tmp/quant_resnet50-calibrated.pth\"))\n",
+    "model.cuda()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "data_path = \"/raid/data/imagenet/imagenet_pytorch\"\n",
+    "\n",
+    "traindir = os.path.join(data_path, 'train')\n",
+    "valdir = os.path.join(data_path, 'val')\n",
+    "dataset, dataset_test, train_sampler, test_sampler = load_data(traindir, valdir, False, False)\n",
+    "\n",
+    "data_loader = torch.utils.data.DataLoader(\n",
+    "    dataset, batch_size=256,\n",
+    "    sampler=train_sampler, num_workers=4, pin_memory=True)\n",
+    "\n",
+    "data_loader_test = torch.utils.data.DataLoader(\n",
+    "    dataset_test, batch_size=256,\n",
+    "    sampler=test_sampler, num_workers=4, pin_memory=True)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Quantized fine tuning\n",
+    "Let's fine tune the model with fake quantization. We only fine tune for 1 epoch as an example."
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "criterion = nn.CrossEntropyLoss()\n",
+    "optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)\n",
+    "lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)\n",
+    "\n",
+    "data_loader = torch.utils.data.DataLoader(\n",
+    "    dataset, batch_size=128,\n",
+    "    sampler=train_sampler, num_workers=16, pin_memory=True)\n",
+    "\n",
+    "# Training takes about one and half hour per epoch on single V100\n",
+    "train_one_epoch(model, criterion, optimizer, data_loader, \"cuda\", 0, 100)"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Evaluate the fine tuned model"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Test:  [  0/196]  eta: 0:09:36  loss: 0.4680 (0.4680)  acc1: 85.9375 (85.9375)  acc5: 98.0469 (98.0469)  time: 2.9406  data: 2.1852  max mem: 16096\n",
+      "Test:  [ 10/196]  eta: 0:01:58  loss: 0.6694 (0.6522)  acc1: 83.2031 (82.9545)  acc5: 96.0938 (96.1293)  time: 0.6346  data: 0.1988  max mem: 16096\n",
+      "Test:  [ 20/196]  eta: 0:01:30  loss: 0.6738 (0.6733)  acc1: 82.0312 (82.4777)  acc5: 95.7031 (95.7961)  time: 0.3928  data: 0.0001  max mem: 16096\n",
+      "Test:  [ 30/196]  eta: 0:01:18  loss: 0.6219 (0.6322)  acc1: 84.3750 (83.9718)  acc5: 95.7031 (96.0181)  time: 0.3859  data: 0.0001  max mem: 16096\n",
+      "Test:  [ 40/196]  eta: 0:01:10  loss: 0.6801 (0.6750)  acc1: 81.6406 (82.5934)  acc5: 95.7031 (95.9604)  time: 0.3861  data: 0.0001  max mem: 16096\n",
+      "Test:  [ 50/196]  eta: 0:01:04  loss: 0.6937 (0.6724)  acc1: 80.0781 (82.3529)  acc5: 96.8750 (96.0938)  time: 0.3834  data: 0.0001  max mem: 16096\n",
+      "Test:  [ 60/196]  eta: 0:00:58  loss: 0.7149 (0.6849)  acc1: 80.0781 (81.9864)  acc5: 96.4844 (96.1066)  time: 0.3854  data: 0.0001  max mem: 16096\n",
+      "Test:  [ 70/196]  eta: 0:00:53  loss: 0.6616 (0.6716)  acc1: 80.8594 (82.2843)  acc5: 96.4844 (96.1983)  time: 0.3859  data: 0.0001  max mem: 16096\n",
+      "Test:  [ 80/196]  eta: 0:00:48  loss: 0.6510 (0.6968)  acc1: 81.2500 (81.7467)  acc5: 95.7031 (95.9201)  time: 0.3860  data: 0.0001  max mem: 16096\n",
+      "Test:  [ 90/196]  eta: 0:00:44  loss: 0.9469 (0.7444)  acc1: 76.1719 (80.6834)  acc5: 92.5781 (95.4370)  time: 0.3868  data: 0.0001  max mem: 16096\n",
+      "Test:  [100/196]  eta: 0:00:39  loss: 1.1594 (0.7964)  acc1: 70.7031 (79.5521)  acc5: 90.6250 (94.8755)  time: 0.3864  data: 0.0001  max mem: 16096\n",
+      "Test:  [110/196]  eta: 0:00:35  loss: 1.1594 (0.8214)  acc1: 72.2656 (79.0365)  acc5: 91.4062 (94.6298)  time: 0.3836  data: 0.0001  max mem: 16096\n",
+      "Test:  [120/196]  eta: 0:00:31  loss: 0.9820 (0.8389)  acc1: 76.1719 (78.7771)  acc5: 92.1875 (94.3634)  time: 0.3856  data: 0.0001  max mem: 16096\n",
+      "Test:  [130/196]  eta: 0:00:26  loss: 1.0825 (0.8705)  acc1: 72.6562 (77.9610)  acc5: 91.4062 (94.0303)  time: 0.3866  data: 0.0001  max mem: 16096\n",
+      "Test:  [140/196]  eta: 0:00:22  loss: 1.1088 (0.8889)  acc1: 72.2656 (77.6125)  acc5: 91.4062 (93.8137)  time: 0.3879  data: 0.0001  max mem: 16096\n",
+      "Test:  [150/196]  eta: 0:00:18  loss: 1.1069 (0.9059)  acc1: 73.0469 (77.2998)  acc5: 91.0156 (93.5586)  time: 0.3914  data: 0.0002  max mem: 16096\n",
+      "Test:  [160/196]  eta: 0:00:14  loss: 1.1360 (0.9197)  acc1: 73.0469 (77.0380)  acc5: 90.2344 (93.3472)  time: 0.3898  data: 0.0002  max mem: 16096\n",
+      "Test:  [170/196]  eta: 0:00:10  loss: 1.2171 (0.9371)  acc1: 71.8750 (76.6265)  acc5: 89.8438 (93.1721)  time: 0.3845  data: 0.0002  max mem: 16096\n",
+      "Test:  [180/196]  eta: 0:00:06  loss: 1.2493 (0.9527)  acc1: 68.7500 (76.2992)  acc5: 90.2344 (93.0205)  time: 0.3815  data: 0.0001  max mem: 16096\n",
+      "Test:  [190/196]  eta: 0:00:02  loss: 1.0816 (0.9511)  acc1: 71.4844 (76.3089)  acc5: 92.1875 (93.0465)  time: 0.3736  data: 0.0001  max mem: 16096\n",
+      "Test: Total time: 0:01:17\n",
+      " * Acc@1 76.426 Acc@5 93.080\n"
+     ]
+    }
+   ],
+   "source": [
+    "with torch.no_grad():\n",
+    "    evaluate(model, criterion, data_loader_test, device=\"cuda\")"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "After only 1 epoch of quantized fine tuning, top-1 improved from ~76.1 to 76.426. Train longer with lr anealing can improve accuracy further"
+   ]
+  }
+ ],
+ "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.6.10"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}

tools/pytorch-quantization/examples/torchvision/classification_flow.py

@@ -0,0 +1,657 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+import datetime
+import inspect
+import os
+import sys
+import time
+import argparse
+import warnings
+import collections
+
+import subprocess
+
+import torch
+import torch.utils.data
+
+from collections import namedtuple
+from torch import nn
+
+from tqdm import tqdm
+
+import torchvision
+from torchvision import transforms
+from torch.hub import load_state_dict_from_url
+
+from pytorch_quantization import nn as quant_nn
+from pytorch_quantization import calib
+from pytorch_quantization.tensor_quant import QuantDescriptor
+from pytorch_quantization import quant_modules
+
+import onnxruntime
+import numpy as np
+import models.classification as models
+
+from prettytable import PrettyTable
+
+# The following path assumes running in nvcr.io/nvidia/pytorch:20.08-py3
+sys.path.insert(0, "/opt/pytorch/vision/references/classification/")
+
+# Import functions from torchvision reference
+try:
+    from train import evaluate, train_one_epoch, load_data, utils
+except Exception as e:
+    raise ModuleNotFoundError(
+        "Add https://github.com/pytorch/vision/blob/master/references/classification/ to PYTHONPATH")
+
+
+def get_parser():
+    """
+    Creates an argument parser.
+    """
+    parser = argparse.ArgumentParser(description='Classification quantization flow script')
+
+    parser.add_argument('--data-dir', '-d', type=str, help='input data folder', required=True)
+    parser.add_argument('--model-name', '-m', default='resnet50', help='model name: default resnet50')
+    parser.add_argument('--disable-pcq',
+                        '-dpcq',
+                        action="store_true",
+                        help='disable per-channel quantization for weights')
+    parser.add_argument('--out-dir', '-o', default='/tmp', help='output folder: default /tmp')
+    parser.add_argument('--print-freq', '-pf', type=int, default=20, help='evaluation print frequency: default 20')
+    parser.add_argument('--threshold',
+                        '-t',
+                        type=float,
+                        default=-1.0,
+                        help='top1 accuracy threshold (less than 0.0 means no comparison): default -1.0')
+
+    parser.add_argument('--fp16', action="store_true", help="Enable FP16 model training, evaluation and export")
+
+    parser.add_argument('--batch-size-train', type=int, default=128, help='batch size for training: default 128')
+    parser.add_argument('--batch-size-test', type=int, default=128, help='batch size for testing: default 128')
+    parser.add_argument('--batch-size-onnx', type=int, default=1, help='batch size for onnx: default 1')
+
+    parser.add_argument('--seed', type=int, default=12345, help='random seed: default 12345')
+
+    checkpoint = parser.add_mutually_exclusive_group(required=True)
+    checkpoint.add_argument('--ckpt-path', default='', type=str, help='path to latest checkpoint (default: none)')
+    checkpoint.add_argument('--ckpt-url', default='', type=str, help='url to latest checkpoint (default: none)')
+    checkpoint.add_argument('--pretrained', action="store_true")
+
+    parser.add_argument('--num-calib-batch',
+                        default=4,
+                        type=int,
+                        help='Number of batches for calibration. 0 will disable calibration. (default: 4)')
+    parser.add_argument('--num-finetune-epochs',
+                        default=0,
+                        type=int,
+                        help='Number of epochs to fine tune. 0 will disable fine tune. (default: 0)')
+    parser.add_argument('--calibrator', type=str, choices=["max", "histogram"], default="max")
+    parser.add_argument('--percentile', nargs='+', type=float, default=[99.9, 99.99, 99.999, 99.9999])
+    parser.add_argument('--sensitivity', action="store_true", help="Build sensitivity profile")
+    parser.add_argument('--evaluate-onnx', action="store_true", help="Evaluate exported ONNX")
+    parser.add_argument('--evaluate-trt', action="store_true", help="Export and evaluate TRT")
+
+    return parser
+
+
+def prepare_model(model_name,
+                  data_dir,
+                  per_channel_quantization,
+                  batch_size_train,
+                  batch_size_test,
+                  batch_size_onnx,
+                  calibrator,
+                  pretrained=True,
+                  ckpt_path=None,
+                  ckpt_url=None,
+                  fp16=False):
+    """
+    Prepare the model for the classification flow.
+    Arguments:
+        model_name: name to use when accessing torchvision model dictionary
+        data_dir: directory with train and val subdirs prepared "imagenet style"
+        per_channel_quantization: iff true use per channel quantization for weights
+                                   note that this isn't currently supported in ONNX-RT/Pytorch
+        batch_size_train: batch size to use when training
+        batch_size_test: batch size to use when testing in Pytorch
+        batch_size_onnx: batch size to use when testing with ONNX-RT
+        calibrator: calibration type to use (max/histogram)
+
+        pretrained: if true a pretrained model will be loaded from torchvision
+        ckpt_path: path to load a model checkpoint from, if not pretrained
+        ckpt_url: url to download a model checkpoint from, if not pretrained and no path was given
+        * at least one of {pretrained, path, url} must be valid
+
+    The method returns a the following list:
+        [
+            Model object,
+            data loader for training,
+            data loader for Pytorch testing,
+            data loader for onnx testing
+        ]
+    """
+    # Use 'spawn' to avoid CUDA reinitialization with forked subprocess
+    torch.multiprocessing.set_start_method('spawn')
+
+    ## Initialize quantization, model and data loaders
+    if per_channel_quantization:
+        quant_desc_input = QuantDescriptor(calib_method=calibrator)
+        quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
+        quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)
+    else:
+        ## Force per tensor quantization for onnx runtime
+        quant_desc_input = QuantDescriptor(calib_method=calibrator, axis=None)
+        quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
+        quant_nn.QuantConvTranspose2d.set_default_quant_desc_input(quant_desc_input)
+        quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)
+
+        quant_desc_weight = QuantDescriptor(calib_method=calibrator, axis=None)
+        quant_nn.QuantConv2d.set_default_quant_desc_weight(quant_desc_weight)
+        quant_nn.QuantConvTranspose2d.set_default_quant_desc_weight(quant_desc_weight)
+        quant_nn.QuantLinear.set_default_quant_desc_weight(quant_desc_weight)
+
+    if model_name in models.__dict__:
+        model = models.__dict__[model_name](pretrained=pretrained, quantize=True)
+    else:
+        quant_modules.initialize()
+        model = torchvision.models.__dict__[model_name](pretrained=pretrained)
+        quant_modules.deactivate()
+
+    if not pretrained:
+        if ckpt_path:
+            checkpoint = torch.load(ckpt_path)
+        else:
+            checkpoint = load_state_dict_from_url(ckpt_url)
+        if 'state_dict' in checkpoint.keys():
+            checkpoint = checkpoint['state_dict']
+        elif 'model' in checkpoint.keys():
+            checkpoint = checkpoint['model']
+        model.load_state_dict(checkpoint)
+    model.eval()
+    model.cuda()
+
+    if fp16:
+        model = model.half()
+
+    ## Prepare the data loaders
+    traindir = os.path.join(data_dir, 'train')
+    valdir = os.path.join(data_dir, 'val')
+    _args = collections.namedtuple("mock_args", [
+        "model", "distributed", "cache_dataset", "val_resize_size", "val_crop_size", "train_crop_size", "interpolation",
+        "ra_magnitude", "augmix_severity", "weights", "backend", "use_v2"
+    ])
+    dataset, dataset_test, train_sampler, test_sampler = load_data(
+        traindir, valdir,
+        _args(model=model_name,
+              distributed=False,
+              cache_dataset=False,
+              val_resize_size=256,
+              val_crop_size=224,
+              train_crop_size=224,
+              interpolation="bilinear",
+              ra_magnitude=9,
+              augmix_severity=3,
+              weights=None,
+              backend="pil",
+              use_v2=False))
+
+    data_loader_train = torch.utils.data.DataLoader(dataset,
+                                                    batch_size=batch_size_train,
+                                                    sampler=train_sampler,
+                                                    num_workers=4,
+                                                    pin_memory=True)
+
+    data_loader_test = torch.utils.data.DataLoader(dataset_test,
+                                                   batch_size=batch_size_test,
+                                                   sampler=test_sampler,
+                                                   num_workers=4,
+                                                   pin_memory=True)
+
+    data_loader_onnx = torch.utils.data.DataLoader(dataset_test,
+                                                   batch_size=batch_size_onnx,
+                                                   sampler=test_sampler,
+                                                   num_workers=4,
+                                                   pin_memory=True)
+
+    return model, data_loader_train, data_loader_test, data_loader_onnx
+
+
+def main(cmdline_args):
+    parser = get_parser()
+    args = parser.parse_args(cmdline_args)
+    print(parser.description)
+    print(args)
+
+    torch.manual_seed(args.seed)
+    np.random.seed(args.seed)
+
+    ## Prepare the pretrained model and data loaders
+    model, data_loader_train, data_loader_test, data_loader_onnx = prepare_model(
+        args.model_name, args.data_dir, not args.disable_pcq, args.batch_size_train, args.batch_size_test,
+        args.batch_size_onnx, args.calibrator, args.pretrained, args.ckpt_path, args.ckpt_url, args.fp16)
+
+    ## Initial accuracy evaluation
+    CrossEntropy = nn.CrossEntropyLoss()
+
+    # nn.CrossEntropyLoss expects float inputs
+    def criterion(output, target):
+        return CrossEntropy(output.float(), target)
+
+    with torch.no_grad():
+        print('Initial evaluation:')
+        top1_initial = evaluate(model, criterion, data_loader_test, device="cuda", print_freq=args.print_freq)
+
+    ## Calibrate the model
+    with torch.no_grad():
+        calibrate_model(model=model,
+                        model_name=args.model_name,
+                        data_loader=data_loader_train,
+                        num_calib_batch=args.num_calib_batch,
+                        calibrator=args.calibrator,
+                        hist_percentile=args.percentile,
+                        out_dir=args.out_dir)
+
+    ## Evaluate after calibration
+    if args.num_calib_batch > 0:
+        with torch.no_grad():
+            print('Calibration evaluation:')
+            top1_calibrated = evaluate(model, criterion, data_loader_test, device="cuda", print_freq=args.print_freq)
+    else:
+        top1_calibrated = -1.0
+
+    ## Build sensitivy profile
+    if args.sensitivity:
+        build_sensitivity_profile(model, criterion, data_loader_test)
+
+    ## Finetune the model
+    criterion = nn.CrossEntropyLoss()
+    optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)
+    lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, args.num_finetune_epochs)
+    for epoch in range(args.num_finetune_epochs):
+        # Training a single epch
+        if "print_freq" in inspect.signature(train_one_epoch).parameters:
+            train_one_epoch(model, criterion, optimizer, data_loader_train, "cuda", 0, 100)
+        else:
+            _args = collections.namedtuple("mock_args",
+                                           ["print_freq", "clip_grad_norm", "model_ema_steps", "lr_warmup_epochs"])
+            train_one_epoch(model, criterion, optimizer, data_loader_train, "cuda", 0,
+                            _args(print_freq=100, clip_grad_norm=None, model_ema_steps=32, lr_warmup_epochs=0))
+        lr_scheduler.step()
+
+    if args.num_finetune_epochs > 0:
+        ## Evaluate after finetuning
+        with torch.no_grad():
+            print('Finetune evaluation:')
+            top1_finetuned = evaluate(model, criterion, data_loader_test, device="cuda")
+    else:
+        top1_finetuned = -1.0
+
+    ## Export to ONNX
+    onnx_filename = args.out_dir + '/' + args.model_name + ".onnx"
+    top1_onnx = -1.0
+    if args.evaluate_onnx and export_onnx(model, onnx_filename, args.batch_size_onnx, not args.disable_pcq):
+        ## Validate ONNX and evaluate
+        top1_onnx = evaluate_onnx(onnx_filename, data_loader_onnx, criterion, args.print_freq)
+
+    trt_filename = args.out_dir + '/' + args.model_name + ".trt"
+    top1_trt = -1.0
+    if args.evaluate_trt and export_trt(model, trt_filename, args.batch_size_onnx, args.fp16):
+        ## Validate TRT and evaluate
+        top1_trt = evaluate_trt(trt_filename, data_loader_onnx, criterion, args.print_freq)
+
+    ## Print summary
+    print("Accuracy summary:")
+    table = PrettyTable(['Stage', 'Top1'])
+    table.align['Stage'] = "l"
+    table.add_row(['Initial', "{:.2f}".format(top1_initial)])
+    table.add_row(['Calibrated', "{:.2f}".format(top1_calibrated)])
+    table.add_row(['Finetuned', "{:.2f}".format(top1_finetuned)])
+    table.add_row(['ONNX', "{:.2f}".format(top1_onnx)])
+    if args.evaluate_trt:
+        table.add_row(['TRT', "{:.2f}".format(top1_trt)])
+    print(table)
+
+    ## Compare results
+    if args.threshold >= 0.0:
+        if args.evaluate_onnx and top1_onnx < 0.0:
+            print("Failed to export/evaluate ONNX!")
+            return 1
+        if args.evaluate_trt and top1_trt < 0.0:
+            print("Failed to export/evaluate TRT!")
+            return 1
+        if args.num_finetune_epochs > 0:
+            if top1_finetuned >= (top1_onnx - args.threshold):
+                print("Accuracy threshold was met!")
+            else:
+                print("Accuracy threshold was missed!")
+                return 1
+
+            if args.evaluate_trt and top1_finetuned >= (top1_trt - args.threshold):
+                print("TRT Accuracy threshold was met!")
+            elif args.evaluate_trt:
+                print("TRT Accuracy threshold was missed!")
+                return 1
+
+    return 0
+
+
+def evaluate_onnx(onnx_filename, data_loader, criterion, print_freq):
+    """Evaluate accuracy on the given ONNX file using the provided data loader and criterion.
+       The method returns the average top-1 accuracy on the given dataset.
+    """
+    print("Loading ONNX file: ", onnx_filename)
+    ort_session = onnxruntime.InferenceSession(onnx_filename, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+    with torch.no_grad():
+        metric_logger = utils.MetricLogger(delimiter="  ")
+        header = 'Test:'
+        with torch.no_grad():
+            for image, target in metric_logger.log_every(data_loader, print_freq, header):
+                image = image.to("cpu", non_blocking=True)
+                image_data = np.array(image)
+                input_data = image_data
+
+                # run the data through onnx runtime instead of torch model
+                input_name = ort_session.get_inputs()[0].name
+                raw_result = ort_session.run([], {input_name: input_data})
+                output = torch.tensor((raw_result[0])).float()
+
+                loss = criterion(output, target)
+                acc1, acc5 = utils.accuracy(output, target, topk=(1, 5))
+                batch_size = image.shape[0]
+                metric_logger.update(loss=loss.item())
+                metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
+                metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
+        # gather the stats from all processes
+        metric_logger.synchronize_between_processes()
+
+        print('  ONNXRuntime: Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f}'.format(top1=metric_logger.acc1,
+                                                                                              top5=metric_logger.acc5))
+        return metric_logger.acc1.global_avg
+
+
+def evaluate_trt(trt_filename, data_loader, criterion, print_freq):
+    print("Loading TRT file: ", trt_filename)
+
+    import pycuda.driver as cuda
+    try:
+        import pycuda.autoprimaryctx
+    except ModuleNotFoundError:
+        import pycuda.autoinit
+
+    import tensorrt as trt
+
+    TRT_LOGGER = trt.Logger()
+
+    TRT_tensor = namedtuple('TRT_tensor', ['binding_idx', 'shape', 'dtype', 'device_memory', 'host_memory'])
+
+    def load_engine(engine_file_path):
+        assert os.path.exists(engine_file_path)
+        print("Reading engine from file {}".format(engine_file_path))
+        with open(engine_file_path, "rb") as f, trt.Runtime(TRT_LOGGER) as runtime:
+            return runtime.deserialize_cuda_engine(f.read())
+
+    def setup_context(engine):
+        return engine.create_execution_context()
+
+    def allocate_buffers(engine, context):
+
+        # Allocate host and device buffers
+        bindings = []
+        inputs = {}
+        outputs = {}
+        for binding_idx in range(engine.num_bindings):
+            binding = engine.get_tensor_name(binding_idx)
+            shape = tuple(context.get_tensor_shape(binding))
+
+            size = trt.volume(context.get_tensor_shape(binding))
+            dtype = np.dtype(trt.nptype(engine.get_tensor_dtype(binding)))
+
+            device_memory = cuda.mem_alloc(size * dtype.itemsize)
+            bindings.append(int(device_memory))
+
+            if engine.get_tensor_mode(binding) == trt.TensorIOMode.INPUT:
+                inputs[binding] = TRT_tensor(binding_idx, shape, dtype, device_memory, None)
+            else:
+                host_memory = cuda.pagelocked_empty(size, dtype)
+                outputs[binding] = TRT_tensor(binding_idx, shape, dtype, device_memory, host_memory)
+
+        stream = cuda.Stream()
+        return bindings, inputs, outputs, stream
+
+    def infer(batch, context, bindings, inputs, outputs, stream):
+
+        # Transfer input data to the GPU.
+        for name, trt_in_t in inputs.items():
+            buffer = np.ascontiguousarray(batch[name])
+            cuda.memcpy_htod_async(trt_in_t.device_memory, buffer, stream)
+
+        # Run inference
+        context.execute_async_v2(bindings=bindings, stream_handle=stream.handle)
+
+        # Transfer predictions back from the GPU.
+        for _, trt_out_t in outputs.items():
+            cuda.memcpy_dtoh_async(trt_out_t.host_memory, trt_out_t.device_memory, stream)
+
+        # Synchronize the stream
+        stream.synchronize()
+
+        return {k: torch.tensor(v.host_memory).reshape(v.shape) for k, v in outputs.items()}
+
+    engine = load_engine(trt_filename)
+    context = setup_context(engine)
+    bindings, inputs, outputs, stream = allocate_buffers(engine, context)
+
+    with torch.no_grad():
+        metric_logger = utils.MetricLogger(delimiter="  ")
+        header = 'Test:'
+        with torch.no_grad():
+            for image, target in metric_logger.log_every(data_loader, print_freq, header):
+                image = image.to("cpu", non_blocking=True)
+                image_data = np.array(image)
+
+                output = infer({"input": image_data}, context, bindings, inputs, outputs, stream)["output"].float()
+
+                loss = criterion(output, target)
+                acc1, acc5 = utils.accuracy(output, target, topk=(1, 5))
+                batch_size = image.shape[0]
+                metric_logger.update(loss=loss.item())
+                metric_logger.meters['acc1'].update(acc1.item(), n=batch_size)
+                metric_logger.meters['acc5'].update(acc5.item(), n=batch_size)
+        # gather the stats from all processes
+        metric_logger.synchronize_between_processes()
+
+        print('  TRTRuntime: Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f}'.format(top1=metric_logger.acc1,
+                                                                                             top5=metric_logger.acc5))
+        return metric_logger.acc1.global_avg
+
+
+def _export_onnx(model, dummy_input, onnx_filename, opset_version):
+    try:
+        if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters:
+            torch.onnx.export(model,
+                              dummy_input,
+                              onnx_filename,
+                              verbose=False,
+                              input_names=["input"],
+                              output_names=["output"],
+                              opset_version=opset_version,
+                              enable_onnx_checker=False,
+                              do_constant_folding=True)
+        else:
+            torch.onnx.export(model,
+                              dummy_input,
+                              onnx_filename,
+                              verbose=False,
+                              input_names=["input"],
+                              output_names=["output"],
+                              opset_version=opset_version,
+                              do_constant_folding=True)
+    except ValueError:
+        print("Failed to export to ONNX")
+        return False
+
+    return True
+
+
+def export_onnx(model, onnx_filename, batch_onnx, per_channel_quantization):
+    model.eval()
+
+    if per_channel_quantization:
+        opset_version = 13
+    else:
+        opset_version = 12
+
+    # Export ONNX for multiple batch sizes
+    print("Creating ONNX file: " + onnx_filename)
+    dummy_input = torch.randn(batch_onnx, 3, 224, 224, device='cuda')  #TODO: switch input dims by model
+    return _export_onnx(model, dummy_input, onnx_filename, opset_version)
+
+
+def export_trt(model, trt_filename, batch_trt, fp16=False):
+    model.eval()
+
+    # Export TRT for multiple batch sizes
+    print("Creating TRT file: " + trt_filename)
+    dummy_input = torch.randn(batch_trt, 3, 224, 224, device='cuda')  #TODO: switch input dims by model
+
+    OPSET = 17
+    onnx_filename = trt_filename.replace(".trt", ".onnx")
+
+    if not _export_onnx(model, dummy_input, onnx_filename, OPSET):
+        return False
+
+    trt_cmd = f"trtexec --onnx={onnx_filename} --saveEngine={trt_filename} --int8"
+
+    if fp16:
+        trt_cmd += " --fp16"
+
+    print(trt_cmd)
+    try:
+        trt_stdout = subprocess.check_output(trt_cmd, shell=True).decode("utf-8")
+    except:
+        print("Failed to export to TRT")
+        return False
+
+    print(trt_stdout)
+    return 'PASSED' in trt_stdout
+
+
+def calibrate_model(model, model_name, data_loader, num_calib_batch, calibrator, hist_percentile, out_dir):
+    """
+        Feed data to the network and calibrate.
+        Arguments:
+            model: classification model
+            model_name: name to use when creating state files
+            data_loader: calibration data set
+            num_calib_batch: amount of calibration passes to perform
+            calibrator: type of calibration to use (max/histogram)
+            hist_percentile: percentiles to be used for historgram calibration
+            out_dir: dir to save state files in
+    """
+
+    if num_calib_batch > 0:
+        print("Calibrating model")
+        with torch.no_grad():
+            collect_stats(model, data_loader, num_calib_batch)
+
+        if not calibrator == "histogram":
+            compute_amax(model, method="max")
+            calib_output = os.path.join(out_dir, F"{model_name}-max-{num_calib_batch*data_loader.batch_size}.pth")
+            torch.save(model.state_dict(), calib_output)
+        else:
+            for percentile in hist_percentile:
+                print(F"{percentile} percentile calibration")
+                compute_amax(model, method="percentile")
+                calib_output = os.path.join(
+                    out_dir, F"{model_name}-percentile-{percentile}-{num_calib_batch*data_loader.batch_size}.pth")
+                torch.save(model.state_dict(), calib_output)
+
+            for method in ["mse", "entropy"]:
+                print(F"{method} calibration")
+                compute_amax(model, method=method)
+                calib_output = os.path.join(out_dir,
+                                            F"{model_name}-{method}-{num_calib_batch*data_loader.batch_size}.pth")
+                torch.save(model.state_dict(), calib_output)
+
+
+def collect_stats(model, data_loader, num_batches):
+    """Feed data to the network and collect statistics"""
+    # 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()
+
+    # Feed data to the network for collecting stats
+    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()
+
+
+def build_sensitivity_profile(model, criterion, data_loader_test):
+    quant_layer_names = []
+    for name, module in model.named_modules():
+        if name.endswith("_quantizer"):
+            module.disable()
+            layer_name = name.replace("._input_quantizer", "").replace("._weight_quantizer", "")
+            if layer_name not in quant_layer_names:
+                quant_layer_names.append(layer_name)
+    for i, quant_layer in enumerate(quant_layer_names):
+        print("Enable", quant_layer)
+        for name, module in model.named_modules():
+            if name.endswith("_quantizer") and quant_layer in name:
+                module.enable()
+                print(F"{name:40}: {module}")
+        with torch.no_grad():
+            evaluate(model, criterion, data_loader_test, device="cuda")
+        for name, module in model.named_modules():
+            if name.endswith("_quantizer") and quant_layer in name:
+                module.disable()
+                print(F"{name:40}: {module}")
+
+
+if __name__ == '__main__':
+    res = main(sys.argv[1:])
+    exit(res)

tools/pytorch-quantization/examples/torchvision/models/__init__.py

@@ -0,0 +1,18 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from . import classification
+

tools/pytorch-quantization/examples/torchvision/models/classification/__init__.py

@@ -0,0 +1,18 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+from .resnet import *
+

tools/pytorch-quantization/examples/torchvision/models/classification/resnet.py

@@ -0,0 +1,495 @@
+#
+# 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.
+#
+
+# BSD 3-Clause License
+#
+# Copyright (c) Soumith Chintala 2016,
+# All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#
+# * Redistributions of source code must retain the above copyright notice, this
+#   list of conditions and the following disclaimer.
+#
+# * Redistributions in binary form must reproduce the above copyright notice,
+#   this list of conditions and the following disclaimer in the documentation
+#   and/or other materials provided with the distribution.
+#
+# * Neither the name of the copyright holder nor the names of its
+#   contributors may be used to endorse or promote products derived from
+#   this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+import torch
+from torch import Tensor
+import torch.nn as nn
+from torch.hub import load_state_dict_from_url
+from typing import Type, Any, Callable, Union, List, Optional
+from pytorch_quantization import quant_modules
+from pytorch_quantization import nn as quant_nn
+
+__all__ = [
+    'ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152', 'resnext50_32x4d', 'resnext101_32x8d',
+    'wide_resnet50_2', 'wide_resnet101_2'
+]
+
+model_urls = {
+    'resnet18': 'https://download.pytorch.org/models/resnet18-f37072fd.pth',
+    'resnet34': 'https://download.pytorch.org/models/resnet34-b627a593.pth',
+    'resnet50': 'https://download.pytorch.org/models/resnet50-0676ba61.pth',
+    'resnet101': 'https://download.pytorch.org/models/resnet101-63fe2227.pth',
+    'resnet152': 'https://download.pytorch.org/models/resnet152-394f9c45.pth',
+    'resnext50_32x4d': 'https://download.pytorch.org/models/resnext50_32x4d-7cdf4587.pth',
+    'resnext101_32x8d': 'https://download.pytorch.org/models/resnext101_32x8d-8ba56ff5.pth',
+    'wide_resnet50_2': 'https://download.pytorch.org/models/wide_resnet50_2-95faca4d.pth',
+    'wide_resnet101_2': 'https://download.pytorch.org/models/wide_resnet101_2-32ee1156.pth',
+}
+
+
+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)
+
+
+class BasicBlock(nn.Module):
+    expansion: int = 1
+
+    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:
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride, quantize=quantize)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes, quantize=quantize)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+        self._quantize = quantize
+        if self._quantize:
+            self.residual_quantizer = quant_nn.TensorQuantizer(quant_nn.QuantConv2d.default_quant_desc_input)
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        if self._quantize:
+            out += self.residual_quantizer(identity)
+        else:
+            out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class Bottleneck(nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+
+    expansion: int = 4
+
+    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:
+        super(Bottleneck, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width, quantize=quantize)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation, quantize=quantize)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, planes * self.expansion, quantize=quantize)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+        self._quantize = quantize
+        if self._quantize:
+            self.residual_quantizer = quant_nn.TensorQuantizer(quant_nn.QuantConv2d.default_quant_desc_input)
+
+    def forward(self, x: Tensor) -> Tensor:
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        if self._quantize:
+            out += self.residual_quantizer(identity)
+        else:
+            out += identity
+        out = self.relu(out)
+
+        return out
+
+
+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
+
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.inplanes = 64
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+
+        if quantize:
+            self.conv1 = quant_nn.QuantConv2d(3,
+                                        self.inplanes,
+                                        kernel_size=7,
+                                        stride=2,
+                                        padding=3,
+                                        bias=False)
+        else:
+            self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
+
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+        self.layer1 = self._make_layer(block, 64, layers[0], quantize=quantize)
+        self.layer2 = self._make_layer(block,
+                                       128,
+                                       layers[1],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[0],
+                                       quantize=quantize)
+        self.layer3 = self._make_layer(block,
+                                       256,
+                                       layers[2],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[1],
+                                       quantize=quantize)
+        self.layer4 = self._make_layer(block,
+                                       512,
+                                       layers[3],
+                                       stride=2,
+                                       dilate=replace_stride_with_dilation[2],
+                                       quantize=quantize)
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+
+        if quantize:
+            self.fc = quant_nn.QuantLinear(512 * block.expansion, num_classes)
+        else:
+            self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)  # type: ignore[arg-type]
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)  # type: ignore[arg-type]
+
+    def _make_layer(self,
+                    block: Type[Union[BasicBlock, Bottleneck]],
+                    planes: int,
+                    blocks: int,
+                    stride: int = 1,
+                    dilate: bool = False,
+                    quantize: bool = False) -> nn.Sequential:
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride, quantize=quantize),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(
+            block(self.inplanes, planes, stride, downsample, self.groups, self.base_width, previous_dilation,
+                  norm_layer, self._quantize))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(
+                block(self.inplanes,
+                      planes,
+                      groups=self.groups,
+                      base_width=self.base_width,
+                      dilation=self.dilation,
+                      norm_layer=norm_layer,
+                      quantize=quantize))
+
+        return nn.Sequential(*layers)
+
+    def _get_dtype(self):
+        return self.conv1.weight.dtype
+
+    def _forward_impl(self, x: Tensor) -> Tensor:
+        # See note [TorchScript super()]
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+        x = self.maxpool(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return x
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self._forward_impl(x.to(self._get_dtype()))
+
+
+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)
+    if pretrained:
+        state_dict = load_state_dict_from_url(model_urls[arch], progress=progress)
+        model.load_state_dict(state_dict)
+    return model
+
+
+def resnet18(pretrained: bool = False, progress: bool = True, quantize: bool = False, **kwargs: Any) -> ResNet:
+    r"""ResNet-18 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet18', BasicBlock, [2, 2, 2, 2], pretrained, progress, quantize, **kwargs)
+
+
+def resnet34(pretrained: bool = False, progress: bool = True, quantize: bool = False, **kwargs: Any) -> ResNet:
+    r"""ResNet-34 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet34', BasicBlock, [3, 4, 6, 3], pretrained, progress, quantize, **kwargs)
+
+
+def resnet50(pretrained: bool = False, progress: bool = True, quantize: bool = False, **kwargs: Any) -> ResNet:
+    r"""ResNet-50 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet50', Bottleneck, [3, 4, 6, 3], pretrained, progress, quantize, **kwargs)
+
+
+def resnet101(pretrained: bool = False, progress: bool = True, quantize: bool = False, **kwargs: Any) -> ResNet:
+    r"""ResNet-101 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet101', Bottleneck, [3, 4, 23, 3], pretrained, progress, quantize, **kwargs)
+
+
+def resnet152(pretrained: bool = False, progress: bool = True, quantize: bool = False, **kwargs: Any) -> ResNet:
+    r"""ResNet-152 model from
+    `"Deep Residual Learning for Image Recognition" <https://arxiv.org/pdf/1512.03385.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    return _resnet('resnet152', Bottleneck, [3, 8, 36, 3], pretrained, progress, quantize, **kwargs)
+
+
+def resnext50_32x4d(pretrained: bool = False, progress: bool = True, quantize: bool = False, **kwargs: Any) -> ResNet:
+    r"""ResNeXt-50 32x4d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['groups'] = 32
+    kwargs['width_per_group'] = 4
+    return _resnet('resnext50_32x4d', Bottleneck, [3, 4, 6, 3], pretrained, progress, quantize, **kwargs)
+
+
+def resnext101_32x8d(pretrained: bool = False, progress: bool = True, quantize: bool = False, **kwargs: Any) -> ResNet:
+    r"""ResNeXt-101 32x8d model from
+    `"Aggregated Residual Transformation for Deep Neural Networks" <https://arxiv.org/pdf/1611.05431.pdf>`_.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['groups'] = 32
+    kwargs['width_per_group'] = 8
+    return _resnet('resnext101_32x8d', Bottleneck, [3, 4, 23, 3], pretrained, progress, quantize, **kwargs)
+
+
+def wide_resnet50_2(pretrained: bool = False, progress: bool = True, quantize: bool = False, **kwargs: Any) -> ResNet:
+    r"""Wide ResNet-50-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.
+
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 2
+    return _resnet('wide_resnet50_2', Bottleneck, [3, 4, 6, 3], pretrained, progress, quantize, **kwargs)
+
+
+def wide_resnet101_2(pretrained: bool = False, progress: bool = True, quantize: bool = False, **kwargs: Any) -> ResNet:
+    r"""Wide ResNet-101-2 model from
+    `"Wide Residual Networks" <https://arxiv.org/pdf/1605.07146.pdf>`_.
+
+    The model is the same as ResNet except for the bottleneck number of channels
+    which is twice larger in every block. The number of channels in outer 1x1
+    convolutions is the same, e.g. last block in ResNet-50 has 2048-512-2048
+    channels, and in Wide ResNet-50-2 has 2048-1024-2048.
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+    """
+    kwargs['width_per_group'] = 64 * 2
+    return _resnet('wide_resnet101_2', Bottleneck, [3, 4, 23, 3], pretrained, progress, quantize, **kwargs)

tools/pytorch-quantization/pytorch_quantization/__init__.py

@@ -0,0 +1,22 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from absl import logging
+from .version import __version__
+from .quant_modules import *
+
+logging.use_absl_handler()

tools/pytorch-quantization/pytorch_quantization/calib/__init__.py

@@ -0,0 +1,24 @@
+#
+# 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.
+#
+
+
+"""``pytorch_quantization.calib`` provides Calibrator classes that
+collect data statistics and determine pytorch_quantization parameters.
+"""
+
+from .max import MaxCalibrator
+from .histogram import *

tools/pytorch-quantization/pytorch_quantization/calib/calibrator.py

@@ -0,0 +1,61 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+"""Abstract base class for calibrators"""
+
+
+class _Calibrator():
+    """Abstract base class of calibrators
+    Args:
+        num_bits: An integer. Number of bits of quantization.
+        axis: A tuple. see QuantDescriptor.
+        unsigned: A boolean. using unsigned quantization.
+
+    Readonly Properties:
+        axis:
+    """
+    def __init__(self, num_bits, axis, unsigned):
+        self._num_bits = num_bits
+        self._axis = axis
+        self._unsigned = unsigned
+
+    def collect(self, x):
+        """Abstract method: collect tensor statistics used to compute amax
+
+        Args:
+            x: A tensor
+        """
+        raise NotImplementedError
+
+    def reset(self):
+        """Abstract method: reset calibrator to initial state"""
+        raise NotImplementedError
+
+    def compute_amax(self, *args, **kwargs):
+        """Abstract method: compute the amax from the collected data
+
+        Returns:
+            amax: a tensor
+        """
+        raise NotImplementedError
+
+    def __repr__(self):
+        s = "num_bits={_num_bits}"
+        s += " axis={_axis}"
+        s += " unsigned={_unsigned}"
+        return s.format(**self.__dict__)

tools/pytorch-quantization/pytorch_quantization/calib/histogram.py

@@ -0,0 +1,386 @@
+#
+# 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.
+#
+"""Histogram based calibrators"""
+from collections import Counter
+import numpy as np
+from scipy.stats import entropy
+
+from absl import logging
+
+import torch
+
+from pytorch_quantization.calib.calibrator import _Calibrator
+from pytorch_quantization.tensor_quant import fake_tensor_quant, scaled_e4m3
+from pytorch_quantization import nn as quant_nn
+from pytorch_quantization import utils as quant_utils
+
+__all__ = ["HistogramCalibrator", "calibrate_weights"]
+
+
+class HistogramCalibrator(_Calibrator):
+    """Unified histogram calibrator
+
+    Histogram will be only collected once. compute_amax() performs entropy, percentile, or mse
+        calibration based on arguments
+
+    Args:
+        num_bits: An integer. Number of bits of quantization.
+        axis: A tuple. see QuantDescriptor.
+        unsigned: A boolean. using unsigned quantization.
+        num_bins: An integer. Number of histograms bins. Default 2048.
+        grow_method: A string. DEPRECATED. default None.
+        skip_zeros: A boolean. If True, skips zeros when collecting data for histogram. Default False.
+        torch_hist: A boolean. If True, collect histogram by torch.histc instead of np.histogram. If input tensor
+            is on GPU, histc will also be running on GPU. Default True.
+    """
+
+    def __init__(self, num_bits, axis, unsigned, num_bins=2048, grow_method=None, skip_zeros=False, torch_hist=True):
+        super(HistogramCalibrator, self).__init__(num_bits, axis, unsigned)
+        self._num_bins = num_bins
+        self._skip_zeros = skip_zeros
+
+        self._calib_bin_edges = None
+        self._calib_hist = None
+
+        self._torch_hist = torch_hist
+
+        if axis is not None:
+            raise NotImplementedError("Calibrator histogram collection only supports per tensor scaling")
+
+        if grow_method is not None:
+            logging.warning("grow_method is deprecated. Got %s, ingored!", grow_method)
+
+    def collect(self, x):
+        """Collect histogram"""
+        if torch.min(x) < 0.:
+            logging.log_first_n(logging.INFO,
+                                ("Calibrator encountered negative values. It shouldn't happen after ReLU. "
+                                 "Make sure this is the right tensor to calibrate."), 1)
+            x = x.abs()
+
+        x = x.float()
+
+        if not self._torch_hist:
+            x_np = x.cpu().detach().numpy()
+
+            if self._skip_zeros:
+                x_np = x_np[np.where(x_np != 0)]
+
+            if self._calib_bin_edges is None and self._calib_hist is None:
+                # first time it uses num_bins to compute histogram.
+                self._calib_hist, self._calib_bin_edges = np.histogram(x_np, bins=self._num_bins)
+            else:
+                temp_amax = np.max(x_np)
+                if temp_amax > self._calib_bin_edges[-1]:
+                    # increase the number of bins
+                    width = self._calib_bin_edges[1] - self._calib_bin_edges[0]
+                    # NOTE: np.arange may create an extra bin after the one containing temp_amax
+                    new_bin_edges = np.arange(self._calib_bin_edges[-1] + width, temp_amax + width, width)
+                    self._calib_bin_edges = np.hstack((self._calib_bin_edges, new_bin_edges))
+                hist, self._calib_bin_edges = np.histogram(x_np, bins=self._calib_bin_edges)
+                hist[:len(self._calib_hist)] += self._calib_hist
+                self._calib_hist = hist
+        else:
+            # This branch of code is designed to match numpy version as close as possible
+            with torch.no_grad():
+                if self._skip_zeros:
+                    x = x[torch.where(x != 0)]
+
+                # Because we collect histogram on absolute value, setting min=0 simplifying the rare case where
+                # minimum value is not exactly 0 and first batch collected has larger min value than later batches
+                x_max = x.max()
+                if self._calib_bin_edges is None and self._calib_hist is None:
+                    self._calib_hist = torch.histc(x, bins=self._num_bins, min=0, max=x_max)
+                    self._calib_bin_edges = torch.linspace(0, x_max, self._num_bins + 1)
+                else:
+                    if x_max > self._calib_bin_edges[-1]:
+                        width = self._calib_bin_edges[1] - self._calib_bin_edges[0]
+                        self._num_bins = int((x_max / width).ceil().item())
+                        self._calib_bin_edges = torch.arange(0, x_max + width, width, device=x.device)
+
+                    hist = torch.histc(x, bins=self._num_bins, min=0, max=self._calib_bin_edges[-1])
+                    hist[:self._calib_hist.numel()] += self._calib_hist
+                    self._calib_hist = hist
+
+    def reset(self):
+        """Reset the collected histogram"""
+        self._calib_bin_edges = None
+        self._calib_hist = None
+
+    def compute_amax(self, method: str, *, stride: int = 1, start_bin: int = 128, percentile: float = 99.99):
+        """Compute the amax from the collected histogram
+
+        Args:
+            method: A string. One of ['entropy', 'mse', 'percentile']
+
+        Keyword Arguments:
+            stride: An integer. Default 1
+            start_bin: An integer. Default 128
+            percentils: A float number between [0, 100]. Default 99.99.
+
+        Returns:
+            amax: a tensor
+        """
+        if isinstance(self._calib_hist, torch.Tensor):
+            calib_hist = self._calib_hist.to(torch.int64).cpu().numpy()
+            calib_bin_edges = self._calib_bin_edges.cpu().numpy()
+        else:
+            calib_hist = self._calib_hist
+            calib_bin_edges = self._calib_bin_edges
+
+        if method == 'entropy':
+            calib_amax = _compute_amax_entropy(calib_hist, calib_bin_edges, self._num_bits, self._unsigned, stride,
+                                               start_bin)
+        elif method == 'mse':
+            calib_amax = _compute_amax_mse(calib_hist, calib_bin_edges, self._num_bits, self._unsigned, stride,
+                                           start_bin)
+        elif method == 'percentile':
+            calib_amax = _compute_amax_percentile(calib_hist, calib_bin_edges, percentile)
+        else:
+            raise TypeError("Unknown calibration method {}".format(method))
+
+        return calib_amax
+
+    # pylint:disable=missing-docstring
+    def __str__(self):
+        s = "HistogramCalibrator("
+        if self._calib_bin_edges is None:
+            bin_edge_str = "None"
+        else:
+            bin_edge_str = "[{:.3f}, ..., {:.3f}]({})".format(self._calib_bin_edges[0], self._calib_bin_edges[-1],
+                                                              len(self._calib_bin_edges))
+        s += "calib_bin_edges={})".format(bin_edge_str)
+        return s
+
+    def __repr__(self):
+        s = "HistogramCalibrator("
+        s += super(HistogramCalibrator, self).__repr__()
+        s += " calib_bin_edges={_calib_bin_edges}"
+        s += " calib_hist={_calib_hist})"
+        return s.format(**self.__dict__)
+
+    # pylint:enable=missing-docstring
+
+
+# Ideally, we want to decouple collector (collect histogram) and calibrator (compute amax) as opposed to
+# the current calibrator design. The following compute amax functions are broken out from the calibrator
+# as first step towards there.
+def _compute_amax_entropy(calib_hist, calib_bin_edges, num_bits, unsigned, stride=1, start_bin=128):
+    """Returns amax that minimizes KL-Divergence of the collected histogram"""
+
+    # If calibrator hasn't collected any data, return none
+    if calib_bin_edges is None and calib_hist is None:
+        return None
+
+    def _normalize_distr(distr):
+        summ = np.sum(distr)
+        if summ != 0:
+            distr = distr / summ
+
+    bins = calib_hist[:]
+    bins[0] = bins[1]
+
+    total_data = np.sum(bins)
+
+    divergences = []
+    arguments = []
+
+    # we are quantizing to 128 values + sign if num_bits=8
+    nbins = 1 << (num_bits - 1 + int(unsigned))
+
+    starting = start_bin
+    stop = len(bins)
+
+    new_density_counts = np.zeros(nbins, dtype=np.float64)
+
+    for i in range(starting, stop + 1, stride):
+        new_density_counts.fill(0)
+        space = np.linspace(0, i, num=nbins + 1)
+        digitized_space = np.digitize(range(i), space) - 1
+
+        digitized_space[bins[:i] == 0] = -1
+
+        for idx, digitized in enumerate(digitized_space):
+            if digitized != -1:
+                new_density_counts[digitized] += bins[idx]
+
+        counter = Counter(digitized_space)
+        for key, val in counter.items():
+            if key != -1:
+                new_density_counts[key] = new_density_counts[key] / val
+
+        new_density = np.zeros(i, dtype=np.float64)
+        for idx, digitized in enumerate(digitized_space):
+            if digitized != -1:
+                new_density[idx] = new_density_counts[digitized]
+
+        total_counts_new = np.sum(new_density) + np.sum(bins[i:])
+        _normalize_distr(new_density)
+
+        reference_density = np.array(bins[:len(digitized_space)])
+        reference_density[-1] += np.sum(bins[i:])
+
+        total_counts_old = np.sum(reference_density)
+        if round(total_counts_new) != total_data or round(total_counts_old) != total_data:
+            raise RuntimeError("Count mismatch! total_counts_new={}, total_counts_old={}, total_data={}".format(
+                total_counts_new, total_counts_old, total_data))
+
+        _normalize_distr(reference_density)
+
+        ent = entropy(reference_density, new_density)
+        divergences.append(ent)
+        arguments.append(i)
+
+    divergences = np.array(divergences)
+    logging.debug("divergences={}".format(divergences))
+    last_argmin = len(divergences) - 1 - np.argmin(divergences[::-1])
+    calib_amax = calib_bin_edges[last_argmin * stride + starting]
+    calib_amax = torch.tensor(calib_amax.item())  #pylint: disable=not-callable
+
+    return calib_amax
+
+
+def _compute_amax_mse(calib_hist, calib_bin_edges, num_bits, unsigned, stride=1, start_bin=128):
+    """Returns amax that minimizes MSE of the collected histogram"""
+
+    # If calibrator hasn't collected any data, return none
+    if calib_bin_edges is None and calib_hist is None:
+        return None
+
+    counts = torch.from_numpy(calib_hist[:]).float().cuda()
+    edges = torch.from_numpy(calib_bin_edges[:]).float().cuda()
+    centers = (edges[1:] + edges[:-1]) / 2
+
+    mses = []
+    arguments = []
+
+    for i in range(start_bin, len(centers), stride):
+
+        amax = centers[i]
+        if isinstance(num_bits, int) and num_bits >= 0:
+            if num_bits == 0:
+                logging.error("num_bits is 0. This will result in the tensor being quantized to all zeros."
+                              " This mode should only be used for debugging purposes.")
+            quant_centers = fake_tensor_quant(centers, amax, num_bits, unsigned)
+        elif num_bits == (4, 3):
+            quant_centers = scaled_e4m3(centers, amax, num_bits[0], num_bits[1])
+        else:
+            raise TypeError("Invalid num_bits. num_bits must be a postivie integer or tuple (4,3).")
+
+        mse = ((quant_centers - centers)**2 * counts).mean()
+
+        mses.append(mse.cpu())
+        arguments.append(i)
+
+    logging.debug("mses={}".format(mses))
+    argmin = np.argmin(mses)
+    calib_amax = centers[arguments[argmin]]
+
+    return calib_amax
+
+
+def _compute_amax_percentile(calib_hist, calib_bin_edges, percentile):
+    """Returns amax that clips the percentile fraction of collected data"""
+
+    if percentile < 0 or percentile > 100:
+        raise ValueError("Invalid percentile. Must be in range 0 <= percentile <= 100.")
+
+    # If calibrator hasn't collected any data, return none
+    if calib_bin_edges is None and calib_hist is None:
+        return None
+
+    total = calib_hist.sum()
+    cdf = np.cumsum(calib_hist / total)
+    idx = np.searchsorted(cdf, percentile / 100)
+    calib_amax = calib_bin_edges[idx]
+    calib_amax = torch.tensor(calib_amax.item())  #pylint: disable=not-callable
+
+    return calib_amax
+
+
+def calibrate_weights(model, method="percentile", perchannel=True, percentile=99.99, num_bins=2048):
+    """Calibrate weights of all child quantized modules
+
+    Ideally, we would split calibration functionality to histogram collector and calibrator which
+    takes histogram and compute amax. But since we haven't decoupled collector and calibrator, it
+    is easier to create a separate function to calibrate weight.
+
+    .. note::
+        This function uses `method` specified by the argument to decide which method to use, NOT the one
+        specified by the calibrator embedded in weight_quantizer.
+        We haven't moved calibration to GPU, so everything is transfered to CPU
+
+    Args:
+        model: A torch.nn.Module.
+        method: A string of calibration method. Supports "mse" and "percentile". Default "percentile"
+        perchannel: A bool. Set channel/neuron axis if True. Default True.
+        percentile: A float. Default 99.99
+        num_bins: A integer. Number of bins of histogram. Default 2048.
+
+    """
+    for name, module in model.named_modules():
+        if hasattr(module, "weight") and hasattr(module, "weight_quantizer"):
+            logging.info("Calibrate weight of %s", name)
+            num_bits = module.weight_quantizer.num_bits
+            unsigned = module.weight_quantizer.unsigned
+            channel_second_modules = (quant_nn.QuantConvTranspose1d, quant_nn.QuantConvTranspose2d,
+                                      quant_nn.QuantConvTranspose3d)
+            if perchannel:
+                axis = 1 if isinstance(module, channel_second_modules) else 0
+            else:
+                axis = None
+            axis_size = module.weight.shape[axis] if axis is not None else 1
+
+            # Histogram is always collected even if method is "max". Although "max" is supported here
+            # but it is not the primary usage of this function
+            if axis is None:
+                input_weights = module.weight.abs().cpu().detach().numpy()
+                calib_hist, calib_bin_edges = np.histogram(input_weights, bins=2048, range=(0, input_weights.max()))
+                calib_hist = [calib_hist]
+                calib_bin_edges = [calib_bin_edges]
+            else:
+                calib_hist = []
+                calib_bin_edges = []
+                for i in range(axis_size):
+                    input_weights = module.weight.index_select(axis, torch.tensor(
+                        i, device=module.weight.device)).abs().cpu().detach().numpy()
+                    hist, bin_edges = np.histogram(input_weights, bins=num_bins, range=(0, input_weights.max()))
+                    calib_hist.append(hist)
+                    calib_bin_edges.append(bin_edges)
+
+            calib_amax = []
+            if method == "max":
+                reduce_axis = list(range(module.weight.dim()))
+                reduce_axis.remove(axis)
+                calib_amax.append(quant_utils.reduce_amax(module.weight, axis=reduce_axis))
+            elif method == 'mse':
+                for i in range(axis_size):
+                    calib_amax.append(_compute_amax_mse(calib_hist[i], calib_bin_edges[i], num_bits, unsigned))
+            elif method == 'percentile':
+                for i in range(axis_size):
+                    calib_amax.append(_compute_amax_percentile(calib_hist[i], calib_bin_edges[i], percentile))
+            else:
+                raise TypeError("Unsupported calibration method {}".format(method))
+
+            if axis is None:
+                calib_amax = calib_amax[0]
+            else:
+                calib_amax_shape = [1] * module.weight.dim()
+                calib_amax_shape[axis] = module.weight.shape[axis]
+                calib_amax = torch.stack(calib_amax).reshape(calib_amax_shape)
+            module.weight_quantizer.amax = calib_amax.detach().cpu().numpy()
+

tools/pytorch-quantization/pytorch_quantization/calib/max.py

@@ -0,0 +1,111 @@
+#
+# 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.
+#
+
+
+"""Calibrator that returns the absolute max of all collected tensors"""
+from absl import logging
+import torch
+
+from pytorch_quantization.calib.calibrator import _Calibrator
+from pytorch_quantization import utils as quant_utils
+
+class MaxCalibrator(_Calibrator):
+    """Max calibrator, tracks the maximum value globally
+
+    Args:
+        calib_desc: A MaxCalibDescriptor.
+        num_bits: An integer. Number of bits of quantization.
+        axis: A tuple. see QuantDescriptor.
+        unsigned: A boolean. using unsigned quantization.
+
+    Readonly Properties:
+        amaxs: A list of amax. Numpy array is saved as it is likely to be used for some plot.
+    """
+    def __init__(self, num_bits, axis, unsigned, track_amax=False):
+        super(MaxCalibrator, self).__init__(num_bits, axis, unsigned)
+        self._track_amax = track_amax
+        if self._track_amax:
+            self._amaxs = []  # shall we have a better name?
+        self._calib_amax = None
+
+    # pylint:disable=missing-docstring
+    @property
+    def amaxs(self):
+        return self._amaxs
+    # pylint:enable=missing-docstring
+
+    def collect(self, x):
+        """Tracks the absolute max of all tensors
+
+        Args:
+            x: A tensor
+
+        Raises:
+            RuntimeError: If amax shape changes
+        """
+        if torch.min(x) < 0.:
+            logging.log_first_n(
+                logging.INFO,
+                ("Calibrator encountered negative values. It shouldn't happen after ReLU. "
+                 "Make sure this is the right tensor to calibrate."),
+                1)
+            x = x.abs()
+
+        # Swap axis to reduce.
+        axis = self._axis if isinstance(self._axis, (list, tuple)) else [self._axis]
+        # Handle negative axis.
+        axis = [x.dim() + i if isinstance(i, int) and i < 0 else i for i in axis]
+        reduce_axis = []
+        for i in range(x.dim()):
+            if not i in axis:
+                reduce_axis.append(i)
+        local_amax = quant_utils.reduce_amax(x, axis=reduce_axis).detach()
+        if self._calib_amax is None:
+            self._calib_amax = local_amax
+        else:
+            if local_amax.shape != self._calib_amax.shape:
+                raise RuntimeError("amax shape changed!")
+            self._calib_amax.copy_(torch.max(self._calib_amax, local_amax).data)
+
+        if self._track_amax:
+            self._amaxs.append(local_amax.cpu().numpy())
+
+    def reset(self):
+        """Reset the collected absolute max"""
+        self._calib_amax = None
+
+    def compute_amax(self):
+        """Return the absolute max of all tensors collected"""
+        return self._calib_amax
+
+    # pylint:disable=missing-docstring
+    def __str__(self):
+        s = "MaxCalibrator("
+        s += "track_amax={_track_amax}"
+        s += ")"
+        return s.format(**self.__dict__)
+
+    def __repr__(self):
+        s = "MaxCalibrator("
+        s += super(MaxCalibrator, self).__repr__()
+        s += " calib_amax={_calib_amax}"
+        s += " track_amax={_track_amax}"
+        if self._track_amax:
+            s += " amaxs={_amaxs}"
+        s += ")"
+        return s.format(**self.__dict__)
+    # pylint:enable=missing-docstring

tools/pytorch-quantization/pytorch_quantization/nn/__init__.py

@@ -0,0 +1,25 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+from pytorch_quantization.nn.modules.tensor_quantizer import *
+from pytorch_quantization.nn.modules.quant_conv import *
+from pytorch_quantization.nn.modules.quant_linear import *
+from pytorch_quantization.nn.modules.quant_pooling import *
+from pytorch_quantization.nn.modules.clip import *
+from pytorch_quantization.nn.modules.quant_rnn import *
+from pytorch_quantization.nn.modules.quant_instancenorm import *

tools/pytorch-quantization/pytorch_quantization/nn/_functions/quant_rnn.py

@@ -0,0 +1,286 @@
+#
+# 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.
+#
+
+
+"""RNN implementation in python
+Originally copied from https://github.com/pytorch/pytorch/blob/v0.4.1/torch/nn/_functions/rnn.py
+with following modification
+    fusedBackend is removed
+    CudnnRNN is removed
+    Hack for ONNX in RNN() is removed
+Only LSTM is quantized. Other paths are excluded in __all__
+"""
+
+import warnings
+from torch.autograd import NestedIOFunction
+from torch.nn import functional as F
+import torch
+import itertools
+from functools import partial
+
+__all__ = ["LSTMCell", "RNN"]
+
+def RNNReLUCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None):
+    hy = F.relu(F.linear(input, w_ih, b_ih) + F.linear(hidden, w_hh, b_hh))
+    return hy
+
+
+def RNNTanhCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None):
+    hy = torch.tanh(F.linear(input, w_ih, b_ih) + F.linear(hidden, w_hh, b_hh))
+    return hy
+
+
+def LSTMCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None, input_quantizer=None, weight_quantizer=None):
+    """Quantized LSTM Cell
+
+    The assumption is at inference time, only one fused gemm will be launched for one time step Weights of 4 gates
+    are fused together, and activation from layer and recurrent paths are fused togather. ``input_quantizer`` will be
+    applied on the fused activation tensor. And ``weight_quantizer`` will be applied on the fused weight tensor.
+    """
+
+    hx, cx = hidden
+    if input_quantizer is not None:
+        input, hx = input_quantizer(torch.cat([input, hx], 1)).split([input.size()[1], hx.size()[1]], 1)
+    if weight_quantizer is not None:
+        w_ih, w_hh = weight_quantizer(torch.cat([w_ih, w_hh], 1)).split([w_ih.size()[1], w_hh.size()[1]], 1)
+    gates = F.linear(input, w_ih, b_ih) + F.linear(hx, w_hh, b_hh)
+
+    ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)
+
+    ingate = torch.sigmoid(ingate)
+    forgetgate = torch.sigmoid(forgetgate)
+    cellgate = torch.tanh(cellgate)
+    outgate = torch.sigmoid(outgate)
+
+    cy = (forgetgate * cx) + (ingate * cellgate)
+    hy = outgate * torch.tanh(cy)
+
+    return hy, cy
+
+
+def GRUCell(input, hidden, w_ih, w_hh, b_ih=None, b_hh=None):
+    gi = F.linear(input, w_ih, b_ih)
+    gh = F.linear(hidden, w_hh, b_hh)
+    i_r, i_i, i_n = gi.chunk(3, 1)
+    h_r, h_i, h_n = gh.chunk(3, 1)
+
+    resetgate = torch.sigmoid(i_r + h_r)
+    inputgate = torch.sigmoid(i_i + h_i)
+    newgate = torch.tanh(i_n + resetgate * h_n)
+    hy = newgate + inputgate * (hidden - newgate)
+
+    return hy
+
+
+def StackedRNN(inners, num_layers, lstm=False, dropout=0, train=True):
+
+    num_directions = len(inners)
+    total_layers = num_layers * num_directions
+
+    def forward(input, hidden, weight, batch_sizes, input_quantizers, weight_quantizers):
+        assert(len(weight) == total_layers)
+        next_hidden = []
+
+        if lstm:
+            hidden = list(zip(*hidden))
+
+        for i in range(num_layers):
+            all_output = []
+            for j, inner in enumerate(inners):
+                l = i * num_directions + j
+
+                hy, output = inner(input, hidden[l], weight[l], batch_sizes,
+                                   input_quantizer=input_quantizers[l], weight_quantizer=weight_quantizers[l])
+                next_hidden.append(hy)
+                all_output.append(output)
+
+            input = torch.cat(all_output, input.dim() - 1)
+
+            if dropout != 0 and i < num_layers - 1:
+                input = F.dropout(input, p=dropout, training=train, inplace=False)
+
+        if lstm:
+            next_h, next_c = zip(*next_hidden)
+            next_hidden = (
+                torch.cat(next_h, 0).view(total_layers, *next_h[0].size()),
+                torch.cat(next_c, 0).view(total_layers, *next_c[0].size())
+            )
+        else:
+            next_hidden = torch.cat(next_hidden, 0).view(
+                total_layers, *next_hidden[0].size())
+
+        return next_hidden, input
+
+    return forward
+
+
+def Recurrent(inner, reverse=False):
+    def forward(input, hidden, weight, batch_sizes, input_quantizer, weight_quantizer):
+        output = []
+        steps = range(input.size(0) - 1, -1, -1) if reverse else range(input.size(0))
+        for i in steps:
+            hidden = inner(input[i], hidden, *weight,
+                           input_quantizer=input_quantizer, weight_quantizer=weight_quantizer)
+            # hack to handle LSTM
+            output.append(hidden[0] if isinstance(hidden, tuple) else hidden)
+
+        if reverse:
+            output.reverse()
+        output = torch.cat(output, 0).view(input.size(0), *output[0].size())
+
+        return hidden, output
+
+    return forward
+
+
+def variable_recurrent_factory(inner, reverse=False):
+    if reverse:
+        return VariableRecurrentReverse(inner)
+    else:
+        return VariableRecurrent(inner)
+
+
+def VariableRecurrent(inner):
+    def forward(input, hidden, weight, batch_sizes, input_quantizer, weight_quantizer):
+
+        output = []
+        input_offset = 0
+        last_batch_size = batch_sizes[0]
+        hiddens = []
+        flat_hidden = not isinstance(hidden, tuple)
+        if flat_hidden:
+            hidden = (hidden,)
+        for batch_size in batch_sizes:
+            step_input = input[input_offset:input_offset + batch_size]
+            input_offset += batch_size
+
+            dec = last_batch_size - batch_size
+            if dec > 0:
+                hiddens.append(tuple(h[-dec:] for h in hidden))
+                hidden = tuple(h[:-dec] for h in hidden)
+            last_batch_size = batch_size
+
+            if flat_hidden:
+                hidden = (inner(step_input, hidden[0], *weight,
+                                input_quantizer=input_quantizer, weight_quantizer=weight_quantizer),)
+            else:
+                hidden = inner(step_input, hidden, *weight,
+                               input_quantizer=input_quantizer, weight_quantizer=weight_quantizer)
+
+            output.append(hidden[0])
+        hiddens.append(hidden)
+        hiddens.reverse()
+
+        hidden = tuple(torch.cat(h, 0) for h in zip(*hiddens))
+        assert hidden[0].size(0) == batch_sizes[0]
+        if flat_hidden:
+            hidden = hidden[0]
+        output = torch.cat(output, 0)
+
+        return hidden, output
+
+    return forward
+
+
+def VariableRecurrentReverse(inner):
+    def forward(input, hidden, weight, batch_sizes, input_quantizer, weight_quantizer):
+        output = []
+        input_offset = input.size(0)
+        last_batch_size = batch_sizes[-1]
+        initial_hidden = hidden
+        flat_hidden = not isinstance(hidden, tuple)
+        if flat_hidden:
+            hidden = (hidden,)
+            initial_hidden = (initial_hidden,)
+        hidden = tuple(h[:batch_sizes[-1]] for h in hidden)
+        for i in reversed(range(len(batch_sizes))):
+            batch_size = batch_sizes[i]
+            inc = batch_size - last_batch_size
+            if inc > 0:
+                hidden = tuple(torch.cat((h, ih[last_batch_size:batch_size]), 0)
+                               for h, ih in zip(hidden, initial_hidden))
+            last_batch_size = batch_size
+            step_input = input[input_offset - batch_size:input_offset]
+            input_offset -= batch_size
+
+            if flat_hidden:
+                hidden = (inner(step_input, hidden[0], *weight,
+                                input_quantizer=input_quantizer, weight_quantizer=weight_quantizer),)
+            else:
+                hidden = inner(step_input, hidden, *weight,
+                               input_quantizer=input_quantizer, weight_quantizer=weight_quantizer)
+            output.append(hidden[0])
+
+        output.reverse()
+        output = torch.cat(output, 0)
+        if flat_hidden:
+            hidden = hidden[0]
+        return hidden, output
+
+    return forward
+
+
+def AutogradRNN(mode, input_size, hidden_size, num_layers=1, batch_first=False,
+                dropout=0, train=True, bidirectional=False, variable_length=False,
+                dropout_state=None, flat_weight=None,
+                input_quantizers=None, weight_quantizers=None):
+
+    if mode == 'RNN_RELU':
+        cell = RNNReLUCell
+    elif mode == 'RNN_TANH':
+        cell = RNNTanhCell
+    elif mode == 'LSTM':
+        cell = LSTMCell
+    elif mode == 'GRU':
+        cell = GRUCell
+    else:
+        raise Exception('Unknown mode: {}'.format(mode))
+
+    rec_factory = variable_recurrent_factory if variable_length else Recurrent
+
+    if bidirectional:
+        layer = (rec_factory(cell), rec_factory(cell, reverse=True))
+    else:
+        layer = (rec_factory(cell),)
+
+    func = StackedRNN(layer,
+                      num_layers,
+                      (mode == 'LSTM'),
+                      dropout=dropout,
+                      train=train)
+
+    def forward(input, weight, hidden, batch_sizes, input_quantizers, weight_quantizers):
+        if batch_first and not variable_length:
+            input = input.transpose(0, 1)
+
+        nexth, output = func(input, hidden, weight, batch_sizes, input_quantizers, weight_quantizers)
+
+        if batch_first and not variable_length:
+            output = output.transpose(0, 1)
+
+        return output, nexth
+
+    return forward
+
+
+def RNN(*args, **kwargs):
+
+    def forward(input, *fargs, **fkwargs):
+        func = AutogradRNN(*args, **kwargs)
+        return func(input, *fargs, **fkwargs)
+
+    return forward

tools/pytorch-quantization/pytorch_quantization/nn/functional.py

@@ -0,0 +1,61 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+"""Some supportive functions"""
+from absl import logging
+
+import torch
+from torch.autograd import Function
+
+
+class ClipFunction(Function):
+    """An universal tensor clip function
+
+    Pytorch's clamp() only supports scalar range and doesn't support broadcast. This implementation uses min/max which
+    is more genaral. The gradient is defined according to IBM's PACT paper https://arxiv.org/abs/1805.06085, which is
+    also the behavior of Tensorflow's clip_by_value()
+    """
+
+    @staticmethod
+    def forward(ctx, input, clip_value_min, clip_value_max):
+        output = torch.min(input, clip_value_max)
+        output = torch.max(output, clip_value_min)
+        ctx.save_for_backward(input, clip_value_min, clip_value_max)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, clip_value_min, clip_value_max = ctx.saved_tensors
+        min_mask = (input > clip_value_min).to(grad_output.dtype)
+        max_mask = (input < clip_value_max).to(grad_output.dtype)
+        grad_input = grad_output * min_mask * max_mask
+
+        if clip_value_min.requires_grad or clip_value_max.requires_grad:
+            logging.log_first_n(logging.WARNING, "Learning clip min/max is experimental, use at your own risk :).", 1)
+        if clip_value_min.numel() != 1 or clip_value_max.numel() != 1:
+            raise ValueError("Learnable min/max can only be scalar, got size %s and %s." % (clip_value_min.size(),
+                                                                                            clip_value_max.size()))
+
+        # Ensure the dtypes of min/max grads matches the input dtype
+        # This might be necessary if running w/ AMP which will cast to fp32 before `sum()`
+        grad_clip_value_min = (grad_output * (1. - min_mask)).sum().to(clip_value_min.dtype) if clip_value_min.requires_grad else None
+        grad_clip_value_max = (grad_output * (1. - max_mask)).sum().to(clip_value_min.dtype) if clip_value_max.requires_grad else None
+
+        return grad_input, grad_clip_value_min, grad_clip_value_max
+
+
+clip = ClipFunction.apply

tools/pytorch-quantization/pytorch_quantization/nn/modules/_utils.py

@@ -0,0 +1,165 @@
+#
+# 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.
+#
+
+
+"""Some helper functions for implementing quantized modules"""
+import copy
+import inspect
+
+from absl import logging
+
+from torch import nn
+
+from pytorch_quantization.nn import TensorQuantizer
+from pytorch_quantization.tensor_quant import QuantDescriptor, QUANT_DESC_8BIT_PER_TENSOR
+
+
+class QuantMixin():
+    """Mixin class for adding basic quantization logic to quantized modules"""
+
+    default_quant_desc_input = QUANT_DESC_8BIT_PER_TENSOR
+    default_quant_desc_weight = QUANT_DESC_8BIT_PER_TENSOR
+
+    @classmethod
+    def set_default_quant_desc_input(cls, value):
+        """
+        Args:
+            value: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`
+        """
+        if not isinstance(value, QuantDescriptor):
+            raise ValueError("{} is not an instance of QuantDescriptor!")
+        cls.default_quant_desc_input = copy.deepcopy(value)
+
+    @classmethod
+    def set_default_quant_desc_weight(cls, value):
+        """
+        Args:
+            value: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`
+        """
+        if not isinstance(value, QuantDescriptor):
+            raise ValueError("{} is not an instance of QuantDescriptor!")
+        cls.default_quant_desc_weight = copy.deepcopy(value)
+
+    def init_quantizer(self, quant_desc_input, quant_desc_weight, num_layers=None):
+        """Helper function for __init__ of quantized module
+
+        Create input and weight quantizer based on quant_desc passed by kwargs, or default of the class.
+
+        Args:
+            quant_desc_input: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`
+            quant_desc_weight: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`
+            num_layers: An integer. Default None. If not None, create a list of quantizers.
+        """
+        if not inspect.stack()[1].function == "__init__":
+            raise TypeError("{} should be only called by __init__ of quantized module.".format(__name__))
+        self._fake_quant = True
+        if (not quant_desc_input.fake_quant) or (not quant_desc_weight.fake_quant):
+            raise ValueError("Only fake quantization is supported!")
+
+        logging.info("Input is %squantized to %d bits in %s with axis %s!", ""
+                     if not quant_desc_input.fake_quant else "fake ",
+                     quant_desc_input.num_bits, self.__class__.__name__, quant_desc_input.axis)
+        logging.info("Weight is %squantized to %d bits in %s with axis %s!", ""
+                     if not quant_desc_weight.fake_quant else "fake ",
+                     quant_desc_weight.num_bits, self.__class__.__name__, quant_desc_weight.axis)
+
+        if num_layers is None:
+            self._input_quantizer = TensorQuantizer(quant_desc_input)
+            self._weight_quantizer = TensorQuantizer(quant_desc_weight)
+        else:
+            self._input_quantizers = nn.ModuleList([TensorQuantizer(quant_desc_input) for _ in range(num_layers)])
+            self._weight_quantizers = nn.ModuleList([TensorQuantizer(quant_desc_weight) for _ in range(num_layers)])
+
+    # pylint:disable=missing-docstring
+    @property
+    def input_quantizer(self):
+        return self._input_quantizer
+
+    @property
+    def weight_quantizer(self):
+        return self._weight_quantizer
+    # pylint:enable=missing-docstring
+
+
+class QuantInputMixin():
+    """Mixin class for adding basic quantization logic to quantized modules"""
+
+    default_quant_desc_input = QUANT_DESC_8BIT_PER_TENSOR
+
+    @classmethod
+    def set_default_quant_desc_input(cls, value):
+        """
+        Args:
+            value: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`
+        """
+        if not isinstance(value, QuantDescriptor):
+            raise ValueError("{} is not an instance of QuantDescriptor!")
+        cls.default_quant_desc_input = copy.deepcopy(value)
+
+    def init_quantizer(self, quant_desc_input):
+        """Helper function for __init__ of simple quantized module
+
+        Create input quantizer based on quant_desc passed by kwargs, or default of the class.
+
+        Args:
+            quant_desc_input: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`
+        """
+        if not inspect.stack()[1].function == "__init__":
+            raise TypeError("{} should be only called by __init__ of quantized module.".format(__name__))
+        self._fake_quant = True
+        if not quant_desc_input.fake_quant:
+            raise ValueError("Only fake quantization is supported!")
+
+        logging.info("Input is %squantized to %d bits in %s with axis %s!", ""
+                     if not quant_desc_input.fake_quant else "fake ",
+                     quant_desc_input.num_bits, self.__class__.__name__, quant_desc_input.axis)
+
+        self._input_quantizer = TensorQuantizer(quant_desc_input)
+
+    # pylint:disable=missing-docstring
+    @property
+    def input_quantizer(self):
+        return self._input_quantizer
+    # pylint:enable=missing-docstring
+
+
+def pop_quant_desc_in_kwargs(quant_cls, input_only=False, **kwargs):
+    """Pop quant descriptors in kwargs
+
+    If there is no descriptor in kwargs, the default one in quant_cls will be used
+
+    Arguments:
+       quant_cls: A class that has default quantization descriptors
+       input_only: A boolean. If True, pop quant_desc_input only, not quant_desc_weight. Default false.
+
+    Keyword Arguments:
+       quant_desc_input: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+           Quantization descriptor of input.
+       quant_desc_weight: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+           Quantization descriptor of weight.
+    """
+    quant_desc_input = kwargs.pop('quant_desc_input', quant_cls.default_quant_desc_input)
+    if not input_only:
+        quant_desc_weight = kwargs.pop('quant_desc_weight', quant_cls.default_quant_desc_weight)
+
+    # Check if anything is left in **kwargs
+    if kwargs:
+        raise TypeError("Unused keys: {}".format(kwargs.keys()))
+
+    if input_only:
+        return quant_desc_input
+    return quant_desc_input, quant_desc_weight

tools/pytorch-quantization/pytorch_quantization/nn/modules/clip.py

@@ -0,0 +1,59 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+"""Implement a clip module as pytorch only has a simple clamp function """
+import torch
+from torch import nn
+from torch.nn.parameter import Parameter
+
+from pytorch_quantization.nn import functional as QF
+
+__all__ = ['Clip']
+
+class Clip(nn.Module):
+    """Clip tensor
+
+    Args:
+        clip_value_min: A number or tensor of lower bound to clip
+        clip_value_max: A number of tensor of upper bound to clip
+        learn_min: A boolean. If True, learn min. clip_value_min will be used to initialize. Default False
+        learn_max: A boolean. Similar as learn_min but for max.
+
+    Raises:
+        ValueError:
+    """
+
+    def __init__(self, clip_value_min, clip_value_max, learn_min=False, learn_max=False):
+        super(Clip, self).__init__()
+        if learn_min:
+            if not isinstance(clip_value_min, float) and clip_value_min.size != 1:
+                raise ValueError("clip_value_min/clip_value_max must be scalar for initilizing learnable range.")
+            self.clip_value_min = Parameter(torch.tensor(clip_value_min))  # pylint: disable=not-callable
+        else:
+            self.clip_value_min = clip_value_min
+
+        if learn_max:
+            if not isinstance(clip_value_max, float) and clip_value_max.size != 1:
+                raise ValueError("clip_value_min/clip_value_max must be scalar for initilizing learnable range.")
+            self.clip_value_max = Parameter(torch.tensor(clip_value_max))  # pylint: disable=not-callable
+        else:
+            self.clip_value_max = clip_value_max
+
+    def forward(self, inputs):
+        outputs = QF.clip(inputs, self.clip_value_min, self.clip_value_max)
+        return outputs

tools/pytorch-quantization/pytorch_quantization/nn/modules/quant_conv.py

@@ -0,0 +1,419 @@
+#
+# 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.
+#
+
+
+"""Quantized convolution
+Base code is from nn.Conv, details of Module and original argument can be found there.
+Module names are intentionally kept same as unquantized version so that they can be dropped into preexisting model
+easily, and load pretrained weight. Aliases with Quant prefix are defined and are encouraged to be used explicitly
+when start scratch.
+"""
+
+import inspect
+import torch
+import torch.nn
+import torch.nn.functional as F
+from torch.nn.modules.utils import _single, _pair, _triple
+from torch.nn.modules.conv import _ConvTransposeNd
+
+from pytorch_quantization import tensor_quant
+
+from . import _utils
+
+__all__ = [
+    "Conv2d", "QuantConv2d", "Conv3d", "QuantConv3d", "Conv1d", "QuantConv1d", "ConvTranspose1d", "ConvTranspose2d",
+    "ConvTranspose3d", "QuantConvTranspose1d", "QuantConvTranspose2d", "QuantConvTranspose3d"
+]
+
+
+class _QuantConvNd(torch.nn.modules.conv._ConvNd, _utils.QuantMixin):
+    """base class of quantized Conv inherited from _ConvNd
+
+    Comments of original arguments can be found in torch.nn.modules.conv
+
+    Arguments:
+        quant_desc_input: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+            Quantization descriptor of input.
+        quant_desc_weight: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+            Quantization descriptor of weight.
+
+    Raises:
+        ValueError: If unsupported arguments are passed in.
+
+    Readonly properties:
+        - input_quantizer:
+        - weight_quantizer:
+
+    Static methods:
+        - set_default_quant_desc_input: Set default_quant_desc_input
+        - set_default_quant_desc_weight: Set default_quant_desc_weight
+    """
+
+    default_quant_desc_input = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+
+    def __init__(self, in_channels, out_channels, kernel_size, stride, padding, dilation, transposed, output_padding,
+                 groups, bias, padding_mode, quant_desc_input, quant_desc_weight):
+        super(_QuantConvNd, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation,
+                                           transposed, output_padding, groups, bias, padding_mode)
+        self.init_quantizer(quant_desc_input, quant_desc_weight)
+
+    def _quant(self, input):
+        """Apply quantization on input and weight
+
+        Function called by the classes lower in the hierarchy, which actually performs the quantization before forward
+        in the derivate class the particular Function.
+
+        Arguments:
+            input: in_features to quantize
+        Returns:
+            A tuple: (quant_in_feature, quant_weight)
+        """
+        quant_input = self._input_quantizer(input)
+        quant_weight = self._weight_quantizer(self.weight)
+
+        return (quant_input, quant_weight)
+
+
+class QuantConv2d(_QuantConvNd):
+    """Quantized 2D conv"""
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 bias=True,
+                 padding_mode='zeros',
+                 **kwargs):
+
+        kernel_size = _pair(kernel_size)
+        stride = _pair(stride)
+        padding = _pair(padding)
+        dilation = _pair(dilation)
+
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        super(QuantConv2d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, False,
+                                          _pair(0), groups, bias, padding_mode,
+                                          quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+
+    def forward(self, input):
+        # the actual quantization happens in the next level of the class hierarchy
+        quant_input, quant_weight = self._quant(input)
+
+        if self.padding_mode == 'circular':
+            expanded_padding = ((self.padding[1] + 1) // 2, self.padding[1] // 2,
+                                (self.padding[0] + 1) // 2, self.padding[0] // 2)
+            output = F.conv2d(F.pad(quant_input, expanded_padding, mode='circular'),
+                              quant_weight, self.bias, self.stride,
+                              _pair(0), self.dilation, self.groups)
+        else:
+            output = F.conv2d(quant_input, quant_weight, self.bias, self.stride, self.padding, self.dilation,
+                              self.groups)
+
+        return output
+
+
+class QuantConv3d(_QuantConvNd):
+    """Quantized 3D Conv"""
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 bias=True,
+                 padding_mode='zeros',
+                 **kwargs):
+
+        kernel_size = _triple(kernel_size)
+        stride = _triple(stride)
+        padding = _triple(padding)
+        dilation = _triple(dilation)
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        super(QuantConv3d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, False,
+                                          _triple(0), groups, bias, padding_mode,
+                                          quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+
+    def forward(self, input):
+        # the actual quantization happens in the next level of the class hierarchy
+        quant_input, quant_weight = self._quant(input)
+
+        if self.padding_mode == 'circular':
+            expanded_padding = ((self.padding[2] + 1) // 2, self.padding[2] // 2,
+                                (self.padding[1] + 1) // 2, self.padding[1] // 2,
+                                (self.padding[0] + 1) // 2, self.padding[0] // 2)
+            output = F.conv3d(F.pad(quant_input, expanded_padding, mode='circular'),
+                              quant_weight, self.bias, self.stride, _triple(0),
+                              self.dilation, self.groups)
+        else:
+            output = F.conv3d(quant_input, quant_weight, self.bias, self.stride, self.padding, self.dilation,
+                              self.groups)
+
+        return output
+
+
+class QuantConv1d(_QuantConvNd):
+    """Quantized 1D Conv"""
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONV1D_WEIGHT_PER_CHANNEL
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 dilation=1,
+                 groups=1,
+                 bias=True,
+                 padding_mode='zeros',
+                 **kwargs):
+
+        kernel_size = _single(kernel_size)
+        stride = _single(stride)
+        padding = _single(padding)
+        dilation = _single(dilation)
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        super(QuantConv1d, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, False,
+                                          _single(0), groups, bias, padding_mode,
+                                          quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+
+    def forward(self, input):
+        # the actual quantization happens in the next level of the class hierarchy
+        quant_input, quant_weight = self._quant(input)
+
+        if self.padding_mode == 'circular':
+            expanded_padding = ((self.padding[0] + 1) // 2, self.padding[0] // 2)
+            output = F.conv1d(F.pad(quant_input, expanded_padding, mode='circular'),
+                              quant_weight, self.bias, self.stride,
+                              _single(0), self.dilation, self.groups)
+        else:
+            output = F.conv1d(quant_input, quant_weight, self.bias, self.stride,
+                              self.padding, self.dilation, self.groups)
+
+        return output
+
+
+class _QuantConvTransposeNd(torch.nn.modules.conv._ConvTransposeNd, _utils.QuantMixin):
+    """base class of quantized Transposed Conv inherited from _ConvTransposeNd
+
+    Comments of original arguments can be found in torch.nn.modules.conv
+
+    Arguments:
+        quant_desc_input: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+            Quantization descriptor of input.
+        quant_desc_weight: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+            Quantization descriptor of weight.
+
+    Raises:
+        ValueError: If unsupported arguments are passed in.
+
+    Readonly properties:
+        - input_quantizer:
+        - weight_quantizer:
+
+    Static methods:
+        - set_default_quant_desc_input: Set default_quant_desc_input
+        - set_default_quant_desc_weight: Set default_quant_desc_weight
+    """
+
+    default_quant_desc_input = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+    def __init__(self, in_channels, out_channels, kernel_size, stride,
+                 padding, dilation, transposed, output_padding,
+                 groups, bias, padding_mode, quant_desc_input, quant_desc_weight):
+        super(_QuantConvTransposeNd, self).__init__(in_channels, out_channels, kernel_size, stride, padding, dilation,
+                                                    transposed, output_padding, groups, bias, padding_mode)
+        self.init_quantizer(quant_desc_input, quant_desc_weight)
+
+    def _quant(self, input):
+        """Apply quantization on input and weight
+
+        Function called by the classes lower in the hierarchy, which actually performs the quantization before forward
+        in the derivate class the particular Function.
+
+        Arguments:
+            input: in_features to quantize
+        Returns:
+            A tuple: (quant_in_feature, quant_weight)
+        """
+        quant_input = self._input_quantizer(input)
+        quant_weight = self._weight_quantizer(self.weight)
+
+        return (quant_input, quant_weight)
+
+    def _output_padding_nd(self,
+                           input,
+                           output_size,
+                           stride,
+                           padding,
+                           kernel_size,
+                           num_spatial_dims,
+                           dilation=None):
+        if "num_spatial_dims" in inspect.signature(self._output_padding).parameters:
+            return self._output_padding(input, output_size, stride, padding, kernel_size, num_spatial_dims)
+        else:
+            return self._output_padding(input, output_size, stride, padding, kernel_size)
+
+
+class QuantConvTranspose1d(_QuantConvTransposeNd):
+    """Quantized ConvTranspose1d"""
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE1D_WEIGHT_PER_CHANNEL
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 output_padding=0,
+                 groups=1,
+                 bias=True,
+                 dilation=1,
+                 padding_mode='zeros',
+                 **kwargs):
+        kernel_size = _single(kernel_size)
+        stride = _single(stride)
+        padding = _single(padding)
+        dilation = _single(dilation)
+        output_padding = _single(output_padding)
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        super(QuantConvTranspose1d, self).__init__(
+            in_channels, out_channels, kernel_size, stride, padding, dilation,
+            True, output_padding, groups, bias, padding_mode,
+            quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+
+    def forward(self, input, output_size=None):
+        if self.padding_mode != 'zeros':
+            raise ValueError('Only `zeros` padding mode is supported for QuantConvTranspose1d')
+
+        num_spatial_dims = 1
+        output_padding = self._output_padding_nd(input, output_size, self.stride, self.padding, self.kernel_size,
+                                                 num_spatial_dims)
+
+        quant_input, quant_weight = self._quant(input)
+        output = F.conv_transpose1d(quant_input, quant_weight, self.bias, self.stride, self.padding, output_padding,
+                                    self.groups, self.dilation)
+        return output
+
+
+class QuantConvTranspose2d(_QuantConvTransposeNd):
+    """Quantized ConvTranspose2d"""
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE2D_WEIGHT_PER_CHANNEL
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 output_padding=0,
+                 groups=1,
+                 bias=True,
+                 dilation=1,
+                 padding_mode='zeros',
+                 **kwargs):
+        kernel_size = _pair(kernel_size)
+        stride = _pair(stride)
+        padding = _pair(padding)
+        dilation = _pair(dilation)
+        output_padding = _pair(output_padding)
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        super(QuantConvTranspose2d, self).__init__(
+            in_channels, out_channels, kernel_size, stride, padding, dilation,
+            True, output_padding, groups, bias, padding_mode,
+            quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+
+    def forward(self, input, output_size=None):
+        if self.padding_mode != 'zeros':
+            raise ValueError('Only `zeros` padding mode is supported for QuantConvTranspose2d')
+
+        num_spatial_dims = 2
+        output_padding = self._output_padding_nd(input, output_size, self.stride, self.padding, self.kernel_size,
+                                                 num_spatial_dims)
+
+        quant_input, quant_weight = self._quant(input)
+        output = F.conv_transpose2d(quant_input, quant_weight, self.bias, self.stride, self.padding, output_padding,
+                                    self.groups, self.dilation)
+
+        return output
+
+
+class QuantConvTranspose3d(_QuantConvTransposeNd):
+    """Quantized ConvTranspose3d"""
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE3D_WEIGHT_PER_CHANNEL
+
+    def __init__(self,
+                 in_channels,
+                 out_channels,
+                 kernel_size,
+                 stride=1,
+                 padding=0,
+                 output_padding=0,
+                 groups=1,
+                 bias=True,
+                 dilation=1,
+                 padding_mode='zeros',
+                 **kwargs):
+        kernel_size = _triple(kernel_size)
+        stride = _triple(stride)
+        padding = _triple(padding)
+        dilation = _triple(dilation)
+        output_padding = _triple(output_padding)
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        super(QuantConvTranspose3d, self).__init__(
+            in_channels, out_channels, kernel_size, stride, padding, dilation,
+            True, output_padding, groups, bias, padding_mode,
+            quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+
+    def forward(self, input, output_size=None):
+        if self.padding_mode != 'zeros':
+            raise ValueError('Only `zeros` padding mode is supported for QuantConvTranspose3d')
+
+        num_spatial_dims = 3
+        output_padding = self._output_padding_nd(input, output_size, self.stride, self.padding, self.kernel_size,
+                                                 num_spatial_dims)
+
+        quant_input, quant_weight = self._quant(input)
+        output = F.conv_transpose3d(quant_input, quant_weight, self.bias, self.stride, self.padding, output_padding,
+                                    self.groups, self.dilation)
+
+        return output
+
+
+# Define alias with Quant prefix
+_ConvNd = _QuantConvNd
+Conv1d = QuantConv1d
+Conv2d = QuantConv2d
+Conv3d = QuantConv3d
+ConvTranspose1d = QuantConvTranspose1d
+ConvTranspose2d = QuantConvTranspose2d
+ConvTranspose3d = QuantConvTranspose3d

tools/pytorch-quantization/pytorch_quantization/nn/modules/quant_instancenorm.py

@@ -0,0 +1,79 @@
+#
+# 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.
+#
+
+"""Quantized instance normalization module
+   Base code is from nn.InstanceNorm, details of the module can be found from the offical repo.
+"""
+
+from torch.nn.modules.batchnorm import _NormBase
+import torch.nn.functional as F
+from torch.nn.modules import instancenorm
+
+from pytorch_quantization.nn import TensorQuantizer
+from pytorch_quantization import tensor_quant
+from . import _utils
+
+__all__ = [
+    "QuantInstanceNorm1d", "QuantInstanceNorm2d", "QuantInstanceNorm3d"
+]
+
+class QuantInstanceNorm1d(instancenorm.InstanceNorm1d, _utils.QuantInputMixin):
+    r"""Applies Quantized Instance Normalization over a 3D input
+    """
+    def __init__(
+            self, num_features: int, eps: float = 1e-5, momentum: float = 0.1, affine: bool = False,
+            track_running_stats: bool = False, **kwargs):
+        super(QuantInstanceNorm1d, self).__init__(
+            num_features, eps, momentum, affine, track_running_stats)
+        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(QuantInstanceNorm1d, self).forward(quant_input)
+
+
+class QuantInstanceNorm2d(instancenorm.InstanceNorm2d, _utils.QuantInputMixin):
+    r"""Applies Quantized Instance Normalization over a 4D input
+    """
+    def __init__(
+            self, num_features: int, eps: float = 1e-5, momentum: float = 0.1, affine: bool = False,
+            track_running_stats: bool = False, **kwargs):
+        super(QuantInstanceNorm2d, self).__init__(
+            num_features, eps, momentum, affine, track_running_stats)
+        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(QuantInstanceNorm2d, self).forward(quant_input)
+
+
+class QuantInstanceNorm3d(instancenorm.InstanceNorm3d, _utils.QuantInputMixin):
+    r"""Applies Quantized Instance Normalization over a 5D input
+    """
+    def __init__(
+            self, num_features: int, eps: float = 1e-5, momentum: float = 0.1, affine: bool = False,
+            track_running_stats: bool = False, **kwargs):
+        super(QuantInstanceNorm3d, self).__init__(
+            num_features, eps, momentum, affine, track_running_stats)
+        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(QuantInstanceNorm3d, self).forward(quant_input)

tools/pytorch-quantization/pytorch_quantization/nn/modules/quant_linear.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.
+#
+
+
+"""Quantized Linear"""
+from torch import nn
+from torch.nn import functional as F
+
+from pytorch_quantization import tensor_quant
+
+from . import _utils
+
+__all__ = ["Linear", "QuantLinear"]
+
+class QuantLinear(nn.Linear, _utils.QuantMixin):
+    """Quantized version of nn.Linear
+
+    Apply quantized linear to the incoming data, y = dequant(quant(x)quant(A)^T + b).
+
+    Keep Module name "Linear" instead of "QuantLinear" so that it can be easily dropped into preexisting model and load
+    pretrained weights. An alias "QuantLinear" is defined below. The base code is a copy of nn.Linear, see detailed
+    comment of original arguments there.
+
+    Quantization descriptors are passed in in kwargs. If not presents, default_quant_desc_input and
+    default_quant_desc_weight are used.
+
+    Keyword Arguments:
+        quant_desc_input: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+            Quantization descriptor of input.
+        quant_desc_wegiht: An instance of :class:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+            Quantization descriptor of weight.
+
+    Raises:
+        ValueError: If unsupported arguments are passed in.
+        KeyError: If unsupported kwargs are passed in.
+
+    Readonly properties:
+        - input_quantizer:
+        - weight_quantizer:
+
+    Static methods:
+        - set_default_quant_desc_input: Set default_quant_desc_input
+        - set_default_quant_desc_weight: Set default_quant_desc_weight
+    """
+
+    default_quant_desc_input = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW
+
+    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
+
+
+Linear = QuantLinear

tools/pytorch-quantization/pytorch_quantization/nn/modules/quant_pooling.py

@@ -0,0 +1,163 @@
+#
+# 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.
+#
+
+
+"""Quantized Pooling
+Base code is from nn.pooling, details of Module and original argument can be found there.
+Module names are intentionally kept same as unquantized version so that they can be dropped into preexisting model
+easily, and load pretrained weight. Aliases with Quant prefix are defined and are encouraged to be used explicitly
+when start scratch.
+"""
+
+from torch.nn.modules import pooling
+
+from . import _utils
+
+__all__ = [
+    "MaxPool1d", "QuantMaxPool1d", "MaxPool2d", "QuantMaxPool2d", "MaxPool3d", "QuantMaxPool3d",
+    "AvgPool1d", "QuantAvgPool1d", "AvgPool2d", "QuantAvgPool2d", "AvgPool3d", "QuantAvgPool3d",
+    "AdaptiveAvgPool1d", "QuantAdaptiveAvgPool1d", "AdaptiveAvgPool2d", "QuantAdaptiveAvgPool2d",
+    "AdaptiveAvgPool3d", "QuantAdaptiveAvgPool3d"
+]
+
+class QuantMaxPool1d(pooling.MaxPool1d, _utils.QuantInputMixin):
+    """Quantized 1D maxpool"""
+    def __init__(self, kernel_size, stride=None, padding=0, dilation=1,
+                 return_indices=False, ceil_mode=False, **kwargs):
+        super(QuantMaxPool1d, 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(QuantMaxPool1d, self).forward(quant_input)
+
+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)
+
+class QuantMaxPool3d(pooling.MaxPool3d, _utils.QuantInputMixin):
+    """Quantized 3D maxpool"""
+    def __init__(self, kernel_size, stride=None, padding=0, dilation=1,
+                 return_indices=False, ceil_mode=False, **kwargs):
+        super(QuantMaxPool3d, 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(QuantMaxPool3d, self).forward(quant_input)
+
+
+class QuantAvgPool1d(pooling.AvgPool1d, _utils.QuantInputMixin):
+    """Quantized 1D average pool"""
+    def __init__(self, kernel_size, stride=None, padding=0, ceil_mode=False,
+                 count_include_pad=True, **kwargs):
+        super(QuantAvgPool1d, self).__init__(kernel_size, stride, padding, ceil_mode,
+                                             count_include_pad)
+        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(QuantAvgPool1d, self).forward(quant_input)
+
+class QuantAvgPool2d(pooling.AvgPool2d, _utils.QuantInputMixin):
+    """Quantized 2D average pool"""
+    def __init__(self, kernel_size, stride=None, padding=0, ceil_mode=False,
+                 count_include_pad=True, divisor_override=None, **kwargs):
+        super(QuantAvgPool2d, self).__init__(kernel_size, stride, padding, ceil_mode,
+                                             count_include_pad, divisor_override)
+        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(QuantAvgPool2d, self).forward(quant_input)
+
+
+class QuantAvgPool3d(pooling.AvgPool3d, _utils.QuantInputMixin):
+    """Quantized 3D average pool"""
+    def __init__(self, kernel_size, stride=None, padding=0, ceil_mode=False,
+                 count_include_pad=True, divisor_override=None, **kwargs):
+        super(QuantAvgPool3d, self).__init__(kernel_size, stride, padding, ceil_mode,
+                                             count_include_pad, divisor_override)
+        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(QuantAvgPool3d, self).forward(quant_input)
+
+
+class QuantAdaptiveAvgPool1d(pooling.AdaptiveAvgPool1d, _utils.QuantInputMixin):
+    """Quantized 1D adaptive average pool"""
+    def __init__(self, output_size, **kwargs):
+        super(QuantAdaptiveAvgPool1d, self).__init__(output_size)
+        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(QuantAdaptiveAvgPool1d, self).forward(quant_input)
+
+
+class QuantAdaptiveAvgPool2d(pooling.AdaptiveAvgPool2d, _utils.QuantInputMixin):
+    """Quantized 2D adaptive average pool"""
+    def __init__(self, output_size, **kwargs):
+        super(QuantAdaptiveAvgPool2d, self).__init__(output_size)
+        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(QuantAdaptiveAvgPool2d, self).forward(quant_input)
+
+
+class QuantAdaptiveAvgPool3d(pooling.AdaptiveAvgPool3d, _utils.QuantInputMixin):
+    """Quantized 3D adaptive average pool"""
+    def __init__(self, output_size, **kwargs):
+        super(QuantAdaptiveAvgPool3d, self).__init__(output_size)
+        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(QuantAdaptiveAvgPool3d, self).forward(quant_input)
+
+
+# Define alias with Quant prefix
+MaxPool1d = QuantMaxPool1d
+MaxPool2d = QuantMaxPool2d
+MaxPool3d = QuantMaxPool3d
+AvgPool1d = QuantAvgPool1d
+AvgPool2d = QuantAvgPool2d
+AvgPool3d = QuantAvgPool3d
+AdaptiveAvgPool1d = QuantAdaptiveAvgPool1d
+AdaptiveAvgPool2d = QuantAdaptiveAvgPool2d
+AdaptiveAvgPool3d = QuantAdaptiveAvgPool3d

tools/pytorch-quantization/pytorch_quantization/nn/modules/quant_rnn.py

@@ -0,0 +1,467 @@
+#
+# 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.
+#
+
+
+"""RNN implementation in python
+originally copied from https://github.com/pytorch/pytorch/blob/v0.4.1/torch/nn/modules/rnn.py
+    backend is changed to _functions/rnn.py
+"""
+import math
+import torch
+import warnings
+import itertools
+import numbers
+
+from torch import nn
+from torch.nn import Parameter
+from torch.nn.utils.rnn import PackedSequence
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization.nn._functions import quant_rnn
+
+from . import _utils
+
+__all__ = ["QuantLSTM", "QuantLSTMCell", "LSTM", "LSTMCell"]
+
+class QuantRNNBase(nn.Module, _utils.QuantMixin):
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW
+
+    def __init__(self, mode, input_size, hidden_size,
+                 num_layers=1, bias=True, batch_first=False,
+                 dropout=0, bidirectional=False, proj_size=0, **kwargs):
+        super(QuantRNNBase, self).__init__()
+        self.mode = mode
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.num_layers = num_layers
+        self.bias = bias
+        self.batch_first = batch_first
+        self.dropout = dropout
+        self.dropout_state = {}
+        self.bidirectional = bidirectional
+        self.proj_size = proj_size
+        num_directions = 2 if bidirectional else 1
+
+        if not isinstance(dropout, numbers.Number) or not 0 <= dropout <= 1 or \
+                isinstance(dropout, bool):
+            raise ValueError("dropout should be a number in range [0, 1] "
+                             "representing the probability of an element being "
+                             "zeroed")
+        if dropout > 0 and num_layers == 1:
+            warnings.warn("dropout option adds dropout after all but last "
+                          "recurrent layer, so non-zero dropout expects "
+                          "num_layers greater than 1, but got dropout={} and "
+                          "num_layers={}".format(dropout, num_layers))
+
+        if proj_size < 0:
+            raise ValueError("proj_size should be a positive integer or zero to disable projections")
+        if proj_size > 0:
+            raise ValueError("proj_size is not supported in pytorch-quantization yet")
+
+        if mode == 'LSTM':
+            gate_size = 4 * hidden_size
+        elif mode == 'GRU':
+            gate_size = 3 * hidden_size
+        else:
+            gate_size = hidden_size
+
+        self._all_weights = []
+        for layer in range(num_layers):
+            for direction in range(num_directions):
+                layer_input_size = input_size if layer == 0 else hidden_size * num_directions
+
+                w_ih = Parameter(torch.Tensor(gate_size, layer_input_size))
+                w_hh = Parameter(torch.Tensor(gate_size, hidden_size))
+                b_ih = Parameter(torch.Tensor(gate_size))
+                b_hh = Parameter(torch.Tensor(gate_size))
+                layer_params = (w_ih, w_hh, b_ih, b_hh)
+
+                suffix = '_reverse' if direction == 1 else ''
+                param_names = ['weight_ih_l{}{}', 'weight_hh_l{}{}']
+                if bias:
+                    param_names += ['bias_ih_l{}{}', 'bias_hh_l{}{}']
+                param_names = [x.format(layer, suffix) for x in param_names]
+
+                for name, param in zip(param_names, layer_params):
+                    setattr(self, name, param)
+                self._all_weights.append(param_names)
+
+        self.flatten_parameters()
+        self.reset_parameters()
+
+        quant_desc_input, quant_desc_weight = _utils.pop_quant_desc_in_kwargs(self.__class__, **kwargs)
+        self.init_quantizer(quant_desc_input, quant_desc_weight, num_layers=num_layers * (1 + bidirectional))
+
+    def flatten_parameters(self):
+        """Resets parameter data pointer so that they can use faster code paths.
+
+        Right now, this works only if the module is on the GPU and cuDNN is enabled.
+        Otherwise, it's a no-op.
+        """
+        any_param = next(self.parameters()).data
+        if not any_param.is_cuda or not torch.backends.cudnn.is_acceptable(any_param):
+            self._data_ptrs = []
+            return
+
+        # If any parameters alias, we fall back to the slower, copying code path. This is
+        # a sufficient check, because overlapping parameter buffers that don't completely
+        # alias would break the assumptions of the uniqueness check in
+        # Module.named_parameters().
+        unique_data_ptrs = set(p.data_ptr() for l in self.all_weights for p in l)
+        if len(unique_data_ptrs) != sum(len(l) for l in self.all_weights):
+            self._data_ptrs = []
+            return
+
+        with torch.cuda.device_of(any_param):
+            import torch.backends.cudnn.rnn as rnn
+
+            weight_arr = list(itertools.chain.from_iterable(self.all_weights))
+            weight_stride0 = len(self.all_weights[0])
+
+            # NB: This is a temporary hack while we still don't have Tensor
+            # bindings for ATen functions
+            with torch.no_grad():
+                # NB: this is an INPLACE function on weight_arr, that's why the
+                # no_grad() is necessary.
+                weight_buf = torch._cudnn_rnn_flatten_weight(weight_arr, weight_stride0, self.input_size,
+                                                             rnn.get_cudnn_mode(self.mode), self.hidden_size,
+                                                             self.proj_size, self.num_layers, self.batch_first,
+                                                             bool(self.bidirectional))
+
+            self._param_buf_size = weight_buf.size(0)
+            self._data_ptrs = list(p.data.data_ptr() for p in self.parameters())
+
+    def _apply(self, fn):
+        ret = super(QuantRNNBase, self)._apply(fn)
+        self.flatten_parameters()
+        return ret
+
+    def reset_parameters(self):
+        stdv = 1.0 / math.sqrt(self.hidden_size)
+        for weight in self.parameters():
+            weight.data.uniform_(-stdv, stdv)
+
+    def check_forward_args(self, input, hidden, batch_sizes):
+        is_input_packed = batch_sizes is not None
+        expected_input_dim = 2 if is_input_packed else 3
+        if input.dim() != expected_input_dim:
+            raise RuntimeError(
+                'input must have {} dimensions, got {}'.format(
+                    expected_input_dim, input.dim()))
+        if self.input_size != input.size(-1):
+            raise RuntimeError(
+                'input.size(-1) must be equal to input_size. Expected {}, got {}'.format(
+                    self.input_size, input.size(-1)))
+
+        if is_input_packed:
+            mini_batch = int(batch_sizes[0])
+        else:
+            mini_batch = input.size(0) if self.batch_first else input.size(1)
+
+        num_directions = 2 if self.bidirectional else 1
+        expected_hidden_size = (self.num_layers * num_directions,
+                                mini_batch, self.hidden_size)
+
+        def check_hidden_size(hx, expected_hidden_size, msg='Expected hidden size {}, got {}'):
+            if tuple(hx.size()) != expected_hidden_size:
+                raise RuntimeError(msg.format(expected_hidden_size, tuple(hx.size())))
+
+        if self.mode == 'LSTM':
+            check_hidden_size(hidden[0], expected_hidden_size,
+                              'Expected hidden[0] size {}, got {}')
+            check_hidden_size(hidden[1], expected_hidden_size,
+                              'Expected hidden[1] size {}, got {}')
+        else:
+            check_hidden_size(hidden, expected_hidden_size)
+
+    def forward(self, input, hx=None):
+        is_packed = isinstance(input, PackedSequence)
+        if is_packed:
+            input, batch_sizes, sorted_indices, unsorted_indices = input
+            max_batch_size = batch_sizes[0]
+            max_batch_size = int(max_batch_size)
+        else:
+            batch_sizes = None
+            max_batch_size = input.size(0) if self.batch_first else input.size(1)
+
+        if hx is None:
+            num_directions = 2 if self.bidirectional else 1
+            hx = input.new_zeros(self.num_layers * num_directions,
+                                 max_batch_size, self.hidden_size,
+                                 requires_grad=False)
+            if self.mode == 'LSTM':
+                hx = (hx, hx)
+
+        has_flat_weights = list(p.data.data_ptr() for p in self.parameters()) == self._data_ptrs
+        if has_flat_weights:
+            first_data = next(self.parameters()).data
+            assert first_data.storage().size() == self._param_buf_size
+            flat_weight = first_data.new().set_(first_data.storage(), 0, torch.Size([self._param_buf_size]))
+        else:
+            flat_weight = None
+
+        self.check_forward_args(input, hx, batch_sizes)
+        func = quant_rnn.RNN(
+            self.mode,
+            self.input_size,
+            self.hidden_size,
+            num_layers=self.num_layers,
+            batch_first=self.batch_first,
+            dropout=self.dropout,
+            train=self.training,
+            bidirectional=self.bidirectional,
+            dropout_state=self.dropout_state,
+            variable_length=is_packed,
+            flat_weight=flat_weight
+        )
+        output, hidden = func(input, self.all_weights, hx, batch_sizes, self._input_quantizers, self._weight_quantizers)
+        if is_packed:
+            output = PackedSequence(output, batch_sizes)
+        return output, hidden
+
+    def extra_repr(self):
+        s = '{input_size}, {hidden_size}'
+        if self.num_layers != 1:
+            s += ', num_layers={num_layers}'
+        if self.bias is not True:
+            s += ', bias={bias}'
+        if self.batch_first is not False:
+            s += ', batch_first={batch_first}'
+        if self.dropout != 0:
+            s += ', dropout={dropout}'
+        if self.bidirectional is not False:
+            s += ', bidirectional={bidirectional}'
+        return s.format(**self.__dict__)
+
+    def __setstate__(self, d):
+        super(QuantRNNBase, self).__setstate__(d)
+        self.__dict__.setdefault('_data_ptrs', [])
+        if 'all_weights' in d:
+            self._all_weights = d['all_weights']
+        if isinstance(self._all_weights[0][0], str):
+            return
+        num_layers = self.num_layers
+        num_directions = 2 if self.bidirectional else 1
+        self._all_weights = []
+        for layer in range(num_layers):
+            for direction in range(num_directions):
+                suffix = '_reverse' if direction == 1 else ''
+                weights = ['weight_ih_l{}{}', 'weight_hh_l{}{}', 'bias_ih_l{}{}', 'bias_hh_l{}{}']
+                weights = [x.format(layer, suffix) for x in weights]
+                if self.bias:
+                    self._all_weights += [weights]
+                else:
+                    self._all_weights += [weights[:2]]
+
+    @property
+    def all_weights(self):
+        return [[getattr(self, weight) for weight in weights] for weights in self._all_weights]
+
+
+class QuantRNN(QuantRNNBase):
+    r"""Applies a multi-layer Elman RNN with `tanh` or `ReLU` non-linearity to an
+    input sequence.
+    """
+
+    def __init__(self, *args, **kwargs):
+        if 'proj_size' in kwargs:
+            raise ValueError("proj_size argument is only supported for LSTM, not RNN or GRU")
+        if 'nonlinearity' in kwargs:
+            if kwargs['nonlinearity'] == 'tanh':
+                mode = 'RNN_TANH'
+            elif kwargs['nonlinearity'] == 'relu':
+                mode = 'RNN_RELU'
+            else:
+                raise ValueError("Unknown nonlinearity '{}'".format(
+                    kwargs['nonlinearity']))
+            del kwargs['nonlinearity']
+        else:
+            mode = 'RNN_TANH'
+
+        super(QuantRNN, self).__init__(mode, *args, **kwargs)
+
+
+class QuantLSTM(QuantRNNBase):
+    r"""Applies a multi-layer long short-term memory (LSTM) RNN to an input
+    sequence.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(QuantLSTM, self).__init__('LSTM', *args, **kwargs)
+
+
+class GRU(QuantRNNBase):
+    r"""Applies a multi-layer gated recurrent unit (GRU) RNN to an input sequence.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super(GRU, self).__init__('GRU', *args, **kwargs)
+
+class QuantRNNCellBase(nn.Module, _utils.QuantMixin):
+
+    default_quant_desc_weight = tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW
+
+    def extra_repr(self):
+        s = '{input_size}, {hidden_size}'
+        if 'bias' in self.__dict__ and self.bias is not True:
+            s += ', bias={bias}'
+        if 'nonlinearity' in self.__dict__ and self.nonlinearity != "tanh":
+            s += ', nonlinearity={nonlinearity}'
+        return s.format(**self.__dict__)
+
+    def check_forward_input(self, input):
+        if input.size(1) != self.input_size:
+            raise RuntimeError(
+                "input has inconsistent input_size: got {}, expected {}".format(
+                    input.size(1), self.input_size))
+
+    def check_forward_hidden(self, input, hx, hidden_label=''):
+        if input.size(0) != hx.size(0):
+            raise RuntimeError(
+                "Input batch size {} doesn't match hidden{} batch size {}".format(
+                    input.size(0), hidden_label, hx.size(0)))
+
+        if hx.size(1) != self.hidden_size:
+            raise RuntimeError(
+                "hidden{} has inconsistent hidden_size: got {}, expected {}".format(
+                    hidden_label, hx.size(1), self.hidden_size))
+
+
+class QuantRNNCell(QuantRNNCellBase):
+    r"""An Elman RNN cell with tanh or ReLU non-linearity.
+    """
+
+    def __init__(self, input_size, hidden_size, bias=True, nonlinearity="tanh"):
+        super(QuantRNNCell, self).__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.bias = bias
+        self.nonlinearity = nonlinearity
+        self.weight_ih = Parameter(torch.Tensor(hidden_size, input_size))
+        self.weight_hh = Parameter(torch.Tensor(hidden_size, hidden_size))
+        if bias:
+            self.bias_ih = Parameter(torch.Tensor(hidden_size))
+            self.bias_hh = Parameter(torch.Tensor(hidden_size))
+        else:
+            self.register_parameter('bias_ih', None)
+            self.register_parameter('bias_hh', None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        stdv = 1.0 / math.sqrt(self.hidden_size)
+        for weight in self.parameters():
+            weight.data.uniform_(-stdv, stdv)
+
+    def forward(self, input, hx=None):
+        self.check_forward_input(input)
+        if hx is None:
+            hx = input.new_zeros(input.size(0), self.hidden_size, requires_grad=False)
+        self.check_forward_hidden(input, hx)
+        if self.nonlinearity == "tanh":
+            func = quant_rnn.RNNTanhCell
+        elif self.nonlinearity == "relu":
+            func = quant_rnn.RNNReLUCell
+        else:
+            raise RuntimeError(
+                "Unknown nonlinearity: {}".format(self.nonlinearity))
+
+        return func(
+            input, hx,
+            self.weight_ih, self.weight_hh,
+            self.bias_ih, self.bias_hh,
+        )
+
+
+class QuantLSTMCell(QuantRNNCellBase):
+    r"""A long short-term memory (LSTM) cell.
+    """
+
+    def __init__(self, input_size, hidden_size, bias=True, **kwargs):
+        super(QuantLSTMCell, self).__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.bias = bias
+        self.weight_ih = Parameter(torch.Tensor(4 * hidden_size, input_size))
+        self.weight_hh = Parameter(torch.Tensor(4 * hidden_size, hidden_size))
+        if bias:
+            self.bias_ih = Parameter(torch.Tensor(4 * hidden_size))
+            self.bias_hh = Parameter(torch.Tensor(4 * hidden_size))
+        else:
+            self.register_parameter('bias_ih', None)
+            self.register_parameter('bias_hh', None)
+        self.reset_parameters()
+        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 reset_parameters(self):
+        stdv = 1.0 / math.sqrt(self.hidden_size)
+        for weight in self.parameters():
+            weight.data.uniform_(-stdv, stdv)
+
+    def forward(self, input, hx=None):
+        self.check_forward_input(input)
+        if hx is None:
+            hx = input.new_zeros(input.size(0), self.hidden_size, requires_grad=False)
+            hx = (hx, hx)
+        self.check_forward_hidden(input, hx[0], '[0]')
+        self.check_forward_hidden(input, hx[1], '[1]')
+        return quant_rnn.LSTMCell(
+            input, hx,
+            self.weight_ih, self.weight_hh,
+            self.bias_ih, self.bias_hh,
+            self._input_quantizer, self._weight_quantizer
+        )
+
+
+class GRUCell(QuantRNNCellBase):
+    r"""A gated recurrent unit (GRU) cell
+    """
+
+    def __init__(self, input_size, hidden_size, bias=True):
+        super(GRUCell, self).__init__()
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.bias = bias
+        self.weight_ih = Parameter(torch.Tensor(3 * hidden_size, input_size))
+        self.weight_hh = Parameter(torch.Tensor(3 * hidden_size, hidden_size))
+        if bias:
+            self.bias_ih = Parameter(torch.Tensor(3 * hidden_size))
+            self.bias_hh = Parameter(torch.Tensor(3 * hidden_size))
+        else:
+            self.register_parameter('bias_ih', None)
+            self.register_parameter('bias_hh', None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        stdv = 1.0 / math.sqrt(self.hidden_size)
+        for weight in self.parameters():
+            weight.data.uniform_(-stdv, stdv)
+
+    def forward(self, input, hx=None):
+        self.check_forward_input(input)
+        if hx is None:
+            hx = input.new_zeros(input.size(0), self.hidden_size, requires_grad=False)
+        self.check_forward_hidden(input, hx)
+        return quant_rnn.GRUCell(
+            input, hx,
+            self.weight_ih, self.weight_hh,
+            self.bias_ih, self.bias_hh,
+        )
+
+LSTM = QuantLSTM
+LSTMCell = QuantLSTMCell

tools/pytorch-quantization/pytorch_quantization/nn/modules/tensor_quantizer.py

@@ -0,0 +1,456 @@
+#
+# 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.
+#
+"""TensorQuantizer Module"""
+import math
+from absl import logging
+
+import torch
+from torch import nn
+
+from pytorch_quantization.tensor_quant import QuantDescriptor, tensor_quant, fake_tensor_quant, scaled_e4m3
+from pytorch_quantization.nn.modules.clip import Clip
+
+from pytorch_quantization import calib
+
+import pytorch_quantization.utils as quant_utils
+
+__all__ = ['TensorQuantizer']
+
+
+class TensorQuantizer(nn.Module):
+    """Tensor quantizer module
+
+    This module uses tensor_quant or fake_tensor_quant function to quantize a tensor. And wrappers variable, moving
+    statistics we'd want when training a quantized network.
+
+    Experimental features:
+        ``clip`` stage learns range before enabling quantization.
+        ``calib`` stage runs calibration
+
+    Args:
+        quant_desc: An instance of :func:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+        disabled: A boolean. If True, by pass the whole module returns input. Default False.
+        if_quant: A boolean. If True, run main quantization body. Default True.
+        if_clip: A boolean. If True, clip before quantization and learn amax. Default False.
+        if_calib: A boolean. If True, run calibration. Not implemented yet. Settings of calibration will probably
+            go to :func:`QuantDescriptor <pytorch_quantization.tensor_quant.QuantDescriptor>`.
+
+    Raises:
+
+    Readonly Properties:
+        - axis:
+        - fake_quant:
+        - scale:
+        - step_size:
+
+    Mutable Properties:
+        - num_bits:
+        - unsigned:
+        - amax:
+    """
+
+    _enable_onnx_export = False
+
+    def __init__(self, quant_desc=QuantDescriptor(), disabled=False, if_quant=True, if_clip=False, if_calib=False):
+        """Initialize quantizer and set up required variables"""
+        super(TensorQuantizer, self).__init__()
+        # Expand quant_desc. Use quant_desc.dict would be eaiser, but adding one-by-one explicitly gives more control
+        self._num_bits = quant_desc.num_bits
+        self._fake_quant = quant_desc.fake_quant
+        self._axis = quant_desc.axis
+        self._scale_amax = quant_desc.scale_amax
+        self._learn_amax = quant_desc.learn_amax
+        self._unsigned = quant_desc.unsigned
+        self._narrow_range = quant_desc.narrow_range
+
+        self._scale = None if not quant_desc.fake_quant else 1.
+        self._disabled = disabled
+        self._if_quant = if_quant
+        self._if_clip = False
+        self._if_calib = if_calib
+
+        if quant_desc.amax is not None:
+            self.register_buffer('_amax', torch.tensor(quant_desc.amax))
+
+        # Clip module consumes a lot of memory, so only create it if learn_amax is True
+        if self._learn_amax:
+            init_amax = quant_desc.amax if quant_desc.amax is not None else 1.
+            self.clip = Clip(-init_amax, init_amax, learn_min=True, learn_max=True)
+            # It makes more sense to enable clip stage (which learns amax) if learn_amax is true
+            self.enable_clip()
+        if if_clip:
+            self.enable_clip()
+
+        if quant_desc.calib_method == "histogram":
+            logging.info("Creating histogram calibrator")
+            self._calibrator = calib.HistogramCalibrator(num_bits=self._num_bits,
+                                                         axis=self._axis,
+                                                         unsigned=self._unsigned)
+        elif quant_desc.calib_method == "max":
+            logging.info("Creating Max calibrator")
+            self._calibrator = calib.MaxCalibrator(num_bits=self._num_bits, axis=self._axis, unsigned=self._unsigned)
+
+    # pylint:disable=missing-docstring
+    @property
+    def num_bits(self):
+        return self._num_bits
+
+    @property
+    def maxbound(self):
+        if self._num_bits == (4, 3):
+            return 448.0
+        return (1 << (self._num_bits - 1 + int(self._unsigned))) - 1
+
+    @property
+    def unsigned(self):
+        return self._unsigned
+
+    @property
+    def scale(self):
+        if self._fake_quant:
+            logging.error("Fake quantize mode doesn't use scale explicitly!")
+        if self._scale is None:
+            logging.critical("Accessing scale before quantizing any tensor!")
+        return self._scale
+
+    @property
+    def pre_quant_scale(self):
+        if not hasattr(self, "_pre_quant_scale"):
+            return None
+        return self._pre_quant_scale
+
+    @property
+    def amax(self):
+        if not hasattr(self, "_amax"):
+            return None
+        return self._amax
+
+    @property
+    def step_size(self):
+        if not hasattr(self, "_amax"):
+            logging.error("step_size is undefined under dynamic amax mode!")
+            return None
+        return self._amax / (2.0**(self._num_bits - 1 + int(self._unsigned)) - 1.0)
+
+    @property
+    def axis(self):
+        return self._axis
+
+    @property
+    def fake_quant(self):
+        return self._fake_quant
+
+    @property
+    def narrow_range(self):
+        return self._narrow_range
+
+    def disable(self):
+        """Bypass the module"""
+        self._disabled = True
+
+    def enable(self):
+        self._disabled = False
+
+    def disable_clip(self):
+        """Disable clip stage"""
+        self._if_clip = False
+        self.clip.clip_value_min.requires_grad = False
+        self.clip.clip_value_max.requires_grad = False
+
+    def enable_clip(self):
+        """Enable clip stage"""
+        logging.warning("Enable `clip` stage for amax learning.")
+        if not self._learn_amax:
+            raise ValueError("learn_amax is False. Cannot enable clip.")
+        self.clip.clip_value_min.requires_grad = True
+        self.clip.clip_value_max.requires_grad = True
+        self._if_clip = True
+
+    def disable_calib(self):
+        logging.warning("Disable {}".format(self._calibrator.__class__.__name__))
+        self._if_calib = False
+
+    def enable_calib(self):
+        if self._calibrator is None:
+            raise ValueError("Calibrator was not created, cannot enable calibration.")
+        logging.info("Enable {}".format(self._calibrator.__class__.__name__))
+        self._if_calib = True
+
+    def disable_quant(self):
+        logging.info("Disable `quant` stage.")
+        self._if_quant = False
+
+    def enable_quant(self):
+        logging.info("Enable `quant` stage.")
+        self._if_quant = True
+
+    @amax.setter
+    def amax(self, value):
+        if value is None:
+            logging.error("Setting amax no None is meaningless.")
+        else:
+            if isinstance(value, torch.Tensor):
+                logging.warning("amax setter is not designed to take tensor.")
+            if not hasattr(self, "_amax"):
+                self.register_buffer('_amax', torch.tensor(value))
+            else:
+                value = torch.tensor(value, device=self._amax.device)
+                if self._amax.shape != value.shape:
+                    raise RuntimeError("Changing shape when setting amax is not allowed.")
+                self._amax.data.copy_(value.data)
+
+    @pre_quant_scale.setter
+    def pre_quant_scale(self, value):
+        if value is None:
+            logging.error("Setting pre_quant_scale no None is meaningless.")
+        else:
+            if not hasattr(self, "_pre_quant_scale"):
+                self.register_buffer('_pre_quant_scale', torch.tensor(value))
+            else:
+                value = torch.tensor(value, device=self._pre_quant_scale.device)
+                if self._pre_quant_scale.shape != value.shape:
+                    raise RuntimeError("Changing shape when setting pre_quant_scale is not allowed.")
+                self._pre_quant_scale.data.copy_(value.data)
+
+    @num_bits.setter
+    def num_bits(self, value):
+        self._num_bits = value
+
+    @unsigned.setter
+    def unsigned(self, value):
+        self._unsigned = value
+
+    @narrow_range.setter
+    def narrow_range(self, value):
+        self._narrow_range = value
+
+    # pylint:enable=missing-docstring
+    def load_calib_amax(self, *args, **kwargs):
+        """Load amax from calibrator.
+
+        Updates the amax buffer with value computed by the calibrator, creating it if necessary.
+        *args and **kwargs are directly passed to compute_amax, except "strict" in kwargs. Refer to
+        compute_amax for more details.
+        """
+        strict = kwargs.pop("strict", True)
+        if getattr(self, '_calibrator', None) is None:
+            raise RuntimeError("Calibrator not created.")
+        calib_amax = self._calibrator.compute_amax(*args, **kwargs)
+        if calib_amax is None:
+            err_msg = "Calibrator returned None. This usually happens when calibrator hasn't seen any tensor."
+            if not strict:
+                logging.warning(err_msg)
+                logging.warning("Set amax to NaN!")
+                calib_amax = torch.tensor(math.nan)
+            else:
+                raise RuntimeError(err_msg + " Passing 'strict=False' to `load_calib_amax()` will ignore the error.")
+        logging.warning("Load calibrated amax, shape={}.".format(calib_amax.shape))
+        logging.log_first_n(logging.WARNING, "Call .cuda() if running on GPU after loading calibrated amax.", 1)
+        if not hasattr(self, '_amax'):
+            self.register_buffer("_amax", calib_amax.data)
+        else:
+            self._amax.copy_(calib_amax)
+
+    def init_learn_amax(self):
+        """Initialize learned amax from fixed amax"""
+        if self._learn_amax is False:
+            raise RuntimeError("Called init_learn_amax with learn_amax=False.")
+        logging.warning("Load amax as initial value for amax learning!")
+        if self._amax.numel() != 1:
+            logging.warning("Per channel learned amax not supported. Initializing with max(amax).")
+            init_amax = torch.max(self._amax)
+        else:
+            init_amax = self._amax
+        self.clip.clip_value_min.data.copy_(-init_amax.data)
+        self.clip.clip_value_max.data.copy_(init_amax.data)
+
+    def _get_amax(self, inputs):
+        """get amax from buffer or compute it dynamically."""
+        if hasattr(self, '_amax'):
+            amax = self._amax
+        else:
+            if self._axis is None:
+                reduce_axis = None
+            else:
+                reduce_axis = []
+                # Swap axis to reduce
+                axis = self._axis if isinstance(self._axis, (list, tuple)) else [self._axis]
+                for i in range(inputs.dim()):
+                    if not i in axis:
+                        reduce_axis.append(i)
+            amax = quant_utils.reduce_amax(inputs, axis=reduce_axis, keepdims=True).detach()
+        if self._scale_amax is not None:
+            amax = amax.detach() * self._scale_amax
+
+        amax = amax.data
+
+        # cast amax to float32 if it is in a lower precision dtype
+        if amax.dtype not in (torch.double, torch.float):
+            amax = amax.float()
+
+        return amax
+
+    def _quant_forward(self, inputs):
+        """Quantized forward pass."""
+        if self._learn_amax:
+            inputs = self.clip(inputs)
+            amax = torch.max(-self.clip.clip_value_min, self.clip.clip_value_max).detach()
+        else:
+            amax = self._get_amax(inputs)
+
+        if self._fake_quant:
+            outputs = fake_tensor_quant(inputs, amax, self._num_bits, self._unsigned, self._narrow_range)
+        else:
+            outputs, self._scale = tensor_quant(inputs, amax, self._num_bits, self._unsigned)
+
+        return outputs
+
+    def _check_onnx_readiness(self, inputs):
+        """Check if quantizer is ready for ONNX export."""
+
+        assert hasattr(
+            self, '_amax'), ("Quantizer has not been calibrated. ONNX export requires the quantizer to be calibrated."
+                             "Calibrate and load amax before exporting to ONNX.")
+
+        if self._if_calib:
+            logging.warning("Quantizer is in calibration mode. "
+                            "Please complete calibration before exporting to ONNX for correct results.")
+
+        amax = self._get_amax(inputs)
+
+        # We only support scalar amax for E4M3 ONNX export
+        if isinstance(self.num_bits, tuple):
+            assert amax.numel() == 1, ("E4M3 supports ONNX export only for per-tensor quantization."
+                                       " Per-tensor quantization requires scalar amax. "
+                                       f"Received non-scalar amax of shape: {amax.shape}")
+
+    def forward(self, inputs):
+        """Apply tensor_quant function to inputs
+
+        Args:
+            inputs: A Tensor of type float32.
+
+        Returns:
+            outputs: A Tensor of type output_dtype
+        """
+
+        if self._enable_onnx_export:
+            self._check_onnx_readiness(inputs)
+
+        # Activation scaling for smoothquant
+        if self.pre_quant_scale is not None:
+            inputs = inputs * self.pre_quant_scale
+
+        if self._disabled:
+            return inputs
+
+        outputs = inputs
+
+        if self._if_calib:
+            if self._calibrator is None:
+                raise RuntimeError("Calibrator was not created.")
+            # Shape is only known when it sees the first tensor
+            self._calibrator.collect(inputs)
+
+        if self._if_clip:
+            if not self._learn_amax:
+                raise RuntimeError("Clip without learning amax is not implemented.")
+            outputs = self.clip(inputs)
+
+        if self._if_quant:
+            if not isinstance(self._num_bits, tuple):
+                outputs = self._quant_forward(inputs)
+            else:
+                E, M = self._num_bits
+                outputs = scaled_e4m3(inputs, self._get_amax(inputs), E, M)
+
+        return outputs
+
+    def _short_amax(self, fmt='.4f'):
+        """Short description of amax
+
+        Returns:
+            'dynamic': if _amax is not registered
+            'amax': if _amax is per-tensor
+            '[min, max](size)': if _amax is per-channel
+        """
+        if not hasattr(self, '_amax'):
+            return 'dynamic'
+        if self._amax is None:
+            return "None"
+        if self._amax.numel() == 1:
+            return '{:{fmt}}'.format(self._amax.item(), fmt=fmt)
+        return '[{:{fmt}}, {:{fmt}}]({})'.format(self._amax.min().item(),
+                                                 self._amax.max().item(),
+                                                 self._amax.numel(),
+                                                 fmt=fmt)
+
+    def extra_repr(self):
+        if self._disabled:
+            return "disabled"
+        s = "{}{}bit".format("unsigned " if self._unsigned else "", self._num_bits)
+        s += " narrow" if (self._narrow_range) else ""
+        s += " fake" if (self._fake_quant) else ""
+        s += " axis={}".format(self._axis) if self._axis is not None else " per-tensor"
+        s += " amax={}".format(self._short_amax())
+        s += " *{}".format(self._scale_amax) if self._scale_amax else ""
+        s += " pre_quant_scale" if self.pre_quant_scale is not None else ""
+        s += " learned" if (self._learn_amax) else ""
+        s += " calibrator={}".format(self._calibrator.__class__.__name__) if (self._calibrator is not None) else ""
+        s += " scale={}".format(self._scale) if self._scale is not None else ""
+        s += " quant" if (self._if_quant) else ""
+        s += " clip" if (self._if_clip) else ""
+        s += " calib" if (self._if_calib) else ""
+        return s
+
+    def _load_from_state_dict(self, state_dict, prefix, *args, **kwargs):
+        """Overloaded module function
+
+        Adds warnings during state_dict loading.
+        A workaround is implemented for loading amax from checkpoint and only supports CUDA.
+
+        Args:
+            state_dict: A dict containing the state of the top level module
+            prefix: A string that prefixes all of this modules state in state_dict, e.g. 'model.conv1.'
+        """
+        dst_has_amax = '_amax' in self._buffers
+        src_has_amax = prefix + '_amax' in state_dict
+
+        if not src_has_amax and dst_has_amax:
+            logging.error("{}: No amax in state_dict.".format(prefix[:-1]))
+        elif src_has_amax and not dst_has_amax:
+            logging.debug(("{}: No '_amax' buffer to load amax into."
+                           " '_amax` will be created as WAR for now. "
+                           "This behavior will change in future.").format(prefix[:-1]))
+            self.register_buffer("_amax", state_dict[prefix + '_amax'].data.cuda())
+        elif src_has_amax and dst_has_amax:
+            logging.warning("{}: Overwriting amax.".format(prefix[:-1]))
+
+        dst_has_pre_quant_scale = '_pre_quant_scale' in self._buffers
+        src_has_pre_quant_scale = prefix + '_pre_quant_scale' in state_dict
+
+        if not src_has_pre_quant_scale and dst_has_pre_quant_scale:
+            logging.error("{}: No pre_quant_scale in state_dict.".format(prefix[:-1]))
+        elif src_has_pre_quant_scale and not dst_has_pre_quant_scale:
+            logging.debug(("{}: No '_pre_quant_scale' buffer to load pre_quant_scale into."
+                           " '_pre_quant_scale` will be created as WAR for now. "
+                           "This behavior will change in future.").format(prefix[:-1]))
+            self.register_buffer("_pre_quant_scale", state_dict[prefix + '_pre_quant_scale'].data.cuda())
+        elif src_has_pre_quant_scale and dst_has_pre_quant_scale:
+            logging.warning("{}: Overwriting pre_quant_scale.".format(prefix[:-1]))
+
+        super(TensorQuantizer, self)._load_from_state_dict(state_dict, prefix, *args, **kwargs)

tools/pytorch-quantization/pytorch_quantization/optim/helper.py

@@ -0,0 +1,131 @@
+#
+# 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.
+#
+
+
+"""Helper functions for quant optimizer/trainer"""
+
+import re
+
+from absl import logging
+
+def match_parameters(model, patterns):
+    """Returns an generator over module parameters if name matches key
+
+    It is useful to group parameters, and apply different functions to different group. This function provides an easy
+    way to group them.
+
+    Args:
+        model: A Module
+        patterns: A list of strings that will be used to match parameter names. If parameter name contains any pattern,
+            it will be yield
+
+    Yields:
+        param: Module parameters
+    """
+    for name, param in model.named_parameters():
+        for pattern in patterns:
+            if re.search(pattern, name):
+                yield param
+
+def group_parameters(model, patterns_list, lrs=None, momentums=None, weight_decays=None):
+    """Group parameters for using per-parameters option in optimizer
+
+    Returns a list of dict that matches Pytorch optimizer fashion, see
+    https://pytorch.org/docs/stable/optim.html#per-parameter-options for more details.
+
+    Example:
+        >>> [
+        >>>    {'params': model.base.parameters()},
+        >>>    {'params': model.classifier.parameters(), 'lr': 1e-3}
+        >>> ]
+    Parameters will be grouped w.r.t first level of the keys_list. e.g. `keys_list=[['conv1', 'conv2'], ['conv3']]` will
+    return 2 groups, one with `conv1` and `conv2` in name, and the other with `conv3` in name.
+
+    If lr, momentum or weight_decay are supplied, they will be added to the group as well.
+
+
+    Args:
+        model: A module
+        patterns_list: A list of list of strings. WARNING: patters must be EXCLUSIVE, the function doesn't
+            perform exclusive check.
+        lrs: A list of float with same length as keys_list or None.
+        momentums: A list of float with same length as keys_list or None.
+        weight_decays: A list of float with same length as keys_list or None.
+
+    Returns:
+        param_group: A list of dict
+
+    """
+    param_groups = []
+    for pattern in patterns_list:
+        if not isinstance(pattern, list):
+            raise TypeError("patterns_list must be list of list of patterns")
+        param_groups.append({'params': match_parameters(model, pattern)})
+
+    if lrs is not None:
+        if len(lrs) != len(patterns_list):
+            raise TypeError("len(lrs) must match len(patterns_list)")
+        for i, lr in enumerate(lrs):
+            param_groups[i]['lr'] = lr
+
+    if momentums is not None:
+        if len(momentums) != len(patterns_list):
+            raise TypeError("len(momentums) must match len(patterns_list)")
+        for i, momentum in enumerate(momentums):
+            param_groups[i]['momentum'] = momentum
+
+    if weight_decays is not None:
+        if len(weight_decays) != len(patterns_list):
+            raise TypeError("len(weight_decays) must match len(patterns_list)")
+        for i, weight_decay in enumerate(weight_decays):
+            param_groups[i]['weight_decay'] = weight_decay
+
+    return param_groups
+
+def freeze_parameters(model, patterns):
+    """Set requires_grad to False if patterns match name
+
+    Args:
+        model: A Module
+        patterns: A list of strings that will be used to match parameter names. If parameter name contains any pattern,
+            it will be frozen.
+    """
+    for name, param in model.named_parameters():
+        for pattern in patterns:
+            if re.search(pattern, name):
+                logging.warning("Freeze %s.", name)
+                param.requires_grad = False
+
+def quant_weight_inplace(model):
+    """Make quantization inplace
+
+    Search for quantized modules including QuantConvNd and QuantLinear, make weight quantization in place using
+    weight_quantizer.
+
+    Most publications of quantization aware training uses STE by default, which is really an approximation of
+    derivative of the nondifferentiable quantization function, which works to some extended but by no means the F=ma of
+    the problem.
+    Inplace quantization can be used to implement relax-and-round, which is a common method in Discrete Optimization's
+    or Integer Programming.
+    """
+    for name, module in model.named_modules():
+        if hasattr(module, '_weight_quantizer') and module.weight_quantizer is not None:
+            if not module.weight_quantizer.fake_quant:
+                logging.warning(("In-place real quantization is VERY dangerous and should be used for inference only. "
+                                 "Make sure that is the desired behavior."))
+            logging.warning("In-place quantize weight of %s", name)
+            module.weight.data.copy_(module.weight_quantizer(module.weight))

tools/pytorch-quantization/pytorch_quantization/quant_modules.py

@@ -0,0 +1,182 @@
+#
+# 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.
+#
+"""Dynamically replace the modules with quantized versions."""
+
+from collections import namedtuple
+from contextlib import contextmanager
+
+import torch
+from pytorch_quantization import nn as quant_nn
+
+__all__ = ['initialize', 'deactivate', 'enable_onnx_export']
+
+# Definition of the named tuple that is used to store mapping of the quantized modules
+_quant_entry = namedtuple('quant_entry', 'orig_mod mod_name replace_mod')
+
+# Global member of the file that contains the mapping of quantized modules
+_DEFAULT_QUANT_MAP = [
+    _quant_entry(torch.nn, "Conv1d", quant_nn.QuantConv1d),
+    _quant_entry(torch.nn, "Conv2d", quant_nn.QuantConv2d),
+    _quant_entry(torch.nn, "Conv3d", quant_nn.QuantConv3d),
+    _quant_entry(torch.nn, "ConvTranspose1d", quant_nn.QuantConvTranspose1d),
+    _quant_entry(torch.nn, "ConvTranspose2d", quant_nn.QuantConvTranspose2d),
+    _quant_entry(torch.nn, "ConvTranspose3d", quant_nn.QuantConvTranspose3d),
+    _quant_entry(torch.nn, "Linear", quant_nn.QuantLinear),
+    _quant_entry(torch.nn, "LSTM", quant_nn.QuantLSTM),
+    _quant_entry(torch.nn, "LSTMCell", quant_nn.QuantLSTMCell),
+    _quant_entry(torch.nn, "AvgPool1d", quant_nn.QuantAvgPool1d),
+    _quant_entry(torch.nn, "AvgPool2d", quant_nn.QuantAvgPool2d),
+    _quant_entry(torch.nn, "AvgPool3d", quant_nn.QuantAvgPool3d),
+    _quant_entry(torch.nn, "AdaptiveAvgPool1d", quant_nn.QuantAdaptiveAvgPool1d),
+    _quant_entry(torch.nn, "AdaptiveAvgPool2d", quant_nn.QuantAdaptiveAvgPool2d),
+    _quant_entry(torch.nn, "AdaptiveAvgPool3d", quant_nn.QuantAdaptiveAvgPool3d),
+]
+
+
+class QuantModuleReplacementHelper():
+    """To help replace torch.nn modules with quantized versions.
+
+    This module is used to replace (by monkey patching) the torch.nn modules with their
+    quantized versions as provided by either tool's internal implementation or any other
+    user provided custom module.
+
+    Attributes:
+        orginal_func_map: A dict. Maintains the original torch.nn module mapping.
+        quant_support_list: A list. Contains the names of modules for which a quantized
+            version is provided by the tool.
+        quant_map: A dict. Contains the map of the module name and its quantized versions.
+        quant_switch_opt: A dict. A map to indicate which modules to be left unreplaced with
+            their quantized versions. This dict is updated by a list provided from the user
+            which indicates the modules to leave out in monkey patching.
+
+    """
+
+    def __init__(self):
+
+        # Will hold the original modules to be replaced back
+        self.orginal_func_map = set()
+
+        # Maintains the list of supported quantized modules by the tool as default
+        self.default_quant_map = _DEFAULT_QUANT_MAP
+
+        # Will hold the final quantized modules after checking if user supplied any
+        # custom quantized functions.
+        self.quant_map = set()
+
+    def prepare_state(self, float_module_list=None, custom_map=None):
+        """
+        Prepare the internal variables that would used in the monkey patching mechanism later.
+        1. Set up the list of quantized modules that are supported by the tool for torch.nn.
+        2. Set up the custom mapping for modules other than torch.nn.
+        3. Use the float_module_list to switch off the monkey patching replacement for user indicated modules
+        """
+
+        # For the default quantized modules supported, generate the quant_map
+        for item in self.default_quant_map:
+            if float_module_list is not None and item.mod_name in float_module_list:
+                # Skip this module if this is present in the float_module_list
+                continue
+            else:
+                # append the modules into the variable that will be used in monkey patching
+                self.quant_map.add(item)
+                # also store the original module to be used in reverse monkey patching
+                self.orginal_func_map.add(
+                    _quant_entry(item.orig_mod, item.mod_name, getattr(item.orig_mod, item.mod_name)))
+
+        # Add custom modules to the quant_map
+        if custom_map is not None:
+            for item in custom_map:
+                # append the custom modules to the list that will be used in monkey patching
+                # Note that we convert a tuple to a named tuple here
+                self.quant_map.add(_quant_entry(item[0], item[1], item[2]))
+                # also store the original module in another list which will be used to reverse monkey patching
+                self.orginal_func_map.add(_quant_entry(item[0], item[1], getattr(item[0], item[1])))
+
+    def apply_quant_modules(self):
+        """
+        For the modules registered in the quant_map, simply monkey patch them and also store the
+        original modules so that they could be later replaced back.
+        """
+        for entry in self.quant_map:
+            setattr(entry.orig_mod, entry.mod_name, entry.replace_mod)
+
+    def restore_float_modules(self):
+        """
+        Reverse the effect of monkey patch by using the orginal_func_map to replace back the
+        original modules.
+        """
+        for entry in self.orginal_func_map:
+            setattr(entry.orig_mod, entry.mod_name, entry.replace_mod)
+
+
+def initialize(float_module_list=None, custom_quant_modules=None):
+    """Dynamic module replacement using monkey patching.
+
+    Dynamically monkey patches the modules with their quantized versions. Internally, the
+    state is maintained by a helper class object which helps in replacing the original
+    modules back.
+
+    Args:
+        float_module_list: A list. User supplied list which indicates which modules to not monkey patch.
+        custom_quant_modules: A dict. A mapping provided by user to indicate any other module apart
+            from torch.nn and its corresponding quantized version.
+
+    Returns:
+        nothing.
+
+    Typical usage example:
+
+        # Define the deny list for torch.nn modules and custom map for modules other than torch.nn.
+        float_module_list = ["Linear"]
+        custom_quant_modules = [(torch.nn, "Linear", quant_nn.QuantLinear)]
+        ## Monkey patch the modules
+        pytorch_quantization.quant_modules.initialize(float_module_list, custom_modules)
+        ## Use the quantized modules
+        pytorch_quantization.quant_modules.deactivate()
+    """
+    _quant_module_helper_object.prepare_state(float_module_list, custom_quant_modules)
+    _quant_module_helper_object.apply_quant_modules()
+
+
+def deactivate():
+    """Dynamic module replacement which reverses the monkey patching.
+
+    Dynamically replaces back the original modules that were monkey patched earlier
+    in the initialize() function call using helper class object which maintains the state.
+    """
+    _quant_module_helper_object.restore_float_modules()
+
+
+# Global object that maintains the state of the modules that are replaced.
+_quant_module_helper_object = QuantModuleReplacementHelper()
+
+
+@contextmanager
+def enable_onnx_export():
+    """Context manager to enable onnx export.
+    
+    .. code-block:: python
+
+        with pytorch_quantization.enable_onnx_export():
+            # export onnx model
+            torch.onnx.export(model, ...)
+    
+    """
+    quant_nn.TensorQuantizer._enable_onnx_export = True
+    yield
+
+    quant_nn.TensorQuantizer._enable_onnx_export = False

tools/pytorch-quantization/pytorch_quantization/tensor_quant.py

@@ -0,0 +1,624 @@
+#
+# 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.
+#
+"""Basic tensor quantization functions"""
+import numpy as np
+import numbers
+import yaml
+
+from absl import logging
+
+import torch
+import torch._C._onnx as _C_onnx
+from torch.autograd import Function
+
+from pytorch_quantization import cuda_ext
+
+from torch.onnx import symbolic_helper
+
+
+class ScaledQuantDescriptor():
+    """Supportive descriptor of quantization
+
+    Describe how a tensor should be quantized. A QuantDescriptor and a tensor defines a quantized tensor.
+
+    Args:
+        num_bits: An integer or a tuple of two integers. 
+            Specifically, `num_bits` can be:
+            
+            #. A positive integer argument for integer qunatization. `num_bits` specify 
+                the number of bits used for integer quantization.
+
+            #. Constant integer tuple (4,3) for E4M3 floating point quantization emulating
+                Nvidia's FP8 quantization. E4M3 quantization only supports per-tensor quantization.
+
+            Default: 8.
+        name: Seems a nice thing to have
+
+    Keyword Arguments:
+        fake_quant: A boolean. If True, use fake quantization mode. Default True.
+        axis: None, int or tuple of int. axes which will have its own max for computing scaling factor.
+            If None (the default), use per tensor scale.
+            Must be in the range [-rank(input_tensor), rank(input_tensor)).
+            e.g. For a KCRS weight tensor, quant_axis=(0) will yield per channel scaling.
+            Default None.
+        amax: A float or list/ndarray of floats of user specified absolute max range. If supplied,
+            ignore quant_axis and use this to quantize. If learn_amax is True, will be used to initialize
+            learnable amax. Default None.
+        learn_amax: A boolean. If True, learn amax. Default False.
+        scale_amax: A float. If supplied, multiply amax by scale_amax. Default None. It is useful for some
+            quick experiment.
+        calib_method: A string. One of ["max", "histogram"] indicates which calibration to use. Except the simple
+            max calibration, other methods are all hisogram based. Default "max".
+        unsigned: A Boolean. If True, use unsigned. Default False.
+
+    Raises:
+        TypeError: If unsupported type is passed in.
+
+    Read-only properties:
+        - fake_quant:
+        - name:
+        - learn_amax:
+        - scale_amax:
+        - axis:
+        - calib_method:
+        - num_bits:
+        - amax:
+        - unsigned:
+    """
+
+    def __init__(self, num_bits=8, name=None, **kwargs):
+        if isinstance(num_bits, int):
+            if num_bits < 0:
+                raise ValueError("num_bits must be > 0, not {}.".format(num_bits))
+            if num_bits == 0:
+                logging.error("num_bits is 0. This will result in the tensor being quantized to all zeros."
+                              " This mode should only be used for debugging purposes.")
+        elif num_bits != (4, 3):
+            raise TypeError("num_bits must be a postive integer or tuple (4,3), not {}.".format(type(num_bits)))
+
+        self._num_bits = num_bits
+        if not isinstance(name, str) and name is not None:
+            raise TypeError("name must be a string or None, not {}.".format(type(name)))
+        self._name = name
+
+        self._fake_quant = kwargs.pop('fake_quant', True)
+        self._axis = kwargs.pop('axis', None)
+        if self._axis is not None:
+            logging.debug("Meaning of axis has changed since v2.0. Make sure to update.")
+        self._learn_amax = kwargs.pop('learn_amax', False)
+        if self._learn_amax and self._axis is not None:
+            raise TypeError("axis is ignored and must be None when learn_amax is true, got {}.".format(type(
+                self._axis)))
+        amax = kwargs.pop('amax', None)
+        if amax is not None:
+            if not isinstance(amax, float) and not isinstance(amax, list) and not isinstance(amax, np.ndarray):
+                raise TypeError("amax must be float, list or ndarray, not {}".format(type(amax)))
+            # Make it single precision array
+            self._amax = np.array(amax, dtype=np.float32)
+        else:
+            self._amax = amax
+
+        self._scale_amax = kwargs.pop('scale_amax', None)
+        self._calib_method = kwargs.pop('calib_method', "max")
+        self._unsigned = kwargs.pop('unsigned', False)
+        self._narrow_range = kwargs.pop('narrow_range', False)
+
+        if kwargs:
+            raise TypeError("Unused keys: {}".format(kwargs.keys()))
+
+    # pylint:disable=missing-docstring
+    @property
+    def num_bits(self):
+        return self._num_bits
+
+    @property
+    def fake_quant(self):
+        return self._fake_quant
+
+    @property
+    def axis(self):
+        return self._axis
+
+    @property
+    def amax(self):
+        return self._amax
+
+    @property
+    def learn_amax(self):
+        return self._learn_amax
+
+    @property
+    def scale_amax(self):
+        return self._scale_amax
+
+    @property
+    def name(self):
+        return self._name
+
+    @property
+    def calib_method(self):
+        return self._calib_method
+
+    @property
+    def unsigned(self):
+        return self._unsigned
+
+    @property
+    def narrow_range(self):
+        return self._narrow_range
+
+    # pylint:enable=missing-docstring
+
+    def __str__(self):
+        s = (self._name + ': ') if self._name is not None else 'QuantDescriptor'
+        s += "({}{}bit".format("unsigned " if self._unsigned else "", self._num_bits)
+        s += " fake" if self._fake_quant else " real"
+        s += " axis={}".format(self._axis if self._axis is not None else " per-tensor")
+        if isinstance(self._amax, torch.Tensor):
+            s += " amax={}".format(
+                np.array2string(self._amax.cpu().numpy().flatten(), edgeitems=1, formatter={'all': "{:.2e}".format}))
+        elif self._amax is not None:
+            s += " amax={_amax}"
+            s += " full_range"
+        if self._learn_amax:
+            s += " learn_amax"
+        if self._scale_amax:
+            s += " scale_amax={_scale_amax}"
+        s += ")"
+        return s.format(**self.__dict__)
+
+    def __eq__(self, rhs):
+        """Compare 2 descriptors"""
+        return self.__dict__ == rhs.__dict__
+
+    def dict(self):
+        """Serialize to dict
+
+        The build-in __dict__ method returns all the attributes, which includes those have default value and have
+        protected prefix "_". This method only returns those have values other than the default one and don't have _ in
+        key. Construct a instance by dict returned by this method should get exactly the same instance.
+        """
+        obj_dict = {}
+        obj_dict['num_bits'] = self._num_bits
+        obj_dict['name'] = self._name
+
+        if not self._fake_quant:
+            obj_dict['fake_quant'] = self._fake_quant
+        if self._axis is not None:
+            obj_dict['axis'] = self._axis
+        if self._amax is not None:
+            obj_dict['amax'] = self._amax.tolist()
+        if self._scale_amax is not None:
+            obj_dict['scale_amax'] = self._scale_amax
+        if self._learn_amax:
+            obj_dict['learn_amax'] = self._learn_amax
+        if self._unsigned:
+            obj_dict['unsigned'] = self._unsigned
+
+        return obj_dict
+
+    def to_yaml(self):
+        """Create yaml serialization
+        Some attributes need special treatment to have human readable form, including amax, axis.
+        """
+        obj_dict = self.dict()
+
+        if "axis" in obj_dict:
+            obj_dict['axis'] = list(obj_dict['axis'])
+        return yaml.dump(obj_dict, width=120)
+
+    @classmethod
+    def from_yaml(cls, yaml_str):
+        """Create descriptor from yaml str"""
+        obj_dict = yaml.safe_load(yaml_str)
+
+        if 'axis' in obj_dict:
+            obj_dict['axis'] = tuple(obj_dict['axis'])
+
+        quant_desc = cls(**obj_dict)
+
+        return quant_desc
+
+
+QuantDescriptor = ScaledQuantDescriptor
+
+# Predefined descriptors
+QUANT_DESC_8BIT_PER_TENSOR = QuantDescriptor(num_bits=8)
+QUANT_DESC_UNSIGNED_8BIT_PER_TENSOR = QuantDescriptor(num_bits=8, unsigned=True)
+QUANT_DESC_8BIT_CONV1D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONV3D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW = QuantDescriptor(num_bits=8, axis=(0))
+QUANT_DESC_8BIT_CONVTRANSPOSE1D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(1))
+QUANT_DESC_8BIT_CONVTRANSPOSE2D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(1))
+QUANT_DESC_8BIT_CONVTRANSPOSE3D_WEIGHT_PER_CHANNEL = QuantDescriptor(num_bits=8, axis=(1))
+
+
+@torch.jit.script
+def _fake_tensor_quant_backward(inputs, amax, grad_outputs):
+    zero = grad_outputs.new_zeros(1)
+    grad_inputs = torch.where(inputs.abs() <= amax, grad_outputs, zero)
+    return grad_inputs
+
+
+def _onnx_int8_helper(g, inputs, amax, num_bits, unsigned, narrow_range):
+    assert num_bits == 8, "Only INT8 ONNX export is supported for now."
+    maxbound = (1 << (num_bits - 1 + int(unsigned))) - 1
+
+    if amax.numel() == 1:
+        zero_point, axis = torch.tensor(0.0, device=amax.device), None
+    else:
+        amax_init_shape = amax.shape
+        amax = amax.squeeze().data
+        assert len(amax.shape) == 1, "ONNX does not support multi-axis quantization."
+        zero_point = torch.zeros_like(amax, dtype=torch.int32).data
+        axis = list(amax_init_shape).index(list(amax.shape)[0])
+
+    zero_point = g.op("Constant", value_t=zero_point)
+
+    if not unsigned:
+        assert not narrow_range, "ONNX does not support unsigned narrow range INT8."
+        zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.INT8)
+    else:
+        zero_point = g.op("Cast", zero_point, to_i=_C_onnx.TensorProtoDataType.UINT8)
+
+    scale = amax / maxbound
+    scale.masked_fill_(scale == 0, 1.0)
+    scale = g.op("Constant", value_t=scale)
+
+    input_type = inputs.type().scalarType()
+
+    # Q inputs are currently constrained to FP32 due to a similar limitation in ORT
+    # custom ops, so cast the input if needed.
+    if input_type == "Half" or input_type == "BFloat16":
+        inputs = g.op("Cast", inputs, to_i=_C_onnx.TensorProtoDataType.FLOAT)
+
+    quantized = g.op("QuantizeLinear", inputs, scale, zero_point, axis_i=axis)
+    out = g.op("DequantizeLinear", quantized, scale, zero_point, axis_i=axis)
+
+    # DQ outputs are currently constrained to FP32 due to a similar limitation in ORT
+    # custom ops, so cast the output if needed.
+    if input_type == "Half":
+        out = g.op("Cast", out, to_i=_C_onnx.TensorProtoDataType.FLOAT16)
+    elif input_type == "BFloat16":
+        out = g.op("Cast", out, to_i=_C_onnx.TensorProtoDataType.BFLOAT16)
+
+    return out
+
+
+class FakeTensorQuantFunction(Function):
+    """Fake version of TensorQuantFunction use CUDA extension"""
+
+    @staticmethod
+    @symbolic_helper.parse_args("v", "t", "i", "b", "b")
+    def symbolic(g, inputs, amax, num_bits=8, unsigned=False, narrow_range=True):
+        return _onnx_int8_helper(g, inputs, amax, num_bits, unsigned, narrow_range)
+
+    @staticmethod
+    def forward(ctx, inputs, amax, num_bits=8, unsigned=False, narrow_range=True):
+        ctx.save_for_backward(inputs, amax)
+
+        def legacy_quant_func():
+            # The LegacyFakeTensorQuantFunction support cpu and amax with any shape that can be broadcasted to inputs.
+            outputs, scale = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range)
+            return outputs / scale.to(inputs.dtype)
+
+        if not inputs.is_cuda:
+            outputs = legacy_quant_func()
+        else:
+            try:
+                if amax.numel() == 1:
+                    outputs = cuda_ext.fake_tensor_quant(inputs, amax, num_bits, unsigned, narrow_range)
+                else:
+                    axis = amax.shape.index(amax.numel())
+
+                    outputs = cuda_ext.fake_tensor_quant_with_axis(inputs, amax.squeeze(), axis, num_bits, unsigned,
+                                                                   narrow_range)
+            except (RuntimeError, ValueError) as error:
+                outputs = legacy_quant_func()
+
+        return outputs
+
+    @staticmethod
+    def backward(ctx, grad_outputs):
+        inputs, amax = ctx.saved_tensors
+        return _fake_tensor_quant_backward(inputs, amax, grad_outputs), None, None, None, None
+
+
+def _onnx_fp8_quantize(g, inputs, scale_inv):
+    """Helper Function for Quantization"""
+    output_shape = torch.onnx.symbolic_helper._get_tensor_sizes(inputs)
+
+    # Q inputs are currently constrained to FP32 due to a similar limitation in ORT
+    # custom ops, so cast the input if needed.
+    if inputs.type().scalarType() == "Half" or inputs.type().scalarType() == "BFloat16":
+        inputs = g.op("Cast", inputs, to_i=_C_onnx.TensorProtoDataType.FLOAT)
+
+    scale = g.op("Constant", value_t=torch.tensor(scale_inv))
+    q_op = g.op("trt::TRT_FP8QuantizeLinear", inputs,
+                scale).setType(inputs.type().with_dtype(torch.uint8).with_sizes(output_shape))
+    return q_op
+
+
+def _onnx_fp8_dequantize(g, inputs, scale_inv, otype=None):
+    """Helper Function for Dequantization"""
+    output_shape = torch.onnx.symbolic_helper._get_tensor_sizes(inputs)
+
+    scale = g.op("Constant", value_t=torch.tensor(scale_inv))
+    out = g.op("trt::TRT_FP8DequantizeLinear", inputs,
+               scale).setType(inputs.type().with_dtype(torch.float32).with_sizes(output_shape))
+
+    # DQ outputs are currently constrained to FP32 due to a similar limitation in ORT
+    # custom ops, so cast the output if needed.
+    if otype == "Half":
+        out = g.op("Cast", out, to_i=_C_onnx.TensorProtoDataType.FLOAT16)
+    elif otype == "BFloat16":
+        out = g.op("Cast", out, to_i=_C_onnx.TensorProtoDataType.BFLOAT16)
+    return out
+
+
+class ScaledE4M3Function(Function):
+    """E4M3fy input with scale"""
+
+    @staticmethod
+    @symbolic_helper.parse_args("v", "t", "i", "b", "b")
+    def symbolic(g, inputs, amax=None, E=4, M=3):
+        if amax is None:
+            scale = 1.0
+        else:
+            scale = 448.0 / amax
+            scale = float(scale.masked_fill_(scale == 0, 1.0))
+
+        input_type = inputs.type().scalarType()
+        q_tensor = _onnx_fp8_quantize(g, inputs, 1.0 / scale)
+        return _onnx_fp8_dequantize(g, q_tensor, 1.0 / scale, input_type)
+
+    @staticmethod
+    def forward(ctx, inputs, amax=None, E=4, M=3):
+        if E != 4 or M != 3:
+            raise NotImplementedError("Only support E=4 & M=3 for now.")
+
+        ctx.save_for_backward(inputs)
+        ctx.amax = amax
+        zero_mask = (inputs.abs() < 1. / (1 << 24))
+
+        if amax is None:
+            outputs = cuda_ext.fake_e4m3fy(inputs)
+        else:
+            # FP8 ONNX export requires scalar `scale`.
+            # To simplify implementation, amax is enforced to be a scalar.
+            scale = 448.0 / amax
+            outputs = cuda_ext.fake_e4m3fy(inputs * scale) / scale
+
+        # Zero out values that are tiny.
+        # Tiny values could lead to tiny amax and then large scale which cause overflow/saturation
+        # and won't go back to normal value after dividing by scale. The right behavior is to mark them
+        # as zero which also get rid of inf/nan
+        outputs[zero_mask] = 0.
+
+        return outputs
+
+    @staticmethod
+    def backward(ctx, grad_outputs):
+        inputs, = ctx.saved_tensors
+        amax = torch.tensor(ctx.amax if ctx.amax is not None else 448.0, dtype=torch.float32, device=inputs.device)
+        grad_inputs = _fake_tensor_quant_backward(inputs, amax, grad_outputs)
+        return grad_inputs, None, None, None
+
+
+class TensorQuantFunction(Function):
+    """A universal tensor quantization function
+
+    Take an input tensor, output an quantized tensor. The granularity of scale can be interpreted from the
+    shape of amax.
+    output_dtype indicates whether the quantized value will be stored in integer or float. The reason we want to store
+    it in float is the pytorch function takes the quantized value may not accept integer input, e.g. Conv2D.
+
+    It uses 2^num_bits -1 values instead of 2^num_bits. e.g., for num_bits=8, it uses [-127, 127] instead of [-128, 127]
+    """
+
+    @staticmethod
+    @symbolic_helper.parse_args("v", "t", "i", "b", "b")
+    def symbolic(g, inputs, amax, num_bits=8, unsigned=False, narrow_range=True):
+        return _onnx_int8_helper(g, inputs, amax, num_bits, unsigned, narrow_range)
+
+    @staticmethod
+    def forward(ctx, inputs, amax, num_bits=8, unsigned=False, narrow_range=True):
+        """
+
+        Follow tensorflow convention, max value is passed in and used to decide scale, instead of inputing scale
+        directly. Though inputing scale directly may be more natural to use.
+
+        Args:
+            ctx: A Context object to store tensors for backward.
+            inputs: A Tensor of type float32.
+            amax: A Tensor of type float32. Inputs will be quantized within range [-amax, amax]
+                amax will be broadcasted to inputs tensor.
+            num_bits: A integer used to calculate scaling factor, scale = (2^(num_bits-1) - 1) / max
+                Effectively, it indicates how many integer bits is used to represent the value. Default 8.
+            output_dtype: A type of Tensor. torch.int32 or torch.float32.
+            unsigned: A boolean. Use unsigned integer range. E.g. [0, 255] for num_bits=8. Default False.
+            narrow_range: A boolean. Use symmetric integer range for signed quantization
+                E.g. [-127,127] instead of [-128,127] for num_bits=8. Default True.
+
+        Returns:
+            outputs: A Tensor of type output_dtype.
+            scale: A Tensor of type float32. outputs / scale will dequantize outputs tensor.
+
+        Raises:
+            ValueError:
+        """
+        ctx.save_for_backward(inputs, amax)
+        outputs, scale = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range)
+        # Check if scale overflows FP16
+        if outputs.dtype == torch.half and scale.max() > 65504:
+            raise ValueError("scale is too large for FP16 with amax={}".format(amax))
+        return outputs, scale.to(inputs.dtype)
+
+    @staticmethod
+    def backward(ctx, grad_outputs, grad_scale):
+        """
+        Implements straight through estimation with clipping. For -amax <= input <= amax
+        the gradient passes straight through, otherwise the gradient is zero.
+
+        Args:
+            ctx: A Context object with saved tensors from forward.
+            grad_outputs: A tensor of gradient of outputs.
+            grad_scale: A tensor of gradient of scale.
+
+        Returns:
+            grad_inputs: A tensor of gradient.
+        """
+        inputs, amax = ctx.saved_tensors
+        zero = grad_outputs.new_zeros(1)  # create a zero tensor with the same type and device
+        grad_inputs = torch.where(inputs.abs() <= amax, grad_outputs, zero)
+        return grad_inputs, None, None, None, None
+
+
+class LegacyFakeTensorQuantFunction(Function):
+    """Fake version of TensorQuantFunction
+    See comments of TensorQuantFunction, arguments are the same.
+    """
+
+    @staticmethod
+    def forward(ctx, inputs, amax, num_bits=8, unsigned=False, narrow_range=True):
+        ctx.save_for_backward(inputs, amax)
+        outputs, scale = _tensor_quant(inputs, amax, num_bits, unsigned, narrow_range)
+        return outputs / scale.to(inputs.dtype)
+
+    @staticmethod
+    def backward(ctx, grad_outputs):
+        inputs, amax = ctx.saved_tensors
+        zero = grad_outputs.new_zeros(1)
+        grad_inputs = torch.where(inputs.abs() <= amax, grad_outputs, zero)
+        return grad_inputs, None, None, None, None
+
+
+def _tensor_quant(inputs, amax, num_bits=8, unsigned=False, narrow_range=True):
+    """Shared function body between TensorQuantFunction and FakeTensorQuantFunction"""
+    # Fine scale, per channel scale will be handled by broadcasting, which could be tricky. Pop a warning.
+    if isinstance(amax, torch.Tensor) and inputs.dim() != amax.dim():
+        logging.debug("amax %s has different shape than inputs %s. Make sure broadcast works as expected!", amax.size(),
+                      inputs.size())
+
+    logging.debug("{} bits quantization on shape {} tensor.".format(num_bits, inputs.size()))
+
+    if unsigned:
+        if inputs.min() < 0.:
+            raise TypeError("Negative values encountered in unsigned quantization.")
+
+    # Computation must be in FP32 to prevent potential over flow.
+    input_dtype = inputs.dtype
+    if inputs.dtype == torch.half:
+        inputs = inputs.float()
+    if amax.dtype == torch.half:
+        amax = amax.float()
+
+    min_amax = amax.min()
+    if min_amax < 0:
+        raise ValueError("Negative values in amax")
+
+    max_bound = torch.tensor((2.0**(num_bits - 1 + int(unsigned))) - 1.0, device=amax.device)
+    if unsigned:
+        min_bound = 0
+    elif narrow_range:
+        min_bound = -max_bound
+    else:
+        min_bound = -max_bound - 1
+    scale = max_bound / amax
+
+    epsilon = 1. / (1 << 24)
+    if min_amax <= epsilon:  # Treat amax smaller than minimum representable of fp16 0
+        zero_amax_mask = (amax <= epsilon)
+        scale[zero_amax_mask] = 0  # Value quantized with amax=0 should all be 0
+
+    outputs = torch.clamp((inputs * scale).round_(), min_bound, max_bound)
+
+    if min_amax <= epsilon:
+        scale[zero_amax_mask] = 1.  # Return 1 makes more sense for values quantized to 0 with amax=0
+
+    if input_dtype == torch.half:
+        outputs = outputs.half()
+
+    return outputs, scale
+
+
+class FakeAffineTensorQuantFunction(Function):
+    """Fake version of affine quantization
+
+    gemmlowp style scale+shift quantization. See more details in
+    https://github.com/google/gemmlowp/blob/master/doc/quantization.md.
+
+    We DO NOT recommend affine quantization on weights for performance reason. There might be value to affine quantize
+    activation as it can be cancelled by bias and comes with no performance penalty. This functionality is only added
+    for experimental purpose.
+    """
+
+    @staticmethod
+    def forward(ctx, inputs, min_range, max_range, num_bits=8):
+        """
+
+        As it will be only applied on activation with per tensor granularity, broadcast is not needed.
+
+        Args:
+            ctx: Pytorch convention.
+            inputs: A Tensor of type float32.
+            min_range: A float.
+            max_range: A float.
+            num_bits: An integer
+
+        Returns:
+            outputs: A Tensor of type output_dtype
+        """
+        logging.debug("{} bits quantization on shape {} tensor.".format(num_bits, inputs.size()))
+        ctx.save_for_backward(inputs, min_range, max_range)
+
+        step_size = (max_range - min_range) / (2.0**num_bits - 1)
+
+        min_bound = -2.0**(num_bits - 1)
+        max_bound = 2.0**(num_bits - 1) - 1
+
+        quant_zero = torch.round(min_range / step_size) - min_bound
+        quantized = torch.round(inputs / step_size) - quant_zero
+        quantized = torch.clamp(quantized, min_bound, max_bound)
+
+        outputs = (quantized + quant_zero) * step_size
+
+        return outputs
+
+    @staticmethod
+    def backward(ctx, grad_outputs):
+        """
+        Args:
+            ctx: Pytorch convention.
+            grad_output: A tensor of gradient of outputs
+
+        Returns:
+            grad_inputs: A tensor of gradient
+        """
+
+        inputs, min_range, max_range = ctx.saved_tensors
+        zero = grad_outputs.new_zeros(1)
+        grad_inputs = torch.where((inputs <= max_range) * (inputs >= min_range), grad_outputs, zero)
+        return grad_inputs, None, None, None
+
+
+tensor_quant = TensorQuantFunction.apply
+legacy_fake_tensor_quant = LegacyFakeTensorQuantFunction.apply
+fake_tensor_quant = FakeTensorQuantFunction.apply
+fake_affine_tensor_quant = FakeAffineTensorQuantFunction.apply
+scaled_e4m3 = ScaledE4M3Function.apply

tools/pytorch-quantization/pytorch_quantization/utils/__init__.py

@@ -0,0 +1,21 @@
+#
+# 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.
+#
+
+
+"""Main entry of all utils"""
+
+from .reduce_amax import reduce_amax

tools/pytorch-quantization/pytorch_quantization/utils/quant_logging.py

@@ -0,0 +1,27 @@
+#
+# 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 WAR for codes that messes up logging format"""
+
+import logging
+
+def reset_logger_handler():
+    """Remove all handler in root logger"""
+    root_logger = logging.getLogger()
+    while root_logger.handlers:
+        root_logger.removeHandler(root_logger.handlers[0])

tools/pytorch-quantization/pytorch_quantization/utils/reduce_amax.py

@@ -0,0 +1,63 @@
+#
+# 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.
+#
+
+
+"""Function to get absolute maximum of a tensor
+Follow numpy fashion, which is more generic as pytorch's
+"""
+
+import torch
+
+def reduce_amax(input, axis=None, keepdims=True):
+    """Compute the absolute maximum value of a tensor.
+
+    Reduces input_tensor along the dimensions given in axis. Unless keepdims is true,
+    the rank of the tensor is reduced by 1 for each entry in axis. If keepdims is true,
+    the reduced dimensions are retained with length 1.
+
+    .. note::
+        Gradient computeation is disabled as this function is never meant learning reduces amax
+
+    Args:
+        input: Input tensor
+        axis: The dimensions to reduce. None or int or tuple of ints. If None (the default),
+            reduces all dimensions. Must be in the range [-rank(input_tensor), rank(input_tensor)).
+        keepdims: A boolean. If true, retains reduced dimensions with length 1. Default True
+        granularity: DEPRECTED. specifies if the statistic has to be calculated at tensor or channel granularity
+
+    Returns:
+        The reduced tensor.
+
+    Raises:
+        ValueError: Any axis which doesn't make sense or is not supported
+        ValueError: If unknown granularity is passed in.
+    """
+    with torch.no_grad():
+        output = input.abs()
+        if axis is None:
+            output = torch.max(output)
+        else:
+            if isinstance(axis, int):
+                output, _ = torch.max(output, dim=axis, keepdim=keepdims)
+            else:
+                if isinstance(axis, tuple) and len(axis) > input.dim():
+                    raise ValueError("Cannot reduce more axes than tensor's dim.")
+                for i in axis:
+                    output, _ = torch.max(output, dim=i, keepdim=True)
+                if not keepdims or output.numel() == 1:
+                    output.squeeze_()
+        return output

tools/pytorch-quantization/requirements.txt

@@ -0,0 +1,6 @@
+numpy
+absl-py>=0.7.0
+scipy
+sphinx_glpi_theme
+prettytable
+pyyaml

tools/pytorch-quantization/setup.cfg

@@ -0,0 +1,6 @@
+[aliases]
+test=pytest
+
+[tool:pytest]
+addopts = --verbose
+python_files = tests/*.py

tools/pytorch-quantization/setup.py

@@ -0,0 +1,80 @@
+#
+# 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.
+#
+
+
+"""Simple setup script"""
+
+import os
+from setuptools import setup, find_packages
+from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension
+
+abspath = os.path.dirname(os.path.realpath(__file__))
+
+with open(os.path.join(abspath, "requirements.txt")) as f:
+    requirements = f.read().splitlines()
+
+license_header = """#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+"""
+
+# Generate version file
+with open(os.path.join(abspath, "VERSION")) as f:
+    version = f.read().strip()
+with open(os.path.join(abspath, "pytorch_quantization/version.py"), "w") as f:
+    f.write(license_header)
+    f.write(F"__version__ = \"{version}\"")
+
+setup(
+    name="pytorch_quantization",
+    version=version,
+    description="NVIDIA Pytorch quantization toolkit",
+    packages=find_packages(exclude=["*.tests", "*.tests.*", "tests.*", "tests"]),
+    setup_requires=["pytest-runner"],
+    tests_require=["pytest"],
+
+    install_requires=requirements,
+    ext_modules=[
+        CUDAExtension(
+            name="pytorch_quantization.cuda_ext",
+            sources=[os.path.join(abspath, "src/tensor_quant.cpp"),
+                     os.path.join(abspath, "src/tensor_quant_gpu.cu")])
+    ],
+    cmdclass={
+        "build_ext": BuildExtension
+    },
+    zip_safe=False,
+
+    long_description=open("README.md", "r", encoding="utf-8").read(),
+    url="https://github.com/nvidia/tensorrt/tools/pytorch-quantization",
+    author="NVIDIA",
+    author_email="nvidia@nvidia.com",
+    license="Apache 2.0",
+)

tools/pytorch-quantization/src/tensor_quant.cpp

@@ -0,0 +1,95 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <cuda.h>
+#include <torch/extension.h>
+
+void fake_tensor_quant_cuda_inplace(at::Tensor, at::Tensor, int, bool, bool);
+at::Tensor fake_tensor_quant_cuda(at::Tensor, at::Tensor, int, bool, bool);
+at::Tensor fake_tensor_quant_with_axis_cuda(at::Tensor, at::Tensor, int, int, bool, bool);
+float bits_to_bound(int, int);
+at::Tensor fake_e4m3fy_cuda(at::Tensor inputs);
+
+void fake_tensor_quant_(at::Tensor inputs, at::Tensor amax, int num_bits = 8,
+                        bool is_unsigned = false, bool narrow_range = true) {
+  TORCH_CHECK(inputs.is_cuda());
+  TORCH_CHECK(inputs.is_contiguous())  // in-place on non-contiguous tensor is more difficult
+  TORCH_CHECK(amax.numel(), 1);
+  fake_tensor_quant_cuda_inplace(inputs, amax, num_bits, is_unsigned, narrow_range);
+}
+
+at::Tensor fake_tensor_quant(at::Tensor inputs, at::Tensor amax, int num_bits = 8,
+                             bool is_unsigned = false, bool narrow_range = true) {
+  TORCH_CHECK(inputs.is_cuda());
+  TORCH_CHECK(amax.numel(), 1);
+  return fake_tensor_quant_cuda(inputs.contiguous(), amax.contiguous(), num_bits, is_unsigned, narrow_range);
+}
+
+at::Tensor fake_tensor_quant_with_axis(at::Tensor inputs, at::Tensor amax, int axis,
+                                       int num_bits = 8, bool is_unsigned = false,
+                                       bool narrow_range = true) {
+  TORCH_CHECK(inputs.is_cuda());
+  TORCH_CHECK(amax.numel(), inputs.size(axis));
+  return fake_tensor_quant_with_axis_cuda(
+      inputs.contiguous(), amax.contiguous(), axis, num_bits, is_unsigned, narrow_range);
+}
+
+
+#if CUDA_VERSION > 11070
+
+  #include <cuda_fp8.h>
+  #include <ATen/ATen.h>
+
+  at::Tensor fake_e4m3fy_cuda(at::Tensor inputs);
+
+  at::Tensor fake_e4m3fy(at::Tensor inputs) {
+    if (inputs.is_cuda()) {
+      return fake_e4m3fy_cuda(inputs.contiguous());
+    } else {
+      TORCH_CHECK(inputs.dtype() == at::ScalarType::Float);
+      TORCH_CHECK(inputs.is_contiguous());
+      auto out = at::zeros_like(inputs);
+      for (int i = 0; i < inputs.numel(); ++i) {
+        out.data_ptr<float>()[i] = static_cast<float>(static_cast<__nv_fp8_e4m3>(inputs.data_ptr<float>()[i]));
+      }
+      return out;
+    }
+  }
+
+#else
+
+  #include <stdexcept>
+
+  at::Tensor fake_e4m3fy(at::Tensor inputs) {
+    throw std::runtime_error("FP8 emulation is not supported on CUDA 11.7 and below");
+  }
+
+#endif
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("fake_tensor_quant_", &fake_tensor_quant_, "Fake Tensor Quant Inplace", py::arg("inputs"),
+        py::arg("amax"), py::arg("num_bits") = 8, py::arg("unsigned") = false,
+        py::arg("narrow_range") = true);
+  m.def("fake_tensor_quant", &fake_tensor_quant, "Fake Tensor Quant", py::arg("inputs"),
+        py::arg("amax"), py::arg("num_bits") = 8, py::arg("unsigned") = false,
+        py::arg("narrow_range") = true);
+  m.def("fake_tensor_quant_with_axis", &fake_tensor_quant_with_axis, "Fake Tensor Quant with axis",
+        py::arg("inputs"), py::arg("amax"), py::arg("axis"), py::arg("num_bits") = 8,
+        py::arg("unsigned") = false, py::arg("narrow_range") = true);
+
+  m.def("fake_e4m3fy", &fake_e4m3fy, "Reduce precision to E4M3",
+        py::arg("inputs"));
+}

tools/pytorch-quantization/src/tensor_quant_gpu.cu

@@ -0,0 +1,159 @@
+/*
+ * 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.
+ */
+
+#include <ATen/ATen.h>
+#include <cuda.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime.h>
+#include <math.h>
+#include <torch/extension.h>
+
+#define BLOCK_SIZE 128
+#define EPSILON (1. / (1<<24))  // Minimum representable of fp16
+
+#define AT_DISPATCH_CASE_FLOATING_TYPES(...)   \
+  AT_DISPATCH_CASE(at::ScalarType::Double, __VA_ARGS__)  \
+  AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__)  \
+  AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__)   \
+  AT_DISPATCH_CASE(at::ScalarType::BFloat16, __VA_ARGS__)
+
+#define AT_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
+  AT_DISPATCH_SWITCH(                                        \
+      TYPE, NAME, AT_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
+
+__host__ __device__ float bits_to_bound(int num_bits, int is_unsigned) {
+  float bound = (1 << (num_bits - 1 + int(is_unsigned))) - 1;
+  return bound;
+}
+
+__device__ float fake_tensor_quant_device(float input, float amax, int min_bound, int max_bound) {
+  CUDA_KERNEL_ASSERT(amax >= 0);
+
+  if (amax < EPSILON) {
+    return 0.f;
+  }
+
+  float scale = max_bound / amax;
+  float output = rint(input * scale);
+  output = output > max_bound ? max_bound : output;
+  output = output < min_bound ? min_bound : output;
+
+  return output / scale;
+}
+
+template <typename T>
+__global__ void fake_tensor_quant_kernel(const T* inputs, size_t n, T* outputs, const float* amax,
+                                         int num_bits = 8, bool is_unsigned = false,
+                                         bool narrow_range = true) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+  if (tid < n) {
+    if (is_unsigned) {
+      CUDA_KERNEL_ASSERT(inputs[tid] >= 0);
+    }
+    float bound = bits_to_bound(num_bits, is_unsigned);
+    float max_bound = bound;
+    float min_bound = -(bound + !narrow_range);
+    outputs[tid] = fake_tensor_quant_device((float)inputs[tid], amax[0], min_bound, max_bound);
+  }
+}
+
+void fake_tensor_quant_cuda_inplace(at::Tensor inputs, at::Tensor amax, int num_bits = 8,
+                                    bool is_unsigned = false, bool narrow_range = true) {
+  size_t numel = inputs.numel();
+  AT_DISPATCH_FLOATING_TYPES(
+      inputs.type().scalarType(), "fake_tensor_quant_cuda_inplace", [&] {
+        fake_tensor_quant_kernel<<<numel / BLOCK_SIZE + 1, BLOCK_SIZE>>>(
+            inputs.data_ptr<scalar_t>(), numel, inputs.data_ptr<scalar_t>(),
+            amax.to(at::ScalarType::Float).data_ptr<float>(), num_bits, is_unsigned);
+      });
+}
+
+at::Tensor fake_tensor_quant_cuda(at::Tensor inputs, at::Tensor amax, int num_bits = 8,
+                                  bool is_unsigned = false, bool narrow_range = true) {
+  size_t numel = inputs.numel();
+  auto outputs = torch::empty_like(inputs);
+  AT_DISPATCH_FLOATING_TYPES(inputs.type().scalarType(), "fake_tensor_quant_cuda", [&] {
+    fake_tensor_quant_kernel<<<numel / BLOCK_SIZE + 1, BLOCK_SIZE>>>(
+        inputs.data_ptr<scalar_t>(), numel, outputs.data_ptr<scalar_t>(),
+        amax.to(at::ScalarType::Float).data_ptr<float>(), num_bits, is_unsigned);
+  });
+
+  return outputs;
+}
+
+template <typename T>
+__global__ void fake_tensor_quant_with_axis_cuda_kernel(const T* inputs, size_t n,
+                                                        T* outputs, const float* amax,
+                                                        int axis_size, int outer_size,
+                                                        int num_bits = 8, bool is_unsigned = false,
+                                                        bool narrow_range = true) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+  float bound = bits_to_bound(num_bits, is_unsigned);
+  float max_bound = bound;
+  float min_bound = -(bound + !narrow_range);
+
+  for (int idx = 4 * tid; idx < 4 * (tid + 1) && idx < n; ++idx) {
+    if (is_unsigned) {
+      CUDA_KERNEL_ASSERT(inputs[idx] >= 0);
+    }
+    int axis_idx = (idx / outer_size) % axis_size;
+
+    outputs[idx] = fake_tensor_quant_device((float)inputs[idx], amax[axis_idx], min_bound, max_bound);
+  }
+}
+
+at::Tensor fake_tensor_quant_with_axis_cuda(at::Tensor inputs, at::Tensor amax, int axis,
+                                            int num_bits = 8, bool is_unsigned = false,
+                                            bool narrow_range = true) {
+  auto outputs = torch::empty_like(inputs);
+  size_t numel = inputs.numel();
+  int axis_size = inputs.size(axis);
+
+  int outer_size = inputs.stride(axis);
+
+  AT_DISPATCH_FLOATING_TYPES(inputs.type().scalarType(), "fake_tensor_quant_cuda_with_axis", [&] {
+    fake_tensor_quant_with_axis_cuda_kernel<<<numel / (BLOCK_SIZE * 4) + 1, BLOCK_SIZE>>>(
+        inputs.data_ptr<scalar_t>(), numel, outputs.data_ptr<scalar_t>(),
+        amax.to(at::ScalarType::Float).data_ptr<float>(), axis_size, outer_size, num_bits, is_unsigned);
+  });
+  return outputs;
+}
+
+#if CUDA_VERSION > 11070
+  #include <cuda_fp8.h>
+
+  template <typename T>
+  __global__ void fake_e4m3fy_kernel(const T* inputs, size_t n, T* outputs) {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+
+    for (int idx = 4 * tid; idx < 4 * (tid + 1) && idx < n; ++idx) {
+      outputs[idx] = static_cast<T>(static_cast<float>(static_cast<__nv_fp8_e4m3>(static_cast<float>(inputs[idx]))));
+    }
+  }
+
+  at::Tensor fake_e4m3fy_cuda(at::Tensor inputs) {
+    size_t numel = inputs.numel();
+    auto outputs = torch::empty_like(inputs);
+    AT_DISPATCH_FLOATING_TYPES(inputs.type().scalarType(), "fake_e4m3fy_cuda", [&] {
+      fake_e4m3fy_kernel<<<numel / (BLOCK_SIZE * 4) + 1, BLOCK_SIZE>>>(
+          inputs.data_ptr<scalar_t>(), numel, outputs.data_ptr<scalar_t>());
+    });
+    return outputs;
+  }
+#endif

tools/pytorch-quantization/tests/calibrator_test.py

@@ -0,0 +1,452 @@
+#
+# 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.
+#
+
+
+"""Tests of calibrators"""
+import pytest
+import numpy as np
+
+import torch
+
+from pytorch_quantization import utils as quant_utils
+from pytorch_quantization import calib
+from pytorch_quantization import nn as quant_nn
+import tests.utils as test_utils
+from tests.fixtures import verbose
+from tests.fixtures.models import QuantLeNet
+
+np.random.seed(12345)
+torch.manual_seed(12345)
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestMaxCalibrator():
+
+    def test_simple_run(self):
+        max_calibrator = calib.MaxCalibrator(8, None, False)
+
+        x_1 = torch.rand(129).cuda()
+        x_2 = torch.rand(127).cuda()
+        max_calibrator.collect(x_1)
+        max_calibrator.collect(x_2)
+
+        test_utils.compare(max_calibrator.compute_amax(), torch.max(x_1.max(), x_2.max()), atol=0, rtol=0, ctol=0)
+
+        # Nothing to test other than creation
+        max_calibrator = calib.MaxCalibrator(8, None, True)
+
+    def test_fine_grain(self):
+        axis = 0
+        reducs_axis = (1, 2, 3)
+        max_calibrator = calib.MaxCalibrator(8, axis, False)
+
+        x_1 = torch.rand(31, 63, 7, 7).cuda()
+        x_2 = torch.rand(31, 63, 7, 7).cuda()
+        max_calibrator.collect(x_1)
+        max_calibrator.collect(x_2)
+
+        assert max_calibrator.compute_amax().shape[0] == 31
+
+        test_utils.compare(max_calibrator.compute_amax(),
+                           quant_utils.reduce_amax(torch.max(x_1, x_2), axis=reducs_axis),
+                           atol=0, rtol=0, ctol=0)
+
+        max_calibrator.reset()
+        assert max_calibrator.compute_amax() is None
+
+    def test_reverse_axis(self):
+        axis = -4
+        reducs_axis = (1, 2, 3)
+        max_calibrator = calib.MaxCalibrator(8, axis, False)
+
+        x_1 = torch.rand(31, 63, 7, 7).cuda()
+        x_2 = torch.rand(31, 63, 7, 7).cuda()
+        max_calibrator.collect(x_1)
+        max_calibrator.collect(x_2)
+
+        assert max_calibrator.compute_amax().shape[0] == 31
+
+        test_utils.compare(max_calibrator.compute_amax(),
+                           quant_utils.reduce_amax(torch.max(x_1, x_2), axis=reducs_axis),
+                           atol=0, rtol=0, ctol=0)
+
+        max_calibrator.reset()
+        assert max_calibrator.compute_amax() is None
+
+    def test_raises(self):
+        axis = 0
+        max_calibrator = calib.MaxCalibrator(8, axis, False)
+
+        x_2 = torch.rand(32, 63, 7, 7).cuda()
+        x_3 = torch.rand(33, 63, 7, 7).cuda()
+        max_calibrator.collect(x_2)
+        with pytest.raises(RuntimeError, match="shape changed"):
+            max_calibrator.collect(x_3)
+
+    def test_track_amax(self):
+        max_calibrator = calib.MaxCalibrator(8, None, False, track_amax=True)
+
+        x_1 = torch.rand(129).cuda()
+        x_2 = torch.rand(127).cuda()
+        max_calibrator.collect(x_1)
+        max_calibrator.collect(x_2)
+
+        test_utils.compare(max_calibrator.compute_amax(), torch.max(x_1.max(), x_2.max()), atol=0, rtol=0, ctol=0)
+        np.testing.assert_array_equal(max_calibrator.amaxs[0], x_1.max().cpu().numpy())
+        np.testing.assert_array_equal(max_calibrator.amaxs[1], x_2.max().cpu().numpy())
+
+    def test_repr(self):
+        max_calibrator = calib.MaxCalibrator(8, None, False, track_amax=True)
+        repr(max_calibrator)
+
+class TestHistogramCalibrator():
+
+    def test_grow(self, verbose):
+        x_1 = torch.tensor([0, 255, 255, 255, 255, 255]).cuda()
+        x_2 = torch.tensor([0, 255, 255, 255, 255, 256]).cuda()
+
+        hist_calibrator = calib.HistogramCalibrator(8, None, False, grow_method='stretch')
+        hist_calibrator.collect(x_1)
+        hist_calibrator.collect(x_2)
+
+        amax = hist_calibrator.compute_amax(method='entropy')
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be closer to 256 because the last bin gets stretched to (~255, 257)
+        assert (amax - 255.).abs() < (amax - 256.).abs()
+
+        hist_calibrator = calib.HistogramCalibrator(8, None, False, grow_method='append')
+        hist_calibrator.collect(x_1)
+        hist_calibrator.collect(x_2)
+
+        amax = hist_calibrator.compute_amax(method='mse')
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be closer to 255
+        assert (amax - 255.).abs() < 0.5
+
+    def test_skip_zeros(self, verbose):
+        x_1 = torch.tensor([0, 0, 0, 0, 0, 1, 2, 3, 4, 5])
+        x_2 = torch.tensor([0, 0, 0, 0, 0, 6, 7, 8, 9, 10])
+
+        calibrator = calib.HistogramCalibrator(8, None, False, skip_zeros=True)
+        calibrator.collect(x_1)
+        calibrator.collect(x_2)
+
+        amax = calibrator.compute_amax("percentile", percentile=50)
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be close to 5
+        assert (amax - 5.).abs() < 10/2048
+
+    def test_torch_hist(self):
+        x_1 = torch.rand(1023, device="cuda")
+        x_1[0] = 0
+        x_2 = torch.rand(1023, device="cuda") + 1  # Make sure histogram bins need to be grown
+        x_2[1] = 0
+
+        calibrator_np = calib.HistogramCalibrator(8, None, False, num_bins=19, torch_hist=False)
+        calibrator_torch = calib.HistogramCalibrator(8, None, False, num_bins=19, torch_hist=True)
+
+        calibrator_np.collect(x_1)
+        calibrator_torch.collect(x_1)
+        assert calibrator_torch._calib_hist.numel() == calibrator_torch._calib_bin_edges.numel() - 1
+        np.testing.assert_array_equal(calibrator_np._calib_hist, calibrator_torch._calib_hist.cpu().numpy())
+        np.testing.assert_array_almost_equal(
+            calibrator_np._calib_bin_edges, calibrator_torch._calib_bin_edges.cpu().numpy())
+
+        # Test multiple collections with some of them needs to expand range
+        for _ in range(3):
+            calibrator_np.collect(x_2)
+            calibrator_torch.collect(x_2)
+            calibrator_np.collect(x_1)
+            calibrator_torch.collect(x_1)
+
+            # Test compute_amax function doesn't convert _calib_hist and _calib_bin_edges unnecessarily
+            calibrator_np.compute_amax("percentile", percentile=99.99)
+            calibrator_torch.compute_amax("percentile", percentile=99.99)
+
+            np.testing.assert_array_equal(calibrator_np._calib_hist, calibrator_torch._calib_hist.cpu().numpy())
+            np.testing.assert_array_almost_equal(
+                calibrator_np._calib_bin_edges, calibrator_torch._calib_bin_edges.cpu().numpy())
+            assert calibrator_torch._calib_hist.numel() == calibrator_torch._calib_bin_edges.numel() - 1
+
+
+class TestEntropyCalibrator():
+
+    def test_one_tensor(self, verbose):
+        hist_calibrator = calib.HistogramCalibrator(8, None, False, grow_method='stretch')
+
+        x_2 = torch.rand(11, 7, 3, 3).cuda() # uniform in (0,1)
+        x_2[1, 1, 1, 1] = 10. # create outlier
+        hist_calibrator.collect(x_2)
+
+        # Don't have a better test metric. One outlier 10 should be discared by KL-divergence
+        amax = hist_calibrator.compute_amax("entropy")
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        assert amax < 1.1
+
+    def test_unsigned(self, verbose):
+        hist_calibrator = calib.HistogramCalibrator(8, None, True, grow_method='stretch')
+
+        x_2 = torch.rand(11, 7, 3, 3).cuda() # uniform in (0,1)
+        x_2[1, 1, 1, 1] = 10. # create outlier
+        hist_calibrator.collect(x_2)
+
+        amax = hist_calibrator.compute_amax("entropy")
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        assert amax < 1.1
+
+    @pytest.mark.parametrize("torch_hist", [False, True])
+    def test_two_tensor(self, torch_hist, verbose):
+        hist_calibrator = calib.HistogramCalibrator(8, None, False, torch_hist=torch_hist)
+
+        x_2 = torch.rand(11, 7, 3, 3).cuda() # uniform in (0,1)
+        x_2[1, 1, 1, 1] = 10. # create outlier
+
+
+        x_2 = torch.rand(11, 7, 3, 3).cuda() # uniform in (0,1)
+        x_2[1, 1, 1, 1] = 10. # create outlier
+        hist_calibrator.collect(x_2)
+        x_3 = torch.rand(11, 7, 3, 3).cuda()
+        hist_calibrator.collect(x_3)
+
+        # Don't have a better test metric. One outlier 10 should be discared by KL-divergence
+        amax = hist_calibrator.compute_amax("entropy")
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        assert amax < 1.1
+
+    def test_repr(self):
+        hist_calibrator = calib.HistogramCalibrator(8, None, True)
+        repr(hist_calibrator)
+
+class TestMSECalibrator():
+
+    def test_one_tensor(self, verbose):
+        calibrator = calib.HistogramCalibrator(8, None, False)
+
+        x_1 = torch.ones(11, 7, 3, 3).cuda() * 255.
+        x_1[1, 1, 1, 1] = 256. # create an outlier
+        calibrator.collect(x_1)
+
+        amax = calibrator.compute_amax("mse")
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be closer to 255
+        assert (amax - 255.).abs() < (amax - 256.).abs()
+
+    def test_unsigned_one_tensor(self, verbose):
+        calibrator = calib.HistogramCalibrator(8, None, True)
+
+        x_1 = torch.ones(11, 7, 3, 3).cuda() * 512.
+        x_1[1, 1, 1, 1] = 513. # create an outlier
+        calibrator.collect(x_1)
+
+        amax = calibrator.compute_amax("mse")
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be closer to 512
+        assert (amax - 512.).abs() < (amax - 513.).abs()
+
+    @pytest.mark.parametrize("torch_hist", [False, True])
+    def test_two_tensor(self, torch_hist, verbose):
+        calibrator = calib.HistogramCalibrator(8, None, False, torch_hist=torch_hist)
+
+        x_1 = torch.ones(11, 7, 3, 3).cuda() * 255.
+        x_1[1, 1, 1, 1] = 256. # create an outlier
+        calibrator.collect(x_1)
+        x_2 = torch.ones(11, 7, 3, 3).cuda() * 255.
+        calibrator.collect(x_2)
+
+        amax = calibrator.compute_amax("mse")
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be closer to 255
+        assert (amax - 255.).abs() < (amax - 256.).abs()
+
+    def test_repr(self):
+        calibrator = calib.HistogramCalibrator(8, None, False)
+        repr(calibrator)
+
+class TestPercentileCalibrator():
+
+    def test_one_tensor(self, verbose):
+        calibrator = calib.HistogramCalibrator(8, None, False)
+
+        x_1 = torch.arange(100)
+        calibrator.collect(x_1)
+
+        amax = calibrator.compute_amax("percentile", percentile=90)
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be approximately 89
+        assert (amax - 89.).abs() < 100/1024
+
+    def test_unsigned_one_tensor(self, verbose):
+        calibrator = calib.HistogramCalibrator( 8, None, True)
+
+        x_1 = torch.arange(100)
+        calibrator.collect(x_1)
+
+        amax = calibrator.compute_amax("percentile", percentile=80)
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be approximately 79
+        assert (amax - 79.).abs() < 100/2048
+
+    @pytest.mark.parametrize("torch_hist", [False, True])
+    def test_two_tensor(self, torch_hist, verbose):
+        calibrator = calib.HistogramCalibrator(8, None, False, torch_hist=torch_hist)
+
+        x_1 = torch.arange(100)
+        calibrator.collect(x_1)
+        x_2 = torch.arange(0, 50, 0.5)
+        calibrator.collect(x_2)
+        amax = calibrator.compute_amax("percentile", percentile=99)
+
+        if verbose:
+            print('amax={:.4f}'.format(amax.item()), end=' ')
+
+        # amax should be approximately 97
+        assert (amax - 97.).abs() < 100/1024
+
+    def test_repr(self):
+        calibrator = calib.HistogramCalibrator(8, None, False)
+        repr(calibrator)
+
+    def test_range(self):
+        calibrator = calib.HistogramCalibrator(8, None, False)
+        x_1 = torch.arange(100)
+        calibrator.collect(x_1)
+        with pytest.raises(ValueError, match="range"):
+            calibrator.compute_amax("percentile", percentile=-10)
+        with pytest.raises(ValueError, match="range"):
+            calibrator.compute_amax("percentile", percentile=200)
+
+class TestCalibrateWeights():
+
+    def test_max(self):
+        torch.manual_seed(12345)
+        ref_lenet = QuantLeNet()
+        torch.manual_seed(12345)
+        test_lenet = QuantLeNet()
+
+        for module in ref_lenet.modules():
+            if isinstance(module, (quant_nn.QuantConv2d, quant_nn.QuantLinear)):
+                module.weight_quantizer.enable_calib()
+                module.weight_quantizer.disable_quant()
+                module.weight_quantizer(module.weight)
+                module.weight_quantizer.load_calib_amax()
+
+        calib.calibrate_weights(test_lenet, method="max")
+
+        for ref_module, test_module in zip(ref_lenet.modules(), test_lenet.modules()):
+            if isinstance(ref_module, (quant_nn.QuantConv2d, quant_nn.QuantLinear)):
+                test_utils.compare(
+                    ref_module.weight_quantizer.amax, test_module.weight_quantizer.amax, rtol=0, atol=0, ctol=0)
+                assert ref_module.weight_quantizer.amax.shape == test_module.weight_quantizer.amax.shape
+
+    def test_shape_with_axis(self):
+        """Check calibrate_weight function returns same shape as TensorQuantizer"""
+        torch.manual_seed(12345)
+        ref_lenet = QuantLeNet()
+        torch.manual_seed(12345)
+        test_lenet = QuantLeNet()
+
+        for module in ref_lenet.modules():
+            if isinstance(module, (quant_nn.QuantConv2d, quant_nn.QuantLinear)):
+                module.weight_quantizer.enable_calib()
+                module.weight_quantizer.disable_quant()
+                module.weight_quantizer(module.weight)
+                module.weight_quantizer.load_calib_amax()
+
+        calib.calibrate_weights(test_lenet, method="percentile")
+
+        for ref_module, test_module in zip(ref_lenet.modules(), test_lenet.modules()):
+            if isinstance(ref_module, (quant_nn.QuantConv2d, quant_nn.QuantLinear)):
+                assert ref_module.weight_quantizer.amax.shape == test_module.weight_quantizer.amax.shape
+
+    def test_percentile(self):
+        torch.manual_seed(12345)
+        test_lenet = QuantLeNet()
+        test_percentile = 99.99
+
+        ref_calibrator = calib.HistogramCalibrator(8, None, False)
+
+        calib.calibrate_weights(test_lenet, method="percentile", perchannel=False, percentile=test_percentile)
+        ref_calibrator.collect(test_lenet.conv1.weight)
+        ref_amax = ref_calibrator.compute_amax("percentile", percentile=test_percentile)
+        test_utils.compare(ref_amax, test_lenet.conv1.weight_quantizer.amax, rtol=0, atol=0, ctol=0)
+
+    def test_percentile_with_axis(self):
+        torch.manual_seed(12345)
+        test_lenet = QuantLeNet()
+        test_percentile = 99.99
+
+        ref_calibrator = calib.HistogramCalibrator(8, None, False)
+
+        calib.calibrate_weights(test_lenet, method="percentile", perchannel=True, percentile=test_percentile)
+        ref_calibrator.collect(test_lenet.conv2.weight[1])
+        ref_amax = ref_calibrator.compute_amax("percentile", percentile=test_percentile)
+        test_utils.compare(ref_amax, test_lenet.conv2.weight_quantizer.amax[1], rtol=0, atol=0, ctol=0)
+
+    def test_mse(self):
+        torch.manual_seed(12345)
+        test_lenet = QuantLeNet()
+
+        ref_calibrator = calib.HistogramCalibrator(8, None, False)
+
+        calib.calibrate_weights(test_lenet, method="mse", perchannel=False)
+        ref_calibrator.collect(test_lenet.conv1.weight)
+        ref_amax = ref_calibrator.compute_amax("mse")
+        test_utils.compare(ref_amax, test_lenet.conv1.weight_quantizer.amax, rtol=0, atol=0, ctol=0)
+
+    def test_mse_with_axis(self):
+        torch.manual_seed(12345)
+        test_lenet = QuantLeNet()
+
+        ref_calibrator = calib.HistogramCalibrator(8, None, False)
+
+        calib.calibrate_weights(test_lenet, method="mse", perchannel=True)
+        ref_calibrator.collect(test_lenet.conv2.weight[1])
+        ref_amax = ref_calibrator.compute_amax("mse")
+        test_utils.compare(ref_amax, test_lenet.conv2.weight_quantizer.amax[1], rtol=0, atol=0, ctol=0)

tools/pytorch-quantization/tests/classification_flow_test.py

@@ -0,0 +1,72 @@
+#
+# 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.
+#
+
+
+"""Tests of the classification flow"""
+import os
+import subprocess
+import sys
+from os import path
+import glob
+import pytest
+# pylint:disable=missing-docstring, no-self-use
+
+class TestClassificationFlow():
+
+    def test_resnet18(self, request, pytestconfig):
+        dir_path = os.path.dirname(os.path.realpath(__file__))
+        dataset_dir = pytestconfig.getoption('--data-dir')
+
+        # skip if the data dir flag was not set
+        if not dataset_dir:
+            pytest.skip("Prepare required dataset and use --data-dir option to enable")
+
+        # Verify data dir exists
+        if not path.exists(dataset_dir):
+            print("Dataset path %s doesn't exist"%(dataset_dir), file=sys.stderr)
+            assert path.exists(dataset_dir)
+
+        # Append required paths to PYTHONPATH
+        test_env = os.environ.copy()
+        if 'PYTHONPATH' not in test_env:
+            test_env['PYTHONPATH'] = ""
+
+        # Add project root and torchvision to the path (assuming running in nvcr.io/nvidia/pytorch:20.08-py3)
+        test_env['PYTHONPATH'] += ":/opt/pytorch/vision/references/classification/:%s/../"%(dir_path)
+
+        # Add requirement egg files manually to path since we're spawning a new process (downloaded by setuptools)
+        for egg in glob.glob(dir_path + "/../.eggs/*.egg"):
+            test_env['PYTHONPATH'] += ":%s"%(egg)
+
+        # Run in a subprocess to avoid contaminating the module state for other test cases
+        ret = subprocess.run(
+            [
+                'python3', dir_path + '/../examples/torchvision/classification_flow.py',
+                '--data-dir', dataset_dir,
+                '--model', 'resnet18', '--pretrained',
+                '-t', '0.5',
+                '--num-finetune-epochs', '2',
+                '--evaluate-onnx',
+            ],
+            env=test_env,
+            check=False, stdout=subprocess.PIPE)
+
+        # If the test failed dump the output to stderr for better logging
+        if ret.returncode != 0:
+            print(ret.stdout, file=sys.stderr)
+
+        assert ret.returncode == 0

tools/pytorch-quantization/tests/clip_test.py

@@ -0,0 +1,69 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+"""tests of Clip module."""
+
+import pytest
+import numpy as np
+
+import torch
+
+from pytorch_quantization.nn.modules import clip
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+np.random.seed(1234)
+torch.manual_seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+
+
+class TestClip():
+
+    def test_simple_run(self):
+        x_np = np.random.rand(1023).astype(np.float32)
+        x_torch = torch.Tensor(x_np)
+        clip_op = clip.Clip(torch.tensor(0.3), torch.tensor(0.7))
+        clip_x_np = np.clip(x_np, 0.3, 0.7)
+        clip_x_torch = clip_op(x_torch)
+        np.testing.assert_array_equal(clip_x_torch.cpu().numpy(), clip_x_np)
+
+    def test_raise(self):
+        with pytest.raises(ValueError, match="must be scalar"):
+            clip_op = clip.Clip(torch.tensor(0.3), torch.tensor(0.7), learn_min=True)
+
+    def test_backward(self):
+        x = torch.randn(3, 7, requires_grad=True)
+        x.retain_grad()
+
+        min_value = 0.3
+        max_value = 0.7
+        clip_op = clip.Clip(min_value, max_value, learn_min=True, learn_max=True)
+        clip_x = clip_op(x)
+        clip_x.retain_grad()
+
+        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
+        criterion = torch.nn.CrossEntropyLoss()
+        loss = criterion(clip_x, labels)
+
+        loss.backward()
+
+        assert x.grad.cpu()[x.cpu() < min_value].sum() == 0
+        assert x.grad.cpu()[x.cpu() > max_value].sum() == 0
+        assert torch.equal(clip_x.grad[(x > min_value) & (x < max_value)], x.grad[(x > min_value) & (x < max_value)])

tools/pytorch-quantization/tests/conftest.py

@@ -0,0 +1,25 @@
+#
+# 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.
+#
+
+
+"""local configuration for pytests"""
+
+import pytest
+
+def pytest_addoption(parser):
+    parser.addoption('--data-dir', type=str, dest="data_dir",
+                     default='', help="set dataset dir for tests")

tools/pytorch-quantization/tests/fixtures/__init__.py

@@ -0,0 +1,22 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+import pytest
+@pytest.fixture
+def verbose(request):
+    return request.config.getoption("verbose")

tools/pytorch-quantization/tests/fixtures/models.py

@@ -0,0 +1,69 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+"""Model used for tests"""
+
+import pytest
+import torch.nn as nn
+import torch.nn.functional as F
+from pytorch_quantization.nn import QuantConv2d, QuantLinear
+from pytorch_quantization.tensor_quant import QuantDescriptor
+
+class LeNet(nn.Module):
+    def __init__(self, **kwargs):
+        super(LeNet, self).__init__()
+        self.conv1 = nn.Conv2d(1, 10, kernel_size=5, **kwargs)
+        self.conv2 = nn.Conv2d(10, 20, kernel_size=5, **kwargs)
+        self.fc1 = nn.Linear(320, 50, **kwargs)
+        self.fc2 = nn.Linear(50, 10, **kwargs)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2(x), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+
+class QuantLeNet(nn.Module):
+    def __init__(self, **kwargs):
+        super(QuantLeNet, self).__init__()
+        self.conv1 = QuantConv2d(1, 10, kernel_size=5, **kwargs)
+        self.conv2 = QuantConv2d(10, 20, kernel_size=5, **kwargs)
+        self.fc1 = QuantLinear(320, 50, **kwargs)
+        self.fc2 = QuantLinear(50, 10, **kwargs)
+
+    def forward(self, x):
+        x = F.relu(F.max_pool2d(self.conv1(x), 2))
+        x = F.relu(F.max_pool2d(self.conv2(x), 2))
+        x = x.view(-1, 320)
+        x = F.relu(self.fc1(x))
+        x = F.dropout(x, training=self.training)
+        x = self.fc2(x)
+        return F.log_softmax(x, dim=1)
+
+@pytest.fixture
+def resnet18():
+    import torchvision
+    return torchvision.models.resnet18()
+
+@pytest.fixture
+def quant_lenet():
+    return QuantLeNet(quant_desc_input=QuantDescriptor(), quant_desc_weight=QuantDescriptor())

tools/pytorch-quantization/tests/functional_test.py

@@ -0,0 +1,94 @@
+#
+# 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.
+#
+
+
+"""tests of supportive functions"""
+
+import pytest
+import numpy as np
+
+import torch
+
+import pytorch_quantization.nn.functional as QF
+
+np.random.seed(1234)
+torch.manual_seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+
+class TestClip():
+
+    def test_simple_run(self):
+        x_np = np.random.rand(1023).astype(np.float32)
+        x_torch = torch.Tensor(x_np)
+        clip_x_np = np.clip(x_np, 0.3, 0.7)
+        clip_x_torch = QF.clip(x_torch, torch.tensor(0.3), torch.tensor(0.7))
+        np.testing.assert_array_equal(clip_x_torch.cpu().numpy(), clip_x_np)
+
+    def test_raise(self):
+        x = torch.randn(3, 7, requires_grad=True)
+
+        min_value = torch.Tensor(3, 7)
+        max_value = torch.Tensor(3, 7)
+        min_value.requires_grad = True
+        max_value.requires_grad = True
+        clip_x = QF.clip(x, min_value, max_value)
+
+        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
+        criterion = torch.nn.CrossEntropyLoss()
+        loss = criterion(clip_x, labels)
+        with pytest.raises(ValueError, match="can only be scalar"):
+            loss.backward()
+
+    def test_broadcast(self):
+        """Test broadcast behavior by randomly picked shuffling of np.random.rand"""
+        x_np = np.random.rand(1023, 4, 5, 6).astype(np.float32) - 0.5
+        x_torch = torch.Tensor(x_np)
+        min_value = np.random.rand(1, 4, 1, 1).astype(np.float32) * 0.1 - 0.2
+        max_value = np.random.rand(1, 4, 1, 1).astype(np.float32) * 10 + 0.5
+        clip_x_np = np.clip(x_np, min_value, max_value)
+        clip_x_torch = QF.clip(x_torch, torch.tensor(min_value), torch.tensor(max_value))
+        np.testing.assert_array_equal(clip_x_torch.cpu().numpy(), clip_x_np)
+
+    def test_backward(self):
+        x = torch.randn(3, 1025, requires_grad=True)
+        x.retain_grad()
+
+        min_value = torch.tensor(0.3)
+        max_value = torch.tensor(0.7)
+        min_value.requires_grad = True
+        max_value.requires_grad = True
+        min_value.retain_grad()
+        max_value.retain_grad()
+        clip_x = QF.clip(x, min_value, max_value)
+        clip_x.retain_grad()
+
+        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
+        criterion = torch.nn.CrossEntropyLoss()
+        loss = criterion(clip_x, labels)
+        loss.backward()
+
+        np.testing.assert_array_almost_equal(
+            clip_x.grad[x < min_value].sum().cpu().numpy(), min_value.grad.cpu().numpy(), decimal=6)
+        np.testing.assert_array_almost_equal(
+            clip_x.grad[x > max_value].sum().cpu().numpy(), max_value.grad.cpu().numpy(), decimal=6)
+        assert x.grad.cpu()[x.cpu() < min_value.cpu()].sum() == 0
+        assert x.grad.cpu()[x.cpu() > max_value.cpu()].sum() == 0
+        assert torch.equal(clip_x.grad[(x > min_value) & (x < max_value)], x.grad[(x > min_value) & (x < max_value)])

tools/pytorch-quantization/tests/integration_test.py

@@ -0,0 +1,216 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+
+"""tests of integrating Quant layers into a network"""
+
+import pytest
+import io
+
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+import torch.optim as optim
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization import quant_modules
+from pytorch_quantization import nn as quant_nn
+from pytorch_quantization.tensor_quant import QuantDescriptor
+from pytorch_quantization.nn.modules import tensor_quantizer
+from tests.fixtures.models import LeNet, QuantLeNet
+from tests.fixtures import verbose
+
+np.random.seed(12345)  # seed 1234 causes 1 number mismatch at 6th decimal in one of the tests
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestNetwork():
+    """test basic operations of quantized network"""
+
+    def test_simple_build(self):
+        """test instantiation"""
+        quant_model = QuantLeNet(quant_desc_input=QuantDescriptor(), quant_desc_weight=QuantDescriptor())
+        for name, module in quant_model.named_modules():
+            if "quantizer" in name:
+                module.disable()
+
+        input_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        weight_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        quant_model = QuantLeNet(quant_desc_input=input_desc, quant_desc_weight=weight_desc)
+
+        input_desc = QuantDescriptor(amax=6.)
+        weight_desc = QuantDescriptor(amax=1.)
+        quant_model = QuantLeNet(quant_desc_input=input_desc, quant_desc_weight=weight_desc)
+
+
+    def test_forward(self):
+        """test forward pass with random data"""
+        input_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        weight_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        quant_model = QuantLeNet(quant_desc_input=input_desc, quant_desc_weight=weight_desc)
+        output = quant_model(torch.empty(16, 1, 28, 28))
+
+    def test_backward(self):
+        """test one iteration with random data and labels"""
+        input_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        weight_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        quant_model = QuantLeNet(quant_desc_input=input_desc, quant_desc_weight=weight_desc)
+        optimizer = optim.SGD(quant_model.parameters(), lr=0.01)
+        optimizer.zero_grad()
+        output = quant_model(torch.empty(16, 1, 28, 28))
+        loss = F.nll_loss(output, torch.randint(10, (16,), dtype=torch.int64))
+        loss.backward()
+        optimizer.step()
+
+    def test_native_amp_fp16(self):
+        """test one iteration with random data and labels"""
+        input_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        weight_desc = tensor_quant.QUANT_DESC_8BIT_PER_TENSOR
+        model = QuantLeNet(quant_desc_input=input_desc, quant_desc_weight=weight_desc)
+        optimizer = optim.SGD(model.parameters(), lr=0.01)
+        optimizer.zero_grad()
+        with torch.cuda.amp.autocast():
+            output = model(torch.empty(16, 1, 28, 28))
+            loss = F.nll_loss(output, torch.randint(10, (16,), dtype=torch.int64))
+        loss.backward()
+        optimizer.step()
+        assert loss.dtype == torch.float32
+
+    def test_asp(self):
+        """test Sparsity (ASP) and QAT toolkits together"""
+        try:
+            from apex.contrib.sparsity import ASP
+        except ImportError:
+            pytest.skip("ASP is not available.")
+
+        quant_modules.initialize()
+        model = LeNet()
+        quant_modules.deactivate()
+
+        optimizer = optim.SGD(model.parameters(), lr=0.01)
+
+        ASP.init_model_for_pruning(
+            model,
+            mask_calculator="m4n2_1d",
+            verbosity=2,
+            whitelist=[torch.nn.Linear, torch.nn.Conv2d, torch.nn.Conv3d, quant_nn.modules.quant_linear.QuantLinear],
+            allow_recompute_mask=False,
+            custom_layer_dict={
+                quant_nn.QuantConv1d: ['weight'],
+                quant_nn.QuantConv2d: ['weight'],
+                quant_nn.QuantConv3d: ['weight'],
+                quant_nn.QuantConvTranspose1d: ['weight'],
+                quant_nn.QuantConvTranspose2d: ['weight'],
+                quant_nn.QuantConvTranspose3d: ['weight'],
+                quant_nn.QuantLinear: ['weight']
+            },
+            allow_permutation=False)
+        ASP.init_optimizer_for_pruning(optimizer)
+        ASP.compute_sparse_masks()
+
+        model = model.to('cuda')
+        output = model(torch.empty(16, 1, 28, 28).to('cuda'))
+        optimizer.zero_grad()
+        loss = F.nll_loss(output, torch.randint(10, (16,), dtype=torch.int64))
+        loss.backward()
+        optimizer.step()
+
+    def test_quant_module_replacement(self):
+        """test monkey patching of modules with their quantized versions"""
+        lenet = LeNet()
+        qlenet = QuantLeNet()
+
+        mod_list = [type(mod) for name, mod in lenet.named_modules()]
+        mod_list = mod_list[1:]    
+        qmod_list = [type(mod) for name, mod in qlenet.named_modules()]
+        qmod_list = qmod_list[1:]  
+
+        # Before any monkey patching, the networks should be different
+        assert(mod_list != qmod_list)
+
+        # Monkey patch the modules
+        no_replace_list = ["Linear"]
+        custom_quant_modules = [(torch.nn, "Linear", quant_nn.QuantLinear)]
+
+        quant_modules.initialize(no_replace_list, custom_quant_modules)
+
+        lenet = LeNet()
+        qlenet = QuantLeNet()
+    
+        mod_list = [type(mod) for name, mod in lenet.named_modules()]
+        mod_list = mod_list[1:]    
+        qmod_list = [type(mod) for name, mod in qlenet.named_modules()]
+        qmod_list = qmod_list[1:]
+
+        # After monkey patching, the networks should be same
+        assert(mod_list == qmod_list)
+
+        # Reverse monkey patching
+        quant_modules.deactivate()
+
+        lenet = LeNet()
+        qlenet = QuantLeNet()
+    
+        mod_list = [type(mod) for name, mod in lenet.named_modules()]
+        mod_list = mod_list[1:]    
+        qmod_list = [type(mod) for name, mod in qlenet.named_modules()]
+        qmod_list = qmod_list[1:]
+
+        # After reversing monkey patching, the networks should again be different
+        assert(mod_list != qmod_list)
+
+    def test_calibration(self):
+        quant_model = QuantLeNet(quant_desc_input=QuantDescriptor(), quant_desc_weight=QuantDescriptor()).cuda()
+
+        for name, module in quant_model.named_modules():
+            if name.endswith("_quantizer"):
+                if module._calibrator is not None:
+                    module.disable_quant()
+                    module.enable_calib()
+                else:
+                    module.disable()
+                print(F"{name:40}: {module}")
+
+        quant_model(torch.rand(16, 1, 224, 224, device="cuda"))
+
+        # Load calib result and disable calibration
+        for name, module in quant_model.named_modules():
+            if name.endswith("_quantizer"):
+                if module._calibrator is not None:
+                    module.load_calib_amax()
+                    module.enable_quant()
+                    module.disable_calib()
+                else:
+                    module.enable()
+        quant_model.cuda()
+
+    def test_state_load(self):
+        quant_desc = tensor_quant.QuantDescriptor(axis=1, num_bits=8, amax=127.0)
+        quantizer = tensor_quantizer.TensorQuantizer(quant_desc).cuda()
+        quantizer2 = tensor_quantizer.TensorQuantizer(quant_desc).cuda()
+        quantizer2.pre_quant_scale = torch.Tensor([[1.0, 2.0, 3.0, 4.0]]).cuda()
+        buffer = io.BytesIO()
+        torch.save(quantizer2.state_dict(), buffer)
+
+        buffer.seek(0)
+        quantizer.load_state_dict(torch.load(buffer))
+
+        assert torch.allclose(quantizer.pre_quant_scale, quantizer2.pre_quant_scale)

tools/pytorch-quantization/tests/license_test.py

@@ -0,0 +1,86 @@
+#
+# 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.
+#
+
+
+"""test the license of source files."""
+
+import pytest
+from pathlib import Path
+from filecmp import cmp
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestLicense():
+
+    def test_license(self):
+        root = Path(__file__).parent.parent.absolute()
+        root_len = len(str(root))
+
+        # Collect files ending with relevant extensions
+        file_list = []
+        file_types = ['*.py', '*.cpp', '*.cu', '*.h', '*.hpp', '*.c', '*.sh']
+        for ft in file_types:
+            file_list += list(root.rglob(ft))
+
+        # Trim files from build folders
+        build_folders = ['build', 'dist', '.eggs', '.vscode']
+        build_files = []
+        for src_file in file_list:
+            local_path = str(src_file.parents[0])[root_len : ]
+            for folder in build_folders:
+                if folder in local_path:
+                    build_files.append(src_file)
+
+        for bf in build_files:
+            file_list.remove(bf)
+
+        print (f"Found {len(file_list)} source files")
+
+        cpp_header = (root / 'tests' / 'license_test_header_cpp.txt').open().readlines()
+        py_header = (root / 'tests' / 'license_test_header_py.txt').open().readlines()
+        sh_header = (root / 'tests' / 'license_test_header_sh.txt').open().readlines()
+
+        invalid_files = []
+        for f in file_list:
+            with open(f) as src_file:
+                src_lines = src_file.readlines()
+
+            # Skip empty files
+            if len(src_lines) == 0:
+                continue
+
+            if f.suffix == '.py':
+                header = py_header
+            elif f.suffix == '.sh':
+                header = sh_header
+            else:
+                header = cpp_header
+
+            num_lines = len(header)
+            if len(src_lines) < num_lines:
+                invalid_files.append(f)
+                continue
+
+            for i in range(num_lines):
+                if src_lines[i] != header[i]:
+                    invalid_files.append(f)
+                    break
+
+        if len(invalid_files) > 0:
+            for f in invalid_files:
+                print(f"The file {f} has an invalid header!")
+            raise AssertionError("%d files have invalid headers!" % (len(invalid_files)))

tools/pytorch-quantization/tests/license_test_header_cpp.txt

@@ -0,0 +1,16 @@
+/*
+ * 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.
+ */

tools/pytorch-quantization/tests/license_test_header_py.txt

@@ -0,0 +1,16 @@
+#
+# 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.
+#

tools/pytorch-quantization/tests/license_test_header_sh.txt

@@ -0,0 +1,17 @@
+#!/bin/bash
+#
+# 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.
+#

tools/pytorch-quantization/tests/model_test.py

@@ -0,0 +1,92 @@
+#
+# 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.
+#
+"""Tests of calibrators"""
+import inspect
+import pytest
+import numpy as np
+
+import torch
+
+from pytorch_quantization import enable_onnx_export
+from pytorch_quantization import utils as quant_utils
+from pytorch_quantization import calib
+from pytorch_quantization import nn as quant_nn
+import tests.utils as test_utils
+from examples.torchvision.models.classification import *
+from tests.fixtures import verbose
+from tests.fixtures.models import QuantLeNet
+
+np.random.seed(12345)
+torch.manual_seed(12345)
+
+# pylint:disable=missing-docstring, no-self-use
+
+
+class TestExampleModels():
+
+    def test_resnet50(self):
+        model = resnet50(pretrained=True, quantize=True)
+        model.eval()
+
+        for name, module in model.named_modules():
+            if name.endswith('_quantizer'):
+                module.amax = 2.50
+
+        model.cuda()
+        dummy_input = torch.randn(1, 3, 224, 224, device='cuda')
+        with enable_onnx_export():
+            if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters:
+                torch.onnx.export(model,
+                                  dummy_input,
+                                  "/tmp/resnet50.onnx",
+                                  verbose=False,
+                                  opset_version=13,
+                                  enable_onnx_checker=False,
+                                  do_constant_folding=True)
+            else:
+                torch.onnx.export(model,
+                                  dummy_input,
+                                  "/tmp/resnet50.onnx",
+                                  verbose=False,
+                                  opset_version=13,
+                                  do_constant_folding=True)
+
+    def test_resnet50_cpu(self):
+        model = resnet50(pretrained=True, quantize=True)
+        model.eval()
+
+        for name, module in model.named_modules():
+            if name.endswith('_quantizer'):
+                module.amax = 2.50
+
+        dummy_input = torch.randn(1, 3, 224, 224)
+        with enable_onnx_export():
+            if "enable_onnx_checker" in inspect.signature(torch.onnx.export).parameters:
+                torch.onnx.export(model,
+                                  dummy_input,
+                                  "/tmp/resnet50_cpu.onnx",
+                                  verbose=False,
+                                  opset_version=13,
+                                  enable_onnx_checker=False,
+                                  do_constant_folding=True)
+            else:
+                torch.onnx.export(model,
+                                  dummy_input,
+                                  "/tmp/resnet50.onnx",
+                                  verbose=False,
+                                  opset_version=13,
+                                  do_constant_folding=True)

tools/pytorch-quantization/tests/optim_helper_test.py

@@ -0,0 +1,111 @@
+#
+# 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.
+#
+
+
+"""Tests of helper functions for quant optimizer"""
+
+import numpy as np
+import pytest
+
+
+import torch.optim as optim
+
+from pytorch_quantization.optim import helper
+from pytorch_quantization.tensor_quant import QuantDescriptor
+from .fixtures.models import QuantLeNet
+from .fixtures.models import resnet18
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestMatchParameters():
+
+    def test_single_key(self, resnet18):
+        param = helper.match_parameters(resnet18, ['downsample.0.weight'])
+        assert len(list(param)) == 3
+
+    def test_multi_keys(self, resnet18):
+        param = list(helper.match_parameters(resnet18, ['conv1', 'downsample']))
+        assert len(param) == 18
+
+    def test_regex(self, resnet18):
+        param = helper.match_parameters(resnet18, ['downsample.*.weight$'])
+        assert len(list(param)) == 6
+
+        param = helper.match_parameters(resnet18, ['downsample.*.wei$'])
+        assert not list(param)
+
+
+class TestGroupParameters():
+
+    def test_single_key(self, resnet18):
+        param_groups = helper.group_parameters(resnet18, [['downsample.1.weight']])
+        assert len(list(param_groups[0]['params'])) == 3
+
+    def test_lr_momentum_decay(self, resnet18):
+        lrs = [0.01, 0.001]
+        momentums = [0.02, 0.002]
+        weight_decays = [0.03, 0.003]
+        param_groups = helper.group_parameters(
+            resnet18, [['conv1.*weight'], ['downsample.*.weight']], lrs, momentums, weight_decays)
+
+        assert param_groups[0]['lr'] == lrs[0]
+        assert param_groups[1]['lr'] == lrs[1]
+        assert param_groups[0]['momentum'] == momentums[0]
+        assert param_groups[1]['momentum'] == momentums[1]
+        assert param_groups[0]['weight_decay'] == weight_decays[0]
+        assert param_groups[1]['weight_decay'] == weight_decays[1]
+
+    def test_optimizer_feed(self, resnet18):
+        """Feed grouped parameters to optimizer, see what happens"""
+        lrs = [0.01, 0.001]
+        momentums = [0.02, 0.002]
+        weight_decays = [0.03, 0.003]
+        param_groups = helper.group_parameters(
+            resnet18, [['conv1.*weight'], ['downsample.*.weight']], lrs, momentums, weight_decays)
+        optimizer = optim.SGD(param_groups)
+        optimizer.step()
+
+    def test_raises(self):
+        with pytest.raises(TypeError, match="must be list of list of patterns"):
+            helper.group_parameters(None, [['downsample.1.weight'], 'conv1'])
+
+        with pytest.raises(TypeError, match="must match"):
+            helper.group_parameters(None, [['downsample.1.weight'], ['conv1']], lrs=[0.1])
+
+        with pytest.raises(TypeError, match="must match"):
+            helper.group_parameters(None, [['downsample.1.weight'], ['conv1']], momentums=[0.1])
+
+        with pytest.raises(TypeError, match="must match"):
+            helper.group_parameters(None, [['downsample.1.weight'], ['conv1']], weight_decays=[0.1])
+
+
+class TestFreezeParameters():
+
+    def test_simple(self, resnet18):
+        helper.freeze_parameters(resnet18, ['downsample.0.weight'])
+        for name, param in resnet18.named_parameters():
+            if 'downsample.0.weight' in name:
+                assert not param.requires_grad
+
+class TestQuantWeightInPlace():
+
+    def test_simple(self):
+        quant_lenet = QuantLeNet(
+            quant_desc_input=QuantDescriptor(),
+            quant_desc_weight=QuantDescriptor())
+        quant_lenet.eval()
+        helper.quant_weight_inplace(quant_lenet)

tools/pytorch-quantization/tests/print_test.py

@@ -0,0 +1,60 @@
+#
+# 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.
+#
+
+
+"""test for str and repr
+Make sure things can print and in a nice form. Put all the print tests together so that running this test file alone
+can inspect all the print messages in the project
+"""
+
+import torch
+from torch import nn
+
+from pytorch_quantization import calib
+from pytorch_quantization import tensor_quant
+from pytorch_quantization import nn as quant_nn
+from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestPrint():
+
+    def test_print_descriptor(self):
+        test_desc = tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL
+        print(test_desc)
+
+    def test_print_tensor_quantizer(self):
+        test_quantizer = TensorQuantizer()
+        print(test_quantizer)
+
+    def test_print_module(self):
+        class _TestModule(nn.Module):
+            def __init__(self):
+                super(_TestModule, self).__init__()
+                self.conv = nn.Conv2d(33, 65, 3)
+                self.quant_conv = quant_nn.Conv2d(33, 65, 3)
+                self.linear = nn.Linear(33, 65)
+                self.quant_linear = quant_nn.Linear(33, 65)
+
+        test_module = _TestModule()
+        print(test_module)
+
+    def test_print_calibrator(self):
+        print(calib.MaxCalibrator(7, 1, False))
+        hist_calibrator = calib.HistogramCalibrator(8, None, True)
+        hist_calibrator.collect(torch.rand(10))
+        print(hist_calibrator)

tools/pytorch-quantization/tests/quant_conv_test.py

@@ -0,0 +1,557 @@
+#
+# 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.
+#
+
+
+"""tests of QuantConv module.
+Mose tests check the functionality of all the combinations in Quant conv against the corresponding functionalities in
+tensor_quant. There are tests for all the three QuantConv1D, QuantConv2D, and QuantConv3D
+"""
+import pytest
+import numpy as np
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization.tensor_quant import QuantDescriptor
+from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer
+from pytorch_quantization import utils as quant_utils
+from pytorch_quantization.nn.modules import quant_conv
+import tests.utils as test_utils
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+torch.backends.cudnn.deterministic = True
+
+np.random.seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+
+_NUM_IN_CHANNELS = 13
+_NUM_OUT_CHANNELS = 17
+
+class TestQuantConv2D():
+    #Quantizing weight
+
+    def test_no_quant(self):
+
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.input_quantizer.disable()
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_copy
+
+        out1 = F.conv2d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_tensor(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor())
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv2d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL)
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(
+            weight_copy,
+            torch.max(torch.abs(weight_copy).view(_NUM_OUT_CHANNELS, -1), dim=1, keepdim=True)[0].view(
+                _NUM_OUT_CHANNELS, 1, 1, 1))
+
+        out1 = F.conv2d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_in_feature_fake_quant(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256, 256)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.conv2d(quant_input, quant_conv_object.weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(
+            _NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=False, quant_desc_weight=QuantDescriptor())
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv2d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=False,
+                                                   quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(
+            weight_copy,
+            torch.max(torch.abs(weight_copy).view(_NUM_OUT_CHANNELS, -1), dim=1, keepdim=True)[0].view(
+                _NUM_OUT_CHANNELS, 1, 1, 1))
+
+        out1 = F.conv2d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_other_prec(self):
+        kernel_size = 3
+
+        quant_desc_input = QuantDescriptor(num_bits=4)
+        quant_desc_weight = QuantDescriptor(num_bits=3)
+
+        quant_conv_object = quant_conv.QuantConv2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+
+        quant_input = test_input_quantizer(test_input)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.conv2d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_bias(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv2d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=True,
+                                                   quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONV2D_WEIGHT_PER_CHANNEL)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(
+            weight_copy,
+            torch.max(torch.abs(weight_copy).view(_NUM_OUT_CHANNELS, -1), dim=1, keepdim=True)[0].view(
+                _NUM_OUT_CHANNELS, 1, 1, 1))
+
+        out1 = F.conv2d(quant_input, quant_weight, bias=quant_conv_object.bias)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_against_unquantized(self):
+        kernel_size = 3
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 24, 24).cuda()
+
+        torch.manual_seed(12345)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(12345)
+        fake_quant_conv2d = quant_conv.QuantConv2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_input=QuantDescriptor(num_bits=16),
+            quant_desc_weight=QuantDescriptor(num_bits=16, axis=(0)))
+
+        # Reset seed. Make sure weight and bias are the same
+        torch.manual_seed(12345)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(12345)
+        conv2d = nn.Conv2d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=True)
+
+        fake_quant_output = fake_quant_conv2d(test_input)
+        output = conv2d(test_input)
+
+        test_utils.compare(fake_quant_output, output, rtol=1e-6, atol=1.5e-4)
+
+
+    def test_set_default_quant_desc(self):
+        quant_conv_layer = quant_conv.Conv2d(32, 257, 3)
+        assert quant_conv_layer.input_quantizer._axis == None
+        assert quant_conv_layer.weight_quantizer._axis == (0)
+
+        # set default to a different one
+        quant_desc_input = QuantDescriptor(num_bits=11)
+        quant_desc_weight = QuantDescriptor(num_bits=13, axis=(1))
+        quant_conv.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
+        quant_conv.QuantConv2d.set_default_quant_desc_weight(quant_desc_weight)
+
+        # Create one with default descriptor
+        quant_conv_layer = quant_conv.Conv2d(32, 257, 3)
+        # Check quant_desc in quantizer created with default descriptor
+        assert quant_conv_layer.input_quantizer._num_bits == quant_desc_input.num_bits
+        assert quant_conv_layer.weight_quantizer._axis == quant_desc_weight.axis
+
+        # Test default is per class
+        quant_conv_layer = quant_conv.Conv3d(31, 255, 5)
+        assert quant_conv_layer.input_quantizer._num_bits != quant_desc_input.num_bits
+        assert quant_conv_layer.weight_quantizer._axis != quant_desc_weight.axis
+
+        # Reset default
+        quant_conv.QuantConv2d.set_default_quant_desc_input(QuantDescriptor())
+        quant_conv.QuantConv2d.set_default_quant_desc_weight(QuantDescriptor(axis=(0)))
+
+    def test_unused_kwargs(self):
+        with pytest.raises(TypeError, match="Unused keys"):
+            quant_conv.Conv2d(32, 257, 3, descriptor='oops')
+
+class TestQuantConv1D():
+
+    def test_no_quant(self):
+        kernel_size = 8
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.input_quantizer.disable()
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_copy
+
+        out1 = F.conv1d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_tensor(self):
+        kernel_size = 8
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor())
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv1d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor(axis=(0)))
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=(1, 2))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv1d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_input(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(20, _NUM_IN_CHANNELS, 50)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.conv1d(quant_input, quant_conv_object.weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=False, quant_desc_weight=QuantDescriptor())
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv1d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor(axis=(0)))
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(
+            weight_copy,
+            torch.max(torch.abs(weight_copy).view(_NUM_OUT_CHANNELS, -1), dim=1, keepdim=True)[0].view(
+                _NUM_OUT_CHANNELS, 1, 1))
+
+        out1 = F.conv1d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_other_prec(self):
+        kernel_size = 3
+
+        quant_desc_input = QuantDescriptor(num_bits=4)
+        quant_desc_weight = QuantDescriptor(num_bits=3, axis=(0))
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+
+        quant_input = test_input_quantizer(test_input)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.conv1d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_bias(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_weight=QuantDescriptor(axis=(0)))
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(
+            weight_copy,
+            torch.max(torch.abs(weight_copy).view(_NUM_OUT_CHANNELS, -1), dim=1, keepdim=True)[0].view(
+                _NUM_OUT_CHANNELS, 1, 1))
+
+        out1 = F.conv1d(quant_input, quant_weight, bias=quant_conv_object.bias)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_against_unquantized(self):
+        kernel_size = 3
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 24).cuda()
+
+        torch.manual_seed(12345)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(12345)
+        fake_quant_conv1d = quant_conv.QuantConv1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_input=QuantDescriptor(num_bits=16),
+            quant_desc_weight=QuantDescriptor(num_bits=16, axis=(0)))
+
+        # Reset seed. Make sure weight and bias are the same
+        torch.manual_seed(12345)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(12345)
+        conv1d = nn.Conv1d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=True)
+
+        fake_quant_output = fake_quant_conv1d(test_input)
+        output = conv1d(test_input)
+
+        test_utils.compare(fake_quant_output, output, rtol=1e-5, atol=1e-4)
+
+
+class TestQuantConv3D():
+    #Quantizing weight
+
+    def test_no_quant(self):
+        kernel_size = 8
+
+        quant_conv_object = quant_conv.QuantConv3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.input_quantizer.disable()
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 8, 8, 8)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_copy
+
+        out1 = F.conv3d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_quant_per_channel_other_prec(self):
+        kernel_size = 3
+
+        quant_desc_input = QuantDescriptor(num_bits=4)
+        quant_desc_weight = QuantDescriptor(num_bits=3, axis=(0))
+
+        quant_conv_object = quant_conv.QuantConv3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 8, 8, 8)
+
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+
+        quant_input = test_input_quantizer(test_input)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.conv3d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_quant_per_channel_bias(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConv3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_weight=QuantDescriptor(axis=(0)))
+        test_input = torch.randn(8, _NUM_IN_CHANNELS, 8, 8, 8)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(
+            weight_copy,
+            torch.max(torch.abs(weight_copy).view(_NUM_OUT_CHANNELS, -1), dim=1, keepdim=True)[0].view(
+                _NUM_OUT_CHANNELS, 1, 1, 1, 1))
+
+        out1 = F.conv3d(quant_input, quant_weight, bias=quant_conv_object.bias)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_against_unquantized(self):
+        kernel_size = 3
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 24, 24, 24).cuda()
+
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        fake_quant_conv3d = quant_conv.QuantConv3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_input=QuantDescriptor(num_bits=16),
+            quant_desc_weight=QuantDescriptor(num_bits=16, axis=(0)))
+
+        # Reset seed. Make sure weight and bias are the same
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        conv3d = nn.Conv3d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=True)
+
+        fake_quant_output = fake_quant_conv3d(test_input)
+        output = conv3d(test_input)
+
+        test_utils.compare(fake_quant_output, output, rtol=1e-6, atol=2e-4)

tools/pytorch-quantization/tests/quant_conv_transposed_test.py

@@ -0,0 +1,522 @@
+#
+# 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.
+#
+
+
+"""tests of QuantConv module.
+Test for QuantConvTransposed
+"""
+import pytest
+import numpy as np
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization.tensor_quant import QuantDescriptor
+from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer
+from pytorch_quantization import utils as quant_utils
+from pytorch_quantization.nn.modules import quant_conv
+import tests.utils as test_utils
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+torch.backends.cudnn.deterministic = True
+
+np.random.seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+
+_NUM_IN_CHANNELS = 13
+_NUM_OUT_CHANNELS = 17
+
+
+class TestQuantConvTranspose2D():
+
+    def test_no_quant(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.input_quantizer.disable()
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 32, 32)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_copy
+
+        out1 = F.conv_transpose2d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_tensor(self):
+        kernel_size = 8
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor())
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(256, _NUM_IN_CHANNELS, 32, 32)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv_transpose2d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE2D_WEIGHT_PER_CHANNEL)
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2, 3))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose2d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_input(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(20, _NUM_IN_CHANNELS, 50, 50)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.conv_transpose2d(quant_input, quant_conv_object.weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=False, quant_desc_weight=QuantDescriptor())
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv_transpose2d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor(axis=(1)))
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2, 3))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose2d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_other_prec(self):
+        kernel_size = 3
+
+        quant_desc_input = QuantDescriptor(num_bits=4)
+        quant_desc_weight = QuantDescriptor(num_bits=3, axis=(1))
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16)
+
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+
+        quant_input = test_input_quantizer(test_input)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.conv_transpose2d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_bias(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_weight=QuantDescriptor(axis=(1)))
+        test_input = torch.randn(2, _NUM_IN_CHANNELS, 2, 2)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2, 3))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose2d(quant_input, quant_weight, bias=quant_conv_object.bias)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_against_unquantized(self):
+        kernel_size = 3
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 32, 32).cuda()
+
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        fake_quant_conv2d = quant_conv.QuantConvTranspose2d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_input=QuantDescriptor(num_bits=16),
+            quant_desc_weight=QuantDescriptor(num_bits=16, axis=(1)))
+
+        # Reset seed. Make sure weight and bias are the same
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        conv2d = nn.ConvTranspose2d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=True)
+
+        fake_quant_output = fake_quant_conv2d(test_input)
+        output = conv2d(test_input)
+
+        test_utils.compare(fake_quant_output, output, rtol=1e-5, atol=2e-4)
+
+
+class TestQuantConvTranspose3D():
+
+    def test_no_quant(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.input_quantizer.disable()
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 32, 32, 32)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_copy
+
+        out1 = F.conv_transpose3d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_other_prec(self):
+        kernel_size = 3
+
+        quant_desc_input = QuantDescriptor(num_bits=4)
+        quant_desc_weight = QuantDescriptor(num_bits=3, axis=(1))
+
+        quant_conv_object = quant_conv.QuantConvTranspose3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16, 16, 16)
+
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+
+        quant_input = test_input_quantizer(test_input)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.conv_transpose3d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_bias(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE3D_WEIGHT_PER_CHANNEL)
+        test_input = torch.randn(2, _NUM_IN_CHANNELS, 2, 2, 2)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2, 3, 4))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose3d(quant_input, quant_weight, bias=quant_conv_object.bias)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_against_unquantized(self):
+        kernel_size = 3
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 32, 32, 32).cuda()
+
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        fake_quant_conv3d = quant_conv.QuantConvTranspose3d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_input=QuantDescriptor(num_bits=16),
+            quant_desc_weight=QuantDescriptor(num_bits=16, axis=(1)))
+
+        # Reset seed. Make sure weight and bias are the same
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        conv3d = nn.ConvTranspose3d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=True)
+
+        fake_quant_output = fake_quant_conv3d(test_input)
+        output = conv3d(test_input)
+
+        test_utils.compare(fake_quant_output, output, rtol=1e-5, atol=2e-4)
+
+
+class TestQuantConvTranspose1D():
+
+    def test_no_quant(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.input_quantizer.disable()
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 32)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_copy
+
+        out1 = F.conv_transpose1d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_tensor(self):
+        kernel_size = 8
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor())
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(256, _NUM_IN_CHANNELS, 32)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv_transpose1d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_CONVTRANSPOSE1D_WEIGHT_PER_CHANNEL)
+        quant_conv_object.input_quantizer.disable()
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 256)
+
+        weight_copy = quant_conv_object.weight.clone()
+
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose1d(test_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_input(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False)
+        quant_conv_object.weight_quantizer.disable()
+        test_input = torch.randn(20, _NUM_IN_CHANNELS, 50)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.conv_transpose1d(quant_input, quant_conv_object.weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=False, quant_desc_weight=QuantDescriptor())
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.conv_transpose1d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_weight=QuantDescriptor(axis=(1)))
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose1d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_other_prec(self):
+        kernel_size = 3
+
+        quant_desc_input = QuantDescriptor(num_bits=4)
+        quant_desc_weight = QuantDescriptor(num_bits=3, axis=(1))
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 16)
+
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+
+        quant_input = test_input_quantizer(test_input)
+
+        weight_copy = quant_conv_object.weight.clone()
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.conv_transpose1d(quant_input, quant_weight)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_bias(self):
+        kernel_size = 3
+
+        quant_conv_object = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_weight=QuantDescriptor(axis=(1)))
+        test_input = torch.randn(2, _NUM_IN_CHANNELS, 2)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        weight_copy = quant_conv_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=(0, 2))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.conv_transpose1d(quant_input, quant_weight, bias=quant_conv_object.bias)
+        out2 = quant_conv_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_against_unquantized(self):
+        kernel_size = 3
+        test_input = torch.randn(16, _NUM_IN_CHANNELS, 24).cuda()
+
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        fake_quant_conv1d = quant_conv.QuantConvTranspose1d(
+            _NUM_IN_CHANNELS,
+            _NUM_OUT_CHANNELS,
+            kernel_size,
+            bias=True,
+            quant_desc_input=QuantDescriptor(num_bits=16),
+            quant_desc_weight=QuantDescriptor(num_bits=16, axis=(1)))
+
+        # Reset seed. Make sure weight and bias are the same
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        conv1d = nn.ConvTranspose1d(_NUM_IN_CHANNELS, _NUM_OUT_CHANNELS, kernel_size, bias=True)
+
+        fake_quant_output = fake_quant_conv1d(test_input)
+        output = conv1d(test_input)
+
+        test_utils.compare(fake_quant_output, output, rtol=1e-5, atol=1e-4)

tools/pytorch-quantization/tests/quant_instancenorm_test.py

@@ -0,0 +1,199 @@
+#
+# 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.
+#
+
+
+"""tests of QuantInstanceNorm module.
+Mose tests check the functionality of all the combinations in Quant instancenorm against the corresponding functionalities in
+tensor_quant. There are tests for all the three QuantInstaceNorm1D, QuantInstanceNorm2D, and QuantInstanceNorm3D
+"""
+import pytest
+import numpy as np
+
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization.tensor_quant import QuantDescriptor
+from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer
+from pytorch_quantization import utils as quant_utils
+from pytorch_quantization.nn.modules import quant_instancenorm
+#import tests.utils as test_utils
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+torch.backends.cudnn.deterministic = True
+
+np.random.seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+NUM_CHANNELS = 15
+
+class TestQuantInstanceNorm1D():
+
+    def test_no_quant(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm1d(NUM_CHANNELS, affine=True)
+        quant_instancenorm_object.input_quantizer.disable()
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128)
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(test_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm1d(NUM_CHANNELS, affine=True,
+                                                                           quant_desc_input=QuantDescriptor())
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128)
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(quant_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm1d(NUM_CHANNELS, affine=True,
+                                                                           quant_desc_input=QuantDescriptor(axis=(1)))
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128)
+        quant_input = tensor_quant.fake_tensor_quant(test_input,
+                                                     torch.abs(test_input).max(0, keepdim=True)[0].max(2, keepdim=True)[0])
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(quant_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+
+class TestQuantInstanceNorm2D():
+
+    def test_no_quant(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm2d(NUM_CHANNELS, affine=True)
+        quant_instancenorm_object.input_quantizer.disable()
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128, 128)
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(test_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm2d(NUM_CHANNELS, affine=True,
+                                                                           quant_desc_input=QuantDescriptor())
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128, 128)
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(quant_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm2d(NUM_CHANNELS, affine=True,
+                                                                           quant_desc_input=QuantDescriptor(axis=(1)))
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128, 128)
+        quant_input = tensor_quant.fake_tensor_quant(test_input,
+                                                     torch.abs(test_input).max(0, keepdim=True)[0].max(2, keepdim=True)[0].max(3, keepdim=True)[0])
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(quant_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+
+
+class TestQuantInstanceNorm3D():
+
+    def test_no_quant(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm3d(NUM_CHANNELS, affine=True)
+        quant_instancenorm_object.input_quantizer.disable()
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128, 128, 128)
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(test_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm3d(NUM_CHANNELS, affine=True,
+                                                                           quant_desc_input=QuantDescriptor())
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128, 128, 128)
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(quant_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+
+        quant_instancenorm_object = quant_instancenorm.QuantInstanceNorm3d(NUM_CHANNELS, affine=True,
+                                                                           quant_desc_input=QuantDescriptor(axis=(1)))
+
+        test_input = torch.randn(8, NUM_CHANNELS, 128, 128, 128)
+        quant_input = tensor_quant.fake_tensor_quant(test_input,
+                                                     torch.abs(test_input).max(0, keepdim=True)[0].max(2, keepdim=True)[0]
+                                                          .max(3, keepdim=True)[0].max(4, keepdim=True)[0])
+
+        out1 = quant_instancenorm_object(test_input)
+        out2 = F.instance_norm(quant_input,
+                               quant_instancenorm_object.running_mean,
+                               quant_instancenorm_object.running_var,
+                               quant_instancenorm_object.weight,
+                               quant_instancenorm_object.bias)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())

tools/pytorch-quantization/tests/quant_linear_test.py

@@ -0,0 +1,231 @@
+#
+# 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.
+#
+
+
+"""tests of QuantLinear module.
+Most tests check the functionality of all the combinations in Quant Linear against the corresponding functionalities
+in tensor_quant.
+"""
+import pytest
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization.nn.modules.tensor_quantizer import TensorQuantizer
+from pytorch_quantization import utils as quant_utils
+from pytorch_quantization.nn.modules import quant_linear
+import tests.utils as test_utils
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+np.random.seed(1234)
+torch.manual_seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+
+
+class TestQuantLinear():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_linear_object = quant_linear.QuantLinear(
+                7, 9, bias=False, quant_desc_weight=tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    #Quantizing weight
+    def test_weight_fake_per_tensor(self):
+        with torch.cuda.device(0):
+            size = 256
+            quant_linear_object = quant_linear.QuantLinear(
+                size,
+                size,
+                bias=False,
+                quant_desc_weight=tensor_quant.QuantDescriptor(axis=None))
+            quant_linear_object.input_quantizer.disable()
+            test_input = torch.randn(size, size)
+
+            weight_copy = quant_linear_object.weight.clone()
+            quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+            out1 = F.linear(test_input, quant_weight)
+            out2 = quant_linear_object(test_input)
+            np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_weight_fake_per_channel(self):
+        size_in = 255
+        size_out = 257
+        quant_linear_object = quant_linear.QuantLinear(
+            size_in, size_out, bias=False,
+            quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW)
+        quant_linear_object.input_quantizer.disable()
+        test_input = torch.randn(32, size_in)
+
+        weight_copy = quant_linear_object.weight.clone()
+        amax = quant_utils.reduce_amax(weight_copy, axis=1, keepdims=True)
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, amax)
+
+        out1 = F.linear(test_input, quant_weight)
+        out2 = quant_linear_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    # Quantizing activations
+    def test_test_input_fake_per_tensor(self):
+        size_in = 255
+        size_out = 257
+        quant_linear_object = quant_linear.QuantLinear(
+            size_in, size_out, bias=False)
+        quant_linear_object.weight_quantizer.disable()
+        test_input = torch.randn(32, size_in)
+
+        weight_copy = quant_linear_object.weight.clone()
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.linear(quant_input, weight_copy)
+        out2 = quant_linear_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor(self):
+        """quantize everything, activations will scaled per tensor in ALL cases"""
+        size_in = 255
+        size_out = 257
+        quant_linear_object = quant_linear.QuantLinear(
+            size_in, size_out, bias=False, quant_desc_weight=tensor_quant.QuantDescriptor())
+        test_input = torch.randn(32, size_in)
+
+        weight_copy = quant_linear_object.weight.clone()
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.linear(quant_input, quant_weight)
+        out2 = quant_linear_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_tensor_with_bias(self):
+        """quantize everything, activations will scaled per tensor in ALL cases"""
+        size_in = 255
+        size_out = 257
+        quant_linear_object = quant_linear.QuantLinear(
+            size_in, size_out, bias=False, quant_desc_weight=tensor_quant.QuantDescriptor())
+        test_input = torch.randn(32, 17, 93, size_in)  # Test input other than 2 dimensional
+
+        weight_copy = quant_linear_object.weight.clone()
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy, torch.max(torch.abs(weight_copy)))
+
+        out1 = F.linear(quant_input, quant_weight, bias=quant_linear_object.bias)
+        out2 = quant_linear_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel(self):
+        """quantize everything, activations will scaled per tensor in ALL cases"""
+        size_in = 255
+        size_out = 257
+        quant_linear_object = quant_linear.QuantLinear(size_in, size_out, bias=False,
+                                                       quant_desc_weight=tensor_quant.QUANT_DESC_8BIT_LINEAR_WEIGHT_PER_ROW)
+        test_input = torch.randn(32, size_in)
+
+        weight_copy = quant_linear_object.weight.clone()
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+        quant_weight = tensor_quant.fake_tensor_quant(weight_copy,
+                                                      torch.max(torch.abs(weight_copy), dim=1, keepdim=True)[0])
+
+        out1 = F.linear(quant_input, quant_weight)
+        out2 = quant_linear_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_per_channel_other_precs(self):
+        """Test some precisions other than 8bit."""
+        size_in = 255
+        size_out = 257
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=4)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=3)
+        quant_linear_object = quant_linear.QuantLinear(
+            size_in,
+            size_out,
+            bias=False,
+            quant_desc_input=quant_desc_input,
+            quant_desc_weight=quant_desc_weight)
+        weight_quantizer = TensorQuantizer(quant_desc_weight)
+        test_input_quantizer = TensorQuantizer(quant_desc_input)
+
+        test_input = torch.randn(32, size_in)
+
+        weight_copy = quant_linear_object.weight.clone()
+        quant_input = test_input_quantizer(test_input)
+        quant_weight = weight_quantizer(weight_copy)
+
+        out1 = F.linear(quant_input, quant_weight)
+        out2 = quant_linear_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_fake_quant_against_unquantized(self):
+        """
+        Quantized Linear should introduce bounded error compare to Linear
+        """
+        size_in = 255
+        size_out = 257
+        test_input = torch.randn(32, size_in).cuda()
+
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        quant_linear_layer = quant_linear.QuantLinear(
+            size_in,
+            size_out,
+            bias=True,
+            quant_desc_input=tensor_quant.QuantDescriptor(num_bits=16),
+            quant_desc_weight=tensor_quant.QuantDescriptor(num_bits=16, axis=0))
+
+        # Reset seed. Make sure weight and bias are the same
+        torch.manual_seed(1234)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(1234)
+        linear_layer = nn.Linear(size_in, size_out, bias=True)
+
+        quant_out_features = quant_linear_layer(test_input)
+        out_features = linear_layer(test_input)
+
+        # The difference between Linear and QuantLinear should be bounded in a range
+        # Small values which become 0 after quantization lead to large relative errors. rtol and atol could be
+        # much smaller without those values
+        np.testing.assert_allclose(
+            quant_out_features.detach().cpu().numpy(), out_features.detach().cpu().numpy(), rtol=0.01, atol=1e-4)
+
+    def test_set_default_quant_desc(self):
+        quant_linear_layer = quant_linear.QuantLinear(32, 257)
+        assert quant_linear_layer.input_quantizer.axis == None
+        assert quant_linear_layer.weight_quantizer.axis == (0)
+
+        # set default to a different one
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=11)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=13, axis=1)
+        quant_linear.Linear.set_default_quant_desc_input(quant_desc_input)
+        quant_linear.Linear.set_default_quant_desc_weight(quant_desc_weight)
+
+        # Create one with default descriptor
+        quant_linear_layer = quant_linear.QuantLinear(32, 257)
+        # Check quant_desc in quantizer created with default descriptor
+        assert quant_linear_layer.input_quantizer.num_bits == quant_desc_input.num_bits
+        assert quant_linear_layer.weight_quantizer.axis == quant_desc_weight.axis
+
+    def test_unused_kwargs(self):
+        with pytest.raises(TypeError, match="Unused keys"):
+            quant_linear_layer = quant_linear.QuantLinear(32, 257, descriptor='oops')

tools/pytorch-quantization/tests/quant_modules_test.py

@@ -0,0 +1,83 @@
+#
+# 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.
+#
+
+
+"""Tests of Quant Module Replacement"""
+import pytest
+import numpy as np
+
+import torch
+
+from pytorch_quantization import nn as quant_nn
+from pytorch_quantization import quant_modules
+from pytorch_quantization.quant_modules import QuantModuleReplacementHelper
+import tests.utils as test_utils
+from tests.fixtures import verbose
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestQuantModuleReplace():
+
+    def test_simple_default_args(self):
+        replacement_helper = QuantModuleReplacementHelper()
+        replacement_helper.prepare_state()
+        replacement_helper.apply_quant_modules()
+
+        # Linear module should not be replaced with its quantized version
+        assert(type(quant_nn.QuantLinear(16, 256, 3)) == type(torch.nn.Linear(16, 256, 3)))
+        assert(type(quant_nn.QuantConv2d(16, 256, 3)) == type(torch.nn.Conv2d(16, 256, 3)))
+
+        replacement_helper.restore_float_modules()
+
+    def test_with_no_replace_list(self):
+        no_replace_list = ["Linear"]
+        custom_quant_modules = None
+        replacement_helper = QuantModuleReplacementHelper()
+        replacement_helper.prepare_state(no_replace_list, custom_quant_modules)
+        replacement_helper.apply_quant_modules()
+
+        # Linear module should not be replaced with its quantized version
+        assert(type(quant_nn.QuantLinear(16, 256, 3)) != type(torch.nn.Linear(16, 256, 3)))
+        assert(type(quant_nn.QuantConv2d(16, 256, 3)) == type(torch.nn.Conv2d(16, 256, 3)))
+
+        replacement_helper.restore_float_modules()
+
+    def test_with_custom_quant_modules(self):
+        no_replace_list = ["Linear"]
+        custom_quant_modules = [(torch.nn, "Linear", quant_nn.QuantLinear)]
+        replacement_helper = QuantModuleReplacementHelper()
+        replacement_helper.prepare_state(no_replace_list, custom_quant_modules)
+        replacement_helper.apply_quant_modules()
+
+        # Although no replace list indicates Linear module should not be replaced with its
+        # quantized version, since the custom_quant_modules still contains the Linear module's
+        # mapping, it will replaced.
+        assert(type(quant_nn.QuantLinear(16, 256, 3)) == type(torch.nn.Linear(16, 256, 3)))
+        assert(type(quant_nn.QuantConv2d(16, 256, 3)) == type(torch.nn.Conv2d(16, 256, 3)))
+
+        replacement_helper.restore_float_modules()
+
+    def test_initialize_deactivate(self):
+        no_replace_list = ["Linear"]
+        custom_quant_modules = [(torch.nn, "Linear", quant_nn.QuantLinear)]
+
+        quant_modules.initialize(no_replace_list, custom_quant_modules)
+
+        assert(type(quant_nn.QuantLinear(16, 256, 3)) == type(torch.nn.Linear(16, 256, 3)))
+        assert(type(quant_nn.QuantConv2d(16, 256, 3)) == type(torch.nn.Conv2d(16, 256, 3)))
+
+        quant_modules.deactivate()

tools/pytorch-quantization/tests/quant_pooling_test.py

@@ -0,0 +1,321 @@
+#
+# 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.
+#
+
+
+"""tests of QuantPooling module.
+Most tests check the functionality of all the combinations in Quant Pooling against the corresponding functionalities
+in tensor_quant.
+"""
+import pytest
+import numpy as np
+
+import torch
+import torch.nn.functional as F
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization.nn.modules import quant_pooling
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+
+np.random.seed(1234)
+torch.manual_seed(1234)
+
+# pylint:disable=missing-docstring, no-self-use
+class TestQuantMaxPool1d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantMaxPool1d(kernel_size=3, stride=1,
+                                                                quant_desc_input=
+                                                                tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantMaxPool1d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(1, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.max_pool1d(quant_input, 3, 1, 0, 1, False, False)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantMaxPool2d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1,
+                                                                quant_desc_input=
+                                                                tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.max_pool2d(quant_input, 3, 1, 0, 1, False, False)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_variable_bits(self):
+        # Repeat checking the output for variable number of bits to QuantDescriptor
+        for bits in [2, 4, 6]:
+            quant_desc_input = tensor_quant.QuantDescriptor(num_bits=bits)
+
+            quant_pooling.QuantMaxPool2d.set_default_quant_desc_input(quant_desc_input)
+            quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1)
+
+            test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+            quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)), bits)
+
+            out1 = F.max_pool2d(quant_input, 3, 1, 0, 1, False, False)
+            out2 = quant_pooling_object(test_input)
+            np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_fake_quant_disable(self):
+        quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+        quant_pooling_object.input_quantizer.disable()
+
+        out1 = F.max_pool2d(test_input, 3, 1, 0, 1, False, False)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_multi_axis(self):
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=8, axis=(0, 1))
+
+        quant_pooling.QuantMaxPool2d.set_default_quant_desc_input(quant_desc_input)
+        quant_pooling_object = quant_pooling.QuantMaxPool2d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(16, 7, 5, 5, dtype=torch.double)
+        input_amax = torch.amax(torch.abs(test_input), dim=(2, 3), keepdim=True)
+        quant_input = tensor_quant.fake_tensor_quant(test_input, input_amax)
+
+        out1 = F.max_pool2d(quant_input, 3, 1, 0, 1, False, False)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantMaxPool3d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantMaxPool3d(kernel_size=3, stride=1,
+                                                                quant_desc_input=
+                                                                tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantMaxPool3d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(5, 5, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.max_pool3d(quant_input, 3, 1, 0, 1, False, False)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantAvgPool1d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantAvgPool1d(kernel_size=3, stride=1,
+                                                                quant_desc_input=
+                                                                tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantAvgPool1d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(1, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.avg_pool1d(quant_input, 3, 1, 0, False, True)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantAvgPool2d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantAvgPool2d(kernel_size=3, stride=1,
+                                                                quant_desc_input=
+                                                                tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantAvgPool2d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.avg_pool2d(quant_input, 3, 1, 0, False, True, None)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_variable_bits(self):
+        # Repeat checking the output for variable number of bits to QuantDescriptor
+        for bits in [2, 4, 6]:
+            quant_desc_input = tensor_quant.QuantDescriptor(num_bits=bits)
+
+            quant_pooling.QuantAvgPool2d.set_default_quant_desc_input(quant_desc_input)
+            quant_pooling_object = quant_pooling.QuantAvgPool2d(kernel_size=3, stride=1)
+
+            test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+            quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)), bits)
+
+            out1 = F.avg_pool2d(quant_input, 3, 1, 0, False, True, None)
+            out2 = quant_pooling_object(test_input)
+            np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_fake_quant_disable(self):
+        quant_pooling_object = quant_pooling.QuantAvgPool2d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+        quant_pooling_object.input_quantizer.disable()
+
+        out1 = F.avg_pool2d(test_input, 3, 1, 0, False, True, None)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantAvgPool3d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantAvgPool3d(kernel_size=3, stride=1,
+                                                                quant_desc_input=
+                                                                tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantAvgPool3d(kernel_size=3, stride=1)
+
+        test_input = torch.randn(5, 5, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.avg_pool3d(quant_input, 3, 1, 0, False, True, None)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantAdaptiveAvgPool1d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool1d(output_size=3,
+                                                                        quant_desc_input=
+                                                                        tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool1d(output_size=3)
+
+        test_input = torch.randn(1, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.adaptive_avg_pool1d(quant_input, 3)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantAdaptiveAvgPool2d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool2d(output_size=3,
+                                                                        quant_desc_input=
+                                                                        tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool2d(output_size=3)
+
+        test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.adaptive_avg_pool2d(quant_input, 3)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_variable_bits(self):
+        # Repeat checking the output for variable number of bits to QuantDescriptor
+        for bits in [2, 4, 6]:
+            quant_desc_input = tensor_quant.QuantDescriptor(num_bits=bits)
+
+            quant_pooling.QuantAdaptiveAvgPool2d.set_default_quant_desc_input(quant_desc_input)
+            quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool2d(output_size=3)
+
+            test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+            quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)), bits)
+
+            out1 = F.adaptive_avg_pool2d(quant_input, 3)
+            out2 = quant_pooling_object(test_input)
+            np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+    def test_input_fake_quant_disable(self):
+        quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool2d(output_size=3)
+
+        test_input = torch.randn(1, 5, 5, 5, dtype=torch.double)
+
+        quant_pooling_object.input_quantizer.disable()
+
+        out1 = F.adaptive_avg_pool2d(test_input, 3)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())
+
+class TestQuantAdaptiveAvgPool3d():
+
+    def test_raise(self):
+        with pytest.raises(ValueError) as excinfo:
+            quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool3d(output_size=3,
+                                                                        quant_desc_input=
+                                                                        tensor_quant.QuantDescriptor(fake_quant=False))
+        assert "Only fake quantization is supported" in str(excinfo.value)
+
+    # Quantizing activations
+    def test_input_fake_quant(self):
+        quant_pooling_object = quant_pooling.QuantAdaptiveAvgPool3d(output_size=3)
+
+        test_input = torch.randn(5, 5, 5, 5, dtype=torch.double)
+
+        quant_input = tensor_quant.fake_tensor_quant(test_input, torch.max(torch.abs(test_input)))
+
+        out1 = F.adaptive_avg_pool3d(quant_input, 3)
+        out2 = quant_pooling_object(test_input)
+        np.testing.assert_array_equal(out1.detach().cpu().numpy(), out2.detach().cpu().numpy())

tools/pytorch-quantization/tests/quant_rnn_test.py

@@ -0,0 +1,520 @@
+#
+# 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.
+#
+
+
+"""tests of QuantRNN module.
+"""
+import pytest
+
+import torch
+from torch import nn
+
+import numpy as np
+
+from pytorch_quantization.nn.modules import quant_rnn
+from pytorch_quantization import tensor_quant
+
+from tests.fixtures import verbose
+
+from . import utils
+
+# make everything run on the GPU
+torch.set_default_tensor_type('torch.cuda.FloatTensor')
+# change default type to double if utils.compare flags a small error, may just be floating point rounding error
+# torch.set_default_tensor_type('torch.cuda.DoubleTensor')
+
+np.random.seed(1234)
+torch.manual_seed(1234)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(1234)
+
+# pylint: disable=no-self-use, missing-docstring, redefined-builtin, bad-continuation
+
+# global state for saving/loading test vectors
+SAVE_VECTORS = 0
+VECTOR_FILE = 'tests/quant_rnn_test_vectors.pt'
+if SAVE_VECTORS:
+    TEST_VECTORS = dict()
+else:
+    TEST_VECTORS = torch.load(VECTOR_FILE)
+
+
+class TestQuantLSTMCell():
+    """
+    tests for quant_rnn.QuantLSTMCell
+    default parameters in QuantLSTMCell:
+        bias=True,
+        num_bits_weight=8, quant_mode_weight='per_channel',
+        num_bits_input=8, quant_mode_input='per_tensor'
+
+    Tests of real quantization mode (nonfake) are disabled as it is not fully supported yet.
+    """
+
+    def test_basic_forward(self, verbose):
+        """Do a forward pass on the cell module and see if anything catches fire."""
+        batch = 7
+        input_size = 11
+        hidden_size = 9
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=8)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=8, axis=(1,))
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        quant_rnn_object._input_quantizer.disable()
+        quant_rnn_object._weight_quantizer.disable()
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_rnn_object(input, hx=(hidden, cell))
+
+    def test_no_quant_input_hidden(self, verbose):
+        """QuantLSTM with quantization disabled vs. pytorch LSTM for input and hidden inputs."""
+        batch = 17
+        input_size = 13
+        hidden_size = 7
+
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=False)
+        quant_rnn_object._input_quantizer.disable()
+        quant_rnn_object._weight_quantizer.disable()
+        ref_rnn_object = nn.LSTMCell(input_size, hidden_size, bias=False)
+
+        # copy weights from one rnn to the other
+        ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_hout, quant_cout = quant_rnn_object(input, hx=(hidden, cell))
+        ref_hout, ref_cout = ref_rnn_object(input, hx=(hidden, cell))
+
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_no_quant_input_hidden_bias(self, verbose):
+        """QuantLSTMCell with quantization disabled vs. pytorch LSTMCell for input, hidden inputs and bias."""
+        batch = 19
+        input_size = 11
+        hidden_size = 3
+
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=True)
+        quant_rnn_object._input_quantizer.disable()
+        quant_rnn_object._weight_quantizer.disable()
+        ref_rnn_object = nn.LSTMCell(input_size, hidden_size, bias=True)
+
+        # copy weights from one rnn to the other
+        ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_hout, quant_cout = quant_rnn_object(input, hx=(hidden, cell))
+        ref_hout, ref_cout = ref_rnn_object(input, hx=(hidden, cell))
+
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_against_unquantized(self, verbose):
+        """Quantization should introduce bounded error utils.compare to pytorch implementation."""
+        batch = 9
+        input_size = 13
+        hidden_size = 7
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=16)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=16, axis=(1,))
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTMCell(input_size, hidden_size, bias=False)
+
+        # copy weights from one rnn to the other
+        ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_hout, quant_cout = quant_rnn_object(input, hx=(hidden, cell))
+        ref_hout, ref_cout = ref_rnn_object(input, hx=(hidden, cell))
+
+        # The difference between reference and quantized should be bounded in a range
+        # Small values which become 0 after quantization lead to large relative errors. rtol and atol could be
+        # much smaller without those values
+        utils.compare(quant_hout, ref_hout, rtol=1e-4, atol=1e-4)
+        utils.compare(quant_cout, ref_cout, rtol=1e-4, atol=1e-4)
+
+        # check that quantization introduces some error
+        utils.assert_min_mse(quant_hout, ref_hout, tol=1e-20)
+        utils.assert_min_mse(quant_cout, ref_cout, tol=1e-20)
+
+    def test_quant_input_hidden(self, verbose):
+        """QuantLSTMCell vs. manual input quantization + pytorchLSTMCell."""
+        batch = 15
+        input_size = 121
+        hidden_size = 51
+        num_bits = 4
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTMCell(input_size, hidden_size, bias=False)
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_hout, quant_cout = quant_rnn_object(input, hx=(hidden, cell))
+
+        quant_input, quant_hidden = utils.quantize_by_range_fused((input, hidden), num_bits)
+
+        utils.copy_state_and_quantize_fused(ref_rnn_object, quant_rnn_object, num_bits)
+
+        ref_hout, ref_cout = ref_rnn_object(quant_input, hx=(quant_hidden, cell))
+
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_quant_input_hidden_bias(self, verbose):
+        """QuantLSTMCell vs. manual input quantization + pytorchLSTMCell
+            bias should not be quantized
+        """
+        batch = 9
+        input_size = 23
+        hidden_size = 31
+        num_bits = 7
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=True,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTMCell(input_size, hidden_size, bias=True)
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_hout, quant_cout = quant_rnn_object(input, hx=(hidden, cell))
+
+        quant_input, quant_hidden = utils.quantize_by_range_fused((input, hidden), num_bits)
+
+        utils.copy_state_and_quantize_fused(ref_rnn_object, quant_rnn_object, num_bits)
+
+        ref_hout, ref_cout = ref_rnn_object(quant_input, hx=(quant_hidden, cell))
+
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_quant_different_prec(self, verbose):
+        """QuantLSTMCell vs. manual input quantization + pytorch LSTMCell
+            different input and weight precisions
+        """
+        batch = 27
+        input_size = 11
+        hidden_size = 10
+        num_bits_weight = 4
+        num_bits_input = 8
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=num_bits_input)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=num_bits_weight)
+        quant_rnn_object = quant_rnn.QuantLSTMCell(input_size, hidden_size, bias=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTMCell(input_size, hidden_size, bias=False)
+
+        input = torch.randn(batch, input_size)
+        hidden = torch.randn(batch, hidden_size)
+        cell = torch.randn(batch, hidden_size)
+
+        quant_hout, quant_cout = quant_rnn_object(input, hx=(hidden, cell))
+
+        quant_input, quant_hidden = utils.quantize_by_range_fused((input, hidden), num_bits_input)
+
+        utils.copy_state_and_quantize_fused(ref_rnn_object, quant_rnn_object, num_bits_weight)
+
+        ref_hout, ref_cout = ref_rnn_object(quant_input, hx=(quant_hidden, cell))
+
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+
+class TestQuantLSTM():
+    """
+    tests for quant_rnn.QuantLSTM
+    default parameters in QuantLSTM:
+        bias=True,
+        quant_weight=True, bits_weight=8, fake_quantTrue, quant_mode_weight='channel',
+        quant_input=True, bits_acts=8, quant_mode_input='tensor'
+
+    Tests of real quantization mode (nonfake) are disabled as it is not fully supported yet.
+    """
+
+    def test_basic_forward(self, verbose):
+        """Do a forward pass on the layer module and see if anything catches fire."""
+        batch = 5
+        input_size = 13
+        hidden_size = 31
+        seq_len = 1
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=8)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=8, axis=(1,))
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+        quant_rnn_object(input, hx=(hidden, cell))
+
+    def test_no_quant(self, verbose):
+        """QuantLSTM with quantization disabled vs. pytorch LSTM."""
+        batch = 11
+        input_size = 14
+        hidden_size = 22
+        seq_len = 1
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False)
+        quant_rnn_object._input_quantizers[0].disable()
+        quant_rnn_object._weight_quantizers[0].disable()
+        ref_rnn_object = nn.LSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False)
+
+        # copy weights from one rnn to the other
+        ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+
+        quant_out, (quant_hout, quant_cout) = quant_rnn_object(input)
+        ref_out, (ref_hout, ref_cout) = ref_rnn_object(input)
+
+        utils.compare(quant_out, ref_out)
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_no_quant_input_hidden(self, verbose):
+        """QuantLSTM with quantization disabled vs. pytorch LSTM for input and hidden inputs."""
+        batch = 13
+        input_size = 19
+        hidden_size = 20
+        seq_len = 1
+
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False)
+        quant_rnn_object._input_quantizers[0].disable()
+        quant_rnn_object._weight_quantizers[0].disable()
+        ref_rnn_object = nn.LSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False)
+
+        # copy weights from one rnn to the other
+        ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+
+        quant_out, (quant_hout, quant_cout) = quant_rnn_object(input, hx=(hidden, cell))
+        ref_out, (ref_hout, ref_cout) = ref_rnn_object(input, hx=(hidden, cell))
+
+        utils.compare(quant_out, ref_out)
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_no_quant_all_modes(self, verbose):
+        """QuantLSTM with quantization disabled vs. pytorch LSTM for all modes."""
+
+        def testcase(input_size, hidden_size, seq_len, batch, num_layers, bias, batch_first, dropout, bidirectional):
+
+            quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size,
+                    num_layers=num_layers, bias=bias, batch_first=batch_first, dropout=dropout,
+                    bidirectional=bidirectional)
+
+            num_quantizers = num_layers * 2 if bidirectional else num_layers
+            for i in range(num_quantizers):
+                quant_rnn_object._input_quantizers[i].disable()
+                quant_rnn_object._weight_quantizers[i].disable()
+
+            ref_rnn_object = nn.LSTM(input_size, hidden_size,
+                    num_layers=num_layers, bias=bias, batch_first=batch_first, dropout=dropout,
+                    bidirectional=bidirectional)
+
+            # copy state from one rnn to the other
+            ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+            input = torch.randn(seq_len, batch, input_size)
+            num_directions = 2 if bidirectional else 1
+            hidden = torch.randn(num_layers*num_directions, batch, hidden_size)
+            cell = torch.randn(num_layers*num_directions, batch, hidden_size)
+
+            quant_out, (quant_hout, quant_cout) = quant_rnn_object(input, hx=(hidden, cell))
+            ref_out, (ref_hout, ref_cout) = ref_rnn_object(input, hx=(hidden, cell))
+
+            utils.compare(quant_out, ref_out)
+            utils.compare(quant_hout, ref_hout)
+            utils.compare(quant_cout, ref_cout)
+
+        # test various permuatations of the following parameters:
+        #   size, num_layers, bias, batch_first, dropout, bidirectional
+        testcase(32, 27, 1, 1, 1, False, False, 0, False)
+        testcase(19, 63, 1, 1, 2, False, False, 0, False)
+        testcase(11, 41, 1, 1, 1, True, False, 0, False)
+        testcase(33, 31, 1, 1, 1, False, True, 0, False)
+        # testcase(32, 32, 1, 1, 2, False, False, 0.5, False) #TODO(pjudd) this fails look into dropout seeding
+        testcase(73, 13, 1, 1, 1, False, False, 0, True)
+
+    def test_against_unquantized(self, verbose):
+        """Quantization should introduce bounded error utils.compare to pytorch implementation."""
+        batch = 21
+        input_size = 33
+        hidden_size = 25
+        seq_len = 1
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=16)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=16, axis=(1,))
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTM(input_size, hidden_size,
+                num_layers=1, bias=False, batch_first=False, dropout=0, bidirectional=False)
+
+        # copy weights from one rnn to the other
+        ref_rnn_object.load_state_dict(quant_rnn_object.state_dict())
+
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+
+        quant_out, (quant_hout, quant_cout) = quant_rnn_object(input, hx=(hidden, cell))
+        ref_out, (ref_hout, ref_cout) = ref_rnn_object(input, hx=(hidden, cell))
+
+        # The difference between reference and quantized should be bounded in a range
+        # Small values which become 0 after quantization lead to large relative errors. rtol and atol could be
+        # much smaller without those values
+        utils.compare(quant_out, ref_out, rtol=1e-4, atol=1e-4)
+        utils.compare(quant_hout, ref_hout, rtol=1e-4, atol=1e-4)
+        utils.compare(quant_cout, ref_cout, rtol=1e-4, atol=1e-4)
+
+        # check that quantization introduces some error
+        utils.assert_min_mse(quant_out, ref_out, tol=1e-20)
+        utils.assert_min_mse(quant_hout, ref_hout, tol=1e-20)
+        utils.assert_min_mse(quant_cout, ref_cout, tol=1e-20)
+
+    def test_quant_input_hidden(self, verbose):
+        """QuantLSTM vs. manual input quantization + pytorchLSTM."""
+        batch = 13
+        input_size = 17
+        hidden_size = 7
+        seq_len = 1
+        num_bits = 6
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size, num_layers=1, bias=False,
+                batch_first=False, dropout=0, bidirectional=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTM(input_size, hidden_size, num_layers=1, bias=False,
+                batch_first=False, dropout=0, bidirectional=False)
+
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+
+        quant_input, quant_hidden = utils.quantize_by_range_fused((input, hidden), num_bits)
+
+        utils.copy_state_and_quantize_fused(ref_rnn_object, quant_rnn_object, num_bits)
+
+        quant_out, (quant_hout, quant_cout) = quant_rnn_object(input, hx=(hidden, cell))
+        ref_out, (ref_hout, ref_cout) = ref_rnn_object(quant_input, hx=(quant_hidden, cell))
+
+        utils.compare(quant_out, ref_out)
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_quant_input_hidden_bias(self, verbose):
+        """QuantLSTM vs. manual input quantization + pytorchLSTM."""
+        batch = 17
+        input_size = 13
+        hidden_size = 7
+        seq_len = 1
+        num_bits = 5
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size, num_layers=1, bias=True,
+                batch_first=False, dropout=0, bidirectional=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTM(input_size, hidden_size, num_layers=1, bias=True,
+                batch_first=False, dropout=0, bidirectional=False)
+
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+
+        quant_input, quant_hidden = utils.quantize_by_range_fused((input, hidden), num_bits)
+
+        utils.copy_state_and_quantize_fused(ref_rnn_object, quant_rnn_object, num_bits)
+
+        quant_out, (quant_hout, quant_cout) = quant_rnn_object(input, hx=(hidden, cell))
+        ref_out, (ref_hout, ref_cout) = ref_rnn_object(quant_input, hx=(quant_hidden, cell))
+
+        utils.compare(quant_out, ref_out)
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+    def test_quant_different_prec(self, verbose):
+        """QuantLSTM vs. manual input quantization + pytorchLSTM."""
+        batch = 22
+        input_size = 23
+        hidden_size = 24
+        seq_len = 1
+        num_bits_weight = 4
+        num_bits_input = 8
+
+        quant_desc_input = tensor_quant.QuantDescriptor(num_bits=num_bits_input)
+        quant_desc_weight = tensor_quant.QuantDescriptor(num_bits=num_bits_weight)
+        quant_rnn_object = quant_rnn.QuantLSTM(input_size, hidden_size, num_layers=1, bias=False,
+                batch_first=False, dropout=0, bidirectional=False,
+                quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight)
+        ref_rnn_object = nn.LSTM(input_size, hidden_size, num_layers=1, bias=False,
+                batch_first=False, dropout=0, bidirectional=False)
+
+        input = torch.randn(seq_len, batch, input_size)
+        hidden = torch.randn(seq_len, batch, hidden_size)
+        cell = torch.randn(seq_len, batch, hidden_size)
+
+        quant_input, quant_hidden = utils.quantize_by_range_fused((input, hidden), num_bits_input)
+
+        utils.copy_state_and_quantize_fused(ref_rnn_object, quant_rnn_object, num_bits_weight)
+
+        quant_out, (quant_hout, quant_cout) = quant_rnn_object(input, hx=(hidden, cell))
+        ref_out, (ref_hout, ref_cout) = ref_rnn_object(quant_input, hx=(quant_hidden, cell))
+
+        utils.compare(quant_out, ref_out)
+        utils.compare(quant_hout, ref_hout)
+        utils.compare(quant_cout, ref_cout)
+
+
+class TestEpilogue():
+    """Run after all tests to save globals."""
+
+    def test_save_vectors(self, verbose):
+        """Save test vectors to file."""
+        if SAVE_VECTORS:
+            torch.save(TEST_VECTORS, VECTOR_FILE)
+            raise Exception('Saved test vectors to {}, for testing set SAVE_VECTORS = 0'.format(VECTOR_FILE))

tools/pytorch-quantization/tests/quant_utils_test.py

@@ -0,0 +1,52 @@
+#
+# 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.
+#
+
+
+"""Test pytorch_quantization.utils"""
+import pytest
+import numpy as np
+
+import torch
+from pytorch_quantization import utils as quant_utils
+from tests.fixtures import verbose
+
+np.random.seed(12345)
+
+# pylint:disable=missing-docstring, no-self-use
+
+class TestQuantUtils():
+
+    def test_reduce_amax(self):
+        x_np = (np.random.rand(3, 7, 11, 13, 17) - 0.1).astype(np.float32)
+        x_torch = torch.tensor(x_np)
+
+        # Test reduce to one value
+        amax_np = np.max(np.abs(x_np))
+        amax_torch = quant_utils.reduce_amax(x_torch)
+        np.testing.assert_array_equal(amax_np, amax_torch.cpu().numpy())
+
+        # Test different axis
+        axes = [(1, 2, 3), (0, 2, 3), (0, 3), (0, 1, 3, 4)]
+        for axis in axes:
+            keepdims = np.random.rand() > 0.5
+            amax_np = np.max(np.abs(x_np), axis=axis, keepdims=keepdims)
+            amax_torch = quant_utils.reduce_amax(x_torch, axis=axis, keepdims=keepdims)
+            np.testing.assert_array_almost_equal(amax_np, amax_torch.cpu().numpy())
+
+        with pytest.raises(ValueError) as excinfo:
+            quant_utils.reduce_amax(x_torch, axis=(0, 1, 2, 3, 4, 5))
+            assert "Cannot reduce more axes" in str(excinfo.value)

tools/pytorch-quantization/tests/tensor_quant_test.py

@@ -0,0 +1,482 @@
+#
+# 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.
+#
+"""tests of tensor quantization function and module"""
+import contextlib
+
+import pytest
+import numpy as np
+
+import torch
+from torch.nn.parameter import Parameter
+
+from pytorch_quantization import cuda_ext
+from pytorch_quantization import tensor_quant
+import pytorch_quantization.utils as quant_utils
+
+import tests.utils as test_utils
+from tests.fixtures import verbose
+
+np.random.seed(123456)  # seed 1234 causes 1 number mismatch at 6th decimal in one of the tests
+
+# pylint:disable=missing-docstring, no-self-use
+
+
+class TestTensorQuant():
+
+    def test_simple_run(self):
+        """ quantizer passes gradcheck
+        """
+        x = Parameter(torch.randn(2, 3, dtype=torch.float64).cuda()) * 100
+        tensor_quant.tensor_quant(x, torch.max(torch.abs(x)), 7)
+
+    def test_per_tensor_scale(self):
+        """ tensor_quant matches numpy quantization
+        """
+        torch.set_default_tensor_type('torch.cuda.FloatTensor')  # Test on GPU
+        x_np = np.random.rand(1023)
+        x_torch = torch.Tensor(x_np)
+        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)))
+        quant_x_torch, _ = tensor_quant.tensor_quant(x_torch, torch.max(torch.abs(x_torch)))
+        np.testing.assert_array_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+        torch.set_default_tensor_type('torch.FloatTensor')
+
+    def test_per_channel_scale(self):
+        """ fake_tensor_quant performs per channel quantization
+        """
+        x_np = np.random.rand(15, 15, 64, 128).astype('float32')
+        x_torch = torch.Tensor(x_np).cuda()
+
+        # Pytorch filter layout seems to be KCRS, reduce max to shape [K, 1, 1, 1] to test per channel scale
+        # Shrink max a little, so that clip behavior is tested
+        amax_x_np = 0.7 * np.max(np.abs(x_np), axis=(1, 2, 3), keepdims=True)
+        # Pytorch's max function doesn't support reduces multiple axis, and returns (max, argmax) tuple,
+        # so it has to be reduced by multiple torch.max
+        amax_x_torch = 0.7 * torch.max(
+            torch.max(torch.max(x_torch, dim=1, keepdim=True)[0], dim=2, keepdim=True)[0], dim=3, keepdim=True)[0]
+
+        quant_x_np = test_utils.quant_np(x_np, amax_x_np)
+        quant_x_torch, _ = tensor_quant.tensor_quant(x_torch, amax_x_torch)
+
+        # np.testing.assert_array_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+        # Pytorch numerics is not the same as numpy, it will be off by 1
+        np.testing.assert_array_less(np.abs(quant_x_torch.cpu().numpy() - quant_x_np), 2)
+        if verbose:
+            mismatches = np.where(np.abs(quant_x_torch.cpu().numpy() - quant_x_np) >= 1)
+            print("Mismatches:")
+            print(" Original: ", x_np[mismatches])
+            print(" numpy: ", quant_x_np[mismatches])
+            print(" Pytorch: ", quant_x_torch.cpu().numpy()[mismatches])
+
+    def test_backward(self):
+        """ tensor_quant implements straight through estimator on the backward pass
+            Note: this does not work for integer output_dtype
+        """
+        x = torch.randn(3, 7, requires_grad=True).cuda()
+        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
+        quant_x, _ = tensor_quant.tensor_quant(x, x.abs().max(), 7)
+        float_quant_x = quant_x.type(torch.FloatTensor).cuda()
+        x.retain_grad()
+        float_quant_x.retain_grad()
+        criterion = torch.nn.CrossEntropyLoss().cuda()
+        loss = criterion(float_quant_x, labels)
+        loss.backward()
+        np.testing.assert_array_equal(float_quant_x.grad.cpu().numpy(), x.grad.cpu().numpy())
+
+    def test_unsigned(self):
+        x_np = np.random.rand(1023).astype('float32')
+        x_torch = torch.Tensor(x_np)
+        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), num_bits=9, fake=False)
+        quant_x_torch, _ = tensor_quant.tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True)
+        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+
+        x_torch = torch.randn(3, 7)
+        with pytest.raises(TypeError, match="Negative values encountered"):
+            tensor_quant.tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True)
+
+    def test_overflow_fp16(self):
+        x_torch = torch.randn(1023).cuda().half()
+        with pytest.raises(ValueError, match="scale is too large for FP16"):
+            quant_x_torch, scale = tensor_quant.tensor_quant(x_torch, torch.tensor(1e-4).cuda().half(), 8, False)
+
+    def test_clip_gradient(self):
+        x = torch.randn(3, 7, requires_grad=True).cuda()
+        x.retain_grad()
+        amax = x.abs().max() / 2
+        x_in_range = (-amax <= x) * (x <= amax)
+        quant_x, _ = tensor_quant.tensor_quant(x, amax, 8)
+        loss = torch.sum((quant_x - 0.5)**2)
+        loss.backward()
+        np.testing.assert_array_equal(x.grad.cpu().numpy() != 0, x_in_range.cpu().numpy())
+
+    def test_full_range(self):
+        """ fake_tensor_quant uses the full integer range when narrow=False
+        """
+        x_np = np.random.rand(1023).astype('float32')
+        x_torch = torch.Tensor(x_np).cuda()
+        amax = np.max(np.abs(x_np))
+        quant_x_np = test_utils.quant_np(x_np, amax, num_bits=9, fake=False, narrow_range=False)
+        quant_x_torch, _ = tensor_quant.tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True, False)
+        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+
+
+class TestFakeTensorQuant():
+
+    def test_simple_run(self):
+        x = Parameter(torch.randn(3, 7).cuda())
+        tensor_quant.fake_tensor_quant(x, torch.max(torch.abs(x)))
+
+    def test_per_tensor_scale(self):
+        """ fake_tensor_quant matches numpy quantization
+        """
+        x_np = np.random.rand(13).astype('float32')
+        print(x_np)
+        x_torch = torch.Tensor(x_np).cuda()
+        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), fake=True)
+        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch)))
+        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+
+    def test_per_channel_scale(self):
+        """ fake_tensor_quant performs per channel quantization
+        """
+        x_np = np.random.rand(15, 15, 64, 128).astype('float32')
+        x_torch = torch.Tensor(x_np).cuda()
+
+        # Pytorch filter layout seems to be KCRS, reduce max to shape [K, 1, 1, 1] to test per channel scale
+        # Shrink max a little, so that clip behavior is tested
+        amax_x_np = 0.9 * np.max(np.abs(x_np), axis=(1, 2, 3), keepdims=True)
+        # Pytorch's max function doesn't support reduces multiple axis, and returns (max, argmax) tuple,
+        # so it has to be reduced by multiple torch.max
+        amax_x_torch = 0.9 * torch.max(
+            torch.max(torch.max(x_torch, dim=1, keepdim=True)[0], dim=2, keepdim=True)[0], dim=3, keepdim=True)[0]
+
+        quant_x_np = test_utils.quant_np(x_np, amax_x_np, fake=True)
+        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, amax_x_torch)
+
+        # Pytorch numerics is not the same as numpy, results will be off a little
+        # np.testing.assert_array_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np, decimal=2)
+        if verbose:
+            mismatches = np.where(np.abs(quant_x_torch.cpu().numpy() - quant_x_np) >= 1e-5)
+            print("Mismatches:")
+            print(" Original: ", x_np[mismatches])
+            print(" numpy: ", quant_x_np[mismatches])
+            print(" Pytorch: ", quant_x_torch.cpu().numpy()[mismatches])
+
+    def test_backward(self):
+        """ fake_tensor_quant implements straight through estimator on the backward pass
+        """
+        x = torch.randn(3, 7, requires_grad=True).cuda()
+        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
+        quant_x = tensor_quant.fake_tensor_quant(x, torch.max(torch.abs(x)), 7)
+        x.retain_grad()
+        quant_x.retain_grad()
+        criterion = torch.nn.CrossEntropyLoss().cuda()
+        loss = criterion(quant_x, labels)
+        loss.backward()
+        np.testing.assert_array_equal(quant_x.grad.cpu().numpy(), x.grad.cpu().numpy())
+
+    def test_unsigned(self):
+        x_np = np.random.rand(1023).astype('float32')
+        x_torch = torch.Tensor(x_np).cuda()
+        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), num_bits=9, fake=True)
+        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True)
+        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+
+    def test_cuda_ext(self):
+        x_np = np.random.rand(1023).astype('float32')
+        x_torch = torch.Tensor(x_np).cuda()
+
+        for num_bits in [3, 4, 5, 7, 8, 11]:
+            for unsigned in [True, False]:
+                test_utils.compare(cuda_ext.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch)), num_bits,
+                                                              unsigned),
+                                   tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch)), num_bits,
+                                                                  unsigned),
+                                   rtol=0,
+                                   atol=0)
+
+        # Test fp16 and bf16
+        for dtype in [torch.float16, torch.bfloat16]:
+            x_np = np.random.rand(1023)
+            x_torch = torch.Tensor(x_np).cuda().to(dtype)
+            cuda_ext_out = cuda_ext.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch))).to(torch.float32)
+            pytorch_out = tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch))).to(torch.float32)
+            test_utils.compare(cuda_ext_out, pytorch_out, rtol=0, atol=0)
+
+    @pytest.mark.parametrize("dtype", [torch.float32, torch.float16, torch.bfloat16])
+    def test_cuda_ext_with_axis(self, dtype):
+        x_np = np.random.rand(3, 4, 5, 6)
+        x_torch = torch.Tensor(x_np).cuda().to(dtype)
+
+        # amax along axis 1
+        amax_torch = torch.tensor([0.8, 0.9, 0.7, 0.6], device="cuda")
+
+        for num_bits in [3, 4, 5, 7, 8, 11]:
+            for unsigned in [True, False]:
+                cuda_ext_out = cuda_ext.fake_tensor_quant_with_axis(x_torch, amax_torch, 1, num_bits, unsigned).to(torch.float32)
+                pytorch_out = tensor_quant.fake_tensor_quant(x_torch, amax_torch.view(1, -1, 1, 1), num_bits, unsigned).to(torch.float32)
+                test_utils.compare(cuda_ext_out, pytorch_out, rtol=0, atol=0)
+
+    def test_cuda_ext_inplace(self):
+        x_np = np.random.rand(1023).astype('float32')
+        x_torch = torch.Tensor(x_np).cuda()
+        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), fake=True)
+        cuda_ext.fake_tensor_quant_(x_torch, torch.max(torch.abs(x_torch)))
+        np.testing.assert_array_equal(x_torch.cpu().numpy(), quant_x_np)
+
+        # Test fp16 and bf16
+        for dtype in [torch.float16, torch.bfloat16]:
+            x_np = np.random.rand(1023)
+            x_torch = torch.Tensor(x_np).cuda().to(dtype)
+            quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)), fake=True)
+            cuda_ext.fake_tensor_quant_(x_torch, torch.max(torch.abs(x_torch)))
+            x_torch = x_torch.to(torch.float32)
+            np.testing.assert_array_almost_equal(x_torch.cpu().numpy(), quant_x_np, decimal=2)
+
+    def test_cuda_ext_tiny_amax(self):
+        x_torch = torch.rand(2, 3, 4, device="cuda")
+        amax = torch.tensor([1., 1.e-26, 1.], device="cuda").unsqueeze(-1).unsqueeze(1)
+        quant_x = cuda_ext.fake_tensor_quant_with_axis(x_torch, amax, axis=1)
+        assert quant_x[:, 1, :].sum() == 0
+
+    def test_overflow_fp16(self):
+        x_torch = torch.randn(1023).cuda().half()
+        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, torch.tensor(1e-4).cuda().half(), 8, False)
+        assert not (torch.isinf(quant_x_torch).any() or torch.isnan(quant_x_torch).any())
+
+    def test_clip_gradient(self):
+        x = torch.randn(3, 7, requires_grad=True).cuda()
+        x.retain_grad()
+        amax = x.abs().max() / 2
+        x_in_range = (-amax <= x) * (x <= amax)
+        quant_x = tensor_quant.fake_tensor_quant(x, amax, 8)
+        loss = torch.sum((quant_x - 0.5)**2)
+        loss.backward()
+        np.testing.assert_array_equal(x.grad.cpu().numpy() != 0, x_in_range.cpu().numpy())
+
+    def test_full_range(self):
+        """ fake_tensor_quant uses the full integer range when narrow=False
+        """
+        x_np = np.random.rand(1023).astype('float32')
+        x_torch = torch.Tensor(x_np).cuda()
+        amax = np.max(np.abs(x_np))
+        quant_x_np = test_utils.quant_np(x_np, amax, num_bits=9, fake=True, narrow_range=False)
+        quant_x_torch = tensor_quant.fake_tensor_quant(x_torch, torch.max(torch.abs(x_torch)), 8, True, False)
+        np.testing.assert_array_almost_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+
+    @pytest.mark.parametrize("dtype", ["float32", "float16"])
+    def test_against_legacy(self, dtype):
+        x_np = np.random.rand(3, 4, 5, 6).astype(dtype)
+        x_torch = torch.Tensor(x_np).cuda()
+
+        amax_torch = torch.tensor(0.7, device="cuda")
+
+        for num_bits in [3, 4, 5, 7, 8, 11]:
+            for unsigned in [True, False]:
+                legacy_out = tensor_quant.legacy_fake_tensor_quant(x_torch, amax_torch, num_bits, unsigned)
+                test_out = tensor_quant.fake_tensor_quant(x_torch, amax_torch, num_bits, unsigned)
+                test_utils.compare(legacy_out, test_out, rtol=0, atol=0)
+
+    def test_against_legacy_noncontiguous(self):
+        x_np = np.random.rand(3, 4, 5, 6)
+        x_torch = torch.Tensor(x_np).cuda()
+
+        amax_torch = torch.tensor(0.7, device="cuda")
+
+        x_torch_noncontiguous = x_torch[:, 2, :, 3]
+        assert not x_torch_noncontiguous.is_contiguous()
+
+        legacy_out = tensor_quant.legacy_fake_tensor_quant(x_torch_noncontiguous, amax_torch)
+        test_out = tensor_quant.fake_tensor_quant(x_torch_noncontiguous, amax_torch)
+        test_utils.compare(legacy_out, test_out, rtol=0, atol=0)
+
+    @pytest.mark.parametrize("dtype", ["float32", "float16"])
+    def test_against_legacy_with_axis(self, dtype):
+        x_np = np.random.rand(3, 4, 5, 6).astype(dtype)
+        x_torch = torch.Tensor(x_np).cuda()
+
+        # amax along axis 1
+        amax_torch = torch.tensor([0.8, 0.9, 0.7, 0.6], device="cuda").view(1, -1, 1, 1)
+
+        for num_bits in [3, 4, 5, 7, 8, 11]:
+            for unsigned in [True, False]:
+                legacy_out = tensor_quant.legacy_fake_tensor_quant(x_torch, amax_torch, num_bits, unsigned)
+                test_out = tensor_quant.fake_tensor_quant(x_torch, amax_torch, num_bits, unsigned)
+                test_utils.compare(legacy_out, test_out, rtol=0, atol=0)
+
+
+class TestQuantDescriptor():
+
+    def test_scaled_mode(self):
+        num_bits = np.random.randint(0, 16)
+
+        test_quant_desc = tensor_quant.QuantDescriptor(num_bits=num_bits)
+        assert test_quant_desc.num_bits == num_bits
+        assert test_quant_desc.axis is None
+        assert test_quant_desc.amax is None
+        assert not test_quant_desc.learn_amax
+
+        axis = (0, 1, 3)
+        test_quant_desc = tensor_quant.QuantDescriptor(axis=axis)
+        assert test_quant_desc.num_bits == 8  # default value
+        assert test_quant_desc.axis == axis
+        assert test_quant_desc.amax is None
+
+        amax = 0.7
+        test_quant_desc = tensor_quant.QuantDescriptor(amax=amax, unsigned=True)
+        assert test_quant_desc.axis is None
+        assert test_quant_desc.amax == np.float32(amax)
+        assert test_quant_desc.unsigned
+
+        amax = 0.7
+        test_quant_desc = tensor_quant.QuantDescriptor(amax=amax, learn_amax=True)
+        assert test_quant_desc.amax == np.float32(amax)
+        assert test_quant_desc.learn_amax
+
+        # Test the print string once if verbose is set.
+        if verbose:
+            print(test_quant_desc)
+
+        with pytest.raises(TypeError, match="must be float, list or ndarray"):
+            tensor_quant.QuantDescriptor(amax='oops')
+
+        with pytest.raises(TypeError, match="amax must be float, list or ndarray"):
+            tensor_quant.QuantDescriptor(amax='oops', learn_amax=True)
+
+        with pytest.raises(TypeError, match="axis is ignored and must be None"):
+            tensor_quant.QuantDescriptor(axis=(1, 2), amax=0.7, learn_amax=True)
+
+    def test_amax(self):
+        test_quant_desc = tensor_quant.QuantDescriptor()
+        assert test_quant_desc.amax is None
+
+        test_quant_desc = tensor_quant.QuantDescriptor(amax=1.2)
+        assert isinstance(test_quant_desc.amax, np.ndarray)
+        np.testing.assert_array_equal(test_quant_desc.amax, np.float32(1.2))
+
+        test_quant_desc = tensor_quant.QuantDescriptor(amax=[1.3, 1.4])
+        assert isinstance(test_quant_desc.amax, np.ndarray)
+        np.testing.assert_array_equal(test_quant_desc.amax, np.float32([1.3, 1.4]))
+
+        with pytest.raises(TypeError, match="must be float, list or ndarray"):
+            tensor_quant.QuantDescriptor(amax='oops')
+
+    def test_from_to_dict(self):
+        quant_desc_1 = tensor_quant.QuantDescriptor(num_bits=2,
+                                                    name='a',
+                                                    fake_quant=True,
+                                                    axis=(1, 2),
+                                                    amax=3.1415926536)
+        quant_desc_2 = tensor_quant.QuantDescriptor(**quant_desc_1.dict())
+        if verbose:
+            print(quant_desc_1.dict())
+        assert quant_desc_1 == quant_desc_2
+
+        quant_desc_1 = tensor_quant.QuantDescriptor(num_bits=2, amax=0.1, unsigned=True)
+        quant_desc_2 = tensor_quant.QuantDescriptor(**quant_desc_1.dict())
+        assert quant_desc_1 == quant_desc_2
+
+    def test_from_to_yaml(self):
+        quant_desc_1 = tensor_quant.QuantDescriptor(num_bits=2,
+                                                    name='a',
+                                                    fake_quant=True,
+                                                    axis=(1, 2),
+                                                    amax=3.1415926536)
+        quant_desc_2 = tensor_quant.QuantDescriptor.from_yaml(quant_desc_1.to_yaml())
+        if verbose:
+            print(quant_desc_1.to_yaml())
+        assert quant_desc_1 == quant_desc_2
+
+        quant_desc_1 = tensor_quant.QuantDescriptor(num_bits=2, amax=0.1)
+        quant_desc_2 = tensor_quant.QuantDescriptor.from_yaml(quant_desc_1.to_yaml())
+        assert quant_desc_1 == quant_desc_2
+
+
+class TestFakeAffineTensorQuant():
+
+    def test_simple_run(self, verbose):
+        x = np.array([-1., -13., -101., -128., 0., 2., 5., 13., 93., 111., 127.], dtype=np.float32)
+        torch_x = torch.tensor(x).cuda()
+        quant_x = tensor_quant.fake_affine_tensor_quant(torch_x, torch.min(torch_x), torch.max(torch_x))
+
+        if verbose:
+            print(quant_x)
+
+        np.testing.assert_array_almost_equal(quant_x.cpu().numpy(), x)
+
+    def test_clip_gradient(self):
+        x = torch.randn(3, 7, requires_grad=True).cuda()
+        x.retain_grad()
+        xmin = x.min() / 2
+        xmax = x.max() / 2
+        x_in_range = (xmin <= x) * (x <= xmax)
+        quant_x = tensor_quant.fake_affine_tensor_quant(x, xmin, xmax, 8)
+        loss = torch.sum((quant_x - 0.5)**2)
+        loss.backward()
+        np.testing.assert_array_equal(x.grad.cpu().numpy() != 0, x_in_range.cpu().numpy())
+
+
+class TestScaledE4M3():
+
+    x = [[-2.0000, -1.8000, -1.6000, -1.4000, -1.2000], [-1.0000, -0.8000, -0.6000, -0.4000, -0.2000],
+         [-0.0000, 0.2000, 0.4000, 0.6000, 0.8000], [1.0000, 1.2000, 1.4000, 1.6000, 1.8000]]
+
+    xq_unscaled = [[-2.0000, -1.7500, -1.6250, -1.3750, -1.2500], [-1.0000, -0.8125, -0.6250, -0.4062, -0.2031],
+                   [0.0000, 0.2031, 0.4062, 0.6250, 0.8125], [1.0000, 1.2500, 1.3750, 1.6250, 1.7500]]
+
+    xq_scaled = [[-2.0000, -1.8571, -1.5714, -1.4286, -1.1429], [-1.0000, -0.7857, -0.5714, -0.3929, -0.1964],
+                 [0.0000, 0.1964, 0.3929, 0.5714, 0.7857], [1.0000, 1.1429, 1.4286, 1.5714, 1.8571]]
+
+    def test_e4m3_no_scale(self):
+        x = torch.tensor(TestScaledE4M3.x, device="cuda")
+        xq_ref = torch.tensor(TestScaledE4M3.xq_unscaled, device="cuda")
+        e4m3_x = tensor_quant.scaled_e4m3(x, None)
+        test_utils.compare(e4m3_x, xq_ref, atol=1e-4, rtol=1e-4)
+
+    def test_e4m3_no_cpu(self):
+        x = torch.tensor(TestScaledE4M3.x)
+        xq_ref = torch.tensor(TestScaledE4M3.xq_unscaled)
+        e4m3_x = tensor_quant.scaled_e4m3(x, None)
+        test_utils.compare(e4m3_x, xq_ref, atol=1e-4, rtol=1e-4)
+
+    def test_with_amax(self):
+        x = torch.tensor(TestScaledE4M3.x, device="cuda").unsqueeze(-1)
+        xq_ref = torch.tensor(TestScaledE4M3.xq_scaled, device="cuda").unsqueeze(-1)
+
+        amax = quant_utils.reduce_amax(x, axis=None, keepdims=True)
+
+        e4m3_x = tensor_quant.scaled_e4m3(x, amax)
+
+        test_utils.compare(e4m3_x, xq_ref, atol=1e-4, rtol=1e-4)
+
+    def test_e4m3_incontiguous(self):
+        x = torch.tensor(TestScaledE4M3.x, device="cuda").transpose(1, 0)
+        xq_ref = torch.tensor(TestScaledE4M3.xq_unscaled, device="cuda").transpose(1, 0)
+        assert not x.is_contiguous()
+        e4m3_x = tensor_quant.scaled_e4m3(x, None)
+        test_utils.compare(e4m3_x, xq_ref, atol=1e-4, rtol=1e-4)
+
+    def test_backward(self):
+        x = torch.randn(3, 7, requires_grad=True).cuda()
+        labels = torch.randint(6, (3,)).type(torch.LongTensor).cuda()
+        quant_x = tensor_quant.scaled_e4m3(x, None)
+        x.retain_grad()
+        quant_x.retain_grad()
+        criterion = torch.nn.CrossEntropyLoss().cuda()
+        loss = criterion(quant_x, labels)
+        loss.backward()
+        np.testing.assert_array_equal(quant_x.grad.cpu().numpy(), x.grad.cpu().numpy())

tools/pytorch-quantization/tests/tensor_quantizer_test.py

@@ -0,0 +1,278 @@
+#
+# 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.
+#
+"""tests of tensor quantizer"""
+import contextlib
+
+import pytest
+import numpy as np
+
+import torch
+
+from pytorch_quantization import tensor_quant
+from pytorch_quantization.nn.modules import tensor_quantizer
+from pytorch_quantization import utils as quant_utils
+import tests.utils as test_utils
+from tests.fixtures import verbose
+
+np.random.seed(12345)
+
+# pylint:disable=missing-docstring, no-self-use
+
+
+class TestTensorQuantizer():
+
+    def test_simple_run(self):
+        """Quantizer calls fake_tensor_quant by default"""
+        x = torch.randn(3, 7).cuda()
+        amax_x = torch.max(torch.abs(x))
+        fn_quant_x = tensor_quant.fake_tensor_quant(x, amax_x)
+        quantizer = tensor_quantizer.TensorQuantizer()
+        module_quant_x = quantizer(x)
+        np.testing.assert_array_equal(fn_quant_x.cpu().numpy(), module_quant_x.cpu().numpy())
+
+    def test_simple_run_no_fake(self):
+        """Quantizer fake_quant=False calls tensor_quant and sets the scale property"""
+        x = torch.randn(3, 7).cuda()
+        amax_x = torch.max(torch.abs(x))
+        fn_quant_x, fn_scale = tensor_quant.tensor_quant(x, amax_x)
+        quantizer = tensor_quantizer.TensorQuantizer(tensor_quant.QuantDescriptor(num_bits=8, fake_quant=False))
+        module_quant_x = quantizer(x)
+        module_scale = quantizer.scale
+        np.testing.assert_array_equal(fn_quant_x.cpu().numpy(), module_quant_x.cpu().numpy())
+        np.testing.assert_array_equal(fn_scale.cpu().numpy(), module_scale.cpu().numpy())
+
+    def test_per_tensor_scale(self):
+        """Quantizer performs expected quantization"""
+        x_np = np.random.rand(1023)
+        x_torch = torch.Tensor(x_np)
+        quant_x_np = test_utils.quant_np(x_np, np.max(np.abs(x_np)))
+        quantizer = tensor_quantizer.TensorQuantizer(tensor_quant.QuantDescriptor(num_bits=8, fake_quant=False))
+        module_quant_x = quantizer(x_torch)
+        np.testing.assert_array_equal(module_quant_x.cpu().numpy(), quant_x_np)
+
+    def test_per_channel_scale(self, verbose):
+        """Quantizer performs per channel scaling"""
+        x_np = np.random.rand(15, 15, 64, 128).astype('float32')
+        x_torch = torch.Tensor(x_np).cuda()
+
+        # Pytorch filter layout seems to be KCRS, reduce max to shape [K, 1, 1, 1] to test per channel scale
+        # Shrink max a little, so that clip behavior is tested
+        amax_x_np = 0.7 * np.max(np.abs(x_np), axis=(1, 2, 3), keepdims=True)
+
+        quant_x_np = test_utils.quant_np(x_np, amax_x_np)
+        quantizer = tensor_quantizer.TensorQuantizer(
+            tensor_quant.QuantDescriptor(num_bits=8, axis=(0), fake_quant=False, scale_amax=0.7))
+        quantizer.cuda()
+        module_quant_x = quantizer(x_torch)
+
+        # np.testing.assert_array_equal(quant_x_torch.cpu().numpy(), quant_x_np)
+        # Pytorch numerics is not the same as numpy, it will be off by 1
+        error = np.abs(module_quant_x.cpu().numpy() - quant_x_np)
+        np.testing.assert_array_less(error, 2)
+        if verbose:
+            mismatches = np.where(error >= 1)
+            print("Mismatches:")
+            print(" Original: ", x_np[mismatches])
+            print(" numpy: ", quant_x_np[mismatches])
+            print(" TensorQuantizer: ", module_quant_x.cpu().numpy()[mismatches])
+
+    def test_learn_amax(self):
+        """Test the clip implied by learn_amax"""
+        x_np = np.random.rand(1023).astype(np.float32)
+        x_torch = torch.Tensor(x_np).cuda()
+        amax = 0.5
+        quant_x_np = test_utils.quant_np(x_np, 0.5, fake=True)
+        quantizer = tensor_quantizer.TensorQuantizer(
+            tensor_quant.QuantDescriptor(num_bits=8, amax=amax, learn_amax=True)).cuda()
+        assert hasattr(quantizer, 'clip')
+        module_quant_x = quantizer(x_torch)
+        np.testing.assert_array_equal(module_quant_x.cpu().detach().numpy(), quant_x_np)
+
+    def test_clip_mode(self):
+        """Test the clip stage only"""
+        x_np = np.random.rand(1023).astype(np.float32)
+        x_torch = torch.Tensor(x_np).cuda()
+        amax = 0.5
+        clip_x_np = np.clip(x_np, -amax, amax)
+        quantizer = tensor_quantizer.TensorQuantizer(tensor_quant.QuantDescriptor(amax=amax, learn_amax=True),
+                                                     if_quant=False,
+                                                     if_clip=True).cuda()
+        assert hasattr(quantizer, 'clip')
+        module_clip_x = quantizer(x_torch)
+        np.testing.assert_array_equal(module_clip_x.cpu().detach().numpy(), clip_x_np)
+
+    def test_scale_amax(self):
+        x_np = np.random.rand(1023).astype(np.float32)
+        x_torch = torch.Tensor(x_np).cuda()
+        amax = 0.5
+        scale_amax = 0.9
+        quant_x_np = test_utils.quant_np(x_np, amax * scale_amax, fake=True)
+        quantizer = tensor_quantizer.TensorQuantizer(
+            tensor_quant.QuantDescriptor(num_bits=8, amax=amax, scale_amax=scale_amax)).cuda()
+        module_quant_x = quantizer(x_torch)
+        np.testing.assert_array_equal(module_quant_x.cpu().detach().numpy(), quant_x_np)
+
+        # Test twice. There was a but in scale amax logic that modify the amax every time
+        module_quant_x = quantizer(x_torch)
+        np.testing.assert_array_equal(module_quant_x.cpu().detach().numpy(), quant_x_np)
+
+    def test_disable(self):
+        x = torch.randn(3, 7).cuda()
+        amax_x = torch.max(torch.abs(x))
+        quantizer = tensor_quantizer.TensorQuantizer(disabled=True).cuda()
+        module_quant_x = quantizer(x)
+        np.testing.assert_array_equal(x.cpu().numpy(), module_quant_x.cpu().numpy())
+
+    def test_state_loading(self):
+        """Test quant_desc loading via state_dict"""
+        amax = [3.142, 2.718]
+        quant_desc1 = tensor_quant.QuantDescriptor(amax=amax)
+        quantizer1 = tensor_quantizer.TensorQuantizer(quant_desc1)
+
+        # copy state
+        quantizer1.load_state_dict(quantizer1.state_dict())
+        np.testing.assert_array_equal(quantizer1.amax.detach().cpu().numpy(), quant_desc1.amax)
+
+    def test_properties(self):
+        quant_desc1 = tensor_quant.QuantDescriptor(amax=3.14)
+        quantizer1 = tensor_quantizer.TensorQuantizer(quant_desc1)
+        quantizer1.amax = 0.577
+
+        assert quantizer1.amax.detach().cpu().numpy() == np.float32(0.577)
+        np.testing.assert_array_equal(quantizer1.amax.detach().cpu().numpy(), quantizer1.amax)
+        assert quantizer1.step_size == 0.577 / 127.
+
+        quant_desc2 = tensor_quant.QuantDescriptor()
+        quantizer2 = tensor_quantizer.TensorQuantizer(quant_desc2)
+        amax_np = np.array([3.142, 2.718], dtype=np.float32)
+        quantizer2.amax = amax_np
+        np.testing.assert_array_equal(quantizer2.amax.detach().cpu().numpy(), amax_np)
+
+        quant_desc3 = tensor_quant.QuantDescriptor()
+        quantizer3 = tensor_quantizer.TensorQuantizer(quant_desc3)
+        assert quantizer3.amax is None
+
+    def test_init_calib(self):
+        quant_desc2 = tensor_quant.QuantDescriptor(axis=(0, 1))
+        quantizer2 = tensor_quantizer.TensorQuantizer(quant_desc2, if_calib=True, if_quant=False).cuda()
+
+        x_2 = torch.rand(127, 63, 7, 7).cuda()
+        quantizer2(x_2)
+        quantizer2.load_calib_amax()
+
+        assert quantizer2.amax.numel() == 127 * 63
+
+    def test_max_calib(self):
+        axis = 0
+        reduce_axis = (1, 2, 3)
+        quant_desc1 = tensor_quant.QuantDescriptor(axis=axis)
+        quantizer1 = tensor_quantizer.TensorQuantizer(quant_desc1).cuda()
+        quantizer1.enable_calib()
+
+        quant_desc1 = tensor_quant.QuantDescriptor(axis=axis)
+        quantizer1 = tensor_quantizer.TensorQuantizer(quant_desc1).cuda()
+        quantizer1.enable_calib()
+
+        with pytest.raises(RuntimeError, match="Calibrator returned None"):
+            quantizer1.load_calib_amax()
+
+        x_1 = torch.rand(127, 63, 7, 7).cuda()
+        x_2 = torch.rand(127, 63, 7, 7).cuda()
+        quantizer1(x_1)
+        quantizer1(x_2)
+        quantizer1.disable_calib()
+
+        global_amax = torch.max(quant_utils.reduce_amax(x_1, axis=reduce_axis, keepdims=True),
+                                quant_utils.reduce_amax(x_2, axis=reduce_axis, keepdims=True))
+        test_utils.compare(quantizer1._calibrator.compute_amax(), global_amax, atol=0, rtol=0, ctol=0)
+
+        quantizer1.load_calib_amax()
+        test_utils.compare(quantizer1.amax, global_amax, atol=0, rtol=0, ctol=0)
+
+        quant_desc2 = tensor_quant.QuantDescriptor(learn_amax=True)
+        quantizer2 = tensor_quantizer.TensorQuantizer(quant_desc2).cuda()
+        quantizer2.enable_calib()
+        quantizer2(x_1)
+        quantizer2(x_2)
+
+        quantizer2.load_calib_amax()
+        quantizer2.init_learn_amax()
+        test_utils.compare(quantizer2.clip.clip_value_min, -torch.max(global_amax), atol=0, rtol=0, ctol=0)
+        test_utils.compare(quantizer2.clip.clip_value_max, torch.max(global_amax), atol=0, rtol=0, ctol=0)
+
+    def test_entropy_and_percentile_calib(self):
+        """Don't really have a good way to test it."""
+        quant_desc1 = tensor_quant.QuantDescriptor(calib_method='histogram')
+        quantizer1 = tensor_quantizer.TensorQuantizer(quant_desc1, if_calib=True, if_quant=False).cuda()
+
+        x_1 = torch.rand(3, 63, 7, 7).cuda()
+        x_2 = torch.rand(3, 63, 7, 7).cuda()
+        quantizer1(x_1)
+        quantizer1(x_2)
+
+        quantizer1.load_calib_amax("entropy")
+        test_utils.compare(quantizer1._calibrator.compute_amax("entropy"), quantizer1.amax, atol=0, rtol=0, ctol=0)
+        quantizer1._calibrator.reset()
+
+        quantizer1(x_1)
+        quantizer1(x_2)
+
+        quantizer1.load_calib_amax("percentile", percentile=99.99)
+        test_utils.compare(quantizer1._calibrator.compute_amax("percentile", percentile=99.99),
+                           quantizer1.amax,
+                           atol=0,
+                           rtol=0,
+                           ctol=0)
+
+    def test_setters(self):
+        quantizer = tensor_quantizer.TensorQuantizer()
+        quantizer.num_bits = 7
+        quantizer.unsigned = True
+
+        assert quantizer.num_bits == 7
+        assert quantizer.unsigned
+
+    def test_pre_quant_scale(self):
+        quant_desc = tensor_quant.QuantDescriptor(axis=1, num_bits=8, amax=127.0)
+        quantizer = tensor_quantizer.TensorQuantizer(quant_desc).cuda()
+        quantizer2 = tensor_quantizer.TensorQuantizer(quant_desc).cuda()
+
+        inputs = torch.Tensor([[0, 0.4, 1.1, 2.0]]).cuda()
+        outputs_gt = torch.Tensor([[0, 0, 1, 2]]).cuda()
+        assert torch.allclose(quantizer(inputs), outputs_gt)
+
+        quantizer.pre_quant_scale = 2.0
+        outputs_gt = torch.Tensor([[0, 1, 2, 4]]).cuda()
+        assert torch.allclose(quantizer(inputs), outputs_gt)
+
+        quantizer2.pre_quant_scale = torch.Tensor([[1.0, 2.0, 3.0, 4.0]]).cuda()
+        outputs_gt = torch.Tensor([[0, 1, 3, 8]]).cuda()
+        assert torch.allclose(quantizer2(inputs), outputs_gt)
+
+    @pytest.mark.parametrize("E, M, axis", [(5, 2, None), (4, 3, None), (4, 3, 1), (7, 3, None)])
+    def test_e4m3(self, E, M, axis):
+        is_error_expected = (E != 4 or M != 3)
+        with (pytest.raises(TypeError)
+              if is_error_expected else contextlib.nullcontext()):
+            e4m3_desc = tensor_quant.QuantDescriptor(num_bits=(E, M), axis=axis)
+            e4m3_quantizer = tensor_quantizer.TensorQuantizer(e4m3_desc).to("cuda")
+
+            x = torch.rand(3, 63, 7, 7, device="cuda")
+
+            e4m3_x = e4m3_quantizer(x)
+            ref = tensor_quant.scaled_e4m3(x, e4m3_quantizer._get_amax(x), E, M)
+            test_utils.compare(e4m3_x, ref, atol=0, rtol=0)

tools/pytorch-quantization/tests/test_onnx_export.py

@@ -0,0 +1,111 @@
+#
+# 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.
+#
+"""Tests for ONNX export."""
+import io
+import onnxruntime
+import pytest
+import torch
+
+# ORT output correctness tests sometimes fails due to random seed.
+# It needs to be investigated closer
+torch.manual_seed(0)
+
+import tests.utils as test_utils
+import torch.nn as nn
+import pytorch_quantization
+from pytorch_quantization.nn import QuantLinear
+from pytorch_quantization.tensor_quant import QuantDescriptor
+
+
+class MyModel(nn.Module):
+    """Test model for ONNX export."""
+
+    def __init__(self, **kwargs):
+        super().__init__()
+        self.net = nn.Sequential(
+            QuantLinear(16, 32, **kwargs),
+            nn.ReLU(),
+            QuantLinear(32, 64, **kwargs),
+            nn.ReLU(),
+            QuantLinear(64, 16, **kwargs),
+        )
+
+    def forward(self, x):
+        return self.net(x)
+
+
+@pytest.mark.parametrize("num_bits, per_channel_quantization, constant_folding, dtype",
+                         [(8, True, True, torch.float32), (8, False, True, torch.float32),
+                          (8, True, False, torch.float32), (8, False, False, torch.float32),
+                          (8, False, False, torch.float16), (8, False, False, torch.bfloat16),
+                          ((4, 3), False, True, torch.float32), ((4, 3), False, False, torch.float32),
+                          ((4, 3), False, False, torch.float16), ((4, 3), False, False, torch.bfloat16)])
+def test_onnx_export(num_bits, per_channel_quantization, constant_folding, dtype, onnx_file_path=None):
+    quant_desc_input = QuantDescriptor(num_bits=num_bits, axis=None)
+    quant_desc_weight = QuantDescriptor(num_bits=num_bits, axis=0 if per_channel_quantization else None)
+
+    model = MyModel(quant_desc_input=quant_desc_input, quant_desc_weight=quant_desc_weight).cuda()
+    model.eval()
+
+    OPSET = 17
+    dummy_input = torch.randn(16, 16).cuda()
+    input_names = ["input"]
+    output_names = ["output"]
+
+    model = model.to(dtype)
+    dummy_input = dummy_input.to(dtype)
+
+    # Calibrate model
+    for name, module in model.named_modules():
+        if name.endswith('_quantizer'):
+            module.enable_calib()
+            module.disable_quant()
+    _ = model(dummy_input)
+    for name, module in model.named_modules():
+        if name.endswith('_quantizer'):
+            module.disable_calib()
+            module.load_calib_amax()
+            module.enable_quant()
+
+    f = io.BytesIO() if onnx_file_path is None else None
+
+    with pytorch_quantization.enable_onnx_export():
+        torch.onnx.export(
+            model,
+            dummy_input,
+            f=f if onnx_file_path is None else onnx_file_path,
+            opset_version=OPSET,
+            input_names=input_names,
+            output_names=output_names,
+            do_constant_folding=constant_folding,
+        )
+
+    # TODO: ort output correctness check for fp8
+    # ONNXRuntime does not seem to be supporting bf16 gemms
+    if num_bits == 8 and dtype != torch.bfloat16:
+        if f is not None:
+            f.seek(0)
+        ort_session = onnxruntime.InferenceSession(f.read() if onnx_file_path is None else onnx_file_path, providers=['CUDAExecutionProvider', 'CPUExecutionProvider'])
+        ort_result = ort_session.run([], {"input": dummy_input.cpu().numpy()})
+        ort_result = torch.tensor(ort_result[0]).cuda()
+        torch_result = model(dummy_input)
+        test_utils.compare(ort_result, torch_result, atol=1e-2, rtol=1e-2)
+
+
+if __name__ == "__main__":
+    test_onnx_export(8, False, False, torch.float16, "/tmp/test_fp16.onnx")
+    test_onnx_export(8, False, False, torch.bfloat16, "/tmp/test_bf16.onnx")

tools/pytorch-quantization/tests/utils.py

@@ -0,0 +1,129 @@
+#
+# 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.
+#
+
+
+"""Utils for testing quantization."""
+import numpy as np
+from scipy.spatial import distance
+
+import torch
+
+from pytorch_quantization import tensor_quant
+
+def quantize_by_range(x, num_bits):
+    """Quantize torch tensor by range to num_bits with symmetric zero-mean quantizer."""
+    amax = x.abs().max()
+    x_q = tensor_quant.fake_tensor_quant(x, amax, num_bits)
+    return x_q
+
+def quantize_by_range_fused(x_tuple, num_bits):
+    """Quantize multiple torch tensors by combined range to num_bits with symmetric zero-mean quantizer."""
+    # compute aggregate amax across all tensors
+    amax = max([x.abs().max() for x in x_tuple])
+    # quantize each tensor with the aggregate amax
+    x_q_tuple = tuple(tensor_quant.fake_tensor_quant(x, amax, num_bits) for x in x_tuple)
+    return x_q_tuple
+
+def copy_state_and_quantize(dst, src, num_bits):
+    """Copy src to dst, quantize all 'weight' entries to num_bits."""
+    src_state_dict = src.state_dict()
+    dst_state_dict = dict()
+    for key in src_state_dict:
+        if 'weight' in key:
+            dst_state_dict[key] = quantize_by_range(src_state_dict[key], num_bits)
+        else:
+            dst_state_dict[key] = src_state_dict[key].clone()
+
+    dst.load_state_dict(dst_state_dict)
+
+def copy_state_and_quantize_fused(dst, src, num_bits):
+    """Copy src to dst, quantize all 'weight' entries to num_bits using the aggregate amax."""
+    src_state_dict = src.state_dict()
+    dst_state_dict = dict()
+
+    # compute aggregate amax across all weight tensors
+    amax = 0
+    for key in src_state_dict:
+        if 'weight' in key:
+            amax = max(amax, src_state_dict[key].abs().max())
+
+    # quantize each weight tensor with the aggregate amax
+    for key in src_state_dict:
+        if 'weight' in key:
+            dst_state_dict[key] = tensor_quant.fake_tensor_quant(src_state_dict[key], amax, num_bits)
+        else:
+            dst_state_dict[key] = src_state_dict[key].clone()
+
+    dst.load_state_dict(dst_state_dict)
+
+def compare(a, b, rtol=1e-7, atol=1e-6, ctol=1e-6):
+    """Compare two tensors and raise AssertionError if their difference is outside of tolerance."""
+    if torch.isinf(a).any():
+        raise ValueError("a contains infs")
+    if torch.isinf(b).any():
+        raise ValueError("b contains infs")
+
+    a = a.detach().cpu().numpy().flatten()
+    b = b.detach().cpu().numpy().flatten()
+
+    # compare elements of a and b relative to the max value in b
+    # large fp32 values may cause quantization errors that propagate to small values
+    rel_diff = np.abs(a-b)/np.linalg.norm(b)
+    abs_diff = np.abs(a-b)
+    cos_diff = distance.cosine(a, b)
+    try:
+        if rel_diff.max() > rtol:
+            raise AssertionError("Tensor relative error > %.2e (%.2e)" % (rtol, rel_diff.max()))
+        if abs_diff.max() > atol:
+            raise AssertionError("Tensor absolute error > %.2e (%.2e)" % (atol, abs_diff.max()))
+        if cos_diff > ctol:
+            raise AssertionError("Tensor cosine distance > %.2e (%.2e)" % (ctol, cos_diff))
+        # np.testing.assert_allclose(a, b, rtol=rtol, atol=atol)
+        # np.testing.assert_array_almost_equal_nulp(a, b)
+    except AssertionError as e:
+        print('norm(a) =', np.linalg.norm(a))
+        print('norm(b) =', np.linalg.norm(b))
+        print('Largest relative difference = %.2e' % rel_diff.max())
+        idx = np.argmax(rel_diff)
+        print('a[%d] = %.10f' % (idx, a[idx]))
+        print('b[%d] = %.10f' % (idx, b[idx]))
+        print('Largest absolute difference = %.2e' % abs_diff.max())
+        idx = np.argmax(abs_diff)
+        print('a[%d] = %.10f' % (idx, a[idx]))
+        print('b[%d] = %.10f' % (idx, b[idx]))
+        print('Cosine distance = %.2e' % cos_diff)
+        raise e
+
+def assert_min_mse(a, b, tol=1e-20):
+    """Assert that the mean squared error between a and b is at least tol."""
+    a = a.detach().cpu().numpy()
+    b = b.detach().cpu().numpy()
+    mse = ((a-b)**2).mean()
+    if mse < tol:
+        raise AssertionError("MSE = %.2e < %.2e" % (mse, tol))
+
+def quant_np(x, amax, num_bits=8, fake=False, narrow_range=True):
+    """Quantize x using numpy."""
+    intmax = 2.0**(num_bits - 1) - 1
+    intmin = -intmax if narrow_range else -intmax - 1
+    scale = intmax / amax
+    x_q = np.round(np.clip(x * scale, intmin, intmax))
+
+    if fake:
+        x_q /= scale
+
+    return x_q