wehub-resource-sync
182 lines
7.8 KiB
Python
182 lines
7.8 KiB
Python
#
|
|
# 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 |