added manual calibration configuration support

examples/onnxruntime_static_quant_vit.yaml

@@ -9,18 +9,15 @@ defaults:
 
 experiment_name: onnxruntime_static_quant_vit
 model: google/vit-base-patch16-224
-device: cuda
+device: cpu
 
 backend:
-  auto_quantization: avx2
-  auto_quantization_config:
+  quantization: true
+  quantization_config:
     is_static: true
     per_channel: false
 
-  auto_calibration: minmax
-
-launcher:
-  device_isolation: true
+  calibration: true
 
 hydra:
   run:

optimum_benchmark/backends/onnxruntime/backend.py

@@ -18,6 +18,7 @@
     AutoCalibrationConfig,
     AutoOptimizationConfig,
     AutoQuantizationConfig,
+    CalibrationConfig,
     OptimizationConfig,
     QuantizationConfig,
 )
@@ -31,7 +32,12 @@
 from ..peft_utils import get_peft_config_class
 from ..pytorch.utils import randomize_weights
 from .config import ORTConfig
-from .utils import TASKS_TO_ORTMODELS, TASKS_TO_ORTSD, format_quantization_config
+from .utils import (
+    TASKS_TO_ORTMODELS,
+    TASKS_TO_ORTSD,
+    format_calibration_config,
+    format_quantization_config,
+)
 
 # disable transformers logging
 set_verbosity_error()
@@ -201,7 +207,7 @@ def is_quantized(self) -> bool:
 
     @property
     def is_calibrated(self) -> bool:
-        return self.config.auto_calibration is not None
+        return (self.config.auto_calibration is not None) or self.config.calibration
 
     @property
     def automodel_kwargs(self) -> Dict[str, Any]:
@@ -290,22 +296,38 @@ def quantize_onnx_files(self) -> None:
 
         LOGGER.info("\t+ Processing quantization config")
         if self.config.auto_quantization is not None:
-            self.config.auto_quantization_config = format_quantization_config(self.config.auto_quantization_config)
-            auto_quantization_config_class = getattr(AutoQuantizationConfig, self.config.auto_quantization)
-            quantization_config = auto_quantization_config_class(**self.config.auto_quantization_config)
+            auto_quantization_config = format_quantization_config(self.config.auto_quantization_config)
+            auto_quantization_class = getattr(AutoQuantizationConfig, self.config.auto_quantization)
+            quantization_config = auto_quantization_class(**auto_quantization_config)
         elif self.config.quantization:
-            self.config.quantization_config = format_quantization_config(self.config.quantization_config)
-            quantization_config = QuantizationConfig(**self.config.quantization_config)
+            quantization_config = format_quantization_config(self.config.quantization_config)
+            quantization_config = QuantizationConfig(**quantization_config)
 
         if self.is_calibrated:
             LOGGER.info("\t+ Generating calibration dataset")
             dataset_shapes = {"dataset_size": 1, "sequence_length": 1, **self.model_shapes}
             calibration_dataset = DatasetGenerator(task=self.task, dataset_shapes=dataset_shapes).generate()
             columns_to_be_removed = list(set(calibration_dataset.column_names) - set(self.inputs_names))
             calibration_dataset = calibration_dataset.remove_columns(columns_to_be_removed)
+
             LOGGER.info("\t+ Processing calibration config")
-            calibration_config_method = getattr(AutoCalibrationConfig, self.config.auto_calibration)
-            calibration_config = calibration_config_method(calibration_dataset, **self.config.auto_calibration_config)
+            if self.config.auto_calibration is not None:
+                LOGGER.info("\t+ Processing calibration config")
+                auto_calibration_method = getattr(AutoCalibrationConfig, self.config.auto_calibration)
+                calibration_config = auto_calibration_method(
+                    calibration_dataset,
+                    **self.config.auto_calibration_config,
+                )
+            elif self.config.calibration:
+                LOGGER.info("\t+ Processing calibration config")
+                calibration_config = format_calibration_config(self.config.calibration_config)
+                calibration_config = CalibrationConfig(
+                    dataset_name="calibration_dataset",
+                    dataset_split=calibration_dataset.split,
+                    dataset_num_samples=calibration_dataset.num_rows,
+                    dataset_config_name=calibration_dataset.config_name,
+                    **self.config.calibration_config,
+                )
 
         for onnx_file_name_to_quantize in self.onnx_files_names_to_quantize:
             LOGGER.info(f"\t+ Creating quantizer for {onnx_file_name_to_quantize}")
