Pass providers directly in aimet-onnx QuantSim init

Signed-off-by: Michael Tuttle <quic_mtuttle@quicinc.com>
Co-authored-by: Michael Tuttle <quic_mtuttle@quicinc.com>

Author: 2 people

Examples/onnx/quantization/AMP.ipynb

@@ -235,10 +235,8 @@
     "# cudnn_conv_algo_search is fixing it to default to avoid changing in accuracies/outputs at every inference\n",
     "if 'CUDAExecutionProvider' in ort.get_available_providers():\n",
     "    providers = [('CUDAExecutionProvider', {'cudnn_conv_algo_search': 'DEFAULT'}), 'CPUExecutionProvider']\n",
-    "    use_cuda = True\n",
     "else:\n",
-    "    providers = ['CPUExecutionProvider']\n",
-    "    use_cuda = False"
+    "    providers = ['CPUExecutionProvider']"
    ]
   },
   {
@@ -341,7 +339,7 @@
     "                           quant_scheme=QuantScheme.post_training_tf_enhanced,\n",
     "                           default_activation_bw=8,\n",
     "                           default_param_bw=8,\n",
-    "                           use_cuda=use_cuda)"
+    "                           providers=providers)"
    ]
   },
   {

Examples/onnx/quantization/adaround.ipynb

@@ -221,10 +221,8 @@
     "# cudnn_conv_algo_search is fixing it to default to avoid changing in accuracies/outputs at every inference\n",
     "if 'CUDAExecutionProvider' in ort.get_available_providers():\n",
     "    providers = [('CUDAExecutionProvider', {'cudnn_conv_algo_search': 'DEFAULT'}), 'CPUExecutionProvider']\n",
-    "    use_cuda = True\n",
     "else:\n",
-    "    providers = ['CPUExecutionProvider']\n",
-    "    use_cuda = False"
+    "    providers = ['CPUExecutionProvider']"
    ]
   },
   {
@@ -306,7 +304,7 @@
     "                           quant_scheme=QuantScheme.post_training_tf_enhanced,\n",
     "                           default_activation_bw=8,\n",
     "                           default_param_bw=8,\n",
-    "                           use_cuda=use_cuda)"
+    "                           providers=providers)"
    ]
   },
   {
@@ -485,7 +483,7 @@
     "                           quant_scheme=QuantScheme.post_training_tf_enhanced,\n",
     "                           default_activation_bw=8,\n",
     "                           default_param_bw=8,\n",
-    "                           use_cuda=use_cuda)\n",
+    "                           providers=providers)\n",
     "\n",
     "sim.set_and_freeze_param_encodings(encoding_path=os.path.join(\"output\", 'adaround.encodings'))\n",
     "\n",

Examples/onnx/quantization/cle.ipynb

@@ -256,10 +256,8 @@
     "# cudnn_conv_algo_search is fixing it to default to avoid changing in accuracies/outputs at every inference\n",
     "if 'CUDAExecutionProvider' in ort.get_available_providers():\n",
     "    providers = [('CUDAExecutionProvider', {'cudnn_conv_algo_search': 'DEFAULT'}), 'CPUExecutionProvider']\n",
-    "    use_cuda = True\n",
     "else:\n",
-    "    providers = ['CPUExecutionProvider']\n",
-    "    use_cuda = False"
+    "    providers = ['CPUExecutionProvider']"
    ]
   },
   {
@@ -355,7 +353,7 @@
     "                           quant_scheme=QuantScheme.post_training_tf_enhanced,\n",
     "                           default_activation_bw=8,\n",
     "                           default_param_bw=8,\n",
-    "                           use_cuda=use_cuda)"
+    "                           providers=providers)"
    ]
   },
   {
@@ -556,7 +554,7 @@
     "                           quant_scheme=QuantScheme.post_training_tf_enhanced,\n",
     "                           default_activation_bw=8,\n",
     "                           default_param_bw=8,\n",
-    "                           use_cuda=use_cuda)\n",
+    "                           providers=providers)\n",
     "\n",
     "sim.compute_encodings(forward_pass_callback=pass_calibration_data,\n",
     "                      forward_pass_callback_args=1000)\n",

Examples/onnx/quantization/quantsim.ipynb

@@ -248,11 +248,9 @@
    "source": [
     "# cudnn_conv_algo_search is fixing it to default to avoid changing in accuracies/outputs at every inference\n",
     "if 'CUDAExecutionProvider' in ort.get_available_providers():\n",
-    "    providers = [('CUDAExecutionProvider', {'cudnn_conv_algo_search': 'DEFAULT'}), 'CPUExecutionProvider']\n",
-    "    use_cuda = True\n",
+    "    providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']\n",
     "else:\n",
-    "    providers = ['CPUExecutionProvider']\n",
-    "    use_cuda = False"
+    "    providers = ['CPUExecutionProvider']"
    ]
   },
   {
@@ -353,7 +351,7 @@
     "                           quant_scheme=QuantScheme.post_training_tf_enhanced,\n",
     "                           default_activation_bw=8,\n",
     "                           default_param_bw=8,\n",
-    "                           use_cuda=use_cuda)"
+    "                           providers=providers)"
    ]
   },
   {

TrainingExtensions/onnx/src/python/aimet_onnx/quantsim.py

@@ -41,7 +41,7 @@
 import tempfile
 from pathlib import Path
 import os
-from typing import Any, Callable, Dict, List, Optional, overload, Tuple, TypeVar, Union
+from typing import Any, Callable, Dict, List, Optional, overload, Tuple, TypeVar, Union, Sequence
 import itertools
 import json
 import warnings
@@ -160,39 +160,92 @@ class QuantizationSimModel:
     :param model: ONNX model
     :param dummy_input: Dummy input to the model. If None, will attempt to auto-generate a dummy input
     :param quant_scheme: Quantization scheme (e.g. QuantScheme.post_training_tf)
-    :param rounding_mode: Rounding mode (e.g. nearest)
+    :param rounding_mode: Deprecated
     :param default_param_bw: Quantization bitwidth for parameter
     :param default_activation_bw: Quantization bitwidth for activation
-    :param use_symmetric_encodings: True if symmetric encoding is used.  False otherwise.
-    :param use_cuda: True if using CUDA to run quantization op. False otherwise.
+    :param use_symmetric_encodings: Deprecated, symmetry is controlled by the config_file
+    :param use_cuda: Deprecated, use `providers` instead
     :param config_file: File path or alias of the configuration file.
                         Alias can be one of {{ {', '.join(_config_file_aliases.keys())} }} (Default: `"default"`)
     :param default_data_type: Default data type to use for quantizing all layer inputs, outputs and parameters.
                              Possible options are QuantizationDataType.int and QuantizationDataType.float.
                              Note that the mode default_data_type=QuantizationDataType.float is only supported with
                              default_output_bw=16 and default_param_bw=16
     :param user_onnx_libs: List of paths to all compiled ONNX custom ops libraries
+    :param providers: Onnxruntime execution providers to use when building InferenceSession. 
+                      If `None`, falls back to `onnxruntime.get_available_providers()`
     :param path: Directory to save the artifacts.
     """
 
     def __init__(self,
                  model: Union[ModelProto, ONNXModel],
-                 dummy_input: Dict[str, np.ndarray] = None,
+                 dummy_input: Optional[Dict[str, np.ndarray]] = None,
                  quant_scheme: QuantScheme = QuantScheme.min_max,
-                 rounding_mode: str = 'nearest',
+                 rounding_mode: str = None, # Deprecated
                  default_param_bw: int = 8,
                  default_activation_bw: int = 8,
-                 use_symmetric_encodings: bool = False, use_cuda: bool = True,
-                 device: int = 0, config_file: str = None,
+                 use_symmetric_encodings: bool = None, # Deprecated
+                 use_cuda: bool = None, # Deprecated
+                 device: int = None, # Deprecated
+                 config_file: Optional[str] = None,
                  default_data_type: QuantizationDataType = QuantizationDataType.int,
-                 user_onnx_libs: List[str] = None, path: str = None):
+                 user_onnx_libs: List[str] = None,
+                 providers: Optional[Sequence[str | Tuple[str, Dict[Any, Any]]]] = None,
+                 path: Optional[str] = None):
+        # pylint: disable = too-many-branches, too-many-statements
+        if rounding_mode is not None:
+            if rounding_mode == 'nearest':
+                warnings.warn(_red("Passing rounding_mode='nearest' is no longer needed " \
+                                   "and will be deprecated soon in the later versions."),
+                              DeprecationWarning, stacklevel=2)
+            else:
+                raise TypeError("'rounding_mode' parameter is no longer supported.")
+
+        if use_symmetric_encodings is not None:
+            warnings.warn(_red("Passing `use_symmetric_encodings` is not needed and will be deprecated in later versions."),
+                          DeprecationWarning, stacklevel=2)
+
+        if device is not None:
+            warnings.warn(_red("Passing `device` will be deprecated in later versions. " \
+                               "Please use the `providers` argument instead to specify cuda device."),
+                          DeprecationWarning, stacklevel=2)
+            if providers is not None:
+                raise RuntimeError("Cannot provide `device` and `providers` at the same time.")
+
+        if use_cuda is not None:
+            warnings.warn(_red("Passing `use_cuda` will be deprecated in later versions. " \
+                               "Please use the `providers` argument instead."),
+                          DeprecationWarning, stacklevel=2)
+            if providers is not None:
+                raise RuntimeError("Cannot provide `use_cuda` and `providers` at the same time.")
+
+            # Legacy behavior of use_cuda
+            if "CUDAExecutionProvider" not in ort.get_available_providers():
+                use_cuda = False
+
+            device = device or 0
+            if use_cuda:
+                providers = [('CUDAExecutionProvider', {'device_id': device}), 'CPUExecutionProvider']
+            else:
+                providers = ['CPUExecutionProvider']
+
+        if not providers:
+            providers = ort.get_available_providers()
+
         if isinstance(quant_scheme, str):
             quant_scheme = QuantScheme.from_str(quant_scheme)
 
         if isinstance(model, ModelProto):
             model = ONNXModel(model)
 
+        op_domain = "aimet.customop.cpu"
+        for provider in providers:
+            if provider == "CUDAExecutionProvider" or provider[0] == "CUDAExecutionProvider":
+                op_domain = "aimet.customop.cuda"
+
         self.model = model
+        self._op_domain = op_domain
+        self.providers = providers
 
         if not dummy_input:
             dummy_input = make_dummy_input(self.model.model)
@@ -204,16 +257,6 @@ def __init__(self,
         self._default_param_bw = default_param_bw
         self._default_activation_bw = default_activation_bw
         self._default_quantization_data_type = default_data_type
-        self._use_symmetric_encodings = use_symmetric_encodings
-        self._use_cuda = use_cuda
-        if 'CUDAExecutionProvider' not in ort.get_available_providers():
-            self._use_cuda = False
-        if self._use_cuda:
-            self._op_domain = "aimet.customop.cuda"
-            self.providers = [('CUDAExecutionProvider', {'device_id': device, 'cudnn_conv_algo_search': 'DEFAULT'}), 'CPUExecutionProvider']
-        else:
-            self._op_domain = "aimet.customop.cpu"
-            self.providers = ['CPUExecutionProvider']
         self._user_onnx_libs = user_onnx_libs
         self.param_names = []
         self.input_quantizers_name = []
@@ -465,7 +508,7 @@ def _insert_param_quantization_nodes(self):
                                                           rounding_mode=self._rounding_mode,
                                                           op_mode=OpMode.oneShotQuantizeDequantize,
                                                           bitwidth=self._default_param_bw,
-                                                          use_symmetric_encodings=self._use_symmetric_encodings,
+                                                          use_symmetric_encodings=False,
                                                           tensor_quantizer_params=tensor_quantizer_params)
 
     def _create_quant_info_object_for_param(self, param_name: str):
@@ -533,7 +576,7 @@ def _insert_activation_quantization_nodes(self):
                                                           rounding_mode=self._rounding_mode,
                                                           op_mode=OpMode.updateStats,
                                                           bitwidth=self._default_activation_bw,
-                                                          use_symmetric_encodings=self._use_symmetric_encodings)
+                                                          use_symmetric_encodings=False)
 
     @staticmethod
     def build_session(model: onnx.ModelProto, providers: List, user_onnx_libs: List[str] = None, path: str = None):

TrainingExtensions/onnx/test/python/test_auto_quant_v2.py

@@ -249,11 +249,16 @@ def __init__(self,
                      rounding_mode: str = 'nearest',
                      default_param_bw: int = 8,
                      default_activation_bw: int = 8,
-                     use_symmetric_encodings: bool = False, use_cuda: bool = True,
-                     device: int = 0, config_file: str = None, default_data_type: QuantizationDataType = QuantizationDataType.int,
-                     user_onnx_libs: List[str] = None):
+                     use_symmetric_encodings: bool = None, # Deprecated
+                     use_cuda: bool = None, # Deprecated
+                     device: int = None, # Deprecated
+                     config_file = None,
+                     default_data_type: QuantizationDataType = QuantizationDataType.int,
+                     user_onnx_libs: List[str] = None,
+                     providers = None,
+                     path = None):
             super(_QuantizationSimModel, self).__init__(model, dummy_input, quant_scheme, rounding_mode, default_param_bw, default_activation_bw,
-                                                        use_symmetric_encodings, use_cuda, device, config_file, default_data_type, user_onnx_libs)
+                                                        use_symmetric_encodings, use_cuda, device, config_file, default_data_type, user_onnx_libs, providers, path)
 
             self.session = {'applied_bn_folding': getattr(model, 'applied_bn_folding'),
                             'applied_cle': getattr(model, 'applied_cle'),

TrainingExtensions/onnx/test/python/test_auto_quant_v2_with_amp.py

@@ -288,11 +288,16 @@ def __init__(self,
                  rounding_mode: str = 'nearest',
                  default_param_bw: int = 8,
                  default_activation_bw: int = 8,
-                 use_symmetric_encodings: bool = False, use_cuda: bool = True,
-                 device: int = 0, config_file: str = None, default_data_type: QuantizationDataType = QuantizationDataType.int,
-                 user_onnx_libs: List[str] = None):
+                 use_symmetric_encodings: bool = None, # Deprecated
+                 use_cuda: bool = None, # Deprecated
+                 device: int = None, # Deprecated
+                 config_file = None,
+                 default_data_type: QuantizationDataType = QuantizationDataType.int,
+                 user_onnx_libs: List[str] = None,
+                 providers = None,
+                 path = None):
         super(_QuantizationSimModel, self).__init__(model, dummy_input, quant_scheme, rounding_mode, default_param_bw, default_activation_bw,
-                                                    use_symmetric_encodings, use_cuda, device, config_file, default_data_type, user_onnx_libs)
+                                                    use_symmetric_encodings, use_cuda, device, config_file, default_data_type, user_onnx_libs, providers, path)
 
         self.session = {'applied_bn_folding': getattr(model, 'applied_bn_folding'),
                         'applied_cle': getattr(model, 'applied_cle'),

TrainingExtensions/onnx/test/python/test_layer_output_utils.py

@@ -71,7 +71,7 @@ def get_quantsim_artifacts():
     def callback(session, input_dict):
         session.run(None, input_dict)
 
-    quantsim = QuantizationSimModel(model=model, dummy_input=input_dict, use_cuda=False)
+    quantsim = QuantizationSimModel(model=model, dummy_input=input_dict, providers=["CPUExecutionProvider"])
     quantsim.compute_encodings(callback, input_dict)
 
     output_names = [node.name for node in quantsim.model.model.graph.input]