[Paddle-TRT] Add Paddle-TRT Yolobox 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(yolo_box);
 USE_TRT_CONVERTER(roi_align);
 USE_TRT_CONVERTER(affine_channel);
 USE_TRT_CONVERTER(multiclass_nms);

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
+                yolo_box_op.cc
                 roi_align_op.cc
                 affine_channel_op.cc
                 multiclass_nms_op.cc

paddle/fluid/inference/tensorrt/convert/yolo_box_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 <vector>
+#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
+#include "paddle/fluid/inference/tensorrt/plugin/yolo_box_op_plugin.h"
+
+namespace paddle {
+namespace framework {
+class Scope;
+namespace proto {
+class OpDesc;
+}  // namespace proto
+}  // namespace framework
+}  // namespace paddle
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+class YoloBoxOpConverter : public OpConverter {
+ public:
+  void operator()(const framework::proto::OpDesc& op,
+                  const framework::Scope& scope, bool test_mode) override {
+    VLOG(3) << "convert a fluid yolo box op to tensorrt plugin";
+
+    framework::OpDesc op_desc(op, nullptr);
+    std::string X = op_desc.Input("X").front();
+    std::string img_size = op_desc.Input("ImgSize").front();
+
+    auto* X_tensor = engine_->GetITensor(X);
+    auto* img_size_tensor = engine_->GetITensor(img_size);
+
+    int class_num = BOOST_GET_CONST(int, op_desc.GetAttr("class_num"));
+    std::vector<int> anchors =
+        BOOST_GET_CONST(std::vector<int>, op_desc.GetAttr("anchors"));
+
+    int downsample_ratio =
+        BOOST_GET_CONST(int, op_desc.GetAttr("downsample_ratio"));
+    float conf_thresh = BOOST_GET_CONST(float, op_desc.GetAttr("conf_thresh"));
+    bool clip_bbox = BOOST_GET_CONST(bool, op_desc.GetAttr("clip_bbox"));
+    float scale_x_y = BOOST_GET_CONST(float, op_desc.GetAttr("scale_x_y"));
+
+    int type_id = static_cast<int>(engine_->WithFp16());
+    auto input_dim = X_tensor->getDimensions();
+    auto* yolo_box_plugin = new plugin::YoloBoxPlugin(
+        type_id ? nvinfer1::DataType::kHALF : nvinfer1::DataType::kFLOAT,
+        anchors, class_num, conf_thresh, downsample_ratio, clip_bbox, scale_x_y,
+        input_dim.d[1], input_dim.d[2]);
+
+    std::vector<nvinfer1::ITensor*> yolo_box_inputs;
+    yolo_box_inputs.push_back(X_tensor);
+    yolo_box_inputs.push_back(img_size_tensor);
+
+    auto* yolo_box_layer = engine_->network()->addPluginV2(
+        yolo_box_inputs.data(), yolo_box_inputs.size(), *yolo_box_plugin);
+
+    std::vector<std::string> output_names;
+    output_names.push_back(op_desc.Output("Boxes").front());
+    output_names.push_back(op_desc.Output("Scores").front());
+
+    RreplenishLayerAndOutput(yolo_box_layer, "yolo_box", output_names,
+                             test_mode);
+  }
+};
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle
+
+REGISTER_TRT_OP_CONVERTER(yolo_box, YoloBoxOpConverter);

paddle/fluid/inference/tensorrt/op_teller.cc

@@ -111,6 +111,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       "flatten2",
       "flatten",
       "gather",
+      "yolo_box",
       "roi_align",
       "affine_channel",
       "multiclass_nms",
@@ -198,6 +199,15 @@ bool OpTeller::Tell(const framework::ir::Node* node, bool use_no_calib_int8,
       if (!with_dynamic_shape || desc.Input("Axis").size() > 0) return false;
     }
 
+    if (op_type == "yolo_box") {
+      if (with_dynamic_shape) return false;
+      bool has_attrs =
+          (desc.HasAttr("class_num") && desc.HasAttr("anchors") &&
+           desc.HasAttr("downsample_ratio") && desc.HasAttr("conf_thresh") &&
+           desc.HasAttr("clip_bbox") && desc.HasAttr("scale_x_y"));
+      return has_attrs;
+    }
+
     if (op_type == "affine_channel") {
       if (!desc.HasAttr("data_layout")) return false;
       auto data_layout = framework::StringToDataLayout(

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
+           yolo_box_op_plugin.cu
            roi_align_op_plugin.cu
            DEPS enforce tensorrt_engine prelu tensor bert_encoder_functor)