Source code for titan_pylib.ahc.optimizer

  1import optuna
  2import time
  3from logging import getLogger, basicConfig
  4import os
  5import multiprocessing
  6from parallel_tester import ParallelTester, build_tester
  7from ahc_settings import AHCSettings
  8
  9logger = getLogger(__name__)
 10basicConfig(
 11    format="%(asctime)s [%(levelname)s] : %(message)s",
 12    datefmt="%H:%M:%S",
 13    level=os.getenv("LOG_LEVEL", "INFO"),
 14)
 15
 16

[docs]
 17def to_red(arg):
 18    return f"\u001b[31m{arg}\u001b[0m"

 19
 20

[docs]
 21def to_blue(arg):
 22    return f"\u001b[94m{arg}\u001b[0m"

 23
 24

[docs]
 25def to_green(arg):
 26    return f"\u001b[32m{arg}\u001b[0m"

 27
 28

[docs]
 29def to_bold(arg):
 30    return f"\u001b[1m{arg}\u001b[0m"

 31
 32

[docs]
 33class Optimizer:
 34
 35    def __init__(self, settings: AHCSettings) -> None:
 36        self.settings: AHCSettings = settings
 37        logger.info(f"------------------------------------------")
 38        logger.info(to_blue(f"Optimizer settings:"))
 39        logger.info(f"- study_name : {to_bold(settings.study_name)}")
 40        logger.info(f"- direction  : {to_bold(settings.direction)}")
 41        logger.info(f"- n_trials   : {to_bold(settings.n_trials)}")
 42        logger.info(f"------------------------------------------")
 43        self.path = f"./optimizer_results/{self.settings.study_name}"
 44

[docs]
 45    def optimize(self) -> None:
 46        if not os.path.exists(self.path):
 47            os.makedirs(self.path)
 48
 49        tester: ParallelTester = build_tester(
 50            self.settings, njobs=self.settings.n_jobs_parallel_tester
 51        )
 52        tester.compile()
 53
 54        start = time.time()
 55
 56        study: optuna.Study = optuna.create_study(
 57            direction=self.settings.direction,
 58            study_name=self.settings.study_name,
 59            storage=f"sqlite:///{self.path}/{self.settings.study_name}.db",
 60            load_if_exists=True,
 61        )
 62
 63        def _objective(trial: optuna.trial.Trial) -> float:
 64            tester: ParallelTester = build_tester(
 65                self.settings, njobs=self.settings.n_jobs_parallel_tester
 66            )
 67            args = self.settings.objective(trial)
 68            tester.append_execute_command(args)
 69            scores = tester.run()
 70            score = tester.get_score(scores)
 71            return score
 72
 73        study.optimize(
 74            _objective,
 75            n_trials=self.settings.n_trials,
 76            n_jobs=min(self.settings.n_jobs_optuna, multiprocessing.cpu_count() - 1),
 77        )
 78
 79        logger.info(study.best_trial)
 80        logger.info("writing results ...")
 81        self.output(study)
 82        logger.info(f"Finish parameter seraching. Time: {time.time() - start:.2f}sec.")

 83

[docs]
 84    def output(self, study: optuna.Study) -> None:
 85        if not os.path.exists(self.path):
 86            os.makedirs(self.path)
 87        with open(f"{self.path}/result.txt", "w", encoding="utf-8") as f:
 88            print(study.best_trial, file=f)
 89
 90        img_path = self.path + "/images"
 91        if not os.path.exists(img_path):
 92            os.makedirs(img_path)
 93
 94        fig = optuna.visualization.plot_contour(study)
 95        fig.write_html(f"{img_path}/contour.html")
 96        fig.write_image(f"{img_path}/contour.png")
 97        fig = optuna.visualization.plot_edf(study)
 98        fig.write_html(f"{img_path}/edf.html")
 99        fig.write_image(f"{img_path}/edf.png")
100        fig = optuna.visualization.plot_optimization_history(study)
101        fig.write_html(f"{img_path}/optimization_history.html")
102        fig.write_image(f"{img_path}/optimization_history.png")
103        fig = optuna.visualization.plot_parallel_coordinate(study)
104        fig.write_html(f"{img_path}/parallel_coordinate.html")
105        fig.write_image(f"{img_path}/parallel_coordinate.png")
106        fig = optuna.visualization.plot_slice(study)
107        fig.write_html(f"{img_path}/slice.html")
108        fig.write_image(f"{img_path}/slice.png")


109
110
111if __name__ == "__main__":
112    optimizer: Optimizer = Optimizer(AHCSettings)
113    optimizer.optimize()