diff --git a/neural_compressor/experimental/strategy/strategy.py b/neural_compressor/experimental/strategy/strategy.py index a14d6a291a2..c1da699f5d3 100644 --- a/neural_compressor/experimental/strategy/strategy.py +++ b/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 diff --git a/neural_compressor/strategy/auto.py b/neural_compressor/strategy/auto.py index 9a77c7645d4..c85c0882a9e 100644 --- a/neural_compressor/strategy/auto.py +++ b/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() \ No newline at end of file diff --git a/neural_compressor/strategy/strategy.py b/neural_compressor/strategy/strategy.py index d7c560154c1..8ed6f943563 100644 --- a/neural_compressor/strategy/strategy.py +++ b/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 diff --git a/test/strategy/test_quant_level.py b/test/strategy/test_quant_level.py index 5c69e0fecff..c7390c585da 100644 --- a/test/strategy/test_quant_level.py +++ b/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()