【SCU】【Paddle TensorRT No.56】Add pd_op.tanh_shrink converter (#69693)

PolaKuma · web-flow · commit 0f7c06184777 · 2025-01-07T11:21:16.000+08:00
* add_tanhshrink

* fix

* fix codestyle

* fix codestyle

* fix codestyle

* Update test_converter_activation.py

* fix codestyle

* add_fp16

* fix

* Update test_converter_activation.py

* Update trt_op_marker_pass.cc

* Update activation.py

* Update activation.py

* update

* 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
@@ -265,6 +265,7 @@ class ActOpPattern : public pir::OpRewritePattern<OpType> {
 };
 using TanhOpPattern = ActOpPattern<paddle::dialect::TanhOp>;
 using CeluOpPattern = ActOpPattern<paddle::dialect::CeluOp>;
+using TanhShrinkOpPattern = ActOpPattern<paddle::dialect::TanhShrinkOp>;
 using LogicalNotOpPattern = ActOpPattern<paddle::dialect::LogicalNotOp>;
 using LogicalNot_OpPattern = ActOpPattern<paddle::dialect::LogicalNot_Op>;
 
@@ -2302,6 +2303,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<ClipPattern>(context));
     ps.Add(std::make_unique<GridSampleOpPattern>(context));
     ps.Add(std::make_unique<StackOpPattern>(context));
+    ps.Add(std::make_unique<TanhShrinkOpPattern>(context));
     ps.Add(std::make_unique<WherePattern>(context));
     ps.Add(std::make_unique<FullLikeOpPattern>(context));
     ps.Add(std::make_unique<FullWithTensorPattern>(context));
diff --git a/python/paddle/tensorrt/impls/activation.py b/python/paddle/tensorrt/impls/activation.py
@@ -188,6 +188,16 @@ def swish_silu_converter(network, paddle_op, inputs):
     return trt_prod(network, inputs[0], layer_output)
 
 
+@converter_registry.register("pd_op.tanh_shrink", trt_version="8.x")
+def tanh_shrink_converter(network, paddle_op, inputs):
+    x = inputs[0]
+    tanh_layer = network.add_activation(x, trt.ActivationType.TANH)
+    subtract_layer = network.add_elementwise(
+        x, tanh_layer.get_output(0), trt.ElementWiseOperation.SUB
+    )
+    return subtract_layer.get_output(0)
+
+
 @converter_registry.register("pd_op.stanh", trt_version="8.x")
 def stanh_converter(network, paddle_op, inputs):
     x = inputs[0]
diff --git a/test/tensorrt/test_converter_activation.py b/test/tensorrt/test_converter_activation.py
@@ -204,6 +204,23 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestTanhShrinkOpFloatTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle._C_ops.tanh_shrink
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 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 TestStanhFloatTRTPattern(TensorRTBaseTest):
     def setUp(self):
         self.python_api = paddle.stanh
