paddlex 3.0.0rc0__py3-none-any.whl → 3.0.1__py3-none-any.whl

paddlex/inference/models/base/predictor/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -13,4 +13,3 @@
 # limitations under the License.
 
 from .base_predictor import BasePredictor
-from .basic_predictor import BasicPredictor

paddlex/inference/models/base/predictor/base_predictor.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,12 +12,29 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Dict, Any, Iterator
+from abc import ABC, abstractmethod
+from copy import deepcopy
 from pathlib import Path
-from abc import abstractmethod, ABC
-
-from ....utils.io import YAMLReader
+from typing import Any, Dict, Iterator, List, Optional, Union
+
+from pydantic import ValidationError
+
+from ..... import constants
+from .....utils import logging
+from .....utils.deps import require_hpip
+from .....utils.device import get_default_device, parse_device
+from .....utils.flags import (
+    INFER_BENCHMARK,
+    INFER_BENCHMARK_ITERS,
+    INFER_BENCHMARK_WARMUP,
+)
+from .....utils.subclass_register import AutoRegisterABCMetaClass
 from ....common.batch_sampler import BaseBatchSampler
+from ....utils.benchmark import ENTRY_POINT_NAME, benchmark
+from ....utils.hpi import HPIConfig, HPIInfo
+from ....utils.io import YAMLReader
+from ....utils.pp_option import PaddlePredictorOption
+from ...common import HPInfer, PaddleInfer
 
 
 class PredictionWrap:
@@ -51,19 +68,45 @@ class PredictionWrap:
         return {key: self._data[key][idx] for key in self._keys}
 
 
-class BasePredictor(ABC):
-    """BasePredictor."""
-
-    MODEL_FILE_PREFIX = "inference"
-
-    def __init__(self, model_dir: str, config: Dict = None) -> None:
+class BasePredictor(
+    ABC,
+    metaclass=AutoRegisterABCMetaClass,
+):
+    MODEL_FILE_PREFIX = constants.MODEL_FILE_PREFIX
+
+    __is_base = True
+
+    def __init__(
+        self,
+        model_dir: str,
+        config: Optional[Dict[str, Any]] = None,
+        *,
+        device: Optional[str] = None,
+        batch_size: int = 1,
+        pp_option: Optional[PaddlePredictorOption] = None,
+        use_hpip: bool = False,
+        hpi_config: Optional[Union[Dict[str, Any], HPIConfig]] = None,
+    ) -> None:
         """Initializes the BasePredictor.
 
         Args:
-            model_dir (str): The directory where the static model files is stored.
-            config (dict, optional): The configuration of model to infer. Defaults to None.
+            model_dir (str): The directory where the model files are stored.
+            config (Optional[Dict[str, Any]], optional): The model configuration
+                dictionary. Defaults to None.
+            device (Optional[str], optional): The device to run the inference
+                engine on. Defaults to None.
+            batch_size (int, optional): The batch size to predict.
+                Defaults to 1.
+            pp_option (Optional[PaddlePredictorOption], optional): The inference
+                engine options. Defaults to None.
+            use_hpip (bool, optional): Whether to use high-performance inference
+                plugin. Defaults to False.
+            hpi_config (Optional[Union[Dict[str, Any], HPIConfig]], optional):
+                The high-performance inference configuration dictionary.
+                Defaults to None.
         """
         super().__init__()
+
         self.model_dir = Path(model_dir)
         self.config = config if config else self.load_config(self.model_dir)
         self.batch_sampler = self._build_batch_sampler()
@@ -71,7 +114,16 @@ class BasePredictor(ABC):
 
         # alias predict() to the __call__()
         self.predict = self.__call__
-        self.benchmark = None
+
+        self.batch_sampler.batch_size = batch_size
+        self._use_hpip = use_hpip
+        if not use_hpip:
+            self._pp_option = self._prepare_pp_option(pp_option, device)
+        else:
+            require_hpip()
+            self._hpi_config = self._prepare_hpi_config(hpi_config, device)
+
+        logging.debug(f"{self.__class__.__name__}: {self.model_dir}")
 
     @property
     def config_path(self) -> str:
@@ -93,40 +145,105 @@ class BasePredictor(ABC):
         """
         return self.config["Global"]["model_name"]
 
-    @classmethod
-    def get_config_path(cls, model_dir) -> str:
-        """Get the path to the configuration file for the given model directory.
+    @property
+    def pp_option(self) -> PaddlePredictorOption:
+        if not hasattr(self, "_pp_option"):
+            raise AttributeError(f"{repr(self)} has no attribute 'pp_option'.")
+        return self._pp_option
 
-        Args:
-            model_dir (Path): The directory where the static model files is stored.
+    @property
+    def hpi_config(self) -> HPIConfig:
+        if not hasattr(self, "_hpi_config"):
+            raise AttributeError(f"{repr(self)} has no attribute 'hpi_config'.")
+        return self._hpi_config
 
