Add resumable exposure autoencoder training

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

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"
             )