new API for export

neural_compressor/adaptor/pytorch.py

@@ -2851,6 +2851,10 @@ 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 
+        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:
@@ -2892,10 +2896,12 @@ def _pre_hook_for_qat(self, dataloader=None):
                                                 example_inputs=example_inputs)
         # This is a flag for reloading
         self.model.q_config = {
+            'calib_sampling_size': 100, # tmp arg for export API
             'is_oneshot': True,
             'framework': 'pytorch_fx',
             'reduce_range': REDUCE_RANGE,
             'quantizable_ops': quantizable_ops,
+            'op': op_config_dict,
             'sub_module_list': self.sub_module_list,
             'approach': 'quant_aware_training'
         }
@@ -2918,6 +2924,11 @@ def _post_hook_for_qat(self):
             PyTorch_FXAdaptor.convert_sub_graph(self.sub_module_list, \
                                                 self.model._model, prefix='')
 
+        if self.approach != 'post_training_dynamic_quant':
+            self._get_scale_zeropoint(self.model._model, self.model.q_config)
+        self._dump_model_op_stats(self.model._model, self.model.q_config, self.approach)
+        torch_utils.util.get_embedding_contiguous(self.model._model)
+
     def train(self, model, dataloader, optimizer_tuple, criterion_tuple, hooks, **kwargs):
         """Execute the train process on the specified model.
 

neural_compressor/adaptor/torch_utils/onnx.py

@@ -30,17 +30,30 @@ def __init__(self, dataloader, sample_size=100):
         self.datasize = self.batch_num * self.batch_size
 
         self.data = []
-        for i, (input, label) in enumerate(self.dataloader):
-            if i * self.batch_size >= self.datasize:
-                break
-            if isinstance(input, dict) or isinstance(input, UserDict):
-                batch = {k: v.detach().cpu().numpy() for k, v in input.items()}
-            elif isinstance(input, list) or isinstance(input, tuple):
-                batch = {'input': [v.detach().cpu().numpy() for v in input]}
-            else:
-                batch = {'input': input.detach().cpu().numpy()}
-            self.data.append(batch)
-        self.data = iter(self.data)
+        try:
+            for i, (input, label) in enumerate(self.dataloader):
+                if i * self.batch_size >= self.datasize:
+                    break
+                if isinstance(input, dict) or isinstance(input, UserDict):
+                    batch = {k: v.detach().cpu().numpy() for k, v in input.items()}
+                elif isinstance(input, list) or isinstance(input, tuple):
+                    batch = {'input': [v.detach().cpu().numpy() for v in input]}
+                else:
+                    batch = {'input': input.detach().cpu().numpy()}
+                self.data.append(batch)
+            self.data = iter(self.data)
+        except:
+            for i, input in enumerate(self.dataloader):
+                if i * self.batch_size >= self.datasize:
+                    break
+                if isinstance(input, dict) or isinstance(input, UserDict):
+                    batch = {k: v.detach().cpu().numpy() for k, v in input.items()}
+                elif isinstance(input, list) or isinstance(input, tuple):
+                    batch = {'input': [v.detach().cpu().numpy() for v in input]}
+                else:
+                    batch = {'input': input.detach().cpu().numpy()}
+                self.data.append(batch)
+            self.data = iter(self.data)
 
     def get_next(self):
         return next(self.data, None)

neural_compressor/config.py

@@ -717,21 +717,19 @@ def __init__(
         self,
         dtype="int8",
         opset_version=14,
-        quant_mode="'QDQ'",
-        sample_inputs=None,
+        quant_format="'QDQ'",
+        example_inputs=None,
         input_names=None,
         output_names=None,
         dynamic_axes=None,
-        **kwargs,
     ):
         self._dtype = dtype
         self._opset_version = opset_version
-        self._quant_mode = quant_mode
-        self._sample_inputs = sample_inputs
+        self._quant_format = quant_format
+        self._example_inputs = example_inputs
         self._input_names = input_names
         self._output_names = output_names
         self._dynamic_axes = dynamic_axes
-        self._kwargs = kwargs
 
     @property
     def dtype(self):
@@ -750,20 +748,20 @@ def opset_version(self, opset_version):
         self._opset_version = opset_version
 
     @property
-    def quant_mode(self):
-        return self._quant_mode
+    def quant_format(self):
+        return self._quant_format
 
-    @quant_mode.setter
-    def quant_mode(self, quant_mode):
-        self._quant_mode = quant_mode
+    @quant_format.setter
+    def quant_format(self, quant_format):
+        self._quant_format = quant_format
 
     @property
-    def sample_inputs(self):
-        return self._sample_inputs
+    def example_inputs(self):
+        return self._example_inputs
 
-    @sample_inputs.setter
-    def sample_inputs(self, sample_inputs):
-        self._sample_inputs = sample_inputs
+    @example_inputs.setter
+    def example_inputs(self, example_inputs):
+        self._example_inputs = example_inputs
 
     @property
     def input_names(self):
@@ -783,7 +781,7 @@ def output_names(self, output_names):
 
     @property
     def dynamic_axes(self):
-        return self._output_names
+        return self._dynamic_axes
 
     @dynamic_axes.setter
     def dynamic_axes(self, dynamic_axes):
@@ -795,8 +793,8 @@ def __init__(
         self,
         dtype="int8",
         opset_version=14,
-        quant_mode="'QDQ'",
-        sample_inputs=None,
+        quant_format="'QDQ'",
+        example_inputs=None,
         input_names=None,
         output_names=None,
         dynamic_axes=None,
@@ -805,22 +803,22 @@ def __init__(
         super().__init__(
             dtype=dtype,
             opset_version=opset_version,
-            quant_mode=quant_mode,
-            sample_inputs=sample_inputs,
+            quant_format=quant_format,
+            example_inputs=example_inputs,
             input_names=input_names,
             output_names=output_names,
             dynamic_axes=dynamic_axes,
-            kwargs=kwargs,
         )
+        self.kwargs = kwargs
 
 
 class TF2ONNXConfig(ExportConfig):
      def __init__(
         self,
         dtype="int8",
         opset_version=14,
-        quant_mode="'QDQ'",
-        sample_inputs=None,
+        quant_format="'QDQ'",
+        example_inputs=None,
         input_names=None,
         output_names=None,
         dynamic_axes=None,
@@ -829,13 +827,13 @@ def __init__(
         super().__init__(
             dtype=dtype,
             opset_version=opset_version,
-            quant_mode=quant_mode,
-            sample_inputs=sample_inputs,
+            quant_format=quant_format,
+            example_inputs=example_inputs,
             input_names=input_names,
             output_names=output_names,
             dynamic_axes=dynamic_axes,
-            kwargs=kwargs,
         )
+        self.kwargs = kwargs
 
 
 def set_random_seed(seed: int):

neural_compressor/model/torch_model.py

@@ -20,7 +20,6 @@
 import inspect
 import sys
 from collections import OrderedDict, UserDict
-from abc import abstractmethod
 from ..adaptor.torch_utils.util import input2tuple
 from neural_compressor.utils.utility import LazyImport, compute_sparsity
 from neural_compressor.utils import logger
@@ -46,8 +45,41 @@ def __init__(self, model, **kwargs):
         self.q_config = None
         self._workspace_path = ''
         self.is_quantized = False
+        try:
+            self.fp32_model = copy.deepcopy(model)
+        except Exception as e:  # pragma: no cover
+            logger.warning("Fail to deep copy the model due to {}, inplace is used now.".format(
+                repr(e)))
+            self.fp32_model = model
         self.kwargs = kwargs if kwargs else None
 
+    def __repr__(self):
+        # rewirte this func to avoid printing fp32_model
+        from torch.nn.modules.module import _addindent
+        # We treat the extra repr like the sub-module, one item per line
+        extra_lines = []
+        extra_repr = self.extra_repr()
+        # empty string will be split into list ['']
+        if extra_repr:
+            extra_lines = extra_repr.split('\n')
+        child_lines = []
+        for key, module in self._modules.items():
+            if key == 'fp32_model':
+                continue
+            mod_str = repr(module)
+            mod_str = _addindent(mod_str, 2)
+            child_lines.append('(' + key + '): ' + mod_str)
+        lines = extra_lines + child_lines
+        main_str = self._get_name() + '('
+        if lines:
+            # simple one-liner info, which most builtin Modules will use
+            if len(extra_lines) == 1 and not child_lines:
+                main_str += extra_lines[0]
+            else:
+                main_str += '\n  ' + '\n  '.join(lines) + '\n'
+        main_str += ')'
+        return main_str
+
     def forward(self, *args, **kwargs):
         return self._model(*args, **kwargs)
 
@@ -355,13 +387,18 @@ def export_to_fp32_onnx(
         opset_version=14,
         dynamic_axes={"input": {0: "batch_size"},
                       "output": {0: "batch_size"}},
+        input_names=None,
+        output_names=None,
         do_constant_folding=True,
         verbose=True,
         fp32_model=None,
     ):
-        example_input_names = ['input']
-        if isinstance(example_inputs, dict) or isinstance(example_inputs, UserDict):
-            example_input_names = list(example_inputs.keys())
+        if input_names:
+            example_input_names = input_names
+        else:
+            example_input_names = ['input']
+            if isinstance(example_inputs, dict) or isinstance(example_inputs, UserDict):
+                example_input_names = list(example_inputs.keys())
         model = self.model
         if fp32_model:
             model = fp32_model
@@ -371,6 +408,7 @@ def export_to_fp32_onnx(
             save_path,
             opset_version=opset_version,
             input_names=example_input_names,
+            output_names=output_names,
             dynamic_axes=dynamic_axes,
             do_constant_folding=do_constant_folding,
         )
@@ -386,6 +424,8 @@ def export_to_bf16_onnx(self,
         opset_version=14, 
         dynamic_axes={"input": {0: "batch_size"},
                       "output": {0: "batch_size"}},
+        input_names=None,
+        output_names=None,
         do_constant_folding=True,
         verbose=True,
     ):
@@ -395,6 +435,8 @@ def export_to_bf16_onnx(self,
             example_inputs = example_inputs,
             opset_version=opset_version,
             dynamic_axes=dynamic_axes,
+            input_names=input_names,
+            output_names=output_names,
             do_constant_folding=do_constant_folding,
             verbose=False,
         )
@@ -437,6 +479,8 @@ def export_to_int8_onnx(
         opset_version=14,
         dynamic_axes={"input": {0: "batch_size"},
                     "output": {0: "batch_size"}},
+        input_names=None,
+        output_names=None,
         do_constant_folding=True,
         quant_format='QDQ',
         dtype='S8S8',
@@ -465,22 +509,13 @@ def export_to_int8_onnx(
             "No quantization configuration found, " + \
             "please use the model generated by INC quantizer"
         if 'dynamic' in self.q_config['approach']:
-            op_types_to_quantize=['MatMul', 'Gather', "LSTM", 'Conv']
-            pytorch_op_types_to_quantize=['Linear', 'Embedding', "LSTM", 
-                                            'Conv1d', 'Conv2d']
-            addition_op_to_quantize = list(ortq.registry.IntegerOpsRegistry.keys())
+            op_types_to_quantize=['MatMul', 'Gather', "LSTM"]
+            pytorch_op_types_to_quantize=['Linear', 'Embedding', "LSTM"]
+            addition_op_to_quantize = []
         else:
             op_types_to_quantize=['MatMul', 'Gather', 'Conv']
             pytorch_op_types_to_quantize=['Linear', 'Embedding', 'Conv1d', 'Conv2d']
-            if quant_format == 'QDQ':
-                addition_op_to_quantize = list(ortq.registry.QDQRegistry.keys())
-                addition_op_to_quantize.remove('Relu') # ValueError: x not in list
-            else:
-                addition_op_to_quantize = list(ortq.registry.QLinearOpsRegistry.keys())
-
-        if 'U8S8' in dtype:
-            op_types_to_quantize.remove('Gather')
-            pytorch_op_types_to_quantize.remove('Embedding')
+            addition_op_to_quantize = []
 
         if quant_format == 'QDQ' and opset_version < 13:   # pragma: no cover 
             opset_version = 13
@@ -495,6 +530,8 @@ def export_to_int8_onnx(
             example_inputs = example_inputs,
             opset_version=opset_version,
             dynamic_axes=dynamic_axes,
+            input_names=input_names,
+            output_names=output_names,
             do_constant_folding=do_constant_folding,
             verbose=False,
             fp32_model=fp32_model
@@ -624,9 +661,33 @@ def export(
         save_path: str,
         conf,
     ):
-        # TODO
-        from neural_compressor.config import Torch2ONNXConfig
-        pass
+        if conf.dtype == 'int8':
+            calib_dataloader = conf.kwargs.pop("calib_dataloader", None)
+            self.export_to_int8_onnx(
+                save_path=save_path,
+                example_inputs=conf.example_inputs,
+                opset_version=conf.opset_version,
+                dynamic_axes=conf.dynamic_axes,
+                input_names=conf.input_names,
+                output_names=conf.output_names,
+                quant_format=conf.quant_format,
+                dtype='U8S8',
+                fp32_model=self.fp32_model,
+                calib_dataloader=calib_dataloader,
+            )
+        elif conf.dtype == 'fp32':
+            self.export_to_fp32_onnx(
+                save_path=save_path,
+                example_inputs=conf.example_inputs,
+                opset_version=conf.opset_version,
+                dynamic_axes=conf.dynamic_axes,
+                input_names=conf.input_names,
+                output_names=conf.output_names,
+                verbose=True,
+                fp32_model=self.fp32_model,
+            )
+        else:   # pragma: no cover
+            assert False, "Not allowed dtype: {}, pleas use 'fp32' or 'int8'.".format(conf.dtype)
 
 
 class PyTorchFXModel(PyTorchModel):