fix bug

Zeref996 · Aug 9, 2021 · d9cd374 · d9cd374
1 parent 04ca2e1
commit d9cd374
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 11 deletions.
diff --git a/inference/python_api_test/test_case/infer_test.py b/inference/python_api_test/test_case/infer_test.py
@@ -159,9 +159,7 @@ def disable_gpu_test(self, input_data_dict: dict, repeat=20):
         gpu_max_mem = max([float(i["used(MB)"]) for i in _gpu_mem_lists])
         assert abs(gpu_max_mem - ori_gpu_mem) < 1, "set disable_gpu(), but gpu activity found"
 
-    def mkldnn_test(
-        self, input_data_dict: dict, output_data_dict: dict, mkldnn_cache_capacity=1, repeat=2, delta=1e-5, gpu_mem=1000
-    ):
+    def mkldnn_test(self, input_data_dict: dict, output_data_dict: dict, mkldnn_cache_capacity=1, repeat=2, delta=1e-5):
         """
         test enable_mkldnn()
         Args:

diff --git a/inference/python_api_test/test_class_model/test_resnet50.py b/inference/python_api_test/test_class_model/test_resnet50.py
@@ -96,8 +96,8 @@ def test_trtfp32_more_bz():
 
     file_path = "./resnet50"
     images_size = 224
-    batch_size_pool = 10
-    for batch_size in range(1, batch_size_pool + 1, 3):
+    batch_size_pool = [1, 5, 10]
+    for batch_size in batch_size_pool:
         test_suite = InferenceTest()
         test_suite.load_config(model_file="./resnet50/inference.pdmodel", params_file="./resnet50/inference.pdiparams")
         images_list, npy_list = test_suite.get_images_npy(file_path, images_size)
@@ -178,8 +178,8 @@ def test_trtfp16_more_bz():
 
     file_path = "./resnet50"
     images_size = 224
-    batch_size_pool = 10
-    for batch_size in range(1, batch_size_pool + 1, 3):
+    batch_size_pool = [1, 5, 10]
+    for batch_size in batch_size_pool:
         test_suite = InferenceTest()
         test_suite.load_config(model_file="./resnet50/inference.pdmodel", params_file="./resnet50/inference.pdiparams")
         images_list, npy_list = test_suite.get_images_npy(file_path, images_size)

diff --git a/inference/python_api_test/test_class_model/test_vgg11.py b/inference/python_api_test/test_class_model/test_vgg11.py
@@ -70,8 +70,8 @@ def test_trtfp32_more_bz():
 
     file_path = "./vgg11"
     images_size = 224
-    batch_size_pool = 10
-    for batch_size in range(1, batch_size_pool + 1, 3):
+    batch_size_pool = [1, 10]
+    for batch_size in batch_size_pool:
         test_suite = InferenceTest()
         test_suite.load_config(model_file="./vgg11/inference.pdmodel", params_file="./vgg11/inference.pdiparams")
         images_list, npy_list = test_suite.get_images_npy(file_path, images_size)
@@ -125,8 +125,8 @@ def test_trtfp16_more_bz():
 
     file_path = "./vgg11"
     images_size = 224
-    batch_size_pool = 10
-    for batch_size in range(1, batch_size_pool + 1, 4):
+    batch_size_pool = [1, 10]
+    for batch_size in batch_size_pool:
         test_suite = InferenceTest()
         test_suite.load_config(model_file="./vgg11/inference.pdmodel", params_file="./vgg11/inference.pdiparams")
         images_list, npy_list = test_suite.get_images_npy(file_path, images_size)