| from __future__ import annotations | |
| import os | |
| import signal | |
| import torch | |
| import pytorch_lightning as pl | |
| from pytorch_lightning.callbacks import ModelCheckpoint | |
| from pytorch_lightning.plugins.environments import SLURMEnvironment | |
| from realfake.callbacks import ConsoleLogger, FeatureExtractorFreezeUnfreeze | |
| from realfake.models import RealFakeParams | |
| from realfake.utils import get_checkpoints_dir, find_latest_checkpoint | |
| def get_existing_checkpoint(job_id: str | None = None) -> tuple: | |
| if job_id is None: | |
| checkpoints_dir = get_checkpoints_dir(timestamp=True) | |
| else: | |
| checkpoints_dir = get_checkpoints_dir(timestamp=False)/job_id | |
| checkpoints_dir.mkdir(parents=True, exist_ok=True) | |
| existing_checkpoint = find_latest_checkpoint(checkpoints_dir) | |
| return checkpoints_dir, existing_checkpoint | |
| def prepare_trainer(args: RealFakeParams) -> pl.Trainer: | |
| job_id = os.environ.get("SLURM_JOB_ID") | |
| checkpoints_dir, existing_checkpoint = get_existing_checkpoint(job_id) | |
| if job_id is None: | |
| print("SLURM job id is not found, running locally.") | |
| if existing_checkpoint is None: | |
| print("No existing checkpoint found, starting from scratch.") | |
| if args.accelerator.override_float32_matmul: | |
| torch.set_float32_matmul_precision(args.accelerator.float32_matmul) | |
| with (checkpoints_dir/"params.json").open("w") as fp: | |
| fp.write(args.json()) | |
| trainer_params = dict( | |
| accelerator=args.accelerator.name, | |
| devices=args.accelerator.devices, | |
| precision=args.accelerator.precision, | |
| max_epochs=args.epochs, | |
| num_nodes=1, | |
| num_sanity_val_steps=0, | |
| enable_progress_bar=args.progress_bar, | |
| callbacks=[ | |
| ConsoleLogger(), | |
| ModelCheckpoint( | |
| monitor="val_acc", | |
| mode="max", | |
| save_last=True, | |
| save_top_k=1, | |
| dirpath=checkpoints_dir, | |
| filename="%s-{epoch:02d}-{val_acc:.4f}" % args.model_name, | |
| ), | |
| FeatureExtractorFreezeUnfreeze(unfreeze_at_epoch=args.freeze_epochs) | |
| ], | |
| resume_from_checkpoint=existing_checkpoint, | |
| ) | |
| if job_id is not None: | |
| trainer_params["enable_progress_bar"] = False | |
| trainer_params["plugins"] = SLURMEnvironment(requeue_signal=signal.SIGHUP) | |
| trainer_params["strategy"] = args.accelerator.strategy | |
| return pl.Trainer(**trainer_params) | |