[Paddle TensorRT] add grid_sample op converter (#68787)

gzy19990617 · web-flow · commit 8c2f6c55e687 · 2024-11-13T17:20:08.000+08:00
* add grid_sample op converter

* merge develop

* fix

* fix review

* 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
@@ -1416,6 +1416,43 @@ class NearestInterV2Pattern
   }
 };
 
+class GridSampleOpPattern
+    : public pir::OpRewritePattern<paddle::dialect::GridSampleOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::GridSampleOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::GridSampleOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op->attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+#if IS_TRT_VERSION_LT(8510)
+    VLOG(3) << "grid_sample is not supported when TensorRT < 8.5.1";
+    return false;
+#else
+    if (!op->HasAttribute("mode") || !op->HasAttribute("padding_mode") ||
+        !op->HasAttribute("align_corners")) {
+      VLOG(3)
+          << "grid_sample need attributes: mode, padding_mode, align_corners";
+      return false;
+    }
+    auto x = op.operand_source(0);
+    auto grid = op.operand_source(1);
+    auto x_shape = pir::GetShapeFromValue(x);
+    auto grid_shape = pir::GetShapeFromValue(grid);
+
+    if (x_shape.size() != 4 || grid_shape.size() != 4) {
+      VLOG(3) << "The input and grid tensors must be shape tensors of rank 4 "
+                 "when using TRT GridSample layer.";
+      return false;
+    }
+#endif
+
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class StackOpPattern : public pir::OpRewritePattern<paddle::dialect::StackOp> {
  public:
   using pir::OpRewritePattern<paddle::dialect::StackOp>::OpRewritePattern;
@@ -1776,6 +1813,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<SumOpPattern>(context));
     ps.Add(std::make_unique<BilinearInterpV2Pattern>(context));
     ps.Add(std::make_unique<NearestInterV2Pattern>(context));
+    ps.Add(std::make_unique<GridSampleOpPattern>(context));
     ps.Add(std::make_unique<StackOpPattern>(context));
     ps.Add(std::make_unique<TanhOpPattern>(context));
     ps.Add(std::make_unique<WherePattern>(context));
diff --git a/python/paddle/tensorrt/converter.py b/python/paddle/tensorrt/converter.py
@@ -44,6 +44,7 @@
 from .impls.pooling import *  # noqa: F403
 from .impls.search import *  # noqa: F403
 from .impls.stat import *  # noqa: F403
+from .impls.vision import *  # noqa: F403
 from .register import converter_registry
 from .util import get_trt_version_list, map_dtype
 
diff --git a/python/paddle/tensorrt/impls/vision.py b/python/paddle/tensorrt/impls/vision.py
@@ -0,0 +1,46 @@
+# Copyright (c) 2024 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.
+
+import tensorrt as trt
+
+from paddle.tensorrt.register import converter_registry
+
+
+@converter_registry.register("pd_op.grid_sample", trt_version="8.x")
+def grid_sample_converter(network, paddle_op, inputs):
+    input_tensor, grid_tensor = inputs
+    padding = paddle_op.attrs().get("paddings", [0, 0])
+
+    mode = paddle_op.attrs().get("mode", "bilinear")
+    padding_mode = paddle_op.attrs().get("padding_mode", "zeros")
+    align_corners = paddle_op.attrs().get("align_corners", True)
+
+    if padding_mode == "zeros":
+        sample_mode = trt.SampleMode.FILL
+    elif padding_mode == "border":
+        sample_mode = trt.SampleMode.CLAMP
+    elif padding_mode == "reflection":
+        sample_mode = trt.SampleMode.REFLECT
+
+    if mode == "nearest":
+        interpolation_mode = trt.InterpolationMode.NEAREST
+    elif mode == "bilinear":
+        interpolation_mode = trt.InterpolationMode.LINEAR
+
+    grid_sample_layer = network.add_grid_sample(input_tensor, grid_tensor)
+
+    grid_sample_layer.interpolation_mode = interpolation_mode
+    grid_sample_layer.align_corners = align_corners
+    grid_sample_layer.sample_mode = sample_mode
+    return grid_sample_layer.get_output(0)
diff --git a/test/tensorrt/CMakeLists.txt b/test/tensorrt/CMakeLists.txt
@@ -23,6 +23,7 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   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_vision PROPERTIES TIMEOUT "500")
   set_tests_properties(test_converter_linalg PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_search PROPERTIES TIMEOUT "300")
   set_tests_properties(test_converter_logic PROPERTIES TIMEOUT "300")
diff --git a/test/tensorrt/test_converter_vision.py b/test/tensorrt/test_converter_vision.py
@@ -0,0 +1,99 @@
+# Copyright (c) 2024 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.
+
+import unittest
+
+import numpy as np
+from tensorrt_test_base import TensorRTBaseTest
+
+import paddle.nn.functional as F
+
+
+class TestGridSampleTRTPatternBase(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = F.grid_sample
+        self.api_args = {
+            "x": np.array(
+                [[[[-0.6, 0.8, -0.5], [-0.5, 0.2, 1.2], [1.4, 0.3, -0.2]]]]
+            ).astype("float32"),
+            "grid": np.array(
+                [
+                    [
+                        [[0.2, 0.3], [-0.4, -0.3], [-0.9, 0.3], [-0.9, -0.6]],
+                        [[0.4, 0.1], [0.9, -0.8], [0.4, 0.5], [0.5, -0.2]],
+                        [[0.1, -0.8], [-0.3, -1.0], [0.7, 0.4], [0.2, 0.8]],
+                    ]
+                ],
+                dtype='float32',
+            ),
+        }
+        self.program_config = {"feed_list": ["x", "grid"]}
+        self.min_shape = {"x": [1, 1, 3, 3], "grid": [1, 3, 4, 2]}
+        self.max_shape = {"x": [5, 1, 3, 3], "grid": [5, 3, 4, 2]}
+
+
+class TestGridSampleTRTPatternCase1(TestGridSampleTRTPatternBase):
+    """default:mode='bilinear', padding_mode='zeros', align_corners=True"""
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestGridSampleTRTPatternCase2(TestGridSampleTRTPatternBase):
+    """default:mode='nearest', padding_mode='reflection', align_corners=False"""
+
+    def setUp(self):
+        super().setUp()
+        self.api_args.update(
+            {
+                "mode": "nearest",
+                "padding_mode": "reflection",
+                "align_corner": False,
+            }
+        )
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestGridSampleTRTPatternCase3(TestGridSampleTRTPatternBase):
+    """default:mode='nearest', padding_mode='border', align_corners=True"""
+
+    def setUp(self):
+        super().setUp()
+        self.api_args.update({"mode": "nearest", "padding_mode": "border"})
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestGridSampleTRTPatternCase4(TestGridSampleTRTPatternBase):
+    """default:mode='bilinear', padding_mode='border', align_corners=False"""
+
+    def setUp(self):
+        super().setUp()
+        self.api_args.update(
+            {
+                "mode": "bilinear",
+                "padding_mode": "border",
+                "align_corner": False,
+            },
+        )
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+if __name__ == '__main__':
+    unittest.main()
