basic_predictor.py 6.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from typing import Union, Tuple, List, Dict, Any, Iterator
  15. from abc import abstractmethod
  16. from .....utils.subclass_register import AutoRegisterABCMetaClass
  17. from .....utils.flags import (
  18. INFER_BENCHMARK,
  19. INFER_BENCHMARK_WARMUP,
  20. )
  21. from .....utils import logging
  22. from ....utils.pp_option import PaddlePredictorOption
  23. from ....utils.benchmark import benchmark
  24. from ....common.batch_sampler import BaseBatchSampler
  25. from .base_predictor import BasePredictor
  26. class PredictionWrap:
  27. """Wraps the prediction data and supports get by index."""
  28. def __init__(self, data: Dict[str, List[Any]], num: int) -> None:
  29. """Initializes the PredictionWrap with prediction data.
  30. Args:
  31. data (Dict[str, List[Any]]): A dictionary where keys are string identifiers and values are lists of predictions.
  32. num (int): The number of predictions, that is length of values per key in the data dictionary.
  33. Raises:
  34. AssertionError: If the length of any list in data does not match num.
  35. """
  36. assert isinstance(data, dict), "data must be a dictionary"
  37. for k in data:
  38. assert len(data[k]) == num, f"{len(data[k])} != {num} for key {k}!"
  39. self._data = data
  40. self._keys = data.keys()
  41. def get_by_idx(self, idx: int) -> Dict[str, Any]:
  42. """Get the prediction by specified index.
  43. Args:
  44. idx (int): The index to get predictions from.
  45. Returns:
  46. Dict[str, Any]: A dictionary with the same keys as the input data, but with the values at the specified index.
  47. """
  48. return {key: self._data[key][idx] for key in self._keys}
  49. class BasicPredictor(
  50. BasePredictor,
  51. metaclass=AutoRegisterABCMetaClass,
  52. ):
  53. """BasicPredictor."""
  54. __is_base = True
  55. def __init__(
  56. self,
  57. model_dir: str,
  58. config: Dict[str, Any] = None,
  59. device: str = None,
  60. pp_option: PaddlePredictorOption = None,
  61. ) -> None:
  62. """Initializes the BasicPredictor.
  63. Args:
  64. model_dir (str): The directory where the model files are stored.
  65. config (Dict[str, Any], optional): The configuration dictionary. Defaults to None.
  66. device (str, optional): The device to run the inference engine on. Defaults to None.
  67. pp_option (PaddlePredictorOption, optional): The inference engine options. Defaults to None.
  68. """
  69. super().__init__(model_dir=model_dir, config=config)
  70. if not pp_option:
  71. pp_option = PaddlePredictorOption(model_name=self.model_name)
  72. if device:
  73. pp_option.device = device
  74. self.pp_option = pp_option
  75. self.batch_sampler = self._build_batch_sampler()
  76. self.result_class = self._get_result_class()
  77. logging.debug(f"{self.__class__.__name__}: {self.model_dir}")
  78. self.benchmark = benchmark
  79. def __call__(self, input: Any, **kwargs: Dict[str, Any]) -> Iterator[Any]:
  80. """
  81. Predict with the input data.
  82. Args:
  83. input (Any): The input data to be predicted.
  84. **kwargs (Dict[str, Any]): Additional keyword arguments to set up predictor.
  85. Returns:
  86. Iterator[Any]: An iterator yielding the prediction output.
  87. """
  88. self.set_predictor(**kwargs)
  89. if self.benchmark:
  90. self.benchmark.start()
  91. if INFER_BENCHMARK_WARMUP > 0:
  92. output = self.apply(input)
  93. warmup_num = 0
  94. for _ in range(INFER_BENCHMARK_WARMUP):
  95. try:
  96. next(output)
  97. warmup_num += 1
  98. except StopIteration:
  99. logging.warning(
  100. f"There are only {warmup_num} batches in input data, but `INFER_BENCHMARK_WARMUP` has been set to {INFER_BENCHMARK_WARMUP}."
  101. )
  102. break
  103. self.benchmark.warmup_stop(warmup_num)
  104. output = list(self.apply(input))
  105. self.benchmark.collect(len(output))
  106. else:
  107. yield from self.apply(input)
  108. def apply(self, input: Any) -> Iterator[Any]:
  109. """
  110. Do predicting with the input data and yields predictions.
  111. Args:
  112. input (Any): The input data to be predicted.
  113. Yields:
  114. Iterator[Any]: An iterator yielding prediction results.
  115. """
  116. for batch_data in self.batch_sampler(input):
  117. prediction = self.process(batch_data)
  118. prediction = PredictionWrap(prediction, len(batch_data))
  119. for idx in range(len(batch_data)):
  120. yield self.result_class(prediction.get_by_idx(idx))
  121. def set_predictor(
  122. self,
  123. batch_size: int = None,
  124. device: str = None,
  125. pp_option: PaddlePredictorOption = None,
  126. ) -> None:
  127. """
  128. Sets the predictor configuration.
  129. Args:
  130. batch_size (int, optional): The batch size to use. Defaults to None.
  131. device (str, optional): The device to run the predictor on. Defaults to None.
  132. pp_option (PaddlePredictorOption, optional): The predictor options to set. Defaults to None.
  133. Returns:
  134. None
  135. """
  136. if batch_size:
  137. self.batch_sampler.batch_size = batch_size
  138. self.pp_option.batch_size = batch_size
  139. if device and device != self.pp_option.device:
  140. self.pp_option.device = device
  141. if pp_option and pp_option != self.pp_option:
  142. self.pp_option = pp_option
  143. @abstractmethod
  144. def _build_batch_sampler(self) -> BaseBatchSampler:
  145. """Build batch sampler.
  146. Returns:
  147. BaseBatchSampler: batch sampler object.
  148. """
  149. raise NotImplementedError
  150. @abstractmethod
  151. def process(self, batch_data: List[Any]) -> Dict[str, List[Any]]:
  152. """process the batch data sampled from BatchSampler and return the prediction result.
  153. Args:
  154. batch_data (List[Any]): The batch data sampled from BatchSampler.
  155. Returns:
  156. Dict[str, List[Any]]: The prediction result.
  157. """
  158. raise NotImplementedError
  159. @abstractmethod
  160. def _get_result_class(self) -> type:
  161. """Get the result class.
  162. Returns:
  163. type: The result class.
  164. """
  165. raise NotImplementedError