Fix hessian criteria statistics collector (#3072)

### Changes

Fix bug with passing subset_size

### Tests

Add a test for backend methods

Author: kshpv

nncf/experimental/common/tensor_statistics/collectors.py

@@ -470,6 +470,21 @@ def _reduce_out_of_place(self, x: List[Tensor]) -> List[Tensor]:
         return [fns.mean(x, reduction_axes, keepdims=self._keepdims)]
 
 
+class MeanVarianceReducer(TensorReducerBase):
+    def _reduce_out_of_place(self, x: List[TensorType]) -> List[TensorType]:
+        raise NotImplementedError()
+
+
+class MaxVarianceReducer(TensorReducerBase):
+    def _reduce_out_of_place(self, x: List[TensorType]) -> List[TensorType]:
+        raise NotImplementedError()
+
+
+class MeanAbsMaxReducer(TensorReducerBase):
+    def _reduce_out_of_place(self, x: List[TensorType]) -> List[TensorType]:
+        raise NotImplementedError()
+
+
 class QuantileReducerBase(TensorReducerBase):
     def __init__(
         self,

nncf/openvino/statistics/collectors.py

@@ -9,17 +9,19 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Optional
+from typing import Optional
 
-from nncf.common.tensor import TensorType
 from nncf.experimental.common.tensor_statistics.collectors import AbsMaxReducer
 from nncf.experimental.common.tensor_statistics.collectors import AbsQuantileReducer
 from nncf.experimental.common.tensor_statistics.collectors import BatchMeanReducer
 from nncf.experimental.common.tensor_statistics.collectors import InplaceInsertionFNType
 from nncf.experimental.common.tensor_statistics.collectors import MaxReducer
+from nncf.experimental.common.tensor_statistics.collectors import MaxVarianceReducer
+from nncf.experimental.common.tensor_statistics.collectors import MeanAbsMaxReducer
 from nncf.experimental.common.tensor_statistics.collectors import MeanAggregator
 from nncf.experimental.common.tensor_statistics.collectors import MeanPerChReducer
 from nncf.experimental.common.tensor_statistics.collectors import MeanReducer
+from nncf.experimental.common.tensor_statistics.collectors import MeanVarianceReducer
 from nncf.experimental.common.tensor_statistics.collectors import MinReducer
 from nncf.experimental.common.tensor_statistics.collectors import NoopAggregator
 from nncf.experimental.common.tensor_statistics.collectors import NoopReducer
@@ -28,7 +30,6 @@
 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.collectors import TensorReducerBase
 from nncf.experimental.common.tensor_statistics.statistics import MeanTensorStatistic
 from nncf.experimental.common.tensor_statistics.statistics import RawTensorStatistic
 from nncf.openvino.graph.node_utils import get_inplace_batch_mean_op
@@ -67,26 +68,17 @@ def get_inplace_fn(self):
         return get_inplace_mean_op(self._reduction_axes)
 
 
-class OVMeanVarianceReducer(TensorReducerBase):
-    def _reduce_out_of_place(self, x: List[TensorType]) -> List[TensorType]:
-        raise NotImplementedError()
-
+class OVMeanVarianceReducer(MeanVarianceReducer):
     def get_inplace_fn(self):
         return get_inplace_mean_var_op(self._reduction_axes)
 
 
-class OVMaxVarianceReducer(TensorReducerBase):
-    def _reduce_out_of_place(self, x: List[TensorType]) -> List[TensorType]:
-        raise NotImplementedError()
-
+class OVMaxVarianceReducer(MaxVarianceReducer):
     def get_inplace_fn(self):
         return get_inplace_max_var_op(self._reduction_axes)
 
 
-class OVMeanAbsMaxReducer(TensorReducerBase):
-    def _reduce_out_of_place(self, x: List[TensorType]) -> List[TensorType]:
-        raise NotImplementedError()
-
+class OVMeanAbsMaxReducer(MeanAbsMaxReducer):
     def get_inplace_fn(self):
         return get_inplace_mean_max_op(self._reduction_axes, True)
 

nncf/quantization/algorithms/weight_compression/mixed_precision.py

@@ -281,7 +281,7 @@ def get_statistic_points(
         return statistic_container
 
     @abstractmethod
-    def _get_statistic_collector():
+    def _get_statistic_collector(self):
         """
         Get statistic collector
         """
@@ -360,7 +360,7 @@ def _calc_weight_sensitivity(
         return fns.linalg.norm(decompressed_weight - weight, ord="fro").item()
 
     def _get_statistic_collector(self):
-        return self._backend_entity.hawq_statistic_collector()
+        return self._backend_entity.hawq_statistic_collector(self._subset_size)
 
 
 @MIXED_PRECISION_CRITERIA.register(SensitivityMetric.MEAN_ACTIVATION_VARIANCE)

tests/cross_fw/test_templates/test_weights_compression_backends.py

@@ -0,0 +1,77 @@
+# Copyright (c) 2024 Intel Corporation
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#      http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from abc import abstractmethod
+
+import pytest
+
+from nncf.experimental.common.tensor_statistics.collectors import HAWQAggregator
+from nncf.experimental.common.tensor_statistics.collectors import MaxVarianceReducer
+from nncf.experimental.common.tensor_statistics.collectors import MeanAbsMaxReducer
+from nncf.experimental.common.tensor_statistics.collectors import MeanAggregator
+from nncf.experimental.common.tensor_statistics.collectors import MeanVarianceReducer
+from nncf.experimental.common.tensor_statistics.collectors import NoopReducer
+from nncf.experimental.common.tensor_statistics.collectors import TensorCollector
+
+
+class TemplateTestMixedPrecisionAlgoBackend:
+    @abstractmethod
+    def get_hawq_with_backend(self, subset_size: int):
+        """Returns a HAWQ instance of the algorithm."""
+
+    @abstractmethod
+    def get_mean_variance_with_backend(self, subset_size: int):
+        """Returns a Mean Variance instance of the algorithm."""
+
+    @abstractmethod
+    def get_max_variance_with_backend(self, subset_size: int):
+        """Returns a Max Variance instance of the algorithm."""
+
+    @abstractmethod
+    def get_mean_max_with_backend(self, subset_size: int):
+        """Returns a Mean Max instance of the algorithm."""
+
+    def check_aggregator(self, collector: TensorCollector, expected_aggregator_type, subset_size: int):
+        assert len(collector.aggregators) == 1, "Collector should have exactly one aggregator."
+        _, aggregator = collector.aggregators.popitem()
+        assert isinstance(
+            aggregator, expected_aggregator_type
+        ), f"Expected aggregator of type {expected_aggregator_type.__name__}, got {type(aggregator).__name__}."
+        assert aggregator.num_samples == subset_size, "Aggregator num_samples does not match the provided subset size."
+
+    def check_reducer(self, collector: TensorCollector, expected_reducer_type):
+        assert len(collector.reducers) == 1
+        reducer = collector.reducers.pop()
+        assert isinstance(
+            reducer, expected_reducer_type
+        ), f"Expected reducer of type {expected_reducer_type.__name__}, got {type(reducer).__name__}."
+
+    @pytest.mark.parametrize("subset_size", [1, 10, None])
+    @pytest.mark.parametrize(
+        "algo_func, aggregator_type, reducer_type",
+        [
+            ("get_hawq_with_backend", HAWQAggregator, NoopReducer),
+            ("get_mean_variance_with_backend", MeanAggregator, MeanVarianceReducer),
+            ("get_max_variance_with_backend", MeanAggregator, MaxVarianceReducer),
+            ("get_mean_max_with_backend", MeanAggregator, MeanAbsMaxReducer),
+        ],
+    )
+    def test_statistic_collector(self, subset_size, algo_func, aggregator_type, reducer_type):
+        """Test function to validate statistic collectors."""
+        algo = getattr(self, algo_func)(subset_size)
+        collector = algo._get_statistic_collector()
+
+        # Verify the collector instance and properties
+        assert isinstance(collector, TensorCollector), "Collector is not an instance of TensorCollector."
+
+        # Validate the aggregator and reducer types
+        self.check_aggregator(collector, aggregator_type, subset_size)
+        self.check_reducer(collector, reducer_type)

tests/openvino/native/quantization/test_weights_compression.py

@@ -899,12 +899,54 @@ def test_compression_for_different_dtypes(activation_dtype, weight_dtype):
 )
 @pytest.mark.parametrize(
     ("compression_args", "multiplier_of_calls"),
-    (
-        (dict(mode=CompressWeightsMode.INT4_ASYM, ratio=1), 0),  # data-free, no reducers
-        (dict(mode=CompressWeightsMode.INT4_ASYM, ratio=0.5), 1),  # 1 reducer for mixed precision
-        (dict(mode=CompressWeightsMode.INT4_ASYM, ratio=1, awq=True), 2),  # mean & shape reducer for AWQ
-        (dict(mode=CompressWeightsMode.INT4_ASYM, ratio=0.5, awq=True), 3),  # 2 - for AWQ + 1 - for Mixed Precision
-    ),
+    [
+        ({"mode": CompressWeightsMode.INT4_ASYM, "ratio": 1}, 0),  # data-free, no reducers
+        ({"mode": CompressWeightsMode.INT4_ASYM, "ratio": 1, "awq": True}, 2),  # mean & shape reducer for AWQ
+        (
+            {"mode": CompressWeightsMode.INT4_ASYM, "ratio": 0.5, "awq": True},
+            3,
+        ),  # 2 - for AWQ + 1 - for Mixed Precision
+        (
+            {
+                "mode": CompressWeightsMode.INT4_ASYM,
+                "ratio": 0.5,
+                "sensitivity_metric": nncf.SensitivityMetric.HESSIAN_INPUT_ACTIVATION,
+            },
+            1,
+        ),  # 1 reducer for mixed precision
+        (
+            {
+                "mode": CompressWeightsMode.INT4_ASYM,
+                "ratio": 0.5,
+                "sensitivity_metric": nncf.SensitivityMetric.MEAN_ACTIVATION_VARIANCE,
+            },
+            1,
+        ),  # 1 reducer for mixed precision
+        (
+            {
+                "mode": CompressWeightsMode.INT4_ASYM,
+                "ratio": 0.5,
+                "sensitivity_metric": nncf.SensitivityMetric.MAX_ACTIVATION_VARIANCE,
+            },
+            1,
+        ),  # 1 reducer for mixed precision
+        (
+            {
+                "mode": CompressWeightsMode.INT4_ASYM,
+                "ratio": 0.5,
+                "sensitivity_metric": nncf.SensitivityMetric.MEAN_ACTIVATION_MAGNITUDE,
+            },
+            1,
+        ),  # 1 reducer for mixed precision
+        (
+            {
+                "mode": CompressWeightsMode.INT4_ASYM,
+                "ratio": 0.5,
+                "sensitivity_metric": nncf.SensitivityMetric.WEIGHT_QUANTIZATION_ERROR,
+            },
+            0,
+        ),  # 0 - data-free method
+    ],
 )
 def test_number_of_reduced_statistics_for_subset_size(
     mocker, dataset_size, subset_size, ref_size, compression_args, multiplier_of_calls

tests/openvino/native/quantization/test_weights_compression_backends.py

@@ -0,0 +1,39 @@
+# Copyright (c) 2024 Intel Corporation
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#      http://www.apache.org/licenses/LICENSE-2.0
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from nncf.quantization.algorithms.weight_compression.mixed_precision import HAWQCriterion
+from nncf.quantization.algorithms.weight_compression.mixed_precision import MaxVarianceCriterion
+from nncf.quantization.algorithms.weight_compression.mixed_precision import MeanMaxCriterion
+from nncf.quantization.algorithms.weight_compression.mixed_precision import MeanVarianceCriterion
+from nncf.quantization.algorithms.weight_compression.openvino_backend import OVMixedPrecisionAlgoBackend
+from tests.cross_fw.test_templates.test_weights_compression_backends import TemplateTestMixedPrecisionAlgoBackend
+from tests.openvino.native.models import IdentityMatmul
+
+
+class TestOVMixedPrecisionAlgoBackend(TemplateTestMixedPrecisionAlgoBackend):
+    def get_hawq_with_backend(self, subset_size):
+        hawq = HAWQCriterion(None, None, subset_size=subset_size)
+        hawq._backend_entity = OVMixedPrecisionAlgoBackend(IdentityMatmul().ov_model)
+        return hawq
+
+    def get_mean_variance_with_backend(self, subset_size: int):
+        mean_variance = MeanVarianceCriterion(None, None, subset_size=subset_size)
+        mean_variance._backend_entity = OVMixedPrecisionAlgoBackend(IdentityMatmul().ov_model)
+        return mean_variance
+
+    def get_max_variance_with_backend(self, subset_size: int):
+        max_variance = MaxVarianceCriterion(None, None, subset_size=subset_size)
+        max_variance._backend_entity = OVMixedPrecisionAlgoBackend(IdentityMatmul().ov_model)
+        return max_variance
+
+    def get_mean_max_with_backend(self, subset_size: int):
+        mean_max_variance = MeanMaxCriterion(None, None, subset_size=subset_size)
+        mean_max_variance._backend_entity = OVMixedPrecisionAlgoBackend(IdentityMatmul().ov_model)
+        return mean_max_variance