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

model_compression_toolkit/core/keras/graph_substitutions/substitutions/shift_negative_activation.py

@@ -227,15 +227,13 @@ def is_padding_node_and_node_has_padding(pad_node_to_consider: BaseNode,
 
 
 def keras_apply_shift_negative_correction(graph: Graph,
-                                          core_config: CoreConfig,
-                                          fw_info: FrameworkInfo) -> Graph:
+                                          core_config: CoreConfig) -> Graph:
     """
     Apply shift negative correction (SNC) on a graph built from a Keras model.
 
     Args:
         graph: Graph to apply SNC on.
         core_config: Quantization configuration.
-        fw_info: FrameworkInfo object with information about the specific framework's module.
 
     Returns:
         Graph after SNC.
@@ -244,7 +242,6 @@ def keras_apply_shift_negative_correction(graph: Graph,
 
     return apply_shift_negative_correction(graph,
                                            core_config,
-                                           fw_info,
                                            snc_node,
                                            linear_node,
                                            bypass_node,

model_compression_toolkit/core/keras/hessian/weights_hessian_scores_calculator_keras.py

@@ -22,7 +22,6 @@ from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS, MIN_HESS
 from model_compression_toolkit.core.common import Graph
 from model_compression_toolkit.core.common.hessian import HessianScoresRequest, HessianScoresGranularity
 from model_compression_toolkit.core.keras.back2framework.float_model_builder import FloatKerasModelBuilder
-from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.core.keras.hessian.hessian_scores_calculator_keras import HessianScoresCalculatorKeras
 from model_compression_toolkit.logger import Logger
 
@@ -95,20 +94,11 @@ class WeightsHessianScoresCalculatorKeras(HessianScoresCalculatorKeras):
                 for i, ipt_node in enumerate(self.hessian_request.target_nodes):  # Per Interest point weights tensor
 
                     # Check if the target node's layer type is supported.
-                    if not DEFAULT_KERAS_INFO.is_kernel_op(ipt_node.type):
+                    if not ipt_node.is_kernel_op:
                         Logger.critical(f"Hessian information with respect to weights is not supported for "
                                         f"{ipt_node.type} layers.")  # pragma: no cover
 
-                    # Get the weight attributes for the target node type
-                    weight_attributes = DEFAULT_KERAS_INFO.get_kernel_op_attributes(ipt_node.type)
-
-                    # Get the weight tensor for the target node
-                    if len(weight_attributes) != 1:  # pragma: no cover
-                        Logger.critical(
-                            f"Hessian-based scoring with respect to weights is currently supported only for nodes with "
-                            f"a single weight attribute. Found {len(weight_attributes)} attributes.")
-
-                    weight_tensor = getattr(model.get_layer(ipt_node.name), weight_attributes[0])
+                    weight_tensor = getattr(model.get_layer(ipt_node.name), ipt_node.kernel_attr)
 
                     if j == 0:
                         # On the first iteration we store the weight_tensor shape for later reshaping the results
@@ -116,7 +106,7 @@ class WeightsHessianScoresCalculatorKeras(HessianScoresCalculatorKeras):
                         tensors_original_shape.append(weight_tensor.shape)
 
                     # Get the output channel index (needed for HessianInfoGranularity.PER_OUTPUT_CHANNEL case)
-                    output_channel_axis, _ = DEFAULT_KERAS_INFO.kernel_channels_mapping.get(ipt_node.type)
+                    output_channel_axis = ipt_node.channel_axis.output
 
                     # Get number of scores that should be calculated by the granularity.
                     num_of_scores = self._get_num_scores_by_granularity(weight_tensor,

model_compression_toolkit/core/keras/keras_implementation.py

@@ -65,7 +65,6 @@ from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
 from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
 from model_compression_toolkit.core.common.node_prior_info import NodePriorInfo
-from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.activation_decomposition import \
     ActivationDecomposition
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.matmul_substitution import \
@@ -175,18 +174,16 @@ class KerasImplementation(FrameworkImplementation):
                       graph: Graph,
                       mode: ModelBuilderMode,
                       append2output: List[Any] = None,
-                      fw_info: FrameworkInfo = DEFAULT_KERAS_INFO,
                       return_float_outputs: bool = False) -> Tuple:
         """
         Build a Keras model from a graph.
