Remove default output axis values from fw_info.

model_compression_toolkit/core/common/collectors/base_collector.py

@@ -13,11 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 
+from abc import ABC, abstractmethod
 import numpy as np
 from model_compression_toolkit.logger import Logger
 
 
-class BaseCollector(object):
+class BaseCollector(ABC):
     """
     Base class for statistics collection object.
     """
@@ -26,6 +27,7 @@ def __init__(self):
         # When manipulation statistics in a granularity they were not collected by, the data is invalid.
         self.is_legal = True
 
+    @abstractmethod
     def scale(self, scale_factor: np.ndarray):
         """
         Scale all statistics in collector by some factor.
@@ -37,6 +39,7 @@ def scale(self, scale_factor: np.ndarray):
         raise NotImplemented(
             f'{self.__class__.__name__} needs to implement scale operation for its state.')  # pragma: no cover
 
+    @abstractmethod
     def shift(self, shift_value: np.ndarray):
         """
         Shift all statistics in collector by some value.

model_compression_toolkit/core/common/collectors/mean_collector.py

@@ -87,7 +87,7 @@ def update(self,
             x: Tensor that goes through the mean collector and needs to be considered in the mean computation.
         """
         self.i += 1  # Update the iteration index
-        axis = (len(x.shape) - 1) if self.axis == LAST_AXIS else self.axis
+        axis = (len(x.shape) - 1) if self.axis == LAST_AXIS or self.axis is None else self.axis
         n = x.shape[axis]
         transpose_index = [axis, *[i for i in range(len(x.shape)) if i != axis]]
         mu = np.mean(np.reshape(np.transpose(x, transpose_index), [n, -1]), axis=-1) # mean per channel for a batch

model_compression_toolkit/core/common/collectors/min_max_per_channel_collector.py

@@ -130,7 +130,7 @@ def update(self,
             x: Tensor that goes through the collector and needs to be considered in the min/max computation.
         """
 
-        axis = (len(x.shape) - 1) if self.axis == LAST_AXIS else self.axis
+        axis = (len(x.shape) - 1) if self.axis == LAST_AXIS or self.axis is None else self.axis
         n = x.shape[axis]
         transpose_index = [axis, *[i for i in range(len(x.shape)) if i != axis]]
         x_reshape = np.reshape(np.transpose(x, transpose_index), [n, -1])

model_compression_toolkit/core/keras/default_framework_info.py

@@ -143,7 +143,7 @@ def get_out_channel_axis(cls, node_type: Any):
             Node's output channel axis.
 
         """
-        return cls.out_channel_axis_mapping.get(node_type, -1)
+        return cls.out_channel_axis_mapping.get(node_type)
 
 
 def set_keras_info(func):

model_compression_toolkit/core/pytorch/default_framework_info.py

@@ -101,7 +101,7 @@ def get_out_channel_axis(cls, node_type: Any):
             Node's output channel axis.
 
         """
-        return cls.out_channel_axis_mapping.get(node_type, 1)
+        return cls.out_channel_axis_mapping.get(node_type)
 
 
 def set_pytorch_info(func):

model_compression_toolkit/gptq/keras/quantization_facade.py

@@ -43,7 +43,6 @@
     import tensorflow as tf
     from model_compression_toolkit.core.keras.default_framework_info import set_keras_info
     from model_compression_toolkit.gptq.keras.gptq_keras_implementation import GPTQKerasImplemantation
-    from model_compression_toolkit.core.keras.keras_model_validation import KerasModelValidation
     from tensorflow.keras.models import Model
     from model_compression_toolkit.gptq.keras.gptq_loss import GPTQMultipleTensorsLoss, sample_layer_attention_loss
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
@@ -235,8 +234,6 @@ def keras_gradient_post_training_quantization(in_model: Model, representative_da
         if core_config.debug_config.bypass:
             return in_model, None
 
-        KerasModelValidation(model=in_model).validate()
-
         if core_config.is_mixed_precision_enabled:
             if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
                 Logger.critical("Given quantization config for mixed-precision is not of type 'MixedPrecisionQuantizationConfig'. "

model_compression_toolkit/ptq/keras/quantization_facade.py

@@ -38,7 +38,6 @@
         AttachTpcToKeras
     from model_compression_toolkit.core.keras.default_framework_info import set_keras_info
     from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
-    from model_compression_toolkit.core.keras.keras_model_validation import KerasModelValidation
     from tensorflow.keras.models import Model
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from model_compression_toolkit.exporter.model_wrapper import get_exportable_keras_model
@@ -129,8 +128,6 @@ def keras_post_training_quantization(in_model: Model,
         if core_config.debug_config.bypass:
             return in_model, None
 
-        KerasModelValidation(model=in_model).validate()
-
         if core_config.is_mixed_precision_enabled:
             if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
                 Logger.critical("Given quantization config to mixed-precision facade is not of type "

model_compression_toolkit/qat/keras/quantization_facade.py

@@ -38,7 +38,6 @@
 
     from model_compression_toolkit.trainable_infrastructure import KerasTrainableQuantizationWrapper
     from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
-    from model_compression_toolkit.core.keras.keras_model_validation import KerasModelValidation
     from model_compression_toolkit.target_platform_capabilities.constants import DEFAULT_TP_MODEL
     from model_compression_toolkit.core.keras.default_framework_info import set_keras_info
 
@@ -175,8 +174,6 @@ def keras_quantization_aware_training_init_experimental(in_model: Model,
                        f"If you encounter an issue, please open an issue in our GitHub "
                        f"project https://github.com/sony/model_optimization")
 
-        KerasModelValidation(model=in_model).validate()
-
         if core_config.is_mixed_precision_enabled:
             if not isinstance(core_config.mixed_precision_config, MixedPrecisionQuantizationConfig):
                 Logger.critical("Given quantization config to mixed-precision facade is not of type "

tests/keras_tests/feature_networks_tests/feature_networks/second_moment_correction_test.py

@@ -33,7 +33,6 @@
 from model_compression_toolkit.core.keras.constants import EPSILON_VAL, GAMMA, BETA, MOVING_MEAN, MOVING_VARIANCE
 from model_compression_toolkit.core.keras.default_framework_info import KerasInfo
 from model_compression_toolkit.core.keras.keras_implementation import KerasImplementation
-from model_compression_toolkit.core.keras.keras_model_validation import KerasModelValidation
 from model_compression_toolkit.core.keras.statistics_correction.apply_second_moment_correction import \
     keras_apply_second_moment_correction
 from model_compression_toolkit.core.runner import core_runner
@@ -269,8 +268,6 @@ def prepare_graph(self,
                       target_platform_capabilities: TargetPlatformCapabilities = DEFAULT_KERAS_TPC) -> \
             Tuple[Graph, Graph]:
 
-        KerasModelValidation(model=in_model).validate()
-
         core_config = CoreConfig(quantization_config=quant_config,
                                  debug_config=DebugConfig(analyze_similarity=analyze_similarity,
                                                           network_editor=network_editor)