@@ -358,15 +380,12 @@ def prepare_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
         if self.library == "diffusers":
             return {"prompt": inputs["prompt"]}
 
-        for key in list(inputs.keys()):
-            # sometimes optimum onnx exported models don't have inputs
-            # that their pytorch counterparts have, for instance token_type_ids
-            if key not in self.inputs_names:
-                inputs.pop(key)
-
         LOGGER.info(f"\t+ Moving inputs tensors to device {self.device}")
         for key, value in inputs.items():
-            inputs[key] = value.to(self.device)
+            if key not in self.inputs_names:
+                inputs.pop(key)
+            else:
+                inputs[key] = value.to(self.device)
 
         return inputs
 

optimum_benchmark/backends/onnxruntime/config.py

@@ -10,13 +10,17 @@
 
 QUANTIZATION_CONFIG = {
     "is_static": False,
-    "format": "QOperator",  # QOperator, QDQ
+    "format": "QOperator",
     # is_static and format are mandatory
 }
 
+CALIBRATION_CONFIG = {
+    "method": "MinMax"
+    # method is mandatory
+}
+
 AUTO_QUANTIZATION_CONFIG = {
     "is_static": False,
-    # full auto quantization config depends on the hardware but
     # is_static is mandatory
 }
 
@@ -78,6 +82,10 @@ class ORTConfig(BackendConfig):
     quantization: bool = False
     quantization_config: Dict[str, Any] = field(default_factory=dict)
 
+    # manual calibration options
+    calibration: bool = False
+    calibration_config: Dict[str, Any] = field(default_factory=dict)
+
     # ort-training is basically a different package so we might need to separate these two backends in the future
     use_inference_session: bool = "${is_inference:${benchmark.name}}"
 
@@ -100,22 +108,25 @@ def __post_init__(self):
                 OmegaConf.merge(QUANTIZATION_CONFIG, self.quantization_config)
             )
             # raise ValueError if the quantization is static but calibration is not enabled
-            if self.quantization_config["is_static"] and self.auto_calibration is None:
+            if self.quantization_config["is_static"] and self.auto_calibration is None and not self.calibration:
                 raise ValueError(
-                    "Quantization is static but auto calibration is not enabled. "
-                    "Please enable auto calibration or disable static quantization."
+                    "Quantization is static but calibration is not enabled. "
+                    "Please enable calibration or disable static quantization."
                 )
 
         if self.auto_quantization is not None:
             self.auto_quantization_config = OmegaConf.to_object(
                 OmegaConf.merge(AUTO_QUANTIZATION_CONFIG, self.auto_quantization_config)
             )
-            if self.auto_quantization_config["is_static"] and self.auto_calibration is None:
+            if self.auto_quantization_config["is_static"] and self.auto_calibration is None and not self.calibration:
                 raise ValueError(
-                    "Quantization is static but auto calibration is not enabled. "
-                    "Please enable auto calibration or disable static quantization."
+                    "Quantization is static but calibration is not enabled. "
+                    "Please enable calibration or disable static quantization."
                 )
 
+        if self.calibration:
+            self.calibration_config = OmegaConf.to_object(OmegaConf.merge(CALIBRATION_CONFIG, self.calibration_config))
+
         if self.peft_strategy is not None:
             if self.peft_strategy not in PEFT_CONFIGS:
                 raise ValueError(

optimum_benchmark/backends/onnxruntime/utils.py

@@ -1,6 +1,11 @@
 from typing import Any, Dict
 
-from onnxruntime.quantization import QuantFormat, QuantizationMode, QuantType
+from onnxruntime.quantization import (
+    CalibrationMethod,
+    QuantFormat,
+    QuantizationMode,
+    QuantType,
+)
 from optimum.pipelines import ORT_SUPPORTED_TASKS
 
 TASKS_TO_ORTSD = {
@@ -11,6 +16,13 @@
 TASKS_TO_ORTMODELS = {task: task_dict["class"][0] for task, task_dict in ORT_SUPPORTED_TASKS.items()}
 
 
+def format_calibration_config(calibration_config: Dict[str, Any]) -> None:
+    if calibration_config.get("method", None) is not None:
+        calibration_config["method"] = CalibrationMethod[calibration_config["method"]]
+
+    return calibration_config
+
+
 def format_quantization_config(quantization_config: Dict[str, Any]) -> None:
     """Format the quantization dictionary for onnxruntime."""
     # the conditionals are here because some quantization strategies don't have all the options