diff --git a/README.md b/README.md index 707ddc5..c564d2b 100644 --- a/README.md +++ b/README.md @@ -83,6 +83,16 @@ torchrun --standalone --nproc_per_node=4 train_exposure_autoencoder.py \ In DDP mode, `--batch_size` is the global batch size and must be divisible by the number of processes. +The trainer also writes `last.pt` after every epoch so interrupted runs can be +continued. Pass the run directory to reuse the original `train_config.json`; +the trainer will load `last.pt` when available and fall back to `best.pt` for +older runs: + +```bash +python train_exposure_autoencoder.py \ + --resume_checkpoint runs/exposure_ae_RUN +``` + Encode every cached exposure window once: ```bash diff --git a/train_exposure_autoencoder.py b/train_exposure_autoencoder.py index 55f9b76..2355b98 100644 --- a/train_exposure_autoencoder.py +++ b/train_exposure_autoencoder.py @@ -53,7 +53,7 @@ def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="Pretrain a lightweight TimesNet exposure autoencoder" ) - parser.add_argument("--exposure_cache_dir", required=True) + parser.add_argument("--exposure_cache_dir", default=None) parser.add_argument("--train_eid_file", default="ukb_train_eid.csv") parser.add_argument("--val_eid_file", default="ukb_val_eid.csv") parser.add_argument( @@ -70,6 +70,14 @@ def parse_args() -> argparse.Namespace: help="Ignore a compatible statistics cache and recompute it.", ) parser.add_argument("--runs_root", default="runs") + parser.add_argument( + "--resume_checkpoint", + default=None, + help=( + "Resume training from a run directory or checkpoint. A directory " + "uses last.pt when present, otherwise best.pt." + ), + ) parser.add_argument("--seed", type=int, default=42) parser.add_argument("--n_embd", type=int, default=120) parser.add_argument("--d_model", type=int, default=64) @@ -112,12 +120,19 @@ def parse_args() -> argparse.Namespace: default=4, help="DataLoader batches prefetched by each worker.", ) - args = parser.parse_args() + return parser.parse_args() + + +def validate_args(args: argparse.Namespace) -> None: + if not args.exposure_cache_dir: + raise ValueError( + "--exposure_cache_dir is required unless loaded from resume config" + ) if not 0.0 <= args.mask_ratio < 1.0: - parser.error("--mask_ratio must be in [0, 1)") + raise ValueError("--mask_ratio must be in [0, 1)") if args.num_workers > 0 and args.prefetch_factor <= 0: - parser.error("--prefetch_factor must be positive") - if args.gpu_ids: + raise ValueError("--prefetch_factor must be positive") + if isinstance(args.gpu_ids, str) and args.gpu_ids: try: args.gpu_ids = [ int(part.strip()) @@ -125,11 +140,45 @@ def parse_args() -> argparse.Namespace: if part.strip() ] except ValueError as exc: - parser.error("--gpu_ids must be a comma-separated list of integers") + raise ValueError( + "--gpu_ids must be a comma-separated list of integers" + ) from exc if not args.gpu_ids: - parser.error("--gpu_ids did not contain any valid CUDA device ids") + raise ValueError("--gpu_ids did not contain any valid CUDA device ids") args.data_parallel = True - return args + + +def resolve_resume_checkpoint(resume_path: str | None) -> Path | None: + if not resume_path: + return None + path = Path(resume_path) + if path.is_dir(): + last_path = path / "last.pt" + if last_path.is_file(): + return last_path + best_path = path / "best.pt" + if best_path.is_file(): + return best_path + raise FileNotFoundError( + f"Resume run directory has neither last.pt nor best.pt: {path}" + ) + if not path.is_file(): + raise FileNotFoundError(f"--resume_checkpoint does not exist: {path}") + return path + + +def apply_resume_config(args: argparse.Namespace, resume_checkpoint: Path) -> None: + config_path = resume_checkpoint.parent / "train_config.json" + if not config_path.is_file(): + raise FileNotFoundError( + f"Resume requires train_config.json next to checkpoint: {config_path}" + ) + config = json.loads(config_path.read_text(encoding="utf-8")) + resume_value = str(resume_checkpoint) + for key, value in config.items(): + if hasattr(args, key): + setattr(args, key, value) + args.resume_checkpoint = resume_value def select_rows(cache: ExposureCache, eids: set[int], split: str) -> np.ndarray: @@ -214,9 +263,14 @@ def distributed_run_dir( ) -> tuple[Path, str]: payload: list[str | None] = [None, None] if rank == 0: - run_dir, run_name = create_unique_run_dir( - lambda stamp: f"exposure_ae_{stamp}", Path(args.runs_root) - ) + if args.resume_checkpoint: + resume_path = Path(args.resume_checkpoint) + run_dir = resume_path.parent + run_name = run_dir.name + else: + run_dir, run_name = create_unique_run_dir( + lambda stamp: f"exposure_ae_{stamp}", Path(args.runs_root) + ) payload = [str(run_dir), run_name] if world_size > 1: dist.broadcast_object_list(payload, src=0) @@ -390,8 +444,145 @@ def learning_rate(epoch: int, args: argparse.Namespace) -> float: return args.base_lr * 0.5 * (1.0 + math.cos(math.pi * progress)) +def autoencoder_checkpoint_payload( + model, + optimizer: AdamW, + scaler: torch.amp.GradScaler, + config: dict, + raw_stats: tuple[np.ndarray, ...], + epoch: int, + val_loss: float, + best_loss: float, + stale_epochs: int, + history: list[dict], + include_training_state: bool, +) -> dict: + checkpoint_model = unwrap_model(model) + payload = { + "model_state_dict": checkpoint_model.state_dict(), + "encoder_state_dict": checkpoint_model.encoder.state_dict(), + "model_config": { + key: config[key] for key in ( + "n_embd", "d_model", "n_layers", "top_k", + "n_backbone_blocks", "backbone_kernel_size", + "backbone_expansion", "dropout", + ) + }, + "normalization": { + "daily_mean": raw_stats[0], + "daily_std": raw_stats[1], + "monthly_mean": raw_stats[2], + "monthly_std": raw_stats[3], + }, + "epoch": epoch, + "val_loss": val_loss, + "best_loss": best_loss, + "stale_epochs": stale_epochs, + "history": history, + } + if include_training_state: + payload["optimizer_state_dict"] = optimizer.state_dict() + payload["scaler_state_dict"] = scaler.state_dict() + return payload + + +def save_autoencoder_checkpoint(payload: dict, checkpoint_path: Path) -> None: + tmp_path = checkpoint_path.with_suffix(checkpoint_path.suffix + ".tmp") + torch.save(payload, tmp_path) + tmp_path.replace(checkpoint_path) + + +def load_resume_checkpoint( + checkpoint_path: Path, + model, + optimizer: AdamW, + scaler: torch.amp.GradScaler, + val_loader: DataLoader, + stats: tuple[torch.Tensor, ...], + device: torch.device, + grad_clip: float, + amp_enabled: bool, + show_progress: bool, + logger, +) -> tuple[int, float, int, list[dict]]: + checkpoint = torch.load(checkpoint_path, map_location=device) + if "model_state_dict" not in checkpoint: + raise KeyError( + f"Checkpoint does not contain model_state_dict: {checkpoint_path}" + ) + unwrap_model(model).load_state_dict(checkpoint["model_state_dict"]) + has_training_state = "optimizer_state_dict" in checkpoint + if has_training_state: + optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) + else: + logger.warning( + "Resume checkpoint has no optimizer state; continuing with a fresh " + "optimizer" + ) + if "scaler_state_dict" in checkpoint: + scaler.load_state_dict(checkpoint["scaler_state_dict"]) + elif scaler.is_enabled(): + logger.warning( + "Resume checkpoint has no AMP scaler state; continuing with a fresh " + "scaler" + ) + start_epoch = int(checkpoint.get("epoch", 0)) + history = list(checkpoint.get("history", [])) + history_path = checkpoint_path.parent / "history.json" + if not history and history_path.is_file(): + history = json.loads(history_path.read_text(encoding="utf-8")) + + if has_training_state: + best_loss = float( + checkpoint.get("best_loss", checkpoint.get("val_loss", float("inf"))) + ) + stale_epochs = int(checkpoint.get("stale_epochs", 0)) + else: + current_val_loss = run_epoch( + model, val_loader, device, stats, 0.0, None, + scaler, grad_clip, amp_enabled, show_progress, + ) + history_best_epoch = start_epoch + historical_best = float(checkpoint.get("val_loss", float("inf"))) + for entry in history: + if "val_loss" not in entry: + continue + entry_loss = float(entry["val_loss"]) + if entry_loss < historical_best: + historical_best = entry_loss + history_best_epoch = int(entry.get("epoch", len(history))) + if math.isfinite(historical_best): + tolerance = max(1e-4, abs(historical_best) * 1e-3) + if abs(current_val_loss - historical_best) > tolerance: + logger.warning( + "Legacy best.pt validation loss differs from history: " + f"recomputed={current_val_loss:.6f}, " + f"history_best={historical_best:.6f}" + ) + best_loss = min(historical_best, current_val_loss) + if history: + start_epoch = max(start_epoch, int(history[-1].get("epoch", len(history)))) + if current_val_loss < historical_best: + stale_epochs = 0 + else: + stale_epochs = max(0, start_epoch - history_best_epoch) + logger.info( + f"Validated legacy checkpoint {checkpoint_path.name}: " + f"val={current_val_loss:.6f}, historical_best={historical_best:.6f}" + ) + logger.info( + f"Resumed from {checkpoint_path} at epoch {start_epoch}; " + f"best_val={best_loss:.6f}, stale_epochs={stale_epochs}" + ) + return start_epoch, best_loss, stale_epochs, history + + def main() -> None: args = parse_args() + resume_checkpoint = resolve_resume_checkpoint(args.resume_checkpoint) + if resume_checkpoint is not None: + apply_resume_config(args, resume_checkpoint) + validate_args(args) device, rank, local_rank, world_size = init_distributed(args) set_seed(args.seed + rank) configure_torch_for_training(device) @@ -499,7 +690,7 @@ def main() -> None: "monthly_mean": raw_stats[2].tolist(), "monthly_std": raw_stats[3].tolist(), } - if rank == 0: + if rank == 0 and not args.resume_checkpoint: (run_dir / "train_config.json").write_text( json.dumps(config, indent=2), encoding="utf-8" ) @@ -507,7 +698,19 @@ def main() -> None: best_loss = float("inf") stale_epochs = 0 history = [] - for epoch in range(args.max_epochs): + start_epoch = 0 + if args.resume_checkpoint: + start_epoch, best_loss, stale_epochs, history = load_resume_checkpoint( + Path(args.resume_checkpoint), model, optimizer, scaler, + val_loader, stats, device, args.grad_clip, amp_enabled, + rank == 0, logger, + ) + if start_epoch >= args.max_epochs: + logger.info( + f"Resume checkpoint is already at epoch {start_epoch}; " + f"--max_epochs={args.max_epochs} leaves no remaining epochs" + ) + for epoch in range(start_epoch, args.max_epochs): if train_sampler is not None: train_sampler.set_epoch(epoch) lr = learning_rate(epoch, args) @@ -533,32 +736,25 @@ def main() -> None: best_loss = val_loss stale_epochs = 0 if rank == 0: - checkpoint_model = unwrap_model(model) - torch.save( - { - "model_state_dict": checkpoint_model.state_dict(), - "encoder_state_dict": checkpoint_model.encoder.state_dict(), - "model_config": { - key: config[key] for key in ( - "n_embd", "d_model", "n_layers", "top_k", - "n_backbone_blocks", "backbone_kernel_size", - "backbone_expansion", "dropout", - ) - }, - "normalization": { - "daily_mean": raw_stats[0], - "daily_std": raw_stats[1], - "monthly_mean": raw_stats[2], - "monthly_std": raw_stats[3], - }, - "epoch": epoch + 1, - "val_loss": val_loss, - }, + save_autoencoder_checkpoint( + autoencoder_checkpoint_payload( + model, optimizer, scaler, config, raw_stats, + epoch + 1, val_loss, best_loss, stale_epochs, + history, include_training_state=False, + ), run_dir / "best.pt", ) else: stale_epochs += 1 if rank == 0: + save_autoencoder_checkpoint( + autoencoder_checkpoint_payload( + model, optimizer, scaler, config, raw_stats, + epoch + 1, val_loss, best_loss, stale_epochs, + history, include_training_state=True, + ), + run_dir / "last.pt", + ) (run_dir / "history.json").write_text( json.dumps(history, indent=2), encoding="utf-8" )