[PTQ] Add support of arbitrary batch size for PTQ

examples/post_training_quantization/torch/mobilenet_v2/main.py

@@ -107,7 +107,8 @@ def get_model_size(ir_path: str, m_type: str = "Mb", verbose: bool = True) -> fl
         ]
     ),
 )
-val_data_loader = torch.utils.data.DataLoader(val_dataset)
+batch_size = 128
+val_data_loader = torch.utils.data.DataLoader(val_dataset, batch_size=batch_size)
 
 torch_model = models.mobilenet_v2(num_classes=DATASET_CLASSES)
 torch_model = load_checkpoint(torch_model)
@@ -140,8 +141,10 @@ def transform_fn(data_item: Tuple[torch.Tensor, int], device: torch.device) -> t
 # item and prepare model input data. The quantize method uses a small subset
 # (default: 300 samples) of the calibration dataset.
 
+# Recalculation default subset_size parameter based on batch_size.
+subset_size = 300 // batch_size
 calibration_dataset = nncf.Dataset(val_data_loader, partial(transform_fn, device=device))
-torch_quantized_model = nncf.quantize(torch_model, calibration_dataset)
+torch_quantized_model = nncf.quantize(torch_model, calibration_dataset, subset_size=subset_size)
 
 ###############################################################################
 # Benchmark performance, calculate compression rate and validate accuracy

nncf/common/graph/utils.py

@@ -10,7 +10,7 @@
 # limitations under the License.
 
 from functools import partial
-from typing import List, Set
+from typing import List, Set, Tuple, Union
 
 from nncf.common.graph import NNCFGraph
 from nncf.common.graph import NNCFNode
