chore: import upstream snapshot with attribution

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