POC histogramm aggregator plugged in to MinMax

Author: daniil-lyakhov

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

@@ -23,6 +23,7 @@
 from nncf.common.tensor_statistics.collectors import ReductionAxes
 from nncf.experimental.common.tensor_statistics.statistical_functions import mean_per_channel
 from nncf.experimental.common.tensor_statistics.statistics import MedianMADTensorStatistic
+from nncf.experimental.common.tensor_statistics.statistics import MinMaxTensorStatistic
 from nncf.experimental.common.tensor_statistics.statistics import TensorStatistic
 from nncf.quantization.advanced_parameters import AggregatorType
 from nncf.quantization.range_estimator import StatisticsType
@@ -1154,8 +1155,9 @@ def _register_reduced_input_impl(self, x: Tensor):
         self.min_val = new_min
         self.max_val = new_max
 
-    def _aggregate_impl(self) -> Any:
-        return self._non_linear_param_search()
+    def _aggregate_impl(self) -> dict[str, float]:
+        min_, max_ = self._non_linear_param_search()
+        return {MinMaxTensorStatistic.MIN_STAT: min_, MinMaxTensorStatistic.MAX_STAT: max_}
 
 
 REDUCERS_MAP = {
@@ -1165,6 +1167,7 @@ def _aggregate_impl(self) -> Any:
     StatisticsType.MEAN: MeanReducer,
     StatisticsType.QUANTILE: QuantileReducer,
     StatisticsType.ABS_QUANTILE: AbsQuantileReducer,
+    StatisticsType.RAW: RawReducer,
 }
 
 AGGREGATORS_MAP = {
@@ -1174,4 +1177,5 @@ def _aggregate_impl(self) -> Any:
     AggregatorType.MEAN_NO_OUTLIERS: MeanNoOutliersAggregator,
     AggregatorType.MEDIAN: MedianAggregator,
     AggregatorType.MEDIAN_NO_OUTLIERS: MedianNoOutliersAggregator,
+    AggregatorType.HISTOGRAM: HistogramAggregator,
 }

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

@@ -77,6 +77,7 @@ def from_config(cls, config: dict[str, Any]) -> TensorStatistic:
 class MinMaxTensorStatistic(TensorStatistic):
     MIN_STAT: ClassVar[str] = "min_values"
     MAX_STAT: ClassVar[str] = "max_values"
+    MIN_MAX_STAT: ClassVar[str] = "min_max_values"
 
     min_values: Tensor
     max_values: Tensor
@@ -86,6 +87,12 @@ def __eq__(self, other: TensorStatistic):
             return fns.allclose(self.min_values, other.min_values) and fns.allclose(self.max_values, other.max_values)
         return False
 
+    @classmethod
+    def from_config(cls, config: dict[str, Any]) -> TensorStatistic:
+        if cls.MIN_MAX_STAT in config:
+            return cls(**config[cls.MIN_MAX_STAT])
+        return cls(min_values=config[cls.MIN_STAT], max_values=config[cls.MAX_STAT])
+
 
 @dataclass
 class AbsMaxTensorStatistic(TensorStatistic):

src/nncf/openvino/statistics/collectors.py

@@ -151,4 +151,5 @@ def get_raw_stat_collector(num_samples: Optional[int] = None) -> TensorCollector
     StatisticsType.MEAN: OVMeanReducer,
     StatisticsType.QUANTILE: OVQuantileReducer,
     StatisticsType.ABS_QUANTILE: OVAbsQuantileReducer,
+    StatisticsType.RAW: RawReducer,
 }

src/nncf/quantization/advanced_parameters.py

@@ -27,6 +27,7 @@
 from nncf.parameters import StrEnum
 from nncf.quantization.range_estimator import AggregatorType
 from nncf.quantization.range_estimator import RangeEstimatorParameters
+from nncf.quantization.range_estimator import StatisticsCollectorParameters
 from nncf.quantization.range_estimator import StatisticsType
 
 TTensor = Any
@@ -272,8 +273,12 @@ class AdvancedQuantizationParameters:
     quantizer_propagation_rule: QuantizerPropagationRule = QuantizerPropagationRule.MERGE_ALL_IN_ONE
 
     # Range estimator parameters
-    activations_range_estimator_params: RangeEstimatorParameters = field(default_factory=RangeEstimatorParameters)
-    weights_range_estimator_params: RangeEstimatorParameters = field(default_factory=RangeEstimatorParameters)
+    activations_range_estimator_params: Union[RangeEstimatorParameters, StatisticsCollectorParameters] = field(
+        default_factory=RangeEstimatorParameters
+    )
+    weights_range_estimator_params: Union[RangeEstimatorParameters, StatisticsCollectorParameters] = field(
+        default_factory=RangeEstimatorParameters
+    )
 
     # Advanced BiasCorrection algorithm parameters
     bias_correction_params: AdvancedBiasCorrectionParameters = field(default_factory=AdvancedBiasCorrectionParameters)

src/nncf/quantization/algorithms/min_max/algorithm.py

@@ -67,6 +67,7 @@
 from nncf.quantization.range_estimator import AggregatorType
 from nncf.quantization.range_estimator import RangeEstimatorParameters
 from nncf.quantization.range_estimator import RangeEstimatorParametersSet
+from nncf.quantization.range_estimator import StatisticsCollectorParameters
 from nncf.quantization.range_estimator import StatisticsType
 from nncf.scopes import IgnoredScope
 from nncf.scopes import get_ignored_node_names_from_ignored_scope
