trt_config.py 3.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111
  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 collections import defaultdict
  15. import lazy_paddle
  16. from ...utils.flags import USE_PIR_TRT
  17. class LazyLoadDict(dict):
  18. def __init__(self, *args, **kwargs):
  19. self._initialized = False
  20. super().__init__(*args, **kwargs)
  21. def _initialize(self):
  22. if not self._initialized:
  23. self.update(self._load())
  24. self._initialized = True
  25. def __getitem__(self, key):
  26. self._initialize()
  27. return super().__getitem__(key)
  28. def __contains__(self, key):
  29. self._initialize()
  30. return super().__contains__(key)
  31. def _load(self):
  32. raise NotImplementedError
  33. class OLD_IR_TRT_PRECISION_MAP_CLASS(LazyLoadDict):
  34. def _load(self):
  35. from lazy_paddle.inference.Config import Precision
  36. return {
  37. "trt_int8": Precision.Int8,
  38. "trt_fp32": Precision.Float32,
  39. "trt_fp16": Precision.Half,
  40. }
  41. class PIR_TRT_PRECISION_MAP_CLASS(LazyLoadDict):
  42. def _load(self):
  43. from lazy_paddle.tensorrt.export import PrecisionMode
  44. return {
  45. "trt_int8": PrecisionMode.INT8,
  46. "trt_fp32": PrecisionMode.FP32,
  47. "trt_fp16": PrecisionMode.FP16,
  48. }
  49. ############ old ir trt ############
  50. OLD_IR_TRT_PRECISION_MAP = OLD_IR_TRT_PRECISION_MAP_CLASS()
  51. OLD_IR_TRT_DEFAULT_CFG = {
  52. "workspace_size": 1 << 30,
  53. "max_batch_size": 32,
  54. "min_subgraph_size": 3,
  55. "use_static": True,
  56. "use_calib_mode": False,
  57. }
  58. OLD_IR_TRT_CFG = {
  59. "SegFormer-B3": {**OLD_IR_TRT_DEFAULT_CFG, "workspace_size": 1 << 31},
  60. "SegFormer-B4": {**OLD_IR_TRT_DEFAULT_CFG, "workspace_size": 1 << 31},
  61. "SegFormer-B5": {**OLD_IR_TRT_DEFAULT_CFG, "workspace_size": 1 << 31},
  62. }
  63. ############ pir trt ############
  64. PIR_TRT_PRECISION_MAP = PIR_TRT_PRECISION_MAP_CLASS()
  65. PIR_TRT_CFG = {
  66. "DETR-R50": {"optimization_level": 4, "workspace_size": 1 << 32},
  67. "SegFormer-B0": {"optimization_level": 4, "workspace_size": 1 << 32},
  68. "SegFormer-B1": {"optimization_level": 4, "workspace_size": 1 << 32},
  69. "SegFormer-B2": {"optimization_level": 4, "workspace_size": 1 << 32},
  70. "SegFormer-B3": {"optimization_level": 4, "workspace_size": 1 << 32},
  71. "SegFormer-B4": {"optimization_level": 4, "workspace_size": 1 << 32},
  72. "SegFormer-B5": {"optimization_level": 4, "workspace_size": 1 << 32},
  73. "LaTeX_OCR_rec": {"disable_ops": ["pd_op.slice"]},
  74. "PP-YOLOE_seg-S": {"disable_ops": ["pd_op.slice", "pd_op.bilinear_interp"]},
  75. "PP-FormulaNet-L": {
  76. "disable_ops": ["pd_op.full_with_tensor"],
  77. "workspace_size": 1 << 32,
  78. },
  79. "PP-FormulaNet-S": {
  80. "disable_ops": ["pd_op.full_with_tensor"],
  81. "workspace_size": 1 << 32,
  82. },
  83. }
  84. if USE_PIR_TRT:
  85. TRT_PRECISION_MAP = PIR_TRT_PRECISION_MAP
  86. TRT_CFG = defaultdict(dict, PIR_TRT_CFG)
  87. else:
  88. TRT_PRECISION_MAP = OLD_IR_TRT_PRECISION_MAP
  89. TRT_CFG = defaultdict(lambda: OLD_IR_TRT_DEFAULT_CFG, OLD_IR_TRT_CFG)