فهرست منبع

change workspace_size to 30 in trt infer to be compatible with more models

gaotingquan 10 ماه پیش
والد
کامیت
8e9b8ddca2
2فایلهای تغییر یافته به همراه2 افزوده شده و 2 حذف شده
  1. 1 1
      paddlex/inference/components/paddle_predictor/predictor.py
  2. 1 1
      paddlex/inference/models/common/static_infer.py

+ 1 - 1
paddlex/inference/components/paddle_predictor/predictor.py

@@ -147,7 +147,7 @@ class BasePaddlePredictor(BaseComponent):
                 }
                 }
                 if self.option.run_mode in precision_map.keys():
                 if self.option.run_mode in precision_map.keys():
                     config.enable_tensorrt_engine(
                     config.enable_tensorrt_engine(
-                        workspace_size=(1 << 25) * self.option.batch_size,
+                        workspace_size=(1 << 30) * self.option.batch_size,
                         max_batch_size=self.option.batch_size,
                         max_batch_size=self.option.batch_size,
                         min_subgraph_size=self.option.min_subgraph_size,
                         min_subgraph_size=self.option.min_subgraph_size,
                         precision_mode=precision_map[self.option.run_mode],
                         precision_mode=precision_map[self.option.run_mode],

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

@@ -234,7 +234,7 @@ class StaticInfer:
                         "trt_fp16": Config.Precision.Half,
                         "trt_fp16": Config.Precision.Half,
                     }
                     }
                     config.enable_tensorrt_engine(
                     config.enable_tensorrt_engine(
-                        workspace_size=(1 << 25) * self.option.batch_size,
+                        workspace_size=(1 << 30) * self.option.batch_size,
                         max_batch_size=self.option.batch_size,
                         max_batch_size=self.option.batch_size,
                         min_subgraph_size=self.option.min_subgraph_size,
                         min_subgraph_size=self.option.min_subgraph_size,
                         precision_mode=precision_map[self.option.run_mode],
                         precision_mode=precision_map[self.option.run_mode],