engine.cc 9.8 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 19
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 <cuda.h>
#include <glog/logging.h>
A
Abhinav Arora 已提交
20
#include <string>
Y
Yan Chunwei 已提交
21
#include "paddle/fluid/inference/analysis/helper.h"
Y
Yan Chunwei 已提交
22 23 24 25 26 27 28
#include "paddle/fluid/inference/tensorrt/helper.h"
#include "paddle/fluid/platform/enforce.h"

namespace paddle {
namespace inference {
namespace tensorrt {

29 30
int TensorRTEngine::runtime_batch_ = 1;

31
void TensorRTEngine::Build(const DescType &paddle_model) {
Y
Yan Chunwei 已提交
32 33 34
  PADDLE_ENFORCE(false, "not implemented");
}

35 36
void TensorRTEngine::Execute(int batch_size, std::vector<void *> *buffers,
                             cudaStream_t stream) {
N
nhzlx 已提交
37
  freshDeviceId();
38
  const std::thread::id tid = std::this_thread::get_id();
N
nhzlx 已提交
39
  batch_size_ = batch_size;
40
  if (infer_context_.find(tid) == infer_context_.end()) {
41
    std::unique_lock<std::mutex> lock(mutex_);
42 43 44 45 46 47
    PADDLE_ENFORCE_NOT_NULL(
        infer_engine_,
        "You should build engine first and then set the context.");
    infer_context_[tid].reset(infer_engine_->createExecutionContext());
  }
  infer_context_[tid]->enqueue(batch_size, buffers->data(), stream, nullptr);
N
nhzlx 已提交
48 49 50
  SetRuntimeBatch(batch_size);
}

Y
Yan Chunwei 已提交
51
void TensorRTEngine::FreezeNetwork() {
N
nhzlx 已提交
52
  freshDeviceId();
53
  VLOG(3) << "TRT to freeze network";
Y
Yan Chunwei 已提交
54 55 56 57 58 59 60
  PADDLE_ENFORCE(infer_builder_ != nullptr,
                 "Call InitNetwork first to initialize network.");
  PADDLE_ENFORCE(infer_network_ != nullptr,
                 "Call InitNetwork first to initialize network.");
  // build engine.
  infer_builder_->setMaxBatchSize(max_batch_);
  infer_builder_->setMaxWorkspaceSize(max_workspace_);
Z
Zhaolong Xing 已提交
61
  bool enable_fp16 = (precision_ == AnalysisConfig::Precision::kHalf);
62
#if IS_TRT_VERSION_GE(5000)
Z
Zhaolong Xing 已提交
63 64 65 66 67 68 69 70
  if (enable_fp16) {
    bool support_fp16 = infer_builder_->platformHasFastFp16();
    infer_builder_->setFp16Mode(support_fp16);
    if (!support_fp16) {
      LOG(INFO) << "You specify FP16 mode, but the hardware do not support "
                   "FP16 speed up, use FP32 instead.";
    }
  }
71
#else
72
  if (enable_fp16)
73
    LOG(INFO) << "Using FP16 in Paddle-TRT must ensure that the version of TRT "
74 75
                 "is at least 5."
                 "So, use FP32 to run.";
76
#endif
Z
Zhaolong Xing 已提交
77 78 79
  bool enable_int8 = (precision_ == AnalysisConfig::Precision::kInt8);

  if (enable_int8) {
N
nhzlx 已提交
80
    infer_builder_->setInt8Mode(true);
81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106
    if (calibrator_) {
      infer_builder_->setInt8Calibrator(calibrator_);
    } else {
      infer_builder_->setInt8Calibrator(nullptr);

#if IS_TRT_VERSION_GE(5000)
      infer_builder_->setStrictTypeConstraints(true);
      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;
      for (int i = 0; i < infer_network_->getNbLayers(); i++) {
        auto layer = infer_network_->getLayer(i);
        for (int j = 0; j < layer->getNbOutputs(); j++) {
          all_t.insert(layer->getOutput(j));
        }
      }
      for (int i = 0; i < infer_network_->getNbInputs(); i++) {
        all_t.insert(infer_network_->getInput(i));
      }

      for (auto &t : all_t) {
        if (!quant_dynamic_range_.count(t)) {
107
          VLOG(3)
108 109 110 111 112
              << "We are in trt int8 mode(not calibration), scale not setted"
              << " for tensor " << t->getName()
              << ", this might be ok when trt does not need this range";
        }
      }
113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131
      std::unordered_set<std::string> all_out_t_name;
      for (int i = 0; i < infer_network_->getNbOutputs(); i++) {
        auto *temp = infer_network_->getOutput(i);
        temp->setDynamicRange(-1, 1);
        all_out_t_name.insert(temp->getName());
      }

      for (int i = 0; i < infer_network_->getNbLayers(); i++) {
        auto layer = infer_network_->getLayer(i);
        for (int j = 0; j < layer->getNbOutputs(); j++) {
          auto *temp_out = layer->getOutput(j);
          if (std::find(all_out_t_name.begin(), all_out_t_name.end(),
                        temp_out->getName()) != all_out_t_name.end()) {
            layer->setPrecision(nvinfer1::DataType::kFLOAT);
            layer->setOutputType(j, nvinfer1::DataType::kFLOAT);
          }
        }
      }

132 133
#endif
    }
N
nhzlx 已提交
134
  }
Y
Yan Chunwei 已提交
135 136 137 138 139

  infer_engine_.reset(infer_builder_->buildCudaEngine(*infer_network_));
  PADDLE_ENFORCE(infer_engine_ != nullptr, "build cuda engine failed!");
}

140
nvinfer1::ITensor *TensorRTEngine::DeclareInput(const std::string &name,
Y
Yan Chunwei 已提交
141
                                                nvinfer1::DataType dtype,
142
                                                const nvinfer1::Dims &dims) {
Y
Yan Chunwei 已提交
143 144 145 146
  PADDLE_ENFORCE_EQ(0, buffer_sizes_.count(name), "duplicate input name %s",
                    name);

  PADDLE_ENFORCE(infer_network_ != nullptr, "should initnetwork first");
147
  auto *input = infer_network_->addInput(name.c_str(), dtype, dims);
Y
Yan Chunwei 已提交
148
  PADDLE_ENFORCE(input, "infer network add input %s failed", name);
Y
Yan Chunwei 已提交
149
  buffer_sizes_[name] = kDataTypeSize[static_cast<int>(dtype)] *
150
                        analysis::AccuDims(dims.d, dims.nbDims) * max_batch_;
151
  PADDLE_ENFORCE(input->isNetworkInput());
L
Luo Tao 已提交
152
  TensorRTEngine::SetITensor(name, input);
Y
Yan Chunwei 已提交
153 154 155
  return input;
}

156 157
void TensorRTEngine::DeclareOutput(const nvinfer1::ILayer *layer, int offset,
                                   const std::string &name) {
Y
Yan Chunwei 已提交
158 159 160
  PADDLE_ENFORCE_EQ(0, buffer_sizes_.count(name), "duplicate output name %s",
                    name);

161
  auto *output = layer->getOutput(offset);
162
  SetITensor(name, output);
Y
Yan Chunwei 已提交
163 164
  PADDLE_ENFORCE(output != nullptr);
  output->setName(name.c_str());
165
  PADDLE_ENFORCE(!output->isNetworkInput());
Y
Yan Chunwei 已提交
166
  infer_network_->markOutput(*output);
167
  PADDLE_ENFORCE(output->isNetworkOutput());
168 169
  // output buffers' size can only be decided later, set zero here to mark this
  // and will reset later.
Y
Yan Chunwei 已提交
170 171 172
  buffer_sizes_[name] = 0;
}

N
nhzlx 已提交
173 174 175 176
bool TensorRTEngine::HasDeclared(const std::string &name) {
  return buffer_sizes_.count(name) > 0;
}

177
void TensorRTEngine::DeclareOutput(const std::string &name) {
L
Luo Tao 已提交
178 179 180
  PADDLE_ENFORCE_EQ(0, buffer_sizes_.count(name), "duplicate output name %s",
                    name);

181
  auto *output = TensorRTEngine::GetITensor(name);
L
Luo Tao 已提交
182 183
  PADDLE_ENFORCE(output != nullptr);
  output->setName(name.c_str());
184
  PADDLE_ENFORCE(!output->isNetworkInput());
L
Luo Tao 已提交
185
  infer_network_->markOutput(*output);
186 187
  // output buffers' size can only be decided later, set zero here to mark this
  // and will reset later.
L
Luo Tao 已提交
188 189 190
  buffer_sizes_[name] = 0;
}

191 192
void TensorRTEngine::SetITensor(const std::string &name,
                                nvinfer1::ITensor *tensor) {
L
Luo Tao 已提交
193
  PADDLE_ENFORCE(tensor != nullptr);
Y
Yan Chunwei 已提交
194
  PADDLE_ENFORCE_EQ(0, itensor_map_.count(name), "duplicate ITensor name %s",
L
Luo Tao 已提交
195 196 197 198
                    name);
  itensor_map_[name] = tensor;
}

199
nvinfer1::ITensor *TensorRTEngine::GetITensor(const std::string &name) {
Y
Yan Chunwei 已提交
200
  PADDLE_ENFORCE(itensor_map_.count(name), "no ITensor %s", name);
L
Luo Tao 已提交
201 202 203
  return itensor_map_[name];
}

204 205 206 207
void TensorRTEngine::SetRuntimeBatch(size_t batch_size) {
  runtime_batch_ = batch_size;
}

208 209 210 211
float *TensorRTEngine::GetWeightCPUData(const std::string &name,
                                        framework::Tensor *weight_tensor,
                                        bool enable_int8,
                                        const std::vector<float> &scale) {
212 213
  static int name_suffix_counter = 0;
  std::string name_suffix = std::to_string(name_suffix_counter);
P
Pei Yang 已提交
214 215
  std::string splitter = "__";
  std::string name_with_suffix = name + splitter + name_suffix;
216 217
  auto w_dims = weight_tensor->dims();
  platform::CPUPlace cpu_place;
218 219 220 221 222 223 224 225 226 227 228
  PADDLE_ENFORCE_EQ(
      weight_map.count(name_with_suffix), 0,
      "During TRT Op converter: We set weight %s with the same name "
      "twice into the weight_map",
      name_with_suffix);
  weight_map[name_with_suffix].reset(new framework::Tensor());
  weight_map[name_with_suffix]->Resize(weight_tensor->dims());
  TensorCopySync(*weight_tensor, cpu_place, weight_map[name_with_suffix].get());
  float *weight_data =
      weight_map[name_with_suffix]->mutable_data<float>(cpu_place);
  name_suffix_counter += 1;
229 230 231

  if (enable_int8) {
    // when the op is fc, scale's size should be 1
232
    // when the op is conv, scale's size should be w_dims[0]
233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250
    bool valid_scale_size =
        (scale.size() == 1 || scale.size() == static_cast<size_t>(w_dims[0]));
    PADDLE_ENFORCE(valid_scale_size, "TRT int8 quant: invalid scale size");
    for (int i = 0; i < weight_tensor->numel(); i++) {
      if (scale.size() == 1) {
        weight_data[i] *= (scale[0] / 127);
      } else {
        PADDLE_ENFORCE(w_dims.size() == 4,
                       "TRT int8 quant : We only use the channel quant for "
                       "conv op, so the weight dims should be 4.");
        int inner_size = w_dims[1] * w_dims[2] * w_dims[3];
        weight_data[i] *= (scale[i / inner_size] / 127);
      }
    }
  }
  return weight_data;
}

251 252
int TensorRTEngine::GetRuntimeBatch() { return runtime_batch_; }

N
nhzlx 已提交
253
nvinfer1::IPluginLayer *TensorRTEngine::AddPlugin(
254 255
    nvinfer1::ITensor *const *inputs, int num_inputs,
    plugin::PluginTensorRT *plugin) {
256
  owned_plugin_.emplace_back(plugin);
257
  return infer_network_.get()->addPluginExt(inputs, num_inputs, *plugin);
258 259
}

N
nhzlx 已提交
260 261 262 263 264 265 266
void TensorRTEngine::freshDeviceId() {
  int count;
  cudaGetDeviceCount(&count);
  PADDLE_ENFORCE_LT(device_id_, count);
  cudaSetDevice(device_id_);
}

Y
Yan Chunwei 已提交
267 268 269
}  // namespace tensorrt
}  // namespace inference
}  // namespace paddle