mct-nightly 1.8.0.8032023.post421__py3-none-any.whl → 1.8.0.8052023.post414__py3-none-any.whl

model_compression_toolkit/core/keras/back2framework/keras_model_builder.py

@@ -13,8 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from abc import abstractmethod
-
 import tensorflow as tf
 from keras.engine.input_layer import InputLayer
 from keras.models import Model, clone_model
@@ -22,7 +20,7 @@ from packaging import version
 
 from model_compression_toolkit.core.common.back2framework.base_model_builder import BaseModelBuilder
 from model_compression_toolkit.core.common.user_info import UserInformation
-from model_compression_toolkit.core.common.constants import INPUT_BASE_NAME
+from model_compression_toolkit.constants import INPUT_BASE_NAME
 
 # As from Tensorflow 2.6, keras is a separate package and some classes should be imported differently.
 if version.parse(tf.__version__) < version.parse("2.6"):
@@ -38,7 +36,6 @@ else:
 from typing import Any, Dict, List, Tuple, Callable
 from tensorflow.python.util.object_identity import Reference as TFReference
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
-from model_compression_toolkit.core.common.logger import Logger
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO

model_compression_toolkit/core/keras/back2framework/mixed_precision_model_builder.py

@@ -36,7 +36,7 @@ else:
     from keras.layers.core import TFOpLambda, SlicingOpLambda
 
 from tensorflow_model_optimization.python.core.quantization.keras.quantize_wrapper import QuantizeWrapper
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO

model_compression_toolkit/core/keras/back2framework/model_gradients.py

@@ -26,13 +26,13 @@ else:
 
 from typing import Any, Dict, List, Tuple
 from tensorflow.python.util.object_identity import Reference as TFReference
