Merge pull request #29 from daniil-lyakhov/dl/quantize_model_ex

Improve quantize_model UX

Author: cavusmustafa

backends/openvino/quantizer/__init__.py

@@ -1,3 +1,3 @@
-from .quantizer import OpenVINOQuantizer
+from .quantizer import OpenVINOQuantizer, quantize_model
 
-__all__ = [OpenVINOQuantizer]
+__all__ = [OpenVINOQuantizer, quantize_model]

backends/openvino/quantizer/quantizer.py

@@ -6,7 +6,7 @@
 
 from collections import defaultdict
 from enum import Enum
-from typing import Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 import nncf
 import nncf.common.quantization as quantization
@@ -345,31 +345,57 @@ def validate(self, model: torch.fx.GraphModule) -> None:
 def quantize_model(
     captured_model: torch.fx.GraphModule,
     calibration_dataset: torch.utils.data.DataLoader,
+    *,
+    mode: QuantizationMode = QuantizationMode.INT8_SYM,
+    subset_size: int = 300,
+    fast_bias_correction: Optional[bool] = True,
+    smooth_quant: bool = False,
+    transform_fn: Optional[Callable[[Any], Any]] = None,
+    extra_quantizer_options: Optional[Dict[str, Any]] = None,
+    **kwargs,
 ) -> torch.fx.GraphModule:
     """
-    Quantizes a model using either NNCF-based or PTQ-based quantization.
+    Quantizes a model using NNCF quantize_pt2e API.
 
     :param captured_model: The model to be quantized, represented as a torch.fx.GraphModule.
     :param calibration_dataset: A DataLoader containing calibration data for quantization.
+    :param mode: Defines special quantization modes.
+        - INT8_SYM: INT8 symmetric quantization for both activations and weights.
+        - INT8_MIXED: INT8 asymmetric quantization for activations, symmetric for weights.
+        - INT8_TRANSFORMER: Optimized INT8 quantization for transformer-based models
+        Default value is INT8_SYM.
+    :param subset_size: Size of a subset to calculate activations
+        statistics used for quantization.
+    :param fast_bias_correction: Setting this option to `False` enables a different
+        bias correction method which is more accurate, in general, and takes
+        more time but requires less memory. None disables the bias correction algorithm.
+    :param smooth_quant: Setting this option to `True` enables the SmoothQuant algorithm.
+    :param extra_quantizer_options: A dictionary containing additional configuration options
+        for the OpenVINOQuantizer.
+    :param kwargs: The keyword arguments for the nncf quantize_pt2e function.
     :return: The quantized model as a torch.fx.GraphModule.
     """
-    quantizer = OpenVINOQuantizer()
+    extra_quantizer_options = extra_quantizer_options or {}
+    if "mode" in extra_quantizer_options:
+        print(
+            f'Ignoring "mode" from the quantizer_config. Using parameter mode = {mode}'
+        )
+        del extra_quantizer_options["mode"]
+
+    quantizer = OpenVINOQuantizer(mode=mode, **extra_quantizer_options)
 
     print("PTQ: Quantize the model")
-    default_subset_size = 300
-    batch_size = calibration_dataset.batch_size
-    subset_size = (default_subset_size // batch_size) + int(
-        default_subset_size % batch_size > 0
-    )
 
-    def transform(x):
-        return x[0]
+    if "fold_quantize" not in kwargs:
+        kwargs["fold_quantize"] = False
 
     quantized_model = nncf_fx.quantize_pt2e(
         captured_model,
         quantizer,
         subset_size=subset_size,
-        calibration_dataset=nncf.Dataset(calibration_dataset, transform_func=transform),
-        fold_quantize=False,
+        calibration_dataset=nncf.Dataset(calibration_dataset, transform_fn),
+        fast_bias_correction=fast_bias_correction,
+        smooth_quant=smooth_quant,
+        **kwargs,
     )
     return quantized_model

examples/openvino/README.md

@@ -57,12 +57,6 @@ python aot_openvino_compiler.py --suite <MODEL_SUITE> --model <MODEL_NAME> --inp
 - **`--quantize`** (optional):
   Enable model quantization. --dataset argument is requred for the quantization. `huggingface` suite  does not supported yet.
 
-- **`--quantization_flow`** (optional):
-  Specifies the way to quantize torch.fx.GraphModule.
-  Supported values:
-  - `nncf`: `nncf quantize_pt2e` API (default)
-  - `pt2e`: torch ao quantization pipeline.
-
 - **`--validate`** (optional):
   Enable model validation. --dataset argument is requred for the validation. `huggingface` suite does not supported yet.
 

examples/openvino/aot_openvino_compiler.py

@@ -13,7 +13,7 @@
 import torch
 import torchvision.models as torchvision_models
 from executorch.backends.openvino.partitioner import OpenvinoPartitioner
-from executorch.backends.openvino.quantizer.quantizer import quantize_model
+from executorch.backends.openvino.quantizer import quantize_model
 from executorch.exir import EdgeProgramManager, to_edge_transform_and_lower
 from executorch.exir.backend.backend_details import CompileSpec
 from executorch.extension.pybindings.portable_lib import (  # @manual
@@ -182,9 +182,19 @@ def main(
         if not dataset_path:
             msg = "Quantization requires a calibration dataset."
             raise ValueError(msg)
+
+        subset_size = 300
+        batch_size = calibration_dataset.batch_size
+        subset_size = (subset_size // batch_size) + int(subset_size % batch_size > 0)
+
+        def transform_fn(x):
+            return x[0]
+
         quantized_model = quantize_model(
             aten_dialect.module(),
             calibration_dataset,
+            subset_size=subset_size,
+            transform_fn=transform_fn,
         )
 
         aten_dialect: ExportedProgram = export(quantized_model, example_args)