mct-nightly 2.2.0.20250116.522__tar.gz → 2.2.0.20250117.527__tar.gz

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: mct-nightly
-Version: 2.2.0.20250116.522
+Version: 2.2.0.20250117.527
 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.20250116.522 → mct_nightly-2.2.0.20250117.527}/mct_nightly.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.2
 Name: mct-nightly
-Version: 2.2.0.20250116.522
+Version: 2.2.0.20250117.527
 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.20250116.522 → mct_nightly-2.2.0.20250117.527}/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.20250116.000522"
+__version__ = "2.2.0.20250117.000527"

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/core/common/graph/memory_graph/compute_graph_max_cut.py

@@ -13,9 +13,9 @@
 # limitations under the License.
 # ==============================================================================
 from collections import namedtuple
-
 from typing import Tuple, List
 
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import OPERATORS_SCHEDULING, MAX_CUT, CUTS, FUSED_NODES_MAPPING
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.graph.memory_graph.cut import Cut
@@ -49,7 +49,17 @@ def compute_graph_max_cut(memory_graph: MemoryGraph,
     it = 0
     while it < n_iter:
         estimate = (u_bound + l_bound) / 2
-        schedule, max_cut_size, cuts = max_cut_astar.solve(estimate=estimate, iter_limit=astar_n_iter)
+        # Add a timeout of 5 minutes to the solver from the 2nd iteration.
+        try:
+            schedule, max_cut_size, cuts = max_cut_astar.solve(estimate=estimate, iter_limit=astar_n_iter,
+                                                               time_limit=None if it == 0 else 300)
+        except TimeoutError:
+            if last_result[0] is None:
+                Logger.critical(f"Max-cut solver stopped on timeout in iteration {it} before finding a solution.")  # pragma: no cover
+            else:
+                Logger.warning(f"Max-cut solver stopped on timeout in iteration {it}.")
+                return last_result
+
         if schedule is None:
             l_bound = estimate
         else:

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/core/common/graph/memory_graph/cut.py

@@ -67,7 +67,7 @@ class Cut:
         return False  # pragma: no cover
 
     def __hash__(self):
-        return hash((frozenset(self.op_order), frozenset(self.op_record), self.mem_elements))
+        return id(self)
 
     def __repr__(self):
-        return f"<Cut: Nodes={[e.node_name for e in self.mem_elements.elements]}, size={self.memory_size()}>"  # pragma: no cover
+        return f"<Cut: Nodes={[e.node_name for e in self.mem_elements.elements]}, size={self.memory_size()}>"  # pragma: no cover

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/core/common/graph/memory_graph/max_cut_astar.py

@@ -13,7 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 import copy
-from typing import List, Tuple, Dict
+from typing import List, Tuple, Dict, Set
+from time import time
 
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.constants import DUMMY_TENSOR, DUMMY_NODE
@@ -122,7 +123,7 @@ class MaxCutAstar:
         self.target_cut = Cut([], set(), MemoryElements(elements={target_dummy_b, target_dummy_b2},
                                                         total_size=0))
 
-    def solve(self, estimate: float, iter_limit: int = 500) -> Tuple[List[BaseNode], float, List[Cut]]:
+    def solve(self, estimate: float, iter_limit: int = 500, time_limit: int = None) -> Tuple[List[BaseNode], float, List[Cut]]:
         """
         The AStar solver function. This method runs an AStar-like search on the memory graph,
         using the given estimate as a heuristic gap for solutions to consider.
@@ -131,6 +132,7 @@ class MaxCutAstar:
             estimate: Cut size estimation to consider larger size of nodes in each
                 expansion step, in order to fasten the algorithm divergence towards a solution.
             iter_limit: An upper limit for the number of expansion steps that the algorithm preforms.
+            time_limit: Optional time limit to the solver. Defaults to None which means no limit.
 
         Returns: A solution (if found within the steps limit) which contains:
         - A schedule for computation of the model (List of nodes).
@@ -139,14 +141,17 @@ class MaxCutAstar:
 
         """
 
-        open_list = [self.src_cut]
-        closed_list = []
+        open_list = {self.src_cut}
+        closed_list = set()
         costs = {self.src_cut: self.src_cut.memory_size()}
         routes = {self.src_cut: [self.src_cut]}
 
         expansion_count = 0
 
+        t1 = time()
         while expansion_count < iter_limit and len(open_list) > 0:
+            if time_limit is not None and time() - t1 > time_limit:
+                raise TimeoutError
             # Choose next node to expand
             next_cut = self._get_cut_to_expand(open_list, costs, routes, estimate)
 
@@ -159,22 +164,21 @@ class MaxCutAstar:
 
             if self.is_pivot(next_cut):
                 # Can clear all search history
-                open_list = []
-                closed_list = []
+                open_list.clear()
+                closed_list.clear()
                 routes = {}
             else:
                 # Can remove only next_cut and put it in closed_list
                 open_list.remove(next_cut)
                 del routes[next_cut]
-                closed_list.append(next_cut)
+                closed_list.add(next_cut)
 
             # Expand the chosen cut
             expanded_cuts = self.expand(next_cut)
             expansion_count += 1
 
             # Only consider nodes that where not already visited
-            expanded_cuts = [_c for _c in expanded_cuts if _c not in closed_list]
-            for c in expanded_cuts:
+            for c in filter(lambda _c: _c not in closed_list, expanded_cuts):
                 cost = self.accumulate(cut_cost, c.memory_size())
                 if c not in open_list:
                     self._update_expanded_node(c, cost, cut_route, open_list, costs, routes)
@@ -192,7 +196,7 @@ class MaxCutAstar:
         return None, 0, None  # pragma: no cover
 
     @staticmethod
-    def _update_expanded_node(cut: Cut, cost: float, route: List[Cut], open_list: List[Cut],
+    def _update_expanded_node(cut: Cut, cost: float, route: List[Cut], open_list: Set[Cut],
                               costs: Dict[Cut, float], routes: Dict[Cut, List[Cut]]):
         """
         An auxiliary method for updating search data structures according to an expanded node.
@@ -201,16 +205,16 @@ class MaxCutAstar:
             cut: A cut to expand the search to.
             cost: The cost of the cut.
             route: The rout to the cut.
-            open_list: The search open list.
+            open_list: The search open set.
             costs: The search utility mapping between cuts and their cost.
             routes: The search utility mapping between cuts and their routes.
 
         """
-        open_list.append(cut)
+        open_list.add(cut)
         costs.update({cut: cost})
         routes.update({cut: [cut] + route})
 
-    def _get_cut_to_expand(self, open_list: List[Cut], costs: Dict[Cut, float], routes: Dict[Cut, List[Cut]],
+    def _get_cut_to_expand(self, open_list: Set[Cut], costs: Dict[Cut, float], routes: Dict[Cut, List[Cut]],
                            estimate: float) -> Cut:
         """
         An auxiliary method for finding a cut for expanding the search out of a set of potential cuts for expansion.

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/core/common/graph/memory_graph/memory_graph.py

@@ -14,6 +14,7 @@
 # ==============================================================================
 from typing import List
 from operator import getitem
+from functools import cache
 
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.graph.edge import EDGE_SOURCE_INDEX
@@ -82,7 +83,6 @@ class MemoryGraph(DirectedBipartiteGraph):
         inputs_tensors_memory = [sum([t.total_size for t in self.operation_node_children(n)])
                                  for n in nodes if n in model_graph.get_inputs()]
 
-        # TODO maxcut: why both inputs and outputs of each nodes, while the A* solves for node outputs only???
         nodes_total_memory = [sum([t.total_size for t in self.operation_node_children(n)] +
                                   [t.total_size for t in self.operation_node_parents(n)])
                               for n in nodes if n not in model_graph.get_inputs()]
@@ -117,6 +117,7 @@ class MemoryGraph(DirectedBipartiteGraph):
         """
         self.sinks_b = [n for n in self.b_nodes if len(list(self.successors(n))) == 0]
 
+    @cache
     def activation_tensor_children(self, activation_tensor: ActivationMemoryTensor) -> List[BaseNode]:
         """
         Returns the children nodes of a side B node (activation tensor) in the bipartite graph.
@@ -129,6 +130,7 @@ class MemoryGraph(DirectedBipartiteGraph):
         """
         return [oe[1] for oe in self.out_edges(activation_tensor)]
 
+    @cache
     def activation_tensor_parents(self, activation_tensor: ActivationMemoryTensor) -> List[BaseNode]:
         """
         Returns the parents nodes of a side B node (activation tensor) in the bipartite graph.
@@ -141,6 +143,7 @@ class MemoryGraph(DirectedBipartiteGraph):
         """
         return [ie[0] for ie in self.in_edges(activation_tensor)]
 
+    @cache
     def operation_node_children(self, op_node: BaseNode) -> List[ActivationMemoryTensor]:
         """
         Returns the children nodes of a side A node (operation) in the bipartite graph.
@@ -153,6 +156,7 @@ class MemoryGraph(DirectedBipartiteGraph):
         """
         return [oe[1] for oe in self.out_edges(op_node)]
 
+    @cache
     def operation_node_parents(self, op_node: BaseNode) -> List[ActivationMemoryTensor]:
         """
         Returns the parents nodes of a side A node (operation) in the bipartite graph.

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/core/common/mixed_precision/resource_utilization_tools/resource_utilization_calculator.py

@@ -17,6 +17,7 @@ from copy import deepcopy
 from enum import Enum, auto
 from typing import Dict, NamedTuple, Optional, Tuple, List, Iterable, Union, Literal, Sequence, Set
 
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import FLOAT_BITWIDTH
 from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core.common import Graph, BaseNode
@@ -169,6 +170,7 @@ class ResourceUtilizationCalculator:
             w_total, *_ = self.compute_weights_utilization(target_criterion, bitwidth_mode, w_qcs)
 
         if {RUTarget.ACTIVATION, RUTarget.TOTAL}.intersection(ru_targets):
+            Logger.warning("Using an experimental feature max-cut for activation memory utilization estimation.")
             a_total = self.compute_activations_utilization(target_criterion, bitwidth_mode, act_qcs)
 
         ru = ResourceUtilization()

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/core/common/mixed_precision/search_methods/linear_programming.py

@@ -16,7 +16,7 @@
 import numpy as np
 from pulp import *
 from tqdm import tqdm
-from typing import Dict, Tuple, Set, Any
+from typing import Dict, Tuple, Any, Optional
 
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization, RUTarget
@@ -182,7 +182,7 @@ def _add_ru_constraints(search_manager: MixedPrecisionSearchManager,
                         target_resource_utilization: ResourceUtilization,
                         indicators_matrix: np.ndarray,
                         lp_problem: LpProblem,
-                        non_conf_ru_dict: Optional[Dict[RUTarget, np.ndarray]]):
+                        non_conf_ru_dict: Dict[RUTarget, np.ndarray]):
     """
     Adding targets constraints for the Lp problem for the given target resource utilization.
     The update to the Lp problem object is done inplace.

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/core/keras/resource_utilization_data_facade.py

@@ -13,13 +13,14 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Callable
+from typing import Callable, Union
 from model_compression_toolkit.core import MixedPrecisionQuantizationConfig, CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization_data import compute_resource_utilization_data
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TF
 
 if FOUND_TF:
@@ -38,7 +39,7 @@ if FOUND_TF:
                                         representative_data_gen: Callable,
                                         core_config: CoreConfig = CoreConfig(
                                             mixed_precision_config=MixedPrecisionQuantizationConfig()),
-                                        target_platform_capabilities: TargetPlatformCapabilities = KERAS_DEFAULT_TPC
+                                        target_platform_capabilities: Union[TargetPlatformCapabilities, str] = KERAS_DEFAULT_TPC
                                         ) -> ResourceUtilization:
         """
         Computes resource utilization data that can be used to calculate the desired target resource utilization
@@ -50,7 +51,7 @@ if FOUND_TF:
             in_model (Model): Keras model to quantize.
             representative_data_gen (Callable): Dataset used for calibration.
             core_config (CoreConfig): CoreConfig containing parameters for quantization and mixed precision of how the model should be quantized.
-            target_platform_capabilities (FrameworkQuantizationCapabilities): FrameworkQuantizationCapabilities to optimize the Keras model according to.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): FrameworkQuantizationCapabilities to optimize the Keras model according to.
 
         Returns:
 
@@ -81,6 +82,7 @@ if FOUND_TF:
 
         fw_impl = KerasImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
         attach2keras = AttachTpcToKeras()
         target_platform_capabilities = attach2keras.attach(

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/core/pytorch/resource_utilization_data_facade.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Callable
+from typing import Callable, Union
 
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import PYTORCH
@@ -23,6 +23,7 @@ from model_compression_toolkit.core.common.mixed_precision.resource_utilization_
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig
 from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 
 if FOUND_TORCH:
@@ -40,7 +41,7 @@ if FOUND_TORCH:
     def pytorch_resource_utilization_data(in_model: Module,
                                           representative_data_gen: Callable,
                                           core_config: CoreConfig = CoreConfig(),
-                                          target_platform_capabilities: TargetPlatformCapabilities= PYTORCH_DEFAULT_TPC
+                                          target_platform_capabilities: Union[TargetPlatformCapabilities, str] = PYTORCH_DEFAULT_TPC
                                           ) -> ResourceUtilization:
         """
         Computes resource utilization data that can be used to calculate the desired target resource utilization for mixed-precision quantization.
@@ -50,7 +51,7 @@ if FOUND_TORCH:
             in_model (Model): PyTorch model to quantize.
             representative_data_gen (Callable): Dataset used for calibration.
             core_config (CoreConfig): CoreConfig containing parameters for quantization and mixed precision
-            target_platform_capabilities (FrameworkQuantizationCapabilities): FrameworkQuantizationCapabilities to optimize the PyTorch model according to.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): FrameworkQuantizationCapabilities to optimize the PyTorch model according to.
 
         Returns:
 
@@ -81,6 +82,7 @@ if FOUND_TORCH:
 
         fw_impl = PytorchImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
         attach2pytorch = AttachTpcToPytorch()
         target_platform_capabilities = (

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -25,6 +25,7 @@ from model_compression_toolkit.constants import TENSORFLOW, ACT_HESSIAN_DEFAULT_
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2keras import \
     AttachTpcToKeras
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfig, GPTQHessianScoresConfig, \
@@ -156,7 +157,8 @@ if FOUND_TF:
                                                   gptq_representative_data_gen: Callable = None,
                                                   target_resource_utilization: ResourceUtilization = None,
                                                   core_config: CoreConfig = CoreConfig(),
-                                                  target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC) -> Tuple[Model, UserInformation]:
+                                                  target_platform_capabilities: Union[TargetPlatformCapabilities, str]
+                                                  = DEFAULT_KERAS_TPC) -> Tuple[Model, UserInformation]:
         """
         Quantize a trained Keras model using post-training quantization. The model is quantized using a
         symmetric constraint quantization thresholds (power of two).
@@ -180,7 +182,7 @@ if FOUND_TF:
             gptq_representative_data_gen (Callable): Dataset used for GPTQ training. If None defaults to representative_data_gen
             target_resource_utilization (ResourceUtilization): ResourceUtilization object to limit the search of the mixed-precision configuration as desired.
             core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
-            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): TargetPlatformCapabilities to optimize the Keras model according to.
 
         Returns:
 
@@ -241,6 +243,7 @@ if FOUND_TF:
 
         fw_impl = GPTQKerasImplemantation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
         attach2keras = AttachTpcToKeras()
         framework_platform_capabilities = attach2keras.attach(

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -32,6 +32,7 @@ from model_compression_toolkit.gptq.runner import gptq_runner
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.metadata import create_model_metadata
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 
 
@@ -145,7 +146,7 @@ if FOUND_TORCH:
                                                     core_config: CoreConfig = CoreConfig(),
                                                     gptq_config: GradientPTQConfig = None,
                                                     gptq_representative_data_gen: Callable = None,
-                                                    target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYTORCH_TPC):
+                                                    target_platform_capabilities: Union[TargetPlatformCapabilities, str] = DEFAULT_PYTORCH_TPC):
         """
         Quantize a trained Pytorch module using post-training quantization.
         By default, the module is quantized using a symmetric constraint quantization thresholds
