engine.cc 23.6 KB
Newer Older
Y
Yan Chunwei 已提交
1 2
/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.

N
nhzlx 已提交
3 4
Licensed under the Apache License, Version 2.0 (the "License"); you may not use
this file except in compliance with the License.
Y
Yan Chunwei 已提交
5 6 7 8 9 10 11 12 13 14 15 16 17 18
You may obtain a copy of the License at

http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License. */

#include "paddle/fluid/inference/tensorrt/engine.h"

#include <NvInfer.h>
#include <glog/logging.h>
19

A
Abhinav Arora 已提交
20
#include <string>
W
wanghuancoder 已提交
21

22
#include "NvInferRuntimeCommon.h"
23
#include "cuda_runtime_api.h"  // NOLINT
Y
Yan Chunwei 已提交
24
#include "paddle/fluid/inference/tensorrt/helper.h"
25
#include "paddle/fluid/platform/device/gpu/gpu_info.h"
Y
Yan Chunwei 已提交
26
#include "paddle/fluid/platform/enforce.h"
27
#include "paddle/phi/common/data_type.h"
Y
Yan Chunwei 已提交
28 29 30 31 32

namespace paddle {
namespace inference {
namespace tensorrt {

33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61
void TensorRTEngine::Weight::SetDataType(phi::DataType type) {
  nvinfer1::DataType nv_type;
  switch (type) {
    case phi::DataType::FLOAT32:
      nv_type = nvinfer1::DataType::kFLOAT;
      break;
    case phi::DataType::FLOAT16:
      nv_type = nvinfer1::DataType::kHALF;
      break;
    case phi::DataType::INT32:
      nv_type = nvinfer1::DataType::kINT32;
      break;
    case phi::DataType::INT8:
      nv_type = nvinfer1::DataType::kINT8;
      break;
#if IS_TRT_VERSION_GE(7000)
    case phi::DataType::BOOL:
      nv_type = nvinfer1::DataType::kBOOL;
      break;
#endif
    default:
      paddle::platform::errors::InvalidArgument(
          "Paddle-TRT loads weighths failed, found not supported data type %s.",
          type);
      break;
  }
  w_.type = nv_type;
}

62 63
int TensorRTEngine::runtime_batch_ = 1;

64 65 66 67 68
void TensorRTEngine::InitNetwork() {
  freshDeviceId();
  infer_builder_.reset(createInferBuilder(&logger_));

  if (with_dynamic_shape_) {
69
    infer_network_.reset(infer_builder_->createNetworkV2(
70 71 72
        1U << static_cast<int>(
            nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH)));
  } else {
73
    infer_network_.reset(infer_builder_->createNetworkV2(0U));
74
  }
75 76

  infer_builder_config_.reset(infer_builder_->createBuilderConfig());
W
wenbin 已提交
77 78 79 80
  // optim_profile_ = infer_builder_->createOptimizationProfile();
  optim_profiles_.resize(max_profile_num_);
  for (int i = 0; i < max_profile_num_; i++)
    optim_profiles_[i] = infer_builder_->createOptimizationProfile();
Y
Yan Chunwei 已提交
81 82
}

83 84
void TensorRTEngine::Execute(int batch_size,
                             std::vector<void *> *buffers,
85
                             cudaStream_t stream) {
N
nhzlx 已提交
86
  freshDeviceId();
87 88 89 90 91 92 93
  auto infer_context = context();
  if (!with_dynamic_shape()) {
    infer_context->enqueue(batch_size, buffers->data(), stream, nullptr);
  } else {
#if IS_TRT_VERSION_GE(6000)
    infer_context->enqueueV2(buffers->data(), stream, nullptr);
#endif
94
  }
N
nhzlx 已提交
95 96 97
  SetRuntimeBatch(batch_size);
}

Y
Yan Chunwei 已提交
98
void TensorRTEngine::FreezeNetwork() {
N
nhzlx 已提交
99
  freshDeviceId();
100
  VLOG(3) << "TRT to freeze network";
101 102 103 104 105 106 107
  PADDLE_ENFORCE_NOT_NULL(infer_builder_,
                          platform::errors::InvalidArgument(
                              "Inference builder of TRT is null. Please make "
                              "sure you call InitNetwork first."));
  PADDLE_ENFORCE_NOT_NULL(network(),
                          platform::errors::InvalidArgument(
                              "Call InitNetwork first to initialize network."));
Y
Yan Chunwei 已提交
108 109
  // build engine.
  infer_builder_->setMaxBatchSize(max_batch_);
110 111
  infer_builder_config_->setMaxWorkspaceSize(max_workspace_);

Z
Zhaolong Xing 已提交
112 113 114
  bool enable_fp16 = (precision_ == AnalysisConfig::Precision::kHalf);
  if (enable_fp16) {
    bool support_fp16 = infer_builder_->platformHasFastFp16();
115
    infer_builder_config_->setFlag(nvinfer1::BuilderFlag::kFP16);
Z
Zhaolong Xing 已提交
116 117 118
    if (!support_fp16) {
      LOG(INFO) << "You specify FP16 mode, but the hardware do not support "
                   "FP16 speed up, use FP32 instead.";
119 120
    } else {
      LOG(INFO) << "Run Paddle-TRT FP16 mode";
Z
Zhaolong Xing 已提交
121 122 123
    }
  }

124
  bool enable_int8 = (precision_ == AnalysisConfig::Precision::kInt8);
Z
Zhaolong Xing 已提交
125
  if (enable_int8) {
C
csy0225 已提交
126 127 128
    if (!use_dla_) {
      infer_builder_config_->setFlag(nvinfer1::BuilderFlag::kFP16);
    }
129 130
    infer_builder_config_->setFlag(nvinfer1::BuilderFlag::kINT8);

131
    if (calibrator_) {
132
      infer_builder_config_->setInt8Calibrator(calibrator_);
133
    } else {
134
      infer_builder_config_->setInt8Calibrator(nullptr);
135 136 137 138 139 140 141 142 143

#if IS_TRT_VERSION_GE(5000)
      for (auto &quant_range : quant_dynamic_range_) {
        auto tensor = quant_range.first;
        float range = quant_range.second;
        tensor->setDynamicRange(-range, range);
      }

      std::unordered_set<nvinfer1::ITensor *> all_t;
144 145
      for (int i = 0; i < network()->getNbLayers(); i++) {
        auto layer = network()->getLayer(i);
146 147 148 149
        for (int j = 0; j < layer->getNbOutputs(); j++) {
          all_t.insert(layer->getOutput(j));
        }
      }
150

151 152
      for (int i = 0; i < network()->getNbInputs(); i++) {
        all_t.insert(network()->getInput(i));
153 154 155 156
      }

      for (auto &t : all_t) {
        if (!quant_dynamic_range_.count(t)) {
T
tianshuo78520a 已提交
157 158 159
          VLOG(3) << "We are in trt int8 mode(not calibration), scale not set"
                  << " for tensor " << t->getName()
                  << ", this might be ok when trt does not need this range";
160 161
        }
      }
162

163
#if IS_TRT_VERSION_GE(5122)
164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182
      auto layer_int8_fallback = [&](nvinfer1::ILayer *layer) -> bool {
        if (layer->getType() == nvinfer1::LayerType::kSHAPE) {
          return false;
        }
        bool all_int = true;
        for (int j = 0; j < layer->getNbInputs(); j++) {
          auto *temp_in = layer->getInput(j);
          if (temp_in->getType() != nvinfer1::DataType::kINT32) {
            all_int = false;
          }
        }
        for (int j = 0; j < layer->getNbOutputs(); j++) {
          auto *temp_out = layer->getOutput(j);
          if (temp_out->getType() != nvinfer1::DataType::kINT32) {
            all_int = false;
          }
        }
        if (all_int) return false;

183 184 185 186 187 188
        for (int j = 0; j < layer->getNbInputs(); j++) {
          auto *temp_in = layer->getInput(j);
          if (!temp_in->dynamicRangeIsSet()) {
            VLOG(1) << "Layer(Name: " << layer->getName()
                    << ") is set to float32 because its input("
                    << temp_in->getName() << ") doesn't have dynamic range.";
189
            return true;
190 191
          }
        }
192 193
        for (int j = 0; j < layer->getNbOutputs(); j++) {
          auto *temp_out = layer->getOutput(j);
194 195 196 197
          if (!temp_out->dynamicRangeIsSet()) {
            VLOG(1) << "Layer(Name: " << layer->getName()
                    << ") is set to float32 because its output("
                    << temp_out->getName() << ") doesn't have dynamic range.";
198
            return true;
199 200
          }
        }
201
        return false;
202 203 204 205 206
      };
      // If a layer's output is the network's output, or not all of its inputs
      // and outputs have scales,
      // this layer's precision and output type are set to float32.
      // This step has no effect if this layer is fused during TRT optimization.
207
      int layers_no_int8 = 0;
208 209
      for (int i = 0; i < network()->getNbLayers(); i++) {
        auto layer = network()->getLayer(i);
210
        if (layer_int8_fallback(layer)) {
211
          layer->setPrecision(nvinfer1::DataType::kFLOAT);
212
          ++layers_no_int8;
213
        }
214
      }
215 216 217 218 219 220 221
      // Disable int8 or build engine failed if all layers aren't int8
      if (layers_no_int8 == network()->getNbLayers()) {
        nvinfer1::BuilderFlags flags = infer_builder_config_->getFlags();
        flags = flags & ~(1U << static_cast<int>(nvinfer1::BuilderFlag::kINT8));
        // reset flags
        infer_builder_config_->setFlags(flags);
      }
222 223 224 225 226
#else
      LOG(WARNING) << "If your TensorRT version is lower than 5.1.2.2, you "
                      "must provide quantization scales for all tensors using "
                      "TRT to run.";
#endif
227 228
#endif
    }
N
nhzlx 已提交
229
  }
Y
Yan Chunwei 已提交
230

231 232 233 234 235 236 237 238 239 240 241 242
  // If model is mixed precision, then we should cast all float output to
  // float32 precision. Otherwise, we can not confirm the output precision of
  // the trt engine.
  if (model_precision_ != phi::DataType::FLOAT32) {
    for (int i = 0; i < network()->getNbOutputs(); ++i) {
      network()->getOutput(i)->setAllowedFormats(
          static_cast<nvinfer1::TensorFormats>(
              1 << static_cast<int>(nvinfer1::TensorFormat::kLINEAR)));
      network()->getOutput(i)->setType(nvinfer1::DataType::kFLOAT);
    }
  }

243 244 245 246 247 248 249 250 251 252 253 254 255 256 257
  if (use_dla_) {
    if (!enable_int8 && !enable_fp16) {
      LOG(WARNING) << "TensorRT DLA must be used with int8 or fp16, but you "
                      "set float32, so DLA is not used.";
    } else if (infer_builder_->getNbDLACores() == 0) {
      LOG(WARNING)
          << "TensorRT DLA is set by config, but your device does not have "
             "DLA, so DLA is not used.";
    } else {
      if (dla_core_ < 0 || dla_core_ >= infer_builder_->getNbDLACores()) {
        dla_core_ = 0;
        LOG(WARNING) << "Invalid DLACore, must be 0 < DLACore < "
                     << infer_builder_->getNbDLACores() << ", but got "
                     << dla_core_ << ", so use use 0 as default.";
      }
258 259 260
      infer_builder_config_->setDefaultDeviceType(nvinfer1::DeviceType::kDLA);
      infer_builder_config_->setDLACore(dla_core_);
      infer_builder_config_->setFlag(nvinfer1::BuilderFlag::kGPU_FALLBACK);
261 262 263 264 265
      LOG(INFO) << "TensorRT DLA enabled in FreezeNetwork(), DLACore "
                << dla_core_;
    }
  }

266 267
  if (with_dynamic_shape_) {
#if IS_TRT_VERSION_GE(6000)
268
    LOG(INFO) << "Run Paddle-TRT Dynamic Shape mode.";
W
wenbin 已提交
269 270
    for (int i = 0; i < max_profile_num_; i++) {
      for (auto &input : min_input_shape_) {
271
#if IS_TRT_VERSION_LT(7000)
W
wenbin 已提交
272
        // trt6 will check all_of input > 0
273 274
        if (!(std::all_of(input.second.begin(),
                          input.second.end(),
W
wenbin 已提交
275 276 277 278 279 280 281 282 283
                          [](int x) { return x > 0; }) &&
              std::all_of(max_input_shape_[input.first].begin(),
                          max_input_shape_[input.first].end(),
                          [](int x) { return x > 0; }) &&
              std::all_of(optim_input_shape_[input.first].begin(),
                          optim_input_shape_[input.first].end(),
                          [](int x) { return x > 0; }))) {
          continue;
        }
284
#endif
W
wenbin 已提交
285 286 287 288 289 290
        VLOG(4) << "TRT dynamic_shape set " << input.first
                << " min: " << Vec2Str(input.second)
                << ", max: " << Vec2Str(max_input_shape_[input.first])
                << ", opt: " << Vec2Str(optim_input_shape_[input.first]);

        optim_profiles_[i]->setDimensions(
291 292
            input.first.c_str(),
            nvinfer1::OptProfileSelector::kMIN,
W
wenbin 已提交
293 294
            Vec2TRT_Dims(input.second, input.first, true));
        optim_profiles_[i]->setDimensions(
295 296
            input.first.c_str(),
            nvinfer1::OptProfileSelector::kMAX,
W
wenbin 已提交
297 298
            Vec2TRT_Dims(max_input_shape_[input.first], input.first, true));
        optim_profiles_[i]->setDimensions(
299 300
            input.first.c_str(),
            nvinfer1::OptProfileSelector::kOPT,
W
wenbin 已提交
301 302 303
            Vec2TRT_Dims(optim_input_shape_[input.first], input.first, true));
      }
      infer_builder_config_->addOptimizationProfile(optim_profiles_[i]);
304
    }
305 306 307 308 309 310
    if (WithFp16() && disable_trt_plugin_fp16()) {
      LOG(INFO) << "NOTE: In order to achieve higher accuracy, you have "
                   "disabled the fp16 mode of TRT Plugin,\n"
                << "you can reopen it with "
                   "'config.SetDynamicShapeInfo(min_shape, max_shape, "
                   "opt_shape, false /*disable_trt_plugin_fp16*/)'";
311
    }
312 313
#endif
  }
314
#if IS_TRT_VERSION_GE(8200)
315 316 317 318
  if (use_inspector_) {
    infer_builder_config_->setProfilingVerbosity(
        nvinfer1::ProfilingVerbosity::kDETAILED);
  }
319 320
#endif

321
#if IS_TRT_VERSION_LT(8000)
322 323
  infer_engine_.reset(infer_builder_->buildEngineWithConfig(
      *network(), *infer_builder_config_));
324
#else
J
JingZhuangzhuang 已提交
325
  infer_builder_config_->setFlag(nvinfer1::BuilderFlag::kSPARSE_WEIGHTS);
Z
zlsh80826 已提交
326
  ihost_memory_.reset(infer_builder_->buildSerializedNetwork(
327 328
      *network(), *infer_builder_config_));
  infer_ptr<nvinfer1::IRuntime> runtime(createInferRuntime(&logger_));
Z
zlsh80826 已提交
329 330
  infer_engine_.reset(runtime->deserializeCudaEngine(ihost_memory_->data(),
                                                     ihost_memory_->size()));
331
#endif
332

333
  PADDLE_ENFORCE_NOT_NULL(
334 335 336 337
      infer_engine_,
      platform::errors::Fatal(
          "Build TensorRT cuda engine failed! Please recheck "
          "you configurations related to paddle-TensorRT."));
338

W
wenbin 已提交
339 340 341 342 343 344 345
  binding_num_ = infer_engine_->getNbBindings();
  // reset status for dynamic shape clone
  if (max_profile_num_ > 1) {
    infer_context_.clear();
    cur_profile_num_ = 0;
  }

346
  GetEngineInfo();
Y
Yan Chunwei 已提交
347 348
}

349
nvinfer1::ITensor *TensorRTEngine::DeclareInput(const std::string &name,
Y
Yan Chunwei 已提交
350
                                                nvinfer1::DataType dtype,
351
                                                const nvinfer1::Dims &dims) {
352 353
  PADDLE_ENFORCE_EQ(network() != nullptr,
                    true,
354 355 356
                    platform::errors::InvalidArgument(
                        "The TRT network should be initialized first."));
  auto *input = network()->addInput(name.c_str(), dtype, dims);
357
  PADDLE_ENFORCE_NOT_NULL(
358 359 360 361 362 363 364
      input,
      platform::errors::InvalidArgument("Adding input %s failed in "
                                        "TensorRT inference network. "
                                        "Please recheck your input.",
                                        name));
  PADDLE_ENFORCE_EQ(input->isNetworkInput(),
                    true,
365 366 367 368
                    platform::errors::InvalidArgument(
                        "Input %s is not the input of TRT inference network. "
                        "Please recheck your input.",
                        name));
L
Luo Tao 已提交
369
  TensorRTEngine::SetITensor(name, input);
Y
Yan Chunwei 已提交
370 371 372
  return input;
}

373 374
void TensorRTEngine::DeclareOutput(const nvinfer1::ILayer *layer,
                                   int offset,
375 376
                                   const std::string &name) {
  auto *output = layer->getOutput(offset);
377
  SetITensor(name, output);
378
  PADDLE_ENFORCE_NOT_NULL(
379 380 381
      output,
      platform::errors::InvalidArgument(
          "The output %s of TRT engine should not be null.", name));
Y
Yan Chunwei 已提交
382
  output->setName(name.c_str());
383 384
  PADDLE_ENFORCE_EQ(output->isNetworkInput(),
                    false,
385 386 387 388
                    platform::errors::InvalidArgument(
                        "The output %s of TRT engine should not be the input "
                        "of the network at the same time.",
                        name));
389
  network()->markOutput(*output);
390
  PADDLE_ENFORCE_EQ(
391 392
      output->isNetworkOutput(),
      true,
393 394 395
      platform::errors::InvalidArgument(
          "The output %s of TRT engine should be the output of the network.",
          name));
N
nhzlx 已提交
396 397
}

398 399
void TensorRTEngine::DeclareOutput(const std::string &name) {
  auto *output = TensorRTEngine::GetITensor(name);
400
  PADDLE_ENFORCE_NOT_NULL(
401 402 403
      output,
      platform::errors::InvalidArgument(
          "The output %s of TRT engine should not be null.", name));
L
Luo Tao 已提交
404
  output->setName(name.c_str());
405 406
  PADDLE_ENFORCE_EQ(output->isNetworkInput(),
                    false,
407 408 409 410
                    platform::errors::InvalidArgument(
                        "The output %s of TRT engine should not be the input "
                        "of the network at the same time.",
                        name));
411
  network()->markOutput(*output);
L
Luo Tao 已提交
412
}
413 414 415 416 417 418 419 420 421 422 423 424 425
void TensorRTEngine::DeleteITensor(const std::string &name,
                                   nvinfer1::ITensor *tensor) {
  PADDLE_ENFORCE_NOT_NULL(
      tensor,
      platform::errors::InvalidArgument(
          "Tensor named %s of TRT engine should not be null.", name));
  PADDLE_ENFORCE_EQ(
      true,
      itensor_map_.count(name),
      platform::errors::InvalidArgument(
          "Tensor named %s of TRT engine should not be null", name));
  itensor_map_.erase(name);
}
L
Luo Tao 已提交
426

427 428
void TensorRTEngine::SetITensor(const std::string &name,
                                nvinfer1::ITensor *tensor) {
429
  PADDLE_ENFORCE_NOT_NULL(
430 431 432
      tensor,
      platform::errors::InvalidArgument(
          "Tensor named %s of TRT engine should not be null.", name));
433
  PADDLE_ENFORCE_EQ(
434 435
      0,
      itensor_map_.count(name),
436 437
      platform::errors::InvalidArgument(
          "Tensor named %s of TRT engine should not be duplicated", name));
L
Luo Tao 已提交
438 439 440
  itensor_map_[name] = tensor;
}

441
nvinfer1::ITensor *TensorRTEngine::GetITensor(const std::string &name) {
442 443
  PADDLE_ENFORCE_EQ(itensor_map_.count(name),
                    true,
444 445
                    platform::errors::NotFound(
                        "Tensor named %s is not found in TRT engine", name));
L
Luo Tao 已提交
446 447 448
  return itensor_map_[name];
}

449 450 451 452 453
std::unordered_map<std::string, nvinfer1::ITensor *>
    *TensorRTEngine::GetITensorMap() {
  return &itensor_map_;
}

454 455 456 457
void TensorRTEngine::SetRuntimeBatch(size_t batch_size) {
  runtime_batch_ = batch_size;
}

458 459 460 461 462 463
TensorRTEngine::Weight TensorRTEngine::GetFp32TrtWeight(
    const std::string &name, const framework::Tensor &weight_tensor) {
  static int name_suffix_counter = 0;
  std::string name_suffix = std::to_string(name_suffix_counter);
  std::string splitter = "__";
  std::string name_with_suffix = name + splitter + name_suffix;
464
  platform::CPUPlace cpu_place;
465 466 467 468 469 470 471 472 473 474 475 476 477 478 479 480 481 482 483 484 485 486 487 488 489 490 491 492 493 494 495 496 497 498 499 500 501 502 503 504 505 506 507 508 509 510 511 512 513 514
  PADDLE_ENFORCE_EQ(weight_map.count(name_with_suffix),
                    0,
                    platform::errors::AlreadyExists(
                        "The weight named %s is set into the weight map "
                        "twice in TRT OP converter.",
                        name_with_suffix));
  weight_map[name_with_suffix].reset(new framework::Tensor());
  weight_map[name_with_suffix]->Resize(weight_tensor.dims());

  TensorRTEngine::Weight weight;
  weight.SetCount(weight_tensor.numel());
  weight.SetDataType(nvinfer1::DataType::kFLOAT);
  // weight_tensor.dims().;

  // if trt not support dtype, we need to cast to  fp32.
  if (weight_tensor.dtype() == phi::DataType::BFLOAT16) {
    framework::Tensor bf16_tensor;
    bf16_tensor.clear();
    paddle::framework::TensorCopySync(
        weight_tensor, platform::CPUPlace(), &bf16_tensor);
    weight_map[name_with_suffix]->set_type(
        paddle::experimental::DataType::FLOAT32);
    weight_map[name_with_suffix]->Resize(weight_tensor.dims());
    auto *fp32_data =
        weight_map[name_with_suffix]->mutable_data<float>(platform::CPUPlace());
    auto *bf16_data = bf16_tensor.mutable_data<bfloat16>(platform::CPUPlace());
    for (int i = 0; i < weight_tensor.numel(); i++) {
      fp32_data[i] = static_cast<float>(bf16_data[i]);
    }
  } else if (weight_tensor.dtype() == phi::DataType::FLOAT16) {
    framework::Tensor fp16_tensor;
    fp16_tensor.clear();
    paddle::framework::TensorCopySync(
        weight_tensor, platform::CPUPlace(), &fp16_tensor);
    weight_map[name_with_suffix]->set_type(
        paddle::experimental::DataType::FLOAT32);
    weight_map[name_with_suffix]->Resize(weight_tensor.dims());
    auto *fp32_data =
        weight_map[name_with_suffix]->mutable_data<float>(platform::CPUPlace());
    auto *fp16_data = fp16_tensor.mutable_data<float16>(platform::CPUPlace());
    for (int i = 0; i < weight_tensor.numel(); i++) {
      fp32_data[i] = static_cast<float>(fp16_data[i]);
    }
  } else {
    paddle::framework::TensorCopySync(
        weight_tensor, cpu_place, weight_map[name_with_suffix].get());
  }
  weight.SetValues(weight_map[name_with_suffix]->data());
  name_suffix_counter += 1;
  return weight;
515 516
}

517 518 519 520 521 522 523 524 525 526 527 528 529 530 531 532 533 534 535 536 537 538 539 540 541 542 543 544 545 546 547 548 549 550 551 552 553 554 555 556 557 558 559 560 561 562 563 564 565 566 567 568 569 570 571 572 573
TensorRTEngine::Weight TensorRTEngine::GetTrtWeight(
    const std::string &name, const framework::Tensor &weight_tensor) {
  static int name_suffix_counter = 0;
  std::string name_suffix = std::to_string(name_suffix_counter);
  std::string splitter = "__";
  std::string name_with_suffix = name + splitter + name_suffix;
  platform::CPUPlace cpu_place;
  PADDLE_ENFORCE_EQ(weight_map.count(name_with_suffix),
                    0,
                    platform::errors::AlreadyExists(
                        "The weight named %s is set into the weight map "
                        "twice in TRT OP converter.",
                        name_with_suffix));

  weight_map[name_with_suffix].reset(new framework::Tensor());
  weight_map[name_with_suffix]->Resize(weight_tensor.dims());

  TensorRTEngine::Weight weight;
  weight.SetCount(weight_tensor.numel());

  // if trt not support dtype, we need to cast to fp32.
  if (weight_tensor.dtype() == phi::DataType::BFLOAT16) {
    framework::Tensor bf16_tensor;
    bf16_tensor.clear();
    paddle::framework::TensorCopySync(
        weight_tensor, platform::CPUPlace(), &bf16_tensor);
    weight_map[name_with_suffix]->set_type(
        paddle::experimental::DataType::FLOAT32);
    auto *fp32_data =
        weight_map[name_with_suffix]->mutable_data<float>(platform::CPUPlace());
    auto *bf16_data = bf16_tensor.mutable_data<bfloat16>(platform::CPUPlace());
    for (int i = 0; i < weight_tensor.numel(); i++) {
      fp32_data[i] = static_cast<float>(bf16_data[i]);
    }
    weight.SetDataType(phi::DataType::FLOAT32);
    weight.SetValues(fp32_data);
  } else if (weight_tensor.dtype() == phi::DataType::INT64) {
    framework::Tensor int64_tensor;
    int64_tensor.clear();
    paddle::framework::TensorCopySync(
        weight_tensor, platform::CPUPlace(), &int64_tensor);
    weight_map[name_with_suffix]->set_type(
        paddle::experimental::DataType::INT32);
    auto *int32_data =
        weight_map[name_with_suffix]->mutable_data<int>(platform::CPUPlace());
    auto *int64_data = int64_tensor.mutable_data<int64_t>(platform::CPUPlace());
    for (int i = 0; i < weight_tensor.numel(); i++) {
      int32_data[i] = int64_data[i];
    }
    weight.SetDataType(phi::DataType::FLOAT32);
    weight.SetValues(int32_data);
  } else {
    paddle::framework::TensorCopySync(
        weight_tensor, cpu_place, weight_map[name_with_suffix].get());
    weight.SetDataType(weight_tensor.dtype());
    weight.SetValues(weight_map[name_with_suffix]->data());
  }
574

575 576 577
  name_suffix_counter += 1;
  return weight;
}
578

579 580
int TensorRTEngine::GetRuntimeBatch() { return runtime_batch_; }

581
nvinfer1::IPluginV2Layer *TensorRTEngine::AddPlugin(
582 583
    nvinfer1::ITensor *const *inputs,
    int num_inputs,
584
    plugin::PluginTensorRT *plugin) {
585
  owned_plugin_.emplace_back(plugin);
586
  return network()->addPluginV2(inputs, num_inputs, *plugin);
587 588
}

589
nvinfer1::IPluginV2Layer *TensorRTEngine::AddPluginV2Ext(
590 591
    nvinfer1::ITensor *const *inputs,
    int num_inputs,
592 593 594 595 596
    plugin::PluginTensorRTV2Ext *plugin) {
  owned_plugin_v2ext_.emplace_back(plugin);
  return network()->addPluginV2(inputs, num_inputs, *plugin);
}

597
nvinfer1::IPluginV2Layer *TensorRTEngine::AddPluginV2IOExt(
598 599
    nvinfer1::ITensor *const *inputs,
    int num_inputs,
600 601 602 603 604
    nvinfer1::IPluginV2IOExt *plugin) {
  owned_plugin_v2ioext_.emplace_back(plugin);
  return network()->addPluginV2(inputs, num_inputs, *plugin);
}

N
nhzlx 已提交
605 606 607
void TensorRTEngine::freshDeviceId() {
  int count;
  cudaGetDeviceCount(&count);
608 609
  PADDLE_ENFORCE_LT(device_id_,
                    count,
610 611
                    platform::errors::OutOfRange(
                        "Device id %d exceeds the current device count: %d.",
612 613
                        device_id_,
                        count));
L
Leo Chen 已提交
614
  platform::SetDeviceId(device_id_);
N
nhzlx 已提交
615 616
}

617 618 619 620 621 622 623 624 625 626 627 628 629 630 631
void TensorRTEngine::GetEngineInfo() {
#if IS_TRT_VERSION_GE(8200)
  LOG(INFO) << "====== engine info ======";
  std::unique_ptr<nvinfer1::IEngineInspector> infer_inspector(
      infer_engine_->createEngineInspector());
  auto infer_context = context();
  infer_inspector->setExecutionContext(infer_context);
  LOG(INFO) << infer_inspector->getEngineInformation(
      nvinfer1::LayerInformationFormat::kONELINE);
  LOG(INFO) << "====== engine info end ======";
#else
  LOG(INFO) << "Inspector needs TensorRT version 8.2 and after.";
#endif
}

Y
Yan Chunwei 已提交
632 633 634
}  // namespace tensorrt
}  // namespace inference
}  // namespace paddle