engine.h 8.3 KB
Newer Older
Y
Yan Chunwei 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18
/* Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
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. */

#pragma once

#include <NvInfer.h>
#include <memory>
19
#include <string>
Y
Yan Chunwei 已提交
20
#include <unordered_map>
21
#include <vector>
N
nhzlx 已提交
22
#include "paddle/fluid/framework/tensor.h"
Y
Yan Chunwei 已提交
23 24
#include "paddle/fluid/inference/engine.h"
#include "paddle/fluid/inference/tensorrt/helper.h"
25
#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
N
nhzlx 已提交
26
#include "paddle/fluid/inference/tensorrt/trt_int8_calibrator.h"
27
#include "paddle/fluid/inference/utils/singleton.h"
Y
Yan Chunwei 已提交
28 29 30 31 32

namespace paddle {
namespace inference {
namespace tensorrt {

N
nhzlx 已提交
33
class TRTInt8Calibrator;
Y
Yan Chunwei 已提交
34 35 36 37 38 39 40 41 42 43 44
/*
 * TensorRT Engine.
 *
 * There are two alternative ways to use it, one is  to build from a paddle
 * protobuf model, another way is to manully construct the network.
 */
class TensorRTEngine : public EngineBase {
 public:
  // Weight is model parameter.
  class Weight {
   public:
45
    Weight() = default;
46
    Weight(nvinfer1::DataType dtype, void* value, size_t num_elem) {
Y
Yan Chunwei 已提交
47 48 49 50
      w_.type = dtype;
      w_.values = value;
      w_.count = num_elem;
    }
51
    const nvinfer1::Weights& get() { return w_; }
Y
Yan Chunwei 已提交
52

53 54
    std::vector<int64_t> dims;

Y
Yan Chunwei 已提交
55 56 57 58
   private:
    nvinfer1::Weights w_;
  };

Y
Yan Chunwei 已提交
59
  TensorRTEngine(int max_batch, int max_workspace,
N
nhzlx 已提交
60
                 cudaStream_t* stream = nullptr, int device = 0,
N
nhzlx 已提交
61 62
                 std::string precision_mode = "FP32",
                 TRTInt8Calibrator* calibrator = nullptr,
Y
Yan Chunwei 已提交
63 64 65
                 nvinfer1::ILogger& logger = NaiveLogger::Global())
      : max_batch_(max_batch),
        max_workspace_(max_workspace),
Y
Yan Chunwei 已提交
66
        stream_(stream ? stream : &default_stream_),
N
nhzlx 已提交
67 68 69 70
        device_(device),
        precision_mode_(precision_mode),
        calibrator_(calibrator),
        logger_(logger) {
N
nhzlx 已提交
71 72
    freshDeviceId();
    cudaStreamCreate(stream_);
73
  }
Y
Yan Chunwei 已提交
74 75 76 77

  virtual ~TensorRTEngine();

  // TODO(Superjomn) implement it later when graph segmentation is supported.
78
  void Build(const DescType& paddle_model) override;
Y
Yan Chunwei 已提交
79

80
  void Execute(int batch_size) override;
Y
Yan Chunwei 已提交
81 82 83 84

  // Initialize the inference network, so that TensorRT layers can add to this
  // network.
  void InitNetwork() {
85
    infer_builder_.reset(createInferBuilder(&logger_));
Y
Yan Chunwei 已提交
86 87 88 89 90 91 92 93 94 95 96 97 98 99
    infer_network_.reset(infer_builder_->createNetwork());
  }
  // After finishing adding ops, freeze this network and creates the executation
  // environment.
  void FreezeNetwork();

  // Add an input and set its name, data type and dimention.
  nvinfer1::ITensor* DeclareInput(const std::string& name,
                                  nvinfer1::DataType dtype,
                                  const nvinfer1::Dims& dim);
  // Set the offset-th output from a layer as the network's output, and set its
  // name.
  void DeclareOutput(const nvinfer1::ILayer* layer, int offset,
                     const std::string& name);
L
Luo Tao 已提交
100 101
  // Set the itensor_map_[name] as the network's output, and set its name.
  void DeclareOutput(const std::string& name);
N
nhzlx 已提交
102 103
  // Check if the ITensor has been declared
  bool HasDeclared(const std::string& name);
Y
Yan Chunwei 已提交
104 105 106 107 108

  // GPU memory address for an ITensor with specific name. One can operate on
  // these memory directly for acceleration, for example, output the converted
  // data directly to the buffer to save data copy overhead.
  // NOTE this should be used after calling `FreezeNetwork`.
Y
Yan Chunwei 已提交
109 110 111
  Buffer& buffer(const std::string& name) override;

  cudaStream_t* stream() { return stream_; }
Y
Yan Chunwei 已提交
112 113

  // Fill an input from CPU memory with name and size.
114
  void SetInputFromCPU(const std::string& name, const void* data, size_t size);
Y
Yan Chunwei 已提交
115 116
  // TODO(Superjomn) is this method necessary given that buffer(xxx) can be
  // accessed directly. Fill an input from GPU memory with name and size.
117
  void SetInputFromGPU(const std::string& name, const void* data, size_t size);
Y
Yan Chunwei 已提交
118
  // Get an output called name, the output of tensorrt is in GPU, so this method
119
  // Return the output's GPU memory address without copy.
Y
Yan Chunwei 已提交
120
  void* GetOutputInGPU(const std::string& name);
121
  // Copy data into dst inside the GPU device.
N
nhzlx 已提交
122
  void GetOutputInGPU(const std::string& name, void* dst, size_t max_size);
Y
Yan Chunwei 已提交
123 124
  // LOW EFFICENCY! Get output to CPU, this will trigger a memory copy from GPU
  // to CPU.
N
nhzlx 已提交
125
  void GetOutputInCPU(const std::string& name, void* dst, size_t max_size);
L
Luo Tao 已提交
126 127 128 129
  // Fill an ITensor into map itensor_map_.
  void SetITensor(const std::string& name, nvinfer1::ITensor* tensor);
  // Get an ITensor called name.
  nvinfer1::ITensor* GetITensor(const std::string& name);
Y
Yan Chunwei 已提交
130 131 132

