From 7cb7535e152d44894c7065f65cc3cfead7fc53a1 Mon Sep 17 00:00:00 2001 From: wenbin Date: Wed, 27 Oct 2021 20:16:34 +0800 Subject: [PATCH] fix ernie serialize problem (#36769) (#36791) Co-authored-by: zlsh80826 --- paddle/fluid/inference/tensorrt/engine.cc | 6 +++--- paddle/fluid/inference/tensorrt/engine.h | 7 +++++++ 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/paddle/fluid/inference/tensorrt/engine.cc b/paddle/fluid/inference/tensorrt/engine.cc index e151ef124a7..f20ef14177a 100644 --- a/paddle/fluid/inference/tensorrt/engine.cc +++ b/paddle/fluid/inference/tensorrt/engine.cc @@ -239,11 +239,11 @@ void TensorRTEngine::FreezeNetwork() { *network(), *infer_builder_config_)); #else infer_builder_config_->setFlag(nvinfer1::BuilderFlag::kSPARSE_WEIGHTS); - infer_ptr plan(infer_builder_->buildSerializedNetwork( + ihost_memory_.reset(infer_builder_->buildSerializedNetwork( *network(), *infer_builder_config_)); infer_ptr runtime(createInferRuntime(&logger_)); - infer_engine_.reset( - runtime->deserializeCudaEngine(plan->data(), plan->size())); + infer_engine_.reset(runtime->deserializeCudaEngine(ihost_memory_->data(), + ihost_memory_->size())); #endif PADDLE_ENFORCE_NOT_NULL( diff --git a/paddle/fluid/inference/tensorrt/engine.h b/paddle/fluid/inference/tensorrt/engine.h index 8cb0d36277d..c7a5eca5a11 100644 --- a/paddle/fluid/inference/tensorrt/engine.h +++ b/paddle/fluid/inference/tensorrt/engine.h @@ -273,7 +273,14 @@ class TensorRTEngine { infer_engine_, platform::errors::InvalidArgument( "The TensorRT engine must be built first before serialization")); +#if IS_TRT_VERSION_LT(8000) ihost_memory_.reset(infer_engine_->serialize()); +#else + PADDLE_ENFORCE_NOT_NULL( + ihost_memory_, + platform::errors::InvalidArgument( + "TensorRT >= 8.0 requires that buildSerializedNetwork is called")); +#endif return ihost_memory_.get(); } -- GitLab