Skip to content

Tensor stats migration #3440

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 5 commits into
base: develop
Choose a base branch
from
Draft

Tensor stats migration #3440

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
c66a223
break circular dependency between old and new tensor statistics imple…
darshil929 Apr 16, 2025
7703d2a
migrate tensor statistics classes to use experimental implementation
darshil929 Apr 16, 2025
99190ab
migrate PyTorch tensor statistics collectors to experimental implemen…
darshil929 Apr 16, 2025
fb861c8
migrate OpenVINO tensor statistics collectors to experimental impleme…
darshil929 Apr 16, 2025
b55711d
update tensor statistics collectors in algorithm backends to use expe…
darshil929 Apr 17, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 6 additions & 5 deletions nncf/experimental/common/tensor_statistics/collectors.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
import nncf
import nncf.tensor.functions as fns
from nncf.common.tensor import TensorType
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 TensorStatistic
Expand All @@ -28,6 +27,8 @@
from nncf.tensor import Tensor

InplaceInsertionFNType = TypeVar("InplaceInsertionFNType")
# AggregationAxes serves the same purpose as ReductionAxes in the old implementation
# but is defined here to break circular dependencies
AggregationAxes = Tuple[int, ...]
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will this help with cyclical dependencies?

Suggested change
AggregationAxes = Tuple[int, ...]
AggregationAxes = Tuple[int, ...]
ReductionAxes = Tuple[int, ...]



Expand All @@ -37,7 +38,7 @@ class TensorReducerBase(ABC):
the specified rule. Could handle tensors inplace or out of place.
"""

def __init__(self, reduction_axes: Optional[ReductionAxes] = None, inplace: bool = False):
def __init__(self, reduction_axes: Optional[AggregationAxes] = None, inplace: bool = False):
"""
:param reduction_axes: Reduction axes for reduction calculation. Equal to list(range(len(input.shape)))
if empty.
Expand Down Expand Up @@ -98,7 +99,7 @@ def __eq__(self, __o: object) -> bool:
def __hash__(self) -> int:
return hash((self.__class__.__name__, self.inplace, self._reduction_axes))

def _get_reduction_axes(self, tensor: Tensor) -> ReductionAxes:
def _get_reduction_axes(self, tensor: Tensor) -> AggregationAxes:
if self._reduction_axes is not None:
return self._reduction_axes
return tuple(range(len(tensor.shape)))
Expand Down Expand Up @@ -489,7 +490,7 @@ def _reduce_out_of_place(self, x: List[Tensor]) -> List[Tensor]:
class QuantileReducerBase(TensorReducerBase):
def __init__(
self,
reduction_axes: Optional[ReductionAxes] = None,
reduction_axes: Optional[AggregationAxes] = None,
quantile: Optional[Union[float, Tuple[float]]] = None,
inplace: bool = False,
):
Expand All @@ -513,7 +514,7 @@ def _reduce_out_of_place(self, x: List[Tensor]) -> List[Tensor]:
class AbsQuantileReducer(QuantileReducerBase):
def __init__(
self,
reduction_axes: Optional[ReductionAxes] = None,
reduction_axes: Optional[AggregationAxes] = None,
quantile: Optional[Union[float, List[float]]] = None,
inplace: bool = False,
):
Expand Down
19 changes: 11 additions & 8 deletions nncf/openvino/graph/node_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -23,7 +23,10 @@
from nncf.common.graph.layer_attributes import GenericWeightedLayerAttributes
from nncf.common.graph.layer_attributes import LinearLayerAttributes
from nncf.common.graph.layer_attributes import WeightedLayerAttributes
from nncf.common.tensor_statistics.collectors import ReductionAxes

