mct-nightly 2.2.0.20250108.523__py3-none-any.whl → 2.2.0.20250109.528__py3-none-any.whl

{mct_nightly-2.2.0.20250108.523.dist-info → mct_nightly-2.2.0.20250109.528.dist-info}/METADATA

@@ -1,6 +1,6 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.2
 Name: mct-nightly
-Version: 2.2.0.20250108.523
+Version: 2.2.0.20250109.528
 Summary: A Model Compression Toolkit for neural networks
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: Apache Software License
@@ -23,6 +23,12 @@ Requires-Dist: scipy
 Requires-Dist: protobuf
 Requires-Dist: mct-quantizers==1.5.2
 Requires-Dist: pydantic<2.0
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
 
 <div align="center" markdown="1">
 <p>

{mct_nightly-2.2.0.20250108.523.dist-info → mct_nightly-2.2.0.20250109.528.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=T9vJvHgnVfRm0QlXM_rI0gN4vrj4U14twllXd5p0Irs,1573
+model_compression_toolkit/__init__.py,sha256=5vDy_YtqV8rtgpIP3dBPtIxugUkeVPFJU05OH1jkhdw,1573
 model_compression_toolkit/constants.py,sha256=i_R6uXBfO1ph_X6DNJych2x59SUojfJbn7dNjs_mZnc,3846
 model_compression_toolkit/defaultdict.py,sha256=LSc-sbZYXENMCw3U9F4GiXuv67IKpdn0Qm7Fr11jy-4,2277
 model_compression_toolkit/logger.py,sha256=3DByV41XHRR3kLTJNbpaMmikL8icd9e1N-nkQAY9oDk,4567
@@ -155,7 +155,7 @@ model_compression_toolkit/core/common/visualization/tensorboard_writer.py,sha256
 model_compression_toolkit/core/keras/__init__.py,sha256=mjbqLD-KcG3eNeCYpu1GBS7VclGVOQ63x2p6mAAuba4,698
 model_compression_toolkit/core/keras/constants.py,sha256=dh4elQWt6Q6NYRht5k5RiiOcnLAq1v0MMBCJqMJzzFk,3225
 model_compression_toolkit/core/keras/custom_layer_validation.py,sha256=f-b14wuiIgitBe7d0MmofYhDCTO3IhwJgwrh-Hq_t_U,1192
-model_compression_toolkit/core/keras/data_util.py,sha256=HQj3-GP5oT5JHpYt80mtKhZjTCvKYs6c3Ll0txEgKHQ,6892
+model_compression_toolkit/core/keras/data_util.py,sha256=-fqhXTzlA3RybWp0M5phPkzVbSJ2vPLrjFcCazWMYHk,7300
 model_compression_toolkit/core/keras/default_framework_info.py,sha256=PYcER89eEXjKtR0T7-2Y4f7cckqoD5OQbpHePoRkMec,5030
 model_compression_toolkit/core/keras/keras_implementation.py,sha256=HwbIR7x4t-TBNbWHVvVNFk8z-KFt6zM0LWAUXQuNZrk,31753
 model_compression_toolkit/core/keras/keras_model_validation.py,sha256=1wNV2clFdC9BzIELRLSO2uKf0xqjLqlkTJudwtCeaJk,1722
@@ -362,9 +362,9 @@ model_compression_toolkit/gptq/common/regularization_factory.py,sha256=hyunpXepV
 model_compression_toolkit/gptq/keras/__init__.py,sha256=cco4TmeIDIh32nj9ZZXVkws4dd9F2UDrmjKzTN8G0V0,697
 model_compression_toolkit/gptq/keras/gptq_keras_implementation.py,sha256=axBwnCSjq5xk-xGymOwSOqjp39It-CVtGcCTRTf0E_4,1248
 model_compression_toolkit/gptq/keras/gptq_loss.py,sha256=2hzWzsbuVd5XcL85NM57YeOyHxRY0qMArKn8NvQ1UWw,7643
-model_compression_toolkit/gptq/keras/gptq_training.py,sha256=0WGiP7Gs4xX3FBs1PNaZ7w3hWRigwQXqYjBrs_-x32o,23241
+model_compression_toolkit/gptq/keras/gptq_training.py,sha256=km9tcuugOkRvprGXQZrsq_GPtA3-7Du_-rnbR_Gyups,23228
 model_compression_toolkit/gptq/keras/graph_info.py,sha256=zwoeHX67nJJ5-zYLjzvMXS9TLsy9BsizARbZiDVjVSA,4473
-model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=meRKqpzZe2Irf21L_rN_mkr5dqPTJHzfSFBeqv4Csp4,18536
+model_compression_toolkit/gptq/keras/quantization_facade.py,sha256=jUAjkIszziedftaQBSmjEL6tYEYpHhlFpSgw2X9OTf4,18672
 model_compression_toolkit/gptq/keras/quantizer/__init__.py,sha256=-DK1CDXvlsnEbki4lukZLpl6Xrbo91_jcqxXlG5Eg6Q,963
 model_compression_toolkit/gptq/keras/quantizer/base_keras_gptq_quantizer.py,sha256=Rbl9urzkmACvVxICSEyJ02qFOBxWK0UQWtysFJzBVZw,4899
 model_compression_toolkit/gptq/keras/quantizer/quant_utils.py,sha256=Vt7Qb8i4JsE4sFtcjpfM4FTXTtfV1t6SwfoNH8a_Iaw,5055
@@ -380,7 +380,7 @@ model_compression_toolkit/gptq/pytorch/gptq_loss.py,sha256=_07Zx_43bnNokwR5S8phI
 model_compression_toolkit/gptq/pytorch/gptq_pytorch_implementation.py,sha256=tECPTavxn8EEwgLaP2zvxdJH6Vg9jC0YOIMJ7857Sdc,1268
 model_compression_toolkit/gptq/pytorch/gptq_training.py,sha256=WtehnyiYXdUXf8-uNpV0mdsalF7YF7eKnL7tcFrzZoE,19549
 model_compression_toolkit/gptq/pytorch/graph_info.py,sha256=4mVM-VvnBaA64ACVdOe6wTGHdMSa2UTLIUe7nACLcdo,4008
-model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=kMSq9mrpcgMBRgrEKfMBHaJG6HhGRYnuiDzF4ofckwo,16581
+model_compression_toolkit/gptq/pytorch/quantization_facade.py,sha256=HSFpx6JgjxGhU-0jA0z85sOOgSjCq6gzDOSkmuksZVE,16713
 model_compression_toolkit/gptq/pytorch/quantizer/__init__.py,sha256=ZHNHo1yzye44m9_ht4UUZfTpK01RiVR3Tr74-vtnOGI,968
 model_compression_toolkit/gptq/pytorch/quantizer/base_pytorch_gptq_quantizer.py,sha256=fKg-PNOhGBiL-4eySS9Fyw0GkA76Pq8jT_HbJuJ8iZU,4143
 model_compression_toolkit/gptq/pytorch/quantizer/quant_utils.py,sha256=OocYYRqvl7rZ37QT0hTzfJnWGiNCPskg7cziTlR7TRk,3893
@@ -525,8 +525,8 @@ model_compression_toolkit/xquant/pytorch/model_analyzer.py,sha256=b93o800yVB3Z-i
 model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py,sha256=3jNiV5Z4BVw9cEWuLKNOlLuLdr0EMuKg6eYnSiAq3LU,3952
 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.2.0.20250108.523.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.2.0.20250108.523.dist-info/METADATA,sha256=riYLks2VpIMjq7W0UIbOGVmX68cYfNGXjl04SFFNSnE,26461
-mct_nightly-2.2.0.20250108.523.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
-mct_nightly-2.2.0.20250108.523.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.2.0.20250108.523.dist-info/RECORD,,
+mct_nightly-2.2.0.20250109.528.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.2.0.20250109.528.dist-info/METADATA,sha256=AYbks8Hsbv8a3bBOMjAhG2oNxrXPfYTn8TEnOUj3KjI,26601
+mct_nightly-2.2.0.20250109.528.dist-info/WHEEL,sha256=In9FTNxeP60KnTkGw7wk6mJPYd_dQSjEZmXdBdMCI-8,91
+mct_nightly-2.2.0.20250109.528.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.2.0.20250109.528.dist-info/RECORD,,

{mct_nightly-2.2.0.20250108.523.dist-info → mct_nightly-2.2.0.20250109.528.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.7.0)
+Generator: setuptools (75.8.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

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.2.0.20250108.000523"
+__version__ = "2.2.0.20250109.000528"

model_compression_toolkit/core/keras/data_util.py

@@ -12,11 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import Generator, Callable
-
-import tensorflow as tf
-
-from model_compression_toolkit.core.keras.tf_tensor_numpy import to_tf_tensor
 
 import tensorflow as tf
 from typing import Callable, Generator, Sequence, Any
@@ -58,7 +53,6 @@ def flat_gen_fn(data_gen_fn: Callable[[], Generator]):
 
     return gen
 
-
 class TFDatasetFromGenerator:
     """
     TensorFlow dataset from a data generator function, batched to a specified size.
@@ -77,15 +71,15 @@ class TFDatasetFromGenerator:
 
         # TFDatasetFromGenerator flattens the dataset, thus we ignore the batch dimension
         output_signature = get_tensor_spec(inputs, ignore_batch_dim=True)
-        self.dataset = tf.data.Dataset.from_generator(flat_gen_fn(data_gen_fn), output_signature=output_signature)
+        self.tf_dataset = tf.data.Dataset.from_generator(flat_gen_fn(data_gen_fn), output_signature=output_signature)
 
     def __iter__(self):
-        return iter(self.dataset)
+        return iter(self.tf_dataset)
 
     def __len__(self):
         """ Returns the number of batches. """
         if self._size is None:
-            self._size = sum(1 for _ in self.dataset)
+            self._size = sum(1 for _ in self.tf_dataset)
         return self._size
 
 
@@ -116,6 +110,12 @@ class FixedTFDataset:
             raise ValueError(f'Not enough samples to create a dataset with {n_samples} samples')
         self.samples = samples
 
+        # Use from_generator to keep tuples intact
+        self.tf_dataset = tf.data.Dataset.from_generator(
+            lambda: iter(self.samples),
+            output_signature=tuple(tf.TensorSpec(shape=sample.shape, dtype=sample.dtype) for sample in self.samples[0])
+        )
+
     def __len__(self):
         return len(self.samples)
 
@@ -134,6 +134,12 @@ class FixedSampleInfoDataset:
         self.samples = samples
         self.sample_info = sample_info
 
+        # Create a TensorFlow dataset that holds (sample, sample_info) tuples
+        self.tf_dataset = tf.data.Dataset.from_tensor_slices((
+            tf.convert_to_tensor(self.samples),
+            tuple(tf.convert_to_tensor(info) for info in self.sample_info)
+        ))
+
     def __len__(self):
         return len(self.samples)
 
@@ -150,18 +156,23 @@ class IterableSampleWithConstInfoDataset:
         self.samples_dataset = samples_dataset
         self.info = info
 
+        # Map to ensure the output is always (sample, info) as a tuple
+        self.tf_dataset = self.samples_dataset.map(
+            lambda *x: ((x,) if not isinstance(x, tuple) else x, *self.info)
+        )
+
     def __iter__(self):
         for sample in self.samples_dataset:
-            yield (sample, *self.info)
+            yield ((sample,) if not isinstance(sample, tuple) else sample, *self.info)
 
 
 def data_gen_to_dataloader(data_gen_fn: Callable[[], Generator], batch_size: int):
     """Create a DataLoader based on samples yielded by data_gen."""
     ds = TFDatasetFromGenerator(data_gen_fn)
-    return create_tf_dataloader(dataset=ds, batch_size=batch_size)
+    return create_tf_dataloader(mct_dataset=ds, batch_size=batch_size)
 
 
-def create_tf_dataloader(dataset, batch_size, shuffle=False, collate_fn=None):
+def create_tf_dataloader(mct_dataset, batch_size, shuffle=False, collate_fn=None):
     """
     Creates a tf.data.Dataset with specified loading options.
 
@@ -174,26 +185,15 @@ def create_tf_dataloader(dataset, batch_size, shuffle=False, collate_fn=None):
     Returns:
         tf.data.Dataset: Configured for batching, shuffling, and custom transformations.
     """
-    def generator():
-        for item in dataset:
-            yield item
-
-    dummy_input_tensors = next(generator())
-
-    output_signature = get_tensor_spec(dummy_input_tensors)
-
-    tf_dataset = tf.data.Dataset.from_generator(
-        generator,
-        output_signature=output_signature
-    )
+    dataset = mct_dataset.tf_dataset
 
     if shuffle:
-        tf_dataset = tf_dataset.shuffle(buffer_size=len(dataset))
+        dataset = dataset.shuffle(buffer_size=len(dataset))
 
-    tf_dataset = tf_dataset.batch(batch_size)
+    dataset = dataset.batch(batch_size)
 
     # Apply collate function if provided
     if collate_fn:
-        tf_dataset = tf_dataset.map(lambda *args: collate_fn(args))
+        dataset = dataset.map(lambda *args: collate_fn(args))
 
-    return tf_dataset
+    return dataset

model_compression_toolkit/gptq/keras/gptq_training.py

@@ -151,7 +151,7 @@ class KerasGPTQTrainer(GPTQTrainer):
 
         # Create final dataset using the new dataloader with collate_fn
         final_dataset = create_tf_dataloader(
-            dataset=sla_train_dataset,
+            sla_train_dataset,
             batch_size=orig_batch_size,
             shuffle=True,
             collate_fn=collate_fn
@@ -176,14 +176,14 @@ class KerasGPTQTrainer(GPTQTrainer):
 
         # Step 2: Compute loss weights
         if self.gptq_config.hessian_weights_config:
-            hessian_dataset = create_tf_dataloader(dataset=dataset, batch_size=self.gptq_config.hessian_weights_config.hessian_batch_size)
+            hessian_dataset = create_tf_dataloader(dataset, batch_size=self.gptq_config.hessian_weights_config.hessian_batch_size)
             hessian_weights = self.compute_hessian_based_weights(hessian_dataset)
             loss_weights = tf.convert_to_tensor(hessian_weights, dtype=tf.float32)
         else:
             loss_weights = tf.ones(num_nodes, dtype=tf.float32) / num_nodes
 
         # Step 3: Create a dataset with samples and loss weights
-        augmented_dataset = IterableSampleWithConstInfoDataset(dataset.dataset, loss_weights)
+        augmented_dataset = IterableSampleWithConstInfoDataset(dataset.tf_dataset, loss_weights)
 
         # Step 4: Add constant regularization weights
         reg_weights = tf.ones(num_nodes, dtype=tf.float32)

model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -115,7 +115,6 @@ if FOUND_TF:
         if regularization_factor is None:
             regularization_factor = REG_DEFAULT_SLA if use_hessian_sample_attention else REG_DEFAULT
 
-        loss = loss or GPTQMultipleTensorsLoss()
         hessian_weights_config = None
         if use_hessian_sample_attention:
             if not use_hessian_based_weights:    # pragma: no cover
@@ -129,7 +128,10 @@ if FOUND_TF:
             hessian_weights_config = GPTQHessianScoresConfig(per_sample=False,
                                                              hessians_num_samples=GPTQ_HESSIAN_NUM_SAMPLES,
                                                              hessian_batch_size=hessian_batch_size)
-
+        
+        # If a loss was not passed (and was not initialized due to use_hessian_sample_attention), use the default loss
+        loss = loss or GPTQMultipleTensorsLoss()
+        
         if isinstance(gradual_activation_quantization, bool):
             gradual_quant_config = GradualActivationQuantizationConfig() if gradual_activation_quantization else None
         elif isinstance(gradual_activation_quantization, GradualActivationQuantizationConfig):

model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -104,7 +104,6 @@ if FOUND_TORCH:
         if regularization_factor is None:
             regularization_factor = REG_DEFAULT_SLA if use_hessian_sample_attention else REG_DEFAULT
 
-        loss = loss or multiple_tensors_mse_loss
         hessian_weights_config = None
         if use_hessian_sample_attention:
             if not use_hessian_based_weights:    # pragma: no cover
@@ -118,6 +117,9 @@ if FOUND_TORCH:
             hessian_weights_config = GPTQHessianScoresConfig(per_sample=False,
                                                              hessians_num_samples=GPTQ_HESSIAN_NUM_SAMPLES,
                                                              hessian_batch_size=hessian_batch_size)
+            
+        # If a loss was not passed (and was not initialized due to use_hessian_sample_attention), use the default loss
+        loss = loss or multiple_tensors_mse_loss
 
         if isinstance(gradual_activation_quantization, bool):
             gradual_quant_config = GradualActivationQuantizationConfig() if gradual_activation_quantization else None