-        The mode determines how the model should be build. append2output is a list of Nodes
+        The mode determines how the model should be built. append2output is a list of Nodes
         to set as the model outputs.
 
         Args:
             graph: Graph to build the model from it.
             mode: Mode for how to build the model.
             append2output: List of Nodes to set as the model's outputs.
-            fw_info: FrameworkInfo object with information about the specific framework's model
             return_float_outputs (bool): whether to return outputs before or after quantization nodes (default)
         Returns:
             A tuple with the model and additional relevant supporting objects.
@@ -195,7 +192,6 @@ class KerasImplementation(FrameworkImplementation):
         keras_model_builder = get_keras_model_builder(mode)
         return keras_model_builder(graph=graph,
                                    append2output=append2output,
-                                   fw_info=fw_info,
                                    return_float_outputs=return_float_outputs).build_model()
 
     def run_model_inference(self,
@@ -227,65 +223,57 @@ class KerasImplementation(FrameworkImplementation):
 
     def shift_negative_correction(self,
                                   graph: Graph,
-                                  core_config: CoreConfig,
-                                  fw_info: FrameworkInfo) -> Graph:
+                                  core_config: CoreConfig) -> Graph:
         """
         Apply shift negative correction (SNC) on a graph.
 
         Args:
             graph: Graph to apply SNC on.
             core_config: Quantization configuration.
-            fw_info: FrameworkInfo object with information about the specific framework's model.
 
         Returns:
             Graph after SNC.
         """
         return keras_apply_shift_negative_correction(graph,
-                                                     core_config,
-                                                     fw_info)
+                                                     core_config)
 
     def compute_activation_bias_correction(self,
                                            graph: Graph,
-                                           quant_config: QuantizationConfig,
-                                           fw_info: FrameworkInfo):
+                                           quant_config: QuantizationConfig):
         """
         Compute activation bias correction on a graph.
 
         Args:
             graph: Graph to apply activation bias correction on.
             quant_config: QuantizationConfig of how the model should be quantized.
-            fw_info: FrameworkInfo object with information about the specific framework's model.
 
         Returns:
             Graph after activation bias correction computing.
         """
         return keras_compute_activation_bias_correction_of_graph(graph=graph,
                                                                  quant_config=quant_config,
-                                                                 fw_info=fw_info,
                                                                  fw_impl=self)
 
     def get_substitutions_channel_equalization(self,
-                                               quant_config: QuantizationConfig,
-                                               fw_info: FrameworkInfo) -> List[common.BaseSubstitution]:
+                                               quant_config: QuantizationConfig) -> List[common.BaseSubstitution]:
         """
         Return a list of the framework substitutions used for channel equalization.
 
         Args:
             quant_config: QuantizationConfig to determine which substitutions to return.
-            fw_info: FrameworkInfo object with information about the specific framework's model.
 
         Returns:
             A list of the framework substitutions used after we collect statistics.
         """
         substitutions_list = []
         if quant_config.activation_channel_equalization:
-            substitutions_list.extend([ScaleEqualization(quant_config, fw_info),
-                                       ScaleEqualizationWithPad(quant_config, fw_info),
-                                       ScaleEqualizationMidActivation(quant_config, fw_info),
-                                       ScaleEqualizationMidActivationWithPad(quant_config, fw_info)])
+            substitutions_list.extend([ScaleEqualization(quant_config),
+                                       ScaleEqualizationWithPad(quant_config),
+                                       ScaleEqualizationMidActivation(quant_config),
+                                       ScaleEqualizationMidActivationWithPad(quant_config)])
         return substitutions_list
 
-    def get_substitutions_prepare_graph(self, fw_info: FrameworkInfo = None) -> List[common.BaseSubstitution]:
+    def get_substitutions_prepare_graph(self) -> List[common.BaseSubstitution]:
         """
 
         Returns: A list of the framework substitutions used to prepare the graph.
@@ -402,22 +390,19 @@ class KerasImplementation(FrameworkImplementation):
 
     def get_node_prior_info(self,
                             node: BaseNode,
-                            fw_info: FrameworkInfo,
                             graph: Graph) -> NodePriorInfo:
         """
         Get a NodePriorInfo object for a node that represents a Keras layer.
 
         Args:
             node: Node to get its prior info.
