From 8d0922ed30a088781983a0a42254f6e2ad624f92 Mon Sep 17 00:00:00 2001
From: Wilber <jiweibo@baidu.com>
Date: Thu, 23 Sep 2021 16:32:24 +0800
Subject: [PATCH] fix trt problem (#35938)

---
 cmake/external/lite.cmake                     |  2 +-
 .../fluid/inference/api/analysis_predictor.cc | 21 ++++++++++++++++---
 paddle/fluid/inference/tensorrt/engine.h      |  8 +++++++
 3 files changed, 27 insertions(+), 4 deletions(-)

diff --git a/cmake/external/lite.cmake b/cmake/external/lite.cmake
index ffadf7bea93..1a5c796f210 100644
--- a/cmake/external/lite.cmake
+++ b/cmake/external/lite.cmake
@@ -50,7 +50,7 @@ if (NOT LITE_SOURCE_DIR OR NOT LITE_BINARY_DIR)
   set(LITE_INSTALL_DIR ${THIRD_PARTY_PATH}/install/lite)
 
   if(NOT LITE_GIT_TAG)
-    set(LITE_GIT_TAG d3a3a6931b6d22d504d21ba32b3ae972770e9204)
+    set(LITE_GIT_TAG 4ab64daecc11fbf74fffdc6a4733f388472e7d5d)
   endif()
 
   if(NOT CUDA_ARCH_NAME)
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
index 6ebb2193e21..804f035a2e2 100644
--- a/paddle/fluid/inference/api/analysis_predictor.cc
+++ b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -686,9 +686,24 @@ void AnalysisPredictor::OptimizeInferenceProgram() {
 // Note, please do NOT use any member variables, because member variables may
 // have been destructed in multiple threads.
 #if PADDLE_WITH_TENSORRT
-        paddle::inference::Singleton<
-            inference::tensorrt::TRTEngineManager>::Global()
-            .DeleteAll();
+        auto &block = prog->Block(0);
+        for (auto &op_desc : block.AllOps()) {
+          if (op_desc->Type() == "tensorrt_engine") {
+            std::string engine_key =
+                BOOST_GET_CONST(std::string, op_desc->GetAttr("engine_key"));
+            int engine_predictor_id =
+                BOOST_GET_CONST(int, op_desc->GetAttr("predictor_id"));
+            std::string engine_name =
+                engine_key + std::to_string(engine_predictor_id);
+            if (paddle::inference::Singleton<
+                    inference::tensorrt::TRTEngineManager>::Global()
+                    .Has(engine_name)) {
+              paddle::inference::Singleton<
+                  inference::tensorrt::TRTEngineManager>::Global()
+                  .DeleteKey(engine_name);
+            }
+          }
+        }
 #endif
         delete prog;
       });
diff --git a/paddle/fluid/inference/tensorrt/engine.h b/paddle/fluid/inference/tensorrt/engine.h
index 29324f29006..e22c2488d3b 100644
--- a/paddle/fluid/inference/tensorrt/engine.h
+++ b/paddle/fluid/inference/tensorrt/engine.h
@@ -631,6 +631,14 @@ class TRTEngineManager {
     }
   }
 
+  void DeleteKey(const std::string& key) {
+    auto iter = engines_.find(key);
+    if (iter != engines_.end()) {
+      iter->second.reset(nullptr);
+      engines_.erase(iter);
+    }
+  }
+
  private:
   std::unordered_map<std::string, std::unique_ptr<TensorRTEngine>> engines_;
 };
-- 
GitLab