engine.cc 23.7 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
int TensorRTEngine::runtime_batch_ = 1;
thread_local int TensorRTEngine::predictor_id_per_thread = -1;

36
void TensorRTEngine::Weight::SetDataType(phi::DataType type) {
37
  nvinfer1::DataType nv_type = nvinfer1::DataType::kFLOAT;
38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64
  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;
}

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

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

  infer_builder_config_.reset(infer_builder_->createBuilderConfig());
W
wenbin 已提交
78 79 80 81
  // 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 已提交
82 83
}

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

Y
Yan Chunwei 已提交
97
void TensorRTEngine::FreezeNetwork() {
N
nhzlx 已提交
98
  freshDeviceId();
99
  VLOG(3) << "TRT to freeze network";
100 101 102 103 104 105 106
  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 已提交
107 108
  // build engine.
  infer_builder_->setMaxBatchSize(max_batch_);
109 110
  infer_builder_config_->setMaxWorkspaceSize(max_workspace_);

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

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

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

      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;
142 143
      for (int i = 0; i < network()->getNbLayers(); i++) {
        auto layer = network()->getLayer(i);
144 145 146 147
        for (int j = 0; j < layer->getNbOutputs(); j++) {
          all_t.insert(layer->getOutput(j));
        }
      }
148

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

      for (auto &t : all_t) {
        if (!quant_dynamic_range_.count(t)) {
T
tianshuo78520a 已提交
155 156 157
          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";
158 159 160
        }
      }
    }
N
nhzlx 已提交
161
  }
Y
Yan Chunwei 已提交
162

163 164 165 166 167 168 169 170 171 172 173 174
  // 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);
    }
  }

175 176 177 178 179 180 181 182 183 184 185 186 187 188 189
  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.";
      }
190 191 192
      infer_builder_config_->setDefaultDeviceType(nvinfer1::DeviceType::kDLA);
      infer_builder_config_->setDLACore(dla_core_);
      infer_builder_config_->setFlag(nvinfer1::BuilderFlag::kGPU_FALLBACK);
193 194 195 196 197
      LOG(INFO) << "TensorRT DLA enabled in FreezeNetwork(), DLACore "
                << dla_core_;
    }
  }

198
  if (with_dynamic_shape_) {
199
    LOG(INFO) << "Run Paddle-TRT Dynamic Shape mode.";
W
wenbin 已提交
200 201
    for (int i = 0; i < max_profile_num_; i++) {
      for (auto &input : min_input_shape_) {
202
#if IS_TRT_VERSION_LT(7000)
W
wenbin 已提交
203
        // trt6 will check all_of input > 0
204 205
        if (!(std::all_of(input.second.begin(),
                          input.second.end(),
W
wenbin 已提交
206 207 208 209 210 211 212 213 214
                          [](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;
        }
215
#endif
W
wenbin 已提交
216 217 218 219 220 221
        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(
222 223
            input.first.c_str(),
            nvinfer1::OptProfileSelector::kMIN,
W
wenbin 已提交
224 225
            Vec2TRT_Dims(input.second, input.first, true));
        optim_profiles_[i]->setDimensions(
226 227
            input.first.c_str(),
            nvinfer1::OptProfileSelector::kMAX,
W
wenbin 已提交
228 229
            Vec2TRT_Dims(max_input_shape_[input.first], input.first, true));
        optim_profiles_[i]->setDimensions(
230 231
            input.first.c_str(),
            nvinfer1::OptProfileSelector::kOPT,
W
wenbin 已提交
232 233 234
            Vec2TRT_Dims(optim_input_shape_[input.first], input.first, true));
      }
      infer_builder_config_->addOptimizationProfile(optim_profiles_[i]);
235
    }
236 237 238 239 240 241
    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*/)'";
242
    }
243
  }
244
#if IS_TRT_VERSION_GE(8200)
245 246 247 248
  if (use_inspector_) {
    infer_builder_config_->setProfilingVerbosity(
        nvinfer1::ProfilingVerbosity::kDETAILED);
  }
249 250
#endif

251
#if IS_TRT_VERSION_LT(8000)
252 253
  infer_engine_.reset(infer_builder_->buildEngineWithConfig(
      *network(), *infer_builder_config_));
254
#else
J
JingZhuangzhuang 已提交
255
  infer_builder_config_->setFlag(nvinfer1::BuilderFlag::kSPARSE_WEIGHTS);
Z
zlsh80826 已提交
256
  ihost_memory_.reset(infer_builder_->buildSerializedNetwork(
257 258
      *network(), *infer_builder_config_));
  infer_ptr<nvinfer1::IRuntime> runtime(createInferRuntime(&logger_));
Z
zlsh80826 已提交
259 260
  infer_engine_.reset(runtime->deserializeCudaEngine(ihost_memory_->data(),
                                                     ihost_memory_->size()));
261
#endif
262

263
  PADDLE_ENFORCE_NOT_NULL(
264 265 266 267
      infer_engine_,
      platform::errors::Fatal(
          "Build TensorRT cuda engine failed! Please recheck "
          "you configurations related to paddle-TensorRT."));
268

W
wenbin 已提交
269 270 271 272 273 274 275
  binding_num_ = infer_engine_->getNbBindings();
  // reset status for dynamic shape clone
  if (max_profile_num_ > 1) {
    infer_context_.clear();
    cur_profile_num_ = 0;
  }

276
  GetEngineInfo();
Y
Yan Chunwei 已提交
277 278
}

