Add option to specify the EP to use, enabling DML EP and others (#17490)

### Description
Add DML EP to the acceptable provider list in the optimizer.

### Motivation and Context
With DML EP, graph optimization was not performed in onnxruntime.
This commit is contained in:
Kaz Nishimura 2023-10-03 15:53:09 +09:00 committed by GitHub
parent 451c02543a
commit d11e053412
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 54 additions and 8 deletions

View file

@ -144,6 +144,7 @@ def _optimize_sd_pipeline(
opt_level=0,
optimization_options=fusion_options,
use_gpu=True,
provider=args.provider,
)
if float16:
@ -168,6 +169,7 @@ def _optimize_sd_pipeline(
optimize_by_onnxruntime(
str(tmp_model_path),
use_gpu=True,
provider=args.provider,
optimized_model_path=str(ort_optimized_model_path),
save_as_external_data=use_external_data_format,
)
@ -324,6 +326,14 @@ def parse_arguments(argv: Optional[List[str]] = None):
)
parser.set_defaults(use_external_data_format=None)
parser.add_argument(
"--provider",
required=False,
type=str,
default=None,
help="Execution provider to use.",
)
FusionOptions.add_arguments(parser)
args = parser.parse_args(argv)

View file

@ -69,6 +69,8 @@ def optimize_by_onnxruntime(
save_as_external_data: bool = False,
external_data_filename: str = "",
external_data_file_threshold: int = 1024,
*,
provider: Optional[str] = None,
) -> str:
"""
Use onnxruntime to optimize model.
@ -82,6 +84,7 @@ def optimize_by_onnxruntime(
save_as_external_data (bool): whether to save external data outside of ONNX model
external_data_filename (str): name of external data file. If not provided, name is automatically created from ONNX model.
external_data_file_threshold (int): threshold to decide whether to save tensor in ONNX model or in external data file
provider (str or None): execution provider to use if use_gpu
Returns:
optimized_model_path (str): the path of optimized model
"""
@ -90,8 +93,12 @@ def optimize_by_onnxruntime(
import onnxruntime
if use_gpu and set(onnxruntime.get_available_providers()).isdisjoint(
["CUDAExecutionProvider", "ROCMExecutionProvider", "MIGraphXExecutionProvider"]
if (
use_gpu
and provider is None
and set(onnxruntime.get_available_providers()).isdisjoint(
["CUDAExecutionProvider", "ROCMExecutionProvider", "MIGraphXExecutionProvider"]
)
):
logger.error("There is no gpu for onnxruntime to do optimization.")
return onnx_model_path
@ -138,17 +145,32 @@ def optimize_by_onnxruntime(
kwargs["disabled_optimizers"] = disabled_optimizers
if not use_gpu:
onnxruntime.InferenceSession(onnx_model_path, sess_options, providers=["CPUExecutionProvider"], **kwargs)
providers = ["CPUExecutionProvider"]
elif provider is not None:
if provider == "dml":
providers = ["DmlExecutionProvider"]
elif provider == "rocm":
providers = ["ROCMExecutionProvider"]
elif provider == "migraphx":
providers = ["MIGraphXExecutionProvider", "ROCMExecutionProvider"]
elif provider == "cuda":
providers = ["CUDAExecutionProvider"]
elif provider == "tensorrt":
providers = ["TensorrtExecutionProvider", "CUDAExecutionProvider"]
else:
providers = ["CUDAExecutionProvider"]
providers.append("CPUExecutionProvider")
else:
gpu_ep = []
providers = []
if torch_version.hip:
gpu_ep.append("MIGraphXExecutionProvider")
gpu_ep.append("ROCMExecutionProvider")
providers.append("MIGraphXExecutionProvider")
providers.append("ROCMExecutionProvider")
else:
gpu_ep.append("CUDAExecutionProvider")
providers.append("CUDAExecutionProvider")
onnxruntime.InferenceSession(onnx_model_path, sess_options, providers=gpu_ep, **kwargs)
onnxruntime.InferenceSession(onnx_model_path, sess_options, providers=providers, **kwargs)
assert os.path.exists(optimized_model_path) and os.path.isfile(optimized_model_path)
logger.debug("Save optimized model by onnxruntime to %s", optimized_model_path)
@ -220,6 +242,8 @@ def optimize_model(
use_gpu: bool = False,
only_onnxruntime: bool = False,
verbose: bool = False,
*,
provider: Optional[str] = None,
):
"""Optimize Model by OnnxRuntime and/or python fusion logic.
@ -257,6 +281,7 @@ def optimize_model(
use_gpu (bool, optional): use gpu or not for onnxruntime. Defaults to False.
only_onnxruntime (bool, optional): only use onnxruntime to optimize model, and no python fusion.
Defaults to False.
provider (str, optional): execution provider to use if use_gpu. Defaults to None.
Returns:
object of an optimizer class.
@ -302,6 +327,7 @@ def optimize_model(
temp_model_path = optimize_by_onnxruntime(
input,
use_gpu=use_gpu,
provider=provider,
optimized_model_path=optimized_model_path,
opt_level=opt_level,
disabled_optimizers=disabled_optimizers,
@ -316,6 +342,7 @@ def optimize_model(
temp_model_path = optimize_by_onnxruntime(
input,
use_gpu=use_gpu,
provider=provider,
optimized_model_path=optimized_model_path,
opt_level=1,
disabled_optimizers=disabled_optimizers,
@ -423,6 +450,14 @@ def _parse_arguments():
)
parser.set_defaults(use_gpu=False)
parser.add_argument(
"--provider",
required=False,
type=str,
default=None,
help="Execution provider to use if use_gpu",
)
parser.add_argument(
"--only_onnxruntime",
required=False,
@ -501,6 +536,7 @@ def main():
opt_level=args.opt_level,
optimization_options=optimization_options,
use_gpu=args.use_gpu,
provider=args.provider,
only_onnxruntime=args.only_onnxruntime,
)