Extract node's channel axis info for statistics collectors from adjacent nodes.

Author: elad-c

model_compression_toolkit/core/common/collectors/mean_collector.py

@@ -87,10 +87,10 @@ 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
-        if axis is None:
+        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

model_compression_toolkit/core/common/collectors/min_max_per_channel_collector.py

@@ -130,10 +130,10 @@ 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
-        if axis is None:
+        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])

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(