static_infer.py 8.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233
  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. import os
  16. import inspect
  17. from abc import abstractmethod
  18. import lazy_paddle as paddle
  19. import numpy as np
  20. from ....utils.flags import FLAGS_json_format_model
  21. from ....utils import logging
  22. from ...utils.pp_option import PaddlePredictorOption
  23. class Copy2GPU:
  24. def __init__(self, input_handlers):
  25. super().__init__()
  26. self.input_handlers = input_handlers
  27. def __call__(self, x):
  28. for idx in range(len(x)):
  29. self.input_handlers[idx].reshape(x[idx].shape)
  30. self.input_handlers[idx].copy_from_cpu(x[idx])
  31. class Copy2CPU:
  32. def __init__(self, output_handlers):
  33. super().__init__()
  34. self.output_handlers = output_handlers
  35. def __call__(self):
  36. output = []
  37. for out_tensor in self.output_handlers:
  38. batch = out_tensor.copy_to_cpu()
  39. output.append(batch)
  40. return output
  41. class Infer:
  42. def __init__(self, predictor):
  43. super().__init__()
  44. self.predictor = predictor
  45. def __call__(self):
  46. self.predictor.run()
  47. class StaticInfer:
  48. """Predictor based on Paddle Inference"""
  49. def __init__(
  50. self, model_dir: str, model_prefix: str, option: PaddlePredictorOption
  51. ) -> None:
  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: PaddlePredictorOption) -> None:
  57. if self.option and option == self.option:
  58. return
  59. self._option = option
  60. self._reset()
  61. @property
  62. def option(self) -> PaddlePredictorOption:
  63. return self._option if hasattr(self, "_option") else None
  64. @option.setter
  65. def option(self, option: Union[None, PaddlePredictorOption]) -> None:
  66. if option:
  67. self._update_option(option)
  68. def _reset(self) -> None:
  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(
  82. self,
  83. ) -> Tuple[
  84. paddle.base.libpaddle.PaddleInferPredictor,
  85. paddle.base.libpaddle.PaddleInferTensor,
  86. paddle.base.libpaddle.PaddleInferTensor,
  87. ]:
  88. """_create"""
  89. from lazy_paddle.inference import Config, create_predictor
  90. model_postfix = ".json" if FLAGS_json_format_model else ".pdmodel"
  91. model_file = (self.model_dir / f"{self.model_prefix}{model_postfix}").as_posix()
  92. params_file = (self.model_dir / f"{self.model_prefix}.pdiparams").as_posix()
  93. config = Config(model_file, params_file)
  94. config.enable_memory_optim()
  95. if self.option.device in ("gpu", "dcu"):
  96. if self.option.device == "gpu":
  97. config.exp_disable_mixed_precision_ops({"feed", "fetch"})
  98. config.enable_use_gpu(100, self.option.device_id)
  99. if self.option.device == "gpu":
  100. # NOTE: The pptrt settings are not aligned with those of FD.
  101. precision_map = {
  102. "trt_int8": Config.Precision.Int8,
  103. "trt_fp32": Config.Precision.Float32,
  104. "trt_fp16": Config.Precision.Half,
  105. }
  106. if self.option.run_mode in precision_map.keys():
  107. config.enable_tensorrt_engine(
  108. workspace_size=(1 << 25) * self.option.batch_size,
  109. max_batch_size=self.option.batch_size,
  110. min_subgraph_size=self.option.min_subgraph_size,
  111. precision_mode=precision_map[self.option.run_mode],
  112. use_static=self.option.trt_use_static,
  113. use_calib_mode=self.option.trt_calib_mode,
  114. )
  115. if self.option.shape_info_filename is not None:
  116. if not os.path.exists(self.option.shape_info_filename):
  117. config.collect_shape_range_info(
  118. self.option.shape_info_filename
  119. )
  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. elif self.option.device == "npu":
  132. config.enable_custom_device("npu")
  133. elif self.option.device == "xpu":
  134. pass
  135. elif self.option.device == "mlu":
  136. config.enable_custom_device("mlu")
  137. else:
  138. assert self.option.device == "cpu"
  139. config.disable_gpu()
  140. if "mkldnn" in self.option.run_mode:
  141. try:
  142. config.enable_mkldnn()
  143. if "bf16" in self.option.run_mode:
  144. config.enable_mkldnn_bfloat16()
  145. except Exception as e:
  146. logging.warning(
  147. "MKL-DNN is not available. We will disable MKL-DNN."
  148. )
  149. config.set_mkldnn_cache_capacity(-1)
  150. else:
  151. if hasattr(config, "disable_mkldnn"):
  152. config.disable_mkldnn()
  153. # Disable paddle inference logging
  154. config.disable_glog_info()
  155. config.set_cpu_math_library_num_threads(self.option.cpu_threads)
  156. if self.option.device in ("cpu", "gpu"):
  157. if not (
  158. self.option.device == "gpu" and self.option.run_mode.startswith("trt")
  159. ):
  160. if hasattr(config, "enable_new_ir"):
  161. config.enable_new_ir(self.option.enable_new_ir)
  162. if hasattr(config, "enable_new_executor"):
  163. config.enable_new_executor()
  164. config.set_optimization_level(3)
  165. for del_p in self.option.delete_pass:
  166. config.delete_pass(del_p)
  167. if self.option.device in ("gpu", "dcu"):
  168. if paddle.is_compiled_with_rocm():
  169. # Delete unsupported passes in dcu
  170. config.delete_pass("conv2d_add_act_fuse_pass")
  171. config.delete_pass("conv2d_add_fuse_pass")
  172. predictor = create_predictor(config)
  173. # Get input and output handlers
  174. input_names = predictor.get_input_names()
  175. input_names.sort()
  176. input_handlers = []
  177. output_handlers = []
  178. for input_name in input_names:
  179. input_handler = predictor.get_input_handle(input_name)
  180. input_handlers.append(input_handler)
  181. output_names = predictor.get_output_names()
  182. for output_name in output_names:
  183. output_handler = predictor.get_output_handle(output_name)
  184. output_handlers.append(output_handler)
  185. return predictor, input_handlers, output_handlers
  186. def __call__(self, x) -> List[Any]:
  187. if self.option.changed:
  188. self._reset()
  189. self.copy2gpu(x)
  190. self.infer()
  191. pred = self.copy2cpu()
  192. return pred
  193. @property
  194. def benchmark(self):
  195. return {
  196. "Copy2GPU": self.copy2gpu,
  197. "Infer": self.infer,
  198. "Copy2CPU": self.copy2cpu,
  199. }