【Paddle TensorRT】pd_op.linear_interp (#71013)

lizexu123 · web-flow · commit 6cf52b949639 · 2025-02-17T15:07:54.000+08:00
* pd_op.linear_interp

* fix

* fix

* fix

* 增加单测覆盖率

* fix

* fix

* fix
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -2501,6 +2501,120 @@ class FullBatchSizeLikeOpPattern
   }
 };
 
+class LinearInterpOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::LinearInterpOp> {
+ public:
+  using pir::OpRewritePattern<
+      paddle::dialect::LinearInterpOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::LinearInterpOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    const std::vector<std::string> required_attrs = {"data_format",
+                                                     "interp_method",
+                                                     "align_corners",
+                                                     "scale",
+                                                     "out_h",
+                                                     "out_w"};
+    for (const auto &attr : required_attrs) {
+      if (!op->HasAttribute(attr)) {
+        VLOG(3) << "pd_op.linear_interp " << attr
+                << " attribute does not exist";
+        return false;
+      }
+    }
+
+    pir::Value size_tensor = op.operand_source(2);
+    if (size_tensor) {
+      auto size_tensor_type = size_tensor.type();
+      if (size_tensor_type.isa<pir::VectorType>()) {
+        auto vector_type = size_tensor.type().dyn_cast<pir::VectorType>();
+        if (vector_type.size() == 1) {
+          op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+          return true;
+        }
+      }
+    }
+
+    if (size_tensor.impl() != nullptr) {
+      VLOG(3) << "The Paddle-TRT doesn't support the SizeTensor for "
+                 "pd_op.linear_interp";
+      return false;
+    }
+
+    auto data_format =
+        op->attribute<pir::StrAttribute>("data_format").AsString();
+    if (data_format != "NCHW" && data_format != "NHWC") {
+      VLOG(3) << "pd_op.linear_interp data_format must be NCHW or NHWC";
+      return false;
+    }
+    auto interp_method =
+        op->attribute<pir::StrAttribute>("interp_method").AsString();
+    if (interp_method != "linear") {
+      VLOG(3) << "The interp_method of pd_op.linear_interp is not linear";
+      return false;
+    }
+
+    pir::Value scale_tensor = op.operand_source(3);
+
+    bool has_scale_input = false;
+    if (scale_tensor) {
+      has_scale_input = true;
+    }
+    if (has_scale_input) {
+      VLOG(3) << "pd_op.linear_interp has scale input can not into trt,support "
+                 "scale attribute into trt";
+      return false;
+    }
+
+    auto scale_tensor_type = scale_tensor.type();
+    int scale_shape = 0;
+    if (scale_tensor_type.isa<pir::VectorType>()) {
+      auto vector_type = scale_tensor.type().dyn_cast<pir::VectorType>();
+      scale_shape = vector_type.size();
+    }
+
+    if (!has_scale_input || (has_scale_input && scale_shape != 1)) {
+      std::vector<float> scale;
+      auto scale_attr = op->attribute<pir::ArrayAttribute>("scale");
+      for (const auto &attr : scale_attr.AsVector()) {
+        scale.push_back(attr.dyn_cast<pir::FloatAttribute>().data());
+      }
+      if (scale.size() == 0) {
+        if (!op->HasAttribute("out_w")) {
+          VLOG(3)
+              << "pd_op.linear_interp doesn't have scale_tensor and the scale "
+                 "size <=1 and without"
+                 "out_w, it will return false";
+          return false;
+        }
+        auto out_w = op->attribute<pir::Int32Attribute>("out_w").data();
+        if (out_w <= 0) {
+          VLOG(3)
+              << "pd_op.linear_interp out_w must be greater than 0 if scale "
+                 "is not set.";
+          return false;
+        }
+      } else {
+        for (size_t i = 0; i < scale.size(); i++) {
+          if (scale[i] <= 0) {
+            VLOG(3)
+                << "pd_op.linear_interp dynamic shape not support Attr(scale["
+                << i << "]" << scale[i]
+                << " less than 1 and Input(Scale) Vector not set.";
+            return false;
+          }
+        }
+      }
+    }
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class TrtOpMarkerPass : public pir::PatternRewritePass {
  public:
   TrtOpMarkerPass() : pir::PatternRewritePass("trt_op_marker_pass", 2) {}
@@ -2674,6 +2788,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
         std::make_unique<FusedBiasDropoutResidualLayerNormOpPattern>(context));
     ps.Add(std::make_unique<YoloBoxOpPattern>(context));
     ps.Add(std::make_unique<FullBatchSizeLikeOpPattern>(context));
+    ps.Add(std::make_unique<LinearInterpOpPattern>(context));
     return ps;
   }
 };
diff --git a/python/paddle/tensorrt/impls/common.py b/python/paddle/tensorrt/impls/common.py
@@ -358,3 +358,84 @@ def nearest_interp_converter(network, paddle_op, inputs):
         resize_layer.scales = scales
 
     return resize_layer.get_output(0)