-            fw_info: Framework specific information needed to create the prior info of the node.
             graph: Graph to check the next node type.
 
         Returns:
             NodePriorInfo with information about the node.
         """
 
-        return create_node_prior_info(node=node,
-                                      fw_info=fw_info, graph=graph)
+        return create_node_prior_info(node=node, graph=graph)
 
     def count_node_for_mixed_precision_interest_points(self, node: BaseNode) -> bool:
         """
@@ -530,23 +515,19 @@ class KerasImplementation(FrameworkImplementation):
         return True
 
     def get_node_mac_operations(self,
-                                node: BaseNode,
-                                fw_info: FrameworkInfo) -> float:
+                                node: BaseNode) -> float:
         """
         Gets the MAC operation count for a given operation.
 
         Args:
             node: A graph node that wraps the operation for which the MAC count is computed.
-            fw_info: FrameworkInfo object with information about the Keras model.
 
         Returns: The MAC count og the operation
         """
-        kernels = fw_info.get_kernel_op_attributes(node.type)
-        if not kernels or kernels[0] is None:
+        if node.kernel_attr is None:
             return 0
 
-        assert len(kernels) == 1
-        kernel_shape = node.get_weights_by_keys(kernels[0]).shape
+        kernel_shape = node.get_weights_by_keys(node.kernel_attr).shape
 
         if node.is_match_type(Conv2D) or node.is_match_type(Conv2DTranspose) or node.is_match_type(DepthwiseConv2D):
             h, w = node.get_output_shapes_list()[0][-3:-1]
@@ -554,8 +535,7 @@ class KerasImplementation(FrameworkImplementation):
 
         if node.is_match_type(Dense):
             # IN * OUT * (all previous dims[:-1])
-            _, input_channel_axis = fw_info.kernel_channels_mapping.get(node.type)
-            return node.get_total_output_params() * kernel_shape[input_channel_axis]
+            return node.get_total_output_params() * kernel_shape[node.channel_axis.input]
 
         return 0
 

model_compression_toolkit/core/keras/keras_model_validation.py

@@ -1,6 +1,6 @@
 from tensorflow.keras.models import Model
 
-from model_compression_toolkit.core import FrameworkInfo
+from model_compression_toolkit.core.common.framework_info import get_fw_info
 from model_compression_toolkit.core.common.framework_info import ChannelAxis
 from model_compression_toolkit.core.common.model_validation import ModelValidation
 from model_compression_toolkit.core.keras.constants import CHANNELS_FORMAT, CHANNELS_FORMAT_LAST, CHANNELS_FORMAT_FIRST
@@ -11,17 +11,15 @@ class KerasModelValidation(ModelValidation):
     Class to define validation methods in order to validate the received Keras model to quantize.
     """
 
-    def __init__(self, model: Model, fw_info: FrameworkInfo):
+    def __init__(self, model: Model):
         """
         Initialize a KerasModelValidation object.
 
         Args:
             model: Keras model to check its validity.
-            fw_info: Information about the framework of the model (Keras).
         """
 
-        super(KerasModelValidation, self).__init__(model=model,
-                                                   fw_info=fw_info)
+        super(KerasModelValidation, self).__init__(model=model)
 
     def validate_output_channel_consistency(self):
         """
@@ -30,9 +28,10 @@ class KerasModelValidation(ModelValidation):
         If the model has layers with different output channels index, an exception is thrown.
 
         """
+        fw_info = get_fw_info()
         for layer in self.model.layers:
             data_format = layer.get_config().get(CHANNELS_FORMAT)
             if data_format is not None:
-                assert (data_format == CHANNELS_FORMAT_LAST and self.fw_info.out_channel_axis_mapping.get(layer) == ChannelAxis.NHWC.value
-                        or data_format == CHANNELS_FORMAT_FIRST and self.fw_info.out_channel_axis_mapping.get(layer) == ChannelAxis.NCHW.value), \
+                assert (data_format == CHANNELS_FORMAT_LAST and fw_info.get_out_channel_axis(layer) == ChannelAxis.NHWC.value
+                        or data_format == CHANNELS_FORMAT_FIRST and fw_info.get_out_channel_axis(layer) == ChannelAxis.NCHW.value), \
                     f'Model can not have layers with different data formats.'

model_compression_toolkit/core/keras/keras_node_prior_info.py

@@ -17,22 +17,19 @@ from model_compression_toolkit.core.common.graph.base_graph import Graph
 
 
 def create_node_prior_info(node: BaseNode,
-                           fw_info: FrameworkInfo,
                            graph: Graph):
     """
     Create a NodePriorInfo object for a given node.
 
     Args:
         node: Node to create its prior info.