-        Returns:
-            Path: The path to the configuration file.
+    @property
+    def use_hpip(self) -> bool:
+        return self._use_hpip
+
+    def __call__(
+        self,
+        input: Any,
+        batch_size: Optional[int] = None,
+        **kwargs: Any,
+    ) -> Iterator[Any]:
         """
-        return model_dir / f"{cls.MODEL_FILE_PREFIX}.yml"
-
-    @classmethod
-    def load_config(cls, model_dir) -> Dict:
-        """Load the configuration from the specified model directory.
+        Predict with the input data.
 
         Args:
-            model_dir (Path): The where the static model files is stored.
+            input (Any): The input data to be predicted.
+            batch_size (int, optional): The batch size to use. Defaults to None.
+            **kwargs (Dict[str, Any]): Additional keyword arguments to set up predictor.
 
         Returns:
-            dict: The loaded configuration dictionary.
+            Iterator[Any]: An iterator yielding the prediction output.
         """
-        yaml_reader = YAMLReader()
-        return yaml_reader.read(cls.get_config_path(model_dir))
+        self.set_predictor(batch_size)
+        if INFER_BENCHMARK:
+            # TODO(zhang-prog): Get metadata of input data
+            @benchmark.timeit_with_options(name=ENTRY_POINT_NAME)
+            def _apply(input, **kwargs):
+                return list(self.apply(input, **kwargs))
+
+            if isinstance(input, list):
+                raise TypeError("`input` cannot be a list in benchmark mode")
+            input = [input] * batch_size
+
+            if not (INFER_BENCHMARK_WARMUP > 0 or INFER_BENCHMARK_ITERS > 0):
+                raise RuntimeError(
+                    "At least one of `INFER_BENCHMARK_WARMUP` and `INFER_BENCHMARK_ITERS` must be greater than zero"
+                )
+
+            if INFER_BENCHMARK_WARMUP > 0:
+                benchmark.start_warmup()
+                for _ in range(INFER_BENCHMARK_WARMUP):
+                    output = _apply(input, **kwargs)
+                benchmark.collect(batch_size)
+                benchmark.stop_warmup()
+
+            if INFER_BENCHMARK_ITERS > 0:
+                for _ in range(INFER_BENCHMARK_ITERS):
+                    output = _apply(input, **kwargs)
+                benchmark.collect(batch_size)
+
+            yield output[0]
+        else:
+            yield from self.apply(input, **kwargs)
+
+    def set_predictor(
+        self,
+        batch_size: Optional[int] = None,
+    ) -> None:
+        """
+        Sets the predictor configuration.
 
-    @abstractmethod
-    def __call__(self, input: Any, **kwargs: Dict[str, Any]) -> Iterator[Any]:
-        """Predict with the given input and additional keyword arguments."""
-        raise NotImplementedError
+        Args:
+            batch_size (Optional[int], optional): The batch size to use. Defaults to None.
 
-    @abstractmethod
-    def set_predictor(self, batch_size: int = None, device: str = None, *args) -> None:
-        """Sets up the predictor."""
-        raise NotImplementedError
+        Returns:
+            None
+        """
+        if batch_size:
+            self.batch_sampler.batch_size = batch_size
+
+    def get_hpi_info(self):
+        if "Hpi" not in self.config:
+            return None
+        try:
+            return HPIInfo.model_validate(self.config["Hpi"])
+        except ValidationError as e:
+            logging.exception("The HPI info in the model config file is invalid.")
+            raise RuntimeError(f"Invalid HPI info: {str(e)}") from e
+
+    def create_static_infer(self):
+        if not self._use_hpip:
+            return PaddleInfer(self.model_dir, self.MODEL_FILE_PREFIX, self._pp_option)
+        else:
+            return HPInfer(
+                self.model_dir,
+                self.MODEL_FILE_PREFIX,
+                self._hpi_config,
+            )
 
     def apply(self, input: Any, **kwargs) -> Iterator[Any]:
         """
@@ -138,7 +255,15 @@ class BasePredictor(ABC):
         Yields:
             Iterator[Any]: An iterator yielding prediction results.
         """
-        for batch_data in self.batch_sampler(input):
+        if INFER_BENCHMARK:
+            if not isinstance(input, list):
+                raise TypeError("In benchmark mode, `input` must be a list")
+            batches = list(self.batch_sampler(input))
+            if len(batches) != 1 or len(batches[0]) != len(input):
+                raise ValueError("Unexpected number of instances")
+        else:
+            batches = self.batch_sampler(input)
+        for batch_data in batches:
             prediction = self.process(batch_data, **kwargs)
             prediction = PredictionWrap(prediction, len(batch_data))
             for idx in range(len(batch_data)):
@@ -156,6 +281,31 @@ class BasePredictor(ABC):
         """
         raise NotImplementedError
 
