predictor.py 9.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274
  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. import os
  15. from abc import abstractmethod
  16. import lazy_paddle as paddle
  17. import numpy as np
  18. from ....utils.flags import FLAGS_json_format_model
  19. from ....utils import logging
  20. from ...utils.pp_option import PaddlePredictorOption
  21. from ..base import BaseComponent
  22. class BasePaddlePredictor(BaseComponent):
  23. """Predictor based on Paddle Inference"""
  24. OUTPUT_KEYS = "pred"
  25. DEAULT_OUTPUTS = {"pred": "pred"}
  26. ENABLE_BATCH = True
  27. def __init__(self, model_dir, model_prefix, option):
  28. super().__init__()
  29. self.model_dir = model_dir
  30. self.model_prefix = model_prefix
  31. self._update_option(option)
  32. def _update_option(self, option):
  33. if option:
  34. if self.option and option == self.option:
  35. return
  36. self._option = option
  37. self._option.attach(self)
  38. self.reset()
  39. @property
  40. def option(self):
  41. return self._option if hasattr(self, "_option") else None
  42. @option.setter
  43. def option(self, option):
  44. self._update_option(option)
  45. def reset(self):
  46. if not self.option:
  47. self.option = PaddlePredictorOption()
  48. (
  49. self.predictor,
  50. self.inference_config,
  51. self.input_names,
  52. self.input_handlers,
  53. self.output_handlers,
  54. ) = self._create()
  55. logging.debug(f"Env: {self.option}")
  56. def _create(self):
  57. """_create"""
  58. from lazy_paddle.inference import Config, create_predictor
  59. model_postfix = ".json" if FLAGS_json_format_model else ".pdmodel"
  60. model_file = (self.model_dir / f"{self.model_prefix}{model_postfix}").as_posix()
  61. params_file = (self.model_dir / f"{self.model_prefix}.pdiparams").as_posix()
  62. config = Config(model_file, params_file)
  63. if self.option.device in ("gpu", "dcu"):
  64. config.enable_use_gpu(200, self.option.device_id)
  65. if hasattr(config, "enable_new_ir"):
  66. config.enable_new_ir(self.option.enable_new_ir)
  67. elif self.option.device == "npu":
  68. config.enable_custom_device("npu")
  69. elif self.option.device == "xpu":
  70. pass
  71. elif self.option.device == "mlu":
  72. config.enable_custom_device("mlu")
  73. else:
  74. assert self.option.device == "cpu"
  75. config.disable_gpu()
  76. if hasattr(config, "enable_new_ir"):
  77. config.enable_new_ir(self.option.enable_new_ir)
  78. if hasattr(config, "enable_new_executor"):
  79. config.enable_new_executor(True)
  80. if "mkldnn" in self.option.run_mode:
  81. try:
  82. config.enable_mkldnn()
  83. config.set_cpu_math_library_num_threads(self.option.cpu_threads)
  84. if "bf16" in self.option.run_mode:
  85. config.enable_mkldnn_bfloat16()
  86. except Exception as e:
  87. logging.warning(
  88. "MKL-DNN is not available. We will disable MKL-DNN."
  89. )
  90. precision_map = {
  91. "trt_int8": Config.Precision.Int8,
  92. "trt_fp32": Config.Precision.Float32,
  93. "trt_fp16": Config.Precision.Half,
  94. }
  95. if self.option.run_mode in precision_map.keys():
  96. config.enable_tensorrt_engine(
  97. workspace_size=(1 << 25) * self.option.batch_size,
  98. max_batch_size=self.option.batch_size,
  99. min_subgraph_size=self.option.min_subgraph_size,
  100. precision_mode=precision_map[self.option.run_mode],
  101. use_static=self.option.trt_use_static,
  102. use_calib_mode=self.option.trt_calib_mode,
  103. )
  104. if self.option.shape_info_filename is not None:
  105. if not os.path.exists(self.option.shape_info_filename):
  106. config.collect_shape_range_info(self.option.shape_info_filename)
  107. logging.info(
  108. f"Dynamic shape info is collected into: {self.option.shape_info_filename}"
  109. )
  110. else:
  111. logging.info(
  112. f"A dynamic shape info file ( {self.option.shape_info_filename} ) already exists. \
  113. No need to generate again."
  114. )
  115. config.enable_tuned_tensorrt_dynamic_shape(
  116. self.option.shape_info_filename, True
  117. )
  118. # Disable paddle inference logging
  119. config.disable_glog_info()
  120. for del_p in self.option.delete_pass:
  121. config.delete_pass(del_p)
  122. if self.option.device in ("gpu", "dcu"):
  123. if paddle.is_compiled_with_rocm():
  124. # Delete unsupported passes in dcu
  125. config.delete_pass("conv2d_add_act_fuse_pass")
  126. config.delete_pass("conv2d_add_fuse_pass")
  127. # Enable shared memory
  128. config.enable_memory_optim()
  129. config.switch_ir_optim(True)
  130. # Disable feed, fetch OP, needed by zero_copy_run
  131. config.switch_use_feed_fetch_ops(False)
  132. predictor = create_predictor(config)
  133. # Get input and output handlers
  134. input_names = predictor.get_input_names()
  135. input_names.sort()
  136. input_handlers = []
  137. output_handlers = []
  138. for input_name in input_names:
  139. input_handler = predictor.get_input_handle(input_name)
  140. input_handlers.append(input_handler)
  141. output_names = predictor.get_output_names()
  142. for output_name in output_names:
  143. output_handler = predictor.get_output_handle(output_name)
  144. output_handlers.append(output_handler)
  145. return predictor, config, input_names, input_handlers, output_handlers
  146. def get_input_names(self):
  147. """get input names"""
  148. return self.input_names
  149. def apply(self, **kwargs):
  150. x = self.to_batch(**kwargs)
  151. for idx in range(len(x)):
  152. self.input_handlers[idx].reshape(x[idx].shape)
  153. self.input_handlers[idx].copy_from_cpu(x[idx])
  154. self.predictor.run()
  155. output = []
  156. for out_tensor in self.output_handlers:
  157. batch = out_tensor.copy_to_cpu()
  158. output.append(batch)
  159. return self.format_output(output)
  160. def format_output(self, pred):
  161. return [{"pred": res} for res in zip(*pred)]
  162. @abstractmethod
  163. def to_batch(self):
  164. raise NotImplementedError
  165. class ImagePredictor(BasePaddlePredictor):
  166. INPUT_KEYS = "img"
  167. DEAULT_INPUTS = {"img": "img"}
  168. def to_batch(self, img):
  169. return [np.stack(img, axis=0).astype(dtype=np.float32, copy=False)]
  170. class ImageDetPredictor(BasePaddlePredictor):
  171. INPUT_KEYS = [
  172. ["img", "scale_factors"],
  173. ["img", "scale_factors", "img_size"],
  174. ["img", "img_size"],
  175. ]
  176. OUTPUT_KEYS = [["boxes"], ["boxes", "masks"]]
  177. DEAULT_INPUTS = {"img": "img", "scale_factors": "scale_factors"}
  178. DEAULT_OUTPUTS = None
  179. def to_batch(self, img, scale_factors=[[1.0, 1.0]], img_size=None):
  180. scale_factors = [scale_factor[::-1] for scale_factor in scale_factors]
  181. if img_size is None:
  182. return [
  183. np.stack(img, axis=0).astype(dtype=np.float32, copy=False),
  184. np.stack(scale_factors, axis=0).astype(dtype=np.float32, copy=False),
  185. ]
  186. else:
  187. img_size = [img_size[::-1] for img_size in img_size]
  188. return [
  189. np.stack(img_size, axis=0).astype(dtype=np.float32, copy=False),
  190. np.stack(img, axis=0).astype(dtype=np.float32, copy=False),
  191. np.stack(scale_factors, axis=0).astype(dtype=np.float32, copy=False),
  192. ]
  193. def format_output(self, pred):
  194. box_idx_start = 0
  195. pred_box = []
  196. if len(pred) == 4:
  197. # Adapt to SOLOv2
  198. pred_class_id = []
  199. pred_mask = []
  200. pred_class_id.append([pred[1], pred[2]])
  201. pred_mask.append(pred[3])
  202. return [
  203. {
  204. "class_id": np.array(pred_class_id[i]),
  205. "masks": np.array(pred_mask[i]),
  206. }
  207. for i in range(len(pred_class_id))
  208. ]
  209. if len(pred) == 3:
  210. # Adapt to Instance Segmentation
  211. pred_mask = []
  212. for idx in range(len(pred[1])):
  213. np_boxes_num = pred[1][idx]
  214. box_idx_end = box_idx_start + np_boxes_num
  215. np_boxes = pred[0][box_idx_start:box_idx_end]
  216. pred_box.append(np_boxes)
  217. if len(pred) == 3:
  218. np_masks = pred[2][box_idx_start:box_idx_end]
  219. pred_mask.append(np_masks)
  220. box_idx_start = box_idx_end
  221. if len(pred) == 3:
  222. return [
  223. {"boxes": np.array(pred_box[i]), "masks": np.array(pred_mask[i])}
  224. for i in range(len(pred_box))
  225. ]
  226. else:
  227. return [{"boxes": np.array(res)} for res in pred_box]
  228. class TSPPPredictor(BasePaddlePredictor):
  229. INPUT_KEYS = "ts"
  230. DEAULT_INPUTS = {"ts": "ts"}
  231. def to_batch(self, ts):
  232. n = len(ts[0])
  233. x = [np.stack([lst[i] for lst in ts], axis=0) for i in range(n)]
  234. return x