mct-nightly 2.0.0.20240509.406__tar.gz → 2.0.0.20240510.421__tar.gz

{mct-nightly-2.0.0.20240509.406 → mct-nightly-2.0.0.20240510.421}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.0.0.20240509.406
+Version: 2.0.0.20240510.421
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN
@@ -47,9 +47,9 @@ Description: # Model Compression Toolkit (MCT)
         Explore the Model Compression Toolkit (MCT) through our tutorials, 
         covering compression techniques for Keras and PyTorch models. Access interactive [notebooks](tutorials/README.md) 
         for hands-on learning. For example:
-        * [Keras MobileNetV2 post training quantization](tutorials/notebooks/keras/ptq/example_keras_imagenet.ipynb)
-        * [Post training quantization with PyTorch](tutorials/notebooks/pytorch/ptq/example_pytorch_quantization_mnist.ipynb)
-        * [Data Generation for ResNet18 with PyTorch](tutorials/notebooks/pytorch/data_generation/example_pytorch_data_generation.ipynb).
+        * [Keras MobileNetV2 post training quantization](tutorials/notebooks/imx500_notebooks/keras/example_keras_mobilenetv2_for_imx500.ipynb)
+        * [Post training quantization with PyTorch](tutorials/notebooks/mct_features_notebooks/pytorch/example_pytorch_ptq_mnist.ipynb)
+        * [Data Generation for ResNet18 with PyTorch](tutorials/notebooks/mct_features_notebooks/pytorch/example_pytorch_data_generation.ipynb).
         
         
         ### Supported Versions
@@ -152,8 +152,6 @@ Description: # Model Compression Toolkit (MCT)
         | ResNet-18 [3]             | 69.86           | 69.63           |69.53|
         | SqueezeNet 1.1 [3]        | 58.128          | 57.678          ||
         