279
nvinfer1::ITensor *TensorRTEngine::DeclareInput(const std::string &name,
Y
Yan Chunwei 已提交
280
                                                nvinfer1::DataType dtype,
281
                                                const nvinfer1::Dims &dims) {
282 283
  PADDLE_ENFORCE_EQ(network() != nullptr,
                    true,
284 285 286
                    platform::errors::InvalidArgument(
                        "The TRT network should be initialized first."));
  auto *input = network()->addInput(name.c_str(), dtype, dims);
287
  PADDLE_ENFORCE_NOT_NULL(
288 289 290 291 292 293 294
      input,
      platform::errors::InvalidArgument("Adding input %s failed in "
                                        "TensorRT inference network. "
                                        "Please recheck your input.",
                                        name));
  PADDLE_ENFORCE_EQ(input->isNetworkInput(),
                    true,
295 296 297 298
                    platform::errors::InvalidArgument(
                        "Input %s is not the input of TRT inference network. "
                        "Please recheck your input.",
                        name));
L
Luo Tao 已提交
299
  TensorRTEngine::SetITensor(name, input);
Y
Yan Chunwei 已提交
300 301 302
  return input;
}

303 304
void TensorRTEngine::DeclareOutput(const nvinfer1::ILayer *layer,
                                   int offset,
305 306
                                   const std::string &name) {
  auto *output = layer->getOutput(offset);
307
  SetITensor(name, output);
308
  PADDLE_ENFORCE_NOT_NULL(
309 310 311
      output,
      platform::errors::InvalidArgument(
          "The output %s of TRT engine should not be null.", name));
Y
Yan Chunwei 已提交
312
  output->setName(name.c_str());
313 314
  PADDLE_ENFORCE_EQ(output->isNetworkInput(),
                    false,
315 316 317 318
                    platform::errors::InvalidArgument(
                        "The output %s of TRT engine should not be the input "
                        "of the network at the same time.",
                        name));
319
  network()->markOutput(*output);
320
  PADDLE_ENFORCE_EQ(
321 322
      output->isNetworkOutput(),
      true,
323 324 325
      platform::errors::InvalidArgument(
          "The output %s of TRT engine should be the output of the network.",
          name));
N
nhzlx 已提交
326 327
}

328 329
void TensorRTEngine::DeclareOutput(const std::string &name) {
  auto *output = TensorRTEngine::GetITensor(name);
330
  PADDLE_ENFORCE_NOT_NULL(
331 332 333
      output,
      platform::errors::InvalidArgument(
          "The output %s of TRT engine should not be null.", name));
L
Luo Tao 已提交
334
  output->setName(name.c_str());
335 336
  PADDLE_ENFORCE_EQ(output->isNetworkInput(),
                    false,
337 338 339 340
                    platform::errors::InvalidArgument(
                        "The output %s of TRT engine should not be the input "
                        "of the network at the same time.",
                        name));
341
  network()->markOutput(*output);
L
Luo Tao 已提交
342
}
343 344 345 346 347 348 349 350 351 352 353 354 355
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 已提交
356

357 358
void TensorRTEngine::SetITensor(const std::string &name,
                                nvinfer1::ITensor *tensor) {
359
  PADDLE_ENFORCE_NOT_NULL(
360 361 362
      tensor,
      platform::errors::InvalidArgument(
          "Tensor named %s of TRT engine should not be null.", name));
363
  PADDLE_ENFORCE_EQ(
364 365
      0,
      itensor_map_.count(name),
366 367
      platform::errors::InvalidArgument(
          "Tensor named %s of TRT engine should not be duplicated", name));
L
Luo Tao 已提交
368 369 370
  itensor_map_[name] = tensor;
}

371
nvinfer1::ITensor *TensorRTEngine::GetITensor(const std::string &name) {
372 373
  PADDLE_ENFORCE_EQ(itensor_map_.count(name),
                    true,
374 375
                    platform::errors::NotFound(
                        "Tensor named %s is not found in TRT engine", name));
L
Luo Tao 已提交
376 377 378
  return itensor_map_[name];
}

379 380 381 382 383
std::unordered_map<std::string, nvinfer1::ITensor *>
    *TensorRTEngine::GetITensorMap() {
  return &itensor_map_;
}