-        fw_info: Information about a specific framework the node was generated from.
         graph: Graph to check the next node type.
 
     Returns:
         NodePriorInfo object with info about the node.
     """
 
-    min_output, max_output = _get_min_max_outputs(node=node,
-                                                  fw_info=fw_info)
+    min_output, max_output = _get_min_max_outputs(node=node)
 
     mean_output, std_output = _get_mean_std_outputs(node=node,
                                                     graph=graph)
@@ -42,14 +39,12 @@ def create_node_prior_info(node: BaseNode,
                          std_output=std_output)
 
 
-def _get_min_max_outputs(node: BaseNode,
-                         fw_info: FrameworkInfo) -> Tuple[Any, Any]:
+def _get_min_max_outputs(node: BaseNode) -> Tuple[Any, Any]:
     """
     Return the min/max output values of a node if known.
     If one of them (or both of them) is unknown - return None instead of a value.
     Args:
         node: Node to create its prior info.
-        fw_info: Information about a specific framework the node was generated from.
 
     Returns:
         Min/max output values if known.
@@ -58,12 +53,8 @@ def _get_min_max_outputs(node: BaseNode,
 
     if node.is_match_type(ReLU):
         min_output = node.framework_attr[THRESHOLD] if node.framework_attr[NEGATIVE_SLOPE] == 0 else None
-
-    elif fw_info.layers_has_min_max(node.type):
-        min_output, max_output = fw_info.layer_min_max_mapping[node.type]
-
-    elif node.is_match_type(Activation) and fw_info.activation_has_min_max(node.framework_attr[ACTIVATION]):
-        min_output, max_output = fw_info.activation_min_max_mapping[node.framework_attr[ACTIVATION]]
+    else:
+        min_output, max_output = node.minmax
 
     return min_output, max_output
 

model_compression_toolkit/core/keras/pruning/pruning_keras_implementation.py

@@ -19,7 +19,6 @@ from model_compression_toolkit.core.common.pruning.pruning_framework_implementat
     PruningFrameworkImplementation
 from model_compression_toolkit.core.common.pruning.pruning_section import PruningSection
 from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.keras.constants import BIAS, GROUPS, FILTERS, UNITS, USE_BIAS
 import keras
@@ -38,27 +37,23 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
 
     def prune_entry_node(self,
                          node: BaseNode,
-                         output_mask: np.ndarray,
-                         fw_info: FrameworkInfo):
+                         output_mask: np.ndarray):
         """
         Prunes the entry node of a model in Keras.
 
         Args:
             node (BaseNode): The entry node to be pruned.
             output_mask (np.ndarray): A numpy array representing the mask to be applied to the output channels.
-            fw_info (FrameworkInfo): Framework-specific information object.
 
         """
         return _prune_keras_edge_node(node=node,
                                       mask=output_mask,
-                                      fw_info=fw_info,
                                       is_exit_node=False)
 
     def prune_intermediate_node(self,
                                 node: BaseNode,
                                 input_mask: np.ndarray,
-                                output_mask: np.ndarray,
-                                fw_info: FrameworkInfo):
+                                output_mask: np.ndarray):
         """
         Prunes an intermediate node in a Keras model.
 
@@ -66,7 +61,6 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
             node (BaseNode): The intermediate node to be pruned.
             input_mask (np.ndarray): A numpy array representing the mask to be applied to the input channels.
             output_mask (np.ndarray): A numpy array representing the mask to be applied to the output channels.
-            fw_info (FrameworkInfo): Framework-specific information object.
 
         """
         _edit_node_input_shape(input_mask, node)
