pytorch_quantization.quant_modules Namespace Reference

Classes
class	QuantModuleReplacementHelper

Functions
def	initialize (float_module_list=None, custom_quant_modules=None)
def	deactivate ()

Variables
	_quant_entry = namedtuple('quant_entry', 'orig_mod mod_name replace_mod')
list	_DEFAULT_QUANT_MAP
	_quant_module_helper_object = QuantModuleReplacementHelper()

Function Documentation

initialize()

def pytorch_quantization.quant_modules.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()

deactivate()

def pytorch_quantization.quant_modules.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.

Variable Documentation

_quant_entry

pytorch_quantization.quant_modules._quant_entry = namedtuple('quant_entry', 'orig_mod mod_name replace_mod')

private

_DEFAULT_QUANT_MAP

list pytorch_quantization.quant_modules._DEFAULT_QUANT_MAP

private

Initial value:

=  [_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_module_helper_object

pytorch_quantization.quant_modules._quant_module_helper_object = QuantModuleReplacementHelper()

private