[PaddleTRT] add linear_interp_v2 converter and fix arg_max converter (#70939)

bukejiyu · web-flow · commit ef649042514b · 2025-01-26T00:34:33.000-08:00
* old trt fix

* update

* update
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -3385,6 +3385,7 @@ USE_TRT_CONVERTER(multiclass_nms3);
 USE_TRT_CONVERTER(nearest_interp);
 USE_TRT_CONVERTER(nearest_interp_v2);
 USE_TRT_CONVERTER(bilinear_interp_v2);
+USE_TRT_CONVERTER(linear_interp_v2);
 USE_TRT_CONVERTER(reshape);
 USE_TRT_CONVERTER(reshape2);
 USE_TRT_CONVERTER(gather_nd);
diff --git a/paddle/fluid/inference/tensorrt/convert/CMakeLists.txt b/paddle/fluid/inference/tensorrt/convert/CMakeLists.txt
@@ -96,6 +96,7 @@ list(
   reverse_roll_op.cc
   tanhshrink_op.cc
   take_along_axis_op.cc
+  linear_interp_v2_op.cc
   logsigmoid_op.cc
   preln_layernorm_shift_partition_op.cc
   trans_layernorm_op.cc
diff --git a/paddle/fluid/inference/tensorrt/convert/arg_max_op.cc b/paddle/fluid/inference/tensorrt/convert/arg_max_op.cc
@@ -42,14 +42,21 @@ class ArgMaxOpConverter : public OpConverter {
                               {output_name + "_value", output_name},
                               test_mode);
     } else {
+      int topk_out_shape_size =
+          topk_layer->getOutput(1)->getDimensions().nbDims;
+      std::vector<bool> should_squeeze(topk_out_shape_size, false);
+      should_squeeze[axis] = true;
+      std::vector<int> gather_indices;
+      for (int i = 0; i < topk_out_shape_size; ++i) {
+        if (!should_squeeze[i]) {
+          gather_indices.push_back(i);
+        }
+      }
       auto squeeze_layer =
           TRT_ENGINE_ADD_LAYER(engine_, Shuffle, *topk_layer->getOutput(1));
-      auto dims = input_dims;
-      dims.nbDims -= 1;
-      for (int i = axis; i < dims.nbDims; i++) {
-        dims.d[i] = dims.d[i + 1];
-      }
-      squeeze_layer->setReshapeDimensions(dims);
+      auto shape_tesnor = Shape(topk_layer->getOutput(1));
+      auto real_shape_tensor = Gather(shape_tesnor, gather_indices);
+      squeeze_layer->setInput(1, *real_shape_tensor);
       ReplenishLayerAndOutput(
           squeeze_layer, "arg_max", {output_name}, test_mode);
     }
diff --git a/paddle/fluid/inference/tensorrt/convert/linear_interp_v2_op.cc b/paddle/fluid/inference/tensorrt/convert/linear_interp_v2_op.cc
@@ -0,0 +1,152 @@
+/* Copyright (c) 2025 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
+
+https://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/framework/data_layout.h"
+#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
+namespace paddle::inference::tensorrt {
+
+class LinearInterpolateV2OpConverter : public OpConverter {
+ public:
+  void operator()(const framework::proto::OpDesc& op,
+                  const framework::Scope& scope,
+                  bool test_mode) override {
+    VLOG(3) << "convert a linear_interp_v2_op op to tensorrt OP";
+
+    framework::OpDesc op_desc(op, nullptr);
+
+    std::string input_name = op_desc.Input("X").front();
+    std::string output_name = op_desc.Output("Out").front();
+
+    auto input = engine_->GetITensor(input_name);
+
+    auto data_layout = common::StringToDataLayout(
+        PADDLE_GET_CONST(std::string, op_desc.GetAttr("data_layout")));
+    auto interp_method =
+        PADDLE_GET_CONST(std::string, op_desc.GetAttr("interp_method"));
+    bool align_corners =
+        PADDLE_GET_CONST(bool, op_desc.GetAttr("align_corners"));
+
+    auto resize_inputs = op_desc.Inputs();
+    auto input_names = op_desc.Input("X");
+
+    auto out_w = PADDLE_GET_CONST(int, op_desc.GetAttr("out_w"));
+
+    auto layer = TRT_ENGINE_ADD_LAYER(engine_, Resize, *input);
+#if IS_TRT_VERSION_GE(8600)
+    layer->setResizeMode(nvinfer1::InterpolationMode::kLINEAR);
+#else
+    layer->setResizeMode(nvinfer1::ResizeMode::kLINEAR);
+#endif
+    // scale > outputsize
+
+#if IS_TRT_VERSION_GE(8000)
+    if (align_corners == true) {
+      layer->setCoordinateTransformation(
+          nvinfer1::ResizeCoordinateTransformation::kALIGN_CORNERS);
+    } else {
+      layer->setCoordinateTransformation(
+          nvinfer1::ResizeCoordinateTransformation::kHALF_PIXEL);
+    }
+#endif
+#if !IS_TRT_VERSION_GE(8000)
+    layer->setAlignCorners(align_corners);
+#endif
+    auto in_dim = input->getDimensions();
+    float scale_w = -1.f;
+
+    // Scales Priority: Scale(tensor) > scale(attr) > out_d/out_h/out_w(attr)
+    bool has_scale_input_attr =
+        (resize_inputs.find("Scale") != resize_inputs.end());
+    bool has_scale_input =
+        has_scale_input_attr && (!op_desc.Input("Scale").empty());
+    if (has_scale_input) {
+      auto* scale_var = scope.FindVar(op_desc.Input("Scale")[0]);
+      auto* scale_tensor = scale_var->GetMutable<phi::DenseTensor>();
+      auto* scale_d = scale_tensor->data<float>();
+      scale_w = scale_d[0];
+    } else {
+      const std::vector<float> scale_attr =
+          PADDLE_GET_CONST(std::vector<float>, op_desc.GetAttr("scale"));
+      if (scale_attr.size() > 0) {
+        scale_w = scale_attr[0];
+      }
+    }
+
+    // axis are different in static/dynamic mode
+    bool with_dynamic = true;
+    int w_axis = (data_layout == phi::DataLayout::kNCHW) + with_dynamic;
+
+    if (scale_w > 0.) {
+      out_w = static_cast<int>(in_dim.d[w_axis] * scale_w);
+    }
+
+    // Priority: Input(OutSize) > Input(SizeTensor)> attr(out_h/out_w) >
+    // attr(scale)
+    nvinfer1::ITensor* outsize_tensor = nullptr;
+    if (resize_inputs.find("OutSize") != resize_inputs.end()) {
+      if (!op_desc.Input("OutSize").empty()) {
+        outsize_tensor = engine_->GetITensor(op_desc.Input("OutSize")[0]);
+      }
+    }
+
+#if IS_TRT_VERSION_GE(8200)
+    if (outsize_tensor == nullptr) {
+      if (resize_inputs.find("SizeTensor") != resize_inputs.end()) {
+        if (op_desc.Input("SizeTensor").size() >= 1) {
+          outsize_tensor = engine_->GetITensor(op_desc.Input("SizeTensor")[0]);
+        }
+      }
+    }
+#endif
+
+    if (out_w > 0 && !(scale_w > 0.)) {
+      scale_w =
+          static_cast<float>(out_w) / static_cast<float>(in_dim.d[w_axis]);
+    }
+
+    std::vector<float> scales;
+    scales.push_back(1.f);
+    if (data_layout == phi::DataLayout::kNCHW) {
+      scales.push_back(1.f);
+      scales.push_back(scale_w);
+    } else if (data_layout == phi::DataLayout::kNHWC) {
+      scales.push_back(scale_w);
+      scales.push_back(1.f);
+    }
+
+    if (outsize_tensor != nullptr) {
+      std::vector<nvinfer1::ITensor*> outsize_itensors;
+      auto* input_shape = Shape(input);
+      outsize_itensors.push_back(GetEleTensorOfShape(input_shape, 0));
+
+      if (data_layout == phi::DataLayout::kNCHW) {
+        outsize_itensors.push_back(GetEleTensorOfShape(input_shape, 1));
+        outsize_itensors.push_back(outsize_tensor);
+      } else if (data_layout == phi::DataLayout::kNHWC) {
+        outsize_itensors.push_back(outsize_tensor);
+        outsize_itensors.push_back(GetEleTensorOfShape(input_shape, 2));
+      }
+      layer->setInput(1, *Concat(outsize_itensors));
+    } else {
+      layer->setScales(scales.data(), scales.size());
+    }
+
+    ReplenishLayerAndOutput(
+        layer, "linear_interp_v2_op", {output_name}, test_mode);
+  }
+};
+
+}  // namespace paddle::inference::tensorrt
+
+REGISTER_TRT_OP_CONVERTER(linear_interp_v2, LinearInterpolateV2OpConverter);
diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -1023,50 +1023,92 @@ struct SimpleOpTypeSetTeller : public Teller {
         }
       }
     }
-
-    if (op_type == "squeeze2") {
-      // If Attribute is Variable(s), HasAttr() will return False
-      if (!desc.HasAttr("axes", /*with_attr_var=*/false)) {
-        VLOG(3) << "Skip to convert into TRT while found Attribute('axes') is "
-                   "Variable type in squeeze2.";
-        return false;
+    if (op_type == "linear_interp_v2") {
+#if IS_TRT_VERSION_LT(7100)
+      return false;
+#endif
+      std::vector<std::string> attrs{"data_layout",
+                                     "interp_method",
+                                     "align_corners",
+                                     "scale",
+                                     "out_h",
+                                     "out_w"};
+      for (auto const& attr : attrs) {
+        if (!desc.HasAttr(attr)) {
+          VLOG(3) << "The op_type " << op_type << " doesn't have the attr "
+                  << attr << " and return false";
+          return false;
+        }
       }
 
-      std::vector<int> axes;
-      if (desc.HasAttr("axes")) {
-        axes = PADDLE_GET_CONST(std::vector<int>, desc.GetAttr("axes"));
-      }
-      if (axes.empty()) {
-        auto* block = desc.Block();
-        if (block) {
-          auto input_var_name = desc.Input("X")[0];
-          auto* input_var_desc = block->FindVarRecursive(input_var_name);
-          const auto input_shape = input_var_desc->GetShape();
-          for (int s : input_shape) {
-            if (s == -1) {
-              VLOG(3) << "The necessary attributes of the squeeze2 operator "
-                         "axes is "
-                         "missing. ss ==== -1";
-              return false;
-            } else if (s == 1) {
-              axes.push_back(s);
-            }
-          }
+      auto resize_inputs = desc.Inputs();
+      if (resize_inputs.find("SizeTensor") != resize_inputs.end()) {
+#if IS_TRT_VERSION_GE(8200)
+        if (desc.Input("SizeTensor").size() == 1) {
+          return true;
         }
-        if (axes.empty()) {
+#else
+        if (!desc.Input("SizeTensor").empty()) {
           VLOG(3)
-              << "The necessary attributes of the squeeze2 operator axes is "
-                 "missing.";
+              << "The Paddle-TRT doesn't support the SizeTensor for op_type "
+              << op_type;
           return false;
         }
+#endif
       }
-      if (!with_dynamic_shape) {
-        if (std::find(axes.begin(), axes.end(), 0) != axes.end()) {
-          VLOG(3) << "Invalid squeeze axes. Axes having batch axis is not "
-                     "supported in static shape";
+      if (resize_inputs.find("OutSize") != resize_inputs.end()) {
+        if (!with_dynamic_shape) {
+          VLOG(3) << "Static shape don't support the OutSize for op_type "
+                  << op_type;
           return false;
         }
       }
+
+      auto data_layout = common::StringToDataLayout(
+          PADDLE_GET_CONST(std::string, desc.GetAttr("data_layout")));
+      if (data_layout != phi::DataLayout::kNCHW &&
+          data_layout != phi::DataLayout::kNHWC) {
+        VLOG(3) << "The op_type " << op_type
+                << " is not NCHW or NHWC return false";
+        return false;
+      }
+      auto interp_method =
+          PADDLE_GET_CONST(std::string, desc.GetAttr("interp_method"));
+      if (interp_method != "linear") {
+        VLOG(3) << "The interp_method of op_type " << op_type
+                << " is not linear";
+        return false;
+      }
+      bool has_scale_input_size =
+          (resize_inputs.find("Scale") != resize_inputs.end());
+      if (!has_scale_input_size ||
+          (has_scale_input_size && desc.Input("Scale").size() != 1)) {
+        const std::vector<float> scale =
+            PADDLE_GET_CONST(std::vector<float>, desc.GetAttr("scale"));
+        if (scale.size() == 0) {
+          if (!desc.HasAttr("out_w")) {
+            VLOG(3) << "The op_type " << op_type
+                    << " doesn't have Scale and the scale size <=1 and without "
+                       " out_w, it will return false";
+            return false;
+          }
+          auto out_w = PADDLE_GET_CONST(int, desc.GetAttr("out_w"));
+          if (out_w <= 0) {
+            VLOG(3) << "The op_type " << op_type
+                    << "'s 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 && with_dynamic_shape) {
+              VLOG(3) << "dynamic shape not support Attr(scale[" << i << "]) "
+                      << scale[i]
+                      << " less than 1 and Input(Scale) vector not set.";
+              return false;
+            }
+          }
+        }
+      }
     }
 
     if (op_type == "unsqueeze2") {
@@ -3102,6 +3144,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       "mish",
       "nearest_interp_v2",
       "bilinear_interp_v2",
+      "linear_interp_v2",
       "pool3d",
       "deformable_conv",
       "relu6",
@@ -3276,6 +3319,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       "conv3d_transpose",
       "mish",
       "bilinear_interp_v2",
+      "linear_interp_v2",
       "nearest_interp_v2",
       "pool3d",
       "deformable_conv",
diff --git a/test/ir/inference/test_trt_convert_linear_interp_v2.py b/test/ir/inference/test_trt_convert_linear_interp_v2.py
