[NNCF]: Optimized memory footprint by removing redundant collected statistics

nncf/common/tensor_statistics/statistic_point.py

@@ -71,6 +71,32 @@ def add_statistic_point(self, statistic_point: StatisticPoint) -> None:
 
             self.data[target_node_name].append(statistic_point)
 
+    def remove_statistic_points(self, algorithm_key: str) -> None:
+        """
+        Method to remove statistic point associated with a given algorithm
+        from statistic point container.
+
+        :param algorithm_key: Algorithm key.
+        """
+        reduced_data = {}
+        for target_node_name, statistic_points in self.data.items():
+            # Reassign every target node name IF it doesn't contain the given algorithm
+            reduced_statistics = []
+            for statistic_point in statistic_points:
+                statistic_point.algorithm_to_tensor_collectors = {
+                    algorithm: value
+                    for algorithm, value in statistic_point.algorithm_to_tensor_collectors.items()
+                    if algorithm != algorithm_key
+                }
+
+                if statistic_point.algorithm_to_tensor_collectors:
+                    reduced_statistics.append(statistic_point)
+
+            if reduced_statistics:
+                reduced_data[target_node_name] = reduced_statistics
+
+        self.data = reduced_data
+
     def iter_through_statistic_points_in_target_node(
         self, target_node_name: str, filter_fn: Callable[[StatisticPoint], bool]
     ) -> StatisticPoint:

nncf/quantization/algorithms/algorithm.py

@@ -26,6 +26,17 @@ class Algorithm(ABC):
     Base class for all Post-Training algorithms.
     """
 
+    def __init__(self):
+        self._algorithm_key = ""
+
+    def algorithm_key(self) -> str:
+        """
+        Returns algorithm key.
+
+        :return: Algorithm key as string.
+        """
+        return self._algorithm_key
+
     @property
     @abstractmethod
     def available_backends(self) -> List[BackendType]:

nncf/quantization/algorithms/pipeline.py

@@ -113,10 +113,14 @@ def run_step(
 
         pipeline_steps = self._remove_unsupported_algorithms(get_backend(model))
         pipeline_step = pipeline_steps[step_index]
+
         for algorithm in pipeline_step[:-1]:
             current_model = algorithm.apply(current_model, current_graph, step_statistics)
             current_graph = NNCFGraphFactory.create(current_model)
+            step_statistics.remove_statistic_points(algorithm.algorithm_key())
+
         current_model = pipeline_step[-1].apply(current_model, current_graph, step_statistics)
+        step_statistics.remove_statistic_points(pipeline_step[-1].algorithm_key())
 
         return current_model
 
@@ -162,7 +166,6 @@ def run_from_step(
 
             # Run current pipeline step
             step_model = self.run_step(step_index, step_statistics, step_model, step_graph)
-
             step_graph = None  # We should rebuild the graph for the next pipeline step
 
         return step_model

tests/common/test_statistic_points.py

@@ -0,0 +1,49 @@
+# 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.
+import pytest
+
+from nncf.common.graph.transformations.commands import TargetPoint
+from nncf.common.graph.transformations.commands import TargetType
+from nncf.common.tensor_statistics.collectors import TensorStatisticCollectorBase
+from nncf.common.tensor_statistics.statistic_point import StatisticPoint
+from nncf.common.tensor_statistics.statistic_point import StatisticPointsContainer
+
+
+class TestStatisticPointsContainer:
+    @pytest.fixture
+    def container_with_statistics(self):
+        container = StatisticPointsContainer()
+        target_type = TargetType.LAYER
+        target_point = TargetPoint(target_type)
+        target_point.target_node_name = "Node"
+        # tensor_collector = TensorStatisticCollectorBase()
+        statistic_point = StatisticPoint(target_point, TensorStatisticCollectorBase, "minmax")
+        statistic_point = StatisticPoint(target_point, TensorStatisticCollectorBase, "post_training")
+        container.add_statistic_point(statistic_point)
+        return container
+
+    def test_remove_statistic_points(self, container_with_statistics):
+        container_with_statistics.remove_statistic_points("minmax")
+
+        # Verify that the statistics for the specified algorithm are removed
+        for target_node_name, statistic_points in container_with_statistics.data.items():
+            for statistic_point in statistic_points:
+                assert (
+                    "minmax" not in statistic_point.algorithm_to_tensor_collectors
+                    and "post_training" in statistic_point.algorithm_to_tensor_collectors
+                )
+
+    def test_remove_statistic_points_empty_container(self):
+        empty_container = StatisticPointsContainer()
+        empty_container.remove_statistic_points("algorithm")
+
+        # Verify that removing statistics from an empty container has no effect
+        assert not empty_container.data