mct-nightly 1.11.0.20240304.post404__py3-none-any.whl → 1.11.0.20240305.post352__py3-none-any.whl

model_compression_toolkit/core/keras/pruning/pruning_keras_implementation.py

@@ -44,9 +44,9 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
         Prunes the entry node of a model in Keras.
 
         Args:
-            node: The entry node to be pruned.
-            output_mask: A numpy array representing the mask to be applied to the output channels.
-            fw_info: Framework-specific information object.
+            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,
@@ -63,10 +63,10 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
         Prunes an intermediate node in a Keras model.
 
         Args:
-            node: The intermediate node to be pruned.
-            input_mask: A numpy array representing the mask to be applied to the input channels.
-            output_mask: A numpy array representing the mask to be applied to the output channels.
-            fw_info: Framework-specific information object.
+            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)
@@ -85,9 +85,9 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
         Prunes the exit node of a model in Keras.
 
         Args:
-            node: The exit node to be pruned.
-            input_mask: A numpy array representing the mask to be applied to the input channels.
-            fw_info: Framework-specific information object.
+            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,
@@ -100,10 +100,10 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
         Determines whether a node is an entry node in a Keras model.
 
         Args:
-            node: The node to be checked.
+            node (BaseNode): The node to be checked.
 
         Returns:
-            Boolean indicating if the node is an entry node.
+            bool: Boolean indicating if the node is an entry node.
         """
         return _is_keras_node_pruning_section_edge(node)
 
@@ -115,26 +115,26 @@ class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementa
         Determines whether a node is an exit node in a Keras model.
 
         Args:
-            node: The node to be checked.
-            corresponding_entry_node: The entry node of the pruning section that is checked.
-            fw_info: Framework-specific information object.
+            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:
-            Boolean indicating if the node is an exit node.
+            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)
 
-    def is_node_intermediate_pruning_section(self, node) -> bool:
+    def is_node_intermediate_pruning_section(self, node: BaseNode) -> bool:
         """
         Determines whether a node is part of the intermediate section in the pruning process of a Keras model.
 
         Args:
-            node: The node to be checked.
+            node (BaseNode): The node to be checked.
 
         Returns:
-            Boolean indicating if the node is part of the intermediate pruning section.
+            bool: Boolean indicating if the node is part of the intermediate pruning section.
         """
         # Nodes that are not Conv2D, Conv2DTranspose, DepthwiseConv2D, or Dense are considered intermediate.
         return node.type not in [keras.layers.DepthwiseConv2D,

model_compression_toolkit/core/pytorch/constants.py

@@ -62,6 +62,9 @@ DIM = 'dim'
 IN_CHANNELS = 'in_channels'
 OUT_CHANNELS = 'out_channels'
 NUM_FEATURES = 'num_features'
+NUM_PARAMETERS = 'num_parameters'
+IN_FEATURES = 'in_features'
+OUT_FEATURES = 'out_features'
 
 # torch devices
 CUDA = 'cuda'

model_compression_toolkit/core/pytorch/kpi_data_facade.py

@@ -22,7 +22,7 @@ from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi_data import compute_kpi_data
 from model_compression_toolkit.core.common.quantization.core_config import CoreConfig
-from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig
 from model_compression_toolkit.constants import FOUND_TORCH
 
 if FOUND_TORCH:
@@ -38,7 +38,7 @@ if FOUND_TORCH:
 
     def pytorch_kpi_data(in_model: Module,
                          representative_data_gen: Callable,
-                         core_config: CoreConfig = CoreConfig(), # TODO: Why pytorch is initilized and keras not?
+                         core_config: CoreConfig = CoreConfig(mixed_precision_config=MixedPrecisionQuantizationConfig()),
                          fw_info: FrameworkInfo = DEFAULT_PYTORCH_INFO,
                          target_platform_capabilities: TargetPlatformCapabilities = PYTORCH_DEFAULT_TPC) -> KPI:
         """
