[Paddle TensorRT] Modify the base class for enabling FP16 unit tests (#70139)

lizexu123 · web-flow · commit f8a219feae24 · 2024-12-12T11:14:01.000+08:00
* fix

* fix

* fix
diff --git a/test/tensorrt/tensorrt_test_base.py b/test/tensorrt/tensorrt_test_base.py
@@ -42,7 +42,6 @@ def __init__(self, methodName='runTest'):
         self.max_shape = None
         self.target_marker_op = ""
         self.dynamic_shape_data = {}
-        self.enable_fp16 = None
 
     def create_fake_program(self):
         if self.python_api is None:
@@ -155,7 +154,7 @@ def prepare_feed(self):
                     new_list_args[sub_arg_name] = self.api_args[arg_name][i]
                 self.api_args[arg_name] = new_list_args
 
-    def check_trt_result(self, rtol=1e-5, atol=1e-5):
+    def check_trt_result(self, rtol=1e-5, atol=1e-5, precision_mode="fp32"):
         paddle.framework.set_flags({"FLAGS_trt_min_group_size": 1})
         with paddle.pir_utils.IrGuard():
             self.prepare_feed()
@@ -255,7 +254,7 @@ def check_trt_result(self, rtol=1e-5, atol=1e-5):
 
             # run TRTConverter(would lower group_op into tensorrt_engine_op)
             trt_config = None
-            if self.enable_fp16:
+            if precision_mode == "fp16":
                 input = Input(
                     min_input_shape=self.min_shape,
                     optim_input_shape=self.min_shape,
diff --git a/test/tensorrt/test_converter_conv.py b/test/tensorrt/test_converter_conv.py
@@ -40,9 +40,11 @@ def setUp(self):
         self.program_config = {"feed_list": ["x"]}
         self.min_shape = {"x": [1, 3, 8, 8]}
         self.max_shape = {"x": [10, 3, 8, 8]}
-        self.enable_fp16 = True
 
-    def test_trt_result(self):
+    def test_trt_result_fp16(self):
+        self.check_trt_result(precision_mode="fp16")
+
+    def test_trt_result_fp32(self):
         self.check_trt_result()
 
 
