predictor.py 12 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352
  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 Copy2GPU(BaseComponent):
  23. def __init__(self, input_handlers):
  24. super().__init__()
  25. self.input_handlers = input_handlers
  26. def apply(self, x):
  27. for idx in range(len(x)):
  28. self.input_handlers[idx].reshape(x[idx].shape)
  29. self.input_handlers[idx].copy_from_cpu(x[idx])
  30. class Copy2CPU(BaseComponent):
  31. def __init__(self, output_handlers):
  32. super().__init__()
  33. self.output_handlers = output_handlers
  34. def apply(self):
  35. output = []
  36. for out_tensor in self.output_handlers:
  37. batch = out_tensor.copy_to_cpu()
  38. output.append(batch)
  39. return output
  40. class Infer(BaseComponent):
  41. def __init__(self, predictor):
  42. super().__init__()
  43. self.predictor = predictor
  44. def apply(self):
  45. self.predictor.run()
  46. class BasePaddlePredictor(BaseComponent):
  47. """Predictor based on Paddle Inference"""
  48. OUTPUT_KEYS = "pred"
  49. DEAULT_OUTPUTS = {"pred": "pred"}
  50. ENABLE_BATCH = True
  51. def __init__(self, model_dir, model_prefix, option):
  52. super().__init__()
  53. self.model_dir = model_dir
  54. self.model_prefix = model_prefix
  55. self._update_option(option)
  56. def _update_option(self, option):
  57. if option:
  58. if self.option and option == self.option:
  59. return
  60. self._option = option
  61. self._reset()
  62. @property
  63. def option(self):
  64. return self._option if hasattr(self, "_option") else None
  65. @option.setter
  66. def option(self, option):
  67. self._update_option(option)
  68. def _reset(self):
  69. if not self.option:
  70. self.option = PaddlePredictorOption()
  71. logging.debug(f"Env: {self.option}")
  72. (
  73. predictor,
  74. input_handlers,
  75. output_handlers,
  76. ) = self._create()
  77. self.copy2gpu = Copy2GPU(input_handlers)
  78. self.copy2cpu = Copy2CPU(output_handlers)
  79. self.infer = Infer(predictor)
  80. self.option.changed = False
  81. def _create(self):
  82. """_create"""
  83. from lazy_paddle.inference import Config, create_predictor
  84. model_postfix = ".json" if FLAGS_json_format_model else ".pdmodel"
  85. model_file = (self.model_dir / f"{self.model_prefix}{model_postfix}").as_posix()
  86. params_file = (self.model_dir / f"{self.model_prefix}.pdiparams").as_posix()
  87. config = Config(model_file, params_file)
  88. config.enable_memory_optim()
  89. if self.option.device in ("gpu", "dcu"):
  90. if self.option.device == "gpu":
  91. config.exp_disable_mixed_precision_ops({"feed", "fetch"})
  92. config.enable_use_gpu(100, self.option.device_id)
  93. if self.option.device == "gpu":
  94. # NOTE: The pptrt settings are not aligned with those of FD.
  95. precision_map = {
  96. "trt_int8": Config.Precision.Int8,
  97. "trt_fp32": Config.Precision.Float32,
  98. "trt_fp16": Config.Precision.Half,
  99. }
  100. if self.option.run_mode in precision_map.keys():
  101. config.enable_tensorrt_engine(
  102. workspace_size=(1 << 25) * self.option.batch_size,
  103. max_batch_size=self.option.batch_size,
  104. min_subgraph_size=self.option.min_subgraph_size,
  105. precision_mode=precision_map[self.option.run_mode],
  106. use_static=self.option.trt_use_static,
  107. use_calib_mode=self.option.trt_calib_mode,
  108. )
  109. if self.option.shape_info_filename is not None:
  110. if not os.path.exists(self.option.shape_info_filename):
  111. config.collect_shape_range_info(
  112. self.option.shape_info_filename
  113. )
  114. logging.info(
  115. f"Dynamic shape info is collected into: {self.option.shape_info_filename}"
  116. )
  117. else:
  118. logging.info(
  119. f"A dynamic shape info file ( {self.option.shape_info_filename} ) already exists. \
  120. No need to generate again."
  121. )
  122. config.enable_tuned_tensorrt_dynamic_shape(
  123. self.option.shape_info_filename, True
  124. )
  125. elif self.option.device == "npu":
  126. config.enable_custom_device("npu")
  127. elif self.option.device == "xpu":
  128. pass
  129. elif self.option.device == "mlu":
  130. config.enable_custom_device("mlu")
  131. elif self.option.device == "gcu":
  132. assert paddle.device.is_compiled_with_custom_device("gcu"), (
  133. "Args device cannot be set as gcu while your paddle "
  134. "is not compiled with gcu!"
  135. )
  136. config.enable_custom_device("gcu")
  137. from paddle_custom_device.gcu import passes as gcu_passes
  138. gcu_passes.setUp()
  139. name = "PaddleX_" + self.option.model_name
  140. if hasattr(config, "enable_new_ir") and self.option.enable_new_ir:
  141. config.enable_new_ir(True)
  142. config.enable_new_executor(True)
  143. kPirGcuPasses = gcu_passes.inference_passes(use_pir=True, name=name)
  144. config.enable_custom_passes(kPirGcuPasses, True)
  145. else:
  146. config.enable_new_ir(False)
  147. config.enable_new_executor(False)
  148. pass_builder = config.pass_builder()
  149. gcu_passes.append_passes_for_legacy_ir(pass_builder, name)
  150. else:
  151. assert self.option.device == "cpu"
  152. config.disable_gpu()
  153. if "mkldnn" in self.option.run_mode:
  154. try:
  155. config.enable_mkldnn()
  156. if "bf16" in self.option.run_mode:
  157. config.enable_mkldnn_bfloat16()
  158. except Exception as e:
  159. logging.warning(
  160. "MKL-DNN is not available. We will disable MKL-DNN."
  161. )
  162. config.set_mkldnn_cache_capacity(-1)
  163. else:
  164. if hasattr(config, "disable_mkldnn"):
  165. config.disable_mkldnn()
  166. # Disable paddle inference logging
  167. config.disable_glog_info()
  168. config.set_cpu_math_library_num_threads(self.option.cpu_threads)
  169. if not (
  170. self.option.device == "gpu" and self.option.run_mode.startswith("trt")
  171. ):
  172. if self.option.device in ("cpu", "gpu"):
  173. if hasattr(config, "enable_new_ir"):
  174. config.enable_new_ir(self.option.enable_new_ir)
  175. config.set_optimization_level(3)
  176. if hasattr(config, "enable_new_executor"):
  177. config.enable_new_executor()
  178. for del_p in self.option.delete_pass:
  179. config.delete_pass(del_p)
  180. if self.option.device in ("gpu", "dcu"):
  181. if paddle.is_compiled_with_rocm():
  182. # Delete unsupported passes in dcu
  183. config.delete_pass("conv2d_add_act_fuse_pass")
  184. config.delete_pass("conv2d_add_fuse_pass")
  185. predictor = create_predictor(config)
  186. # Get input and output handlers
  187. input_names = predictor.get_input_names()
  188. input_names.sort()
  189. input_handlers = []
  190. output_handlers = []
  191. for input_name in input_names:
  192. input_handler = predictor.get_input_handle(input_name)
  193. input_handlers.append(input_handler)
  194. output_names = predictor.get_output_names()
  195. for output_name in output_names:
  196. output_handler = predictor.get_output_handle(output_name)
  197. output_handlers.append(output_handler)
  198. return predictor, input_handlers, output_handlers
  199. def apply(self, **kwargs):
  200. if self.option.changed:
  201. self._reset()
  202. batches = self.to_batch(**kwargs)
  203. self.copy2gpu.apply(batches)
  204. self.infer.apply()
  205. pred = self.copy2cpu.apply()
  206. return self.format_output(pred)
  207. @property
  208. def sub_cmps(self):
  209. return {
  210. "Copy2GPU": self.copy2gpu,
  211. "Infer": self.infer,
  212. "Copy2CPU": self.copy2cpu,
  213. }
  214. @abstractmethod
  215. def to_batch(self):
  216. raise NotImplementedError
  217. @abstractmethod
  218. def format_output(self, pred):
  219. return [{"pred": res} for res in zip(*pred)]
  220. class ImagePredictor(BasePaddlePredictor):
  221. INPUT_KEYS = "img"
  222. OUTPUT_KEYS = "pred"
  223. DEAULT_INPUTS = {"img": "img"}
  224. DEAULT_OUTPUTS = {"pred": "pred"}
  225. def to_batch(self, img):
  226. return [np.stack(img, axis=0).astype(dtype=np.float32, copy=False)]
  227. def format_output(self, pred):
  228. return [{"pred": res} for res in zip(*pred)]
  229. class ImageDetPredictor(BasePaddlePredictor):
  230. INPUT_KEYS = [
  231. ["img", "scale_factors"],
  232. ["img", "scale_factors", "img_size"],
  233. ["img", "img_size"],
  234. ]
  235. OUTPUT_KEYS = [["boxes"], ["boxes", "masks"]]
  236. DEAULT_INPUTS = {"img": "img", "scale_factors": "scale_factors"}
  237. DEAULT_OUTPUTS = None
  238. def to_batch(self, img, scale_factors=[[1.0, 1.0]], img_size=None):
  239. scale_factors = [scale_factor[::-1] for scale_factor in scale_factors]
  240. if img_size is None:
  241. return [
  242. np.stack(img, axis=0).astype(dtype=np.float32, copy=False),
  243. np.stack(scale_factors, axis=0).astype(dtype=np.float32, copy=False),
  244. ]
  245. else:
  246. img_size = [img_size[::-1] for img_size in img_size]
  247. return [
  248. np.stack(img_size, axis=0).astype(dtype=np.float32, copy=False),
  249. np.stack(img, axis=0).astype(dtype=np.float32, copy=False),
  250. np.stack(scale_factors, axis=0).astype(dtype=np.float32, copy=False),
  251. ]
  252. def format_output(self, pred):
  253. box_idx_start = 0
  254. pred_box = []
  255. if len(pred) == 4:
  256. # Adapt to SOLOv2
  257. pred_class_id = []
  258. pred_mask = []
  259. pred_class_id.append([pred[1], pred[2]])
  260. pred_mask.append(pred[3])
  261. return [
  262. {
  263. "class_id": np.array(pred_class_id[i]),
  264. "masks": np.array(pred_mask[i]),
  265. }
  266. for i in range(len(pred_class_id))
  267. ]
  268. if len(pred) == 3:
  269. # Adapt to Instance Segmentation
  270. pred_mask = []
  271. for idx in range(len(pred[1])):
  272. np_boxes_num = pred[1][idx]
  273. box_idx_end = box_idx_start + np_boxes_num
  274. np_boxes = pred[0][box_idx_start:box_idx_end]
  275. pred_box.append(np_boxes)
  276. if len(pred) == 3:
  277. np_masks = pred[2][box_idx_start:box_idx_end]
  278. pred_mask.append(np_masks)
  279. box_idx_start = box_idx_end
  280. if len(pred) == 3:
  281. return [
  282. {"boxes": np.array(pred_box[i]), "masks": np.array(pred_mask[i])}
  283. for i in range(len(pred_box))
  284. ]
  285. else:
  286. return [{"boxes": np.array(res)} for res in pred_box]
  287. class TSPPPredictor(BasePaddlePredictor):
  288. INPUT_KEYS = "ts"
  289. OUTPUT_KEYS = "pred"
  290. DEAULT_INPUTS = {"ts": "ts"}
  291. DEAULT_OUTPUTS = {"pred": "pred"}
  292. def to_batch(self, ts):
  293. n = len(ts[0])
  294. x = [np.stack([lst[i] for lst in ts], axis=0) for i in range(n)]
  295. return x
  296. def format_output(self, pred):
  297. return [{"pred": res} for res in zip(*pred)]