From 4f9a167f96758767ec6b8ea662ecf65155e9b661 Mon Sep 17 00:00:00 2001 From: nghiavt1 Date: Tue, 16 Jul 2024 09:50:20 +0700 Subject: [PATCH] feat(quantization): add extra_options --- optimum/onnxruntime/quantization.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/optimum/onnxruntime/quantization.py b/optimum/onnxruntime/quantization.py index 056123f8d8..d744e80340 100644 --- a/optimum/onnxruntime/quantization.py +++ b/optimum/onnxruntime/quantization.py @@ -19,6 +19,7 @@ from collections import defaultdict from pathlib import Path from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Union +from typing import Any import onnx from datasets import Dataset, load_dataset @@ -286,6 +287,7 @@ def quantize( calibration_tensors_range: Optional[Dict[str, Tuple[float, float]]] = None, use_external_data_format: bool = False, preprocessor: Optional[QuantizationPreprocessor] = None, + extra_options: Optional[Dict[str, Any]] = {} ) -> Path: """ Quantizes a model given the optimization specifications defined in `quantization_config`. @@ -382,6 +384,7 @@ def quantize( "AddQDQPairToWeight": quantization_config.qdq_add_pair_to_weight, "DedicatedQDQPair": quantization_config.qdq_dedicated_pair, "QDQOpTypePerChannelSupportToAxis": quantization_config.qdq_op_type_per_channel_support_to_axis, + **extra_options, }, }