mct-nightly 2.0.0.20240508.121838__py3-none-any.whl → 2.0.0.20240508.122218__py3-none-any.whl

{mct_nightly-2.0.0.20240508.121838.dist-info → mct_nightly-2.0.0.20240508.122218.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mct-nightly
-Version: 2.0.0.20240508.121838
+Version: 2.0.0.20240508.122218
 Summary: A Model Compression Toolkit for neural networks
 Home-page: UNKNOWN
 License: UNKNOWN

{mct_nightly-2.0.0.20240508.121838.dist-info → mct_nightly-2.0.0.20240508.122218.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=TDig5WOAQRhJrpoVfD05x8ZxU7iOSUl3OtaVQ1acZMM,1573
+model_compression_toolkit/__init__.py,sha256=UwOX-ld2fBD1wK2ssmJUlOZGDaW6NDVqN1PQO_5RR6A,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
@@ -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=r2pOWFK-mSG8OzRiKGVOG4skzX0ZiM0eiRuBsL-ThoI,6067
+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_torchscript_pytorch_exporter.py,sha256=ksWV2A-Njo-wAxQ_Ye2sLIZXBWJ_WNyjT7-qFFwvV2o,2897
-model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py,sha256=yz5dPMX5r1d9LJV4rYFS1pXqCbVUxvUmV4LELWcRinQ,6350
+model_compression_toolkit/exporter/model_exporter/pytorch/pytorch_export_facade.py,sha256=5G9dikFY4A66XVjpaOWVWX81Qr6ZdwnoyBzFDL_abi8,6242
 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.121838.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.0.0.20240508.121838.dist-info/METADATA,sha256=pYFfSFSP2HsNDAum_dYBTotnT_hTBo3sPLJUZXAkrkE,18798
-mct_nightly-2.0.0.20240508.121838.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-mct_nightly-2.0.0.20240508.121838.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.0.0.20240508.121838.dist-info/RECORD,,
+mct_nightly-2.0.0.20240508.122218.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.0.0.20240508.122218.dist-info/METADATA,sha256=la_dJPFTwsLA7eYayNIaXNrtV6umEQjM8O7uyHz45Aw,18798
+mct_nightly-2.0.0.20240508.122218.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+mct_nightly-2.0.0.20240508.122218.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.0.0.20240508.122218.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.121838"
+__version__ = "2.0.0.20240508.122218"

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -16,117 +16,124 @@ 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)
 
-        self._use_onnx_custom_quantizer_ops = use_onnx_custom_quantizer_ops
-        self._onnx_opset_version = onnx_opset_version
+if FOUND_ONNX:
+    import onnx
+    from mct_quantizers.pytorch.metadata import add_onnx_metadata
 
-    def export(self) -> None:
+    class FakelyQuantONNXPyTorchExporter(BasePyTorchExporter):
         """
-        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.
+        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.
         """
-        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()
+        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

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
+from model_compression_toolkit.constants import FOUND_TORCH, FOUND_ONNX
 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
-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 = {