The main entry of auto-tune.
(
self, model: FakeModel, tune_config: TuningConfig, eval_fn: Callable, eval_args=None, *args, **kwargs
)
| 382 | return model |
| 383 | |
| 384 | def run( |
| 385 | self, model: FakeModel, tune_config: TuningConfig, eval_fn: Callable, eval_args=None, *args, **kwargs |
| 386 | ) -> Optional[FakeModel]: |
| 387 | """The main entry of auto-tune.""" |
| 388 | best_quant_model = None |
| 389 | eval_func_wrapper = EvaluationFuncWrapper(eval_fn, eval_args) |
| 390 | config_loader, tuning_logger, tuning_monitor = init_tuning(tuning_config=tune_config) |
| 391 | baseline: float = eval_func_wrapper.evaluate(model) |
| 392 | tuning_monitor.set_baseline(baseline) |
| 393 | tuning_logger.tuning_start() |
| 394 | for trial_index, quant_config in enumerate(config_loader): |
| 395 | tuning_logger.trial_start(trial_index=trial_index) |
| 396 | tuning_logger.execution_start() |
| 397 | logger.info(quant_config.to_dict()) |
| 398 | q_model = self._quantize(copy.deepcopy(model), quant_config, *args, **kwargs) |
| 399 | tuning_logger.execution_end() |
| 400 | tuning_logger.evaluation_start() |
| 401 | eval_result: float = eval_func_wrapper.evaluate(q_model) |
| 402 | tuning_logger.evaluation_end() |
| 403 | tuning_monitor.add_trial_result(trial_index, eval_result, quant_config) |
| 404 | tuning_logger.trial_end(trial_index) |
| 405 | if tuning_monitor.need_stop(): |
| 406 | logger.info("Stopped tuning.") |
| 407 | del q_model # maybe gc.collect() is needed for memory release |
| 408 | best_quant_config: BaseConfig = tuning_monitor.get_best_quant_config() |
| 409 | q_model = self._quantize(copy.deepcopy(model), best_quant_config, *args, **kwargs) |
| 410 | best_quant_model = q_model # quantize model inplace |
| 411 | break |
| 412 | tuning_logger.tuning_end() |
| 413 | return best_quant_model |
| 414 | |
| 415 | config_set = [FakeAlgoConfig(weight_bits=4), FakeAlgoConfig(weight_bits=8)] |
| 416 | tuning_config = TuningConfig(config_set=config_set) |
no test coverage detected