mct-nightly 2.0.0.20240508.145608__py3-none-any.whl → 2.0.0.20240509.406__py3-none-any.whl

{mct_nightly-2.0.0.20240508.145608.dist-info → mct_nightly-2.0.0.20240509.406.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.0.0.20240508.145608
+Version: 2.0.0.20240509.406
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN
@@ -72,9 +72,6 @@ for hands-on learning. For example:
 * [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).
 
-Additionally, for quick quantization of a variety of models from well-known collections,
-visit the [quick-start page](tutorials/quick_start/README.md) and the
-[results CSV](tutorials/quick_start/results/model_quantization_results.csv).
 
 ### Supported Versions
 

{mct_nightly-2.0.0.20240508.145608.dist-info → mct_nightly-2.0.0.20240509.406.dist-info}/RECORD

@@ -1,5 +1,5 @@
-model_compression_toolkit/__init__.py,sha256=BUMAk1BC9peC1TCFqoOqMrOyFT_LMr7kLtuYqQZTp8A,1573
-model_compression_toolkit/constants.py,sha256=b63Jk_bC7VXEX3Qn9TZ3wUvrNKD8Mkz8zIuayoyF5eU,3828
+model_compression_toolkit/__init__.py,sha256=ChRp1KQR5GCWN6Py2srdeALrqbdAT6VHfoe2LCWSqJc,1573
+model_compression_toolkit/constants.py,sha256=yIJyJ-e1WrDeKD9kG15qkqfYnoj7J1J2CxnJDt008ik,3756
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
 model_compression_toolkit/metadata.py,sha256=IyoON37lBv3TI0rZGCP4K5t3oYI4TOmYy-LRXOwHGpE,1136
@@ -145,7 +145,7 @@ model_compression_toolkit/core/common/substitutions/weights_activation_split.py,
 model_compression_toolkit/core/common/visualization/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7VclGVOQ63x2p6mAAuba4,698
 model_compression_toolkit/core/common/visualization/final_config_visualizer.py,sha256=6I10jKLesB-RQKaXA75Xgz2wPvylQUrnPtCcQZIynGo,6371
 model_compression_toolkit/core/common/visualization/nn_visualizer.py,sha256=HOq7AObkmEZiDSZXUMJDAEJzUY-fSXUT0AMgwiyH7dg,7388
-model_compression_toolkit/core/common/visualization/tensorboard_writer.py,sha256=lkQ5B3YKcojPfNdkPCZ9ViJ0zMOSsWmZ-ELmiBcNcqI,22510
+model_compression_toolkit/core/common/visualization/tensorboard_writer.py,sha256=4E4ZXZmqusGIJ4XQNH8FFt07htAHgT3gy5E7wPIaVBI,21951
 model_compression_toolkit/core/keras/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7VclGVOQ63x2p6mAAuba4,698
 model_compression_toolkit/core/keras/constants.py,sha256=Uv3c0UdW55pIVQNW_1HQlgl-dHXREkltOLyzp8G1mTQ,3163
 model_compression_toolkit/core/keras/custom_layer_validation.py,sha256=f-b14wuiIgitBe7d0MmofYhDCTO3IhwJgwrh-Hq_t_U,1192
@@ -266,7 +266,7 @@ model_compression_toolkit/core/pytorch/reader/node_holders.py,sha256=TaolORuwBZE
 model_compression_toolkit/core/pytorch/reader/reader.py,sha256=GEJE0QX8XJFWbYCkbRBtzttZtmmuoACLx8gw9KyAQCE,6015
 model_compression_toolkit/core/pytorch/statistics_correction/__init__.py,sha256=Rf1RcYmelmdZmBV5qOKvKWF575ofc06JFQSq83Jz99A,696
 model_compression_toolkit/core/pytorch/statistics_correction/apply_second_moment_correction.py,sha256=VgU24J3jf7QComHH7jonOXSkg6mO4TOch3uFkOthZvM,3261
-model_compression_toolkit/data_generation/__init__.py,sha256=S8pRUqlRvpM5AFHpFWs3zb0H0rtY5nUwmeCQij01oi4,1507
+model_compression_toolkit/data_generation/__init__.py,sha256=R_RnB8Evj4uq0WKiPWvBWfeePrbake7Z03ugJgK7jLo,1466
 model_compression_toolkit/data_generation/common/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/data_generation/common/constants.py,sha256=21e3ZX9WVYojexG2acTgklrBk8ZO9DjJnKpP4KHZC44,1018
 model_compression_toolkit/data_generation/common/data_generation.py,sha256=fccGG6cTMScZwjnJDQKMugOLdgm9dKg5rRfcBD4EFYQ,6415
@@ -292,7 +292,7 @@ model_compression_toolkit/data_generation/pytorch/constants.py,sha256=QWyreMImcf
 model_compression_toolkit/data_generation/pytorch/image_pipeline.py,sha256=6g7OpOuO3cU4TIuelaRjBKpCPgiMbe1a3iy9bZtdZUo,6617
 model_compression_toolkit/data_generation/pytorch/model_info_exctractors.py,sha256=sO9tA03nIaeYnzOL4Egec5sVcSGU8H8k9-nNjhaLEbk,9690
 model_compression_toolkit/data_generation/pytorch/optimization_utils.py,sha256=AjYsO-lm06JOUMoKkS6VbyF4O_l_ffWXrgamqJm1ofE,19085
-model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py,sha256=UGX0J0lU1bY4ZI6qE1K0AnFWsDFs3clYPBC4GZf9KxA,21219
+model_compression_toolkit/data_generation/pytorch/pytorch_data_generation.py,sha256=rZ-4YAcgEc9qZEs5FrK0OJaNtSsQC57Y61UdbXbQcE4,20937
 model_compression_toolkit/data_generation/pytorch/optimization_functions/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/data_generation/pytorch/optimization_functions/batchnorm_alignment_functions.py,sha256=dMc4zz9XfYfAT4Cxns57VgvGZWPAMfaGlWLFyCyl8TA,1968
 model_compression_toolkit/data_generation/pytorch/optimization_functions/bn_layer_weighting_functions.py,sha256=i3ePEI8xDE3xZEtmzT5lCkLn9wpObUi_OgqnVDf7nj8,2597
@@ -315,9 +315,9 @@ model_compression_toolkit/exporter/model_exporter/keras/mctq_keras_exporter.py,s
 model_compression_toolkit/exporter/model_exporter/pytorch/__init__.py,sha256=uZ2RigbY9O2PJ0Il8wPpS_s7frgg9WUGd_SHeKGyl1A,699
 model_compression_toolkit/exporter/model_exporter/pytorch/base_pytorch_exporter.py,sha256=UPVkEUQCMZ4Lld6CRnEOPEmlfe5vcQZG0Q3FwRBodD4,4021
 model_compression_toolkit/exporter/model_exporter/pytorch/export_serialization_format.py,sha256=bPevy6OBqng41PqytBR55e6cBEuyrUS0H8dWX4zgjQ4,967
-model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py,sha256=b-qC60LiRtc52gIXdUbrdTBKUgCIaResDLXFE8zt_F4,6732
+model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py,sha256=r2pOWFK-mSG8OzRiKGVOG4skzX0ZiM0eiRuBsL-ThoI,6067
 model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_torchscript_pytorch_exporter.py,sha256=ksWV2A-Njo-wAxQ_Ye2sLIZXBWJ_WNyjT7-qFFwvV2o,2897
-model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py,sha256=5G9dikFY4A66XVjpaOWVWX81Qr6ZdwnoyBzFDL_abi8,6242
+model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py,sha256=yz5dPMX5r1d9LJV4rYFS1pXqCbVUxvUmV4LELWcRinQ,6350
 model_compression_toolkit/exporter/model_wrapper/__init__.py,sha256=7CF2zvpTrIEm8qnbuHnLZyTZkwBBxV24V8QA0oxGbh0,1187
 model_compression_toolkit/exporter/model_wrapper/fw_agnostic/__init__.py,sha256=pKAdbTCFM_2BrZXUtTIw0ouKotrWwUDF_hP3rPwCM2k,696
 model_compression_toolkit/exporter/model_wrapper/fw_agnostic/get_inferable_quantizers.py,sha256=Bd3QhAR__YC9Xmobd5qHv9ofh_rPn_eTFV0sXizcBnY,2297
@@ -483,8 +483,8 @@ model_compression_toolkit/trainable_infrastructure/keras/quantize_wrapper.py,sha
 model_compression_toolkit/trainable_infrastructure/keras/quantizer_utils.py,sha256=MVwXNymmFRB2NXIBx4e2mdJ1RfoHxRPYRgjb1MQP5kY,1797
 model_compression_toolkit/trainable_infrastructure/pytorch/__init__.py,sha256=huHoBUcKNB6BnY6YaUCcFvdyBtBI172ZoUD8ZYeNc6o,696
 model_compression_toolkit/trainable_infrastructure/pytorch/base_pytorch_quantizer.py,sha256=MxylaVFPgN7zBiRBy6WV610EA4scLgRJFbMucKvvNDU,2896
-mct_nightly-2.0.0.20240508.145608.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.0.0.20240508.145608.dist-info/METADATA,sha256=5XKDpGXcxCMKbL1iHeQRupmzRYmVT1gTB1hAUPPSTJU,18798
-mct_nightly-2.0.0.20240508.145608.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-mct_nightly-2.0.0.20240508.145608.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.0.0.20240508.145608.dist-info/RECORD,,
+mct_nightly-2.0.0.20240509.406.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.0.0.20240509.406.dist-info/METADATA,sha256=CnhcTwwsr7Ks92s0saVvRr0npvEJoGDpeUCCH3OcWfU,18559
+mct_nightly-2.0.0.20240509.406.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+mct_nightly-2.0.0.20240509.406.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.0.0.20240509.406.dist-info/RECORD,,

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.20240508.145608"
+__version__ = "2.0.0.20240509.000406"

model_compression_toolkit/constants.py

@@ -20,7 +20,6 @@ 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

model_compression_toolkit/core/common/visualization/tensorboard_writer.py

@@ -72,19 +72,10 @@ def get_node_properties(node_dict_to_log: dict,
     # Create protobuf for the node's output shapes
     if output_shapes is not None:
         tshape_protos = []
-        is_tf_combined_non_max_suppression = len(output_shapes) == 1 and 'function' in node_dict_to_log and node_dict_to_log['function'] == 'image.combined_non_max_suppression'
-
-        if is_tf_combined_non_max_suppression:
-            combined_nms_output = output_shapes[0]
-            output_shapes = [combined_nms_output.nmsed_boxes,
-                             combined_nms_output.nmsed_scores,
-                             combined_nms_output.nmsed_classes,
-                             combined_nms_output.valid_detections]
-
         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)) # dim shold ne an integer
