Переглянути джерело

Default device aligns with paddle (#4151)

Lin Manhui 5 місяців тому
батько
коміт
95ecadaf94
2 змінених файлів з 5 додано та 20 видалено
  1. 4 19
      paddlex/utils/device.py
  2. 1 1
      setup.py

+ 4 - 19
paddlex/utils/device.py

@@ -15,8 +15,6 @@
 import os
 from contextlib import ContextDecorator
 
-import GPUtil
-
 from . import logging
 from .custom_device_list import (
     DCU_WHITELIST,
@@ -41,25 +39,12 @@ def constr_device(device_type, device_ids):
 
 
 def get_default_device():
-    try:
-        gpu_list = GPUtil.getGPUs()
-    except Exception:
-        logging.debug(
-            "Failed to query GPU devices. Falling back to CPU.", exc_info=True
-        )
-        has_gpus = False
+    import paddle
+
+    if paddle.device.is_compiled_with_cuda() and paddle.device.cuda.device_count() > 0:
+        return constr_device("gpu", [0])
     else:
-        has_gpus = bool(gpu_list)
-    if not has_gpus:
-        # HACK
-        if os.path.exists("/etc/nv_tegra_release"):
-            logging.debug(
-                "The current device appears to be an NVIDIA Jetson. GPU 0 will be used as the default device."
-            )
-    if not has_gpus:
         return "cpu"
-    else:
-        return constr_device("gpu", [0])
 
 
 def parse_device(device):

+ 1 - 1
setup.py

@@ -80,7 +80,6 @@ REQUIRED_DEPS = [
     "chardet",
     "colorlog",
     "filelock",
-    "GPUtil",
     "huggingface_hub",
     "numpy",
     "packaging",
@@ -111,6 +110,7 @@ EXTRAS = {
         "multimodal": [
             "einops",
             "ftfy",
+            "GPUtil",
             "Jinja2",
             "opencv-contrib-python",
             # For the same reason as in `cv`