@@ -169,7 +170,7 @@ if FOUND_TORCH:
             core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
             gptq_config (GradientPTQConfig): Configuration for using gptq (e.g. optimizer).
             gptq_representative_data_gen (Callable): Dataset used for GPTQ training. If None defaults to representative_data_gen
-            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the PyTorch model according to.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): TargetPlatformCapabilities to optimize the PyTorch model according to.
 
         Returns:
             A quantized module and information the user may need to handle the quantized module.
@@ -214,6 +215,7 @@ if FOUND_TORCH:
 
         fw_impl = GPTQPytorchImplemantation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
         attach2pytorch = AttachTpcToPytorch()
         framework_quantization_capabilities = attach2pytorch.attach(target_platform_capabilities,

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/pruning/keras/pruning_facade.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Callable, Tuple
+from typing import Callable, Tuple, Union
 
 from model_compression_toolkit import get_target_platform_capabilities
 from model_compression_toolkit.constants import TENSORFLOW
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.pruning.pruner import Pruner
@@ -43,7 +44,8 @@ if FOUND_TF:
                                    target_resource_utilization: ResourceUtilization,
                                    representative_data_gen: Callable,
                                    pruning_config: PruningConfig = PruningConfig(),
-                                   target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC) -> Tuple[Model, PruningInfo]:
+                                   target_platform_capabilities: Union[TargetPlatformCapabilities, str]
+                                   = DEFAULT_KERAS_TPC) -> Tuple[Model, PruningInfo]:
         """
         Perform structured pruning on a Keras model to meet a specified target resource utilization.
         This function prunes the provided model according to the target resource utilization by grouping and pruning
@@ -61,7 +63,7 @@ if FOUND_TF:
             target_resource_utilization (ResourceUtilization): The target Key Performance Indicators to be achieved through pruning.
             representative_data_gen (Callable): A function to generate representative data for pruning analysis.
             pruning_config (PruningConfig): Configuration settings for the pruning process. Defaults to standard config.
-            target_platform_capabilities (FrameworkQuantizationCapabilities): Platform-specific constraints and capabilities. Defaults to DEFAULT_KERAS_TPC.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): Platform-specific constraints and capabilities. Defaults to DEFAULT_KERAS_TPC.
 
         Returns:
             Tuple[Model, PruningInfo]: A tuple containing the pruned Keras model and associated pruning information.
@@ -112,6 +114,7 @@ if FOUND_TF:
         # Instantiate the Keras framework implementation.
         fw_impl = PruningKerasImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
         attach2keras = AttachTpcToKeras()
         target_platform_capabilities = attach2keras.attach(target_platform_capabilities)

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/pruning/pytorch/pruning_facade.py

@@ -13,10 +13,11 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Callable, Tuple
+from typing import Callable, Tuple, Union
 from model_compression_toolkit import get_target_platform_capabilities
 from model_compression_toolkit.constants import PYTORCH
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.pruning.pruner import Pruner
@@ -47,7 +48,8 @@ if FOUND_TORCH:
                                      target_resource_utilization: ResourceUtilization,
                                      representative_data_gen: Callable,
                                      pruning_config: PruningConfig = PruningConfig(),
-                                     target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYOTRCH_TPC) -> \
+                                     target_platform_capabilities: Union[TargetPlatformCapabilities, str]
+                                     = DEFAULT_PYOTRCH_TPC) -> \
             Tuple[Module, PruningInfo]:
         """
         Perform structured pruning on a Pytorch model to meet a specified target resource utilization.
@@ -66,7 +68,7 @@ if FOUND_TORCH:
             target_resource_utilization (ResourceUtilization): Key Performance Indicators specifying the pruning targets.
             representative_data_gen (Callable): A function to generate representative data for pruning analysis.
             pruning_config (PruningConfig): Configuration settings for the pruning process. Defaults to standard config.
-            target_platform_capabilities (TargetPlatformCapabilities): Platform-specific constraints and capabilities.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): Platform-specific constraints and capabilities.
                 Defaults to DEFAULT_PYTORCH_TPC.
 
         Returns:
@@ -118,6 +120,7 @@ if FOUND_TORCH:
         # Instantiate the Pytorch framework implementation.
         fw_impl = PruningPytorchImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach TPC to framework
         attach2pytorch = AttachTpcToPytorch()
         framework_platform_capabilities = attach2pytorch.attach(target_platform_capabilities)

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/ptq/keras/quantization_facade.py

@@ -23,6 +23,7 @@ from model_compression_toolkit.core.common.visualization.tensorboard_writer impo
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import TENSORFLOW
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
@@ -70,7 +71,7 @@ if FOUND_TF:
              representative_data_gen (Callable): Dataset used for calibration.
              target_resource_utilization (ResourceUtilization): ResourceUtilization object to limit the search of the mixed-precision configuration as desired.
              core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
