Add script to rerun best trials.

tuning/rerun_best_trials.py

@@ -0,0 +1,98 @@
+import argparse
+from typing import List, Tuple
+
+import optuna
+
+import hp_search_spaces
+
+
+def make_parser() -> argparse.ArgumentParser:
+
+    parser = argparse.ArgumentParser(
+        description=
+        "Re-run the best trials from a previous tuning run.",
+        epilog=f"Example usage:\n"
+               f"python rerun_best_trials.py --journal-log tuning_run.json --tok-k 3\n",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument(
+        "algo",
+        type=str,
+        default=None,
+        choices=hp_search_spaces.objectives_by_algo.keys(),
+        help="The algorithm that has been tuned. "
+             "Can usually be deduced from the study name.",
+    )
+    parser.add_argument(
+        "--top-k",
+        type=int,
+        default=1,
+        help="How many of the top configurations to re-run."
+    )
+    parser.add_argument(
+        "--num-reruns",
+        type=int,
+        default=5,
+        help="How many times to re-run each of the top configurations."
+    )
+    parser.add_argument(
+        "-j",
+        "--journal-log",
+        type=str,
+        help="The journal file of the previous tuning run."
+    )
+    return parser
+
+
+def infer_algo_name(study: optuna.Study) -> Tuple[str, List[str]]:
+    """Infer the algo name from the study name.
+
+    Assumes that the study name is of the form "tuning_{algo}_with_{named_configs_}".
+
+    Args:
+        study: The optuna study.
+
+    Returns:
+        The algo name.
+    """
+    study_name_parts = study.study_name.split("_")
+    assert len(study_name_parts) >= 3
+    assert study_name_parts[0] == "tuning"
+    assert study_name_parts[2] == "with"
+    return study_name_parts[1]
+
+
+def main():
+    parser = make_parser()
+    args = parser.parse_args()
+    study: optuna.Study = optuna.load_study(storage=optuna.storages.JournalStorage(
+            optuna.storages.JournalFileStorage(args.journal_log)
+    ))
+
+    top_k_trials = sorted(
+        study.trials, key=lambda t: t.value, reverse=True
+    )[:args.top_k]
+
+    print("Best trials:")
+    for trial in top_k_trials:
+        print(trial)
+
+    algo_name = args.algo if args.algo is not None else infer_algo_name(study)
+    sacred_experiment = hp_search_spaces.objectives_by_algo[algo_name].sacred_ex
+
+    for trial in top_k_trials:
+        for i in range(args.num_reruns):
+            print(f"Rerunning trial {trial.number} for the {i+1}-th time.")
+            result = sacred_experiment.run(
+                config_updates=trial.user_attrs["config_updates"],
+                named_configs=trial.user_attrs["named_configs"],
+                options={"--name": study.study_name, "--file_storage": "sacred"},
+            )
+            if result.status != "COMPLETED":
+                raise RuntimeError(
+                    f"Trial failed with {result.fail_trace()} and status {result.status}."
+                )
+
+
+if __name__ == '__main__':
+    main()