mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-08 17:17:15 +00:00
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:
parent
451c02543a
commit
d11e053412
2 changed files with 54 additions and 8 deletions
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
|
|
|||
Loading…
Reference in a new issue