“0ba63475659822bd146f1f1dcfc7eabca8b7047d”上不存在“paddle/fluid/inference/analysis/device.h”
op_converter.h 6.0 KB
Newer Older
L
Luo Tao 已提交
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
/* 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 <string>
#include <unordered_map>
#include "paddle/fluid/framework/block_desc.h"
20
#include "paddle/fluid/framework/op_registry.h"
L
Luo Tao 已提交
21 22
#include "paddle/fluid/framework/scope.h"
#include "paddle/fluid/inference/tensorrt/engine.h"
L
Luo Tao 已提交
23
#include "paddle/fluid/inference/utils/singleton.h"
L
Luo Tao 已提交
24 25 26 27 28 29 30 31 32 33 34

namespace paddle {
namespace inference {
namespace tensorrt {

/*
 * Convert Op from Fluid to TensorRT Engine.
 */
class OpConverter {
 public:
  OpConverter() {}
L
Luo Tao 已提交
35

36 37
  // Converter logic for an op.
  virtual void operator()(const framework::proto::OpDesc& op,
38 39
                          const framework::Scope& scope,
                          bool test_mode = false) {}
40

41 42
  // Convert a single fluid operator and add the corresponding layer to TRT.
  // test_mode: whether the instance executes in an unit test.
43 44
  void ConvertOp(const framework::proto::OpDesc& op,
                 const std::unordered_set<std::string>& parameters,
45 46
                 const framework::Scope& scope, TensorRTEngine* engine,
                 bool test_mode = false) {
Y
Yan Chunwei 已提交
47
    framework::OpDesc op_desc(op, nullptr);
48 49

    OpConverter* it{nullptr};
L
Luo Tao 已提交
50

51 52 53 54 55 56 57
    if (op_desc.Type() == "mul") {
      PADDLE_ENFORCE_EQ(op_desc.Input("Y").size(), 1UL);
      std::string Y = op_desc.Input("Y")[0];
      if (parameters.count(Y)) {
        it = Registry<OpConverter>::Lookup("fc");
      }
    }
N
nhzlx 已提交
58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73
    if (op_desc.Type().find("elementwise") != std::string::npos) {
      static std::unordered_set<std::string> add_tensor_op_set{
          "add", "mul", "sub", "div", "max", "min", "pow"};
      // TODO(xingzhaolong): all mul, sub, div
      // static std::unordered_set<std::string> add_weight_op_set {"add", "mul",
      // "sub", "div"};
      static std::unordered_set<std::string> add_weight_op_set{"add"};
      PADDLE_ENFORCE_EQ(op_desc.Input("Y").size(), 1UL);
      int op_type_len = op_desc.Type().size();
      std::string op_type = op_desc.Type().substr(op_type_len - 3, op_type_len);
      std::string Y = op_desc.Input("Y")[0];
      if (parameters.count(Y)) {
        PADDLE_ENFORCE(add_weight_op_set.count(op_type) > 0,
                       "Unsupported elementwise type" + op_type);
        it =
            Registry<OpConverter>::Lookup("elementwise_" + op_type + "_weight");
74 75
        PADDLE_ENFORCE_NOT_NULL(it, "no OpConverter for optype [%s]",
                                op_desc.Type());
N
nhzlx 已提交
76 77 78 79 80 81
      } else {
        PADDLE_ENFORCE(add_tensor_op_set.count(op_type) > 0,
                       "Unsupported elementwise type" + op_type);
        it =
            Registry<OpConverter>::Lookup("elementwise_" + op_type + "_tensor");
      }
N
nhzlx 已提交
82 83 84 85 86 87 88 89
      PADDLE_ENFORCE_NOT_NULL(it, "no OpConverter for optype [%s]",
                              op_desc.Type());
    }

    if (op_desc.Type() == "depthwise_conv2d") {
      it = Registry<OpConverter>::Lookup("conv2d");
      PADDLE_ENFORCE_NOT_NULL(it, "no OpConverter for optype [%s]",
                              op_desc.Type());
N
nhzlx 已提交
90 91
    }

92 93 94 95 96 97
    if (!it) {
      it = Registry<OpConverter>::Lookup(op_desc.Type());
    }
    PADDLE_ENFORCE_NOT_NULL(it, "no OpConverter for optype [%s]",
                            op_desc.Type());
    it->SetEngine(engine);
98
    (*it)(op, scope, test_mode);
L
Luo Tao 已提交
99 100
  }

Y
Yan Chunwei 已提交
101 102
  // Convert a fluid block to tensorrt network, NOTE it just convert operators,
  // the INetwork's inputs and outputs should specified in some other modules.
103
  void ConvertBlock(const framework::proto::BlockDesc& block,
104 105
                    const std::unordered_set<std::string>& parameters,
                    const framework::Scope& scope, TensorRTEngine* engine) {
K
Kexin Zhao 已提交
106
    for (int i = 0; i < block.ops_size(); i++) {
107
      const auto& op = block.ops(i);
108
      ConvertOp(op, parameters, scope, engine);
L
Luo Tao 已提交
109 110 111 112 113
    }
  }

  void SetEngine(TensorRTEngine* engine) { engine_ = engine; }

L
Luo Tao 已提交
114 115
  virtual ~OpConverter() {}

L
Luo Tao 已提交
116 117 118
  // TensorRT engine
  TensorRTEngine* engine_{nullptr};

119 120 121
 protected:
  bool test_mode_;

L
Luo Tao 已提交
122 123 124 125 126
 private:
  // registered op converter map, whose key is the fluid op type, and value is
  // the pointer position of corresponding OpConverter class.
  std::unordered_map<std::string, OpConverter*> converters_;
  // fluid inference scope
L
Luo Tao 已提交
127
  framework::Scope* scope_{nullptr};
L
Luo Tao 已提交
128 129
};

130 131 132 133
}  // namespace tensorrt
}  // namespace inference
}  // namespace paddle

134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151
#define REGISTER_TRT_OP_CONVERTER(op_type__, Converter__)                      \
  struct trt_##op_type__##_converter : public ::paddle::framework::Registrar { \
    trt_##op_type__##_converter() {                                            \
      ::paddle::inference::                                                    \
          Registry<paddle::inference::tensorrt::OpConverter>::Register<        \
              ::paddle::inference::tensorrt::Converter__>(#op_type__);         \
    }                                                                          \
  };                                                                           \
  trt_##op_type__##_converter trt_##op_type__##_converter__;                   \
  int TouchConverterRegister_##op_type__() {                                   \
    trt_##op_type__##_converter__.Touch();                                     \
    return 0;                                                                  \
  }

#define USE_TRT_CONVERTER(op_type__)                                    \
  extern int TouchConverterRegister_##op_type__();                      \
  static int use_op_converter_trt_##op_type__ __attribute__((unused)) = \
      TouchConverterRegister_##op_type__();