@@ -79,20 +73,17 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
 
     def prune_exit_node(self,
                         node: BaseNode,
-                        input_mask: np.ndarray,
-                        fw_info: FrameworkInfo):
+                        input_mask: np.ndarray):
         """
         Prunes the exit node of a model in Keras.
 
         Args:
             node (BaseNode): The exit node to be pruned.
             input_mask (np.ndarray): A numpy array representing the mask to be applied to the input channels.
-            fw_info (FrameworkInfo): Framework-specific information object.
 
         """
         return _prune_keras_edge_node(node=node,
                                       mask=input_mask,
-                                      fw_info=fw_info,
                                       is_exit_node=True)
 
     def is_node_entry_node(self, node: BaseNode) -> bool:
@@ -109,22 +100,19 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
 
     def is_node_exit_node(self,
                           node: BaseNode,
-                          corresponding_entry_node: BaseNode,
-                          fw_info: FrameworkInfo) -> bool:
+                          corresponding_entry_node: BaseNode) -> bool:
         """
         Determines whether a node is an exit node in a Keras model.
 
         Args:
             node (BaseNode): The node to be checked.
             corresponding_entry_node (BaseNode): The entry node of the pruning section that is checked.
-            fw_info (FrameworkInfo): Framework-specific information object.
 
         Returns:
             bool: Boolean indicating if the node is an exit node.
         """
         return _is_keras_node_pruning_section_edge(node) and PruningSection.has_matching_channel_count(node,
-                                                                                                       corresponding_entry_node,
-                                                                                                       fw_info)
+                                                                                                       corresponding_entry_node)
 
     def is_node_intermediate_pruning_section(self, node: BaseNode) -> bool:
         """
@@ -143,8 +131,7 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
                                  keras.layers.Dense]
 
     def attrs_oi_channels_info_for_pruning(self,
-                                           node: BaseNode,
-                                           fw_info: FrameworkInfo) -> Dict[str, Tuple[int, int]]:
+                                           node: BaseNode) -> Dict[str, Tuple[int, int]]:
         """
         Retrieves the attributes of a given node along with the output/input (OI) channel axis
         for each attribute used to prune these attributes.
@@ -161,7 +148,6 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
 
         Args:
             node (BaseNode): The node from the computational graph.
-            fw_info (FrameworkInfo): Contains framework-specific information and utilities.
 
         Returns:
             Dict[str, Tuple[int, int]]: A dictionary where each key is an attribute name (like 'kernel' or 'bias')
@@ -169,13 +155,8 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
         """
 
         attributes_with_axis = {}
-        if fw_info.is_kernel_op(node.type):
-            kernel_attributes = fw_info.get_kernel_op_attributes(node.type)
-            if kernel_attributes is None or len(kernel_attributes)==0:
-                Logger.critical(f"Expected kernel attributes for operation for node type {node.type}, found None or empty.")
-
-            for attr in kernel_attributes:
-                attributes_with_axis[attr] = fw_info.kernel_channels_mapping.get(node.type)
+        if node.is_kernel_op:
+            attributes_with_axis[node.kernel_attr] = (node.channel_axis.output, node.channel_axis.input)
 
             # Bias is a vector at the length of the number of output channels.
             # For this reason, input channel axis is irrelevant to the bias attribute.
@@ -216,7 +197,6 @@ def _is_keras_node_pruning_section_edge(node: BaseNode) -> bool:
 
 def _prune_keras_edge_node(node: BaseNode,
                            mask: np.ndarray,
-                           fw_info: FrameworkInfo,
                            is_exit_node: bool):
     """
     Prunes the given Keras node by applying the mask to the node's weights (kernels and biases).