+                proto_dims_list.append(TensorShapeProto.Dim(size=dim))
             tshape_proto = TensorShapeProto(dim=proto_dims_list)
             tshape_protos.append(tshape_proto)
         node_properties['_output_shapes'] = AttrValue(list=AttrValue.ListValue(shape=tshape_protos))

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, FOUND_TORCHVISION
+from model_compression_toolkit.constants import FOUND_TORCH, FOUND_TF
 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 and FOUND_TORCHVISION:
+if FOUND_TORCH:
     from model_compression_toolkit.data_generation.pytorch.pytorch_data_generation import (
         pytorch_data_generation_experimental, get_pytorch_data_generation_config)

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, FOUND_TORCHVISION
+from model_compression_toolkit.constants import FOUND_TORCH
 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 and FOUND_TORCHVISION:
+if FOUND_TORCH:
     # Importing necessary libraries
     import torch
     from torch import Tensor
@@ -354,9 +354,10 @@ else:
     # If torch is not installed,
     # we raise an exception when trying to use these functions.
     def get_pytorch_data_generation_config(*args, **kwargs):
-        msg = f"torch and torchvision must be installed to use get_pytorch_data_generation_config. " + ("" if FOUND_TORCH else "'torch' package is missing. ") + ("" if FOUND_TORCHVISION else "'torchvision' package is missing. ") # pragma: no cover
-        Logger.critical(msg)  # pragma: no cover
+        Logger.critical('PyTorch must be installed to use get_pytorch_data_generation_config. '
+                        "The 'torch' package is missing.")  # pragma: no cover
+
 
     def pytorch_data_generation_experimental(*args, **kwargs):
