【Paddle TensorRT】Modified the serialization save path for TensorRT and added an attribute name to the Input class (#71772) (#71767)

* fix

* merge

Author: lizexu123

python/paddle/tensorrt/converter.py

@@ -496,7 +496,7 @@ def convert_subgraph_to_trt(self, program, group_op):
             int(hashlib.sha256(group_str.encode('utf-8')).hexdigest(), 16)
             % 10**8
         )
-        CACHE_ROOT = get_cache_path()
+        CACHE_ROOT = get_cache_path(self.trt_config.save_model_dir)
         CACHE_FILE = f"{CACHE_ROOT}/engine_{engine_name}_{self.engine_num}.trt"
         with open(CACHE_FILE, "wb") as f:
             f.write(trt_engine)

python/paddle/tensorrt/export.py

@@ -51,56 +51,6 @@
 
 
 class Input:
-    """
-    A class used to configure input data for models. This class serves two purposes:
-
-    1. Random Data Generation: When no input data is supplied, it automatically generates random input data based on the specified minimum, optimal, and maximum shapes. In this mode,you can configure the data type (e.g., 'float32', 'int64', etc.) and the range of values (e.g.,(0.0, 1.0) for floats or (1, 10) for integers).
-
-    2. User-Provided Input: Alternatively, you can supply your own input data via the `warmup_data` argument. In this case, the provided data will be used directly, and the`input_data_type` and `input_range` settings will be ignored.
-
-    Args:
-        warmup_data (tuple):
-            The tuple of actual input data (for the automatic shape collection mechanism).
-        min_input_shape (tuple):
-            The shape of the minimum input tensor.
-        max_input_shape (tuple):
-            The shape of the maximum input tensor.
-        optim_input_shape (tuple, optional):
-            The shape of the optimal input tensor (default is None).
-        input_data_type (str, optional):
-            The data type for the input tensors, such as 'float32' or 'int64' or 'float32' or 'int32'  (default is float32).
-            This option only applies when min_input_shape, optim_input_shape, and max_input_shape are provided; it does not apply to warmup_data.
-        input_range (tuple, optional):
-            The range of values used to generate input data. For floats, the default range is (0.0, 1.0). For integers, the default range is (1, 10).
-            This option only applies when min_input_shape, optim_input_shape, and max_input_shape are provided; it does not apply to warmup_data.
-    Returns:
-        None
-
-    Examples:
-        .. code-block:: python
-
-            >>> # example 1:
-            >>> from paddle.tensorrt.export import Input
-            >>> input_config = Input(
-            >>>     min_input_shape=(1,100),
-            >>>     optim_input_shape=(4,100),
-            >>>     max_input_shape=(8,100),
-            >>> )
-            >>> input_config.input_data_type='int64'
-            >>> input_config.input_range=(1,10)
-
-            >>> # example 2:
-            >>> from paddle.tensorrt.export import Input
-            >>> import numpy as np
-            >>> input_config = Input(
-            >>>     warmup_data=(
-            >>>         np.random.rand(1,100).astype(np.float32),
-            >>>         np.random.rand(4,100).astype(np.float32),
-            >>>         np.random.rand(8,100).astype(np.float32),
-            >>>     )
-            >>> )
-    """
-
     def __init__(
         self,
         warmup_data: tuple[np.ndarray, ...] | None = None,
@@ -109,7 +59,59 @@ def __init__(
         optim_input_shape: tuple | None = None,
         input_data_type: str | None = 'float32',
         input_range: tuple | None = None,
+        name: str | None = None,
     ) -> None:
+        """
+        A class used to configure input data for models. This class serves two purposes:
+
+        1. Random Data Generation: When no input data is supplied, it automatically generates random input data based on the specified minimum, optimal, and maximum shapes. In this mode,you can configure the data type (e.g., 'float32', 'int64', etc.) and the range of values (e.g.,(0.0, 1.0) for floats or (1, 10) for integers).
+
+        2. User-Provided Input: Alternatively, you can supply your own input data via the `warmup_data` argument. In this case, the provided data will be used directly, and the`input_data_type` and `input_range` settings will be ignored.
+
+        Args:
+            warmup_data (tuple):
+                The tuple of actual input data (for the automatic shape collection mechanism).
+            min_input_shape (tuple):
+                The shape of the minimum input tensor.
+            max_input_shape (tuple):
+                The shape of the maximum input tensor.
+            optim_input_shape (tuple):
+                The shape of the optimal input tensor.
+            input_data_type (str, optional):
+                The data type for the input tensors, such as 'float32' or 'int64' or 'float32' or 'int32'  (default is float32).
+                This option only applies when min_input_shape, optim_input_shape, and max_input_shape are provided; it does not apply to warmup_data.
+            input_range (tuple, optional):
+                The range of values used to generate input data. For floats, the default range is (0.0, 1.0). For integers, the default range is (1, 10).
+                This option only applies when min_input_shape, optim_input_shape, and max_input_shape are provided; it does not apply to warmup_data.
+            name:(str,optional):
+                The name of the input to the model.
+        Returns:
+            None
+
+        Examples:
+            .. code-block:: python
+
+                >>> # example 1:
+                >>> from paddle.tensorrt.export import Input
+                >>> input_config = Input(
+                >>>     min_input_shape=(1,100),
+                >>>     optim_input_shape=(4,100),
+                >>>     max_input_shape=(8,100),
+                >>> )
+                >>> input_config.input_data_type='int64'
+                >>> input_config.input_range=(1,10)
+
+                >>> # example 2:
+                >>> from paddle.tensorrt.export import Input
+                >>> import numpy as np
+                >>> input_config = Input(
+                >>>     warmup_data=(
+                >>>         np.random.rand(1,100).astype(np.float32),
+                >>>         np.random.rand(4,100).astype(np.float32),
+                >>>         np.random.rand(8,100).astype(np.float32),
+                >>>     )
+                >>> )
+        """
         if warmup_data is not None:
             if min_input_shape or max_input_shape or optim_input_shape:
                 raise ValueError(
@@ -132,6 +134,7 @@ def __init__(
         self.optim_input_shape = optim_input_shape
         self.input_data_type = input_data_type
         self.input_range = input_range
+        self.name = name
 
     def generate_input_data(self):
         """
@@ -331,17 +334,19 @@ def convert_to_trt(program, trt_config, scope):
             assert len({len(t) for t in input_tuples}) == 1
             num_samples = len(input_tuples[0])
             for sample_idx in range(num_samples):
-                feed_dict = {
-                    name: input_tuples[i][sample_idx]
-                    for i, name in enumerate(feed_name)
-                }
+                feed_dict = {}
+                for i, inp in enumerate(trt_config.inputs):
+                    name = inp.name if inp.name is not None else feed_name[i]
+                    feed_dict[name] = input_tuples[i][sample_idx]
                 feeds.append(feed_dict)
         else:
             input_tuples = [i.generate_input_data() for i in trt_config.inputs]
-            feeds = [
-                {name: t[i] for t, name in zip(input_tuples, feed_name)}
-                for i in range(len(input_tuples[0]))
-            ]
+            for i in range(len(input_tuples[0])):
+                feed_dict = {}
+                for j, inp in enumerate(trt_config.inputs):
+                    name = inp.name if inp.name is not None else feed_name[j]
+                    feed_dict[name] = input_tuples[j][i]
+                feeds.append(feed_dict)
         # run pir pass (including trt_op_marker_pass)
         program_with_pir = run_pir_pass(
             program,
@@ -640,7 +645,8 @@ def convert(model_path, config):
             >>> input_config = Input(
             >>>     min_input_shape=[1, input_dim],
             >>>     optim_input_shape=[2, input_dim],
-            >>>     max_input_shape=[4, input_dim]
+            >>>     max_input_shape=[4, input_dim],
+            >>>     name='x',
             >>> )
 
             >>> trt_config = TensorRTConfig(inputs=[input_config])
@@ -693,7 +699,8 @@ def convert(model_path, config):
             >>>         np.random.rand(1,3).astype(np.float32),
             >>>         np.random.rand(2,3).astype(np.float32),
             >>>         np.random.rand(4,3).astype(np.float32),
-            >>>     )
+            >>>     ),
+            >>>     name='x',
             >>> )
 
             >>> trt_config = TensorRTConfig(inputs=[input_config])

python/paddle/tensorrt/util.py

@@ -334,9 +334,12 @@ def is_shape_tensor(value):
     return total_elements <= 8 and total_elements >= 1 and is_int_dtype
 
 
-def get_cache_path():
-    home_path = os.path.expanduser("~")
-    cache_path = os.path.join(home_path, ".pp_trt_cache")
+def get_cache_path(cache_path):
+    if cache_path is not None:
+        cache_path = cache_path
+    else:
+        home_path = os.path.expanduser("~")
+        cache_path = os.path.join(home_path, ".pp_trt_cache")
 
     if not os.path.exists(cache_path):
         os.makedirs(cache_path)

test/ir/inference/test_trt_convert_range.py

@@ -144,7 +144,7 @@ def generate_trt_nodes_num(attrs, dynamic_shape):
         ), 1e-2
 
     def test(self):
-        self.run_test(run_pir=True)
+        self.run_test()
 
 
 class TrtConvertRangeStaticTest(TrtLayerAutoScanTest):

test/tensorrt/test_converter_model_resnet50.py

@@ -60,6 +60,7 @@ def test_paddle_to_tensorrt_conversion_r50(self):
             optim_input_shape=(1, 3, 224, 224),
             max_input_shape=(4, 3, 224, 224),
             input_data_type='float32',
+            name='input',
         )
         _, input_optim_data, _ = input_config.generate_input_data()