enhance export API and add 3 recipes

Signed-off-by: Xin He <[email protected]>

Author: xin3he

neural_compressor/adaptor/pytorch.py

@@ -2738,6 +2738,8 @@ def quantize(self, tune_cfg, model, dataloader, q_func=None):
             # q_func can be created by neural_compressor internal or passed by user. It's critical to
             # distinguish how q_func is passed since neural_compressor built-in functions accept
             # neural_compressor model and user defined func should accept framework model.
+            # For export API
+            hook_list = torch_utils.util._set_input_scale_hook(q_model._model, op_cfgs)
             q_model._model = q_func(
                 q_model if getattr(q_func, 'builtin', None) else q_model._model)
             assert q_model._model is not None, "Please return a trained model in train function!"
@@ -2766,6 +2768,8 @@ def quantize(self, tune_cfg, model, dataloader, q_func=None):
                                                     prefix='',
                                                     example_inputs=example_inputs)
             if self.approach in ['post_training_static_quant', 'post_training_auto_quant']:
+                # For export API
+                hook_list = torch_utils.util._set_input_scale_hook(q_model._model, op_cfgs)
                 iterations = tune_cfg.get('calib_iteration', 1)
                 if q_func is not None:
                     q_func(q_model._model)
@@ -2774,6 +2778,11 @@ def quantize(self, tune_cfg, model, dataloader, q_func=None):
                                            dataloader,
                                            iterations,
                                            calib_sampling_size=tune_cfg.get('calib_sampling_size', 1))
+
+        if self.approach != 'post_training_dynamic_quant':
+            # For export API
+            scale_info = torch_utils.util._get_input_scale(q_model._model, hook_list)
+
         if self.sub_module_list is None:
             if self.version > Version("1.12.1"):  # pragma: no cover
                 # pylint: disable=E1123
@@ -2796,6 +2805,7 @@ def quantize(self, tune_cfg, model, dataloader, q_func=None):
         q_model.q_config = copy.deepcopy(self.tune_cfg)
         if self.approach != 'post_training_dynamic_quant':
             self._get_scale_zeropoint(q_model._model, q_model.q_config)
+            q_model.q_config['scale_info'] = scale_info
 
         self._dump_model_op_stats(q_model._model, q_model.q_config, self.approach)
         torch_utils.util.get_embedding_contiguous(q_model._model)
@@ -2873,10 +2883,11 @@ def _pre_hook_for_qat(self, dataloader=None):
                 quantized_ops[op[0]] = torch.quantization.default_dynamic_qconfig
             else:
                 quantized_ops[op[0]] = q_cfgs
-        # build for fetching scale and zeropoint
+        # build op_config_dict to save module scale and zeropoint
         op_config_dict = {}
         for op in quantizable_ops:
             op_config_dict[op] = {'weight': {'dtype': 'int8'}, 'activation': {'dtype': 'uint8'}}
+
         if self.version.release < Version("1.11.0").release:
             quantized_ops["default_qconfig"] = None
         else:
@@ -2928,8 +2939,14 @@ def _pre_hook_for_qat(self, dataloader=None):
             'sub_module_list': self.sub_module_list,
             'approach': 'quant_aware_training'
         }
+        # For export API
+        global hook_list
+        hook_list = torch_utils.util._set_input_scale_hook(self.model._model, quantized_ops)
 
     def _post_hook_for_qat(self):
+        # For export API
+        scale_info = torch_utils.util._get_input_scale(self.model._model, hook_list)
+        self.model.q_config['scale_info'] = scale_info
         from torch.quantization.quantize_fx import convert_fx
         if self.sub_module_list is None:
             if self.version > Version("1.12.1"):  # pragma: no cover

neural_compressor/adaptor/torch_utils/util.py

@@ -44,6 +44,125 @@ def contiguous_hook(module, input):
             child.register_forward_pre_hook(contiguous_hook)
 
 
