Generate layer-outputs of a pytorch model (original/quantsim)

Signed-off-by: Raj Gite <quic_rgite@quicinc.com>

TrainingExtensions/common/src/python/aimet_common/default_logging_config.json

@@ -124,6 +124,7 @@
             "handlers": ["console"],
             "propagate": false
         }
+
     },
 
     "root": {

TrainingExtensions/torch/src/python/aimet_torch/onnx_utils.py

@@ -61,7 +61,6 @@
 
 _logger = AimetLogger.get_area_logger(AimetLogger.LogAreas.Utils)
 
-
 # runs the second pass of markers for non-leaf torch module and updates names of onnx ops belonging to
 # non-leaf pytorch module
 update_all_onnx_nodes_name = True
@@ -317,6 +316,30 @@ class OnnxSaver:
     Utilities to save/load onnx models
     """
 
+    @classmethod
+    def create_onnx_model_with_pytorch_layer_names(cls, onnx_model_path: str, pytorch_model: torch.nn.Module,
+                                                   dummy_input: Union[torch.Tensor, Tuple, List], is_conditional=False,
+                                                   module_marker_map=None, onnx_export_args: Optional[Union[OnnxExportApiArgs, dict]] = None):
+        """
+        This utility does some pre-processing on the pytorch model and then uses it to obtain an equivalent onnx model
+        with node names same as that in pytorch model. Whatever pre-processing/post-processing steps to be done on
+        the resultant onnx model must be done here.
+
+        :param onnx_model_path: Path to the ONNX model file
+        :param pytorch_model: Equivalent PyTorch model instance
+        :param dummy_input: Dummy input to the model. Used to parse model graph.
+        :param is_conditional: True if model is a conditional model, False otherwise
+        :param module_marker_map: Maps module names to traced custom markers (only used for conditional models)
+        :param onnx_export_args:  override options for torch.onnx.export call
+        :return:
+        """
+        # Pre-processing pytorch model
+        for dropout_type in aimet_torch.utils.DROPOUT_TYPES:
+            aimet_torch.utils.replace_modules_of_type1_with_type2(pytorch_model, dropout_type, torch.nn.Identity)
+
+        # Obtaining equivalent onnx model
+        cls.set_node_names(onnx_model_path, pytorch_model, dummy_input, is_conditional, module_marker_map, onnx_export_args)
+
     @classmethod
     def set_node_names(cls, onnx_model_path: str, pytorch_model: torch.nn.Module,
                        dummy_input: Union[torch.Tensor, Tuple], is_conditional=False, module_marker_map=None,
@@ -497,7 +520,6 @@ def _map_onnx_nodes_to_pytorch_modules(cls, pt_model, dummy_input, onnx_model_pa
         cls._fix_param_names(onnx_model)
         cls._fix_initializer_names(onnx_model, pt_model)
 
-
         return onnx_model
 
     @classmethod

TrainingExtensions/torch/src/python/aimet_torch/quantsim.py

@@ -88,9 +88,6 @@
 
 SUPPORTED_KERNELS_ACTION = SupportedKernelsAction.warn_on_error
 
-DROPOUT_TYPES = (torch.nn.Dropout, torch.nn.Dropout2d, torch.nn.Dropout3d)
-
-
 def _get_encoding_by_quantizer(quantizer: Union[StaticGridTensorQuantizer, LearnedGridTensorQuantizer]) \
         -> Optional[Union[libpymo.TfEncoding, List[libpymo.TfEncoding]]]:
     """
@@ -394,9 +391,7 @@ def export(self, path: str, filename_prefix: str, dummy_input: Union[torch.Tenso
         model_path = os.path.join(path, model_filename)
 
         # Create a version of the model without any quantization ops
-        model_to_export = copy.deepcopy(self.model).cpu()
-        all_modules_in_model_to_export = [module for module in model_to_export.modules()]
-        self._remove_quantization_wrappers(model_to_export, all_modules_in_model_to_export)
+        model_to_export = QuantizationSimModel.get_original_model(self.model)
 
         torch.save(model_to_export, model_path)
 
@@ -430,7 +425,7 @@ def export_torch_script_model_and_encodings(path: str, filename_prefix: str,
                                                 dummy_input: Union[torch.Tensor, Tuple],
                                                 excluded_layer_names: List = None):
         """
-        This method exports  a onnx mode and the corresponding encodings
+        This method exports a torchscript mode and the corresponding encodings
 
         :param path: path where to store model pth and encodings
         :param filename_prefix: Prefix to use for filenames of the model pth and encodings files
@@ -440,12 +435,11 @@ def export_torch_script_model_and_encodings(path: str, filename_prefix: str,
         :param excluded_layer_names: List of names of layers that have been excluded from quantization.
         :return: None
         """
+        # Create torchscript model and obtain node to i/o tensor name map
+        ts_path = os.path.join(path, filename_prefix + '.torchscript.pth')
         with utils.in_eval_mode(original_model), torch.no_grad():
-            trace = torch.jit.trace(original_model, dummy_input)
-            ts_path = os.path.join(path, filename_prefix + '.torchscript.pth')
-            trace.save(ts_path)
+            torchscript_utils.create_torch_script_model(ts_path, original_model, dummy_input)
 
-            # reload the trace from the saved trace file
             trace = torch.jit.load(ts_path)
             torch_script_node_io_tensor_map, valid_param_set = \
                 torchscript_utils.get_node_to_io_tensor_names_map(original_model, trace, dummy_input)
@@ -481,16 +475,10 @@ def export_onnx_model_and_encodings(path: str, filename_prefix: str, original_mo
 
         """
         # pylint: disable=too-many-locals
-        if module_marker_map is None:
-            module_marker_map = {}
-        # Save model to onnx
+        # Create onnx model and obtain node to i/o tensor name map
         onnx_path = os.path.join(path, filename_prefix + '.onnx')
-
-        for dropout_type in DROPOUT_TYPES:
-            utils.replace_modules_of_type1_with_type2(original_model, dropout_type, torch.nn.Identity)
-
-        OnnxSaver.set_node_names(onnx_path, original_model, dummy_input, is_conditional, module_marker_map,
-                                 onnx_export_args)
+        OnnxSaver.create_onnx_model_with_pytorch_layer_names(onnx_path, original_model, dummy_input, is_conditional,
+                                                             module_marker_map, onnx_export_args)
 
         onnx_model = onnx.load(onnx_path)
         onnx_node_to_io_tensor_map, valid_param_set = OnnxSaver.get_onnx_node_to_io_tensor_names_map(onnx_model)
@@ -745,7 +733,7 @@ def _export_encodings_to_files(sim_model: torch.nn.Module, path: str, filename_p
             # TODO: specifically call out dropout layers here since they are specifically switched out during export.
             # These ops should eventually be reworked as part of math invariant ops to ignore quantization altogether.
             # pylint: disable=protected-access
-            if isinstance(layer, QcQuantizeWrapper) and isinstance(layer._module_to_wrap, DROPOUT_TYPES):
+            if isinstance(layer, QcQuantizeWrapper) and isinstance(layer._module_to_wrap, utils.DROPOUT_TYPES):
                 continue
 
             if layer_name not in layers_in_io_tensor:
@@ -1245,6 +1233,17 @@ def _remove_quantization_wrappers(cls, starting_module, list_of_modules_to_exclu
             if not utils.is_leaf_module(module_ref):
                 cls._remove_quantization_wrappers(module_ref, list_of_modules_to_exclude)
 
+    @staticmethod
+    def get_original_model(model: torch.nn.Module):
+        """
+        This function returns the model with all quantization wrappers removed.
+        :return: Model without quantization wrappers.
+        """
+        original_model = copy.deepcopy(model)
+        all_modules_in_original_model = [module for module in original_model.modules()]
+        QuantizationSimModel._remove_quantization_wrappers(original_model, all_modules_in_original_model)
+        return original_model
+
     def _add_inputs_hook(self, hooks):
         module_to_name_map = {}
         for name, module in self.model.named_modules():
@@ -1570,7 +1569,6 @@ def _validate_torchquantizer(quant_sim_model):
             OnnxSaver._export_model_to_onnx(quant_sim_model, dummy_input, model_path, is_conditional, onnx_export_args) # pylint: disable=protected-access
 
 
-
 def save_checkpoint(quant_sim_model: QuantizationSimModel, file_path: str):
     """
     This API provides a way for the user to save a checkpoint of the quantized model which can

TrainingExtensions/torch/src/python/aimet_torch/torchscript_utils.py

@@ -37,7 +37,7 @@
 # =============================================================================
 
 """ Implementation for simulating models running on Quantized hardware """
-from typing import List, Union, Dict
+from typing import List, Union, Dict, Tuple
 import torch
 
 from aimet_common.utils import AimetLogger
@@ -257,3 +257,17 @@ def forward_hook(curr_module: torch.nn.Module, *_):
     #assert index == len(modules)
 
     return node_to_io_tensor_name_map, valid_param_set
+
+
+def create_torch_script_model(ts_path: str, original_model: torch.nn.Module, dummy_input: Union[torch.Tensor, Tuple]):
+    """
+    This utility obtains an equivalent torchscript model for the given pytorch model. Whatever pre-processing/post-processing
+    steps to be done on the resultant torchscript model must be done here.
+
+    :param ts_path: Path to the torchscript model file
+    :param original_model: Equivalent PyTorch model instance
+    :param dummy_input: Dummy input to the model. Used to parse model graph
+    :return:
+    """
+    trace = torch.jit.trace(original_model, dummy_input)
+    trace.save(ts_path)

TrainingExtensions/torch/src/python/aimet_torch/utils.py

@@ -58,7 +58,7 @@
 dtypes_to_ignore_for_quantization = (int, float, bool, str, tuple, type(None))
 torch_dtypes_to_ignore_for_quantization = [torch.int, torch.int8, torch.int16, torch.int32, torch.int64, torch.bool]
 allowed_output_types = (torch.Tensor, *dtypes_to_ignore_for_quantization)
-
+DROPOUT_TYPES = (torch.nn.Dropout, torch.nn.Dropout2d, torch.nn.Dropout3d)
 
 class IterFirstX:
     """ Iterator for the first x samples in a given data-loader """