Initial commit

Author: Nikita Savelyev

src/nncf/experimental/common/tensor_statistics/statistics.py

@@ -12,12 +12,14 @@
 from __future__ import annotations
 
 from collections import Counter
+from collections import deque
 from dataclasses import dataclass
 from dataclasses import fields
 from typing import Any, ClassVar
 
 import nncf
 from nncf.tensor import Tensor
+from nncf.tensor import TensorDataType
 from nncf.tensor import functions as fns
 
 
@@ -107,23 +109,27 @@ class MeanTensorStatistic(TensorStatistic):
     mean_values: Tensor
     shape: tuple[int, ...]
 
+    def __post_init__(self):
+        if isinstance(self.shape, (deque, list)):
+            # If NoopAggregator was used, shape can be a sequence containing a single tensor
+            self.shape = tuple(self.shape[0].data.tolist())
+
     def __eq__(self, other: TensorStatistic):
         if isinstance(other, MeanTensorStatistic):
             return self.shape == other.shape and fns.allclose(self.mean_values, other.mean_values)
         return False
 
     def _get_serialized_data(self) -> dict[str, Tensor]:
         backend = self.mean_values.backend
-        dtype = self.mean_values.dtype
         device = self.mean_values.device
         return {
             self.MEAN_STAT: self.mean_values,
-            self.SHAPE_STAT: fns.tensor(self.shape, backend=backend, dtype=dtype, device=device),
+            self.SHAPE_STAT: fns.tensor(self.shape, backend=backend, dtype=TensorDataType.int32, device=device),
         }
 
     def load_data(self, loaded_data: dict[str, Tensor]) -> None:
         self.mean_values = loaded_data[self.MEAN_STAT]
-        self.shape_values = tuple(loaded_data[self.SHAPE_STAT].tolist())
+        self.shape = tuple(loaded_data[self.SHAPE_STAT].data.tolist())
 
 
 @dataclass
@@ -261,14 +267,13 @@ def __eq__(self, other: Any) -> bool:
 
     def _get_serialized_data(self) -> dict[str, Tensor]:
         backend = self.mean_values[0].backend
-        dtype = self.mean_values[0].dtype
         device = self.mean_values[0].device
         return {
             self.MEAN_STAT: fns.stack(self.mean_values),
             self.SHAPE_STAT: fns.tensor(
-                [[dim.data for dim in shape] for shape in self.shape_values],
+                self.shape_values,
                 backend=backend,
-                dtype=dtype,
+                dtype=TensorDataType.int32,
                 device=device,
             ),
         }
@@ -283,5 +288,5 @@ def from_config(cls, config: dict[str, Any]) -> TensorStatistic:
         if cls.MEAN_STAT in config and config[cls.MEAN_STAT] is not None:
             mean_values = [fns.squeeze(it) for it in config[cls.MEAN_STAT]]
         if cls.SHAPE_STAT in config and config[cls.SHAPE_STAT] is not None:
-            shape_values = [tuple(it) for it in config[cls.SHAPE_STAT]]
+            shape_values = [tuple(it.data.tolist()) for it in config[cls.SHAPE_STAT]]
         return cls(mean_values=mean_values, shape_values=shape_values)

src/nncf/onnx/statistics/collectors.py

@@ -16,7 +16,7 @@
 from nncf.experimental.common.tensor_statistics.collectors import MeanPerChReducer
 from nncf.experimental.common.tensor_statistics.collectors import NoopAggregator
 from nncf.experimental.common.tensor_statistics.collectors import RawReducer
-from nncf.experimental.common.tensor_statistics.collectors import ShapeAggregator
+from nncf.experimental.common.tensor_statistics.collectors import ShapeReducer
 from nncf.experimental.common.tensor_statistics.collectors import TensorCollector
 from nncf.experimental.common.tensor_statistics.statistics import MeanTensorStatistic
 from nncf.experimental.common.tensor_statistics.statistics import RawTensorStatistic
