【SCU】【Paddle TensorRT No.43】Add pd_op.leaky_relu 、pd_op.prelu converter (#70591)

PolaKuma · web-flow · commit 341ac15b67ff · 2025-02-20T10:48:12.000+08:00
* add

* Update test_converter_activation.py

* fix codestyle

* update

* Update test_converter_activation.py

* Update test_converter_activation.py

* update prelu

* fix codestyle

* Update activation.py

* Update trt_op_marker_pass.cc

* Update test_converter_activation.py

* Update test_converter_activation.py

* Update test_converter_activation.py
diff --git a/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc b/paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc
@@ -94,6 +94,8 @@ DEFINE_GENERAL_PATTERN(Flip, paddle::dialect::FlipOp)
 DEFINE_GENERAL_PATTERN(Mish, paddle::dialect::MishOp)
 DEFINE_GENERAL_PATTERN(AssignValue, paddle::dialect::AssignValueOp)
 DEFINE_GENERAL_PATTERN(AssignValue_, paddle::dialect::AssignValue_Op)
+DEFINE_GENERAL_PATTERN(LeakyRelu, paddle::dialect::LeakyReluOp)
+DEFINE_GENERAL_PATTERN(LeakyRelu_, paddle::dialect::LeakyRelu_Op)
 DEFINE_GENERAL_PATTERN(Anchor_Generator, paddle::dialect::AnchorGeneratorOp)
 DEFINE_GENERAL_PATTERN(Exp, paddle::dialect::ExpOp)
 DEFINE_GENERAL_PATTERN(Abs, paddle::dialect::AbsOp)
@@ -2538,6 +2540,25 @@ class AffineChannelOpPattern
   }
 };
 
+class PreluOpPattern : public pir::OpRewritePattern<paddle::dialect::PreluOp> {
+ public:
+  using pir::OpRewritePattern<paddle::dialect::PreluOp>::OpRewritePattern;
+  bool MatchAndRewrite(paddle::dialect::PreluOp op,
+                       pir::PatternRewriter &rewriter) const override {
+    if (op->HasAttribute(kCanRunTrtAttr) &&
+        op.attribute<pir::BoolAttribute>(kCanRunTrtAttr).data()) {
+      return false;
+    }
+    pir::Value alpha_var = op.operand_source(1);
+    if (!alpha_var) {
+      VLOG(3) << "Variable Alpha of prelu TRT converter not found.";
+      return false;
+    }
+    op->set_attribute(kCanRunTrtAttr, rewriter.bool_attr(true));
+    return true;
+  }
+};
+
 class YoloBoxOpPattern
     : public pir::OpRewritePattern<paddle::dialect::YoloBoxOp> {
  public:
@@ -2762,6 +2783,8 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ADD_PATTERN(Mish)
     ADD_PATTERN(AssignValue)
     ADD_PATTERN(AssignValue_)
+    ADD_PATTERN(LeakyRelu)
+    ADD_PATTERN(LeakyRelu_)
     ADD_PATTERN(Anchor_Generator)
     ADD_PATTERN(Exp)
     ADD_PATTERN(Abs)
@@ -2880,6 +2903,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<EinsumOpPattern>(context));
     ps.Add(std::make_unique<PNormOpPattern>(context));
     ps.Add(std::make_unique<AffineChannelOpPattern>(context));
+    ps.Add(std::make_unique<PreluOpPattern>(context));
     ps.Add(
         std::make_unique<FusedBiasDropoutResidualLayerNormOpPattern>(context));
     ps.Add(std::make_unique<YoloBoxOpPattern>(context));
