mct-nightly 2.2.0.20250107.164940__py3-none-any.whl → 2.2.0.20250108.130055__py3-none-any.whl

{mct_nightly-2.2.0.20250107.164940.dist-info → mct_nightly-2.2.0.20250108.130055.dist-info}/METADATA

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

{mct_nightly-2.2.0.20250107.164940.dist-info → mct_nightly-2.2.0.20250108.130055.dist-info}/RECORD

@@ -1,4 +1,4 @@
-model_compression_toolkit/__init__.py,sha256=A5js8ho45WYpU_ZDFPmwshjYKWICZlIgdFrShfdEJGo,1573
+model_compression_toolkit/__init__.py,sha256=Lgj_ZJtNyscR_iDVY0C4uiaxzLuXK-jUDeqFW3fTxjI,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
@@ -361,8 +361,8 @@ model_compression_toolkit/gptq/common/gradual_activation_quantization.py,sha256=
 model_compression_toolkit/gptq/common/regularization_factory.py,sha256=hyunpXepVeHyoAFJw6zNLK-3ZHBmiut3lmNisJN_L3E,2514
 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=k5s7D4CTqbYCHgydyevw1c2p3S2TZCECHNvK79QGE2U,7797
-model_compression_toolkit/gptq/keras/gptq_training.py,sha256=0WGiP7Gs4xX3FBs1PNaZ7w3hWRigwQXqYjBrs_-x32o,23241
+model_compression_toolkit/gptq/keras/gptq_loss.py,sha256=2hzWzsbuVd5XcL85NM57YeOyHxRY0qMArKn8NvQ1UWw,7643
+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=jUAjkIszziedftaQBSmjEL6tYEYpHhlFpSgw2X9OTf4,18672
 model_compression_toolkit/gptq/keras/quantizer/__init__.py,sha256=-DK1CDXvlsnEbki4lukZLpl6Xrbo91_jcqxXlG5Eg6Q,963
@@ -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.20250107.164940.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
-mct_nightly-2.2.0.20250107.164940.dist-info/METADATA,sha256=9IxOhXh5nRY6Ck-0wINZ4huD8os58zC-X_15pdsxHCo,26464
-mct_nightly-2.2.0.20250107.164940.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
-mct_nightly-2.2.0.20250107.164940.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
-mct_nightly-2.2.0.20250107.164940.dist-info/RECORD,,
+mct_nightly-2.2.0.20250108.130055.dist-info/LICENSE.md,sha256=aYSSIb-5AFPeITTvXm1UAoe0uYBiMmSS8flvXaaFUks,10174
+mct_nightly-2.2.0.20250108.130055.dist-info/METADATA,sha256=kkALabHdxcfL12ANa1E1K_iS5fS1KluedrBW4p39YcA,26464
+mct_nightly-2.2.0.20250108.130055.dist-info/WHEEL,sha256=A3WOREP4zgxI0fKrHUG8DC8013e3dK3n7a6HDbcEIwE,91
+mct_nightly-2.2.0.20250108.130055.dist-info/top_level.txt,sha256=gsYA8juk0Z-ZmQRKULkb3JLGdOdz8jW_cMRjisn9ga4,26
+mct_nightly-2.2.0.20250108.130055.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.2.0.20250107.164940"
+__version__ = "2.2.0.20250108.130055"

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_loss.py

@@ -144,9 +144,8 @@ def activation_mse(flp_act_list,
         loss_values_list.append(point_loss)
         bias_loss_list.append(bias_loss)
     if weights_for_average_loss is not None:
-        print(f"weights_for_average_loss.shape: {weights_for_average_loss.shape}")
-        print(f"tf.stack(loss_values_list).shape: {tf.stack(loss_values_list).shape}")
-        return tf.reduce_sum(weights_for_average_loss * tf.stack(loss_values_list)), tf.reduce_mean(tf.stack(bias_loss_list))
+        return tf.reduce_sum(weights_for_average_loss * tf.stack(loss_values_list)), \
+               tf.reduce_mean(tf.stack(bias_loss_list))
     else:
         return tf.reduce_mean(tf.stack(loss_values_list)), tf.reduce_mean(tf.stack(bias_loss_list))
 
@@ -188,4 +187,4 @@ class GPTQMultipleTensorsLoss:
                                                         weights_for_average_loss=weights_for_average_loss,
                                                         norm_loss=self.norm_loss)
 
-        return loss_act
+        return loss_act

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)