@@ -75,9 +75,9 @@ if FOUND_TORCH:
 
         """
 
-        if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfigV2):
-            Logger.error("KPI data computation can't be executed without MixedPrecisionQuantizationConfigV2 object."
-                         "Given quant_config is not of type MixedPrecisionQuantizationConfigV2.")
+        if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
+            Logger.error("KPI data computation can't be executed without MixedPrecisionQuantizationConfig object."
+                         "Given quant_config is not of type MixedPrecisionQuantizationConfig.")
 
         fw_impl = PytorchImplementation()
 

model_compression_toolkit/core/pytorch/pruning/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================

model_compression_toolkit/core/pytorch/pruning/pruning_pytorch_implementation.py

@@ -0,0 +1,315 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from typing import Tuple, Dict
+
+from model_compression_toolkit.core.common.pruning.pruning_framework_implementation import \
+    PruningFrameworkImplementation
+from model_compression_toolkit.core.common.pruning.pruning_section import PruningSection
+from model_compression_toolkit.core.pytorch.pytorch_implementation import PytorchImplementation
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+from model_compression_toolkit.core.common import BaseNode
+from model_compression_toolkit.core.pytorch.constants import BIAS, GROUPS, OUT_CHANNELS, OUT_FEATURES, NUM_FEATURES, \
+    IN_CHANNELS, IN_FEATURES, NUM_PARAMETERS
+import torch
+
+import numpy as np
+
+from model_compression_toolkit.logger import Logger
+
+
+class PruningPytorchImplementation(PytorchImplementation, PruningFrameworkImplementation):
+    """
+    Implementation of the PruningFramework for the Pytorch framework. This class provides
+    concrete implementations of the abstract methods defined in PruningFrameworkImplementation
+    for the Pytorch framework.
+    """
+
+    def prune_entry_node(self,
+                         node: BaseNode,
+                         output_mask: np.ndarray,
+                         fw_info: FrameworkInfo):
+        """
+        Prunes the entry node of a model in Pytorch.
+
+        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_pytorch_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):
+        """
+        Prunes an intermediate node in a Pytorch model.
+
+        Args:
+            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.
+
+        """
+        # TODO (reuvenp/liord): Address handling of node parameters that can be either a single value across all channels or distinct per channel, e.g., PReLU. Consider developing a structured approach.
+        pruning_en = True
+        _edit_node_input_shape(node, input_mask, fw_info)
+        pruned_parameters = {}
+        mask_bool = output_mask.astype(bool)
+        node.weights = pruned_parameters
+        if node.type == torch.nn.BatchNorm2d:
+            node.framework_attr[NUM_FEATURES] = int(np.sum(input_mask))
+        elif node.type == torch.nn.PReLU:
+            if node.framework_attr[NUM_PARAMETERS] > 1:
+                node.framework_attr[NUM_PARAMETERS] = int(np.sum(input_mask))
+            else:
+                pruning_en = False
+
+        if pruning_en:
+            for k, v in node.weights.items():
+                # Apply the mask to the weights.
+                pruned_parameters[k] = v.compress(mask_bool, axis=-1)
+
+    def prune_exit_node(self,
+                        node: BaseNode,
+                        input_mask: np.ndarray,
+                        fw_info: FrameworkInfo):
+        """
+        Prunes the exit node of a model in Pytorch.
+
+        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_pytorch_edge_node(node=node,
+                                        mask=input_mask,
+                                        fw_info=fw_info,
+                                        is_exit_node=True)
+
+    def is_node_entry_node(self, node: BaseNode) -> bool:
+        """
+        Determines whether a node is an entry node in a Pytorch model.
+
+        Args:
+            node (BaseNode): The node to be checked.
+
+        Returns:
+            bool: Boolean indicating if the node is an entry node.
+        """
+        return _is_pytorch_node_pruning_section_edge(node)
+
+    def is_node_exit_node(self,
+                          node: BaseNode,
+                          corresponding_entry_node: BaseNode,
+                          fw_info: FrameworkInfo) -> bool:
+        """
+        Determines whether a node is an exit node in a Pytorch 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_pytorch_node_pruning_section_edge(node) and PruningSection.has_matching_channel_count(node,
+                                                                                                         corresponding_entry_node,
+                                                                                                         fw_info)
+
+    def is_node_intermediate_pruning_section(self, node: BaseNode) -> bool:
+        """
+        Determines whether a node is part of the intermediate section in the pruning process of a Pytorch model.
+
+        Args:
+            node (BaseNode): The node to be checked.
+
+        Returns:
+            bool: Boolean indicating if the node is part of the intermediate pruning section.
+        """
+        # Nodes that are not Conv2d, ConvTranspose2d, or Linear are considered intermediate.
+        # For PReLU prune attributes only if there is a parameter per channel
+        return node.type not in [torch.nn.Conv2d,
+                                 torch.nn.ConvTranspose2d,
+                                 torch.nn.Linear]
+
+    def attrs_oi_channels_info_for_pruning(self,
+                                           node: BaseNode,
+                                           fw_info: FrameworkInfo) -> 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.
+
+        Not all attributes of a node are directly associated with both input and output channels.
+        For example, bias vectors in convolutional layers are solely related to the number of output
+        channels and do not have a corresponding input channel dimension.
+        In cases like that, None is returned in the tuple of axis for such attributes.
+
+        For kernel operations (like convolutions), the function identifies the output and input
+        channel axis based on framework-specific information.
+        For non-kernel operations, it defaults to setting the last axis as the output
+        channel axis, assuming no specific input channel axis.
+
+        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 'weight' or 'bias')
+            and each value is a tuple representing the output and input channel axis indices respectively.
+        """
+
+        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.error(f"Expected to find attributes but found {kernel_attributes}")
+
+            for attr in kernel_attributes:
+                attributes_with_axis[attr] = fw_info.kernel_channels_mapping.get(node.type)
+
+            # 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.
+            attributes_with_axis[BIAS] = (0, None)
+        else:
+            # We have several assumptions here:
+            # 1. For intermediate nodes, we prune all nodes' weights.
+            # 2. The output channel axis is the last axis of this attribute.
+            # 3. The input channel axis is irrelevant since these attributes are pruned only by
+            #    their output channels.
+            for attr in list(node.weights.keys()):
+                # If the number of float parameters is 1 or less - is the case where
+                # we have one parameter for all channels. For this case, we don't
+                # want to prune the parameter.
+                if node.get_num_parameters(fw_info)[1] <= 1:
+                    attributes_with_axis[attr] = (None, None)
+                else:
+                    attributes_with_axis[attr] = (-1, None)
+
+        return attributes_with_axis
+
+
+def _is_pytorch_node_pruning_section_edge(node: BaseNode) -> bool:
+    """
+    Determines if a Pytorch node is an edge of a pruning section.
+
+    In the context of pruning, an 'edge' node is a layer that can potentially be pruned.
+    This function identifies such nodes based on their type and attributes. Specifically,
+    Conv2d and ConvTranspose2d layers with 'groups' attribute set to 1, and Linear layers
+    are considered as edges for pruning sections.
+
+    Args:
+        node (BaseNode): The node to be evaluated.
+
+    Returns:
+        bool: True if the node is an edge of a pruning section, False otherwise.
+    """
+
+    # Check if the node is a Conv2D or Conv2DTranspose layer with groups set to 1.
+    if node.type in [torch.nn.Conv2d, torch.nn.ConvTranspose2d]:
+        return node.framework_attr[GROUPS] == 1
+    return node.type == torch.nn.Linear
+
+
+def _prune_pytorch_edge_node(node: BaseNode,
+                             mask: np.ndarray,
+                             fw_info: FrameworkInfo,
+                             is_exit_node: bool):
+    """
+    Prunes the given Pytorch node by applying the mask to the node's weights (weights and biases).
+    This function can handle both entry and exit nodes by specifying the is_exit_node parameter.
+
+    Args:
+        node (BaseNode): The node to be pruned.
+        mask (np.ndarray): The pruning mask to be applied.
+        fw_info (FrameworkInfo): Framework-specific information object.
+        is_exit_node (bool): 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)
+    # 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)
+
+    if not is_exit_node and node.framework_attr[BIAS]:
+        # Prune the bias if applicable and it's an entry node.
+        bias = node.get_weights_by_keys(BIAS)
+        pruned_bias = bias.compress(mask_bool)
+        node.set_weights_by_keys(name=BIAS, tensor=pruned_bias)
+
+    if not is_exit_node:
+        # Update 'out_channels' or 'out_features' attributes for entry nodes
+        # Conv2d,ConvTranspose2d / Linear layers.
+        if node.type in [torch.nn.Conv2d, torch.nn.ConvTranspose2d]:
+            node.framework_attr[OUT_CHANNELS] = int(np.sum(mask))
+        elif node.type == torch.nn.Linear:
+            node.framework_attr[OUT_FEATURES] = int(np.sum(mask))
+        else:
+            Logger.exception(f"{node.type} is currently not supported"
+                             f"as an edge node in a pruning section")
+
+    if is_exit_node:
+        if node.type in [torch.nn.Conv2d, torch.nn.ConvTranspose2d]:
+            node.framework_attr[IN_CHANNELS] = int(np.sum(mask))
+        elif node.type == torch.nn.Linear:
+            node.framework_attr[IN_FEATURES] = int(np.sum(mask))
+        else:
+            Logger.exception(f"{node.type} is currently not supported"
+                             f"as an edge node in a pruning section")
+        # Adjust the input shape for the last node in the section.
+        _edit_node_input_shape(node, mask_bool, fw_info)
+
+
+def _edit_node_input_shape(node: BaseNode,
+                           input_mask: np.ndarray,
+                           fw_info: FrameworkInfo):
+    """
+    Adjusts the input shape of a node based on the given input mask.
+
+    This function modifies the input shape of the given node to reflect the pruning
+    that has taken place. It updates the last dimension of the node's input shape
+    to match the number of channels that remain after pruning.
+
+    Args:
+        node (BaseNode): The node whose input shape needs to be adjusted.
+        input_mask (np.ndarray): A binary array where 1 indicates the channel is kept and 0 means pruned.
+        fw_info (FrameworkInfo): Framework-specific information object.
+    """
+    # Start with the current input shape of the node.
+    new_input_shape = list(node.input_shape)
+
+    # Adjust the last dimension of the shape to match the number of unpruned (retained) channels.
+    # This is done by summing the mask, as each '1' in the mask represents a retained channel.
+    channel_axis = fw_info.out_channel_axis_mapping.get(node.type)
+    new_input_shape[0][channel_axis] = int(np.sum(input_mask))
+
+    # Update the node's input shape with the new dimensions.
+    node.input_shape = tuple(new_input_shape)