@@ -114,3 +114,19 @@ def get_number_of_quantized_ops(
             else:
                 nodes_to_see.extend(graph.get_next_nodes(node))
     return len(quantized_ops)
+
+
+def get_reduction_axes(
+    channel_axes: Union[List[int], Tuple[int, ...]], shape: Union[List[int], Tuple[int, ...]]
+) -> Tuple[int, ...]:
+    """
+    Returns filtered reduction axes without axes that corresponds channels.
+
+    :param channel_axes: Channel axes.
+    :param shape: Shape that need to be filtered.
+    :return: Reduction axes.
+    """
+    reduction_axes = list(range(len(shape)))
+    for channel_axis in sorted(channel_axes, reverse=True):
+        del reduction_axes[channel_axis]
+    return tuple(reduction_axes)

nncf/common/quantization/initialization/range.py

@@ -9,12 +9,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Tuple, Union
 
+from nncf.common.graph.utils import get_reduction_axes
 from nncf.common.initialization.dataloader import NNCFDataLoader
 from nncf.common.quantization.structs import QuantizationScheme
 from nncf.common.quantization.structs import QuantizerGroup
+from nncf.common.tensor_statistics.collectors import ReductionAxes
 from nncf.config.schemata.defaults import NUM_INIT_SAMPLES
+from nncf.experimental.common.tensor_statistics.collectors import AggregationAxes
 
 
 class RangeInitConfig:
@@ -204,3 +207,51 @@ def use_means_of_mins(self) -> bool:
     @property
     def use_means_of_maxs(self) -> bool:
         return not self._is_weights and not self._is_per_channel
+
+    def _get_reduction_axes(
+        self,
+        shape_to_reduce: Union[Tuple[int, ...], List[int]],
+        quantization_axes: Union[Tuple[int, ...], List[int]],
+        aggregation_axes: Union[Tuple[int, ...], List[int]],
+    ):
+        """
+        Returns axes for a reducer regarding aggregation axes. As aggregator takes axes counting from stacked tensors,
+        from these axes only tensor related axes should be used for reducer.
+
+        :param shape_to_reduce: Shape of a reduced tensor.
+        :param quantization_axes: Axes of quantization.
+        :param aggregation_axes: Axes of aggregator which is applied onto reduced tensor.
+        :return: Axes for reducer.
+        """
+        axes_to_keep = set(el - 1 for el in aggregation_axes if el != 0)
+        axes_to_keep.update(quantization_axes)
+        return get_reduction_axes(axes_to_keep, shape_to_reduce)
+
+    def _get_aggregation_axes(self, batchwise_statistics: bool) -> Tuple[int, ...]:
+        """
+        Returns axes for aggregator.
+
+        :param batchwise_statistics: Determines whether quantizer statistics should be calculated
+            for each item of the batch or for the entire batch.
+        :return Tuple[int]: Aggregation axes.
+        """
+        return (0, 1) if batchwise_statistics else (0,)
+
+    def get_reduction_aggregation_axes(
+        self,
+        shape_to_reduce: Union[Tuple[int, ...], List[int]],
+        quantization_axes: Union[Tuple[int, ...], List[int]],
+        batchwise_statistics: bool,
+    ) -> Tuple[ReductionAxes, AggregationAxes]:
+        """
+        Calculates the reduction axes, aggregation axes for the tensor.
+
+        :param shape_to_reduce: Shape of the tensor.
+        :param quantization_axes: Quantization axes if per-channel quantization.
+        :param batchwise_statistics: Determines whether quantizer statistics should be calculated
+            for each item of the batch or for the entire batch.
+        :return: Reduction axes and aggregation axes.
+        """
+        aggregation_axes = self._get_aggregation_axes(batchwise_statistics)
+        reduction_axes = self._get_reduction_axes(shape_to_reduce, quantization_axes, aggregation_axes)
+        return reduction_axes, aggregation_axes

nncf/common/tensor_statistics/aggregator.py

@@ -11,12 +11,13 @@
 from abc import ABC
 from abc import abstractmethod
 from itertools import islice
-from typing import Any, Dict, TypeVar
+from typing import Any, Dict, Optional, TypeVar
 
 import nncf
 from nncf.common import factory
 from nncf.common.graph.graph import NNCFGraph
 from nncf.common.graph.transformations.layout import TransformationLayout
+from nncf.common.logging.logger import nncf_logger
 from nncf.common.logging.track_progress import track
 from nncf.common.tensor import NNCFTensor
 from nncf.common.tensor_statistics.statistic_point import StatisticPointsContainer
@@ -25,6 +26,13 @@
 TensorType = TypeVar("TensorType")
 TModel = TypeVar("TModel")
 
+EMPTY_DATASET_ERROR = (
+    "Calibration dataset must not be empty. Please provide calibration dataset with at least one sample."
+)
+ITERATIONS_NUMBER_WARNING = (
+    "The number of iterations for statistics collection is bigger than the length of the dataset."
+)
+
 
 class StatisticsAggregator(ABC):
     """
@@ -33,9 +41,23 @@ class StatisticsAggregator(ABC):
 
     def __init__(self, dataset: Dataset):
         self.dataset = dataset
-        self.stat_subset_size = None
+        self.iterations_number = None
         self.statistic_points = StatisticPointsContainer()
 
+    def _get_iterations_number(self) -> Optional[int]:
+        """
+        Returns number of iterations which in min(self.iterations_number, dataset_length).
+
+        :return: Number of iterations for statistics collection.
+        """
+        dataset_length = self.dataset.get_length()
+        if dataset_length and self.iterations_number:
+            if self.iterations_number > dataset_length:
+                nncf_logger.warning(ITERATIONS_NUMBER_WARNING)
+                return dataset_length
+            return self.iterations_number
+        return dataset_length or self.iterations_number
+
     def collect_statistics(self, model: TModel, graph: NNCFGraph) -> None:
         """
         Collects statistics for registered StatisticPoints.
@@ -46,34 +68,25 @@ def collect_statistics(self, model: TModel, graph: NNCFGraph) -> None:
         """
         if not self.statistic_points:
             return
-
         model_transformer = factory.ModelTransformerFactory.create(model)
-
         merged_statistics = self._get_merged_statistic_points(self.statistic_points, model, graph)
         transformation_layout = self._get_transformation_layout_extra_outputs(merged_statistics)
         model_with_outputs = model_transformer.transform(transformation_layout)
         engine = factory.EngineFactory.create(model_with_outputs)
 
-        dataset_length = self.dataset.get_length()
-        total = (
-            min(dataset_length or self.stat_subset_size, self.stat_subset_size)
-            if self.stat_subset_size is not None
-            else None
-        )
+        iterations_number = self._get_iterations_number()
         empty_statistics = True
         for input_data in track(
-            islice(self.dataset.get_inference_data(), self.stat_subset_size),
-            total=total,
+            islice(self.dataset.get_inference_data(), iterations_number),
+            total=self.iterations_number,
             description="Statistics collection",
         ):
             outputs = engine.infer(input_data)
             processed_outputs = self._process_outputs(outputs)
             self._register_statistics(processed_outputs, merged_statistics)
             empty_statistics = False
         if empty_statistics:
-            raise nncf.ValidationError(
-                "Calibration dataset must not be empty. Please provide calibration dataset with at least one sample."
-            )
+            raise nncf.ValidationError(EMPTY_DATASET_ERROR)
 
     def register_statistic_points(self, statistic_points: StatisticPointsContainer) -> None:
         """
@@ -90,10 +103,10 @@ def register_statistic_points(self, statistic_points: StatisticPointsContainer)
             for _statistic_point in _statistic_points:
                 for _, tensor_collectors in _statistic_point.algorithm_to_tensor_collectors.items():
                     for tensor_collector in tensor_collectors:
-                        if self.stat_subset_size is None:
-                            self.stat_subset_size = tensor_collector.num_samples
+                        if self.iterations_number is None:
+                            self.iterations_number = tensor_collector.num_samples
                         elif tensor_collector.num_samples is not None:
-                            self.stat_subset_size = max(self.stat_subset_size, tensor_collector.num_samples)
+                            self.iterations_number = max(self.iterations_number, tensor_collector.num_samples)
 
     @abstractmethod
     def _register_statistics(self, outputs: Dict[str, NNCFTensor], statistic_points: StatisticPointsContainer) -> None:

nncf/data/dataset.py

@@ -81,6 +81,17 @@ def get_length(self) -> Optional[int]:
             return self._data_source.__len__()
         return None
 
+    def get_batch_size(self) -> Optional[int]:
+        """
+        Tries to fetch batch size of the underlying dataset.
+        :return: The value of batch_size or _batch_size attributes of the data_source if exist, and None otherwise.
+        """
+        if hasattr(self._data_source, "batch_size"):  # Torch dataloader
+            return self._data_source.batch_size
+        if hasattr(self._data_source, "_batch_size"):  # TF dataloader
+            return self._data_source._batch_size
+        return None
+
 
 class DataProvider(Generic[DataItem, ModelInput]):
     def __init__(

nncf/experimental/common/tensor_statistics/statistical_functions.py

@@ -9,22 +9,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from typing import Optional
+
 from nncf.experimental.tensor import Tensor
+from nncf.experimental.tensor import TensorDataType
 from nncf.experimental.tensor.functions import numeric as fns
 
 
-def mean_per_channel(x: Tensor, axis: int) -> Tensor:
+def mean_per_channel(x: Tensor, axis: int, dtype: Optional[TensorDataType] = None) -> Tensor:
     """
     Computes the mean of elements across given channel dimension of Tensor.
 
     :param x: Tensor to reduce.
     :param axis: The channel dimensions to reduce.
+    :param dtype: Type to use in computing the mean.
     :return: Reduced Tensor.
     """
     if len(x.shape) < 3:
-        return fns.mean(x, axis=0)
+        return fns.mean(x, axis=0, dtype=dtype)
+
     pos_axis = axis + x.ndim if axis < 0 else axis
     if pos_axis < 0 or pos_axis >= x.ndim:
         raise ValueError(f"axis {axis} is out of bounds for array of dimension {x.ndim}")
     axis = tuple(i for i in range(x.ndim) if i != pos_axis)
-    return fns.mean(x, axis=axis)
+    return fns.mean(x, axis=axis, dtype=dtype)

nncf/experimental/tensor/functions/numeric.py

@@ -355,16 +355,19 @@ def moveaxis(a: Tensor, source: Union[int, Tuple[int, ...]], destination: Union[
 
 @functools.singledispatch
 @tensor_guard
-def mean(a: Tensor, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False) -> Tensor:
+def mean(
+    a: Tensor, axis: Optional[Union[int, Tuple[int, ...]]] = None, keepdims: bool = False, dtype: TensorDataType = None
+) -> Tensor:
     """
     Compute the arithmetic mean along the specified axis.
 
     :param a: Array containing numbers whose mean is desired.
     :param axis: Axis or axes along which the means are computed.
     :param keepdims: Destination positions for each of the original axes. These must also be unique.
+    :param dtype: Type to use in computing the mean.
     :return: Array with moved axes.
     """
-    return Tensor(mean(a.data, axis, keepdims))
+    return Tensor(mean(a.data, axis, keepdims, dtype))
 
 
 @functools.singledispatch

nncf/experimental/tensor/functions/numpy_numeric.py

@@ -170,8 +170,14 @@ def _(a: np.ndarray, source: Union[int, Tuple[int, ...]], destination: Union[int
 
 
 @register_numpy_types(numeric.mean)
-def _(a: Union[np.ndarray, np.generic], axis: Union[int, Tuple[int, ...]] = None, keepdims: bool = False) -> np.ndarray:
-    return np.array(np.mean(a, axis=axis, keepdims=keepdims))
+def _(
+    a: Union[np.ndarray, np.generic],
+    axis: Union[int, Tuple[int, ...]] = None,
+    keepdims: bool = False,
+    dtype: Optional[TensorDataType] = None,
+) -> np.ndarray:
+    dtype = DTYPE_MAP[dtype] if dtype else None
+    return np.array(np.mean(a, axis=axis, keepdims=keepdims, dtype=dtype))
 
 
 @register_numpy_types(numeric.round)

nncf/experimental/tensor/functions/torch_numeric.py

@@ -183,8 +183,14 @@ def _(a: torch.Tensor, source: Union[int, Tuple[int, ...]], destination: Union[i
 
 
 @numeric.mean.register(torch.Tensor)
-def _(a: torch.Tensor, axis: Union[int, Tuple[int, ...]] = None, keepdims: bool = False) -> torch.Tensor:
-    return torch.mean(a, dim=axis, keepdim=keepdims)
+def _(
+    a: torch.Tensor,
+    axis: Union[int, Tuple[int, ...]] = None,
+    keepdims: bool = False,
+    dtype: Optional[TensorDataType] = None,
+) -> torch.Tensor:
+    dtype = DTYPE_MAP[dtype] if dtype else None
+    return torch.mean(a, dim=axis, keepdim=keepdims, dtype=dtype)
 
 
 @numeric.round.register(torch.Tensor)

nncf/onnx/graph/metatypes/groups.py

@@ -127,3 +127,12 @@
     onnx_metatypes.ONNXQuantizeLinearMetatype,
     onnx_metatypes.ONNXDequantizeLinearMetatype,
 ]
+
+# These metatypes mix outputs for different samples into one axis.
+# If reducers and aggregators collect statistics at the output of the following operations,
+# assuming that 0-axis is batch axis, they get only 1 value instead of batch_size values.
+# It could lead to inaccurate/incorrect statistics result.
+OPERATIONS_OUTPUT_HAS_NO_BATCH_AXIS = [
+    onnx_metatypes.ONNXROIAlignMetatype,
+    onnx_metatypes.ONNXEmbeddingMetatype,
+]

nncf/onnx/graph/metatypes/onnx_metatypes.py

@@ -113,7 +113,7 @@ class ONNXGemmMetatype(ONNXOpWithWeightsMetatype):
     name = "GemmOp"
     op_names = ["Gemm"]
     hw_config_names = [HWConfigOpName.MATMUL]
-    weight_channel_axis = -1
+    weight_channel_axis = -1  # For port_id=1
     weight_port_ids = None
     bias_port_id = 2
     possible_weight_ports = [0, 1]
@@ -125,7 +125,7 @@ class ONNXMatMulMetatype(ONNXOpMetatype):
     name = "MatMulOp"
     op_names = ["MatMul"]
     hw_config_names = [HWConfigOpName.MATMUL]
-    weight_channel_axis = -1
+    weight_channel_axis = -1  # For port_id=1
     weight_port_ids = None
     bias_port_id = 2
     possible_weight_ports = [0, 1]
@@ -446,8 +446,8 @@ class ONNXScatterNDMetatype(ONNXOpMetatype):
 
 
 @ONNX_OPERATION_METATYPES.register()
-class ONNXRoiAlignMetatype(ONNXOpMetatype):
-    name = "RoiAlignOp"
+class ONNXROIAlignMetatype(ONNXOpMetatype):
+    name = "ROIAlignOp"
     op_names = ["RoiAlign"]