Add the ability to choose the ONNX runtime execution provider in ORTModel

examples/onnxruntime/optimization/question-answering/run_qa.py

@@ -81,6 +81,10 @@ class ModelArguments:
             "with private models)."
         },
     )
+    execution_provider: str = field(
+        default="CPUExecutionProvider",
+        metadata={"help": "ONNX Runtime execution provider to use for inference."},
+    )
 
 
 @dataclass
@@ -257,6 +261,28 @@ def main():
     transformers.utils.logging.enable_default_handler()
     transformers.utils.logging.enable_explicit_format()
 
+    if (
+        optim_args.optimization_level > 1
+        and optim_args.optimize_for_gpu
+        and model_args.execution_provider == "CPUExecutionProvider"
+    ):
+        raise ValueError(
+            f"Optimization level is set at {optim_args.optimization_level} and "
+            f"GPU optimization will be done, although the CPU execution provider "
+            f"was selected. Use --execution_provider CUDAExecutionProvider."
+        )
+
+    if (
+        optim_args.optimization_level > 1
+        and optim_args.optimize_for_gpu is False
+        and model_args.execution_provider == "CUDAExecutionProvider"
+    ):
+        raise ValueError(
+            f"Optimization level is set at {optim_args.optimization_level} and "
+            f"CPU optimization will be done, although the GPU execution provider "
+            f"was selected. Remove the argument --execution_provider CUDAExecutionProvider."
+        )
+
     logger.info(f"Optimization with the following parameters {optim_args}")
 
     if os.path.isdir(training_args.output_dir) and not training_args.overwrite_output_dir:
@@ -441,7 +467,12 @@ def compute_metrics(p: EvalPrediction):
             # During Feature creation dataset samples might increase, we will select required samples again
             eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
 
-        ort_model = ORTModel(optimized_model_path, optimizer._onnx_config, compute_metrics=compute_metrics)
+        ort_model = ORTModel(
+            optimized_model_path,
+            optimizer._onnx_config,
+            execution_provider=model_args.execution_provider,
+            compute_metrics=compute_metrics,
+        )
         outputs = ort_model.evaluation_loop(eval_dataset)
         predictions = post_processing_function(eval_examples, eval_dataset, outputs.predictions)
         metrics = compute_metrics(predictions)
@@ -472,7 +503,9 @@ def compute_metrics(p: EvalPrediction):
             # During Feature creation dataset samples might increase, we will select required samples again
             predict_dataset = predict_dataset.select(range(data_args.max_predict_samples))
 
-        ort_model = ORTModel(optimized_model_path, optimizer._onnx_config)
+        ort_model = ORTModel(
+            optimized_model_path, optimizer._onnx_config, execution_provider=model_args.execution_provider
+        )
         outputs = ort_model.evaluation_loop(predict_dataset)
         predictions = post_processing_function(predict_examples, predict_dataset, outputs.predictions)
         metrics = compute_metrics(predictions)

examples/onnxruntime/optimization/text-classification/run_glue.py

@@ -130,6 +130,10 @@ class ModelArguments:
         default=None,
         metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
     )
+    execution_provider: str = field(
+        default="CPUExecutionProvider",
+        metadata={"help": "ONNX Runtime execution provider to use for inference."},
+    )
 
 
 @dataclass
@@ -195,6 +199,28 @@ def main():
     transformers.utils.logging.enable_default_handler()
     transformers.utils.logging.enable_explicit_format()
 
+    if (
+        optim_args.optimization_level > 1
+        and optim_args.optimize_for_gpu
+        and model_args.execution_provider == "CPUExecutionProvider"
+    ):
+        raise ValueError(
+            f"Optimization level is set at {optim_args.optimization_level} and "
+            f"GPU optimization will be done, although the CPU execution provider "
+            f"was selected. Use --execution_provider CUDAExecutionProvider."
+        )
+
+    if (
+        optim_args.optimization_level > 1
+        and optim_args.optimize_for_gpu is False
+        and model_args.execution_provider == "CUDAExecutionProvider"
+    ):
+        raise ValueError(
+            f"Optimization level is set at {optim_args.optimization_level} and "
+            f"CPU optimization will be done, although the GPU execution provider "
+            f"was selected. Remove the argument --execution_provider CUDAExecutionProvider."
+        )
+
     logger.info(f"Optimization with the following parameters {optim_args}")
 
     if os.path.isdir(training_args.output_dir) and not training_args.overwrite_output_dir:
@@ -354,7 +380,11 @@ def compute_metrics(p: EvalPrediction):
         )
 
         ort_model = ORTModel(
-            optimized_model_path, optimizer._onnx_config, compute_metrics=compute_metrics, label_names=["label"]
+            optimized_model_path,
+            optimizer._onnx_config,
+            execution_provider=model_args.execution_provider,
+            compute_metrics=compute_metrics,
+            label_names=["label"],
         )
         outputs = ort_model.evaluation_loop(eval_dataset)
         # Save metrics
