Remove default output axis values from fw_info. (#1489)

* Remove default output axis values from fw_info.
* Handle missing defaults:
    1. Collectors for nodes that don't have a known output channel axis: Extract node's channel axis info for statistics collectors from adjacent nodes, or collect for the whole tensor.
    2. Pruning retain defaults that are now defined in the PruningImplementation.

Author: elad-c

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,10 +87,13 @@ 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
-        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
+        if self.axis is None:
+            mu = np.mean(np.reshape(x, [1, -1]), axis=-1)  # mean per channel for a batch
+        else:
+            axis = (len(x.shape) - 1) if self.axis == LAST_AXIS 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
         self.current_sum += mu # sum of all batches
         self.current_mean = self.current_sum / self.i # mean of all batches
 

model_compression_toolkit/core/common/collectors/min_max_per_channel_collector.py

@@ -130,10 +130,13 @@ 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
-        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])
+        if self.axis is None:
+            x_reshape = np.reshape(x, [1, -1])
+        else:
+            axis = (len(x.shape) - 1) if self.axis == LAST_AXIS 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])
         if self.state is None:
             x_max = np.max(x_reshape, axis=-1)
             x_min = np.min(x_reshape, axis=-1)

model_compression_toolkit/core/common/model_collector.py

@@ -157,6 +157,17 @@ def __init__(self, graph: Graph,
         for n in graph.get_topo_sorted_nodes():
             quant_node_in_fln = n.is_fln_quantization() and graph.fusing_info.is_quantized_node_in_fln(n)
             sc = create_stats_collector_for_node(n, quant_node_in_fln=quant_node_in_fln)  # Get static collector for the node
+            if isinstance(sc, common.StatsCollector) and (sc.mc.axis is None or sc.mpcc.axis is None):
+                # Missing output channel axis info, so try to extract it from previous and next nodes output channel axis.
+                possible_output_channel_axis_set = {nn.out_channel_axis for nn in graph.get_next_nodes(n) + graph.get_prev_nodes(n)}
+                # Filter out None values.
+                possible_output_channel_axis_list = list(filter(lambda x: x is not None, possible_output_channel_axis_set))
+                if len(possible_output_channel_axis_list) > 0:
+                    if len(possible_output_channel_axis_list) > 1:
+                        Logger.warning(f'Ambiguous input channel data from next nodes for {n.name}.')
+                    sc.mc.axis = possible_output_channel_axis_list[0]
+                    sc.mpcc.axis = possible_output_channel_axis_list[0]
+
             # If we use bias correction, and the node has kernel weights to quantize, we need to make sure
             # its previous nodes' tensors are consistent with this node.
             if qc.weights_bias_correction and n.kernel_attr is not None and n.is_weights_quantization_enabled(

model_compression_toolkit/core/common/model_validation.py

@@ -303,7 +303,7 @@ def get_pruned_node_num_params(self,
             num_oc = np.sum(output_mask)
         else:
             # Get the node channel axis from framework info
-            channel_axis = node.out_channel_axis
+            channel_axis = self.fw_impl.default_output_channel_axis if node.out_channel_axis is None else node.out_channel_axis
             if channel_axis is None:
                 Logger.critical(f"The channel axis is undefined. Please ensure the channel axis is explicitly defined for node {node.type} in the framework info.")
 

model_compression_toolkit/core/common/pruning/memory_calculator.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/keras/default_framework_info.py

@@ -28,6 +28,10 @@
 from model_compression_toolkit.logger import Logger
 
 
+# default output channel axis to use when it's not defined in node's fw_info.
+_default_output_channel_axis = -1
+
+
 class PruningKerasImplementation(KerasImplementation, PruningFrameworkImplementation):
     """
     Implementation of the PruningFramework for the Keras framework. This class provides
@@ -172,6 +176,10 @@ def attrs_oi_channels_info_for_pruning(self,
 
         return attributes_with_axis
 
+    @property
+    def default_output_channel_axis(self):
+        return _default_output_channel_axis
+
 
 def _is_keras_node_pruning_section_edge(node: BaseNode) -> bool:
     """

model_compression_toolkit/core/keras/keras_model_validation.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):