【SCU】[Paddle TensorRT No.63] Add pd_op.mish converter

paddle/fluid/pir/transforms/tensorrt/trt_op_marker_pass.cc

@@ -258,6 +258,7 @@ class ActOpPattern : public pir::OpRewritePattern<OpType> {
 };
 using TanhOpPattern = ActOpPattern<paddle::dialect::TanhOp>;
 using CeluOpPattern = ActOpPattern<paddle::dialect::CeluOp>;
+using MishOpPattern = ActOpPattern<paddle::dialect::MishOp>;
 
 class Pool2dOpPattern
     : public pir::OpRewritePattern<paddle::dialect::Pool2dOp> {
@@ -2204,6 +2205,7 @@ class TrtOpMarkerPass : public pir::PatternRewritePass {
     ps.Add(std::make_unique<NotEqualOpPattern>(context));
     ps.Add(std::make_unique<TanhOpPattern>(context));
     ps.Add(std::make_unique<CeluOpPattern>(context));
+    ps.Add(std::make_unique<MishOpPattern>(context));
     ps.Add(std::make_unique<OneHotOpPattern>(context));
     ps.Add(std::make_unique<AssignValueOpPattern>(context));
     ps.Add(std::make_unique<AssignValue_OpPattern>(context));

python/paddle/tensorrt/impls/activation.py

@@ -135,6 +135,20 @@ def swish_silu_converter(network, paddle_op, inputs):
     return trt_prod(network, inputs[0], layer_output)
 
 
+@converter_registry.register("pd_op.mish", trt_version="8.x")
+def mish_converter(network, paddle_op, inputs):
+    x = inputs[0]
+    softplus_layer = network.add_activation(x, trt.ActivationType.SOFTPLUS)
+    softplus_output = softplus_layer.get_output(0)
+
+    tanh_layer = network.add_activation(
+        softplus_output, trt.ActivationType.TANH
+    )
+    tanh_output = tanh_layer.get_output(0)
+
+    return trt_prod(network, x, tanh_output)
+
+
 @converter_registry.register("pd_op.celu", trt_version="8.x")
 def celu_converter(network, paddle_op, inputs):
     input_tensor = inputs[0]

test/tensorrt/test_converter_activation.py

@@ -158,5 +158,61 @@ def test_trt_result(self):
         self.check_trt_result()
 
 
+class TestMishCase1TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.nn.functional.mish
+        self.api_args = {
+            "x": np.random.randn(2).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1]}
+        self.max_shape = {"x": [5]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestMishCase2TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.nn.functional.mish
+        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(self):
+        self.check_trt_result()
+
+
+class TestMishCase3TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.nn.functional.mish
+        self.api_args = {
+            "x": np.random.randn(2, 3, 4).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 4]}
+        self.max_shape = {"x": [5, 3, 4]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
+class TestMishCase4TRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.nn.functional.mish
+        self.api_args = {
+            "x": np.random.randn(2, 3, 4, 2).astype("float32"),
+        }
+        self.program_config = {"feed_list": ["x"]}
+        self.min_shape = {"x": [1, 3, 4, 2]}
+        self.max_shape = {"x": [5, 3, 4, 2]}
+
+    def test_trt_result(self):
+        self.check_trt_result()
+
+
 if __name__ == '__main__':
     unittest.main()