From 7726883196da00ce5c24a7220fba632d93d9b2dd Mon Sep 17 00:00:00 2001 From: Pei Yang Date: Fri, 29 Nov 2019 18:55:07 +0800 Subject: [PATCH] fix trt weight bug (#21231) (#21443) 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 f806069b476..85722c94b25 100644 --- a/paddle/fluid/inference/tensorrt/engine.cc +++ b/paddle/fluid/inference/tensorrt/engine.cc @@ -192,7 +192,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