384 385 386 387
void TensorRTEngine::SetRuntimeBatch(size_t batch_size) {
  runtime_batch_ = batch_size;
}

388 389 390 391 392 393 394 395 396 397 398 399 400 401 402 403 404 405 406 407 408 409 410 411 412 413 414 415 416 417 418 419 420 421 422 423 424 425 426 427 428 429 430 431 432 433 434 435 436 437 438 439 440 441 442 443 444 445 446 447 448
// Note: Only for support plugin.
TensorRTEngine::Weight TensorRTEngine::GetFp16TrtWeight(
    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());
  weight.SetDataType(nvinfer1::DataType::kHALF);
  // weight_tensor.dims().;

  // if trt not support dtype, we need to cast to  fp16.
  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::FLOAT16);
    weight_map[name_with_suffix]->Resize(weight_tensor.dims());
    auto *fp16_data = weight_map[name_with_suffix]->mutable_data<float16>(
        platform::CPUPlace());
    auto *bf16_data = bf16_tensor.mutable_data<bfloat16>(platform::CPUPlace());
    for (int i = 0; i < weight_tensor.numel(); i++) {
      fp16_data[i] = static_cast<float16>(bf16_data[i]);
    }
  } else if (weight_tensor.dtype() == phi::DataType::FLOAT32) {
    framework::Tensor fp32_tensor;
    fp32_tensor.clear();
    paddle::framework::TensorCopySync(
        weight_tensor, platform::CPUPlace(), &fp32_tensor);
    weight_map[name_with_suffix]->set_type(
        paddle::experimental::DataType::FLOAT16);
    weight_map[name_with_suffix]->Resize(weight_tensor.dims());
    auto *fp16_data = weight_map[name_with_suffix]->mutable_data<float16>(
        platform::CPUPlace());
    auto *fp32_data = fp32_tensor.mutable_data<float>(platform::CPUPlace());
    for (int i = 0; i < weight_tensor.numel(); i++) {
      fp16_data[i] = static_cast<float16>(fp32_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;
}

// Note: Only for support plugin.
449 450 451 452 453 454
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;
455
  platform::CPUPlace cpu_place;
456 457 458 459 460 461 462 463 464 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
  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;
506 507
}

508 509 510 511 512 513 514 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
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());
  }
565

566 567 568
  name_suffix_counter += 1;
  return weight;
}
569

570 571
int TensorRTEngine::GetRuntimeBatch() { return runtime_batch_; }

572
nvinfer1::IPluginV2Layer *TensorRTEngine::AddPlugin(
573 574
    nvinfer1::ITensor *const *inputs,
    int num_inputs,
575
    plugin::PluginTensorRT *plugin) {
576
  owned_plugin_.emplace_back(plugin);
577
  return network()->addPluginV2(inputs, num_inputs, *plugin);
578 579
}

580
nvinfer1::IPluginV2Layer *TensorRTEngine::AddPluginV2Ext(
581 582
    nvinfer1::ITensor *const *inputs,
    int num_inputs,
583 584 585 586 587
    plugin::PluginTensorRTV2Ext *plugin) {
  owned_plugin_v2ext_.emplace_back(plugin);
  return network()->addPluginV2(inputs, num_inputs, *plugin);
}

588
nvinfer1::IPluginV2Layer *TensorRTEngine::AddPluginV2IOExt(
589 590
    nvinfer1::ITensor *const *inputs,
    int num_inputs,
591 592 593 594 595
    nvinfer1::IPluginV2IOExt *plugin) {
  owned_plugin_v2ioext_.emplace_back(plugin);
  return network()->addPluginV2(inputs, num_inputs, *plugin);
}

N
nhzlx 已提交
596 597 598
void TensorRTEngine::freshDeviceId() {
  int count;
  cudaGetDeviceCount(&count);
599 600
  PADDLE_ENFORCE_LT(device_id_,
                    count,
601 602
                    platform::errors::OutOfRange(
                        "Device id %d exceeds the current device count: %d.",
603 604
                        device_id_,
                        count));
L
Leo Chen 已提交
605
  platform::SetDeviceId(device_id_);
N
nhzlx 已提交
606 607
}

608 609 610 611 612
void TensorRTEngine::GetEngineInfo() {
#if IS_TRT_VERSION_GE(8200)
  LOG(INFO) << "====== engine info ======";
  std::unique_ptr<nvinfer1::IEngineInspector> infer_inspector(
      infer_engine_->createEngineInspector());
613 614 615
  auto infer_context = infer_ptr<nvinfer1::IExecutionContext>(
      infer_engine_->createExecutionContextWithoutDeviceMemory());
  infer_inspector->setExecutionContext(infer_context.get());
616 617 618 619 620 621 622 623
  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 已提交
624 625 626
}  // namespace tensorrt
}  // namespace inference
}  // namespace paddle