mct-nightly 2.4.0.20250617.613__py3-none-any.whl → 2.4.0.20250619.621__py3-none-any.whl

model_compression_toolkit/pruning/keras/pruning_facade.py

@@ -35,11 +35,12 @@ if FOUND_TF:
         AttachTpcToKeras
     from model_compression_toolkit.core.keras.back2framework.float_model_builder import FloatKerasModelBuilder
     from model_compression_toolkit.core.keras.pruning.pruning_keras_implementation import PruningKerasImplementation
-    from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
+    from model_compression_toolkit.core.keras.default_framework_info import set_keras_info
     from tensorflow.keras.models import Model
 
     DEFAULT_KERAS_TPC = get_target_platform_capabilities(TENSORFLOW, DEFAULT_TP_MODEL)
 
+    @set_keras_info
     def keras_pruning_experimental(model: Model,
                                    target_resource_utilization: ResourceUtilization,
                                    representative_data_gen: Callable,
@@ -123,7 +124,6 @@ if FOUND_TF:
         float_graph = read_model_to_graph(model,
                                           representative_data_gen,
                                           target_platform_capabilities,
-                                          DEFAULT_KERAS_INFO,
                                           fw_impl)
 
         # Apply quantization configuration to the graph. This step is necessary even when not quantizing,
@@ -134,7 +134,6 @@ if FOUND_TF:
 
         # Create a Pruner object with the graph and configuration.
         pruner = Pruner(float_graph_with_compression_config,
-                        DEFAULT_KERAS_INFO,
                         fw_impl,
                         target_resource_utilization,
                         representative_data_gen,

model_compression_toolkit/pruning/pytorch/pruning_facade.py

@@ -36,7 +36,7 @@ if FOUND_TORCH:
     from model_compression_toolkit.core.pytorch.back2framework.float_model_builder import FloatPyTorchModelBuilder
     from model_compression_toolkit.core.pytorch.pruning.pruning_pytorch_implementation import \
         PruningPytorchImplementation
-    from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
+    from model_compression_toolkit.core.pytorch.default_framework_info import set_pytorch_info
     from torch.nn import Module
     from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2pytorch import \
         AttachTpcToPytorch
@@ -44,6 +44,7 @@ if FOUND_TORCH:
     # Set the default Target Platform Capabilities (TPC) for PyTorch.
     DEFAULT_PYOTRCH_TPC = get_target_platform_capabilities(PYTORCH, DEFAULT_TP_MODEL)
 
+    @set_pytorch_info
     def pytorch_pruning_experimental(model: Module,
                                      target_resource_utilization: ResourceUtilization,
                                      representative_data_gen: Callable,
@@ -129,7 +130,6 @@ if FOUND_TORCH:
         float_graph = read_model_to_graph(model,
                                           representative_data_gen,
                                           framework_platform_capabilities,
-                                          DEFAULT_PYTORCH_INFO,
                                           fw_impl)
 
         # Apply quantization configuration to the graph. This step is necessary even when not quantizing,
@@ -140,7 +140,6 @@ if FOUND_TORCH:
 
         # Create a Pruner object with the graph and configuration.
         pruner = Pruner(float_graph_with_compression_config,
-                        DEFAULT_PYTORCH_INFO,
                         fw_impl,
                         target_resource_utilization,
                         representative_data_gen,

model_compression_toolkit/ptq/keras/quantization_facade.py

@@ -36,7 +36,7 @@ from model_compression_toolkit.metadata import create_model_metadata
 if FOUND_TF:
     from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2keras import \
         AttachTpcToKeras
-    from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
+    from model_compression_toolkit.core.keras.default_framework_info import set_keras_info
     from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
     from model_compression_toolkit.core.keras.keras_model_validation import KerasModelValidation
     from tensorflow.keras.models import Model
@@ -49,6 +49,7 @@ if FOUND_TF:
     DEFAULT_KERAS_TPC = get_target_platform_capabilities(TENSORFLOW, DEFAULT_TP_MODEL)
 
 
+    @set_keras_info
     def keras_post_training_quantization(in_model: Model,
                                          representative_data_gen: Callable,
                                          target_resource_utilization: ResourceUtilization = None,
@@ -128,10 +129,7 @@ if FOUND_TF:
         if core_config.debug_config.bypass:
             return in_model, None
 
-        fw_info = DEFAULT_KERAS_INFO
-
-        KerasModelValidation(model=in_model,
-                             fw_info=fw_info).validate()
+        KerasModelValidation(model=in_model).validate()
 
         if core_config.is_mixed_precision_enabled:
             if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
@@ -139,7 +137,7 @@ if FOUND_TF:
                                 "MixedPrecisionQuantizationConfig. Please use keras_post_training_quantization "
                                 "API, or pass a valid mixed precision configuration.")  # pragma: no cover
 
-        tb_w = init_tensorboard_writer(fw_info)
+        tb_w = init_tensorboard_writer()
 
         fw_impl = KerasImplementation()
 
@@ -153,7 +151,6 @@ if FOUND_TF:
         tg, bit_widths_config, _, scheduling_info = core_runner(in_model=in_model,
                                                                 representative_data_gen=representative_data_gen,
                                                                 core_config=core_config,
-                                                                fw_info=fw_info,
                                                                 fw_impl=fw_impl,
                                                                 fqc=framework_platform_capabilities,
                                                                 target_resource_utilization=target_resource_utilization,
@@ -169,7 +166,6 @@ if FOUND_TF:
         graph_with_stats_correction = ptq_runner(tg,
                                                  representative_data_gen,
                                                  core_config,
-                                                 fw_info,
                                                  fw_impl,
                                                  tb_w)
 
@@ -179,8 +175,7 @@ if FOUND_TF:
                                         tb_w,
                                         similarity_baseline_graph,
                                         quantized_graph,
-                                        fw_impl,
-                                        fw_info)
+                                        fw_impl)
 
         exportable_model, user_info = get_exportable_keras_model(graph_with_stats_correction)
         if framework_platform_capabilities.tpc.add_metadata:

model_compression_toolkit/ptq/pytorch/quantization_facade.py

@@ -34,7 +34,7 @@ from model_compression_toolkit.core.common.quantization.quantize_graph_weights i
 from model_compression_toolkit.metadata import create_model_metadata
 
 if FOUND_TORCH:
-    from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
+    from model_compression_toolkit.core.pytorch.default_framework_info import set_pytorch_info
     from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from torch.nn import Module
@@ -46,6 +46,7 @@ if FOUND_TORCH:
 
     DEFAULT_PYTORCH_TPC = get_target_platform_capabilities(PYTORCH, DEFAULT_TP_MODEL)
 
+    @set_pytorch_info
     def pytorch_post_training_quantization(in_module: Module,
                                            representative_data_gen: Callable,
                                            target_resource_utilization: ResourceUtilization = None,
@@ -102,8 +103,6 @@ if FOUND_TORCH:
         if core_config.debug_config.bypass:
             return in_module, None
 
-        fw_info = DEFAULT_PYTORCH_INFO
-
         if core_config.is_mixed_precision_enabled:
             if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
                 Logger.critical("Given quantization config to mixed-precision facade is not of type "
@@ -111,7 +110,7 @@ if FOUND_TORCH:
                                 "pytorch_post_training_quantization API, or pass a valid mixed precision "
                                 "configuration.")  # pragma: no cover
 
-        tb_w = init_tensorboard_writer(fw_info)
+        tb_w = init_tensorboard_writer()
 
         fw_impl = PytorchImplementation()
 
@@ -125,7 +124,6 @@ if FOUND_TORCH:
         tg, bit_widths_config, _, scheduling_info = core_runner(in_model=in_module,
                                                                 representative_data_gen=representative_data_gen,
                                                                 core_config=core_config,
-                                                                fw_info=fw_info,
                                                                 fw_impl=fw_impl,
                                                                 fqc=framework_platform_capabilities,
                                                                 target_resource_utilization=target_resource_utilization,
@@ -141,7 +139,6 @@ if FOUND_TORCH:
         graph_with_stats_correction = ptq_runner(tg,
                                                  representative_data_gen,
                                                  core_config,
-                                                 fw_info,
                                                  fw_impl,
                                                  tb_w)
 
@@ -151,8 +148,7 @@ if FOUND_TORCH:
                                         tb_w,
                                         similarity_baseline_graph,
                                         quantized_graph,
-                                        fw_impl,
-                                        fw_info)
+                                        fw_impl)
 
         exportable_model, user_info = get_exportable_pytorch_model(graph_with_stats_correction)
         if framework_platform_capabilities.tpc.add_metadata:

model_compression_toolkit/ptq/runner.py

@@ -16,7 +16,6 @@
 
 from typing import Callable
 
-from model_compression_toolkit.core.common import FrameworkInfo
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
@@ -28,7 +27,6 @@ from model_compression_toolkit.core.common.visualization.tensorboard_writer impo
 def ptq_runner(tg: Graph,
                representative_data_gen: Callable,
                core_config: CoreConfig,
-               fw_info: FrameworkInfo,
                fw_impl: FrameworkImplementation,
                tb_w: TensorboardWriter) -> Graph:
     """
@@ -38,7 +36,6 @@ def ptq_runner(tg: Graph,
         tg: Graph to apply PTQ and to quantize.
         representative_data_gen (Callable): Dataset used for calibration.
         core_config: CoreConfig containing parameters of how the model should be quantized.
-        fw_info: Information needed for quantization about the specific framework (e.g., kernel channels indices,
         groups of layers by how they should be quantized, etc.)
         fw_impl: FrameworkImplementation object with a specific framework methods implementation.
         tb_w: A TensorBoardWriter object initialized with the logger dir path if it was set, or None otherwise.
@@ -50,5 +47,5 @@ def ptq_runner(tg: Graph,
     #############################################
     # Statistics Correction
     #############################################
-    tg = apply_statistics_correction(tg, representative_data_gen, core_config, fw_info, fw_impl, tb_w)
+    tg = apply_statistics_correction(tg, representative_data_gen, core_config, fw_impl, tb_w)
     return tg

model_compression_toolkit/qat/common/qat_config.py

@@ -19,21 +19,17 @@ from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.trainable_infrastructure import TrainingMethod
 
 
-def is_qat_applicable(node: common.BaseNode,
-                      fw_info: FrameworkInfo) -> bool:
+def is_qat_applicable(node: common.BaseNode) -> bool:
     """
     A function for deciding if a layer should be fine-tuned during QAT
 
     Args:
         node (BaseNode): Node for quantization decision
-        fw_info (FrameworkInfo): Pytorch quantization information
 
     Returns:
         A boolean whether the layer is to be wrapped with a QuantizeWrapper
     """
-
-    kernel_attr = fw_info.get_kernel_op_attributes(node.type)[0]
-    return (kernel_attr is not None and node.is_weights_quantization_enabled(kernel_attr)) \
+    return (node.kernel_attr is not None and node.is_weights_quantization_enabled(node.kernel_attr)) \
             or node.is_activation_quantization_enabled()
 
 

model_compression_toolkit/qat/keras/quantization_facade.py

@@ -37,10 +37,10 @@ if FOUND_TF:
     from tensorflow.keras.models import Model
 
     from model_compression_toolkit.trainable_infrastructure import KerasTrainableQuantizationWrapper
-    from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
     from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
     from model_compression_toolkit.core.keras.keras_model_validation import KerasModelValidation
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
+    from model_compression_toolkit.core.keras.default_framework_info import set_keras_info
 
     from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder
 
@@ -52,7 +52,6 @@ if FOUND_TF:
     from model_compression_toolkit.constants import TENSORFLOW
     from model_compression_toolkit.qat.common.qat_config import is_qat_applicable
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
-    from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
     from model_compression_toolkit.qat.keras.quantizer.quantization_builder import quantization_builder, \
     get_activation_quantizer_holder
     from model_compression_toolkit.qat.common.qat_config import QATConfig
@@ -73,11 +72,11 @@ if FOUND_TF:
         Returns: Wrapped layer
 
         """
-        if is_qat_applicable(n, DEFAULT_KERAS_INFO):
+        if is_qat_applicable(n):
             # If we are here, then the node has a kernel attribute to quantize and training during QAT
             weights_quantizers, _ = quantization_builder(n,
                                                          qat_config,
-                                                         DEFAULT_KERAS_INFO.get_kernel_op_attributes(n.type)[0])
+                                                         n.kernel_attr)
             if len(weights_quantizers) > 0:
                 layer.trainable = True
                 return KerasTrainableQuantizationWrapper(layer, weights_quantizers)
@@ -87,6 +86,7 @@ if FOUND_TF:
         return layer
 
 
+    @set_keras_info
     def keras_quantization_aware_training_init_experimental(in_model: Model,
                                                             representative_data_gen: Callable,
                                                             target_resource_utilization: ResourceUtilization = None,
@@ -175,8 +175,7 @@ if FOUND_TF:
                        f"If you encounter an issue, please open an issue in our GitHub "
                        f"project https://github.com/sony/model_optimization")
 
-        KerasModelValidation(model=in_model,
-                             fw_info=DEFAULT_KERAS_INFO).validate()
+        KerasModelValidation(model=in_model).validate()
 
         if core_config.is_mixed_precision_enabled:
             if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
@@ -184,7 +183,7 @@ if FOUND_TF:
                              "MixedPrecisionQuantizationConfig. Please use keras_post_training_quantization API,"
                              "or pass a valid mixed precision configuration.")
 
-        tb_w = init_tensorboard_writer(DEFAULT_KERAS_INFO)
+        tb_w = init_tensorboard_writer()
 
         fw_impl = KerasImplementation()
 
@@ -198,17 +197,15 @@ if FOUND_TF:
         tg, bit_widths_config, _, _ = core_runner(in_model=in_model,
                                                   representative_data_gen=representative_data_gen,
                                                   core_config=core_config,
-                                                  fw_info=DEFAULT_KERAS_INFO,
                                                   fw_impl=fw_impl,
                                                   fqc=target_platform_capabilities,
                                                   target_resource_utilization=target_resource_utilization,
                                                   tb_w=tb_w)
 
-        tg = ptq_runner(tg, representative_data_gen, core_config, DEFAULT_KERAS_INFO, fw_impl, tb_w)
+        tg = ptq_runner(tg, representative_data_gen, core_config, fw_impl, tb_w)
 
         _qat_wrapper = partial(qat_wrapper, qat_config=qat_config)
         qat_model, user_info = KerasModelBuilder(graph=tg,
-                                                 fw_info=DEFAULT_KERAS_INFO,
                                                  wrapper=_qat_wrapper,
                                                  get_activation_quantizer_holder_fn=partial(get_activation_quantizer_holder,
                                                                                             qat_config=qat_config)).build_model()

model_compression_toolkit/qat/pytorch/quantization_facade.py

@@ -36,7 +36,7 @@ if FOUND_TORCH:
     import torch.nn as nn
     from torch.nn import Module
     from mct_quantizers import PytorchActivationQuantizationHolder
-    from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
+    from model_compression_toolkit.core.pytorch.default_framework_info import set_pytorch_info
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
     from model_compression_toolkit.qat.common.qat_config import is_qat_applicable
@@ -62,10 +62,10 @@ if FOUND_TORCH:
         Returns: Wrapped layer
 
         """
-        if is_qat_applicable(n, DEFAULT_PYTORCH_INFO):
+        if is_qat_applicable(n):
             # If we are here, then the node has a kernel attribute to quantize and training during QAT
             weights_quantizers, _ = quantization_builder(n, qat_config,
-                                                         DEFAULT_PYTORCH_INFO.get_kernel_op_attributes(n.type)[0])
+                                                         n.kernel_attr)
             if len(weights_quantizers) > 0:
                 return PytorchQuantizationWrapper(module, weights_quantizers)
 
@@ -74,6 +74,7 @@ if FOUND_TORCH:
         return module
 
 
+    @set_pytorch_info
     def pytorch_quantization_aware_training_init_experimental(in_model: Module,
                                                               representative_data_gen: Callable,
                                                               target_resource_utilization: ResourceUtilization = None,
@@ -149,7 +150,7 @@ if FOUND_TORCH:
                                 "MixedPrecisionQuantizationConfig. Please use pytorch_post_training_quantization API,"
                                 "or pass a valid mixed precision configuration.")
 
-        tb_w = init_tensorboard_writer(DEFAULT_PYTORCH_INFO)
+        tb_w = init_tensorboard_writer()
         fw_impl = PytorchImplementation()
 
         target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
@@ -162,18 +163,16 @@ if FOUND_TORCH:
         tg, bit_widths_config, _, _ = core_runner(in_model=in_model,
                                                   representative_data_gen=representative_data_gen,
                                                   core_config=core_config,
-                                                  fw_info=DEFAULT_PYTORCH_INFO,
                                                   fw_impl=fw_impl,
                                                   fqc=framework_platform_capabilities,
                                                   target_resource_utilization=target_resource_utilization,
                                                   tb_w=tb_w)
 
-        tg = ptq_runner(tg, representative_data_gen, core_config, DEFAULT_PYTORCH_INFO, fw_impl, tb_w)
+        tg = ptq_runner(tg, representative_data_gen, core_config, fw_impl, tb_w)
 
         _qat_wrapper = partial(qat_wrapper, qat_config=qat_config)
 
         qat_model, user_info = PyTorchModelBuilder(graph=tg,
-                                                   fw_info=DEFAULT_PYTORCH_INFO,
                                                    wrapper=_qat_wrapper,
                                                    get_activation_quantizer_holder_fn=partial(
                                                        get_activation_quantizer_holder,
@@ -181,9 +180,6 @@ if FOUND_TORCH:
 
         user_info.mixed_precision_cfg = bit_widths_config
 
-        # Remove fw_info from graph to enable saving the pytorch model (fw_info can not be pickled)
-        delattr(qat_model.graph, 'fw_info')
-
         return qat_model, user_info
 
 

model_compression_toolkit/xquant/common/core_report_generator.py

@@ -50,7 +50,7 @@ def core_report_generator(float_model: Any,
 
     # Collect histograms on the float model.
     float_graph = fw_report_utils.model_folding_utils.create_float_folded_graph(float_model, repr_dataset)
-    mi = ModelCollector(float_graph, fw_report_utils.fw_impl, fw_report_utils.fw_info)
+    mi = ModelCollector(float_graph, fw_report_utils.fw_impl)
     for _data in tqdm(repr_dataset(), desc="Collecting Histograms"):
         mi.infer(_data)
 

model_compression_toolkit/xquant/common/framework_report_utils.py

@@ -34,7 +34,6 @@ class FrameworkReportUtils:
     """
 
     def __init__(self,
-                 fw_info: FrameworkInfo,
                  fw_impl: FrameworkImplementation,
                  similarity_calculator: SimilarityCalculator,
                  dataset_utils: DatasetUtils,
@@ -45,7 +44,6 @@ class FrameworkReportUtils:
         Initializes the FrameworkReportUtils class with various utility components required for generating the report.
 
         Args:
-            fw_info (FrameworkInfo): Information about the framework being used.
             fw_impl (FrameworkImplementation): The implemented functions of the framework.
             similarity_calculator (SimilarityCalculator): A utility for calculating similarity metrics.
             dataset_utils (DatasetUtils): Utilities for handling datasets.
@@ -53,7 +51,6 @@ class FrameworkReportUtils:
             tb_utils (TensorboardUtils): Utilities for TensorBoard operations.
             get_metadata_fn (Callable): Function to retrieve the metadata from the quantized model.
         """
-        self.fw_info = fw_info
         self.fw_impl = fw_impl
         self.similarity_calculator = similarity_calculator
         self.dataset_utils = dataset_utils

model_compression_toolkit/xquant/common/model_folding_utils.py

@@ -34,7 +34,6 @@ class ModelFoldingUtils:
     """
 
     def __init__(self,
-                 fw_info: FrameworkInfo,
                  fw_impl: FrameworkImplementation,
                  fw_default_fqc: FrameworkQuantizationCapabilities):
         """
@@ -42,11 +41,9 @@ class ModelFoldingUtils:
         and default FQC.
 
         Args:
-            fw_info: Framework-specific information.
             fw_impl: Implementation functions for the framework.
             fw_default_fqc: Default target platform capabilities for the handled framework.
         """
-        self.fw_info = fw_info
         self.fw_impl = fw_impl
         self.fw_default_fqc = fw_default_fqc
 
@@ -69,8 +66,7 @@ class ModelFoldingUtils:
         float_folded_model, _ = self.fw_impl.model_builder(
             float_graph,
             mode=ModelBuilderMode.FLOAT,
-            append2output=None,
-            fw_info=self.fw_info
+            append2output=None
         )
         return float_folded_model
 
@@ -100,7 +96,6 @@ class ModelFoldingUtils:
         graph = graph_preparation_runner(in_model=model,
                                          representative_data_gen=repr_dataset,
                                          fw_impl=self.fw_impl,
-                                         fw_info=self.fw_info,
                                          quantization_config=DEFAULTCONFIG,
                                          fqc=self.fw_default_fqc)
         return graph

model_compression_toolkit/xquant/common/tensorboard_utils.py

@@ -36,19 +36,16 @@ class TensorboardUtils:
 
     def __init__(self,
                  report_dir: str,
-                 fw_info: FrameworkInfo,
                  fw_impl: FrameworkImplementation):
         """
         Initialize the TensorboardUtils.
 
         Args:
             report_dir (str): Directory where Tensorboard logs will be stored.
-            fw_info (FrameworkInfo): Framework-specific information.
             fw_impl (FrameworkImplementation): Framework-specific implementation.
         """
         self.fw_impl = fw_impl
-        self.fw_info = fw_info
-        self.tb_writer = TensorboardWriter(report_dir, fw_info)
+        self.tb_writer = TensorboardWriter(report_dir)
         Logger.info(f"Please run: tensorboard --logdir {self.tb_writer.dir_path}")
 
     def get_graph_for_tensorboard_display(self,

model_compression_toolkit/xquant/keras/keras_report_utils.py

@@ -15,7 +15,6 @@
 
 from model_compression_toolkit import get_target_platform_capabilities
 from model_compression_toolkit.constants import TENSORFLOW
-from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2keras import \
     AttachTpcToKeras
@@ -40,7 +39,6 @@ class KerasReportUtils(FrameworkReportUtils):
         Args:
             report_dir: Logging dir path.
         """
-        fw_info = DEFAULT_KERAS_INFO
         fw_impl = KerasImplementation()
 
         # Set the default Target Platform Capabilities (TPC) for Keras.
@@ -49,8 +47,7 @@ class KerasReportUtils(FrameworkReportUtils):
         framework_platform_capabilities = attach2pytorch.attach(default_tpc)
 
         dataset_utils = KerasDatasetUtils()
-        model_folding = ModelFoldingUtils(fw_info=fw_info,
-                                          fw_impl=fw_impl,
+        model_folding = ModelFoldingUtils(fw_impl=fw_impl,
                                           fw_default_fqc=framework_platform_capabilities)
 
         similarity_calculator = SimilarityCalculator(dataset_utils=dataset_utils,
@@ -59,10 +56,8 @@ class KerasReportUtils(FrameworkReportUtils):
                                                      model_analyzer_utils=KerasModelAnalyzer())
 
         tb_utils = KerasTensorboardUtils(report_dir=report_dir,
-                                         fw_impl=fw_impl,
-                                         fw_info=fw_info)
-        super().__init__(fw_info,
-                         fw_impl,
+                                         fw_impl=fw_impl)
+        super().__init__(fw_impl,
                          similarity_calculator,
                          dataset_utils,
                          model_folding,

model_compression_toolkit/xquant/keras/tensorboard_utils.py

@@ -40,18 +40,15 @@ class KerasTensorboardUtils(TensorboardUtils):
     """
 
     def __init__(self, report_dir: str,
-                 fw_info: FrameworkInfo,
                  fw_impl: FrameworkImplementation):
         """
         Initialize the KerasTensorboardUtils class with the given parameters.
 
         Args:
             report_dir (str): Directory where the TensorBoard files will be stored.
-            fw_info (FrameworkInfo): Information about the framework being used.
             fw_impl (FrameworkImplementation): Implementation functions for the framework.
         """
         super().__init__(report_dir,
-                         fw_info,
                          fw_impl)
 
     def get_graph_for_tensorboard_display(self,

model_compression_toolkit/xquant/pytorch/pytorch_report_utils.py

@@ -20,7 +20,7 @@ from model_compression_toolkit.target_platform_capabilities.targetplatform2frame
     AttachTpcToPytorch
 
 from model_compression_toolkit.xquant.common.framework_report_utils import FrameworkReportUtils
-from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
+from model_compression_toolkit.core.pytorch.default_framework_info import PyTorchInfo
 from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
 from model_compression_toolkit.xquant.common.model_folding_utils import ModelFoldingUtils
 from model_compression_toolkit.xquant.common.similarity_calculator import SimilarityCalculator
@@ -30,6 +30,7 @@ from model_compression_toolkit.xquant.pytorch.similarity_functions import Pytorc
 from model_compression_toolkit.xquant.pytorch.tensorboard_utils import PytorchTensorboardUtils
 from mct_quantizers.pytorch.metadata import get_metadata
 
+
 class PytorchReportUtils(FrameworkReportUtils):
     """
     Class with various utility components required for generating the report for a Pytorch model.
@@ -39,7 +40,6 @@ class PytorchReportUtils(FrameworkReportUtils):
         Args:
             report_dir: Logging dir path.
         """
-        fw_info = DEFAULT_PYTORCH_INFO
         fw_impl = PytorchImplementation()
         # Set the default Target Platform Capabilities (TPC) for PyTorch.
         default_tpc = get_target_platform_capabilities(PYTORCH, DEFAULT_TP_MODEL)
@@ -47,8 +47,7 @@ class PytorchReportUtils(FrameworkReportUtils):
         framework_quantization_capabilities = attach2pytorch.attach(default_tpc)
 
         dataset_utils = PytorchDatasetUtils()
-        model_folding = ModelFoldingUtils(fw_info=fw_info,
-                                          fw_impl=fw_impl,
+        model_folding = ModelFoldingUtils(fw_impl=fw_impl,
                                           fw_default_fqc=framework_quantization_capabilities)
 
         similarity_calculator = SimilarityCalculator(dataset_utils=dataset_utils,
@@ -58,11 +57,9 @@ class PytorchReportUtils(FrameworkReportUtils):
                                                      device=get_working_device())
 
         tb_utils = PytorchTensorboardUtils(report_dir=report_dir,
-                                           fw_impl=fw_impl,
-                                           fw_info=fw_info)
+                                           fw_impl=fw_impl)
 
-        super().__init__(fw_info=fw_info,
-                         fw_impl=fw_impl,
+        super().__init__(fw_impl=fw_impl,
                          tb_utils=tb_utils,
                          dataset_utils=dataset_utils,
                          similarity_calculator=similarity_calculator,

model_compression_toolkit/xquant/pytorch/tensorboard_utils.py

@@ -41,18 +41,15 @@ class PytorchTensorboardUtils(TensorboardUtils):
 
     def __init__(self,
                  report_dir: str,
-                 fw_info: FrameworkInfo,
                  fw_impl: FrameworkImplementation):
         """
         Initialize the PytorchTensorboardUtils instance.
 
         Args:
             report_dir: Directory where the reports are stored.
-            fw_info: Information about the framework being used.
             fw_impl: Implementation methods for the framework.
         """
         super().__init__(report_dir,
-                         fw_info,
                          fw_impl)
 
     def get_graph_for_tensorboard_display(self,