mct-nightly 2.3.0.20250429.622__py3-none-any.whl → 2.3.0.20250430.538__py3-none-any.whl

{mct_nightly-2.3.0.20250429.622.dist-info → mct_nightly-2.3.0.20250430.538.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: mct-nightly
-Version: 2.3.0.20250429.622
+Version: 2.3.0.20250430.538
 Summary: A Model Compression Toolkit for neural networks
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: Apache Software License

{mct_nightly-2.3.0.20250429.622.dist-info → mct_nightly-2.3.0.20250430.538.dist-info}/RECORD

@@ -1,5 +1,5 @@
-mct_nightly-2.3.0.20250429.622.dist-info/licenses/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-model_compression_toolkit/__init__.py,sha256=PrThtvqTwsbfiKBPdckZFWhbBXaRfAATT21j_WbV8pA,1557
+mct_nightly-2.3.0.20250430.538.dist-info/licenses/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+model_compression_toolkit/__init__.py,sha256=_uN22zhCLe0oS_Ex9lxOL4jmaLZUIfp9zJ82XXgBrqQ,1557
 model_compression_toolkit/constants.py,sha256=iJ6vfTjC2oFIZWt8wvHoxEw5YJi3yl0Hd4q30_8q0Zc,3958
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=L3q7tn3Uht0i_7phnlOWMR2Te2zvzrt2HOz9vYEInts,4529
@@ -335,7 +335,7 @@ 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=G2X_lDx6u12U4ErEuEHCdNczh0qSGWObySw3upEys6Q,7506
+model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py,sha256=5DHg8ch8_DtSQ7M5Aw3LcZLSVFqFH556cKcosp3Ik-I,8720
 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=8vYGKa58BkasvoHejYaPwubOJPcW0s-RY79_Kkw0Hy8,6236
 model_compression_toolkit/exporter/model_wrapper/__init__.py,sha256=7CF2zvpTrIEm8qnbuHnLZyTZkwBBxV24V8QA0oxGbh0,1187
@@ -528,7 +528,7 @@ model_compression_toolkit/xquant/pytorch/model_analyzer.py,sha256=b93o800yVB3Z-i
 model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py,sha256=UVN_S9ULHBEldBpShCOt8-soT8YTQ5oE362y96qF_FA,3950
 model_compression_toolkit/xquant/pytorch/similarity_functions.py,sha256=CERxq5K8rqaiE-DlwhZBTUd9x69dtYJlkHOPLB54vm8,2354
 model_compression_toolkit/xquant/pytorch/tensorboard_utils.py,sha256=mkoEktLFFHtEKzzFRn_jCnxjhJolK12TZ5AQeDHzUO8,9767
-mct_nightly-2.3.0.20250429.622.dist-info/METADATA,sha256=R9WTA_IVw4cvLIbQbGpmzTl_ujwHc4RNElUajWHnSNE,25101
-mct_nightly-2.3.0.20250429.622.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
-mct_nightly-2.3.0.20250429.622.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.3.0.20250429.622.dist-info/RECORD,,
+mct_nightly-2.3.0.20250430.538.dist-info/METADATA,sha256=8l9P7bLANjVXo-Z5_aR8x8VtcIWHhdYNNHQmoOS6cig,25101
+mct_nightly-2.3.0.20250430.538.dist-info/WHEEL,sha256=ck4Vq1_RXyvS4Jt6SI0Vz6fyVs4GWg7AINwpsaGEgPE,91
+mct_nightly-2.3.0.20250430.538.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.3.0.20250430.538.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.3.0.20250429.000622"
+__version__ = "2.3.0.20250430.000538"

model_compression_toolkit/exporter/model_exporter/pytorch/fakely_quant_onnx_pytorch_exporter.py

@@ -24,11 +24,11 @@ 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.
@@ -63,7 +63,7 @@ if FOUND_ONNX:
             self._use_onnx_custom_quantizer_ops = use_onnx_custom_quantizer_ops
             self._onnx_opset_version = onnx_opset_version
 
-        def export(self) -> None:
+        def export(self, output_names=None) -> 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.
@@ -95,6 +95,28 @@ if FOUND_ONNX:
                 Logger.info(f"Exporting fake-quant onnx model: {self.save_model_path}")
 
             model_input = to_torch_tensor(next(self.repr_dataset()))
+            model_output = self.model(*model_input) if isinstance(model_input, (list, tuple)) else self.model(
+                model_input)
+
+            if output_names is None:
+                # Determine number of outputs and prepare output_names and dynamic_axes
+                if isinstance(model_output, (list, tuple)):
+                    output_names = [f"output_{i}" for i in range(len(model_output))]
+                    dynamic_axes = {'input': {0: 'batch_size'}}
+                    dynamic_axes.update({name: {0: 'batch_size'} for name in output_names})
+                else:
+                    output_names = ['output']
+                    dynamic_axes = {'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}}
+            else:
+                if isinstance(model_output, (list, tuple)):
+                    num_of_outputs = len(model_output)
+                else:
+                    num_of_outputs = 1
+                assert len(output_names) == num_of_outputs, (f"Mismatch between number of requested output names "
+                                                             f"({output_names}) and model output count "
+                                                             f"({num_of_outputs}):\n")
+                dynamic_axes = {'input': {0: 'batch_size'}}
+                dynamic_axes.update({name: {0: 'batch_size'} for name in output_names})
 
             if hasattr(self.model, 'metadata'):
                 onnx_bytes = BytesIO()
@@ -104,9 +126,8 @@ if FOUND_ONNX:
                                   opset_version=self._onnx_opset_version,
                                   verbose=False,
                                   input_names=['input'],
-                                  output_names=['output'],
-                                  dynamic_axes={'input': {0: 'batch_size'},
-                                                'output': {0: 'batch_size'}})
+                                  output_names=output_names,
+                                  dynamic_axes=dynamic_axes)
                 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)
@@ -117,9 +138,8 @@ if FOUND_ONNX:
                                   opset_version=self._onnx_opset_version,
                                   verbose=False,
                                   input_names=['input'],
-                                  output_names=['output'],
-                                  dynamic_axes={'input': {0: 'batch_size'},
-                                                'output': {0: 'batch_size'}})
+                                  output_names=output_names,
+                                  dynamic_axes=dynamic_axes)
 
             for layer in self.model.children():
                 # Set disable for reuse for weight quantizers if quantizer was reused