[Paddle-TRT] Add Paddle-TRT anchor generator support

paddle/fluid/inference/api/analysis_predictor.cc

@@ -1192,6 +1192,7 @@ USE_TRT_CONVERTER(scale);
 USE_TRT_CONVERTER(stack);
 USE_TRT_CONVERTER(clip);
 USE_TRT_CONVERTER(gather);
+USE_TRT_CONVERTER(anchor_generator);
 USE_TRT_CONVERTER(yolo_box);
 USE_TRT_CONVERTER(roi_align);
 USE_TRT_CONVERTER(affine_channel);

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

@@ -6,6 +6,7 @@ nv_library(tensorrt_converter
                 shuffle_channel_op.cc swish_op.cc instance_norm_op.cc stack_op.cc transpose_op.cc flatten_op.cc
                 emb_eltwise_layernorm.cc skip_layernorm.cc scale_op.cc slice_op.cc hard_sigmoid_op.cc hard_swish_op.cc clip_op.cc
                 gather_op.cc
+                anchor_generator_op.cc
                 yolo_box_op.cc
                 roi_align_op.cc
                 affine_channel_op.cc

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

@@ -0,0 +1,79 @@
+/* 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. */
+
+#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
+#include "paddle/fluid/inference/tensorrt/plugin/anchor_generator_op_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+/* Anchor Generator Op */
+class AnchorGeneratorOpConverter : public OpConverter {
+ public:
+  void operator()(const paddle::framework::proto::OpDesc& op,
+                  const paddle::framework::Scope& scope,
+                  bool test_mode) override {
+    VLOG(3) << "convert a fluid anchor generator op to tensorrt plugin";
+    framework::OpDesc op_desc(op, nullptr);
+    std::string input_name = op_desc.Input("Input").front();
+    std::string anchor_name = op_desc.Output("Anchors").front();
+    std::string variance_name = op_desc.Output("Variances").front();
+
+    auto* input = engine_->GetITensor(input_name);
+    const auto input_dims = input->getDimensions();  // C, H, W
+    std::vector<std::string> output_names{anchor_name, variance_name};
+
+    const auto anchor_sizes =
+        BOOST_GET_CONST(std::vector<float>, op_desc.GetAttr("anchor_sizes"));
+    const auto aspect_ratios =
+        BOOST_GET_CONST(std::vector<float>, op_desc.GetAttr("aspect_ratios"));
+    const auto stride =
+        BOOST_GET_CONST(std::vector<float>, op_desc.GetAttr("stride"));
+    const auto variances =
+        BOOST_GET_CONST(std::vector<float>, op_desc.GetAttr("variances"));
+    const auto offset = BOOST_GET_CONST(float, op_desc.GetAttr("offset"));
+    const int num_anchors = aspect_ratios.size() * anchor_sizes.size();
+    bool is_dynamic = engine_->with_dynamic_shape();
+    const auto height = input_dims.d[1];
+    const auto width = input_dims.d[2];
+    const int box_num = width * height * num_anchors;
+    const nvinfer1::DataType data_type = nvinfer1::DataType::kFLOAT;
+
+    nvinfer1::IPluginV2* anchor_generator_plugin = nullptr;
+    if (is_dynamic) {
+      anchor_generator_plugin = new plugin::AnchorGeneratorPluginDynamic(
+          data_type, anchor_sizes, aspect_ratios, stride, variances, offset,
+          num_anchors);
+    } else {
+      anchor_generator_plugin = new plugin::AnchorGeneratorPlugin(
+          data_type, anchor_sizes, aspect_ratios, stride, variances, offset,
+          height, width, num_anchors, box_num);
+    }
+
+    std::vector<nvinfer1::ITensor*> anchor_generator_inputs{input};
+    auto* anchor_generator_layer = engine_->network()->addPluginV2(
+        anchor_generator_inputs.data(), anchor_generator_inputs.size(),
+        *anchor_generator_plugin);
+
+    RreplenishLayerAndOutput(anchor_generator_layer, "anchor_generator",
+                             output_names, test_mode);
+  }
+};
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle
+
+REGISTER_TRT_OP_CONVERTER(anchor_generator, AnchorGeneratorOpConverter);

paddle/fluid/inference/tensorrt/op_teller.cc

@@ -116,6 +116,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       "affine_channel",
       "multiclass_nms",
       "nearest_interp",
+      "anchor_generator",
   };
 };
 

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

@@ -5,6 +5,7 @@ nv_library(tensorrt_plugin
            instance_norm_op_plugin.cu emb_eltwise_layernorm_plugin.cu
            qkv_to_context_plugin.cu skip_layernorm_op_plugin.cu slice_op_plugin.cu
            hard_swish_op_plugin.cu stack_op_plugin.cu special_slice_plugin.cu
+           anchor_generator_op_plugin.cu
            yolo_box_op_plugin.cu
            roi_align_op_plugin.cu
            DEPS enforce tensorrt_engine prelu tensor bert_encoder_functor)