[Inference]Modify TensorRT Input data name to collect shape in PIR (#71281)

PolaKuma · web-flow · commit b1e9134aba24 · 2025-02-27T19:28:28.000+08:00
* Update export.py

* fix codestyle

* Update export.py

* update
diff --git a/python/paddle/tensorrt/export.py b/python/paddle/tensorrt/export.py
@@ -52,7 +52,7 @@ class Input:
     This class supports generating random input data for minimum, optimal, and maximum shapes, with configurable data types (e.g., 'int' or 'float') and value ranges.
 
     Args:
-        input_data (tuple):
+        warmup_data (tuple):
             The tuple of Input data arrays (possibly different shapes).
         min_input_shape (tuple):
             The shape of the minimum input tensor.
@@ -89,14 +89,14 @@ class Input:
 
     def __init__(
         self,
-        input_data: tuple[np.ndarray, ...] | None = None,
+        warmup_data: tuple[np.ndarray, ...] | None = None,
         min_input_shape: tuple | None = None,
         max_input_shape: tuple | None = None,
         optim_input_shape: tuple | None = None,
         input_data_type: str | None = 'float32',
         input_range: tuple | None = None,
     ) -> None:
-        if input_data is not None:
+        if warmup_data is not None:
             if min_input_shape or max_input_shape or optim_input_shape:
                 logging.warning(
                     "Input data provided; min/max/optim shapes are ignored."
@@ -107,7 +107,7 @@ def __init__(
                     "When input_data is None, min/max/optim shapes must be specified."
                 )
 
-        self.input_data = input_data
+        self.warmup_data = warmup_data
         self.min_input_shape = min_input_shape
         self.max_input_shape = max_input_shape
         self.optim_input_shape = optim_input_shape
@@ -141,8 +141,8 @@ def generate_input_data(self):
             >>> input.input_range=(1,10)
             >>> input_min_data, input_optim_data, input_max_data = input_config.generate_input_data()
         """
-        if self.input_data is not None:
-            return self.input_data
+        if self.warmup_data is not None:
+            return self.warmup_data
         else:
             if self.input_data_type is None:
                 self.input_data_type = 'float32'
@@ -272,7 +272,7 @@ def __init__(
             >>> trt_config.workspace_size = 2 << 30
         """
         # Checking Input Consistency
-        has_input_data = [i.input_data is not None for i in inputs]
+        has_input_data = [i.warmup_data is not None for i in inputs]
         if any(has_input_data):
             if not all(has_input_data):
                 raise ValueError("All Inputs must have input_data if any does.")
@@ -306,6 +306,10 @@ def convert_to_trt(program, trt_config, scope):
 
     with paddle.pir_utils.IrGuard():
         input_tuples = [i.generate_input_data() for i in trt_config.inputs]
+        # Check all inputs have same number of warmup_data samples
+        assert (
+            len({len(t) for t in input_tuples}) == 1
+        ), "All inputs must have the same number of warmup_data samples."
         feeds = [
             {name: t[i] for t, name in zip(input_tuples, feed_name)}
             for i in range(len(input_tuples[0]))
diff --git a/test/tensorrt/test_converter_model_dummy.py b/test/tensorrt/test_converter_model_dummy.py
@@ -80,7 +80,7 @@ def test_paddle_to_tensorrt_collect_shape(self):
             np.random.rand(n, 64).astype(np.float32) for n in (1, 4, 8)
         )
         input_optim_data = input_data[1]
-        input_config = Input(input_data=input_data)
+        input_config = Input(warmup_data=input_data)
 
         # Create a TensorRTConfig with inputs as a required field.
         trt_config = TensorRTConfig(inputs=[input_config])
diff --git a/test/tensorrt/test_converter_model_resnet50.py b/test/tensorrt/test_converter_model_resnet50.py
@@ -103,7 +103,7 @@ def test_paddle_to_tensorrt_conversion_r50_collect_shape(self):
             np.random.rand(n, 3, 224, 224).astype(np.float32) for n in (1, 2, 4)
         )
         input_optim_data = input_data[1]
-        input_config = Input(input_data=input_data)
+        input_config = Input(warmup_data=input_data)
 
         # Create a TensorRTConfig with inputs as a required field.
         trt_config = TensorRTConfig(inputs=[input_config])

Original file line number	Diff line number	Diff line change
`@@ -80,7 +80,7 @@ def test_paddle_to_tensorrt_collect_shape(self):`
`80`	`80`	`np.random.rand(n, 64).astype(np.float32) for n in (1, 4, 8)`
`81`	`81`	`)`
`82`	`82`	`input_optim_data = input_data[1]`
`83`		`- input_config = Input(input_data=input_data)`
	`83`	`+ input_config = Input(warmup_data=input_data)`
`84`	`84`
`85`	`85`	`# Create a TensorRTConfig with inputs as a required field.`
`86`	`86`	`trt_config = TensorRTConfig(inputs=[input_config])`
Original file line number	Diff line number	Diff line change
`@@ -103,7 +103,7 @@ def test_paddle_to_tensorrt_conversion_r50_collect_shape(self):`
`103`	`103`	`np.random.rand(n, 3, 224, 224).astype(np.float32) for n in (1, 2, 4)`
`104`	`104`	`)`
`105`	`105`	`input_optim_data = input_data[1]`
`106`		`- input_config = Input(input_data=input_data)`
	`106`	`+ input_config = Input(warmup_data=input_data)`
`107`	`107`
`108`	`108`	`# Create a TensorRTConfig with inputs as a required field.`
`109`	`109`	`trt_config = TensorRTConfig(inputs=[input_config])`