predictor.py 10 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287
  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 paddle.is_compiled_with_rocm():
  66. os.environ["FLAGS_conv_workspace_size_limit"] = "2000"
  67. elif hasattr(config, "enable_new_ir"):
  68. config.enable_new_ir(self.option.enable_new_ir)
  69. elif self.option.device == "npu":
  70. config.enable_custom_device("npu")
  71. os.environ["FLAGS_npu_jit_compile"] = "0"
  72. os.environ["FLAGS_use_stride_kernel"] = "0"
  73. os.environ["FLAGS_allocator_strategy"] = "auto_growth"
  74. os.environ["CUSTOM_DEVICE_BLACK_LIST"] = (
  75. "pad3d,pad3d_grad,set_value,set_value_with_tensor"
  76. )
  77. os.environ["FLAGS_npu_scale_aclnn"] = "True"
  78. os.environ["FLAGS_npu_split_aclnn"] = "True"
  79. elif self.option.device == "xpu":
  80. os.environ["BKCL_FORCE_SYNC"] = "1"
  81. os.environ["BKCL_TIMEOUT"] = "1800"
  82. os.environ["FLAGS_use_stride_kernel"] = "0"
  83. elif self.option.device == "mlu":
  84. config.enable_custom_device("mlu")
  85. os.environ["FLAGS_use_stride_kernel"] = "0"
  86. else:
  87. assert self.option.device == "cpu"
  88. config.disable_gpu()
  89. if hasattr(config, "enable_new_ir"):
  90. config.enable_new_ir(self.option.enable_new_ir)
  91. if hasattr(config, "enable_new_executor"):
  92. config.enable_new_executor(True)
  93. if "mkldnn" in self.option.run_mode:
  94. try:
  95. config.enable_mkldnn()
  96. config.set_cpu_math_library_num_threads(self.option.cpu_threads)
  97. if "bf16" in self.option.run_mode:
  98. config.enable_mkldnn_bfloat16()
  99. except Exception as e:
  100. logging.warning(
  101. "MKL-DNN is not available. We will disable MKL-DNN."
  102. )
  103. precision_map = {
  104. "trt_int8": Config.Precision.Int8,
  105. "trt_fp32": Config.Precision.Float32,
  106. "trt_fp16": Config.Precision.Half,
  107. }
  108. if self.option.run_mode in precision_map.keys():
  109. config.enable_tensorrt_engine(
  110. workspace_size=(1 << 25) * self.option.batch_size,
  111. max_batch_size=self.option.batch_size,
  112. min_subgraph_size=self.option.min_subgraph_size,
  113. precision_mode=precision_map[self.option.run_mode],
  114. use_static=self.option.trt_use_static,
  115. use_calib_mode=self.option.trt_calib_mode,
  116. )
  117. if self.option.shape_info_filename is not None:
  118. if not os.path.exists(self.option.shape_info_filename):
  119. config.collect_shape_range_info(self.option.shape_info_filename)
  120. logging.info(
  121. f"Dynamic shape info is collected into: {self.option.shape_info_filename}"
  122. )
  123. else:
  124. logging.info(
  125. f"A dynamic shape info file ( {self.option.shape_info_filename} ) already exists. \
  126. No need to generate again."
  127. )
  128. config.enable_tuned_tensorrt_dynamic_shape(
  129. self.option.shape_info_filename, True
  130. )
  131. # Disable paddle inference logging
  132. config.disable_glog_info()
  133. for del_p in self.option.delete_pass:
  134. config.delete_pass(del_p)
  135. if self.option.device in ("gpu", "dcu"):
  136. if paddle.is_compiled_with_rocm():
  137. # Delete unsupported passes in dcu
  138. config.delete_pass("conv2d_add_act_fuse_pass")
  139. config.delete_pass("conv2d_add_fuse_pass")
  140. # Enable shared memory
  141. config.enable_memory_optim()
  142. config.switch_ir_optim(True)
  143. # Disable feed, fetch OP, needed by zero_copy_run
  144. config.switch_use_feed_fetch_ops(False)
  145. predictor = create_predictor(config)
  146. # Get input and output handlers
  147. input_names = predictor.get_input_names()
  148. input_names.sort()
  149. input_handlers = []
  150. output_handlers = []
  151. for input_name in input_names:
  152. input_handler = predictor.get_input_handle(input_name)
  153. input_handlers.append(input_handler)
  154. output_names = predictor.get_output_names()
  155. for output_name in output_names:
  156. output_handler = predictor.get_output_handle(output_name)
  157. output_handlers.append(output_handler)
  158. return predictor, config, input_names, input_handlers, output_handlers
  159. def get_input_names(self):
  160. """get input names"""
  161. return self.input_names
  162. def apply(self, **kwargs):
  163. x = self.to_batch(**kwargs)
  164. for idx in range(len(x)):
  165. self.input_handlers[idx].reshape(x[idx].shape)
  166. self.input_handlers[idx].copy_from_cpu(x[idx])
  167. self.predictor.run()
  168. output = []
  169. for out_tensor in self.output_handlers:
  170. batch = out_tensor.copy_to_cpu()
  171. output.append(batch)
  172. return self.format_output(output)
  173. def format_output(self, pred):
  174. return [{"pred": res} for res in zip(*pred)]
  175. @abstractmethod
  176. def to_batch(self):
  177. raise NotImplementedError
  178. class ImagePredictor(BasePaddlePredictor):
  179. INPUT_KEYS = "img"
  180. DEAULT_INPUTS = {"img": "img"}
  181. def to_batch(self, img):
  182. return [np.stack(img, axis=0).astype(dtype=np.float32, copy=False)]
  183. class ImageDetPredictor(BasePaddlePredictor):
  184. INPUT_KEYS = [
  185. ["img", "scale_factors"],
  186. ["img", "scale_factors", "img_size"],
  187. ["img", "img_size"],
  188. ]
  189. OUTPUT_KEYS = [["boxes"], ["boxes", "masks"]]
  190. DEAULT_INPUTS = {"img": "img", "scale_factors": "scale_factors"}
  191. DEAULT_OUTPUTS = None
  192. def to_batch(self, img, scale_factors=[[1.0, 1.0]], img_size=None):
  193. scale_factors = [scale_factor[::-1] for scale_factor in scale_factors]
  194. if img_size is None:
  195. return [
  196. np.stack(img, axis=0).astype(dtype=np.float32, copy=False),
  197. np.stack(scale_factors, axis=0).astype(dtype=np.float32, copy=False),
  198. ]
  199. else:
  200. img_size = [img_size[::-1] for img_size in img_size]
  201. return [
  202. np.stack(img_size, axis=0).astype(dtype=np.float32, copy=False),
  203. np.stack(img, axis=0).astype(dtype=np.float32, copy=False),
  204. np.stack(scale_factors, axis=0).astype(dtype=np.float32, copy=False),
  205. ]
  206. def format_output(self, pred):
  207. box_idx_start = 0
  208. pred_box = []
  209. if len(pred) == 4:
  210. # Adapt to SOLOv2
  211. pred_class_id = []
  212. pred_mask = []
  213. pred_class_id.append([pred[1], pred[2]])
  214. pred_mask.append(pred[3])
  215. return [
  216. {
  217. "class_id": np.array(pred_class_id[i]),
  218. "masks": np.array(pred_mask[i]),
  219. }
  220. for i in range(len(pred_class_id))
  221. ]
  222. if len(pred) == 3:
  223. # Adapt to Instance Segmentation
  224. pred_mask = []
  225. for idx in range(len(pred[1])):
  226. np_boxes_num = pred[1][idx]
  227. box_idx_end = box_idx_start + np_boxes_num
  228. np_boxes = pred[0][box_idx_start:box_idx_end]
  229. pred_box.append(np_boxes)
  230. if len(pred) == 3:
  231. np_masks = pred[2][box_idx_start:box_idx_end]
  232. pred_mask.append(np_masks)
  233. box_idx_start = box_idx_end
  234. if len(pred) == 3:
  235. return [
  236. {"boxes": np.array(pred_box[i]), "masks": np.array(pred_mask[i])}
  237. for i in range(len(pred_box))
  238. ]
  239. else:
  240. return [{"boxes": np.array(res)} for res in pred_box]
  241. class TSPPPredictor(BasePaddlePredictor):
  242. INPUT_KEYS = "ts"
  243. DEAULT_INPUTS = {"ts": "ts"}
  244. def to_batch(self, ts):
  245. n = len(ts[0])
  246. x = [np.stack([lst[i] for lst in ts], axis=0) for i in range(n)]
  247. return x