-from model_compression_toolkit.core.common.constants import EPS, MIN_JACOBIANS_ITER, JACOBIANS_COMP_TOLERANCE
+from model_compression_toolkit.constants import EPS, MIN_JACOBIANS_ITER, JACOBIANS_COMP_TOLERANCE
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode, Graph
 from model_compression_toolkit.core.common.graph.edge import EDGE_SINK_INDEX
 from model_compression_toolkit.core.keras.back2framework.instance_builder import OperationHandler
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 
 
 def build_input_tensors_list(node: BaseNode,
@@ -171,8 +171,9 @@ def keras_iterative_approx_jacobian_trace(graph_float: common.Graph,
 
                     # If the change to the mean Jacobian approximation is insignificant we stop the calculation
                     if j > MIN_JACOBIANS_ITER:
-                        delta = np.mean([jac_trace_approx, *trace_jv]) - np.mean(trace_jv)
-                        if np.abs(delta) / (np.abs(np.mean(trace_jv)) + 1e-6) < JACOBIANS_COMP_TOLERANCE:
+                        new_mean = np.mean([jac_trace_approx, *trace_jv])
+                        delta = new_mean - np.mean(trace_jv)
+                        if np.abs(delta) / (np.abs(new_mean) + 1e-6) < JACOBIANS_COMP_TOLERANCE:
                             trace_jv.append(jac_trace_approx)
                             break
 

model_compression_toolkit/core/keras/back2framework/quantized_model_builder.py

@@ -14,7 +14,7 @@
 # ==============================================================================
 from typing import List
 
-from model_compression_toolkit import FrameworkInfo
+from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.keras.back2framework.keras_model_builder import KerasModelBuilder

model_compression_toolkit/core/keras/constants.py

@@ -99,13 +99,6 @@ OUTPUT_BIAS = '/attention_output/bias'
 # ReLU bound constants
 RELU_POT_BOUND = 8.0
 
-# Supported TP models names for Tensorflow:
-DEFAULT_TP_MODEL = 'default'
-IMX500_TP_MODEL = 'imx500'
-TFLITE_TP_MODEL = 'tflite'
-QNNPACK_TP_MODEL = 'qnnpack'
-
-
 # TFOpLambda functions:
 ADD = 'add'
 PAD = 'pad'

model_compression_toolkit/core/keras/default_framework_info.py

@@ -25,9 +25,9 @@ else:
     from keras.layers import Conv2D, DepthwiseConv2D, Dense, Conv2DTranspose, Softmax, ELU
 
 from model_compression_toolkit.core.common.defaultdict import DefaultDict
-from model_compression_toolkit.core.common.framework_info import FrameworkInfo, ChannelAxis
-from model_compression_toolkit.core.common.target_platform import QuantizationMethod
-from model_compression_toolkit.core.common.constants import SOFTMAX_THRESHOLD
+from model_compression_toolkit.core.common.framework_info import FrameworkInfo
+from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
+from model_compression_toolkit.constants import SOFTMAX_THRESHOLD
 from model_compression_toolkit.core.keras.constants import SOFTMAX, LINEAR, RELU, SWISH, SIGMOID, IDENTITY, TANH, SELU, \
     KERNEL, DEPTHWISE_KERNEL
 from model_compression_toolkit.core.keras.quantizer.fake_quant_builder import power_of_two_quantization, symmetric_quantization, uniform_quantization

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

@@ -17,7 +17,7 @@
 from tensorflow.keras.layers import Dense, DepthwiseConv2D, Conv2D, Conv2DTranspose, Activation, SeparableConv2D
 
 from model_compression_toolkit.core import common
-from model_compression_toolkit.core.common.constants import FLOAT_32, DATA_TYPE
+from model_compression_toolkit.constants import FLOAT_32, DATA_TYPE
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher, \
     NodeFrameworkAttrMatcher

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

@@ -23,7 +23,7 @@ from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher, EdgeMatcher, WalkMatcher
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
 from model_compression_toolkit.core.common.quantization.quantization_config import QuantizationConfig
-from model_compression_toolkit.core.common.constants import THRESHOLD
+from model_compression_toolkit.constants import THRESHOLD
 from model_compression_toolkit.core.keras.constants import KERNEL
 
 input_node = NodeOperationMatcher(InputLayer)

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

@@ -21,7 +21,7 @@ from model_compression_toolkit.core.common.graph.graph_matchers import NodeOpera
 from model_compression_toolkit.core.common.substitutions.linear_collapsing import Conv2DCollapsing
 from model_compression_toolkit.core.keras.constants import KERNEL, KERNEL_SIZE, STRIDES, DILATIONS, LINEAR, \
     ACTIVATION, BIAS, USE_BIAS, LAYER_NAME, FILTERS, PADDING, GROUPS, DATA_FORMAT
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 
 
 def linear_collapsing_node_matchers() -> Tuple[NodeOperationMatcher, NodeOperationMatcher]:

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

@@ -23,17 +23,16 @@ else:
     from keras.layers.core import TFOpLambda
     from keras.layers import MultiHeadAttention, Conv2D, Softmax, Concatenate, Reshape, Permute
 
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.graph.base_graph import Graph, BaseNode, OutTensor
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
-from model_compression_toolkit.core.common.constants import REUSE, REUSE_GROUP
-from model_compression_toolkit.core.keras.reader.node_builder import REUSED_IDENTIFIER
+from model_compression_toolkit.constants import REUSE, REUSE_GROUP
 from model_compression_toolkit.core.keras.constants import KERNEL, BIAS, USE_BIAS, NUM_HEADS, KEY_DIM, VALUE_DIM, \
     QUERY_SHAPE, KEY_SHAPE, VALUE_SHAPE, OUTPUT_SHAPE, ATTENTION_AXES, ACTIVATION, LINEAR, FILTERS, \
     FUNCTION, DIMS, TARGET_SHAPE, F_STRIDED_SLICE, F_STACK, Q_KERNEL, Q_BIAS, K_KERNEL, K_BIAS, V_KERNEL, V_BIAS, \
-    OUTPUT_KERNEL, OUTPUT_BIAS, F_MATMUL, TRANSPOSE_B, KERNEL_SIZE, AXIS, F_STRIDED_SLICE_BEGIN, F_STRIDED_SLICE_END
+    OUTPUT_KERNEL, OUTPUT_BIAS, F_MATMUL, KERNEL_SIZE, AXIS, F_STRIDED_SLICE_BEGIN, F_STRIDED_SLICE_END
 
 
 class MHAParams:

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

@@ -23,6 +23,7 @@ from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher, WalkMatcher
 from model_compression_toolkit.core.keras.constants import KERNEL, BIAS, RELU_MAX_VALUE, RELU_POT_BOUND
+from model_compression_toolkit.logger import Logger
 
 
 class ReLUBoundToPowerOfTwo(common.BaseSubstitution):
@@ -81,7 +82,7 @@ class ReLUBoundToPowerOfTwo(common.BaseSubstitution):
         scale_factor = max_value / self.threshold
 
         non_linear_node.framework_attr[RELU_MAX_VALUE] = np.float32(self.threshold)
-        common.Logger.debug(
+        Logger.debug(
             f"Node named:{non_linear_node.name} max value change "
             f"to:{non_linear_node.framework_attr[RELU_MAX_VALUE]}")
 

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

@@ -20,7 +20,8 @@ from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher,NodeFrameworkAttrMatcher
 from model_compression_toolkit.core.keras.constants import RELU_MAX_VALUE
-from model_compression_toolkit.core.common.constants import THRESHOLD
+from model_compression_toolkit.constants import THRESHOLD
+from model_compression_toolkit.logger import Logger
 
 MATCHER = NodeOperationMatcher(ReLU) & NodeFrameworkAttrMatcher(RELU_MAX_VALUE, None).logic_not()
 
@@ -56,5 +57,5 @@ class RemoveReLUUpperBound(common.BaseSubstitution):
                 node.final_activation_quantization_cfg.activation_quantization_params.get(THRESHOLD) == \
                 node.framework_attr.get(RELU_MAX_VALUE):
             node.framework_attr[RELU_MAX_VALUE] = None
-            common.Logger.info(f'Removing upper bound of {node.name}. Threshold and upper bound are equal.')
+            Logger.info(f'Removing upper bound of {node.name}. Threshold and upper bound are equal.')
         return graph

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

@@ -21,7 +21,7 @@ from model_compression_toolkit.core.common.graph.graph_matchers import NodeOpera
     NodeFrameworkAttrMatcher
 from model_compression_toolkit.core.common.substitutions.residual_collapsing import ResidualCollapsing
 from model_compression_toolkit.core.keras.constants import KERNEL, LINEAR, ACTIVATION, LAYER_NAME
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 
 
 def residual_collapsing_node_matchers() -> Tuple[NodeOperationMatcher, NodeOperationMatcher]:

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

@@ -22,7 +22,7 @@ from tensorflow.python.keras.layers.core import TFOpLambda
 from tensorflow.keras.layers import Activation, Conv2D, Dense, DepthwiseConv2D, ZeroPadding2D, Reshape, \
     GlobalAveragePooling2D, Dropout, ReLU, PReLU, ELU
 
-from model_compression_toolkit import CoreConfig, FrameworkInfo
+from model_compression_toolkit.core import CoreConfig, FrameworkInfo
 from model_compression_toolkit.core.common import BaseNode, Graph
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher, \

model_compression_toolkit/core/keras/keras_implementation.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from typing import List, Any, Tuple, Callable, Type, Dict
+from typing import List, Any, Tuple, Callable, Dict
 
 import numpy as np
 import tensorflow as tf
@@ -43,7 +43,7 @@ else:
         Concatenate, Add
     from keras.layers.core import TFOpLambda
 
-from model_compression_toolkit import QuantizationConfig, FrameworkInfo, CoreConfig, MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core import QuantizationConfig, FrameworkInfo, CoreConfig, MixedPrecisionQuantizationConfigV2
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import Graph, BaseNode
 from model_compression_toolkit.core.common.collectors.statistics_collector import BaseStatsCollector
@@ -52,8 +52,6 @@ from model_compression_toolkit.core.common.model_builder_mode import ModelBuilde
 from model_compression_toolkit.core.common.node_prior_info import NodePriorInfo
 from model_compression_toolkit.core.common.user_info import UserInformation
 from model_compression_toolkit.core.keras.default_framework_info import DEFAULT_KERAS_INFO
-from model_compression_toolkit.gptq.common.gptq_training import GPTQTrainer
-from model_compression_toolkit.gptq.keras.gptq_training import KerasGPTQTrainer
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.activation_decomposition import \
     ActivationDecomposition
 from model_compression_toolkit.core.keras.graph_substitutions.substitutions.softmax_shift import \
@@ -348,12 +346,6 @@ class KerasImplementation(FrameworkImplementation):
             substitutions_list.append(keras_batchnorm_refusing())
         return substitutions_list
 
-    def get_gptq_trainer_obj(self) -> Type[GPTQTrainer]:
-        """
-        Returns:  Keras object of GPTQTrainer
-        """
-        return KerasGPTQTrainer
-
     def get_sensitivity_evaluator(self,
                                   graph: Graph,
                                   quant_config: MixedPrecisionQuantizationConfigV2,

model_compression_toolkit/core/keras/keras_model_validation.py

@@ -1,6 +1,6 @@
 from tensorflow.keras.models import Model
 
-from model_compression_toolkit import FrameworkInfo
+from model_compression_toolkit.core import FrameworkInfo
 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

model_compression_toolkit/core/keras/keras_node_prior_info.py

@@ -8,7 +8,7 @@ if version.parse(tf.__version__) < version.parse("2.6"):
 else:
     from keras.layers import Activation, ReLU, BatchNormalization
 
-from model_compression_toolkit import FrameworkInfo
+from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.node_prior_info import NodePriorInfo
 from model_compression_toolkit.core.keras.constants import ACTIVATION, RELU_MAX_VALUE, NEGATIVE_SLOPE, THRESHOLD, \

model_compression_toolkit/core/keras/kpi_data_facade.py

@@ -15,19 +15,19 @@
 
 from typing import Callable
 
-from model_compression_toolkit import MixedPrecisionQuantizationConfig, CoreConfig, MixedPrecisionQuantizationConfigV2
+from model_compression_toolkit.core import MixedPrecisionQuantizationConfig, CoreConfig, MixedPrecisionQuantizationConfigV2
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi import KPI
-from model_compression_toolkit.core.common import Logger
-from model_compression_toolkit.core.common.constants import TENSORFLOW
-from model_compression_toolkit.core.common.target_platform import TargetPlatformCapabilities
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import TENSORFLOW
+from model_compression_toolkit.target_platform_capabilities.target_platform import TargetPlatformCapabilities
 from model_compression_toolkit.core.common.mixed_precision.kpi_tools.kpi_data import compute_kpi_data
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo
 from model_compression_toolkit.core.common.mixed_precision.mixed_precision_quantization_config import \
     DEFAULT_MIXEDPRECISION_CONFIG
-from model_compression_toolkit.core.common.constants import FOUND_TF
+from model_compression_toolkit.constants import FOUND_TF
 
 if FOUND_TF:
-    from model_compression_toolkit.core.keras.constants import DEFAULT_TP_MODEL
+    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
@@ -51,7 +51,7 @@ if FOUND_TF:
             representative_data_gen (Callable): Dataset used for calibration.
             quant_config (MixedPrecisionQuantizationConfig): MixedPrecisionQuantizationConfig containing parameters of how the model should be quantized.
             fw_info (FrameworkInfo): Information needed for quantization about the specific framework (e.g., kernel channels indices, groups of layers by how they should be quantized, etc.). `Default Keras info <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/core/keras/default_framework_info.py>`_
-            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to. `Default Keras TPC <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/core/tpc_models/keras_tp_models/keras_default.py>`_
+            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
 
         Returns:
             A KPI object with total weights parameters sum, max activation tensor and total kpi.
@@ -75,7 +75,7 @@ if FOUND_TF:
             Import MCT and call for KPI data calculation:
 
             >>> import model_compression_toolkit as mct
-            >>> kpi_data = mct.keras_kpi_data(model, repr_datagen)
+            >>> kpi_data = mct.core.keras_kpi_data(model, repr_datagen)
 
 
         """
@@ -112,7 +112,7 @@ if FOUND_TF:
             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.
             fw_info (FrameworkInfo): Information needed for quantization about the specific framework (e.g., kernel channels indices, groups of layers by how they should be quantized, etc.). `Default Keras info <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/core/keras/default_framework_info.py>`_
-            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to. `Default Keras TPC <https://github.com/sony/model_optimization/blob/main/model_compression_toolkit/core/tpc_models/keras_tp_models/keras_default.py>`_
+            target_platform_capabilities (TargetPlatformCapabilities): TargetPlatformCapabilities to optimize the Keras model according to.
 
         Returns:
 
@@ -133,7 +133,7 @@ if FOUND_TF:
             Import MCT and call for KPI data calculation:
 
             >>> import model_compression_toolkit as mct
-            >>> kpi_data = mct.keras_kpi_data(model, repr_datagen)
+            >>> kpi_data = mct.core.keras_kpi_data(model, repr_datagen)
 
         """
 

model_compression_toolkit/core/keras/quantizer/fake_quant_builder.py

@@ -20,8 +20,8 @@ import tensorflow as tf
 import numpy as np
 from tensorflow.python.util.object_identity import Reference as TFReference
 
-from model_compression_toolkit.core.common.logger import Logger
-from model_compression_toolkit.core.common.constants import THRESHOLD, SIGNED, RANGE_MIN, RANGE_MAX
+from model_compression_toolkit.logger import Logger
+from model_compression_toolkit.constants import THRESHOLD, SIGNED, RANGE_MIN, RANGE_MAX
 from model_compression_toolkit.core.common.quantization.quantizers.uniform_quantizers import threshold_is_power_of_two
 
 

model_compression_toolkit/core/keras/quantizer/input_layer_quantize_transform.py

@@ -23,7 +23,7 @@ from tensorflow_model_optimization.python.core.quantization.keras.quantize_confi
 
 
 from model_compression_toolkit.core.common import BaseNode
-from model_compression_toolkit.core.common.constants import INPUT_BASE_NAME
+from model_compression_toolkit.constants import INPUT_BASE_NAME
 
 
 class InputLayerWrapperTransform(InputLayerQuantize):

model_compression_toolkit/core/keras/quantizer/lut_fake_quant.py

@@ -1,11 +1,11 @@
-from typing import Tuple, Dict, Any, Callable
+from typing import Tuple, Dict, Callable
 
 import numpy as np
 import tensorflow as tf
 from keras.layers import Layer
 from tensorflow.python.util.object_identity import Reference as TFReference
 
-from model_compression_toolkit.core.common.constants import SIGNED, CLUSTER_CENTERS, EPS, \
+from model_compression_toolkit.constants import SIGNED, CLUSTER_CENTERS, EPS, \
     MULTIPLIER_N_BITS, THRESHOLD
 
 

model_compression_toolkit/core/keras/quantizer/mixed_precision/selective_quantize_config.py

@@ -24,7 +24,7 @@ from model_compression_toolkit.core.common.quantization.candidate_node_quantizat
 from model_compression_toolkit.core.keras.quantizer.mixed_precision.selective_activation_quantizer import \
     SelectiveActivationQuantizer
 from packaging import version
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 
 if version.parse(tf.__version__) < version.parse("2.6"):
     from tensorflow.python.keras.layers import Layer  # pragma: no cover

model_compression_toolkit/core/keras/reader/common.py

@@ -29,7 +29,7 @@ else:
     from keras.engine.functional import Functional
     from keras.engine.sequential import Sequential
 
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
 
 

model_compression_toolkit/core/keras/statistics_correction/apply_second_moment_correction.py

@@ -19,7 +19,7 @@ from tensorflow.keras.layers import BatchNormalization
 from tqdm import tqdm
 
 import model_compression_toolkit.core.keras.constants as keras_constants
-from model_compression_toolkit import CoreConfig
+from model_compression_toolkit.core import CoreConfig
 from model_compression_toolkit.core import common
 
 

model_compression_toolkit/core/pytorch/back2framework/factory_model_builder.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-from model_compression_toolkit.core.common import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.model_builder_mode import ModelBuilderMode
 from model_compression_toolkit.core.pytorch.back2framework.float_model_builder import FloatPyTorchModelBuilder
 from model_compression_toolkit.core.pytorch.back2framework.mixed_precision_model_builder import \

model_compression_toolkit/core/pytorch/back2framework/float_model_builder.py

@@ -17,7 +17,7 @@ from typing import List, Tuple
 
 import torch
 
-from model_compression_toolkit import FrameworkInfo
+from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.user_info import UserInformation

model_compression_toolkit/core/pytorch/back2framework/mixed_precision_model_builder.py

@@ -17,7 +17,7 @@ from typing import List, Any, Tuple
 
 import torch
 
-from model_compression_toolkit import FrameworkInfo
+from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode

model_compression_toolkit/core/pytorch/back2framework/model_gradients.py

@@ -22,13 +22,14 @@ import numpy as np
 
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode, Graph
-from model_compression_toolkit.core.common.constants import EPS, MIN_JACOBIANS_ITER, JACOBIANS_COMP_TOLERANCE
+from model_compression_toolkit.constants import EPS, MIN_JACOBIANS_ITER, JACOBIANS_COMP_TOLERANCE
 from model_compression_toolkit.core.common.graph.edge import EDGE_SINK_INDEX
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode
 from model_compression_toolkit.core.pytorch.back2framework.instance_builder import node_builder
-from model_compression_toolkit.core.pytorch.reader.node_holders import DummyPlaceHolder
-from model_compression_toolkit.core.pytorch.utils import torch_tensor_to_numpy
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.core.pytorch.constants import BUFFER
+from model_compression_toolkit.core.pytorch.reader.node_holders import DummyPlaceHolder, BufferHolder
+from model_compression_toolkit.core.pytorch.utils import torch_tensor_to_numpy, get_working_device
+from model_compression_toolkit.logger import Logger
 
 
 def build_input_tensors_list(node: BaseNode,
@@ -133,7 +134,13 @@ class PytorchModelGradients(torch.nn.Module):
 
         for n in self.node_sort:
             if not isinstance(n, FunctionalNode):
-                self.add_module(n.name, node_builder(n))
+                if n.type == BufferHolder:
+                    self.add_module(n.name, node_builder(n))
+                    self.get_submodule(n.name). \
+                        register_buffer(n.name,
+                                        torch.Tensor(n.get_weights_by_keys(BUFFER)).to(get_working_device()))
+                else:
+                    self.add_module(n.name, node_builder(n))
 
     def forward(self,
                 *args: Any) -> Any:
@@ -289,9 +296,9 @@ def pytorch_iterative_approx_jacobian_trace(graph_float: common.Graph,
 
             # If the change to the mean Jacobian approximation is insignificant we stop the calculation
             if j > MIN_JACOBIANS_ITER:
-                delta = torch.mean(torch.stack([jac_trace_approx, *trace_jv])) - torch.mean(
-                    torch.stack(trace_jv))
-                if torch.abs(delta) / (torch.abs(torch.mean(torch.stack(trace_jv))) + 1e-6) < JACOBIANS_COMP_TOLERANCE:
+                new_mean = torch.mean(torch.stack([jac_trace_approx, *trace_jv]))
+                delta = new_mean - torch.mean(torch.stack(trace_jv))
+                if torch.abs(delta) / (torch.abs(new_mean) + 1e-6) < JACOBIANS_COMP_TOLERANCE:
                     trace_jv.append(jac_trace_approx)
                     break
 

model_compression_toolkit/core/pytorch/back2framework/pytorch_model_builder.py

@@ -18,7 +18,7 @@ from typing import Tuple, Any, Dict, List, Union, Callable
 import torch
 from networkx import topological_sort
 
-from model_compression_toolkit import FrameworkInfo
+from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode, Graph
 from model_compression_toolkit.core.common.back2framework.base_model_builder import BaseModelBuilder

model_compression_toolkit/core/pytorch/back2framework/quantized_model_builder.py

@@ -17,7 +17,7 @@ from typing import List, Tuple
 
 import torch
 
-from model_compression_toolkit import FrameworkInfo
+from model_compression_toolkit.core import FrameworkInfo
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.user_info import UserInformation

model_compression_toolkit/core/pytorch/constants.py

@@ -69,12 +69,6 @@ CPU = 'cpu'
 # ReLU bound constants
 RELU_POT_BOUND = 8.0
 
-# Supported TP models names for Pytorch:
-DEFAULT_TP_MODEL = 'default'
-IMX500_TP_MODEL = 'imx500'
-TFLITE_TP_MODEL = 'tflite'
-QNNPACK_TP_MODEL = 'qnnpack'
-
 # MultiHeadAttention layer attributes:
 EMBED_DIM = 'embed_dim'
 NUM_HEADS = 'num_heads'
@@ -92,3 +86,7 @@ IN_PROJ_WEIGHT = 'in_proj_weight'
 IN_PROJ_BIAS = 'in_proj_bias'
 BIAS_K = 'bias_k'
 BIAS_V = 'bias_v'
+
+# # Batch size value for 'reshape' and 'view' operators,
+# # the value is -1 so the batch size is inferred from the length of the array and remaining dimensions.
+BATCH_DIM_VALUE = -1

model_compression_toolkit/core/pytorch/default_framework_info.py

@@ -19,8 +19,8 @@ from torch import sigmoid
 
 from model_compression_toolkit.core.common.defaultdict import DefaultDict
 from model_compression_toolkit.core.common.framework_info import FrameworkInfo, ChannelAxis
-from model_compression_toolkit.core.common.target_platform import QuantizationMethod
-from model_compression_toolkit.core.common.constants import SOFTMAX_THRESHOLD
+from model_compression_toolkit.target_platform_capabilities.target_platform import QuantizationMethod
+from model_compression_toolkit.constants import SOFTMAX_THRESHOLD
 from model_compression_toolkit.core.pytorch.constants import KERNEL
 from model_compression_toolkit.core.pytorch.quantizer.fake_quant_builder import power_of_two_quantization, \
     symmetric_quantization, uniform_quantization

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/linear_collapsing.py

@@ -22,7 +22,7 @@ from model_compression_toolkit.core.common.graph.graph_matchers import NodeOpera
 from model_compression_toolkit.core.common import BaseNode
 from model_compression_toolkit.core.common.substitutions.linear_collapsing import Conv2DCollapsing
 from model_compression_toolkit.core.pytorch.constants import KERNEL, KERNEL_SIZE, STRIDES, DILATIONS, BIAS, USE_BIAS, FILTERS, PADDING, GROUPS
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 
 
 def linear_collapsing_node_matchers() -> Tuple[NodeOperationMatcher, NodeOperationMatcher]:

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/multi_head_attention_decomposition.py

@@ -20,7 +20,7 @@ import torch.nn as nn
 import operator
 from typing import List
 
-from model_compression_toolkit.core.common.logger import Logger
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.graph.base_graph import Graph, BaseNode, OutTensor
 from model_compression_toolkit.core.common.graph.functional_node import FunctionalNode

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/relu_bound_to_power_of_2.py

@@ -25,6 +25,7 @@ from model_compression_toolkit.core.common.graph.graph_matchers import NodeOpera
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
 from model_compression_toolkit.core.pytorch.constants import KERNEL, BIAS, INPLACE, HARDTANH_MIN_VAL, HARDTANH_MAX_VAL, \
     RELU_POT_BOUND
+from model_compression_toolkit.logger import Logger
 
 
 class ReLUBoundToPowerOfTwo(common.BaseSubstitution):
@@ -102,8 +103,8 @@ class ReLUBoundToPowerOfTwo(common.BaseSubstitution):
             else:
                 return graph
         else:
-            common.Logger.error(f"In substitution with wrong matched pattern")
-        common.Logger.debug(
+            Logger.error(f"In substitution with wrong matched pattern")
+        Logger.debug(
             f"Node named:{non_linear_node.name} changed "
             f"to:{non_linear_node.type}")
 

model_compression_toolkit/core/pytorch/graph_substitutions/substitutions/reshape_with_static_shapes.py

@@ -14,10 +14,13 @@
 # ==============================================================================
 from torch import reshape
 import torch
+
+from model_compression_toolkit.logger import Logger
 from model_compression_toolkit.core.common.graph.graph_matchers import NodeOperationMatcher
 from model_compression_toolkit.core import common
 from model_compression_toolkit.core.common.graph.base_graph import Graph
 from model_compression_toolkit.core.common.graph.base_node import BaseNode
+from model_compression_toolkit.core.pytorch.constants import BATCH_DIM_VALUE
 
 
 class ReshapeWithStaticShapes(common.BaseSubstitution):
@@ -47,14 +50,25 @@ class ReshapeWithStaticShapes(common.BaseSubstitution):
         Returns:
             Graph after applying the substitution.
         """
+        # we want the batch size value to infer from the length of the array and remaining dimensions
+        if len(node.output_shape) == 1:
+            node.output_shape[0][0] = BATCH_DIM_VALUE
+        else:
+            Logger.error('Reshape or view nodes should have a single output shape')  # pragma: no cover
+
         # configure the new static output shape attribute
         node.op_call_args = node.output_shape
 
         # modify the node input info
         node.input_shape = [node.input_shape[0]]
+
+        # the first input is the tensor to be reshaped, we want his batch size value to infer
+        # from the length of the array and remaining dimensions
+        node.input_shape[0][0] = BATCH_DIM_VALUE
+
         nodes_to_check = []
         for in_edge in graph.incoming_edges(node):
-            if in_edge.sink_index > 0: # the first input is the tensor to be reshaped
+            if in_edge.sink_index > 0:  # the first input is the tensor to be reshaped
                 nodes_to_check.append(in_edge.source_node)
                 graph.remove_edge(in_edge.source_node, node)
         for n in nodes_to_check:
@@ -80,4 +94,4 @@ def clean_graph_from_nodes_without_out_edges(graph: Graph,
             graph.remove_edge(in_edge.source_node, node)
         graph.remove_node(node)
         for n in nodes_to_check:
-            clean_graph_from_nodes_without_out_edges(graph, n)
+            clean_graph_from_nodes_without_out_edges(graph, n)