@@ -378,7 +408,9 @@ def compute_metrics(p: EvalPrediction):
             desc="Running tokenizer on the test dataset",
         )
 
-        ort_model = ORTModel(optimized_model_path, optimizer._onnx_config)
+        ort_model = ORTModel(
+            optimized_model_path, optimizer._onnx_config, execution_provider=model_args.execution_provider
+        )
         outputs = ort_model.evaluation_loop(predict_dataset)
         predictions = np.squeeze(outputs.predictions) if is_regression else np.argmax(outputs.predictions, axis=1)
 

examples/onnxruntime/optimization/token-classification/run_ner.py

@@ -81,6 +81,10 @@ class ModelArguments:
             "with private models)."
         },
     )
+    execution_provider: str = field(
+        default="CPUExecutionProvider",
+        metadata={"help": "ONNX Runtime execution provider to use for inference."},
+    )
 
 
 @dataclass
@@ -226,6 +230,28 @@ def main():
     transformers.utils.logging.enable_default_handler()
     transformers.utils.logging.enable_explicit_format()
 
+    if (
+        optim_args.optimization_level > 1
+        and optim_args.optimize_for_gpu
+        and model_args.execution_provider == "CPUExecutionProvider"
+    ):
+        raise ValueError(
+            f"Optimization level is set at {optim_args.optimization_level} and "
+            f"GPU optimization will be done, although the CPU execution provider "
+            f"was selected. Use --execution_provider CUDAExecutionProvider."
+        )
+
+    if (
+        optim_args.optimization_level > 1
+        and optim_args.optimize_for_gpu is False
+        and model_args.execution_provider == "CUDAExecutionProvider"
+    ):
+        raise ValueError(
+            f"Optimization level is set at {optim_args.optimization_level} and "
+            f"CPU optimization will be done, although the GPU execution provider "
+            f"was selected. Remove the argument --execution_provider CUDAExecutionProvider."
+        )
+
     logger.info(f"Optimization with the following parameters {optim_args}")
 
     if os.path.isdir(training_args.output_dir) and not training_args.overwrite_output_dir:
@@ -425,7 +451,12 @@ def compute_metrics(p):
             desc="Running tokenizer on the validation dataset",
         )
 
-        ort_model = ORTModel(optimized_model_path, optimizer._onnx_config, compute_metrics=compute_metrics)
+        ort_model = ORTModel(
+            optimized_model_path,
+            optimizer._onnx_config,
+            execution_provider=model_args.execution_provider,
+            compute_metrics=compute_metrics,
+        )
         outputs = ort_model.evaluation_loop(eval_dataset)
 
         # Save evaluation metrics
@@ -450,7 +481,12 @@ def compute_metrics(p):
             desc="Running tokenizer on the prediction dataset",
         )
 
-        ort_model = ORTModel(optimized_model_path, optimizer._onnx_config, compute_metrics=compute_metrics)
+        ort_model = ORTModel(
+            optimized_model_path,
+            optimizer._onnx_config,
+            execution_provider=model_args.execution_provider,
+            compute_metrics=compute_metrics,
+        )
         outputs = ort_model.evaluation_loop(predict_dataset)
         predictions = np.argmax(outputs.predictions, axis=2)
 

examples/onnxruntime/quantization/question-answering/run_qa.py

@@ -88,6 +88,10 @@ class ModelArguments:
             "with private models)."
         },
     )
+    execution_provider: str = field(
+        default="CPUExecutionProvider",
+        metadata={"help": "ONNX Runtime execution provider to use for inference."},
+    )
 
 
 @dataclass
@@ -651,7 +655,12 @@ def compute_metrics(p: EvalPrediction):
     if training_args.do_eval:
         logger.info("*** Evaluate ***")
 
-        ort_model = ORTModel(quantized_model_path, quantizer._onnx_config, compute_metrics=compute_metrics)
+        ort_model = ORTModel(
+            quantized_model_path,
+            quantizer._onnx_config,
+            execution_provider=model_args.execution_provider,
+            compute_metrics=compute_metrics,
+        )
         outputs = ort_model.evaluation_loop(eval_dataset)
         predictions = post_processing_function(eval_examples, eval_dataset, outputs.predictions)
         metrics = compute_metrics(predictions)
@@ -664,7 +673,9 @@ def compute_metrics(p: EvalPrediction):
     if training_args.do_predict:
         logger.info("*** Predict ***")
 
-        ort_model = ORTModel(quantized_model_path, quantizer._onnx_config)
+        ort_model = ORTModel(
+            quantized_model_path, quantizer._onnx_config, execution_provider=model_args.execution_provider
+        )
         outputs = ort_model.evaluation_loop(predict_dataset)
         predictions = post_processing_function(predict_examples, predict_dataset, outputs.predictions)
         metrics = compute_metrics(predictions)