model_compression_toolkit/core/pytorch/pytorch_implementation.py

@@ -26,7 +26,7 @@ from torch.nn import Module, Sigmoid, Softmax
 
 import model_compression_toolkit.core.pytorch.constants as pytorch_constants
 from model_compression_toolkit.constants import HESSIAN_NUM_ITERATIONS
-from model_compression_toolkit.core import QuantizationConfig, FrameworkInfo, CoreConfig, MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core import QuantizationConfig, FrameworkInfo, CoreConfig, MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.framework_implementation import FrameworkImplementation
@@ -332,7 +332,7 @@ class PytorchImplementation(FrameworkImplementation):
 
     def get_sensitivity_evaluator(self,
                                   graph: Graph,
-                                  quant_config: MixedPrecisionQuantizationConfigV2,
+                                  quant_config: MixedPrecisionQuantizationConfig,
                                   representative_data_gen: Callable,
                                   fw_info: FrameworkInfo,
                                   disable_activation_for_metric: bool = False,

model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -24,7 +24,7 @@ from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.gptq.common.gptq_config import GradientPTQConfigV2
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
-from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import MixedPrecisionQuantizationConfig
 from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.runner import core_runner
 from model_compression_toolkit.gptq.runner import gptq_runner
@@ -177,7 +177,7 @@ if FOUND_TF:
             with different bitwidths for different layers.
             The candidates bitwidth for quantization should be defined in the target platform model:
 
-            >>> config = mct.core.CoreConfig(mixed_precision_config=mct.core.MixedPrecisionQuantizationConfigV2(num_of_images=1))
+            >>> config = mct.core.CoreConfig(mixed_precision_config=mct.core.MixedPrecisionQuantizationConfig(num_of_images=1))
 
             For mixed-precision set a target KPI object:
             Create a KPI object to limit our returned model's size. Note that this value affects only coefficients
@@ -199,9 +199,9 @@ if FOUND_TF:
                              fw_info=fw_info).validate()
 
         if core_config.mixed_precision_enable:
-            if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfigV2):
+            if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
                 Logger.error("Given quantization config to mixed-precision facade is not of type "
-                             "MixedPrecisionQuantizationConfigV2. Please use keras_post_training_quantization "
+                             "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)

model_compression_toolkit/gptq/pytorch/quantization_facade.py

@@ -29,7 +29,7 @@ from model_compression_toolkit.core.exporter import export_model
 from model_compression_toolkit.core.analyzer import analyzer_model_quantization
 from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
-    MixedPrecisionQuantizationConfigV2
+    MixedPrecisionQuantizationConfig
 
 LR_DEFAULT = 1e-4
 LR_REST_DEFAULT = 1e-4
@@ -157,9 +157,9 @@ if FOUND_TORCH:
         """
 
         if core_config.mixed_precision_enable:
-            if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfigV2):
+            if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
                 Logger.error("Given quantization config to mixed-precision facade is not of type "
-                             "MixedPrecisionQuantizationConfigV2. Please use keras_post_training_quantization "
+                             "MixedPrecisionQuantizationConfig. Please use keras_post_training_quantization "
                              "API, or pass a valid mixed precision configuration.")  # pragma: no cover
 
         tb_w = init_tensorboard_writer(DEFAULT_PYTORCH_INFO)

model_compression_toolkit/pruning/__init__.py

@@ -16,4 +16,5 @@
 from model_compression_toolkit.core.common.pruning.pruning_info import PruningInfo
 from model_compression_toolkit.core.common.pruning.pruning_config import ImportanceMetric, PruningConfig, ChannelsFilteringStrategy
 from model_compression_toolkit.pruning.keras.pruning_facade import keras_pruning_experimental
+from model_compression_toolkit.pruning.pytorch.pruning_facade import pytorch_pruning_experimental
 

model_compression_toolkit/pruning/pytorch/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2024 Sony Semiconductor Israel, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================