瀏覽代碼

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

gaotingquan 10 月之前
父節點
當前提交
8e9b8ddca2

+ 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():
                     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,
                         min_subgraph_size=self.option.min_subgraph_size,
                         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,
                     }
                     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,
                         min_subgraph_size=self.option.min_subgraph_size,
                         precision_mode=precision_map[self.option.run_mode],