From 9cf3aa6159dd73ecccd51f51d57be7850987f2a8 Mon Sep 17 00:00:00 2001 From: Zhang Jun Date: Wed, 16 Nov 2022 14:46:15 +0800 Subject: [PATCH] trt memory set change from setMaxWorkspaceSize to setMemoryPoolLimit since trt 8.3+ (#47795) --- paddle/fluid/inference/tensorrt/engine.cc | 6 +++++- paddle/fluid/inference/tensorrt/test_tensorrt.cc | 4 ++++ 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/paddle/fluid/inference/tensorrt/engine.cc b/paddle/fluid/inference/tensorrt/engine.cc index c6c0d0479c9..0b3c099934c 100644 --- a/paddle/fluid/inference/tensorrt/engine.cc +++ b/paddle/fluid/inference/tensorrt/engine.cc @@ -150,8 +150,12 @@ void TensorRTEngine::FreezeNetwork() { "Call InitNetwork first to initialize network.")); // build engine. infer_builder_->setMaxBatchSize(max_batch_); +#if IS_TRT_VERSION_GE(8300) + infer_builder_config_->setMemoryPoolLimit( + nvinfer1::MemoryPoolType::kWORKSPACE, max_workspace_); +#else infer_builder_config_->setMaxWorkspaceSize(max_workspace_); - +#endif bool enable_fp16 = (precision_ == AnalysisConfig::Precision::kHalf); if (enable_fp16) { bool support_fp16 = infer_builder_->platformHasFastFp16(); diff --git a/paddle/fluid/inference/tensorrt/test_tensorrt.cc b/paddle/fluid/inference/tensorrt/test_tensorrt.cc index 4a8bcb25656..571a2b56c56 100644 --- a/paddle/fluid/inference/tensorrt/test_tensorrt.cc +++ b/paddle/fluid/inference/tensorrt/test_tensorrt.cc @@ -95,7 +95,11 @@ nvinfer1::IHostMemory* CreateNetwork() { network->markOutput(*output); // Build the engine. builder->setMaxBatchSize(1); +#if IS_TRT_VERSION_GE(8300) + config->setMemoryPoolLimit(nvinfer1::MemoryPoolType::kWORKSPACE, 1 << 10); +#else config->setMaxWorkspaceSize(1 << 10); +#endif auto engine = builder->buildEngineWithConfig(*network, *config); EXPECT_NE(engine, nullptr); // Serialize the engine to create a model, then close. -- GitLab