Enhance strategy exit policy (#1105)

neural_compressor/experimental/strategy/strategy.py

@@ -674,8 +674,10 @@ def _tune_cfg_converter(self, op_tuning_cfg):
                 tune_cfg[op_name_type] = op_config
         tune_cfg['calib_sampling_size'] = op_tuning_cfg['calib_sampling_size']
         if self.calib_dataloader is not None:
-            tune_cfg['calib_iteration'] =  math.ceil(int(tune_cfg['calib_sampling_size']) / \
-                                                    self.calib_dataloader.batch_size)
+            # For the accelerate's DataLoaderShard, use total_batch_size instead of batch_size
+            bs = getattr(self.calib_dataloader, 'batch_size') or getattr(self.calib_dataloader, 'total_batch_size')
+            assert bs > 0, f"Calibration dataloader's batch size should be greater than one but got {bs}"
+            tune_cfg['calib_iteration'] =  math.ceil(int(tune_cfg['calib_sampling_size']) / bs)
         else:
             tune_cfg['calib_iteration'] = 1
         tune_cfg['advance'] = self.cfg.quantization.advance
@@ -704,8 +706,10 @@ def set_tuning_space(self, conf):
         calib_sampling_size_lst = self.cfg.quantization.calibration.sampling_size
         calib_sampling_size_lst = [int(calib_sampling_size) for calib_sampling_size in calib_sampling_size_lst]
         if self.calib_dataloader:
-            self.calib_iter = [math.ceil(int(x) / self.calib_dataloader.batch_size) \
-                               for x in calib_sampling_size_lst]
+            # For the accelerate's DataLoaderShard, use total_batch_size instead of batch_size
+            bs = getattr(self.calib_dataloader, 'batch_size') or getattr(self.calib_dataloader, 'total_batch_size')
+            assert bs > 0, f"Calibration dataloader's batch size should be greater than one but got {bs}"
+            self.calib_iter = [math.ceil(int(x) / bs) for x in calib_sampling_size_lst]
         else:
             self.calib_iter = 1
         # create tuning space

neural_compressor/strategy/auto.py

@@ -107,7 +107,11 @@ def traverse(self):
         # Quantize model with default config
         super().traverse()
         if self.best_qmodel:
+            logger.info("[Strategy] Found the model meets accuracy requirements, ending the tuning process.")
             return
+        elif self.config.tuning_criterion.max_trials == 1:
+            logger.info("[Strategy] Not found the model meets accuracy requirements,\
+                but the max trial is 1, ending the tuning process.")
         else:
             # Start to try different strategies sequentially
             self.sequential_traverse()

neural_compressor/strategy/strategy.py

@@ -1130,8 +1130,10 @@ def _tune_cfg_converter(self, op_tuning_cfg):
                 tune_cfg[op_name_type] = op_config
         tune_cfg['calib_sampling_size'] = op_tuning_cfg['calib_sampling_size']
         if self.calib_dataloader is not None:
-            tune_cfg['calib_iteration'] =  math.ceil(int(tune_cfg['calib_sampling_size']) / \
-                                                    self.calib_dataloader.batch_size)
+            # For the accelerate's DataLoaderShard, use total_batch_size instead of batch_size
+            bs = getattr(self.calib_dataloader, 'batch_size') or getattr(self.calib_dataloader, 'total_batch_size')
+            assert bs > 0, f"Calibration dataloader's batch size should be greater than one but got {bs}"
+            tune_cfg['calib_iteration'] =  math.ceil(int(tune_cfg['calib_sampling_size']) / bs)
         else:
             tune_cfg['calib_iteration'] = 1
         tune_cfg['approach'] = self.config.approach
@@ -1160,8 +1162,10 @@ def build_tuning_space(self, config):
         calib_sampling_size_lst = self.config.calibration_sampling_size
         calib_sampling_size_lst = [int(calib_sampling_size) for calib_sampling_size in calib_sampling_size_lst]
         if self.calib_dataloader:
-            self.calib_iter = [math.ceil(int(x) / self.calib_dataloader.batch_size) \
-                               for x in calib_sampling_size_lst]
+            # For the accelerate's DataLoaderShard, use total_batch_size instead of batch_size
+            bs = getattr(self.calib_dataloader, 'batch_size') or getattr(self.calib_dataloader, 'total_batch_size')
+            assert bs > 0, f"Calibration dataloader's batch size should be greater than one but got {bs}"
+            self.calib_iter = [math.ceil(int(x) / bs) for x in calib_sampling_size_lst]
         else:
             self.calib_iter = 1
         # create tuning space

test/strategy/test_quant_level.py

@@ -472,6 +472,21 @@ def fake_eval3(model):
                 found_fp32_conv = True
         self.assertTrue(found_fp32_conv)
 
+    def test_quant_level_auto_with_max_trial(self):
+        # maxt_trails = 1: even if the accuracy does not meet the requirements,
+        # the tuning process ends after the first trial.
+        from neural_compressor.config import PostTrainingQuantConfig, TuningCriterion
+        acc_lst = [1.0, 0.9, 1.1, 1.2]
+        def fake_eval3(model):
+            result = acc_lst[0]
+            del acc_lst[0]
+            return result
+        tuning_criterion = TuningCriterion(max_trials=1)
+        conf = PostTrainingQuantConfig(approach='static', tuning_criterion=tuning_criterion)
+        q_model = fit(model=deepcopy(self.ort_resnet18), conf=conf, \
+            calib_dataloader=self.ort_cv_dataloader, eval_func=fake_eval3)
+        self.assertIsNone(q_model)
+
 
 if __name__ == "__main__":
     unittest.main()