tensorrt_engine_op.h 4.5 KB
Newer Older
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

#ifdef PADDLE_WITH_CUDA

G
gongweibao 已提交
19 20 21
#include <string>
#include <vector>

22 23 24 25 26
#include "paddle/fluid/framework/operator.h"
#include "paddle/fluid/inference/analysis/helper.h"
#include "paddle/fluid/inference/tensorrt/engine.h"

namespace paddle {
27 28 29

DECLARE_int32(tensorrt_engine_batch_size);

30 31
namespace operators {

Y
Yan Chunwei 已提交
32 33 34
using inference::Singleton;
using inference::tensorrt::TRT_EngineManager;

35 36 37 38 39 40 41 42 43
class TensorRTEngineOp : public framework::OperatorWithKernel {
 public:
  using framework::OperatorWithKernel::OperatorWithKernel;

 protected:
  void InferShape(framework::InferShapeContext* ctx) const override {}

  framework::OpKernelType GetExpectedKernelType(
      const framework::ExecutionContext& ctx) const override {
44
    auto input0 = ctx.Inputs("Xs").front();
45
    framework::OpKernelType kt = framework::OpKernelType(
46 47 48 49
        framework::ToDataType(ctx.scope()
                                  .FindVar(input0)
                                  ->GetMutable<framework::LoDTensor>()
                                  ->type()),
50 51 52 53 54 55 56 57 58
        platform::CPUPlace());
    return kt;
  }
};

template <typename DeviceContext, typename T>
class TensorRTEngineKernel : public framework::OpKernel<T> {
 public:
  void Compute(const framework::ExecutionContext& context) const override {
Y
Yan Chunwei 已提交
59 60
    auto engine_name = context.Attr<std::string>("engine_uniq_key");
    if (!Singleton<TRT_EngineManager>::Global().HasEngine(engine_name)) {
61 62
      Prepare(context);
    }
Y
Yan Chunwei 已提交
63
    auto* engine = Singleton<TRT_EngineManager>::Global().Get(engine_name);
64 65
    auto input_names = context.op().Inputs("Xs");
    PADDLE_ENFORCE(!input_names.empty(), "should pass more than one inputs");
66 67
    PADDLE_ENFORCE_LE(FLAGS_tensorrt_engine_batch_size,
                      context.Attr<int>("max_batch"));
68 69 70 71

    // Convert input tensor from fluid to engine.
    for (const auto& x : context.Inputs("Xs")) {
      // convert input and copy to TRT engine's buffer
72 73
      auto& t = inference::analysis::GetFromScope<framework::LoDTensor>(
          context.scope(), x);
74
      if (platform::is_cpu_place(t.place())) {
Y
Yan Chunwei 已提交
75 76
        engine->SetInputFromCPU(x, static_cast<const void*>(t.data<void>()),
                                t.memory_size());
77
      } else {
Y
Yan Chunwei 已提交
78 79
        engine->SetInputFromGPU(x, static_cast<const void*>(t.data<void>()),
                                t.memory_size());
80 81 82
      }
    }
    // Execute the engine.
83 84
    PADDLE_ENFORCE_GT(FLAGS_tensorrt_engine_batch_size, 0);
    engine->Execute(FLAGS_tensorrt_engine_batch_size);
85 86 87
    // Convert output tensor from engine to fluid
    for (const auto& y : context.Outputs("Ys")) {
      // convert output and copy to fluid.
Y
Yan Chunwei 已提交
88
      nvinfer1::ITensor* trt_t = engine->GetITensor(y);
89 90 91 92 93 94 95
      auto dims = trt_t->getDimensions();
      // Use the output ITensor's dims to reshape the Fluid Tensor.
      std::vector<int> ddim(dims.d, dims.d + dims.nbDims);

      auto* fluid_v = context.scope().FindVar(y);
      PADDLE_ENFORCE_NOT_NULL(fluid_v, "no output variable called %s", y);
      auto* fluid_t = fluid_v->GetMutable<framework::LoDTensor>();
N
nhzlx 已提交
96

97
      fluid_t->Resize(framework::make_ddim(ddim));
98

99 100 101 102
      // TODO(Superjomn) find some way to determine which device to output the
      // tensor.
      // if (platform::is_cpu_place(fluid_t->place())) {
      // TODO(Superjomn) change this float to dtype size.
N
nhzlx 已提交
103 104 105 106 107
      auto size = inference::analysis::AccuDims(dims.d, dims.nbDims) *
                  FLAGS_tensorrt_engine_batch_size;
      engine->GetOutputInCPU(y,
                             fluid_t->mutable_data<float>(platform::CPUPlace()),
                             size * sizeof(float));
108 109 110 111 112
      //} else {
      // engine->GetOutputInGPU(
      // y, fluid_t->mutable_data<float>(platform::CUDAPlace()),
      // size * sizeof(float));
      //}
113
    }
114

Y
Yan Chunwei 已提交
115
    cudaStreamSynchronize(*engine->stream());
116 117 118 119 120 121 122 123 124 125 126
  }

 protected:
  // Build the engine.
  void Prepare(const framework::ExecutionContext& context) const;
};

}  // namespace operators
}  // namespace paddle

#endif  // PADDLE_WITH_CUDA