auto registray op converters

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

 nv_library(tensorrt_convert SRCS convert.cc DEPS dynload_cuda)
-nv_test(tensorrt_convert_test SRCS convert_test.cc DEPS tensorrt paddle_fluid)
+nv_test(test_tensorrt_convert SRCS test_convert.cc DEPS tensorrt paddle_fluid)

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

@@ -13,39 +13,23 @@ 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/convert_conv2d.h"
+#include "paddle/fluid/inference/tensorrt/convert/convert_mul.h"
 
 namespace paddle {
 namespace inference {
 namespace tensorrt {
 
-void TensorRTConverter::ConvertOp(const framework::OpDesc& op) {
-  std::string type = op.Type();
-  PADDLE_ENFORCE(op_registry_.count(type), "No converter registered for op: %s",
-                 type);
-  std::function<void(const framework::OpDesc&)> op_converter =
-      op_registry_.at(type);
-  op_converter(op);
-}
-
 void TensorRTConverter::ConvertBlock(const framework::BlockDesc& block) {
   for (auto op : block.AllOps()) {
-    ConvertOp(*op);
+    std::string type = op->Type();
+    PADDLE_ENFORCE(GetOpConverter().count(type),
+                   "No converter registered for op: %s", type);
+    auto op_converter = GetOpConverter()[type];
+    op_converter->Convert(*op);
   }
 }
 
-void TensorRTConverter::RegisterOpConverters() {
-  op_registry_["mul"] = ConvertMul;
-  op_registry_["conv2d"] = ConvertConv2D;
-}
-
-void TensorRTConverter::ConvertMul(const framework::OpDesc& op) {
-  LOG(INFO) << "convert a fluid mul op to tensorrt fc layer without bias";
-}
-
-void TensorRTConverter::ConvertConv2D(const framework::OpDesc& op) {
-  LOG(INFO) << "convert a fluid Conv2d op to tensorrt conv layer without bias";
-}
-
 }  // namespace tensorrt
 }  // namespace inference
 }  // namespace paddle

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

@@ -26,37 +26,46 @@ namespace paddle {
 namespace inference {
 namespace tensorrt {
 
-class TensorRTConverter {
+class ConverterBase {
  public:
-  explicit TensorRTConverter(const framework::Scope& scope) : scope_(scope) {
-    this->RegisterOpConverters();
-  }
+  ConverterBase() {}
 
-  // convert fluid op to tensorrt layer
-  void ConvertOp(const framework::OpDesc& op);
-
-  // convert fluid block to tensorrt network
-  void ConvertBlock(const framework::BlockDesc& block);
-
- private:
-  // convert op registry, whose key is the fluid op type, and value is the
-  // convert tensorrt function name
-  std::unordered_map<std::string, std::function<void(const framework::OpDesc&)>>
-      op_registry_;
   // fluid inference scope
-  const framework::Scope& scope_;
+  framework::Scope* scope_;
   // tensorrt input/output tensor list, 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_;
+};
 
-  // register different op converters
-  void RegisterOpConverters();
+class OpConverter : public ConverterBase {
+ public:
+  OpConverter() {}
+  virtual ~OpConverter() {}
 
-  // convert a fluid Mul op to tensorrt fc layer without bias
-  static void ConvertMul(const framework::OpDesc& op);
+  // convert fluid op to tensorrt layer
+  virtual void Convert(const framework::OpDesc& op) = 0;
+};
+
+static std::unordered_map<std::string, OpConverter*>& GetOpConverter() {
+  static std::unordered_map<std::string, OpConverter*> register_op_converter;
+  return register_op_converter;
+}
 
-  // convert a fluid Conv2d op to tensorrt conv layer without bias
-  static void ConvertConv2D(const framework::OpDesc& op);
+#define REGISTER_TRT_OP_CONVETER(op_type, convert_class) \
+  class convert_class##Register {                        \
+   public:                                               \
+    convert_class##Register() {                          \
+      GetOpConverter()[#op_type] = new convert_class;    \
+    }                                                    \
+  };                                                     \
+  convert_class##Register convert_class##reg;
+
+class TensorRTConverter : public ConverterBase {
+ public:
+  TensorRTConverter() {}
+
+  // convert fluid block to tensorrt network
+  void ConvertBlock(const framework::BlockDesc& block);
 };
 
 }  // namespace tensorrt

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

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

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

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

paddle/fluid/inference/tensorrt/convert/convert_test.cc → paddle/fluid/inference/tensorrt/convert/test_convert.cc

@@ -12,9 +12,9 @@ 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 <gtest/gtest.h>
 #include "paddle/fluid/framework/program_desc.h"
+#include "paddle/fluid/inference/tensorrt/convert/convert.h"
 
 namespace paddle {
 namespace inference {
@@ -28,8 +28,7 @@ TEST(tensorrt, ConvertBlock) {
   auto* conv2d_op = block->AppendOp();
   conv2d_op->SetType("conv2d");
 
-  framework::Scope scope;
-  TensorRTConverter converter(scope);
+  TensorRTConverter converter;
   converter.ConvertBlock(*block);
 }