From cf68eafa5639ad358c6203a1842c277c37ec29bb Mon Sep 17 00:00:00 2001
From: littletomatodonkey <2120160898@bit.edu.cn>
Date: Tue, 15 Dec 2020 19:22:13 +0800
Subject: [PATCH] add cpp trt (#480)

---
 deploy/cpp_infer/include/cls.h        | 7 ++++++-
 deploy/cpp_infer/include/cls_config.h | 6 ++++++
 deploy/cpp_infer/readme.md            | 2 +-
 deploy/cpp_infer/readme_en.md         | 2 +-
 deploy/cpp_infer/src/cls.cpp          | 7 +++++++
 deploy/cpp_infer/src/main.cpp         | 1 +
 deploy/cpp_infer/tools/config.txt     | 2 ++
 7 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/deploy/cpp_infer/include/cls.h b/deploy/cpp_infer/include/cls.h
index f9549573..f55d3952 100644
--- a/deploy/cpp_infer/include/cls.h
+++ b/deploy/cpp_infer/include/cls.h
@@ -40,13 +40,16 @@ public:
                       const std::string &params_path, const bool &use_gpu,
                       const int &gpu_id, const int &gpu_mem,
                       const int &cpu_math_library_num_threads,
-                      const bool &use_mkldnn, const int &resize_short_size,
+                      const bool &use_mkldnn, const bool &use_tensorrt,
+                      const bool &use_fp16, const int &resize_short_size,
                       const int &crop_size) {
     this->use_gpu_ = use_gpu;
     this->gpu_id_ = gpu_id;
     this->gpu_mem_ = gpu_mem;
     this->cpu_math_library_num_threads_ = cpu_math_library_num_threads;
     this->use_mkldnn_ = use_mkldnn;
+    this->use_tensorrt_ = use_tensorrt;
+    this->use_fp16_ = use_fp16;
 
     this->resize_short_size_ = resize_short_size;
     this->crop_size_ = crop_size;
@@ -68,6 +71,8 @@ private:
   int gpu_mem_ = 4000;
   int cpu_math_library_num_threads_ = 4;
   bool use_mkldnn_ = false;
+  bool use_tensorrt_ = false;
+  bool use_fp16_ = false;
 
   std::vector<float> mean_ = {0.485f, 0.456f, 0.406f};
   std::vector<float> scale_ = {1 / 0.229f, 1 / 0.224f, 1 / 0.225f};
diff --git a/deploy/cpp_infer/include/cls_config.h b/deploy/cpp_infer/include/cls_config.h
index 4d40347f..7377573c 100644
--- a/deploy/cpp_infer/include/cls_config.h
+++ b/deploy/cpp_infer/include/cls_config.h
@@ -41,6 +41,9 @@ public:
 
     this->use_mkldnn = bool(stoi(config_map_["use_mkldnn"]));
 
+    this->use_tensorrt = bool(stoi(config_map_["use_tensorrt"]));
+    this->use_fp16 = bool(stoi(config_map_["use_fp16"]));
+
     this->cls_model_path.assign(config_map_["cls_model_path"]);
 
     this->cls_params_path.assign(config_map_["cls_params_path"]);
@@ -60,6 +63,9 @@ public:
 
   bool use_mkldnn = false;
 
+  bool use_tensorrt = false;
+  bool use_fp16 = false;
+
   std::string cls_model_path;
 
   std::string cls_params_path;
diff --git a/deploy/cpp_infer/readme.md b/deploy/cpp_infer/readme.md
index aad8d87d..db045adf 100644
--- a/deploy/cpp_infer/readme.md
+++ b/deploy/cpp_infer/readme.md
@@ -175,7 +175,7 @@ cmake .. \
     -DDEMO_NAME=clas_system \
     -DWITH_GPU=OFF \
     -DWITH_STATIC_LIB=OFF \
-    -DUSE_TENSORRT=OFF \
+    -DWITH_TENSORRT=OFF \
     -DOPENCV_DIR=${OPENCV_DIR} \
     -DCUDNN_LIB=${CUDNN_LIB_DIR} \
     -DCUDA_LIB=${CUDA_LIB_DIR} \
diff --git a/deploy/cpp_infer/readme_en.md b/deploy/cpp_infer/readme_en.md
index 56be7549..1d706afd 100644
--- a/deploy/cpp_infer/readme_en.md
+++ b/deploy/cpp_infer/readme_en.md
@@ -183,7 +183,7 @@ cmake .. \
     -DDEMO_NAME=ocr_system \
     -DWITH_GPU=OFF \
     -DWITH_STATIC_LIB=OFF \
-    -DUSE_TENSORRT=OFF \
+    -DWITH_TENSORRT=OFF \
     -DOPENCV_DIR=${OPENCV_DIR} \
     -DCUDNN_LIB=${CUDNN_LIB_DIR} \
     -DCUDA_LIB=${CUDA_LIB_DIR} \
diff --git a/deploy/cpp_infer/src/cls.cpp b/deploy/cpp_infer/src/cls.cpp
index 0755eed1..d304ff5e 100644
--- a/deploy/cpp_infer/src/cls.cpp
+++ b/deploy/cpp_infer/src/cls.cpp
@@ -23,6 +23,13 @@ void Classifier::LoadModel(const std::string &model_path,
 
   if (this->use_gpu_) {
     config.EnableUseGpu(this->gpu_mem_, this->gpu_id_);
+    if (this->use_tensorrt_) {
+      config.EnableTensorRtEngine(
+          1 << 20, 1, 3,
+          this->use_fp16_ ? paddle_infer::Config::Precision::kHalf
+                          : paddle_infer::Config::Precision::kFloat32,
+          false, false);
+    }
   } else {
     config.DisableGpu();
     if (this->use_mkldnn_) {
diff --git a/deploy/cpp_infer/src/main.cpp b/deploy/cpp_infer/src/main.cpp
index 9b562adc..8885cc23 100644
--- a/deploy/cpp_infer/src/main.cpp
+++ b/deploy/cpp_infer/src/main.cpp
@@ -62,6 +62,7 @@ int main(int argc, char **argv) {
   Classifier classifier(config.cls_model_path, config.cls_params_path,
                         config.use_gpu, config.gpu_id, config.gpu_mem,
                         config.cpu_math_library_num_threads, config.use_mkldnn,
+                        config.use_tensorrt, config.use_fp16,
                         config.resize_short_size, config.crop_size);
 
   double elapsed_time = 0.0;
diff --git a/deploy/cpp_infer/tools/config.txt b/deploy/cpp_infer/tools/config.txt
index 5277f366..b0becad6 100755
--- a/deploy/cpp_infer/tools/config.txt
+++ b/deploy/cpp_infer/tools/config.txt
@@ -4,6 +4,8 @@ gpu_id  0
 gpu_mem  4000
 cpu_math_library_num_threads  10
 use_mkldnn 1
+use_tensorrt 0
+use_fp16 0
 
 # cls config
 cls_model_path  /PaddleClas/inference/cls_infer.pdmodel
-- 
GitLab