paddlex 3.0.0rc1__py3-none-any.whl → 3.0.2__py3-none-any.whl

paddlex/inference/utils/hpi.py

@@ -17,6 +17,7 @@ import importlib.resources
 import importlib.util
 import json
 import platform
+from collections import defaultdict
 from functools import lru_cache
 from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
@@ -30,6 +31,7 @@ from ...utils.env import (
     get_paddle_version,
 )
 from ...utils.flags import USE_PIR_TRT
+from .misc import is_mkldnn_available
 from .model_paths import ModelPaths
 
 
@@ -163,8 +165,12 @@ def suggest_inference_backend_and_config(
         # TODO: Is it better to also check the runtime versions of CUDA and
         # cuDNN, and the versions of CUDA and cuDNN used to build `ultra-infer`?
         cuda_version = get_paddle_cuda_version()
+        if not cuda_version:
+            return None, "No CUDA version was found."
         cuda_version = "".join(map(str, cuda_version))
         cudnn_version = get_paddle_cudnn_version()
+        if not cudnn_version:
+            return None, "No cuDNN version was found."
         cudnn_version = "".join(map(str, cudnn_version[:-1]))
         key = f"gpu_cuda{cuda_version}_cudnn{cudnn_version}"
     else:
@@ -182,24 +188,23 @@ def suggest_inference_backend_and_config(
         hpi_config.pdx_model_name
     ].copy()
 
+    if not is_mkldnn_available():
+        if "paddle_mkldnn" in supported_pseudo_backends:
+            supported_pseudo_backends.remove("paddle_mkldnn")
+
     # XXX
     if not (
         USE_PIR_TRT
         and importlib.util.find_spec("tensorrt")
         and ctypes.util.find_library("nvinfer")
     ):
-        if (
-            "paddle_tensorrt" in supported_pseudo_backends
-            or "paddle_tensorrt_fp16" in supported_pseudo_backends
-        ):
-            supported_pseudo_backends.append("paddle")
         if "paddle_tensorrt" in supported_pseudo_backends:
             supported_pseudo_backends.remove("paddle_tensorrt")
         if "paddle_tensorrt_fp16" in supported_pseudo_backends:
             supported_pseudo_backends.remove("paddle_tensorrt_fp16")
 
-    candidate_backends = []
-    backend_to_pseudo_backend = {}
+    supported_backends = []
+    backend_to_pseudo_backends = defaultdict(list)
     for pb in supported_pseudo_backends:
         if pb.startswith("paddle"):
             backend = "paddle"
@@ -209,38 +214,47 @@ def suggest_inference_backend_and_config(
             backend = pb
         if available_backends is not None and backend not in available_backends:
             continue
-        candidate_backends.append(backend)
-        backend_to_pseudo_backend[backend] = pb
+        supported_backends.append(backend)
+        backend_to_pseudo_backends[backend].append(pb)
 
-    if not candidate_backends:
+    if not supported_backends:
         return None, "No inference backend can be selected."
 
     if hpi_config.backend is not None:
-        if hpi_config.backend not in candidate_backends:
+        if hpi_config.backend not in supported_backends:
             return (
                 None,
                 f"{repr(hpi_config.backend)} is not a supported inference backend.",
             )
         suggested_backend = hpi_config.backend
+        pseudo_backends = backend_to_pseudo_backends[suggested_backend]
+        pseudo_backend = pseudo_backends[0]
     else:
-        # The first backend is the preferred one.
-        suggested_backend = candidate_backends[0]
+        # Prefer the first one.
+        suggested_backend = supported_backends[0]
+        pseudo_backend = supported_pseudo_backends[0]
 
     suggested_backend_config = {}
     if suggested_backend == "paddle":
-        pseudo_backend = backend_to_pseudo_backend["paddle"]
         assert pseudo_backend in (
             "paddle",
+            "paddle_fp16",
+            "paddle_mkldnn",
             "paddle_tensorrt",
             "paddle_tensorrt_fp16",
         ), pseudo_backend
-        if pseudo_backend == "paddle_tensorrt":
+        if pseudo_backend == "paddle":
+            suggested_backend_config.update({"run_mode": "paddle"})
+        elif pseudo_backend == "paddle_fp16":
+            suggested_backend_config.update({"run_mode": "paddle_fp16"})
+        elif pseudo_backend == "paddle_mkldnn":
+            suggested_backend_config.update({"run_mode": "mkldnn"})
+        elif pseudo_backend == "paddle_tensorrt":
             suggested_backend_config.update({"run_mode": "trt_fp32"})
         elif pseudo_backend == "paddle_tensorrt_fp16":
             # TODO: Check if the target device supports FP16.
             suggested_backend_config.update({"run_mode": "trt_fp16"})
     elif suggested_backend == "tensorrt":
-        pseudo_backend = backend_to_pseudo_backend["tensorrt"]
         assert pseudo_backend in ("tensorrt", "tensorrt_fp16"), pseudo_backend
         if pseudo_backend == "tensorrt_fp16":
             suggested_backend_config.update({"precision": "fp16"})