pp_option.py 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160
  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 ...utils.func_register import FuncRegister
  15. from ...utils.device import parse_device, set_env_for_device, get_default_device
  16. from ...utils import logging
  17. from .new_ir_blacklist import NEWIR_BLOCKLIST
  18. class PaddlePredictorOption(object):
  19. """Paddle Inference Engine Option"""
  20. SUPPORT_RUN_MODE = (
  21. "paddle",
  22. "trt_fp32",
  23. "trt_fp16",
  24. "trt_int8",
  25. "mkldnn",
  26. "mkldnn_bf16",
  27. )
  28. SUPPORT_DEVICE = ("gpu", "cpu", "npu", "xpu", "mlu")
  29. _FUNC_MAP = {}
  30. register = FuncRegister(_FUNC_MAP)
  31. def __init__(self, model_name=None, **kwargs):
  32. super().__init__()
  33. self.model_name = model_name
  34. self._cfg = {}
  35. self._init_option(**kwargs)
  36. def _init_option(self, **kwargs):
  37. for k, v in kwargs.items():
  38. if k not in self._FUNC_MAP:
  39. raise Exception(
  40. f"{k} is not supported to set! The supported option is: \
  41. {list(self._FUNC_MAP.keys())}"
  42. )
  43. self._FUNC_MAP.get(k)(self, v)
  44. for k, v in self._get_default_config().items():
  45. self._cfg.setdefault(k, v)
  46. def _get_default_config(self):
  47. """get default config"""
  48. device_type, device_id = parse_device(get_default_device())
  49. return {
  50. "run_mode": "paddle",
  51. "device": device_type,
  52. "device_id": device_id[0],
  53. "min_subgraph_size": 3,
  54. "shape_info_filename": None,
  55. "trt_calib_mode": False,
  56. "cpu_threads": 1,
  57. "trt_use_static": False,
  58. "delete_pass": [],
  59. "enable_new_ir": True if self.model_name not in NEWIR_BLOCKLIST else False,
  60. }
  61. @register("run_mode")
  62. def set_run_mode(self, run_mode: str):
  63. """set run mode"""
  64. if run_mode not in self.SUPPORT_RUN_MODE:
  65. support_run_mode_str = ", ".join(self.SUPPORT_RUN_MODE)
  66. raise ValueError(
  67. f"`run_mode` must be {support_run_mode_str}, but received {repr(run_mode)}."
  68. )
  69. self._cfg["run_mode"] = run_mode
  70. @register("device")
  71. def set_device(self, device: str):
  72. """set device"""
  73. if not device:
  74. return
  75. device_type, device_ids = parse_device(device)
  76. self._cfg["device"] = device_type
  77. if device_type not in self.SUPPORT_DEVICE:
  78. support_run_mode_str = ", ".join(self.SUPPORT_DEVICE)
  79. raise ValueError(
  80. f"The device type must be one of {support_run_mode_str}, but received {repr(device_type)}."
  81. )
  82. device_id = device_ids[0] if device_ids is not None else 0
  83. self._cfg["device_id"] = device_id
  84. set_env_for_device(device)
  85. if device_type not in ("cpu"):
  86. if device_ids is None or len(device_ids) > 1:
  87. logging.warning(f"The device ID has been set to {device_id}.")
  88. @register("min_subgraph_size")
  89. def set_min_subgraph_size(self, min_subgraph_size: int):
  90. """set min subgraph size"""
  91. if not isinstance(min_subgraph_size, int):
  92. raise Exception()
  93. self._cfg["min_subgraph_size"] = min_subgraph_size
  94. @register("shape_info_filename")
  95. def set_shape_info_filename(self, shape_info_filename: str):
  96. """set shape info filename"""
  97. self._cfg["shape_info_filename"] = shape_info_filename
  98. @register("trt_calib_mode")
  99. def set_trt_calib_mode(self, trt_calib_mode):
  100. """set trt calib mode"""
  101. self._cfg["trt_calib_mode"] = trt_calib_mode
  102. @register("cpu_threads")
  103. def set_cpu_threads(self, cpu_threads):
  104. """set cpu threads"""
  105. if not isinstance(cpu_threads, int) or cpu_threads < 1:
  106. raise Exception()
  107. self._cfg["cpu_threads"] = cpu_threads
  108. @register("trt_use_static")
  109. def set_trt_use_static(self, trt_use_static):
  110. """set trt use static"""
  111. self._cfg["trt_use_static"] = trt_use_static
  112. @register("delete_pass")
  113. def set_delete_pass(self, delete_pass):
  114. self._cfg["delete_pass"] = delete_pass
  115. @register("enable_new_ir")
  116. def set_enable_new_ir(self, enable_new_ir: bool):
  117. """set run mode"""
  118. self._cfg["enable_new_ir"] = enable_new_ir
  119. def get_support_run_mode(self):
  120. """get supported run mode"""
  121. return self.SUPPORT_RUN_MODE
  122. def get_support_device(self):
  123. """get supported device"""
  124. return self.SUPPORT_DEVICE
  125. def get_device(self):
  126. """get device"""
  127. return f"{self._cfg['device']}:{self._cfg['device_id']}"
  128. def __str__(self):
  129. return ", ".join([f"{k}: {v}" for k, v in self._cfg.items()])
  130. def __getattr__(self, key):
  131. if key not in self._cfg:
  132. raise Exception(f"The key ({key}) is not found in cfg: \n {self._cfg}")
  133. return self._cfg.get(key)
  134. def __eq__(self, obj):
  135. if isinstance(obj, PaddlePredictorOption):
  136. return obj._cfg == self._cfg
  137. return False