-        For more results, please refer to [quick start](https://github.com/sony/model_optimization/tree/main/tutorials/quick_start).
-        
         
         #### Pruning Results
         

{mct-nightly-2.0.0.20240509.406 → mct-nightly-2.0.0.20240510.421}/README.md

@@ -41,9 +41,9 @@ For installing the nightly version or installing from source, refer to the [inst
 Explore the Model Compression Toolkit (MCT) through our tutorials, 
 covering compression techniques for Keras and PyTorch models. Access interactive [notebooks](tutorials/README.md) 
 for hands-on learning. For example:
-* [Keras MobileNetV2 post training quantization](tutorials/notebooks/keras/ptq/example_keras_imagenet.ipynb)
-* [Post training quantization with PyTorch](tutorials/notebooks/pytorch/ptq/example_pytorch_quantization_mnist.ipynb)
-* [Data Generation for ResNet18 with PyTorch](tutorials/notebooks/pytorch/data_generation/example_pytorch_data_generation.ipynb).
+* [Keras MobileNetV2 post training quantization](tutorials/notebooks/imx500_notebooks/keras/example_keras_mobilenetv2_for_imx500.ipynb)
+* [Post training quantization with PyTorch](tutorials/notebooks/mct_features_notebooks/pytorch/example_pytorch_ptq_mnist.ipynb)
+* [Data Generation for ResNet18 with PyTorch](tutorials/notebooks/mct_features_notebooks/pytorch/example_pytorch_data_generation.ipynb).
 
 
 ### Supported Versions
@@ -146,8 +146,6 @@ In the following table we present the ImageNet validation results for these mode
 | ResNet-18 [3]             | 69.86           | 69.63           |69.53|
 | SqueezeNet 1.1 [3]        | 58.128          | 57.678          ||
 
-For more results, please refer to [quick start](https://github.com/sony/model_optimization/tree/main/tutorials/quick_start).
-
 
 #### Pruning Results
 

{mct-nightly-2.0.0.20240509.406 → mct-nightly-2.0.0.20240510.421}/mct_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.0.0.20240509.406
+Version: 2.0.0.20240510.421
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN
@@ -47,9 +47,9 @@ Description: # Model Compression Toolkit (MCT)
         Explore the Model Compression Toolkit (MCT) through our tutorials, 
         covering compression techniques for Keras and PyTorch models. Access interactive [notebooks](tutorials/README.md) 
         for hands-on learning. For example:
-        * [Keras MobileNetV2 post training quantization](tutorials/notebooks/keras/ptq/example_keras_imagenet.ipynb)
-        * [Post training quantization with PyTorch](tutorials/notebooks/pytorch/ptq/example_pytorch_quantization_mnist.ipynb)
-        * [Data Generation for ResNet18 with PyTorch](tutorials/notebooks/pytorch/data_generation/example_pytorch_data_generation.ipynb).
+        * [Keras MobileNetV2 post training quantization](tutorials/notebooks/imx500_notebooks/keras/example_keras_mobilenetv2_for_imx500.ipynb)
+        * [Post training quantization with PyTorch](tutorials/notebooks/mct_features_notebooks/pytorch/example_pytorch_ptq_mnist.ipynb)
+        * [Data Generation for ResNet18 with PyTorch](tutorials/notebooks/mct_features_notebooks/pytorch/example_pytorch_data_generation.ipynb).
         
         
         ### Supported Versions
@@ -152,8 +152,6 @@ Description: # Model Compression Toolkit (MCT)
         | ResNet-18 [3]             | 69.86           | 69.63           |69.53|
         | SqueezeNet 1.1 [3]        | 58.128          | 57.678          ||
         
-        For more results, please refer to [quick start](https://github.com/sony/model_optimization/tree/main/tutorials/quick_start).
-        
         
         #### Pruning Results
         

{mct-nightly-2.0.0.20240509.406 → mct-nightly-2.0.0.20240510.421}/model_compression_toolkit/__init__.py

@@ -27,4 +27,4 @@ from model_compression_toolkit import data_generation
 from model_compression_toolkit import pruning
 from model_compression_toolkit.trainable_infrastructure.keras.load_model import keras_load_quantized_model
 
-__version__ = "2.0.0.20240509.000406"
+__version__ = "2.0.0.20240510.000421"

{mct-nightly-2.0.0.20240509.406 → mct-nightly-2.0.0.20240510.421}/model_compression_toolkit/constants.py

@@ -20,6 +20,7 @@ TENSORFLOW = 'tensorflow'
 PYTORCH = 'pytorch'
 FOUND_TF = importlib.util.find_spec(TENSORFLOW) is not None
 FOUND_TORCH = importlib.util.find_spec("torch") is not None
+FOUND_TORCHVISION = importlib.util.find_spec("torchvision") is not None
 FOUND_ONNX = importlib.util.find_spec("onnx") is not None
 FOUND_ONNXRUNTIME = importlib.util.find_spec("onnxruntime") is not None
 FOUND_SONY_CUSTOM_LAYERS = importlib.util.find_spec('sony_custom_layers') is not None

{mct-nightly-2.0.0.20240509.406 → mct-nightly-2.0.0.20240510.421}/model_compression_toolkit/core/common/visualization/tensorboard_writer.py

@@ -75,7 +75,7 @@ def get_node_properties(node_dict_to_log: dict,
         for output_shape in output_shapes:  # create protobuf for each output shape
             proto_dims_list = []
             for dim in output_shape:
-                proto_dims_list.append(TensorShapeProto.Dim(size=dim))
+                proto_dims_list.append(TensorShapeProto.Dim(size=dim))  # dim is expected to be an integer
             tshape_proto = TensorShapeProto(dim=proto_dims_list)
             tshape_protos.append(tshape_proto)
         node_properties['_output_shapes'] = AttrValue(list=AttrValue.ListValue(shape=tshape_protos))
@@ -263,7 +263,13 @@ class TensorboardWriter(object):
 
             # For nodes with an "empty" output shape.
             output_shape = (None,) if n.output_shape == () else n.output_shape
-
+            if 'CombinedNonMaxSuppression' in str(output_shape):
+                # output_shapes is expected to be a list of tuples where each tuple is an output shape.
+                # For NMS layers, we need to align the node's output shapes before creating the node's properties.
+                output_shape = [output_shape.nmsed_boxes,
+                                output_shape.nmsed_scores,
+                                output_shape.nmsed_classes,
+                                output_shape.valid_detections]
             dims = []
             if isinstance(output_shape, list):
                 for o in output_shape:

{mct-nightly-2.0.0.20240509.406 → mct-nightly-2.0.0.20240510.421}/model_compression_toolkit/data_generation/__init__.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from model_compression_toolkit.constants import FOUND_TORCH, FOUND_TF
+from model_compression_toolkit.constants import FOUND_TORCH, FOUND_TF, FOUND_TORCHVISION
 from model_compression_toolkit.data_generation.common.data_generation_config import DataGenerationConfig
 from model_compression_toolkit.data_generation.common.enums import ImageGranularity, DataInitType, SchedulerType, BNLayerWeightingType, OutputLossType, BatchNormAlignemntLossType, ImagePipelineType, ImageNormalizationType
 
@@ -21,6 +21,6 @@ if FOUND_TF:
     from model_compression_toolkit.data_generation.keras.keras_data_generation import (
         keras_data_generation_experimental, get_keras_data_generation_config)
 
-if FOUND_TORCH:
+if FOUND_TORCH and FOUND_TORCHVISION:
     from model_compression_toolkit.data_generation.pytorch.pytorch_data_generation import (
         pytorch_data_generation_experimental, get_pytorch_data_generation_config)

{mct-nightly-2.0.0.20240509.406 → mct-nightly-2.0.0.20240510.421}/model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py

@@ -17,7 +17,7 @@ from typing import Callable, Any, Tuple, List
 
 from tqdm import tqdm
 
-from model_compression_toolkit.constants import FOUND_TORCH
+from model_compression_toolkit.constants import FOUND_TORCH, FOUND_TORCHVISION
 from model_compression_toolkit.core.pytorch.utils import set_model
 from model_compression_toolkit.data_generation.common.constants import DEFAULT_N_ITER, DEFAULT_DATA_GEN_BS
 from model_compression_toolkit.data_generation.common.data_generation import get_data_generation_classes
@@ -44,7 +44,7 @@ from model_compression_toolkit.data_generation.pytorch.optimization_functions.sc
 from model_compression_toolkit.data_generation.pytorch.optimization_utils import PytorchImagesOptimizationHandler
 from model_compression_toolkit.logger import Logger
 
-if FOUND_TORCH:
+if FOUND_TORCH and FOUND_TORCHVISION:
     # Importing necessary libraries
     import torch
     from torch import Tensor
@@ -354,10 +354,9 @@ else:
     # If torch is not installed,
     # we raise an exception when trying to use these functions.
     def get_pytorch_data_generation_config(*args, **kwargs):
-        Logger.critical('PyTorch must be installed to use get_pytorch_data_generation_config. '
-                        "The 'torch' package is missing.")  # pragma: no cover
-
+        msg = f"PyTorch and torchvision must be installed to use get_pytorch_data_generation_config. " + ("" if FOUND_TORCH else "The 'torch' package is missing. ") + ("" if FOUND_TORCHVISION else "The 'torchvision' package is missing. ") # pragma: no cover
+        Logger.critical(msg)  # pragma: no cover
 
     def pytorch_data_generation_experimental(*args, **kwargs):
-        Logger.critical("PyTorch must be installed to use 'pytorch_data_generation_experimental'. "
-                        "The 'torch' package is missing.")  # pragma: no cover
+        msg = f"PyTorch and torchvision must be installed to use pytorch_data_generation_experimental. " + ("" if FOUND_TORCH else "The 'torch' package is missing. ") + ("" if FOUND_TORCHVISION else "The 'torchvision' package is missing. ") # pragma: no cover
+        Logger.critical(msg)  # pragma: no cover

mct-nightly-2.0.0.20240510.421/model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -0,0 +1,138 @@
+# Copyright 2022 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# 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.
+# ==============================================================================
+from typing import Callable
+from io import BytesIO
+
+import torch.nn
+
+from mct_quantizers import PytorchActivationQuantizationHolder, PytorchQuantizationWrapper
+from model_compression_toolkit.constants import FOUND_ONNX
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
+from model_compression_toolkit.exporter.model_exporter.pytorch.base_pytorch_exporter import BasePyTorchExporter
+from mct_quantizers import pytorch_quantizers
+
+
+if FOUND_ONNX:
+    import onnx
+    from mct_quantizers.pytorch.metadata import add_onnx_metadata
+
+    class FakelyQuantONNXPyTorchExporter(BasePyTorchExporter):
+        """
+        Exporter for fakely-quant PyTorch models.
+        The exporter expects to receive an exportable model (where each layer's full quantization parameters
+        can be retrieved), and convert it into a fakely-quant model (namely, weights that are in fake-quant
+        format) and fake-quant layers for the activations.
+        """
+
+        def __init__(self,
+                     model: torch.nn.Module,
+                     is_layer_exportable_fn: Callable,
+                     save_model_path: str,
+                     repr_dataset: Callable,
+                     onnx_opset_version: int,
+                     use_onnx_custom_quantizer_ops: bool = False):
+            """
+
+            Args:
+                model: Model to export.
+                is_layer_exportable_fn: Callable to check whether a layer can be exported or not.
+                save_model_path: Path to save the exported model.
+                repr_dataset: Representative dataset (needed for creating torch script).
+                onnx_opset_version: ONNX opset version to use for exported ONNX model.
+                use_onnx_custom_quantizer_ops: Whether to export quantizers custom ops in ONNX or not.
+            """
+
+            super().__init__(model,
+                             is_layer_exportable_fn,
+                             save_model_path,
+                             repr_dataset)
+
+            self._use_onnx_custom_quantizer_ops = use_onnx_custom_quantizer_ops
+            self._onnx_opset_version = onnx_opset_version
+
+        def export(self) -> None:
+            """
+            Convert an exportable (fully-quantized) PyTorch model to a fakely-quant model
+            (namely, weights that are in fake-quant format) and fake-quant layers for the activations.
+
+            Returns:
+                Fake-quant PyTorch model.
+            """
+            for layer in self.model.children():
+                self.is_layer_exportable_fn(layer)
+
+            # Set forward that is used during onnx export.
+            # If _use_onnx_custom_quantizer_ops is set to True, the quantizer forward function will use
+            # the custom implementation when exporting the operator into onnx model. If not, it removes the
+            # wraps and quantizes the ops in place (for weights, for activation torch quantization function is
+            # exported since it's used during forward).
+            if self._use_onnx_custom_quantizer_ops:
+                self._enable_onnx_custom_ops_export()
+            else:
+                self._substitute_fully_quantized_model()
+
+            if self._use_onnx_custom_quantizer_ops:
+                Logger.info(f"Exporting onnx model with MCTQ quantizers: {self.save_model_path}")
+            else:
+                Logger.info(f"Exporting fake-quant onnx model: {self.save_model_path}")
+
+            model_input = to_torch_tensor(next(self.repr_dataset())[0])
+
+            if hasattr(self.model, 'metadata'):
+                onnx_bytes = BytesIO()
+                torch.onnx.export(self.model,
+                                  model_input,
+                                  onnx_bytes,
+                                  opset_version=self._onnx_opset_version,
+                                  verbose=False,
+                                  input_names=['input'],
+                                  output_names=['output'],
+                                  dynamic_axes={'input': {0: 'batch_size'},
+                                                'output': {0: 'batch_size'}})
+                onnx_model = onnx.load_from_string(onnx_bytes.getvalue())
+                onnx_model = add_onnx_metadata(onnx_model, self.model.metadata)
+                onnx.save_model(onnx_model, self.save_model_path)
+            else:
+                torch.onnx.export(self.model,
+                                  model_input,
+                                  self.save_model_path,
+                                  opset_version=self._onnx_opset_version,
+                                  verbose=False,
+                                  input_names=['input'],
+                                  output_names=['output'],
+                                  dynamic_axes={'input': {0: 'batch_size'},
+                                                'output': {0: 'batch_size'}})
+
+        def _enable_onnx_custom_ops_export(self):
+            """
+            Enable the custom implementation forward in quantizers, so it is exported
+            with custom quantizers.
+            """
+
+            for n, m in self.model.named_modules():
+                if isinstance(m, PytorchActivationQuantizationHolder):
+                    assert isinstance(m.activation_holder_quantizer, pytorch_quantizers.BasePyTorchInferableQuantizer)
+                    m.activation_holder_quantizer.enable_custom_impl()
+
+                if isinstance(m, PytorchQuantizationWrapper):
+                    for wq in m.weights_quantizers.values():
+                        assert isinstance(wq, pytorch_quantizers.BasePyTorchInferableQuantizer)
+                        wq.enable_custom_impl()
+
+else:
+    def FakelyQuantONNXPyTorchExporter(*args, **kwargs):
+        Logger.critical("ONNX must be installed to use 'FakelyQuantONNXPyTorchExporter'. "
+                        "The 'onnx' package is missing.")  # pragma: no cover

{mct-nightly-2.0.0.20240509.406 → mct-nightly-2.0.0.20240510.421}/model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py

@@ -19,15 +19,15 @@ from model_compression_toolkit.exporter.model_exporter.fw_agonstic.quantization_
 from model_compression_toolkit.exporter.model_exporter.pytorch.export_serialization_format import \
     PytorchExportSerializationFormat
 from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
+
+
+DEFAULT_ONNX_OPSET_VERSION = 15
 
 
 if FOUND_TORCH:
     import torch.nn
-    from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_onnx_pytorch_exporter import \
-    FakelyQuantONNXPyTorchExporter, DEFAULT_ONNX_OPSET_VERSION
-    from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_torchscript_pytorch_exporter import \
-        FakelyQuantTorchScriptPyTorchExporter
+    from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_onnx_pytorch_exporter import FakelyQuantONNXPyTorchExporter
+    from model_compression_toolkit.exporter.model_exporter.pytorch.fakely_quant_torchscript_pytorch_exporter import FakelyQuantTorchScriptPyTorchExporter
     from model_compression_toolkit.exporter.model_wrapper.pytorch.validate_layer import is_pytorch_layer_exportable
 
     supported_serialization_quantization_export_dict = {

{mct-nightly-2.0.0.20240509.406 → mct-nightly-2.0.0.20240510.421}/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2/tp_model.py

@@ -152,19 +152,6 @@ def generate_tp_model(default_config: OpQuantizationConfig,
     # this configuration will be used for the operation quantization:
     default_configuration_options = tp.QuantizationConfigOptions([default_config])
 
-    # Create a QuantizationConfigOptions for quantizing constants in functional ops.
-    # Constant configuration is similar to the default eight bit configuration except for PoT
-    # quantization method for the constant.
-    # Since the constants are not named attributes of the layer, we use the default_weight_attr_config to
-    # define the desired quantization properties for them.
-    const_config = default_config.clone_and_edit(
-        default_weight_attr_config=default_config.default_weight_attr_config.clone_and_edit(
-            enable_weights_quantization=True))
-    if not (const_config.default_weight_attr_config.weights_quantization_method == tp.QuantizationMethod.POWER_OF_TWO and
-            const_config.default_weight_attr_config.weights_per_channel_threshold is False):
-        mct.logger.Logger.error('Constant quantization config should be per-tensor PoT.')
-    const_configuration_options = tp.QuantizationConfigOptions([const_config])
-
     # Create a TargetPlatformModel and set its default quantization config.
     # This default configuration will be used for all operations
     # unless specified otherwise (see OperatorsSet, for example):
@@ -198,10 +185,10 @@ def generate_tp_model(default_config: OpQuantizationConfig,
         # Define operations sets without quantization configuration
         # options (useful for creating fusing patterns, for example):
         any_relu = tp.OperatorsSet("AnyReLU")
-        add = tp.OperatorsSet("Add", const_configuration_options)
-        sub = tp.OperatorsSet("Sub", const_configuration_options)
-        mul = tp.OperatorsSet("Mul", const_configuration_options)
-        div = tp.OperatorsSet("Div", const_configuration_options)
+        add = tp.OperatorsSet("Add")
+        sub = tp.OperatorsSet("Sub")
+        mul = tp.OperatorsSet("Mul")
+        div = tp.OperatorsSet("Div")
         prelu = tp.OperatorsSet("PReLU")
         swish = tp.OperatorsSet("Swish")
         sigmoid = tp.OperatorsSet("Sigmoid")

{mct-nightly-2.0.0.20240509.406 → mct-nightly-2.0.0.20240510.421}/model_compression_toolkit/target_platform_capabilities/tpc_models/imx500_tpc/v2_lut/tp_model.py

@@ -151,19 +151,6 @@ def generate_tp_model(default_config: OpQuantizationConfig,
     # this configuration will be used for the operation quantization:
     default_configuration_options = tp.QuantizationConfigOptions([default_config])
 
-    # Create a QuantizationConfigOptions for quantizing constants in functional ops.
-    # Constant configuration is similar to the default eight bit configuration except for PoT
-    # quantization method for the constant.
-    # Since the constants are not named attributes of the layer, we use the default_weight_attr_config to
-    # define the desired quantization properties for them.
-    const_config = default_config.clone_and_edit(
-        default_weight_attr_config=default_config.default_weight_attr_config.clone_and_edit(
-            enable_weights_quantization=True))
-    if not (const_config.default_weight_attr_config.weights_quantization_method == tp.QuantizationMethod.POWER_OF_TWO and
-            const_config.default_weight_attr_config.weights_per_channel_threshold is False):
-        mct.logger.Logger.error('Constant quantization config should be per-tensor PoT.')
-    const_configuration_options = tp.QuantizationConfigOptions([const_config])
-
     # Create a TargetPlatformModel and set its default quantization config.
     # This default configuration will be used for all operations
     # unless specified otherwise (see OperatorsSet, for example):
@@ -195,10 +182,10 @@ def generate_tp_model(default_config: OpQuantizationConfig,
         # Define operations sets without quantization configuration
         # options (useful for creating fusing patterns, for example):
         any_relu = tp.OperatorsSet("AnyReLU")
-        add = tp.OperatorsSet("Add", const_configuration_options)
-        sub = tp.OperatorsSet("Sub", const_configuration_options)
-        mul = tp.OperatorsSet("Mul", const_configuration_options)
-        div = tp.OperatorsSet("Div", const_configuration_options)
+        add = tp.OperatorsSet("Add")
+        sub = tp.OperatorsSet("Sub")
+        mul = tp.OperatorsSet("Mul")
+        div = tp.OperatorsSet("Div")
         prelu = tp.OperatorsSet("PReLU")
         swish = tp.OperatorsSet("Swish")
         sigmoid = tp.OperatorsSet("Sigmoid")

mct-nightly-2.0.0.20240509.406/model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -1,132 +0,0 @@
-# Copyright 2022 Sony Semiconductor Israel, Inc. All rights reserved.
-#
-# 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.
-# ==============================================================================
-from typing import Callable
-from io import BytesIO
-
-import torch.nn
-import onnx
-
-from mct_quantizers import PytorchActivationQuantizationHolder, PytorchQuantizationWrapper
-from model_compression_toolkit.logger import Logger
-from model_compression_toolkit.core.pytorch.utils import to_torch_tensor
-from model_compression_toolkit.exporter.model_exporter.pytorch.base_pytorch_exporter import BasePyTorchExporter
-from mct_quantizers import pytorch_quantizers
-from mct_quantizers.pytorch.metadata import add_onnx_metadata
-
-DEFAULT_ONNX_OPSET_VERSION=15
-
-
-class FakelyQuantONNXPyTorchExporter(BasePyTorchExporter):
-    """
-    Exporter for fakely-quant PyTorch models.
-    The exporter expects to receive an exportable model (where each layer's full quantization parameters
-    can be retrieved), and convert it into a fakely-quant model (namely, weights that are in fake-quant
-    format) and fake-quant layers for the activations.
-    """
-
-    def __init__(self,
-                 model: torch.nn.Module,
-                 is_layer_exportable_fn: Callable,
-                 save_model_path: str,
-                 repr_dataset: Callable,
-                 use_onnx_custom_quantizer_ops: bool = False,
-                 onnx_opset_version=DEFAULT_ONNX_OPSET_VERSION):
-        """
-
-        Args:
-            model: Model to export.
-            is_layer_exportable_fn: Callable to check whether a layer can be exported or not.
-            save_model_path: Path to save the exported model.
-            repr_dataset: Representative dataset (needed for creating torch script).
-            use_onnx_custom_quantizer_ops: Whether to export quantizers custom ops in ONNX or not.
-            onnx_opset_version: ONNX opset version to use for exported ONNX model.
-        """
-
-        super().__init__(model,
-                         is_layer_exportable_fn,
-                         save_model_path,
-                         repr_dataset)
-
-        self._use_onnx_custom_quantizer_ops = use_onnx_custom_quantizer_ops
-        self._onnx_opset_version = onnx_opset_version
-
-    def export(self) -> None:
-        """
-        Convert an exportable (fully-quantized) PyTorch model to a fakely-quant model
-        (namely, weights that are in fake-quant format) and fake-quant layers for the activations.
-
-        Returns:
-            Fake-quant PyTorch model.
-        """
-        for layer in self.model.children():
-            self.is_layer_exportable_fn(layer)
-
-        # Set forward that is used during onnx export.
-        # If _use_onnx_custom_quantizer_ops is set to True, the quantizer forward function will use
-        # the custom implementation when exporting the operator into onnx model. If not, it removes the
-        # wraps and quantizes the ops in place (for weights, for activation torch quantization function is
-        # exported since it's used during forward).
-        if self._use_onnx_custom_quantizer_ops:
-            self._enable_onnx_custom_ops_export()
-        else:
-            self._substitute_fully_quantized_model()
-
-        if self._use_onnx_custom_quantizer_ops:
-            Logger.info(f"Exporting onnx model with MCTQ quantizers: {self.save_model_path}")
-        else:
-            Logger.info(f"Exporting fake-quant onnx model: {self.save_model_path}")
-
-        model_input = to_torch_tensor(next(self.repr_dataset())[0])
-
-        if hasattr(self.model, 'metadata'):
-            onnx_bytes = BytesIO()
-            torch.onnx.export(self.model,
-                              model_input,
-                              onnx_bytes,
-                              opset_version=self._onnx_opset_version,
-                              verbose=False,
-                              input_names=['input'],
-                              output_names=['output'],
-                              dynamic_axes={'input': {0: 'batch_size'},
-                                            'output': {0: 'batch_size'}})
-            onnx_model = onnx.load_from_string(onnx_bytes.getvalue())
-            onnx_model = add_onnx_metadata(onnx_model, self.model.metadata)
-            onnx.save_model(onnx_model, self.save_model_path)
-        else:
-            torch.onnx.export(self.model,
-                              model_input,
-                              self.save_model_path,
-                              opset_version=self._onnx_opset_version,
-                              verbose=False,
-                              input_names=['input'],
-                              output_names=['output'],
-                              dynamic_axes={'input': {0: 'batch_size'},
-                                            'output': {0: 'batch_size'}})
-
-    def _enable_onnx_custom_ops_export(self):
-        """
-        Enable the custom implementation forward in quantizers, so it is exported
-        with custom quantizers.
-        """
-
-        for n, m in self.model.named_modules():
-            if isinstance(m, PytorchActivationQuantizationHolder):
-                assert isinstance(m.activation_holder_quantizer, pytorch_quantizers.BasePyTorchInferableQuantizer)
-                m.activation_holder_quantizer.enable_custom_impl()
-
-            if isinstance(m, PytorchQuantizationWrapper):
-                for wq in m.weights_quantizers.values():
-                    assert isinstance(wq, pytorch_quantizers.BasePyTorchInferableQuantizer)
-                    wq.enable_custom_impl()