-        msg = f"torch and torchvision must be installed to use pytorch_data_generation_experimental. " + ("" if FOUND_TORCH else "'torch' package is missing. ") + ("" if FOUND_TORCHVISION else "'torchvision' package is missing. ") # pragma: no cover
-        Logger.critical(msg)  # pragma: no cover
+        Logger.critical("PyTorch must be installed to use 'pytorch_data_generation_experimental'. "
+                        "The 'torch' package is missing.")  # pragma: no cover

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -16,124 +16,117 @@ from typing import Callable
 from io import BytesIO
 
 import torch.nn
+import onnx
 
 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
+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)
 
-if FOUND_ONNX:
-    import onnx
-    from mct_quantizers.pytorch.metadata import add_onnx_metadata
+        self._use_onnx_custom_quantizer_ops = use_onnx_custom_quantizer_ops
+        self._onnx_opset_version = onnx_opset_version
 
-    class FakelyQuantONNXPyTorchExporter(BasePyTorchExporter):
+    def export(self) -> None:
         """
-        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.
+        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()
 
-        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('Installing onnx is mandatory '
-                            'when using FakelyQuantONNXPyTorchExporter. '
-                            'Could not find onnx package.')  # pragma: no cover
+            if isinstance(m, PytorchQuantizationWrapper):
+                for wq in m.weights_quantizers.values():
+                    assert isinstance(wq, pytorch_quantizers.BasePyTorchInferableQuantizer)
+                    wq.enable_custom_impl()

model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py

@@ -14,20 +14,20 @@
 # ==============================================================================
 from typing import Callable
 
-from model_compression_toolkit.constants import FOUND_TORCH, FOUND_ONNX
+from model_compression_toolkit.constants import FOUND_TORCH
 from model_compression_toolkit.exporter.model_exporter.fw_agonstic.quantization_format import QuantizationFormat
 from model_compression_toolkit.exporter.model_exporter.pytorch.export_serialization_format import \
     PytorchExportSerializationFormat
 from model_compression_toolkit.logger import Logger
-
-
-DEFAULT_ONNX_OPSET_VERSION = 15
+from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
 
 
 if FOUND_TORCH:
     import torch.nn
-    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_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_wrapper.pytorch.validate_layer import is_pytorch_layer_exportable
 
     supported_serialization_quantization_export_dict = {