未验证 提交 77268831 编写于 作者: P Pei Yang 提交者: GitHub

fix trt weight bug (#21231) (#21443)

added splitter "__" between weight name and suffix number to avoid conflicts.
上级 6ce49eea
...@@ -192,7 +192,8 @@ float *TensorRTEngine::GetWeightCPUData(const std::string &name, ...@@ -192,7 +192,8 @@ float *TensorRTEngine::GetWeightCPUData(const std::string &name,
const std::vector<float> &scale) { const std::vector<float> &scale) {
static int name_suffix_counter = 0; static int name_suffix_counter = 0;
std::string name_suffix = std::to_string(name_suffix_counter); 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(); auto w_dims = weight_tensor->dims();
platform::CPUPlace cpu_place; platform::CPUPlace cpu_place;
PADDLE_ENFORCE_EQ( PADDLE_ENFORCE_EQ(
......
...@@ -159,7 +159,8 @@ class TensorRTEngine { ...@@ -159,7 +159,8 @@ class TensorRTEngine {
std::unique_ptr<framework::Tensor> w_tensor) { std::unique_ptr<framework::Tensor> w_tensor) {
static int suffix_counter = 0; static int suffix_counter = 0;
std::string suffix = std::to_string(suffix_counter); 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; suffix_counter += 1;
} }
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册