Merge pull request PaddlePaddle#9 from lizexu123/add_trt

fix
lizexu123 · Jul 29, 2024 · faeb360 · faeb360
2 parents 52f1419 + d0a724f
commit faeb360
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 8 deletions.
diff --git a/python/paddle/pp_tensorrt/test_converter.py b/python/paddle/pp_tensorrt/test_converter.py
@@ -15,18 +15,16 @@
 import numpy as np
 from converter import PaddleToTensorRTConverter
 from util import (
-    forbid_op_lower_trt,
     enforce_op_lower_trt,
     get_bert_program,
     get_dummy_program,
-    get_r50_program,
     get_idg_program,
+    get_r50_program,
     predict_program,
     run_pir_pass,
     warmup_shape_infer,
 )
 
-import paddle.static as static
 import paddle
 
 
@@ -187,21 +185,32 @@ def test_paddle_to_tensorrt_conversion_r50():
 def test_paddle_to_tensorrt_conversion_idg():
     # Step1: get program and init fake inputs
     program, scope, param_dict = get_idg_program()
-    
+
     map_vector_features_data = np.random.rand(1, 1400, 11, 17).astype('float32')
     polyline_mask_data = np.random.randint(0, 2, size=(1, 1400)).astype('bool')
 
     # Step1.1: get original results(for tests only)
     output_var = program.list_vars()[-1]
     output_expected = predict_program(
-        program, {"map_vector_features": map_vector_features_data, "polyline_mask": polyline_mask_data}, [output_var]
+        program,
+        {
+            "map_vector_features": map_vector_features_data,
+            "polyline_mask": polyline_mask_data,
+        },
+        [output_var],
     )
 
     # Step2: run warmup for collecting shape
     warmup_shape_infer(
         program,
-        min_shape_feed={"map_vector_features": map_vector_features_data, "polyline_mask": polyline_mask_data},
-        max_shape_feed={"map_vector_features": map_vector_features_data, "polyline_mask": polyline_mask_data},
+        min_shape_feed={
+            "map_vector_features": map_vector_features_data,
+            "polyline_mask": polyline_mask_data,
+        },
+        max_shape_feed={
+            "map_vector_features": map_vector_features_data,
+            "polyline_mask": polyline_mask_data,
+        },
     )
 
     # Step3: run pir pass(including some fusion pass and trt_op_marker_pass)
@@ -222,7 +231,12 @@ def test_paddle_to_tensorrt_conversion_idg():
 
     # Step6: run inference(converted_program)
     output_converted = predict_program(
-        program_with_pir, {"map_vector_features": map_vector_features_data, "polyline_mask": polyline_mask_data}, [output_var]
+        program_with_pir,
+        {
+            "map_vector_features": map_vector_features_data,
+            "polyline_mask": polyline_mask_data,
+        },
+        [output_var],
     )
 
     # Check that the results are close to each other within a tolerance of 1e-3

diff --git a/python/paddle/pp_tensorrt/util.py b/python/paddle/pp_tensorrt/util.py
@@ -98,6 +98,7 @@ def run_pir_pass(program, partition_mode=False):
     return program
 
 
+
 def forbid_op_lower_trt(program, op_name):
     for op in program.global_block().ops:
         if op.name() == op_name:
@@ -238,6 +239,7 @@ def forward(self, input_ids):
         return encoded_output
 
 
+
 def get_bert_program():
     paddle.enable_static()