  nvinfer1::ICudaEngine* engine() { return infer_engine_.get(); }
  nvinfer1::INetworkDefinition* network() { return infer_network_.get(); }
133 134
  void SetRuntimeBatch(size_t batch_size);
  int GetRuntimeBatch();
N
nhzlx 已提交
135
  int GetDevice() { return device_; }
N
nhzlx 已提交
136
  nvinfer1::IPluginLayer* AddPlugin(nvinfer1::ITensor* const* inputs,
137
                                    int num_inputs, plugin::PluginTensorRT*);
N
nhzlx 已提交
138 139 140 141 142 143 144 145

  // A pointer to CPU memory is needed of the TRT weight.
  // Before TRT runs, fluid loads weight into GPU storage.
  // so we need to copy the weights from GPU to CPU in our op converter.
  // We use a map to store these weights for the weight memory is not released
  // in advance, which affecting the construction of TRT Op.
  std::unordered_map<std::string /*name*/, std::unique_ptr<framework::Tensor>>
      weight_map;
Y
Yan Chunwei 已提交
146

147
  // TODO(NHZLX)
N
nhzlx 已提交
148 149 150
  // In the normal case, the paddle-trt exists bug when runing the googlenet.
  // When there are more than two convolutions of 1 * 1 with the same input, the
  // paddle-tensorrt will do the merging optimization, which fuse those conv
N
nhzlx 已提交
151 152
  // into one conv, and then trigger bug. So,  We should use strategy to avoid
  // this
N
nhzlx 已提交
153 154 155 156
  // optimization for the time being. This bug will be fixed in the future.
  std::unordered_map<std::string /*name*/, int /*ITensor_quote_num*/>
      itensor_quote_num;

Y
Yan Chunwei 已提交
157 158 159
 private:
  // the max batch size
  int max_batch_;
160 161
  // the runtime batch size
  static int runtime_batch_;
Y
Yan Chunwei 已提交
162 163
  // the max memory size the engine uses
  int max_workspace_;
164

Y
Yan Chunwei 已提交
165
  cudaStream_t* stream_;
Y
Yan Chunwei 已提交
166 167
  // If stream_ is not set from outside, hold its own stream.
  cudaStream_t default_stream_;
N
nhzlx 已提交
168 169 170 171 172 173 174
  // The specific GPU id that the TensorRTEngine bounded to.
  int device_;

  std::string precision_mode_;
  TRTInt8Calibrator* calibrator_;
  // batch size of the current data, will be updated each Executation.
  int batch_size_{-1};
Y
Yan Chunwei 已提交
175 176
  nvinfer1::ILogger& logger_;

Y
Yan Chunwei 已提交
177
  std::vector<Buffer> buffers_;
Y
Yan Chunwei 已提交
178 179
  // max data size for the buffers.
  std::unordered_map<std::string /*name*/, size_t /*max size*/> buffer_sizes_;
L
Luo Tao 已提交
180 181
  std::unordered_map<std::string /*name*/, nvinfer1::ITensor* /*ITensor*/>
      itensor_map_;
182

183
  std::vector<std::unique_ptr<plugin::PluginTensorRT>> owned_plugin_;
Y
Yan Chunwei 已提交
184 185 186 187

  // TensorRT related internal members
  template <typename T>
  struct Destroyer {
188 189 190 191 192
    void operator()(T* x) {
      if (x) {
        x->destroy();
      }
    }
Y
Yan Chunwei 已提交
193 194 195 196 197 198 199
  };
  template <typename T>
  using infer_ptr = std::unique_ptr<T, Destroyer<T>>;
  infer_ptr<nvinfer1::IBuilder> infer_builder_;
  infer_ptr<nvinfer1::INetworkDefinition> infer_network_;
  infer_ptr<nvinfer1::ICudaEngine> infer_engine_;
  infer_ptr<nvinfer1::IExecutionContext> infer_context_;
N
nhzlx 已提交
200 201 202 203
  // Each ICudaEngine object is bound to a specific GPU when it is instantiated,
  // ensure that the thread is associated with the correct device by calling
  // freshDeviceId().
  void freshDeviceId();
Y
Yan Chunwei 已提交
204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222
};  // class TensorRTEngine

// Add an layer__ into engine__ with args ARGS.
// For example:
//   TRT_ENGINE_ADD_LAYER(xxx, FullyConnected, input, dim, weights, bias)
//
// Reference
// https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#charRNN_define_network
//
// will add a fully connected layer into the engine.
// TensorRT has too many layers, so that is not wise to add member functions for
// them, and an macro like this is more extensible when underlying TensorRT
// library add new layer supports.
#define TRT_ENGINE_ADD_LAYER(engine__, layer__, ARGS...) \
  engine__->network()->add##layer__(ARGS);

}  // namespace tensorrt
}  // namespace inference
}  // namespace paddle