From fc465f94bc69b1f85d4aa4c4a269bab70d2b1b9d Mon Sep 17 00:00:00 2001 From: Seunghoon Lee Date: Tue, 23 Jan 2024 23:17:36 +0900 Subject: [PATCH] ONNX ROCm Linux. --- installer.py | 26 ++++++++++++++++++++++++++ modules/onnx.py | 6 +++++- modules/onnx_ep.py | 15 +++++---------- 3 files changed, 36 insertions(+), 11 deletions(-) diff --git a/installer.py b/installer.py index 3c9c880c3..7318cf8b3 100644 --- a/installer.py +++ b/installer.py @@ -440,6 +440,14 @@ def check_torch(): # ROCm 5.5 is oldest for PyTorch 2.1 torch_command = os.environ.get('TORCH_COMMAND', 'torch torchvision --index-url https://download.pytorch.org/whl/rocm5.5') xformers_package = os.environ.get('XFORMERS_PACKAGE', 'none') + if rocm_ver is not None: + install(os.environ.get('ONNXRUNTIME_PACKAGE', get_onnxruntime_source_for_rocm(arr)), "onnxruntime-training built with ROCm", ignore=True) + try: + import onnxruntime + if "ROCMExecutionProvider" not in onnxruntime.get_available_providers(): + log.warn('Failed to automatically install onxnruntime package for ROCm. Please manually install it if you need.') + except Exception: + pass elif allow_ipex and (args.use_ipex or shutil.which('sycl-ls') is not None or shutil.which('sycl-ls.exe') is not None or os.environ.get('ONEAPI_ROOT') is not None or os.path.exists('/opt/intel/oneapi') or os.path.exists("C:/Program Files (x86)/Intel/oneAPI") or os.path.exists("C:/oneAPI")): args.use_ipex = True # pylint: disable=attribute-defined-outside-init log.info('Intel OneAPI Toolkit detected') @@ -833,6 +841,24 @@ def get_version(): return version +def get_onnxruntime_source_for_rocm(rocm_ver): + ort_version = "1.16.3" + + try: + import onnxruntime + ort_version = onnxruntime.__version__ + except ImportError: + pass + + cp_str = f"{sys.version_info.major}{sys.version_info.minor}" + + if rocm_ver is None: + command = subprocess.run('hipconfig --version', shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + rocm_ver = command.stdout.decode(encoding="utf8", errors="ignore").split('.') + + return f"https://download.onnxruntime.ai/onnxruntime_training-{ort_version}%2Brocm{rocm_ver[0]}{rocm_ver[1]}-cp{cp_str}-cp{cp_str}-manylinux_2_17_x86_64.manylinux2014_x86_64.whl" + + # check version of the main repo and optionally upgrade it def check_version(offline=False, reset=True): # pylint: disable=unused-argument if args.skip_all: diff --git a/modules/onnx.py b/modules/onnx.py index a30f1fa5a..00af55ddc 100644 --- a/modules/onnx.py +++ b/modules/onnx.py @@ -141,7 +141,11 @@ def initialize(): if initialized: return - from modules import onnx_pipelines as pipelines + from modules import onnx_pipelines as pipelines, devices + from modules.onnx_ep import ExecutionProvider, TORCH_DEVICE_TO_EP + + if devices.backend == "rocm": + TORCH_DEVICE_TO_EP["cuda"] = ExecutionProvider.ROCm # OnnxRuntimeModel Hijack. OnnxRuntimeModel.__module__ = 'diffusers' diff --git a/modules/onnx_ep.py b/modules/onnx_ep.py index ddc696881..b50e05ee6 100644 --- a/modules/onnx_ep.py +++ b/modules/onnx_ep.py @@ -10,6 +10,7 @@ class ExecutionProvider(str, Enum): DirectML = "DmlExecutionProvider" CUDA = "CUDAExecutionProvider" ROCm = "ROCMExecutionProvider" + MIGraphX = "MIGraphXExecutionProvider" OpenVINO = "OpenVINOExecutionProvider" @@ -19,6 +20,7 @@ EP_TO_NAME = { ExecutionProvider.DirectML: "gpu-dml", ExecutionProvider.CUDA: "gpu-cuda", # test required ExecutionProvider.ROCm: "gpu-rocm", # test required + ExecutionProvider.MIGraphX: "gpu-migraphx", # test required ExecutionProvider.OpenVINO: "gpu-openvino??", # test required } TORCH_DEVICE_TO_EP = { @@ -42,7 +44,7 @@ def get_default_execution_provider() -> ExecutionProvider: if ExecutionProvider.ROCm in available_execution_providers: return ExecutionProvider.ROCm else: - log.warning("Currently, there's no pypi release for onnxruntime-rocm. Please download and install .whl file from https://download.onnxruntime.ai/") + return ExecutionProvider.CPU elif devices.backend == "ipex" or devices.backend == "openvino": return ExecutionProvider.OpenVINO return ExecutionProvider.CPU @@ -59,8 +61,6 @@ def get_execution_provider_options(): if ExecutionProvider.ROCm in available_execution_providers: execution_provider_options["tunable_op_enable"] = 1 execution_provider_options["tunable_op_tuning_enable"] = 1 - else: - log.warning("Currently, there's no pypi release for onnxruntime-rocm. Please download and install .whl file from https://download.onnxruntime.ai/ The inference will be fall back to CPU.") elif opts.onnx_execution_provider == ExecutionProvider.OpenVINO: from modules.intel.openvino import get_device as get_raw_openvino_device raw_openvino_device = get_raw_openvino_device() @@ -79,7 +79,7 @@ def get_provider() -> Tuple: def install_execution_provider(ep: ExecutionProvider): - from installer import pip, uninstall, installed + from installer import pip, uninstall, installed, get_onnxruntime_source_for_rocm if installed("onnxruntime"): uninstall("onnxruntime") @@ -103,12 +103,7 @@ def install_execution_provider(ep: ExecutionProvider): log.warn("ROCMExecutionProvider is not supported on Windows.") return - try: - major, minor = sys.version_info - cp_str = f"{major}{minor}" - packages.append(f"https://download.onnxruntime.ai/onnxruntime_training-1.16.3%2Brocm56-cp{cp_str}-cp{cp_str}-manylinux_2_17_x86_64.manylinux2014_x86_64.whl") - except Exception: - log.warn("Failed to install onnxruntime for ROCm.") + packages.append(get_onnxruntime_source_for_rocm()) elif ep == ExecutionProvider.OpenVINO: if installed("openvino"): uninstall("openvino")