Explorar o código

Fix reshape dim out kMAX_TENSOR_SIZE

zhink hai 8 meses
pai
achega
5e90bf09e8
Modificáronse 1 ficheiros con 40 adicións e 3 borrados
  1. 40 3
      paddlex/inference/utils/trt_config.py

+ 40 - 3
paddlex/inference/utils/trt_config.py

@@ -73,22 +73,40 @@ OLD_IR_TRT_CFG_DEFAULT_SETTING = {
 }
 
 OLD_IR_TRT_CFG_SETTING = {
+    "SegFormer-B0": {
+        "enable_tensorrt_engine": {
+            **OLD_IR_TRT_CFG_DEFAULT_SETTING,
+            "workspace_size": 1 << 32,
+        }
+    },
+    "SegFormer-B1": {
+        "enable_tensorrt_engine": {
+            **OLD_IR_TRT_CFG_DEFAULT_SETTING,
+            "workspace_size": 1 << 32,
+        }
+    },
+    "SegFormer-B2": {
+        "enable_tensorrt_engine": {
+            **OLD_IR_TRT_CFG_DEFAULT_SETTING,
+            "workspace_size": 1 << 32,
+        }
+    },
     "SegFormer-B3": {
         "enable_tensorrt_engine": {
             **OLD_IR_TRT_CFG_DEFAULT_SETTING,
-            "workspace_size": 1 << 31,
+            "workspace_size": 1 << 32,
         }
     },
     "SegFormer-B4": {
         "enable_tensorrt_engine": {
             **OLD_IR_TRT_CFG_DEFAULT_SETTING,
-            "workspace_size": 1 << 31,
+            "workspace_size": 1 << 32,
         }
     },
     "SegFormer-B5": {
         "enable_tensorrt_engine": {
             **OLD_IR_TRT_CFG_DEFAULT_SETTING,
-            "workspace_size": 1 << 31,
+            "workspace_size": 1 << 32,
         }
     },
     "SLANeXt_wired": {
@@ -129,6 +147,25 @@ OLD_IR_TRT_CFG_SETTING = {
             ["bilinear_interp_v2_1.tmp_0", "bilinear_interp_v2_1.tmp_0_slice_0"]
         ],
     },
+    "TiDE": {
+        "enable_tensorrt_engine": OLD_IR_TRT_CFG_DEFAULT_SETTING,
+        "exp_disable_tensorrt_ops": [
+            [
+                "reshape2_3.tmp_0",
+                "reshape2_2.tmp_0",
+                "reshape2_1.tmp_0",
+                "reshape2_0.tmp_0",
+            ]
+        ],
+    },
+    "Nonstationary": {
+        "enable_tensorrt_engine": OLD_IR_TRT_CFG_DEFAULT_SETTING,
+        "exp_disable_tensorrt_ops": [
+            [
+                "reshape2_13.tmp_0",
+            ]
+        ],
+    },
 }
 
 ############ pir trt ############