# Using experimental tensor statistics implementation as part of the migration
# from old tensor statistics to experimental tensor statistics
from nncf.experimental.common.tensor_statistics.collectors import AggregationAxes
from nncf.openvino.graph.layout import OVLayoutElem
from nncf.openvino.graph.layout import get_conv_weights_layout
from nncf.openvino.graph.layout import get_conv_weights_layout_from_node
Expand Down Expand Up @@ -228,7 +231,7 @@ def get_ov_model_reduce_node_name(output_name: str, reduce_node_name: str, port_

def get_inplace_reduce_op(
op: Type[ov.Node],
reduction_axes: Optional[ReductionAxes],
reduction_axes: Optional[AggregationAxes],
use_abs: bool,
keep_dims: bool = True,
) -> InplaceInsertionFnType:
Expand Down Expand Up @@ -265,7 +268,7 @@ def get_reduce_op(node: ov.Node, output_port_id: int, output_node_name: str) ->
return get_reduce_op


def get_inplace_min_op(reduction_axes: Optional[ReductionAxes]) -> InplaceInsertionFnType:
def get_inplace_min_op(reduction_axes: Optional[AggregationAxes]) -> InplaceInsertionFnType:
"""
Returns inplace min function that adds reduce min node to a passed node.

Expand All @@ -277,7 +280,7 @@ def get_inplace_min_op(reduction_axes: Optional[ReductionAxes]) -> InplaceInsert


def get_inplace_max_op(
reduction_axes: Optional[ReductionAxes], use_abs_max: bool, keep_dims: bool = True
reduction_axes: Optional[AggregationAxes], use_abs_max: bool, keep_dims: bool = True
) -> InplaceInsertionFnType:
"""
Returns inplace max function that adds reduce max node to a passed node.
Expand All @@ -291,7 +294,7 @@ def get_inplace_max_op(
return get_inplace_reduce_op(opset.reduce_max, reduction_axes, use_abs_max, keep_dims)


def get_inplace_mean_op(reduction_axes: Optional[ReductionAxes]) -> InplaceInsertionFnType:
def get_inplace_mean_op(reduction_axes: Optional[AggregationAxes]) -> InplaceInsertionFnType:
"""
Returns inplace mean function that adds reduce mean node to a passed node.

Expand Down Expand Up @@ -329,7 +332,7 @@ def var_op(
return variance


def get_inplace_mean_var_op(reduction_axes: Optional[ReductionAxes] = None) -> InplaceInsertionFnType:
def get_inplace_mean_var_op(reduction_axes: Optional[AggregationAxes] = None) -> InplaceInsertionFnType:
"""
Return an operation getter function that computes variance across given axes and then mean-reduces the result across
the remaining axes.
Expand Down Expand Up @@ -358,7 +361,7 @@ def get_mean_var_reduce_op(node: ov.Node, output_port_id: int, output_node_name:
return get_mean_var_reduce_op


def get_inplace_max_var_op(reduction_axes: Optional[ReductionAxes] = None) -> InplaceInsertionFnType:
def get_inplace_max_var_op(reduction_axes: Optional[AggregationAxes] = None) -> InplaceInsertionFnType:
"""
Return an operation getter function that computes variance across given axes and then max-reduces the result across
the remaining axes.
Expand Down Expand Up @@ -387,7 +390,7 @@ def get_max_var_reduce_op(node: ov.Node, output_port_id: int, output_node_name:
return get_max_var_reduce_op


def get_inplace_mean_max_op(reduction_axes: Optional[ReductionAxes], use_abs_max: bool) -> InplaceInsertionFnType:
def get_inplace_mean_max_op(reduction_axes: Optional[AggregationAxes], use_abs_max: bool) -> InplaceInsertionFnType:
"""
Return an operation getter function that computes maximum across given axes and then mean-reduces the result across
the remaining axes.
Expand Down
13 changes: 8 additions & 5 deletions nncf/quantization/algorithms/bias_correction/backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,10 @@
from nncf.common.graph.transformations.commands import TargetType
from nncf.common.graph.transformations.commands import TransformationCommand
from nncf.common.tensor import NNCFTensor
from nncf.common.tensor_statistics.collectors import TensorStatisticCollectorBase

# Using experimental tensor statistics implementation as part of the migration
# from old tensor statistics to experimental tensor statistics
from nncf.experimental.common.tensor_statistics.collectors import TensorCollector
from nncf.tensor import Tensor

TModel = TypeVar("TModel")
Expand Down Expand Up @@ -85,26 +88,26 @@ def mean_statistic_collector(
inplace: bool,
num_samples: Optional[int] = None,
window_size: Optional[int] = None,
) -> TensorStatisticCollectorBase:
) -> TensorCollector:
"""
Returns backend-specific mean statistic collector.

:param channel_axis: Channel axis for the statistics aggregation.
:param inplace: Whether to calculate statistic inplace or not.
:param num_samples: Maximum number of samples to collect.
:param window_size: The maximum size of the samples queue.
:return: Backend-specific TensorStatisticCollectorBase for the statistics calculation.
:return: Backend-specific TensorCollector for the statistics calculation.
"""

@staticmethod
@abstractmethod
def raw_statistic_collector(num_samples: Optional[int] = None) -> TensorStatisticCollectorBase:
def raw_statistic_collector(num_samples: Optional[int] = None) -> TensorCollector:
"""
Returns backend-specific raw statistic collector.
This statistic collector is used for raw data calculation, without aggregating.

:param num_samples: Maximum number of samples to collect.
:return: Backend-specific TensorStatisticCollectorBase for the statistics calculation.
:return: Backend-specific TensorCollector for the statistics calculation.
"""

@staticmethod
Expand Down
9 changes: 5 additions & 4 deletions nncf/quantization/algorithms/channel_alignment/backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,10 @@
from nncf.common.graph.layer_attributes import ConvolutionLayerAttributes
from nncf.common.graph.transformations.commands import TargetPoint
from nncf.common.graph.transformations.commands import TargetType
from nncf.common.tensor_statistics.collectors import TensorStatisticCollectorBase

# Using experimental tensor statistics implementation as part of the migration
# from old tensor statistics to experimental tensor statistics
from nncf.experimental.common.tensor_statistics.collectors import TensorCollector

TModel = TypeVar("TModel")

Expand Down Expand Up @@ -95,9 +98,7 @@ def get_weights_port_ids_for_node(node: NNCFNode) -> Tuple[int, int]:

@staticmethod
@abstractmethod
def get_statistic_collector(
reduction_axes, q: float, num_samples: int, inplace: bool
) -> TensorStatisticCollectorBase:
def get_statistic_collector(reduction_axes, q: float, num_samples: int, inplace: bool) -> TensorCollector:
"""
Get backend-specific tensor collector that collects medians of minimal and maximal quantiles.

Expand Down
5 changes: 1 addition & 4 deletions nncf/quantization/algorithms/channel_alignment/openvino_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@
from nncf.common.graph import NNCFNode
from nncf.common.graph.layer_attributes import ConvolutionLayerAttributes
from nncf.common.graph.transformations.commands import TargetType
from nncf.common.tensor_statistics.collectors import TensorStatisticCollectorBase
from nncf.experimental.common.tensor_statistics.collectors import MedianAggregator
from nncf.experimental.common.tensor_statistics.collectors import TensorCollector
from nncf.experimental.common.tensor_statistics.statistics import MinMaxTensorStatistic
Expand Down Expand Up @@ -77,9 +76,7 @@ def get_add_metatypes():
return [OVAddMetatype, OVSubtractMetatype]

@staticmethod
def get_statistic_collector(
reduction_axes, q: float, num_samples: int, inplace: bool
) -> TensorStatisticCollectorBase:
def get_statistic_collector(reduction_axes, q: float, num_samples: int, inplace: bool) -> TensorCollector:
tensor_collector = TensorCollector(MinMaxTensorStatistic)
quantile_reducer = OVQuantileReducer(reduction_axes, (q, 1 - q), inplace)

Expand Down
9 changes: 6 additions & 3 deletions nncf/quantization/algorithms/fast_bias_correction/backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,10 @@
from nncf.common.graph.transformations.commands import TargetType
from nncf.common.graph.transformations.commands import TransformationCommand
from nncf.common.graph.transformations.layout import TransformationLayout
from nncf.common.tensor_statistics.collectors import TensorStatisticCollectorBase

# Using experimental tensor statistics implementation as part of the migration
# from old tensor statistics to experimental tensor statistics
from nncf.experimental.common.tensor_statistics.collectors import TensorCollector
from nncf.tensor import Tensor

TModel = TypeVar("TModel")
Expand Down Expand Up @@ -77,15 +80,15 @@ def mean_statistic_collector(
inplace: bool,
num_samples: Optional[int] = None,
window_size: Optional[int] = None,
) -> TensorStatisticCollectorBase:
) -> TensorCollector:
"""
Returns backend-specific mean statistic collector.

:param channel_axis: Channel axes for the statistics aggregation.
:param inplace: Whether to calculate statistic inplace or not.
:param num_samples: Maximum number of samples to collect.
:param window_size: The maximum size of the samples queue.
:return: Backend-specific TensorStatisticCollectorBase for the statistics calculation.
:return: Backend-specific TensorCollector for the statistics calculation.
"""

@staticmethod
Expand Down
7 changes: 5 additions & 2 deletions nncf/quantization/algorithms/layerwise/backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,8 +16,11 @@
from nncf.common.graph import NNCFGraph
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.data.dataset import Dataset

# Using experimental tensor statistics implementation as part of the migration
# from old tensor statistics to experimental tensor statistics
from nncf.experimental.common.tensor_statistics.collectors import TensorCollector
from nncf.quantization.algorithms.layerwise.iterator import LayerwiseIterator
from nncf.quantization.algorithms.layerwise.scheduler import LayerwiseStep
from nncf.quantization.algorithms.layerwise.scheduler import NodeOutputPort
Expand Down Expand Up @@ -67,7 +70,7 @@ def target_point(target_type: TargetType, target_node_name: str, port_id: int) -

@staticmethod
@abstractmethod
def raw_statistic_collector(num_samples: Optional[int] = None) -> TensorStatisticCollectorBase:
def raw_statistic_collector(num_samples: Optional[int] = None) -> TensorCollector:
"""
Returns backend-specific raw statistic collector.
This statistic collector is used for raw data calculation, without aggregating.
Expand Down
7 changes: 5 additions & 2 deletions nncf/quantization/algorithms/layerwise/openvino_backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15,8 +15,11 @@

from nncf.common.graph import NNCFGraph
from nncf.common.graph.transformations.commands import TargetType
from nncf.common.tensor_statistics.collectors import TensorStatisticCollectorBase
from nncf.data.dataset import Dataset

# Using experimental tensor statistics implementation as part of the migration
# from old tensor statistics to experimental tensor statistics
from nncf.experimental.common.tensor_statistics.collectors import TensorCollector
from nncf.openvino.graph.transformations.commands import OVTargetPoint
from nncf.openvino.statistics.collectors import get_raw_stat_collector
from nncf.quantization.algorithms.layerwise.backend import LayerwiseEngineBackend
Expand All @@ -43,5 +46,5 @@ def target_point(target_type: TargetType, target_node_name: str, port_id: int) -
return OVTargetPoint(target_type, target_node_name, port_id)

@staticmethod
def raw_statistic_collector(num_samples: Optional[int] = None) -> TensorStatisticCollectorBase:
def raw_statistic_collector(num_samples: Optional[int] = None) -> TensorCollector:
return get_raw_stat_collector(num_samples)
3 changes: 1 addition & 2 deletions nncf/quantization/algorithms/weight_compression/backend.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,7 +18,6 @@
from nncf.common.graph.operator_metatypes import OperatorMetatype
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.experimental.common.tensor_statistics.collectors import HAWQAggregator
from nncf.experimental.common.tensor_statistics.collectors import RawReducer
Expand Down Expand Up @@ -213,7 +212,7 @@ def target_point(target_type: TargetType, target_node_name: str, port_id: int) -
@abstractmethod
def mean_statistic_collector(
self, reduction_axes: Tuple[int], subset_size: Optional[int] = None
) -> TensorStatisticCollectorBase:
) -> TensorCollector:
"""
Return mean statistic collector

Expand Down
10 changes: 6 additions & 4 deletions nncf/tensorflow/tensor_statistics/statistics.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,12 @@

import tensorflow as tf

from nncf.common.tensor_statistics.statistics import MedianMADTensorStatistic
from nncf.common.tensor_statistics.statistics import MinMaxTensorStatistic
from nncf.common.tensor_statistics.statistics import PercentileTensorStatistic
from nncf.common.tensor_statistics.statistics import TensorStatistic
# Using experimental tensor statistics implementation as part of the migration
# from old tensor statistics to experimental tensor statistics
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 PercentileTensorStatistic
from nncf.experimental.common.tensor_statistics.statistics import TensorStatistic


class TFMinMaxTensorStatistic(MinMaxTensorStatistic):
Expand Down
17 changes: 10 additions & 7 deletions nncf/torch/quantization/algo.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -59,7 +59,6 @@
from nncf.common.quantization.structs import WeightQuantizerId
from nncf.common.schedulers import BaseCompressionScheduler
from nncf.common.statistics import NNCFStatistics
from nncf.common.tensor_statistics.collectors import ReductionAxes
from nncf.common.utils.api_marker import api
from nncf.common.utils.backend import BackendType
from nncf.common.utils.backend import copy_model
Expand All @@ -75,6 +74,10 @@
from nncf.config.schemata.defaults import QUANTIZATION_PRESET
from nncf.config.schemata.defaults import QUANTIZE_INPUTS
from nncf.config.schemata.defaults import QUANTIZE_OUTPUTS

# Using experimental tensor statistics implementation as part of the migration
# from old tensor statistics to experimental tensor statistics
from nncf.experimental.common.tensor_statistics.collectors import AggregationAxes
from nncf.experimental.common.tensor_statistics.statistics import MinMaxTensorStatistic
from nncf.experimental.common.tensor_statistics.statistics import TensorStatistic
from nncf.parameters import StripFormat
Expand Down Expand Up @@ -592,7 +595,7 @@ def _parse_precision_init_params(self, initializer_config: Dict) -> Tuple[str, B
def _get_minmax_values_for_quantizer_locations(
self,
quantizer_setup: SingleConfigQuantizerSetup,
tensor_statistics: Dict[PTTargetPoint, Dict[ReductionAxes, TensorStatistic]],
tensor_statistics: Dict[PTTargetPoint, Dict[AggregationAxes, TensorStatistic]],
target_model_graph: PTNNCFGraph,
) -> Dict[QuantizationPointId, MinMaxTensorStatistic]:
retval = {}
Expand Down Expand Up @@ -669,7 +672,7 @@ def _get_transformation_layout(self, target_model: NNCFNetwork) -> PTTransformat
@staticmethod
def get_statistics_for_quantizer_setup(
target_model: NNCFNetwork, quantizer_setup: QuantizerSetupBase, range_init_params: PTRangeInitParams
) -> Dict[PTTargetPoint, Dict[ReductionAxes, TensorStatistic]]:
) -> Dict[PTTargetPoint, Dict[AggregationAxes, TensorStatistic]]:
if range_init_params is None:
return {}
observation_points_vs_collectors_dict = (
Expand All @@ -695,7 +698,7 @@ def get_statistics_for_quantizer_setup(

def _get_statistics_for_final_range_init(
self, target_model: NNCFNetwork, quantizer_setup: QuantizerSetupBase, range_init_params: PTRangeInitParams
) -> Dict[PTTargetPoint, Dict[ReductionAxes, TensorStatistic]]:
) -> Dict[PTTargetPoint, Dict[AggregationAxes, TensorStatistic]]:
return self.get_statistics_for_quantizer_setup(target_model, quantizer_setup, range_init_params)

def _get_single_config_quantizer_setup(self, target_model) -> SingleConfigQuantizerSetup:
Expand Down Expand Up @@ -1493,7 +1496,7 @@ def __init__(
self,
quantizer_setup: MultiConfigQuantizerSetup,
initial_quantizer_setup: SingleConfigQuantizerSetup,
tensor_stats_for_all_setup_variations: Dict[PTTargetPoint, Dict[ReductionAxes, TensorStatistic]],
tensor_stats_for_all_setup_variations: Dict[PTTargetPoint, Dict[AggregationAxes, TensorStatistic]],
hw_config: HWConfig = None,
):
should_init = bool(tensor_stats_for_all_setup_variations)
Expand All @@ -1512,7 +1515,7 @@ def _get_single_config_quantizer_setup(self, target_model) -> SingleConfigQuanti

def _get_statistics_for_final_range_init(
self, target_model: NNCFNetwork, quantizer_setup: QuantizerSetupBase, range_init_params: PTRangeInitParams
) -> Dict[PTTargetPoint, Dict[ReductionAxes, TensorStatistic]]:
) -> Dict[PTTargetPoint, Dict[AggregationAxes, TensorStatistic]]:
return self._tensor_stats

def _build_controller(self, model: NNCFNetwork) -> "ExperimentalQuantizationController":
Expand Down Expand Up @@ -1568,7 +1571,7 @@ def __init__(
quantizer_setup: MultiConfigQuantizerSetup,
initial_quantizer_setup: SingleConfigQuantizerSetup,
setup_to_module_id_translation_dict: Dict[QuantizationPointId, QuantizerId],
tensor_stats: Dict[PTTargetPoint, Dict[ReductionAxes, TensorStatistic]],
tensor_stats: Dict[PTTargetPoint, Dict[AggregationAxes, TensorStatistic]],
build_time_metric_info: QuantizationShareBuildTimeInfo,
should_setup_adjust_pad_ops=False,
hw_config: HWConfig = None,
Expand Down
Loading