@@ -433,6 +434,15 @@ def _get_range_estimator_parameters(
         if user_params is None:
             return deepcopy(params)
 
+        if isinstance(user_params, StatisticsCollectorParameters):
+            if quantizer_config.per_channel:
+                msg = (
+                    f"Could not create signle aggregator with parameters {user_params}",
+                    " Per channel statistic collection is not supported for the single aggregator case yet.",
+                )
+                raise nncf.InternalError(msg)
+            return deepcopy(user_params)
+
         min_changes = changes_asdict(user_params.min)
         min_statistic_collector = dataclasses.replace(params.min, **min_changes)
 
@@ -495,7 +505,7 @@ def _get_stat_collector(
 
     def _get_statistic_collector(
         self,
-        range_estimator_params: RangeEstimatorParameters,
+        range_estimator_params: Union[RangeEstimatorParameters, StatisticsCollectorParameters],
         use_abs_max: bool,
         reduction_axes: Optional[tuple[int, ...]],
         aggregation_axes: Optional[tuple[int, ...]],
@@ -513,10 +523,24 @@ def _get_statistic_collector(
         :param num_samples: Maximum number of samples to collect.
         :return: TensorCollector for the statistics calculation.
         """
+        collector = TensorCollector(MinMaxTensorStatistic)
+        if isinstance(range_estimator_params, StatisticsCollectorParameters):
+            if range_estimator_params.statistics_type is not StatisticsType.RAW:
+                msg = "Only RAW statistic type is suppored for single aggregator case."
+                raise nncf.InternalError(msg)
+
+            if range_estimator_params.aggregator_type is not AggregatorType.HISTOGRAM:
+                msg = "Only HISTOGRAM aggregator type is suppored for single aggregator case."
+                raise nncf.InternalError(msg)
+
+            reducer = self._backend_entity.reducer_map[StatisticsType.RAW]()
+            aggregator = AGGREGATORS_MAP[AggregatorType.HISTOGRAM]()
+            collector.register_statistic_branch(MinMaxTensorStatistic.MIN_MAX_STAT, reducer, aggregator)
+            return collector
+
         if not self._backend_entity.supports_inplace_statistics:
             inplace = False
 
-        collector = TensorCollector(MinMaxTensorStatistic)
         for params, container_key in zip(
             [range_estimator_params.min, range_estimator_params.max],
             [MinMaxTensorStatistic.MIN_STAT, MinMaxTensorStatistic.MAX_STAT],

src/nncf/quantization/range_estimator.py

@@ -37,6 +37,7 @@ class StatisticsType(Enum):
     QUANTILE = "quantile"
     ABS_QUANTILE = "abs_quantile"
     MEAN = "mean"
+    RAW = "raw"
 
 
 @api()
@@ -59,6 +60,7 @@ class AggregatorType(Enum):
     MEDIAN = "median"
     MEAN_NO_OUTLIERS = "mean_no_outliers"
     MEDIAN_NO_OUTLIERS = "median_no_outliers"
+    HISTOGRAM = "histogram"
 
 
 @api()

src/nncf/version.py

@@ -12,6 +12,6 @@
 __version__ = "2.18.0"
 
 
-BKC_TORCH_SPEC = "==2.7.*"
+BKC_TORCH_SPEC = "==2.8.*"
 BKC_TF_SPEC = "==2.15.*"
 STRICT_TF_SPEC = ">=2.9.3,<2.16.0"

tests/post_training/model_scope.py

@@ -23,6 +23,9 @@
 from nncf.quantization.advanced_parameters import AdvancedQuantizationParameters
 from nncf.quantization.advanced_parameters import AdvancedScaleEstimationParameters
 from nncf.quantization.advanced_parameters import AdvancedSmoothQuantParameters
+from nncf.quantization.range_estimator import AggregatorType
+from nncf.quantization.range_estimator import StatisticsCollectorParameters
+from nncf.quantization.range_estimator import StatisticsType
 from tests.post_training.pipelines.base import ALL_PTQ_BACKENDS
 from tests.post_training.pipelines.base import FX_BACKENDS
 from tests.post_training.pipelines.base import NNCF_PTQ_BACKENDS
@@ -89,6 +92,11 @@
         "pipeline_cls": ImageClassificationTorchvision,
         "compression_params": {
             "subset_size": 2,
+            "advanced_parameters": AdvancedQuantizationParameters(
+                activations_range_estimator_params=StatisticsCollectorParameters(
+                    statistics_type=StatisticsType.RAW, aggregator_type=AggregatorType.HISTOGRAM
+                )
+            ),
         },
         "backends": FX_BACKENDS
         + [
@@ -117,7 +125,12 @@
         "pipeline_cls": ImageClassificationTorchvision,
         "compression_params": {
             "model_type": ModelType.TRANSFORMER,
-            "advanced_parameters": AdvancedQuantizationParameters(smooth_quant_alpha=0.15),
+            "advanced_parameters": AdvancedQuantizationParameters(
+                smooth_quant_alpha=0.15,
+                activations_range_estimator_params=StatisticsCollectorParameters(
+                    statistics_type=StatisticsType.RAW, aggregator_type=AggregatorType.HISTOGRAM
+                ),
+            ),
         },
         "backends": FX_BACKENDS + [BackendType.OV],
         "batch_size": 1,
@@ -128,7 +141,12 @@
         "pipeline_cls": ImageClassificationTorchvision,
         "compression_params": {
             "model_type": ModelType.TRANSFORMER,
-            "advanced_parameters": AdvancedQuantizationParameters(smooth_quant_alpha=0.5),
+            "advanced_parameters": AdvancedQuantizationParameters(
+                smooth_quant_alpha=0.5,
+                activations_range_estimator_params=StatisticsCollectorParameters(
+                    statistics_type=StatisticsType.RAW, aggregator_type=AggregatorType.HISTOGRAM
+                ),
+            ),
         },
         "backends": FX_BACKENDS + [BackendType.OV],
         "batch_size": 1,