PaddlePaddle
diff --git a/‎test/ir/inference/test_trt_convert_expand_as_v2.py
Copy file name to clipboard
+86-171Lines changed: 86 additions & 171 deletions b/‎test/ir/inference/test_trt_convert_expand_as_v2.py
Copy file name to clipboard
+86-171Lines changed: 86 additions & 171 deletions
@@ -16,7 +16,6 @@
 
 import unittest
 from functools import partial
-from typing import Any
 
 import numpy as np
 from program_config import ProgramConfig, TensorConfig
@@ -25,146 +24,79 @@
 import paddle.inference as paddle_infer
 
 
-class TrtConvertExpandASV2Test(TrtLayerAutoScanTest):
+class TrtConvertExpandAsV2Test(TrtLayerAutoScanTest):
     def is_program_valid(self, program_config: ProgramConfig) -> bool:
         attrs = [
             program_config.ops[i].attrs for i in range(len(program_config.ops))
         ]
-        if len(attrs[0]['target_shape']) < self.dims:
-            return False
-        if self.dims == 1:
-            if len(attrs[0]['target_shape']) == 4:
-                return False
         return True
 
     def sample_program_configs(self):
-        def generate_input1(attrs: list[dict[str, Any]]):
-            if self.dims == 4:
-                self.input_shape = [1, 8, 1, 32]
-                return np.random.random([1, 8, 1, 32]).astype(np.float32)
-            elif self.dims == 3:
-                self.input_shape = [1, 32, 32]
-                return np.random.random([1, 32, 32]).astype(np.float32)
-            elif self.dims == 2:
-                self.input_shape = [1, 32]
-                return np.random.random([1, 32]).astype(np.float32)
-            elif self.dims == 1:
-                self.input_shape = [32]
-                return np.random.random([32]).astype(np.float32)
-            elif self.dims == 0:
-                self.input_shape = []
-                return np.random.random([]).astype(np.float32)
-
-        for dims in [0, 1, 2, 3, 4]:
-            for shape in [
-                [10, 8, 32, 32],
-                [2, 8, 32, 32],
-                [8, 32, 32],
-                [2, 32],
-                [32],
-            ]:
-                dics = [
-                    {
-                        "target_shape": shape,
-                    },
-                ]
-                self.dims = dims
-
-                ops_config = [
-                    {
-                        "op_type": "expand_as_v2",
-                        "op_inputs": {"X": ["expand_v2_input"]},
-                        "op_outputs": {"Out": ["expand_v2_out"]},
-                        "op_attrs": dics[0],
-                    }
-                ]
-                ops = self.generate_op_config(ops_config)
-                program_config = ProgramConfig(
-                    ops=ops,
-                    weights={},
-                    inputs={
-                        "expand_v2_input": TensorConfig(
-                            data_gen=partial(generate_input1, dics)
-                        )
-                    },
-                    outputs=["expand_v2_out"],
-                )
-
-                yield program_config
+        def generate_input1():
+            return np.random.random([1, 4, 1]).astype(np.float32)
+
+        def generate_input2():
+            return np.random.random([1, 4, 10]).astype(np.float32)
+
+        ops_config = [
+            {
+                "op_type": "expand_as_v2",
+                "op_inputs": {
+                    "X": ["expand_as_v2_input1"],
+                    "Y": ["expand_as_v2_input2"],
+                },
+                "op_outputs": {"Out": ["expand_as_v2_out"]},
+                "op_attrs": {"target_shape": [1, 4, 10]},
+            },
+        ]
+        ops = self.generate_op_config(ops_config)
+        program_config = ProgramConfig(
+            ops=ops,
+            weights={},
+            inputs={
+                "expand_as_v2_input1": TensorConfig(
+                    data_gen=partial(generate_input1)
+                ),
+                "expand_as_v2_input2": TensorConfig(
+                    data_gen=partial(generate_input2)
+                ),
+            },
+            outputs=["expand_as_v2_out"],
+        )
+
+        yield program_config
 
     def sample_predictor_configs(
         self, program_config
     ) -> tuple[paddle_infer.Config, list[int], float]:
-        def generate_dynamic_shape(attrs):
-            if self.dims == 4:
-                self.dynamic_shape.min_input_shape = {
-                    "expand_v2_input": [1, 8, 1, 32]
-                }
-                self.dynamic_shape.max_input_shape = {
-                    "expand_v2_input": [10, 8, 1, 32]
-                }
-                self.dynamic_shape.opt_input_shape = {
-                    "expand_v2_input": [1, 8, 1, 32]
-                }
-            elif self.dims == 3:
-                self.dynamic_shape.min_input_shape = {
-                    "expand_v2_input": [1, 32, 32]
-                }
-                self.dynamic_shape.max_input_shape = {
-                    "expand_v2_input": [8, 32, 32]
-                }
-                self.dynamic_shape.opt_input_shape = {
-                    "expand_v2_input": [1, 32, 32]
-                }
-            elif self.dims == 2:
-                self.dynamic_shape.min_input_shape = {
-                    "expand_v2_input": [1, 32]
-                }
-                self.dynamic_shape.max_input_shape = {
-                    "expand_v2_input": [4, 32]
-                }
-                self.dynamic_shape.opt_input_shape = {
-                    "expand_v2_input": [1, 32]
-                }
-            elif self.dims == 1:
-                self.dynamic_shape.min_input_shape = {"expand_v2_input": [32]}
-                self.dynamic_shape.max_input_shape = {"expand_v2_input": [64]}
-                self.dynamic_shape.opt_input_shape = {"expand_v2_input": [32]}
-            elif self.dims == 0:
-                self.dynamic_shape.min_input_shape = {"expand_v2_input": []}
-                self.dynamic_shape.max_input_shape = {"expand_v2_input": []}
-                self.dynamic_shape.opt_input_shape = {"expand_v2_input": []}
+        def generate_dynamic_shape():
+            self.dynamic_shape.min_input_shape = {
+                "expand_as_v2_input1": [1, 4, 1],
+                "expand_as_v2_input2": [1, 4, 10],
+            }
+            self.dynamic_shape.max_input_shape = {
+                "expand_as_v2_input1": [1, 4, 1],
+                "expand_as_v2_input2": [1, 4, 10],
+            }
+            self.dynamic_shape.opt_input_shape = {
+                "expand_as_v2_input1": [1, 4, 1],
+                "expand_as_v2_input2": [1, 4, 10],
+            }
 
         def clear_dynamic_shape():
             self.dynamic_shape.min_input_shape = {}
             self.dynamic_shape.max_input_shape = {}
             self.dynamic_shape.opt_input_shape = {}
 
-        def generate_trt_nodes_num(attrs, dynamic_shape):
-            ver = paddle_infer.get_trt_compile_version()
-            ver_num = ver[0] * 1000 + ver[1] * 100 + ver[2] * 10
-            if dynamic_shape and (ver_num > 8000 or self.dims > 0):
-                return 1, 2
-            else:
-                return 0, 3
-
-        attrs = [
-            program_config.ops[i].attrs for i in range(len(program_config.ops))
-        ]
-
         clear_dynamic_shape()
         # for dynamic_shape
-        generate_dynamic_shape(attrs)
+        generate_dynamic_shape()
         self.trt_param.precision = paddle_infer.PrecisionType.Float32
         program_config.set_input_type(np.float32)
-        yield self.create_inference_config(), generate_trt_nodes_num(
-            attrs, True
-        ), 1e-5
+        yield self.create_inference_config(), (1, 3), 1e-5
         self.trt_param.precision = paddle_infer.PrecisionType.Half
         program_config.set_input_type(np.float16)
-        yield self.create_inference_config(), generate_trt_nodes_num(
-            attrs, True
-        ), 1e-3
+        yield self.create_inference_config(), (1, 3), 1e-3
 
     def add_skip_trt_case(self):
         pass
@@ -174,70 +106,54 @@ def test(self):
         self.run_test()
 
 
-class TrtConvertExpandV2Test2(TrtLayerAutoScanTest):
+class TrtConvertExpandAsV2Test2(TrtLayerAutoScanTest):
     def is_program_valid(self, program_config: ProgramConfig) -> bool:
         attrs = [
             program_config.ops[i].attrs for i in range(len(program_config.ops))
         ]
         return True
 
     def sample_program_configs(self):
-        def generate_input1(attrs: list[dict[str, Any]]):
-            if self.dims == 1:
-                self.input_shape = [1]
-                return np.random.random([1]).astype(np.float32)
-
-        for dims in [1]:
-            for shape in [[10]]:
-                dics = [
-                    {
-                        "target_shape": shape,
-                    },
-                ]
-                self.dims = dims
-                dics_input = [
-                    {"X": ["expand_v2_input"], "Y": ["shapeT1_data"]},
-                ]
-                ops_config = [
-                    {
-                        "op_type": "fill_constant",
-                        "op_inputs": {},
-                        "op_outputs": {"Out": ["shapeT1_data"]},
-                        "op_attrs": {
-                            "dtype": 2,
-                            "str_value": "10",
-                            "shape": [1],
-                        },
-                    },
-                    {
-                        "op_type": "expand_as_v2",
-                        "op_inputs": dics_input[0],
-                        "op_outputs": {"Out": ["expand_v2_out"]},
-                        "op_attrs": dics[0],
-                    },
-                ]
-                ops = self.generate_op_config(ops_config)
-                program_config = ProgramConfig(
-                    ops=ops,
-                    weights={},
-                    inputs={
-                        "expand_v2_input": TensorConfig(
-                            data_gen=partial(generate_input1, dics)
-                        )
-                    },
-                    outputs=["expand_v2_out"],
-                )
-
-                yield program_config
+        def generate_input():
+            return np.random.random([1, 4, 1]).astype(np.float32)
+
+        ops_config = [
+            {
+                "op_type": "expand_as_v2",
+                "op_inputs": {
+                    "X": ["expand_as_v2_input1"],
+                },
+                "op_outputs": {"Out": ["expand_as_v2_out"]},
+                "op_attrs": {"target_shape": [1, 4, 10]},
+            },
+        ]
+        ops = self.generate_op_config(ops_config)
+        program_config = ProgramConfig(
+            ops=ops,
+            weights={},
+            inputs={
+                "expand_as_v2_input1": TensorConfig(
+                    data_gen=partial(generate_input)
+                ),
+            },
+            outputs=["expand_as_v2_out"],
+        )
+
+        yield program_config
 
     def sample_predictor_configs(
         self, program_config
     ) -> tuple[paddle_infer.Config, list[int], float]:
         def generate_dynamic_shape():
-            if self.dims == 1:
-                self.dynamic_shape.min_input_shape = {"expand_v2_input": [1]}
-                self.dynamic_shape.max_input_shape = {"expand_v2_input": [1]}
-                self.dynamic_shape.opt_input_shape = {"expand_v2_input": [1]}
+            self.dynamic_shape.min_input_shape = {
+                "expand_as_v2_input1": [1, 4, 1],
+            }
+            self.dynamic_shape.max_input_shape = {
+                "expand_as_v2_input1": [1, 4, 1],
+            }
+            self.dynamic_shape.opt_input_shape = {
+                "expand_as_v2_input1": [1, 4, 1],
+            }
 
         def clear_dynamic_shape():
             self.dynamic_shape.min_input_shape = {}
@@ -249,7 +165,6 @@ def clear_dynamic_shape():
         generate_dynamic_shape()
         self.trt_param.precision = paddle_infer.PrecisionType.Float32
         program_config.set_input_type(np.float32)
-        # fill_constant will be folded by constant folding pass!
         yield self.create_inference_config(), (1, 2), 1e-5
         self.trt_param.precision = paddle_infer.PrecisionType.Half
         program_config.set_input_type(np.float16)