+    @classmethod
+    def get_config_path(cls, model_dir) -> str:
+        """Get the path to the configuration file for the given model directory.
+
+        Args:
+            model_dir (Path): The directory where the static model files is stored.
+
+        Returns:
+            Path: The path to the configuration file.
+        """
+        return model_dir / f"{cls.MODEL_FILE_PREFIX}.yml"
+
+    @classmethod
+    def load_config(cls, model_dir) -> Dict:
+        """Load the configuration from the specified model directory.
+
+        Args:
+            model_dir (Path): The where the static model files is stored.
+
+        Returns:
+            dict: The loaded configuration dictionary.
+        """
+        yaml_reader = YAMLReader()
+        return yaml_reader.read(cls.get_config_path(model_dir))
+
     @abstractmethod
     def _build_batch_sampler(self) -> BaseBatchSampler:
         """Build batch sampler.
@@ -173,3 +323,92 @@ class BasePredictor(ABC):
             type: The result class.
         """
         raise NotImplementedError
+
+    def _prepare_pp_option(
+        self,
+        pp_option: Optional[PaddlePredictorOption],
+        device: Optional[str],
+    ) -> PaddlePredictorOption:
+        if pp_option is None or device is not None:
+            device_info = self._get_device_info(device)
+        else:
+            device_info = None
+        if pp_option is None:
+            pp_option = PaddlePredictorOption(model_name=self.model_name)
+        elif pp_option.model_name is None:
+            pp_option.model_name = self.model_name
+        if device_info:
+            pp_option.device_type = device_info[0]
+            pp_option.device_id = device_info[1]
+        hpi_info = self.get_hpi_info()
+        if hpi_info is not None:
+            hpi_info = hpi_info.model_dump(exclude_unset=True)
+            if pp_option.trt_dynamic_shapes is None:
+                trt_dynamic_shapes = (
+                    hpi_info.get("backend_configs", {})
+                    .get("paddle_infer", {})
+                    .get("trt_dynamic_shapes", None)
+                )
+                if trt_dynamic_shapes is not None:
+                    logging.debug(
+                        "TensorRT dynamic shapes set to %s", trt_dynamic_shapes
+                    )
+                    pp_option.trt_dynamic_shapes = trt_dynamic_shapes
+            if pp_option.trt_dynamic_shape_input_data is None:
+                trt_dynamic_shape_input_data = (
+                    hpi_info.get("backend_configs", {})
+                    .get("paddle_infer", {})
+                    .get("trt_dynamic_shape_input_data", None)
+                )
+                if trt_dynamic_shape_input_data is not None:
+                    logging.debug(
+                        "TensorRT dynamic shape input data set to %s",
+                        trt_dynamic_shape_input_data,
+                    )
+                    pp_option.trt_dynamic_shape_input_data = (
+                        trt_dynamic_shape_input_data
+                    )
+        return pp_option
+
+    def _prepare_hpi_config(
+        self,
+        hpi_config: Optional[Union[Dict[str, Any], HPIConfig]],
+        device: Optional[str],
+    ) -> HPIConfig:
+        if hpi_config is None:
+            hpi_config = {}
+        elif isinstance(hpi_config, HPIConfig):
+            hpi_config = hpi_config.model_dump(exclude_unset=True)
+        else:
+            hpi_config = deepcopy(hpi_config)
+
+        if "model_name" not in hpi_config:
+            hpi_config["model_name"] = self.model_name
+
+        if device is not None or "device_type" not in hpi_config:
+            device_type, device_id = self._get_device_info(device)
+            hpi_config["device_type"] = device_type
+            if device is not None or "device_id" not in hpi_config:
+                hpi_config["device_id"] = device_id
+
+        if "hpi_info" not in hpi_config:
+            hpi_info = self.get_hpi_info()
+            if hpi_info is not None:
+                hpi_config["hpi_info"] = hpi_info
+
+        hpi_config = HPIConfig.model_validate(hpi_config)
+
+        return hpi_config
+
+    # Should this be static?
+    def _get_device_info(self, device):
+        if device is None:
+            device = get_default_device()
+        device_type, device_ids = parse_device(device)
+        if device_ids is not None:
+            device_id = device_ids[0]
+        else:
+            device_id = None
+        if device_ids and len(device_ids) > 1:
+            logging.debug("Got multiple device IDs. Using the first one: %d", device_id)
+        return device_type, device_id

paddlex/inference/models/common/__init__.py

@@ -1,4 +1,4 @@
-# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+# Copyright (c) 2024 PaddlePaddle Authors. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -12,24 +12,15 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from .static_infer import HPInfer, PaddleInfer, StaticInfer
+from .tokenizer import PretrainedTokenizer
+from .ts import BuildTSDataset, TimeFeature, TSCutOff, TSNormalize, TStoArray, TStoBatch
 from .vision import (
     F,
+    Normalize,
     Resize,
     ResizeByLong,
     ResizeByShort,
-    Normalize,
-    ToCHWImage,
     ToBatch,
+    ToCHWImage,
 )
-
-from .ts import (
-    BuildTSDataset,
-    TSCutOff,
-    TSNormalize,
-    TimeFeature,
-    TStoArray,
-    TStoBatch,
-)
-from .static_infer import StaticInfer
-
-from .tokenizer import PretrainedTokenizer, PretrainedTokenizer