From b53e110921074e6860520c0b6cbab3c4e7441534 Mon Sep 17 00:00:00 2001 From: brokenhammer <415516853@qq.com> Date: Sun, 22 Dec 2024 17:04:59 -0800 Subject: [PATCH 1/5] enable selecting gpu --- desc/__init__.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/desc/__init__.py b/desc/__init__.py index 840b9985b..be8b0c27f 100644 --- a/desc/__init__.py +++ b/desc/__init__.py @@ -61,7 +61,7 @@ def __getattr__(name): config = {"device": None, "avail_mem": None, "kind": None} -def set_device(kind="cpu"): +def set_device(kind="cpu",gpuid=None): """Sets the device to use for computation. If kind==``'gpu'``, checks available GPUs and selects the one with the most @@ -127,11 +127,15 @@ def set_device(kind="cpu"): set_device(kind="cpu") return devices = [dev for dev in devices if dev["index"] in gpu_ids] - for dev in devices: - mem = dev["mem_total"] - dev["mem_used"] - if mem > maxmem: - maxmem = mem - selected_gpu = dev + + if (not (gpuid == None)) and (str(gpuid) in gpu_ids): + selected_gpu = [dev for dev in devices if dev["index"] == str(gpuid)][0] + else: + for dev in devices: + mem = dev["mem_total"] - dev["mem_used"] + if mem > maxmem: + maxmem = mem + selected_gpu = dev config["device"] = selected_gpu["type"] + " (id={})".format( selected_gpu["index"] ) From 23e06070534c8e0457dcaf763e05d2035784236c Mon Sep 17 00:00:00 2001 From: Xishuo Wei <415516853@qq.com> Date: Mon, 23 Dec 2024 12:54:18 -0800 Subject: [PATCH 2/5] Update desc/__init__.py formatting Co-authored-by: Yigit Gunsur Elmacioglu <102380275+YigitElma@users.noreply.github.com> --- desc/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/__init__.py b/desc/__init__.py index be8b0c27f..010aa4beb 100644 --- a/desc/__init__.py +++ b/desc/__init__.py @@ -61,7 +61,7 @@ def __getattr__(name): config = {"device": None, "avail_mem": None, "kind": None} -def set_device(kind="cpu",gpuid=None): +def set_device(kind="cpu", gpuid=None): """Sets the device to use for computation. If kind==``'gpu'``, checks available GPUs and selects the one with the most From 60963e84aee423c8d41ed432030867b62c2061c0 Mon Sep 17 00:00:00 2001 From: brokenhammer <415516853@qq.com> Date: Mon, 23 Dec 2024 13:24:56 -0800 Subject: [PATCH 3/5] Update docs; print warning message for wrong gpuid. --- desc/__init__.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/desc/__init__.py b/desc/__init__.py index 010aa4beb..4fbc171a5 100644 --- a/desc/__init__.py +++ b/desc/__init__.py @@ -64,8 +64,9 @@ def __getattr__(name): def set_device(kind="cpu", gpuid=None): """Sets the device to use for computation. - If kind==``'gpu'``, checks available GPUs and selects the one with the most - available memory. + If kind==``'gpu'`` and a gpuid is specified, uses the specified GPU. If + gpud==``None`` or a wrong GPU id is given, checks available GPUs and selects the + one with the most available memory. Respects environment variable CUDA_VISIBLE_DEVICES for selecting from multiple available GPUs @@ -128,7 +129,7 @@ def set_device(kind="cpu", gpuid=None): return devices = [dev for dev in devices if dev["index"] in gpu_ids] - if (not (gpuid == None)) and (str(gpuid) in gpu_ids): + if gpuid is not None and (str(gpuid) in gpu_ids): selected_gpu = [dev for dev in devices if dev["index"] == str(gpuid)][0] else: for dev in devices: @@ -139,6 +140,9 @@ def set_device(kind="cpu", gpuid=None): config["device"] = selected_gpu["type"] + " (id={})".format( selected_gpu["index"] ) + if gpuid is not None and not (str(gpuid) in gpu_ids): + warnings.warn(colored("Specified gpuid {} not found, falling back to ".format( + str(gpuid)) + config["device"], "yellow")) config["avail_mem"] = ( selected_gpu["mem_total"] - selected_gpu["mem_used"] ) / 1024 # in GB From a3ddbaaed131db5a040e19c2ea5654688c9c8de8 Mon Sep 17 00:00:00 2001 From: Xishuo Wei <415516853@qq.com> Date: Mon, 23 Dec 2024 15:55:20 -0800 Subject: [PATCH 4/5] fix typo --- desc/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/__init__.py b/desc/__init__.py index 4fbc171a5..d2537bd6a 100644 --- a/desc/__init__.py +++ b/desc/__init__.py @@ -65,7 +65,7 @@ def set_device(kind="cpu", gpuid=None): """Sets the device to use for computation. If kind==``'gpu'`` and a gpuid is specified, uses the specified GPU. If - gpud==``None`` or a wrong GPU id is given, checks available GPUs and selects the + gpuid==``None`` or a wrong GPU id is given, checks available GPUs and selects the one with the most available memory. Respects environment variable CUDA_VISIBLE_DEVICES for selecting from multiple available GPUs From f90973b57d008722543d193a8f45edcf86277c22 Mon Sep 17 00:00:00 2001 From: brokenhammer <415516853@qq.com> Date: Mon, 23 Dec 2024 21:30:43 -0800 Subject: [PATCH 5/5] apply black formatter --- desc/__init__.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/desc/__init__.py b/desc/__init__.py index d2537bd6a..2e5f9003f 100644 --- a/desc/__init__.py +++ b/desc/__init__.py @@ -64,7 +64,7 @@ def __getattr__(name): def set_device(kind="cpu", gpuid=None): """Sets the device to use for computation. - If kind==``'gpu'`` and a gpuid is specified, uses the specified GPU. If + If kind==``'gpu'`` and a gpuid is specified, uses the specified GPU. If gpuid==``None`` or a wrong GPU id is given, checks available GPUs and selects the one with the most available memory. Respects environment variable CUDA_VISIBLE_DEVICES for selecting from multiple @@ -141,8 +141,13 @@ def set_device(kind="cpu", gpuid=None): selected_gpu["index"] ) if gpuid is not None and not (str(gpuid) in gpu_ids): - warnings.warn(colored("Specified gpuid {} not found, falling back to ".format( - str(gpuid)) + config["device"], "yellow")) + warnings.warn( + colored( + "Specified gpuid {} not found, falling back to ".format(str(gpuid)) + + config["device"], + "yellow", + ) + ) config["avail_mem"] = ( selected_gpu["mem_total"] - selected_gpu["mem_used"] ) / 1024 # in GB