@@ -40,19 +40,19 @@ def get_mean_statistic_collector(
         reducer = BatchMeanReducer(inplace)
     else:
         reducer = MeanPerChReducer(channel_axis=channel_axis, inplace=inplace)
-    raw_reducer = RawReducer()
+    shape_reducer = ShapeReducer(inplace=inplace)
 
     kwargs = {
         "num_samples": num_samples,
         "window_size": window_size,
     }
 
     aggregate_mean = MeanAggregator(**kwargs)
-    aggregate_shape = ShapeAggregator()
+    aggregate_noop = NoopAggregator(num_samples=1)
 
     collector = TensorCollector(MeanTensorStatistic)
     collector.register_statistic_branch(MeanTensorStatistic.MEAN_STAT, reducer, aggregate_mean)
-    collector.register_statistic_branch(MeanTensorStatistic.SHAPE_STAT, raw_reducer, aggregate_shape)
+    collector.register_statistic_branch(MeanTensorStatistic.SHAPE_STAT, shape_reducer, aggregate_noop)
     return collector
 
 

src/nncf/openvino/statistics/collectors.py

@@ -26,7 +26,6 @@
 from nncf.experimental.common.tensor_statistics.collectors import NoopAggregator
 from nncf.experimental.common.tensor_statistics.collectors import QuantileReducer
 from nncf.experimental.common.tensor_statistics.collectors import RawReducer
-from nncf.experimental.common.tensor_statistics.collectors import ShapeAggregator
 from nncf.experimental.common.tensor_statistics.collectors import ShapeReducer
 from nncf.experimental.common.tensor_statistics.collectors import TensorCollector
 from nncf.experimental.common.tensor_statistics.statistics import MeanTensorStatistic
@@ -120,18 +119,18 @@ def get_mean_statistic_collector(
         reducer = OVBatchMeanReducer(inplace)
     else:
         reducer = OVMeanPerChanelReducer(channel_axis=channel_axis, inplace=inplace)
-    raw_reducer = RawReducer()
+    shape_reducer = OVShapeReducer(inplace=inplace)
 
     kwargs = {
         "num_samples": num_samples,
         "window_size": window_size,
     }
     aggregate_mean = MeanAggregator(**kwargs)
-    aggregate_shape = ShapeAggregator()
+    aggregate_noop = NoopAggregator(num_samples=1)
 
     collector = TensorCollector(MeanTensorStatistic)
     collector.register_statistic_branch(MeanTensorStatistic.MEAN_STAT, reducer, aggregate_mean)
-    collector.register_statistic_branch(MeanTensorStatistic.SHAPE_STAT, raw_reducer, aggregate_shape)
+    collector.register_statistic_branch(MeanTensorStatistic.SHAPE_STAT, shape_reducer, aggregate_noop)
     return collector
 
 

src/nncf/torch/tensor_statistics/collectors.py

@@ -28,7 +28,7 @@
 from nncf.experimental.common.tensor_statistics.collectors import PercentileAggregator
 from nncf.experimental.common.tensor_statistics.collectors import QuantileReducer
 from nncf.experimental.common.tensor_statistics.collectors import RawReducer
-from nncf.experimental.common.tensor_statistics.collectors import ShapeAggregator
+from nncf.experimental.common.tensor_statistics.collectors import ShapeReducer
 from nncf.experimental.common.tensor_statistics.collectors import TensorCollector
 from nncf.experimental.common.tensor_statistics.statistics import MeanTensorStatistic
 from nncf.experimental.common.tensor_statistics.statistics import MedianMADTensorStatistic
@@ -306,18 +306,18 @@ def get_mean_statistic_collector(
         reducer = BatchMeanReducer()
     else:
         reducer = MeanPerChReducer(channel_axis=channel_axis)
-    raw_reducer = RawReducer()
+    shape_reducer = ShapeReducer()
 
     kwargs = {
         "num_samples": num_samples,
         "window_size": window_size,
     }
     aggregate_mean = MeanAggregator(**kwargs)
-    aggregate_shape = ShapeAggregator()
+    aggregate_noop = NoopAggregator(num_samples=1)
 
     collector = TensorCollector(MeanTensorStatistic)
     collector.register_statistic_branch(MeanTensorStatistic.MEAN_STAT, reducer, aggregate_mean)
-    collector.register_statistic_branch(MeanTensorStatistic.SHAPE_STAT, raw_reducer, aggregate_shape)
+    collector.register_statistic_branch(MeanTensorStatistic.SHAPE_STAT, shape_reducer, aggregate_noop)
     return collector