diff --git a/python/paddle/tensorrt/impls/activation.py b/python/paddle/tensorrt/impls/activation.py
@@ -16,7 +16,9 @@
 import tensorrt as trt
 
 from paddle.tensorrt.converter_utils import (
+    add_1D_constant_layer,
     add_constant_layer,
+    trt_concat,
     trt_div,
     trt_min,
     trt_pow,
@@ -276,6 +278,16 @@ def thresholded_relu_converter(network, paddle_op, inputs):
     return thresholded_relu_layer.get_output(0)
 
 
+@converter_registry.register("pd_op.leaky_relu", trt_version="8.x")
+@converter_registry.register("pd_op.leaky_relu_", trt_version="8.x")
+def leaky_relu_converter(network, paddle_op, inputs):
+    x = inputs[0]
+    negative_slope = paddle_op.attrs()["negative_slope"]
+    leaky_relu_layer = network.add_activation(x, trt.ActivationType.LEAKY_RELU)
+    leaky_relu_layer.alpha = negative_slope
+    return leaky_relu_layer.get_output(0)
+
+
 @converter_registry.register("pd_op.selu", trt_version="8.x")
 def selu_converter(network, paddle_op, inputs):
     x = inputs[0]
@@ -285,3 +297,44 @@ def selu_converter(network, paddle_op, inputs):
     selu_layer.alpha = alpha
     selu_layer.beta = scale
     return selu_layer.get_output(0)
+
+
+@converter_registry.register("pd_op.prelu", trt_version="8.x")
+def prelu_converter(network, paddle_op, inputs):
+    input, alpha_data = inputs
+    input_dims = input.shape
+    mode = paddle_op.attrs()["mode"]
+    data_format = paddle_op.attrs().get("data_format", "NCHW")
+    w_dims = trt.Dims(alpha_data.numpy().shape)
+    trt_w_dims = w_dims
+    alpha_tensor = network.add_constant(trt_w_dims, alpha_data).get_output(0)
+    alpha_dims = alpha_tensor.shape
+    real_alpha_tensor = alpha_tensor
+    if len(alpha_dims) != len(input_dims):
+        reshape_layer = network.add_shuffle(alpha_tensor)
+        c = alpha_dims[0]
+        n_tensor = add_1D_constant_layer(network, [1])
+        c_tensor = add_1D_constant_layer(network, [c])
+        hw_tensor = None
+        if len(input_dims) - 2 > 0:
+            hw_tensor = add_1D_constant_layer(
+                network, [1] * (len(input_dims) - 2)
+            )
+        if data_format == "NCHW":
+            if hw_tensor:
+                shape_tensor = trt_concat(
+                    network, [n_tensor, c_tensor, hw_tensor]
+                )
+            else:
+                shape_tensor = trt_concat(network, [n_tensor, c_tensor])
+        else:
+            if hw_tensor:
+                shape_tensor = trt_concat(
+                    network, [n_tensor, hw_tensor, c_tensor]
+                )
+            else:
+                shape_tensor = trt_concat(network, [n_tensor, c_tensor])
+        reshape_layer.set_input(1, shape_tensor)
+        real_alpha_tensor = reshape_layer.get_output(0)
+    layer = network.add_parametric_relu(input, real_alpha_tensor)
+    return layer.get_output(0)
diff --git a/python/paddle/tensorrt/util.py b/python/paddle/tensorrt/util.py
@@ -299,6 +299,7 @@ def weight_to_tensor(network, paddle_value, trt_tensor, use_op_name):
         "pd_op.depthwise_conv2d_transpose",
         "pd_op.fused_conv2d_add_act",
         "pd_op.affine_channel",
+        "pd_op.prelu",
         "pd_op.fused_bias_dropout_residual_layer_norm",
         "pd_op.deformable_conv",
     ]
diff --git a/test/tensorrt/test_converter_activation.py b/test/tensorrt/test_converter_activation.py
@@ -369,9 +369,169 @@ def setUp(self):
     def test_trt_result(self):
         self.check_trt_result()
 
+
+class TestLeakyReluCas1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.nn.functional.leaky_relu
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+            "negative_slope": 0.5,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.opt_shape = {"x": [2, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestLeakyReluCase2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.nn.functional.leaky_relu
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+            "negative_slope": -0.5,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.opt_shape = {"x": [2, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestLeakyRelu_Cas1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.nn.functional.leaky_relu_
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+            "negative_slope": 0.5,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.opt_shape = {"x": [2, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestLeakyRelu_Case2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.nn.functional.leaky_relu_
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+            "negative_slope": -0.5,
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.opt_shape = {"x": [2, 3]}
+        self.max_shape = {"x": [5, 3]}
+
     def test_trt_result_fp16(self):
         self.check_trt_result(precision_mode="fp16")
 
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+def prelu_wrapper(x, alpha_shape, data_format='NCHW'):
+    alpha = paddle.create_parameter(
+        shape=alpha_shape, dtype='float32', name="alpha"
+    )
+    return paddle.nn.functional.prelu(x, alpha, data_format)
+
+
+class TestPReluCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = prelu_wrapper
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+            "alpha_shape": [3],
+            "data_format": "NCHW",
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.opt_shape = {"x": [2, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestPReluCase2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = prelu_wrapper
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+            "alpha_shape": [3],
+            "data_format": "NHWC",
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3]}
+        self.opt_shape = {"x": [2, 3]}
+        self.max_shape = {"x": [5, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestPReluCase3TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = prelu_wrapper
+        self.api_args = {
+            "x": np.random.randn(2, 3, 3).astype("float32"),
+            "alpha_shape": [3],
+            "data_format": "NCHW",
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 3]}
+        self.opt_shape = {"x": [2, 3, 3]}
+        self.max_shape = {"x": [5, 3, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
+
+class TestPReluCase4TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = prelu_wrapper
+        self.api_args = {
+            "x": np.random.randn(2, 3, 3).astype("float32"),
+            "alpha_shape": [3],
+            "data_format": "NHWC",
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 3]}
+        self.opt_shape = {"x": [2, 3, 3]}
+        self.max_shape = {"x": [5, 3, 3]}
+
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
+        self.check_trt_result()
+
 
 if __name__ == '__main__':
     unittest.main()

Original file line number	Diff line number	Diff line change
`@@ -299,6 +299,7 @@ def weight_to_tensor(network, paddle_value, trt_tensor, use_op_name):`
`299`	`299`	`"pd_op.depthwise_conv2d_transpose",`
`300`	`300`	`"pd_op.fused_conv2d_add_act",`
`301`	`301`	`"pd_op.affine_channel",`
	`302`	`+ "pd_op.prelu",`
`302`	`303`	`"pd_op.fused_bias_dropout_residual_layer_norm",`
`303`	`304`	`"pd_op.deformable_conv",`
`304`	`305`	`]`