examples/onnxruntime/quantization/text-classification/run_glue.py

@@ -151,6 +151,10 @@ class ModelArguments:
         default=None,
         metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
     )
+    execution_provider: str = field(
+        default="CPUExecutionProvider",
+        metadata={"help": "ONNX Runtime execution provider to use for inference."},
+    )
 
 
 @dataclass
@@ -476,7 +480,11 @@ def compute_metrics(p: EvalPrediction):
             eval_dataset = eval_dataset.select(range(data_args.max_eval_samples))
 
         ort_model = ORTModel(
-            quantized_model_path, quantizer._onnx_config, compute_metrics=compute_metrics, label_names=["label"]
+            quantized_model_path,
+            quantizer._onnx_config,
+            execution_provider=model_args.execution_provider,
+            compute_metrics=compute_metrics,
+            label_names=["label"],
         )
         outputs = ort_model.evaluation_loop(eval_dataset)
         # Save metrics
@@ -493,7 +501,9 @@ def compute_metrics(p: EvalPrediction):
         if data_args.max_predict_samples is not None:
             predict_dataset = predict_dataset.select(range(data_args.max_predict_samples))
 
-        ort_model = ORTModel(quantized_model_path, quantizer._onnx_config)
+        ort_model = ORTModel(
+            quantized_model_path, quantizer._onnx_config, execution_provider=model_args.execution_provider
+        )
         outputs = ort_model.evaluation_loop(predict_dataset)
         predictions = np.squeeze(outputs.predictions) if is_regression else np.argmax(outputs.predictions, axis=1)
 

examples/onnxruntime/quantization/token-classification/run_ner.py

@@ -87,6 +87,10 @@ class ModelArguments:
             "with private models)."
         },
     )
+    execution_provider: str = field(
+        default="CPUExecutionProvider",
+        metadata={"help": "ONNX Runtime execution provider to use for inference."},
+    )
 
 
 @dataclass
@@ -542,7 +546,12 @@ def compute_metrics(p):
                 desc="Running tokenizer on the validation dataset",
             )
 
-        ort_model = ORTModel(quantized_model_path, quantizer._onnx_config, compute_metrics=compute_metrics)
+        ort_model = ORTModel(
+            quantized_model_path,
+            quantizer._onnx_config,
+            execution_provider=model_args.execution_provider,
+            compute_metrics=compute_metrics,
+        )
         outputs = ort_model.evaluation_loop(eval_dataset)
 
         # Save evaluation metrics
@@ -568,7 +577,12 @@ def compute_metrics(p):
                 desc="Running tokenizer on the prediction dataset",
             )
 
-        ort_model = ORTModel(quantized_model_path, quantizer._onnx_config, compute_metrics=compute_metrics)
+        ort_model = ORTModel(
+            quantized_model_path,
+            quantizer._onnx_config,
+            execution_provider=model_args.execution_provider,
+            compute_metrics=compute_metrics,
+        )
         outputs = ort_model.evaluation_loop(predict_dataset)
         predictions = np.argmax(outputs.predictions, axis=2)
 

optimum/onnxruntime/model.py

@@ -37,6 +37,7 @@ def __init__(
         self,
         model_path: Union[str, os.PathLike],
         onnx_config: OnnxConfig,
+        execution_provider: Optional[str] = "CPUExecutionProvider",
         compute_metrics: Optional[Callable[[EvalPrediction], Dict]] = None,
         label_names: Optional[List[str]] = None,
     ):
@@ -47,6 +48,8 @@ def __init__(
             onnx_config (`OnnxConfig`):
                 An ONNX configuration associated to the ONNX model describing metadata on how to export the model
                 through the ONNX format.
+            execution_provider (:obj:`str`, `optional`):
+                ONNX Runtime execution provider to use.
             compute_metrics (`Callable[[EvalPrediction], Dict]`, `optional`):
                 The function that will be used to compute metrics at evaluation. Must take an `EvalPrediction` and
                 return a dictionary string to metric values.
@@ -62,6 +65,7 @@ def __init__(
         self.onnx_named_inputs = list(onnx_config.inputs.keys())
         self.onnx_named_outputs = list(onnx_config.outputs.keys())
         self.onnx_config = onnx_config
+        self.execution_provider = execution_provider
         self.model_path = Path(model_path)
         self.compute_metrics = compute_metrics
         default_label_names = (
@@ -81,7 +85,7 @@ def evaluation_loop(self, dataset: Dataset):
         all_preds = None
         all_labels = None
         options = SessionOptions()
-        session = InferenceSession(self.model_path.as_posix(), options)
+        session = InferenceSession(self.model_path.as_posix(), options, providers=[self.execution_provider])
         for step, inputs in enumerate(dataset):
             has_labels = all(inputs.get(k) is not None for k in self.label_names)
             if has_labels: