[Bug fixes] fix qat tests

test/quantization/test_post_training_quantization_mobilenetv1.py

@@ -25,6 +25,7 @@
 
 import paddle
 from paddle.dataset.common import download
+from paddle.io import Dataset
 from paddle.static.log_helper import get_logger
 from paddle.static.quantization import PostTrainingQuantization
 
@@ -116,6 +117,33 @@ def val(data_dir=DATA_DIR):
     return _reader_creator(file_list, 'val', shuffle=False, data_dir=data_dir)
 
 
+class ImageNetDataset(Dataset):
+    def __init__(self, data_dir=DATA_DIR, shuffle=False, need_label=False):
+        super().__init__()
+        self.need_label = need_label
+        self.data_dir = data_dir
+        val_file_list = os.path.join(data_dir, 'val_list.txt')
+        with open(val_file_list) as flist:
+            lines = [line.strip() for line in flist]
+            if shuffle:
+                np.random.shuffle(lines)
+            self.data = [line.split() for line in lines]
+
+    def __getitem__(self, index):
+        sample = self.data[index]
+        data_path = os.path.join(self.data_dir, sample[0])
+        data, label = process_image(
+            [data_path, sample[1]], mode='val', color_jitter=False, rotate=False
+        )
+        if self.need_label:
+            return data, np.array([label]).astype('int64')
+        else:
+            return data
+
+    def __len__(self):
+        return len(self.data)
+
+
 class TestPostTrainingQuantization(unittest.TestCase):
     def setUp(self):
         self.int8_download = 'int8/download'
@@ -267,7 +295,7 @@ def run_program(
         throughput = cnt / np.sum(periods)
         latency = np.average(periods)
         acc1 = np.sum(test_info) / cnt
-        return (throughput, latency, acc1)
+        return (throughput, latency, acc1, feed_dict)
 
     def generate_quantized_model(
         self,
@@ -284,6 +312,7 @@ def generate_quantized_model(
         batch_nums=1,
         onnx_format=False,
         deploy_backend=None,
+        feed_name="inputs",
     ):
         try:
             os.system("mkdir " + self.int8_model)
@@ -293,11 +322,30 @@ def generate_quantized_model(
 
         place = paddle.CPUPlace()
         exe = paddle.static.Executor(place)
-        val_reader = val()
+        image = paddle.static.data(
+            name=feed_name[0], shape=[None, 3, 224, 224], dtype='float32'
+        )
+        feed_list = [image]
+        if len(feed_name) == 2:
+            label = paddle.static.data(
+                name='label', shape=[None, 1], dtype='int64'
+            )
+            feed_list.append(label)
+
+        val_dataset = ImageNetDataset(need_label=len(feed_list) == 2)
+        data_loader = paddle.io.DataLoader(
+            val_dataset,
+            places=place,
+            feed_list=feed_list,
+            drop_last=False,
+            return_list=False,
+            batch_size=2,
+            shuffle=False,
+        )
 
         ptq = PostTrainingQuantization(
             executor=exe,
-            sample_generator=val_reader,
+            data_loader=data_loader,
             model_dir=model_path,
             model_filename=model_filename,
             params_filename=params_filename,
@@ -348,7 +396,12 @@ def run_test(
                 model, infer_iterations * batch_size
             )
         )
-        (fp32_throughput, fp32_latency, fp32_acc1) = self.run_program(
+        (
+            fp32_throughput,
+            fp32_latency,
+            fp32_acc1,
+            feed_name,
+        ) = self.run_program(
             model_path,
             model_filename,
             params_filename,
@@ -370,14 +423,15 @@ def run_test(
             batch_nums,
             onnx_format,
             deploy_backend,
+            feed_name,
         )
 
         _logger.info(
             "Start INT8 inference for {} on {} images ...".format(
                 model, infer_iterations * batch_size
             )
         )
-        (int8_throughput, int8_latency, int8_acc1) = self.run_program(
+        (int8_throughput, int8_latency, int8_acc1, _) = self.run_program(
             self.int8_model,
             model_filename,
             params_filename,
@@ -421,7 +475,7 @@ def test_post_training_kl_mobilenetv1(self):
         is_use_cache_file = False
         is_optimize_model = True
         diff_threshold = 0.025
-        batch_nums = 1
+        batch_nums = 2
         self.run_test(
             model,
             'inference.pdmodel',
@@ -607,7 +661,7 @@ def test_post_training_onnx_format_mobilenetv1_tensorrt(self):
         is_optimize_model = False
         onnx_format = True
         diff_threshold = 0.05
-        batch_nums = 2
+        batch_nums = 12
         deploy_backend = "tensorrt"
         self.run_test(
             model,
@@ -650,7 +704,7 @@ def test_post_training_onnx_format_mobilenetv1_mkldnn(self):
         is_optimize_model = False
         onnx_format = True
         diff_threshold = 0.05
-        batch_nums = 1
+        batch_nums = 12
         deploy_backend = "mkldnn"
         self.run_test(
             model,

test/quantization/test_post_training_quantization_resnet50.py

@@ -113,7 +113,7 @@ def run_program(
         throughput = cnt / np.sum(periods)
         latency = np.average(periods)
         acc1 = np.sum(test_info) / cnt
-        return (throughput, latency, acc1)
+        return (throughput, latency, acc1, feed_dict)
 
 
 class TestPostTrainingForResnet50ONNXFormat(TestPostTrainingForResnet50):