From 2e2f92a5b10d1c5cf7b1d5384bc3c7db5e6ed25b Mon Sep 17 00:00:00 2001 From: Pei Yang Date: Wed, 20 Nov 2019 12:18:37 +0800 Subject: [PATCH] fix trt weight bug (#21231) added splitter "__" between weight name and suffix number to avoid conflicts. --- paddle/fluid/inference/tensorrt/engine.cc | 3 ++- paddle/fluid/inference/tensorrt/engine.h | 3 ++- 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/paddle/fluid/inference/tensorrt/engine.cc b/paddle/fluid/inference/tensorrt/engine.cc index 098aa5ccfdf..6f66e8d972c 100644 --- a/paddle/fluid/inference/tensorrt/engine.cc +++ b/paddle/fluid/inference/tensorrt/engine.cc @@ -211,7 +211,8 @@ float *TensorRTEngine::GetWeightCPUData(const std::string &name, const std::vector &scale) { static int name_suffix_counter = 0; std::string name_suffix = std::to_string(name_suffix_counter); - std::string name_with_suffix = name + name_suffix; + std::string splitter = "__"; + std::string name_with_suffix = name + splitter + name_suffix; auto w_dims = weight_tensor->dims(); platform::CPUPlace cpu_place; PADDLE_ENFORCE_EQ( diff --git a/paddle/fluid/inference/tensorrt/engine.h b/paddle/fluid/inference/tensorrt/engine.h index c1d950035ca..50857674fc7 100644 --- a/paddle/fluid/inference/tensorrt/engine.h +++ b/paddle/fluid/inference/tensorrt/engine.h @@ -159,7 +159,8 @@ class TensorRTEngine { std::unique_ptr w_tensor) { static int suffix_counter = 0; std::string suffix = std::to_string(suffix_counter); - weight_map[w_name + suffix] = std::move(w_tensor); + std::string splitter = "__"; + weight_map[w_name + splitter + suffix] = std::move(w_tensor); suffix_counter += 1; } -- GitLab