update the register method

Author: luotao1

paddle/fluid/inference/tensorrt/CMakeLists.txt

@@ -1,4 +1,3 @@
 nv_test(test_tensorrt SRCS test_tensorrt.cc DEPS dynload_cuda device_context dynamic_loader)
 nv_test(test_tensorrt_engine SRCS test_engine.cc engine.cc DEPS dynload_cuda)
-cc_library(tensorrt DEPS tensorrt_convert)
 add_subdirectory(convert)

paddle/fluid/inference/tensorrt/convert/CMakeLists.txt

@@ -1,2 +1,2 @@
-nv_library(tensorrt_convert SRCS convert.cc mul_op.cc conv2d_op.cc DEPS dynload_cuda)
-nv_test(test_tensorrt_convert SRCS test_convert.cc DEPS tensorrt paddle_fluid)
+file(GLOB TENSORRT_OPS RELATIVE "${CMAKE_CURRENT_SOURCE_DIR}" "*_op.cc")
+nv_test(test_tensorrt_op_converter SRCS test_op_converter.cc ${TENSORRT_OPS} DEPS ${FLUID_CORE_MODULES})

paddle/fluid/inference/tensorrt/convert/conv2d_op.cc

@@ -4,25 +4,30 @@ 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
+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/convert/convert.h"
+#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
 
 namespace paddle {
 namespace inference {
 namespace tensorrt {
 
-REGISTER_TRT_OP_CONVETER(conv2d, Conv2dOpConverter);
+class Conv2dOpConverter : public OpConverter {
+ public:
+  Conv2dOpConverter() {}
+  void operator()(const framework::OpDesc& op) override {
+    LOG(INFO)
+        << "convert a fluid conv2d op to tensorrt conv layer without bias";
+  }
+};
 
-void Conv2dOpConverter::Convert(const framework::OpDesc& op) {
-  LOG(INFO) << "convert a fluid conv2d op to tensorrt conv layer without bias";
-}
+REGISTER_TRT_OP_CONVERTER(conv2d, Conv2dOpConverter);
 
 }  // namespace tensorrt
 }  // namespace inference

paddle/fluid/inference/tensorrt/convert/convert.cc

@@ -4,25 +4,29 @@ 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
+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/convert/convert.h"
+#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
 
 namespace paddle {
 namespace inference {
 namespace tensorrt {
 
-REGISTER_TRT_OP_CONVETER(mul, MulOpConverter);
+class MulOpConverter : public OpConverter {
+ public:
+  MulOpConverter() {}
+  void operator()(const framework::OpDesc& op) override {
+    LOG(INFO) << "convert a fluid mul op to tensorrt fc layer without bias";
+  }
+};
 
-void MulOpConverter::Convert(const framework::OpDesc& op) {
-  LOG(INFO) << "convert a fluid mul op to tensorrt fc layer without bias";
-}
+REGISTER_TRT_OP_CONVERTER(mul, MulOpConverter);
 
 }  // namespace tensorrt
 }  // namespace inference

paddle/fluid/inference/tensorrt/convert/convert.h

@@ -0,0 +1,89 @@
+/* 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"
+#include "paddle/fluid/framework/scope.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+/*
+ * Convert Op from Fluid to TensorRT Engine.
+ */
+class OpConverter {
+ public:
+  OpConverter() {}
+
+  virtual void operator()(const framework::OpDesc& op) {}
+  void Execute(const framework::OpDesc& op) {
+    std::string type = op.Type();
+    auto it = converters_.find(type);
+    PADDLE_ENFORCE(it != converters_.end(), "no OpConverter for optype [%s]",
+                   type);
+    (*it->second)(op);
+  }
+
+  static OpConverter& Global() {
+    static auto* x = new OpConverter;
+    return *x;
+  }
+
+  template <typename T>
+  void Register(const std::string& key) {
+    converters_[key] = new T;
+  }
+
+  virtual ~OpConverter() {}
+
+ 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
+  framework::Scope* scope_;
+  // tensorrt input/output tensor map, whose key is the fluid variable name,
+  // and value is the pointer position of tensorrt tensor
+  std::unordered_map<std::string, nvinfer1::ITensor*> tr_tensors_;
+};
+
+#define REGISTER_TRT_OP_CONVERTER(op_type__, Converter__)      \
+  struct trt_##op_type__##_converter {                         \
+    trt_##op_type__##_converter() {                            \
+      OpConverter::Global().Register<Converter__>(#op_type__); \
+    }                                                          \
+  };                                                           \
+  trt_##op_type__##_converter trt_##op_type__##_converter__;
+
+class BlockConverter {
+ public:
+  BlockConverter() {}
+
+  // convert fluid block to tensorrt network
+  void ConvertBlock(const framework::BlockDesc& block) {
+    for (auto op : block.AllOps()) {
+      OpConverter::Global().Execute(*op);
+    }
+  }
+};
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/convert/mul_op.cc

@@ -14,21 +14,21 @@ limitations under the License. */
 
 #include <gtest/gtest.h>
 #include "paddle/fluid/framework/program_desc.h"
-#include "paddle/fluid/inference/tensorrt/convert/convert.h"
+#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
 
 namespace paddle {
 namespace inference {
 namespace tensorrt {
 
-TEST(tensorrt, ConvertBlock) {
+TEST(BlockConverter, ConvertBlock) {
   framework::ProgramDesc prog;
   auto* block = prog.MutableBlock(0);
   auto* mul_op = block->AppendOp();
   mul_op->SetType("mul");
   auto* conv2d_op = block->AppendOp();
   conv2d_op->SetType("conv2d");
 
-  TensorRTConverter converter;
+  BlockConverter converter;
   converter.ConvertBlock(*block);
 }