+def is_fused_module(module):
+    """This is a helper function for `_propagate_qconfig_helper` to detecte
+        if this module is fused.
+
+    Args:
+        module (object): input module
+
+    Returns:
+        (bool): is fused or not
+    """
+    op_type = str(type(module))
+    if 'fused' in op_type:
+        return True
+    else:
+        return False
+
+
+def _set_input_scale_hook(model, op_cfgs):
+    """Insert hooks to observer input scale and zeropoint.
+
+    Args:
+        model (object): input model
+        op_cfgs (dict): dictionary of quantization configure for each op
+
+    Returns:
+        hook_list (list): input observer hooks
+    """
+    def input_scale_hook(module, input):
+        module.input_observer = module.qconfig.activation()
+        module.input_observer(input[0])
+        return input
+
+    def output_scale_hook(module, input, output):
+        module.output_observer = module.qconfig.activation()
+        module.output_observer(output)
+        return output
+
+    def ConvReLU2d_scale_hook(module, input):
+        module.input_observer = module.qconfig.activation()
+        module.input_observer(input[0])
+        output = module._conv_forward(input[0], module.weight_fake_quant(module.weight), module.bias)
+        module.output_observer = module.qconfig.activation()
+        module.output_observer(output)
+        return input
+
+    def LinearReLU_scale_hook(module, input):
+        import torch.nn.functional as F
+        module.input_observer = module.qconfig.activation()
+        module.input_observer(input[0])
+        output = F.linear(input[0], module.weight_fake_quant(module.weight), module.bias)
+        module.output_observer = module.qconfig.activation()
+        module.output_observer(output)
+        return input
+
+    hook_list = []
+    for name, module in model.named_modules():
+        if 'Conv' in str(module.__class__.__name__) or \
+          'Linear' in str(module.__class__.__name__):
+            if not hasattr(module, 'qconfig') or not module.qconfig:
+                continue
+            from torch.nn.intrinsic.qat import ConvBn2d, ConvReLU2d, ConvBnReLU2d, LinearReLU
+            if type(module) in [ConvBn2d, ConvBnReLU2d]:
+                handle_in = module.register_forward_pre_hook(input_scale_hook)
+                # module[0] == torch.nn.BatchNorm2d
+                module[0].qconfig = module.qconfig
+                handle_out = module[0].register_forward_hook(output_scale_hook)
+                hook_list.extend([handle_in, handle_out])
+            elif type(module) in [ConvReLU2d]:
+                handle_in_out = module.register_forward_pre_hook(ConvReLU2d_scale_hook)
+                hook_list.extend([handle_in_out])
+            elif type(module) in [LinearReLU]:
+                handle_in_out = module.register_forward_pre_hook(LinearReLU_scale_hook)
+                hook_list.extend([handle_in_out])
+            else:
+                if is_fused_module(module):
+                    continue
+                handle_in = module.register_forward_pre_hook(input_scale_hook)
+                handle_out = module.register_forward_hook(output_scale_hook)
+                hook_list.extend([handle_in, handle_out])
+    return hook_list
+
+
+def _get_input_scale(model, hook_list):
+    """Fetch input scale and zeropoint from observer.
+
+    Args:
+        model (object): input model
+        hook_list (list): input observer hooks
+
+    Returns:
+        input_scale_info (dict): input scale and zero_point of each modules
+    """
+    scale_info = {}
+    for name, module in model.named_modules():
+        from torch.nn.intrinsic.qat import ConvBn2d, ConvBnReLU2d
+        if type(module) in [ConvBn2d, ConvBnReLU2d]:
+            if hasattr(module, "input_observer") and hasattr(module[0], "output_observer"):
+                scale_in, zero_point_in = module.input_observer.calculate_qparams()
+                scale_out, zero_point_out = module[0].output_observer.calculate_qparams()
+                scale_info[name] = {
+                    'input_scale': float(scale_in),
+                    'input_zeropoint': int(zero_point_in),
+                    'output_scale': float(scale_out),
+                    'output_zeropoint': int(zero_point_out)
+                }
+        elif hasattr(module, "input_observer") and hasattr(module, "output_observer"):
+            scale_in, zero_point_in = module.input_observer.calculate_qparams()
+            scale_out, zero_point_out = module.output_observer.calculate_qparams()
+            scale_info[name] = {
+                'input_scale': float(scale_in),
+                'input_zeropoint': int(zero_point_in),
+                'output_scale': float(scale_out),
+                'output_zeropoint': int(zero_point_out)
+            }
+    for h in hook_list:
+        h.remove()
+    return scale_info
+
+
 def collate_torch_preds(results):
     batch = results[0]
     if isinstance(batch, list):

neural_compressor/experimental/__init__.py

@@ -27,6 +27,8 @@
 from .model_conversion import ModelConversion
 from .distillation import Distillation
 from .nas import NAS
+from . import export
 
 __all__ = ['Component', 'Quantization', 'Pruning', 'Benchmark', 'Graph_Optimization', \
-           'GraphOptimization', 'ModelConversion', 'Distillation', 'NAS', 'MixedPrecision']
+           'GraphOptimization', 'ModelConversion', 'Distillation', 'NAS', 'MixedPrecision', \
+           'export']

neural_compressor/experimental/export/__init__.py

@@ -0,0 +1,21 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+#
+# Copyright (c) 2021 Intel 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.
+
+
+"""Intel Neural Compressor Export."""
+
+from .torch2onnx import torch_to_fp32_onnx, torch_to_int8_onnx