-             target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
+             target_platform_capabilities (Union[TargetPlatformCapabilities, str]): TargetPlatformCapabilities to optimize the Keras model according to.
 
          Returns:
 
@@ -137,6 +138,7 @@ if FOUND_TF:
 
         fw_impl = KerasImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         attach2keras = AttachTpcToKeras()
         framework_platform_capabilities = attach2keras.attach(
             target_platform_capabilities,

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/ptq/pytorch/quantization_facade.py

@@ -14,12 +14,13 @@
 # ==============================================================================
 import copy
 
-from typing import Callable
+from typing import Callable, Union
 
 from model_compression_toolkit.core.common.visualization.tensorboard_writer import init_tensorboard_writer
 from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.constants import PYTORCH
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TORCH
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core import CoreConfig
@@ -48,7 +49,7 @@ if FOUND_TORCH:
                                            representative_data_gen: Callable,
                                            target_resource_utilization: ResourceUtilization = None,
                                            core_config: CoreConfig = CoreConfig(),
-                                           target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_PYTORCH_TPC):
+                                           target_platform_capabilities: Union[TargetPlatformCapabilities, str] = DEFAULT_PYTORCH_TPC):
         """
         Quantize a trained Pytorch module using post-training quantization.
         By default, the module is quantized using a symmetric constraint quantization thresholds
@@ -67,7 +68,7 @@ if FOUND_TORCH:
             representative_data_gen (Callable): Dataset used for calibration.
             target_resource_utilization (ResourceUtilization): ResourceUtilization object to limit the search of the mixed-precision configuration as desired.
             core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
-            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the PyTorch model according to.
+            target_platform_capabilities (Union[TargetPlatformCapabilities, str]): TargetPlatformCapabilities to optimize the PyTorch model according to.
 
         Returns:
             A quantized module and information the user may need to handle the quantized module.
@@ -109,6 +110,7 @@ if FOUND_TORCH:
 
         fw_impl = PytorchImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         # Attach tpc model to framework
         attach2pytorch = AttachTpcToPytorch()
         framework_platform_capabilities = attach2pytorch.attach(target_platform_capabilities,

{mct_nightly-2.2.0.20250116.522 → mct_nightly-2.2.0.20250117.527}/model_compression_toolkit/qat/keras/quantization_facade.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from typing import Callable
+from typing import Callable, Union
 from functools import partial
 
 from model_compression_toolkit.core import CoreConfig
@@ -22,6 +22,7 @@ from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.target_platform_capabilities.schema.mct_current_schema import TargetPlatformCapabilities
 from model_compression_toolkit.target_platform_capabilities.targetplatform2framework.attach2keras import \
     AttachTpcToKeras
+from model_compression_toolkit.target_platform_capabilities.tpc_io_handler import load_target_platform_capabilities
 from model_compression_toolkit.verify_packages import FOUND_TF
 from model_compression_toolkit.core.common.mixed_precision.resource_utilization_tools.resource_utilization import ResourceUtilization
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
@@ -49,7 +50,6 @@ if FOUND_TF:
     from model_compression_toolkit.core import common
     from model_compression_toolkit.core.common import BaseNode
     from model_compression_toolkit.constants import TENSORFLOW
-    from model_compression_toolkit.core.common.framework_info import FrameworkInfo
     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
@@ -92,7 +92,7 @@ if FOUND_TF:
                                                             target_resource_utilization: ResourceUtilization = None,
                                                             core_config: CoreConfig = CoreConfig(),
                                                             qat_config: QATConfig = QATConfig(),
-                                                            target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC):
+                                                            target_platform_capabilities: Union[TargetPlatformCapabilities, str] = DEFAULT_KERAS_TPC):
         """
          Prepare a trained Keras model for quantization aware training. First the model quantization is optimized
          with post-training quantization, then the model layers are wrapped with QuantizeWrappers. The model is
@@ -114,7 +114,7 @@ if FOUND_TF:
              target_resource_utilization (ResourceUtilization): ResourceUtilization object to limit the search of the mixed-precision configuration as desired.
              core_config (CoreConfig): Configuration object containing parameters of how the model should be quantized, including mixed precision parameters.
              qat_config (QATConfig): QAT configuration
-             target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
+             target_platform_capabilities (Union[TargetPlatformCapabilities, str]): TargetPlatformCapabilities to optimize the Keras model according to.
 
          Returns:
 
@@ -188,6 +188,7 @@ if FOUND_TF:
 
         fw_impl = KerasImplementation()
 
+        target_platform_capabilities = load_target_platform_capabilities(target_platform_capabilities)
         attach2keras = AttachTpcToKeras()
         target_platform_capabilities = attach2keras.attach(
             target_platform_capabilities,