+
+
+@converter_registry.register(
+    "pd_op.linear_interp", trt_version="trt_version_ge=8.0"
+)
+def linear_interp_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    data_layout = paddle_op.attrs().get("data_format")
+    interp_method = paddle_op.attrs().get("interp_method")
+    align_corners = paddle_op.attrs().get("align_corners")
+    out_w = paddle_op.attrs().get("out_w")
+    scale_attr = paddle_op.attrs().get("scale")
+    layer = network.add_resize(input_tensor)
+    trt_major = get_trt_version_list()[0]
+    trt_minor = get_trt_version_list()[1]
+    trt_version_float = float(f"{trt_major}.{trt_minor}")
+
+    if trt_version_float >= 8.6:
+        layer.resize_mode = trt.InterpolationMode.LINEAR
+    else:
+        layer.resize_mode = trt.ResizeMode.LINEAR
+
+    if align_corners:
+        layer.coordinate_transformation = (
+            trt.ResizeCoordinateTransformation.ALIGN_CORNERS
+        )
+    else:
+        layer.coordinate_transformation = (
+            trt.ResizeCoordinateTransformation.HALF_PIXEL
+        )
+
+    in_dim = input_tensor.shape
+    scale_w = -1.0
+
+    if scale_attr and len(scale_attr) > 0:
+        scale_w = scale_attr[0]
+
+    w_axis = 2 if data_layout == "NCHW" else 1
+
+    if float(scale_w) > 0.0:
+        out_w = int(in_dim[w_axis] * scale_w)
+
+    outsize_tensor = None
+    if len(inputs) > 1 and inputs[1] is not None:
+        outsize_tensor = inputs[1]
+
+    if outsize_tensor is None:
+        if len(inputs) > 2 and inputs[2] is not None:
+            outsize_tensor = inputs[2][0]
+
+    if out_w > 0 and scale_w <= 0:
+        scale_w = float(out_w) / float(in_dim[w_axis])
+
+    scales = [1.0]
+    if data_layout == "NCHW":
+        scales.append(1.0)
+        scales.append(scale_w)
+    elif data_layout == "NHWC":
+        scales.append(scale_w)
+        scales.append(1.0)
+
+    if outsize_tensor is not None:
+        outsize_itensors = []
+        input_shape = trt_shape(network, input_tensor)
+        batch_dim = get_shape_tensor_element(network, input_shape, 0)
+        outsize_itensors.append(batch_dim)
+
+        if data_layout == "NCHW":
+            channel_dim = get_shape_tensor_element(network, input_shape, 1)
+            outsize_itensors.append(channel_dim)
+            outsize_itensors.append(outsize_tensor)
+        elif data_layout == "NHWC":
+            outsize_itensors.append(outsize_tensor)
+            channel_dim = get_shape_tensor_element(network, input_shape, 2)
+            outsize_itensors.append(channel_dim)
+
+        layer.set_input(1, trt_concat(network, outsize_itensors))
+    else:
+        layer.scales = scales
+
+    return layer.get_output(0)
diff --git a/python/paddle/tensorrt/impls/ops.py b/python/paddle/tensorrt/impls/ops.py
@@ -139,34 +139,42 @@ def YoloBoxOpConverter(network, paddle_op, inputs):
         trt.PluginField(
             "conf_thresh",
             np.array(conf_thresh, dtype=np.float32),
+            trt.PluginFieldType.FLOAT32,
         ),
         trt.PluginField(
             "downsample_ratio",
             np.array(downsample_ratio, dtype=np.int32),
+            trt.PluginFieldType.INT32,
         ),
         trt.PluginField(
             "clip_bbox",
-            np.array(clip_bbox, dtype=np.bool),
+            np.array(clip_bbox, dtype=np.bool_),
+            trt.PluginFieldType.INT32,
         ),
         trt.PluginField(
             "scale_x_y",
             np.array(scale_x_y, dtype=np.float32),
+            trt.PluginFieldType.FLOAT32,
         ),
         trt.PluginField(
             "iou_aware",
-            np.array(iou_aware, dtype=np.bool),
+            np.array(iou_aware, dtype=np.bool_),
+            trt.PluginFieldType.INT32,
         ),
         trt.PluginField(
             "iou_aware_factor",
             np.array(iou_aware_factor, dtype=np.float32),
+            trt.PluginFieldType.FLOAT32,
         ),
         trt.PluginField(
             "h",
             np.array(input_h, dtype=np.int32),
+            trt.PluginFieldType.INT32,
         ),
         trt.PluginField(
             "w",
             np.array(input_w, dtype=np.int32),
+            trt.PluginFieldType.INT32,
         ),
     ]
     plugin_field_collection = trt.PluginFieldCollection(plugin_fields)
diff --git a/test/tensorrt/CMakeLists.txt b/test/tensorrt/CMakeLists.txt
@@ -22,7 +22,7 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   set_tests_properties(test_converter_manipulation PROPERTIES TIMEOUT "600")
   set_tests_properties(test_converter_creation PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_attribute PROPERTIES TIMEOUT "300")
-  set_tests_properties(test_converter_common PROPERTIES TIMEOUT "300")
+  set_tests_properties(test_converter_common PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_input PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_vision PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_linalg PROPERTIES TIMEOUT "100")
diff --git a/test/tensorrt/test_converter_common.py b/test/tensorrt/test_converter_common.py
diff --git a/test/tensorrt/test_converter_ops.py b/test/tensorrt/test_converter_ops.py
