diff options
| -rw-r--r-- | training/artifacts.py | 115 |
1 files changed, 115 insertions, 0 deletions
diff --git a/training/artifacts.py b/training/artifacts.py new file mode 100644 index 0000000..9f970f8 --- /dev/null +++ b/training/artifacts.py @@ -0,0 +1,115 @@ +"""Fetches model artifacts from wandb.""" +from datetime import datetime +from pathlib import Path +import shutil +import sys +from typing import Optional, Union + +import click +from loguru import logger as log +import wandb +from wandb.apis.public import Run + +TRAINING_DIR = Path(__file__).parents[0].resolve() +ARTIFACTS_DIR = TRAINING_DIR.parent / "text_recognizer" / "artifacts" +RUNS_DIR = TRAINING_DIR / "logs" / "runs" + + +def _get_run_dir(run: Run) -> Optional[Path]: + created_at = datetime.fromisoformat(run.created_at).astimezone() + date = created_at.date() + hour = (created_at + created_at.utcoffset()).hour + runs = list((RUNS_DIR / f"{date}").glob(f"{hour}-*")) + if not runs: + return None + return runs[0] + + +def _get_best_weights(run_dir: Path) -> Optional[Path]: + checkpoints = list(run_dir.glob("**/epoch=*.ckpt")) + if not checkpoints: + return None + return checkpoints[0] + + +def _copy_config(run_dir: Path, dst_dir: Path) -> None: + log.info(f"Copying config to artifact folders ({dst_dir})") + shutil.copyfile(src=run_dir / "config.yaml", dst=dst_dir / "config.yaml") + + +def _copy_checkpoint(checkpoint: Path, dst_dir: Path) -> None: + """Copy model checkpoint from local directory.""" + log.info(f"Copying best run ({checkpoint}) to artifact folders ({dst_dir})") + shutil.copyfile(src=checkpoint, dst=dst_dir / "model.pt") + + +def save_model(run: Run, tag: str) -> None: + """Save model to artifacts.""" + dst_dir = ARTIFACTS_DIR / f"{tag}_text_recognizer" + dst_dir.mkdir(parents=True, exist_ok=True) + run_dir = _get_run_dir(run) + if not run_dir: + log.error("Could not find experiment locally!") + best_weights = _get_best_weights(run_dir) + if not best_weights: + log.error("Could not find checkpoint locally!") + _copy_config(run_dir, dst_dir) + _copy_checkpoint(best_weights, dst_dir) + log.info("Successfully moved model and config to artifacts directory") + # TODO: be able to download from w&b + + +def find_best_model(entity: str, project: str, tag: str, metric: str, mode: str) -> Run: + """Find the best model on wandb.""" + if mode == "min": + default_metric_value = sys.maxsize + sort_reverse = False + else: + default_metric_value = 0 + sort_reverse = True + api = wandb.Api() + runs = api.runs(f"{entity}/{project}", filters={"tags": {"$in": [tag]}}) + runs = sorted( + runs, + key=lambda run: run.summary.get(metric, default_metric_value), + reverse=sort_reverse, + ) + best_run = runs[0] + summary = best_run.summary + log.info( + f"Best run is ({best_run.name}, {best_run.id}) picked from {len(runs)} runs with the following metric" + ) + log.info( + f"{metric}: {summary[metric]}" + # , {metric.replace('val', 'test')}: {summary[metric.replace('val', 'test')]}" + ) + return best_run + + +@click.command() +@click.option("--entity", type=str, default="aktersnurra", help="Name of the author") +@click.option( + "--project", type=str, default="text-recognizer", help="The wandb project name" +) +@click.option( + "--tag", + type=click.Choice(["paragraphs", "lines"]), + default="paragraphs", + help="Tag to filter by", +) +@click.option( + "--metric", type=str, default="val_loss", help="Which metric to filter on" +) +@click.option( + "--mode", + type=click.Choice(["min", "max"]), + default="min", + help="Min or max value of metric", +) +def main(entity: str, project: str, tag: str, metric: str, mode: str) -> None: + best_run = find_best_model(entity, project, tag, metric, mode) + save_model(best_run, tag) + + +if __name__ == "__main__": + main() |