diff --git a/nncf/experimental/common/tensor_statistics/collectors.py b/nncf/experimental/common/tensor_statistics/collectors.py index e6e2a9ecee7..0bf6230b742 100644 --- a/nncf/experimental/common/tensor_statistics/collectors.py +++ b/nncf/experimental/common/tensor_statistics/collectors.py @@ -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 @@ -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, ...] @@ -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. @@ -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))) @@ -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, ): @@ -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, ): diff --git a/nncf/openvino/graph/node_utils.py b/nncf/openvino/graph/node_utils.py index 7370be9f7a6..b5c99a92829 100644 --- a/nncf/openvino/graph/node_utils.py +++ b/nncf/openvino/graph/node_utils.py @@ -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 @@ -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: @@ -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. @@ -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. @@ -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. @@ -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. @@ -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. @@ -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. diff --git a/nncf/quantization/algorithms/bias_correction/backend.py b/nncf/quantization/algorithms/bias_correction/backend.py index d1d664ff101..14bd6626434 100644 --- a/nncf/quantization/algorithms/bias_correction/backend.py +++ b/nncf/quantization/algorithms/bias_correction/backend.py @@ -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") @@ -85,7 +88,7 @@ def mean_statistic_collector( inplace: bool, num_samples: Optional[int] = None, window_size: Optional[int] = None, - ) -> TensorStatisticCollectorBase: + ) -> TensorCollector: """ Returns backend-specific mean statistic collector. @@ -93,18 +96,18 @@ def mean_statistic_collector( :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 diff --git a/nncf/quantization/algorithms/channel_alignment/backend.py b/nncf/quantization/algorithms/channel_alignment/backend.py index a3032b8f4a2..7fed65bef11 100644 --- a/nncf/quantization/algorithms/channel_alignment/backend.py +++ b/nncf/quantization/algorithms/channel_alignment/backend.py @@ -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") @@ -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. diff --git a/nncf/quantization/algorithms/channel_alignment/openvino_backend.py b/nncf/quantization/algorithms/channel_alignment/openvino_backend.py index 343ed0416e7..294f54bfa5d 100644 --- a/nncf/quantization/algorithms/channel_alignment/openvino_backend.py +++ b/nncf/quantization/algorithms/channel_alignment/openvino_backend.py @@ -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 @@ -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) diff --git a/nncf/quantization/algorithms/fast_bias_correction/backend.py b/nncf/quantization/algorithms/fast_bias_correction/backend.py index 7256446069c..f57881b4773 100644 --- a/nncf/quantization/algorithms/fast_bias_correction/backend.py +++ b/nncf/quantization/algorithms/fast_bias_correction/backend.py @@ -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") @@ -77,7 +80,7 @@ def mean_statistic_collector( inplace: bool, num_samples: Optional[int] = None, window_size: Optional[int] = None, - ) -> TensorStatisticCollectorBase: + ) -> TensorCollector: """ Returns backend-specific mean statistic collector. @@ -85,7 +88,7 @@ def mean_statistic_collector( :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 diff --git a/nncf/quantization/algorithms/layerwise/backend.py b/nncf/quantization/algorithms/layerwise/backend.py index 031c9ad59e9..2060201ad0e 100644 --- a/nncf/quantization/algorithms/layerwise/backend.py +++ b/nncf/quantization/algorithms/layerwise/backend.py @@ -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 @@ -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. diff --git a/nncf/quantization/algorithms/layerwise/openvino_backend.py b/nncf/quantization/algorithms/layerwise/openvino_backend.py index e82242807ba..7766ab59a78 100644 --- a/nncf/quantization/algorithms/layerwise/openvino_backend.py +++ b/nncf/quantization/algorithms/layerwise/openvino_backend.py @@ -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 @@ -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) diff --git a/nncf/quantization/algorithms/weight_compression/backend.py b/nncf/quantization/algorithms/weight_compression/backend.py index bae917b0289..4d27b124711 100644 --- a/nncf/quantization/algorithms/weight_compression/backend.py +++ b/nncf/quantization/algorithms/weight_compression/backend.py @@ -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 @@ -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 diff --git a/nncf/tensorflow/tensor_statistics/statistics.py b/nncf/tensorflow/tensor_statistics/statistics.py index 06c3acff440..59ea4d1086f 100644 --- a/nncf/tensorflow/tensor_statistics/statistics.py +++ b/nncf/tensorflow/tensor_statistics/statistics.py @@ -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): diff --git a/nncf/torch/quantization/algo.py b/nncf/torch/quantization/algo.py index 28bcb3ccaff..b2051c26a6f 100644 --- a/nncf/torch/quantization/algo.py +++ b/nncf/torch/quantization/algo.py @@ -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 @@ -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 @@ -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 = {} @@ -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 = ( @@ -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: @@ -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) @@ -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": @@ -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, diff --git a/nncf/torch/quantization/init_range.py b/nncf/torch/quantization/init_range.py index fa00b82899d..2bdc68296c8 100644 --- a/nncf/torch/quantization/init_range.py +++ b/nncf/torch/quantization/init_range.py @@ -31,10 +31,12 @@ from nncf.common.quantization.structs import QuantizerId from nncf.common.quantization.structs import WeightQuantizerId from nncf.common.scopes import should_consider_scope -from nncf.common.tensor_statistics.collectors import ReductionAxes -from nncf.common.tensor_statistics.collectors import TensorStatisticCollectorBase from nncf.config.schemata.algo.quantization import RANGE_INIT_TYPES_VS_DESCRIPTIONS + +# 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.collectors import TensorCollector from nncf.torch.graph.graph import PTNNCFGraph from nncf.torch.initialization import DataLoaderBaseRunner from nncf.torch.nncf_network import NNCFNetwork @@ -107,7 +109,7 @@ def __init__( self._input_shape = input_shape self._channel_idx = channel_idx - def get_reduction_aggregation_axes(self, is_per_sample: bool) -> Tuple[ReductionAxes, AggregationAxes]: + def get_reduction_aggregation_axes(self, is_per_sample: bool) -> Tuple[AggregationAxes, AggregationAxes]: if self.is_per_channel: return super().get_reduction_aggregation_axes(self._input_shape, (self._channel_idx,), is_per_sample) return super().get_reduction_aggregation_axes(self._input_shape, (), is_per_sample) @@ -117,7 +119,7 @@ class StatCollectorGenerator: @staticmethod def generate_collectors_for_range_init_statistics_collection( target_model_graph: PTNNCFGraph, quantizer_setup: QuantizerSetupBase, range_init_params: PTRangeInitParams - ) -> Dict[TensorStatisticObservationPoint, Dict[ReductionAxes, TensorStatisticCollectorBase]]: + ) -> Dict[TensorStatisticObservationPoint, Dict[AggregationAxes, TensorCollector]]: retval = {} for qp in quantizer_setup.quantization_points.values(): init_config = range_init_params.get_init_config_for_quantization_point(qp) @@ -147,10 +149,10 @@ def generate_collectors_for_range_init_statistics_collection( @staticmethod def generate_stat_collector_for_range_init_config( init_config: RangeInitConfig, - scale_shape: ReductionAxes = None, + scale_shape: AggregationAxes = None, collector_params: PTRangeInitCollectorParams = None, num_samples_to_collect_override: int = None, - ) -> TensorStatisticCollectorBase: + ) -> TensorCollector: num_samples = init_config.num_init_samples if num_samples_to_collect_override is not None: num_samples = num_samples_to_collect_override @@ -222,7 +224,7 @@ def generate_stat_collector_for_range_init_config( @classmethod def get_all_scale_shapes_with_params( cls, qp: QuantizationPointBase, target_nncf_graph: PTNNCFGraph - ) -> Dict[ReductionAxes, PTRangeInitCollectorParams]: + ) -> Dict[AggregationAxes, PTRangeInitCollectorParams]: qconfigs = qp.get_all_configs_list() if qp.is_weight_quantization_point(): module_node = target_nncf_graph.get_node_by_name(qp.insertion_point.target_node_name) @@ -262,14 +264,12 @@ def __init__( self.modules_to_init = modules_to_init_vs_init_configs self.progressbar_description = "Range parameters initialization" - self.collectors_and_modules_to_init: Dict[str, Tuple[TensorStatisticCollectorBase, BaseQuantizer]] = ( - OrderedDict() - ) + self.collectors_and_modules_to_init: Dict[str, Tuple[TensorCollector, BaseQuantizer]] = OrderedDict() self.hook_handles = [] self.batch_size = batch_size def _get_fwd_hook( - self, collector: TensorStatisticCollectorBase + self, collector: TensorCollector ) -> Callable[["torch.Module", torch.Tensor, torch.Tensor], torch.Tensor]: hook = create_register_input_hook(collector=collector) diff --git a/nncf/torch/tensor_statistics/algo.py b/nncf/torch/tensor_statistics/algo.py index a3c70d51097..67ca55b0d8a 100644 --- a/nncf/torch/tensor_statistics/algo.py +++ b/nncf/torch/tensor_statistics/algo.py @@ -15,9 +15,11 @@ from nncf.api.compression import CompressionStage from nncf.common.schedulers import StubCompressionScheduler from nncf.common.statistics import NNCFStatistics -from nncf.common.tensor_statistics.collectors import ReductionAxes -from nncf.common.tensor_statistics.collectors import TensorStatisticCollectorBase from nncf.config import NNCFConfig + +# 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.collectors import TensorCollector from nncf.tensor import Tensor from nncf.torch.algo_selector import ZeroCompressionLoss @@ -33,7 +35,7 @@ class TensorStatisticObservationPoint: - def __init__(self, target_point: PTTargetPoint, reduction_shapes: Set[ReductionAxes] = None): + def __init__(self, target_point: PTTargetPoint, reduction_shapes: Set[AggregationAxes] = None): self.target_point = target_point self.reduction_shapes = reduction_shapes @@ -71,7 +73,7 @@ class TensorStatisticsCollectionBuilder(PTCompressionAlgorithmBuilder): def __init__( self, config: NNCFConfig, - observation_points_vs_collectors: Dict[TensorStatisticObservationPoint, TensorStatisticCollectorBase], + observation_points_vs_collectors: Dict[TensorStatisticObservationPoint, TensorCollector], ): super().__init__(config) self._observation_points_vs_collectors = observation_points_vs_collectors @@ -107,9 +109,7 @@ def _get_algo_specific_config_section(self) -> Dict: class TensorStatisticsCollectionController(PTCompressionAlgorithmController): - def __init__( - self, target_model: NNCFNetwork, ip_vs_collector_dict: Dict[PTTargetPoint, TensorStatisticCollectorBase] - ): + def __init__(self, target_model: NNCFNetwork, ip_vs_collector_dict: Dict[PTTargetPoint, TensorCollector]): super().__init__(target_model) self.ip_vs_collector_dict = ip_vs_collector_dict self._scheduler = StubCompressionScheduler() diff --git a/tests/cross_fw/test_templates/test_channel_alignment.py b/tests/cross_fw/test_templates/test_channel_alignment.py index ed5d92fac88..d46a4feaa6c 100644 --- a/tests/cross_fw/test_templates/test_channel_alignment.py +++ b/tests/cross_fw/test_templates/test_channel_alignment.py @@ -24,10 +24,13 @@ from nncf.common.graph.transformations.commands import TransformationType from nncf.common.tensor_statistics.statistic_point import StatisticPoint from nncf.common.tensor_statistics.statistic_point import StatisticPointsContainer -from nncf.common.tensor_statistics.statistics import MinMaxTensorStatistic from nncf.experimental.common.tensor_statistics.collectors import MedianAggregator from nncf.experimental.common.tensor_statistics.collectors import QuantileReducer from nncf.experimental.common.tensor_statistics.collectors import TensorCollector + +# 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 MinMaxTensorStatistic from nncf.quantization.algorithms.channel_alignment.algorithm import ChannelAlignment from nncf.quantization.algorithms.channel_alignment.backend import ChannelAlignmentAlgoBackend from nncf.quantization.algorithms.channel_alignment.backend import LayoutDescriptor diff --git a/tests/tensorflow/tensor_statistics/test_tensor_statistics.py b/tests/tensorflow/tensor_statistics/test_tensor_statistics.py index 337e6519d20..3c26feedde9 100644 --- a/tests/tensorflow/tensor_statistics/test_tensor_statistics.py +++ b/tests/tensorflow/tensor_statistics/test_tensor_statistics.py @@ -16,10 +16,13 @@ import tensorflow as tf from nncf.common.tensor_statistics.collectors import OfflineTensorStatisticCollector -from nncf.common.tensor_statistics.collectors import ReductionAxes from nncf.common.tensor_statistics.collectors import StatisticsNotCollectedError -from nncf.common.tensor_statistics.collectors import TensorStatisticCollectorBase -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.collectors import AggregationAxes +from nncf.experimental.common.tensor_statistics.collectors import TensorCollector +from nncf.experimental.common.tensor_statistics.statistics import TensorStatistic from nncf.tensorflow.tensor_statistics.collectors import TFMeanMinMaxStatisticCollector from nncf.tensorflow.tensor_statistics.collectors import TFMeanPercentileStatisticCollector from nncf.tensorflow.tensor_statistics.collectors import TFMedianMADStatisticCollector @@ -98,8 +101,8 @@ class TestCollectedStatistics: ) def test_collected_statistics_with_shape_convert( self, - collector: Type[TensorStatisticCollectorBase], - reduction_shapes_vs_ref_statistic: Dict[Tuple[ReductionAxes, ReductionAxes], TensorStatistic], + collector: Type[TensorCollector], + reduction_shapes_vs_ref_statistic: Dict[Tuple[AggregationAxes, AggregationAxes], TensorStatistic], ): for reduction_shape in reduction_shapes_vs_ref_statistic: collector_obj = collector(use_abs_max=True, reduction_shape=reduction_shape) @@ -176,8 +179,8 @@ def test_collected_statistics_with_shape_convert( ) def test_collected_statistics( self, - collector: Type[TensorStatisticCollectorBase], - reduction_shapes_vs_ref_statistic: Dict[ReductionAxes, TensorStatistic], + collector: Type[TensorCollector], + reduction_shapes_vs_ref_statistic: Dict[AggregationAxes, TensorStatistic], ): for reduction_shape in reduction_shapes_vs_ref_statistic: collector_obj = collector(reduction_shape=reduction_shape) @@ -206,12 +209,12 @@ def collector_for_interface_test(self, request): collector_type = request.param return collector_type(reduction_shape=(1,)) - def test_collected_samples(self, collector_for_interface_test: TensorStatisticCollectorBase): + def test_collected_samples(self, collector_for_interface_test: TensorCollector): for input_ in TestCollectedStatistics.REF_INPUTS: collector_for_interface_test.register_input(input_) assert collector_for_interface_test.collected_samples() == len(TestCollectedStatistics.REF_INPUTS) - def test_reset(self, collector_for_interface_test: TensorStatisticCollectorBase): + def test_reset(self, collector_for_interface_test: TensorCollector): for input_ in TestCollectedStatistics.REF_INPUTS: collector_for_interface_test.register_input(input_) collector_for_interface_test.reset() @@ -219,7 +222,7 @@ def test_reset(self, collector_for_interface_test: TensorStatisticCollectorBase) with pytest.raises(StatisticsNotCollectedError): collector_for_interface_test.get_statistics() - def test_enable_disable(self, collector_for_interface_test: TensorStatisticCollectorBase): + def test_enable_disable(self, collector_for_interface_test: TensorCollector): for input_ in TestCollectedStatistics.REF_INPUTS: collector_for_interface_test.register_input(input_) diff --git a/tests/torch/tensor_statistics/test_tensor_statistics.py b/tests/torch/tensor_statistics/test_tensor_statistics.py index f2b596106f4..1c2a6c05a08 100644 --- a/tests/torch/tensor_statistics/test_tensor_statistics.py +++ b/tests/torch/tensor_statistics/test_tensor_statistics.py @@ -15,12 +15,14 @@ import pytest import torch -from nncf.common.tensor_statistics.collectors import ReductionAxes -from nncf.common.tensor_statistics.collectors import TensorStatisticCollectorBase -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.collectors import AggregationAxes +from nncf.experimental.common.tensor_statistics.collectors import TensorCollector 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 from nncf.tensor import Tensor from nncf.tensor import functions as fns from nncf.torch.tensor_statistics.collectors import get_mean_percentile_statistic_collector @@ -112,8 +114,8 @@ class TestCollectedStatistics: ) def test_collected_statistics_with_shape_convert( self, - collector: Type[TensorStatisticCollectorBase], - reduction_axes_vs_ref_statistic: Dict[Tuple[ReductionAxes, ReductionAxes], TensorStatistic], + collector: Type[TensorCollector], + reduction_axes_vs_ref_statistic: Dict[Tuple[AggregationAxes, AggregationAxes], TensorStatistic], ): for shapes in reduction_axes_vs_ref_statistic: scale_shape, reducer_axes = shapes @@ -203,8 +205,8 @@ def test_collected_statistics_with_shape_convert( ) def test_collected_statistics( self, - collector: Type[TensorStatisticCollectorBase], - reduction_axes_vs_ref_statistic: Dict[ReductionAxes, TensorStatistic], + collector: Type[TensorCollector], + reduction_axes_vs_ref_statistic: Dict[AggregationAxes, TensorStatistic], ): for reduction_axes in reduction_axes_vs_ref_statistic: if len(reduction_axes) > 1: