[Paddle TensorRT] add pd_op.shape converter (#68395)

lizexu123 · web-flow · commit cfc6faf076bd · 2024-09-26T16:06:54.000+08:00
* shape converter

* 适配多输入

* ci bug

* 命名

* fix
diff --git a/python/paddle/tensorrt/converter.py b/python/paddle/tensorrt/converter.py
@@ -30,6 +30,7 @@
 from paddle.base.log_helper import get_logger
 
 from .impls.activation import *  # noqa: F403
+from .impls.attribute import *  # noqa: F403
 from .impls.conv import *  # noqa: F403
 from .impls.creation import *  # noqa: F403
 from .impls.linalg import *  # noqa: F403
diff --git a/python/paddle/tensorrt/export.py b/python/paddle/tensorrt/export.py
@@ -184,7 +184,7 @@ def __init__(
             >>> trt_config = TensorRTConfig(inputs=[input])
             >>> trt_config.disable_ops = "pd_op.dropout"
         """
-        self.inputs = (inputs,)
+        self.inputs = inputs
         self.min_subgraph_size = min_subgraph_size
         self.save_model_dir = save_model_dir
         self.disable_ops = disable_ops
@@ -207,18 +207,22 @@ def convert_to_trt(program, trt_config, scope):
             feed_name.append(param_name)
 
     with paddle.pir_utils.IrGuard():
+        min_shape_feed = {}
+        max_shape_feed = {}
         for i, input_instance in enumerate(trt_config.inputs):
             # get fake inputs
-            min_data, _, max_data = input_instance[i].generate_input_data()
+            min_data, _, max_data = input_instance.generate_input_data()
             program_with_output = program.list_vars()[-1]
+            min_shape_feed[feed_name[i]] = min_data
+            max_shape_feed[feed_name[i]] = max_data
 
             # run warmup for collecting shape
-            program = warmup_shape_infer(
-                program,
-                min_shape_feed={feed_name[0]: min_data},
-                max_shape_feed={feed_name[0]: max_data},
-                scope=scope,
-            )
+        program = warmup_shape_infer(
+            program,
+            min_shape_feed=min_shape_feed,
+            max_shape_feed=max_shape_feed,
+            scope=scope,
+        )
 
         # run pir pass (including trt_op_marker_pass)
         program_with_pir = run_pir_pass(program, partition_mode=False)
@@ -314,7 +318,7 @@ def convert(function=None, input_spec=None, config=None, **kwargs):
         ...         )
         ...         trt_config = TensorRTConfig(inputs=[input_config])
         ...         for i, input_instrance in enumerate(trt_config.inputs):
-        ...             min_data, _, max_data = input_instrance[i].generate_input_data()
+        ...             min_data, _, max_data = input_instrance.generate_input_data()
         ...             paddle.disable_static()
         ...             x = paddle.to_tensor(min_data)
         ...             net = CumsumModel(input_dim=min_data.shape[-1])
@@ -608,7 +612,7 @@ def convert_loaded_model(model_dir, config):
 
             ...    paddle.disable_static()
             ...    for i, input_instrance in enumerate(trt_config.inputs):
-            ...        min_data, _, max_data = input_instrance[i].generate_input_data()
+            ...        min_data, _, max_data = input_instrance.generate_input_data()
             ...        model_inputs = paddle.to_tensor(min_data)
             ...        output_converted = predictor.run([model_inputs])
 
diff --git a/python/paddle/tensorrt/impls/attribute.py b/python/paddle/tensorrt/impls/attribute.py
@@ -0,0 +1,22 @@
+# 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.
+
+from paddle.tensorrt.register import converter_registry
+
+
+@converter_registry.register("pd_op.shape", trt_version="8.x")
+def shape_converter(network, paddle_op, inputs):
+    input_tensor = inputs[0]
+    shape_layer = network.add_shape(input_tensor)
+    return shape_layer.get_output(0)
diff --git a/test/tensorrt/CMakeLists.txt b/test/tensorrt/CMakeLists.txt
@@ -21,4 +21,5 @@ if(NOT WIN32 AND TENSORRT_FOUND)
   set_tests_properties(test_converter_others PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_manipulation PROPERTIES TIMEOUT "100")
   set_tests_properties(test_converter_creation PROPERTIES TIMEOUT "100")
+  set_tests_properties(test_converter_attribute PROPERTIES TIMEOUT "100")
 endif()
diff --git a/test/tensorrt/test_converter_attribute.py b/test/tensorrt/test_converter_attribute.py
@@ -0,0 +1,38 @@
+# 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
+
+
+class TestShapeTRTPattern(TensorRTBaseTest):
+    def setUp(self):
+        self.python_api = paddle.shape
+        self.api_args = {
+            "x": np.random.randn(2, 3).astype(np.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()
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/test/tensorrt/test_export.py b/test/tensorrt/test_export.py
@@ -20,7 +20,7 @@
 import paddle
 import paddle.inference as paddle_infer
 import paddle.nn.functional as F
-from paddle import nn
+from paddle import Tensor, nn
 from paddle.static import InputSpec
 from paddle.tensorrt.export import (
     Input,
@@ -33,6 +33,59 @@
 )
 
 
+class LeNetMultiInput(nn.Layer):
+    """LeNet model modified to accept two inputs."""
+
+    def __init__(self, num_classes: int = 10) -> None:
+        super().__init__()
+        self.num_classes = num_classes
+
+        # Convolution layers for the first input
+        self.features1 = nn.Sequential(
+            nn.Conv2D(1, 6, 3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2D(2, 2),
+            nn.Conv2D(6, 16, 5, stride=1, padding=0),
+            nn.ReLU(),
+            nn.MaxPool2D(2, 2),
+        )
+
+        # Convolution layers for the second input
+        self.features2 = nn.Sequential(
+            nn.Conv2D(1, 6, 3, stride=1, padding=1),
+            nn.ReLU(),
+            nn.MaxPool2D(2, 2),
+            nn.Conv2D(6, 16, 5, stride=1, padding=0),
+            nn.ReLU(),
+            nn.MaxPool2D(2, 2),
+        )
+
+        # Fully connected layers
+        if num_classes > 0:
+            self.fc = nn.Sequential(
+                nn.Linear(400 * 2, 120),  # Adjusted for two inputs
+                nn.Linear(120, 84),
+                nn.Linear(84, num_classes),
+            )
+
+    def forward(self, input1: Tensor, input2: Tensor) -> Tensor:
+        # Apply feature extraction on both inputs
+        x1 = self.features1(input1)
+        x2 = self.features2(input2)
+
+        # Flatten both feature maps
+        x1 = paddle.flatten(x1, 1)
+        x2 = paddle.flatten(x2, 1)
+
+        # Concatenate the features from both inputs
+        x = paddle.concat([x1, x2], axis=1)
+
+        if self.num_classes > 0:
+            x = self.fc(x)
+
+        return x
+
+
 class CumsumModel(nn.Layer):
     def __init__(self, input_dim):
         super().__init__()
@@ -121,7 +174,7 @@ def test_paddle_to_tensorrt_conversion_cumsum(self):
 
         paddle.disable_static()
         for i, input_instrance in enumerate(trt_config.inputs):
-            min_data, _, max_data = input_instrance[i].generate_input_data()
+            min_data, _, max_data = input_instrance.generate_input_data()
             model_inputs = paddle.to_tensor(min_data)
             output_converted = predictor.run([model_inputs])
 
@@ -136,7 +189,7 @@ def test_run(self):
             )
             trt_config = TensorRTConfig(inputs=[input_config])
             for i, input_instrance in enumerate(trt_config.inputs):
-                min_data, _, max_data = input_instrance[i].generate_input_data()
+                min_data, _, max_data = input_instrance.generate_input_data()
                 paddle.disable_static()
                 x = paddle.to_tensor(min_data)
                 net = CumsumModel(input_dim=min_data.shape[-1])
@@ -169,5 +222,69 @@ def test_run(self):
                 )
 
 
+class TestConvertMultipleInputs(unittest.TestCase):
+    def test_run(self):
+        with paddle.pir_utils.IrGuard():
+            input_config = Input(
+                min_input_shape=(1, 1, 28, 28),
+                optim_input_shape=(1, 1, 28, 28),
+                max_input_shape=(1, 1, 28, 28),
+            )
+            input_config2 = Input(
+                min_input_shape=(1, 1, 28, 28),
+                optim_input_shape=(1, 1, 28, 28),
+                max_input_shape=(1, 1, 28, 28),
+            )
+            trt_config = TensorRTConfig(inputs=[input_config, input_config2])
+
+            min_data_list = []
+            max_data_list = []
+            for i, input_instrance in enumerate(trt_config.inputs):
+                min_data, _, max_data = input_instrance.generate_input_data()
+
+                min_data_list.append(min_data)
+                max_data_list.append(max_data)
+                paddle.disable_static()
+
+            x = [paddle.to_tensor(md) for md in min_data_list]
+            net = LeNetMultiInput()
+            out = net(*x)
+
+            input_spec = [
+                InputSpec(
+                    shape=min_data_list[0].shape, dtype='float32', name='input1'
+                ),
+                InputSpec(
+                    shape=min_data_list[1].shape, dtype='float32', name='input2'
+                ),
+            ]
+
+            program_with_trt, scope = convert(
+                net,
+                input_spec=input_spec,
+                config=trt_config,
+                full_graph=True,
+            )
+            output_var = program_with_trt.list_vars()[-1]
+
+            output_converted = predict_program(
+                program_with_trt,
+                {"input1": min_data_list[0], "input2": min_data_list[1]},
+                [output_var],
+                scope=scope,
+            )
+            output_expected = out.numpy()
+            output_converted_np = output_converted[0]
+
+            # Check that the results are close to each other within a tolerance of 1e-2
+            np.testing.assert_allclose(
+                output_expected,
+                output_converted_np,
+                rtol=1e-2,
+                atol=1e-2,
+                err_msg="Outputs are not within the 1e-2 tolerance",
+            )
+
+
 if __name__ == "__main__":
     unittest.main()
