add trt shape file to table and layout model

Frida-a · Sep 19, 2022 · 43838e3 · 43838e3
1 parent a166e54
commit 43838e3
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 2 deletions.
diff --git a/deploy/cpp_infer/include/structure_layout.h b/deploy/cpp_infer/include/structure_layout.h
@@ -88,7 +88,6 @@ class StructureLayoutRecognizer {
 
   // post-process
   PicodetPostProcessor post_processor_;
-
-}; // class StructureTableRecognizer
+};
 
 } // namespace PaddleOCR
diff --git a/deploy/cpp_infer/src/structure_layout.cpp b/deploy/cpp_infer/src/structure_layout.cpp
@@ -120,6 +120,11 @@ void StructureLayoutRecognizer::LoadModel(const std::string &model_dir) {
         precision = paddle_infer::Config::Precision::kInt8;
       }
       config.EnableTensorRtEngine(1 << 20, 10, 3, precision, false, false);
+      if (!Utility::PathExists("./trt_layout_shape.txt")) {
+        config.CollectShapeRangeInfo("./trt_layout_shape.txt");
+      } else {
+        config.EnableTunedTensorRtDynamicShape("./trt_layout_shape.txt", true);
+      }
     }
   } else {
     config.DisableGpu();

diff --git a/deploy/cpp_infer/src/structure_table.cpp b/deploy/cpp_infer/src/structure_table.cpp
@@ -133,6 +133,11 @@ void StructureTableRecognizer::LoadModel(const std::string &model_dir) {
         precision = paddle_infer::Config::Precision::kInt8;
       }
       config.EnableTensorRtEngine(1 << 20, 10, 3, precision, false, false);
+      if (!Utility::PathExists("./trt_table_shape.txt")) {
+        config.CollectShapeRangeInfo("./trt_table_shape.txt");
+      } else {
+        config.EnableTunedTensorRtDynamicShape("./trt_table_shape.txt", true);
+      }
     }
   } else {
     config.DisableGpu();