@@ -225,21 +205,18 @@ def _prune_keras_edge_node(node: BaseNode,
     Args:
         node: The node to be pruned.
         mask: The pruning mask to be applied.
-        fw_info: Framework-specific information object.
         is_exit_node: A boolean indicating whether the node is an exit node.
 
     """
 
     # Retrieve the kernel attribute and the axes to prune.
-    kernel_attr = fw_info.get_kernel_op_attributes(node.type)[0]
-    io_axis = fw_info.kernel_channels_mapping.get(node.type)
-    axis_to_prune = io_axis[int(is_exit_node)]
-    kernel = node.get_weights_by_keys(kernel_attr)
+    axis_to_prune = node.channel_axis.input if is_exit_node else node.channel_axis.output
+    kernel = node.get_weights_by_keys(node.kernel_attr)
     # Convert mask to boolean.
     mask_bool = mask.astype(bool)
 
     pruned_kernel = kernel.compress(mask_bool, axis=axis_to_prune)
-    node.set_weights_by_keys(name=kernel_attr, tensor=pruned_kernel)
+    node.set_weights_by_keys(name=node.kernel_attr, tensor=pruned_kernel)
 
     if not is_exit_node and node.framework_attr[USE_BIAS]:
         # Prune the bias if applicable and it's an entry node.

model_compression_toolkit/core/keras/resource_utilization_data_facade.py

@@ -27,7 +27,6 @@ if FOUND_TF:
     from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2keras import \
         AttachTpcToKeras
     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.core.keras.keras_implementation import KerasImplementation
     from tensorflow.keras.models import Model
 
@@ -93,7 +92,6 @@ if FOUND_TF:
                                                  representative_data_gen,
                                                  core_config,
                                                  target_platform_capabilities,
-                                                 DEFAULT_KERAS_INFO,
                                                  fw_impl)
 
 else:

model_compression_toolkit/core/keras/statistics_correction/keras_compute_activation_bias_correction_of_graph.py

@@ -43,7 +43,6 @@ def activation_bias_correction_node_matchers():
 
 def keras_compute_activation_bias_correction_of_graph(graph: Graph,
                                                       quant_config: QuantizationConfig,
-                                                      fw_info: FrameworkInfo,
                                                       fw_impl: FrameworkImplementation) -> Graph:
     """
     Compute the activation bias correction term for graph based on a Keras model.
@@ -51,7 +50,6 @@ def keras_compute_activation_bias_correction_of_graph(graph: Graph,
     Args:
         graph: Graph with nodes to compute the activation bias correction.
         quant_config: QuantizationConfig of how the model should be quantized.
-        fw_info: Framework info like lists of nodes their kernel should quantized.
         fw_impl: FrameworkImplementation object with a specific framework methods implementation.
 
     Returns:
@@ -59,7 +57,6 @@ def keras_compute_activation_bias_correction_of_graph(graph: Graph,
     """
     graph = compute_activation_bias_correction_of_graph(graph=graph,
                                                         quant_config=quant_config,
-                                                        fw_info=fw_info,
                                                         fw_impl=fw_impl,
                                                         activation_bias_correction_node_matchers=
                                                         activation_bias_correction_node_matchers,

model_compression_toolkit/core/pytorch/back2framework/float_model_builder.py

@@ -24,8 +24,6 @@ from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.core.pytorch.back2framework.pytorch_model_builder import PyTorchModelBuilder, \
     PytorchModel
 
-from model_compression_toolkit.core.pytorch.default_framework_info import DEFAULT_PYTORCH_INFO
-
 
 class FloatPyTorchModel(PytorchModel):
     """
@@ -34,19 +32,16 @@ class FloatPyTorchModel(PytorchModel):
 
     def __init__(self,
                  graph: common.Graph,
-                 append2output=None,
-                 fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO):
+                 append2output=None):
         """
 
         Args:
             graph: Graph to build its corresponding Pytorch model.
             append2output: List of nodes or OutTensor objects.
-            fw_info: Framework information (e.g., mapping from layers to their attributes to quantize).
         """
 
         super().__init__(graph,
-                         append2output,
-                         fw_info)
+                         append2output)
 
     def _quantize_node_activations(self,
                                    node: BaseNode,
@@ -71,20 +66,17 @@ class FloatPyTorchModelBuilder(PyTorchModelBuilder):
     def __init__(self,
                  graph: common.Graph,
                  append2output=None,
-                 fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO,
                  return_float_outputs: bool = False):
         """
 
         Args:
             graph: Graph to build the model from.
             append2output: Nodes to append to model's output.
-            fw_info: Information about the specific framework of the model that is built.
             return_float_outputs: Whether the model returns float tensors or not.
         """
 
         super().__init__(graph,
                          append2output,
-                         fw_info,
                          return_float_outputs)
 
     def build_model(self) -> Tuple[PytorchModel, UserInformation]:
@@ -94,5 +86,4 @@ class FloatPyTorchModelBuilder(PyTorchModelBuilder):
 
         """
         return FloatPyTorchModel(self.graph,
-                                 self.append2output,
-                                 self.fw_info), self.graph.user_info
+                                 self.append2output), self.graph.user_info