소스 검색

set trt config for PP-YOLOE_seg-S

gaotingquan 8 달 전
부모
커밋
5abc0bfdca
2개의 변경된 파일1개의 추가작업 그리고 1개의 파일을 삭제
  1. 0 1
      paddlex/inference/models/common/static_infer.py
  2. 1 0
      paddlex/inference/utils/trt_config.py

+ 0 - 1
paddlex/inference/models/common/static_infer.py

@@ -182,7 +182,6 @@ def _convert_trt(
         "trt_fp16": PrecisionMode.FP16,
     }
     trt_inputs = []
-    # for name in input_names:
     for name, candidate_shapes in dynamic_shapes.items():
         # XXX: Currently we have no way to get the data type of the tensor
         # without creating an input handle.

+ 1 - 0
paddlex/inference/utils/trt_config.py

@@ -21,4 +21,5 @@ TRT_CFG = {
     "SegFormer-B4": {"optimization_level": 4, "workspace_size": 1 << 32},
     "SegFormer-B5": {"optimization_level": 4, "workspace_size": 1 << 32},
     "LaTeX_OCR_rec": {"disable_ops": ["pd_op.slice"]},
+    "PP-YOLOE_seg-S": {"disable_ops": ["pd_op.slice", "pd_op.bilinear_interp"]},
 }