#!/usr/bin/env python3 """ Training script for Spatial JEPA on The Well datasets. Usage: python train_jepa.py --dataset turbulent_radiative_layer_2D --batch_size 16 python train_jepa.py --dataset active_matter --streaming --epochs 50 """ import argparse import logging import math import os import time import torch import torch.nn as nn from torch.amp import GradScaler, autocast from tqdm import tqdm from data_pipeline import create_dataloader, prepare_batch, get_channel_info from jepa import JEPA logging.basicConfig(level=logging.WARNING) # suppress noisy library logs logger = logging.getLogger("train_jepa") logger.setLevel(logging.INFO) _handler = logging.StreamHandler() _handler.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S")) logger.addHandler(_handler) logger.propagate = False def cosine_lr(step, warmup, total, base_lr, min_lr=1e-6): if step < warmup: return base_lr * step / max(warmup, 1) progress = (step - warmup) / max(total - warmup, 1) return min_lr + 0.5 * (base_lr - min_lr) * (1 + math.cos(progress * math.pi)) def cosine_ema(step, total, start=0.996, end=1.0): """EMA decay schedule: ramps from start to end over training.""" progress = step / max(total, 1) return end - (end - start) * (1 + math.cos(progress * math.pi)) / 2 def train(args): device = torch.device("cuda" if torch.cuda.is_available() else "cpu") logger.info(f"Device: {device}") # ---- Data ---- logger.info(f"Loading dataset: {args.dataset} (streaming={args.streaming})") train_loader, train_dataset = create_dataloader( dataset_name=args.dataset, split="train", batch_size=args.batch_size, n_steps_input=args.n_input, n_steps_output=args.n_output, num_workers=args.workers, streaming=args.streaming, local_path=args.local_path, ) ch_info = get_channel_info(train_dataset) logger.info(f"Channel info: {ch_info}") c_in = ch_info["input_channels"] c_out = ch_info["output_channels"] # JEPA uses same channel count for input and target # If they differ, we use max and pad in forward assert c_in == c_out, ( f"JEPA expects same input/output channels, got {c_in} vs {c_out}. " "Set n_input == n_output or use different architecture." ) # ---- Model ---- model = JEPA( in_channels=c_in, latent_channels=args.latent_ch, base_ch=args.base_ch, pred_hidden=args.pred_hidden, ema_decay=args.ema_start, ).to(device) n_params = sum(p.numel() for p in model.parameters() if p.requires_grad) logger.info(f"Trainable parameters: {n_params:,}") # ---- Optimizer ---- # Only optimize online encoder + predictor (target is EMA) trainable = list(model.online_encoder.parameters()) + list(model.predictor.parameters()) optimizer = torch.optim.AdamW(trainable, lr=args.lr, weight_decay=args.wd) scaler = GradScaler("cuda", enabled=args.amp) # ---- Resume ---- start_epoch = 0 global_step = 0 if args.resume and os.path.exists(args.resume): ckpt = torch.load(args.resume, map_location=device, weights_only=False) model.load_state_dict(ckpt["model"]) optimizer.load_state_dict(ckpt["optimizer"]) scaler.load_state_dict(ckpt["scaler"]) start_epoch = ckpt["epoch"] + 1 global_step = ckpt["global_step"] logger.info(f"Resumed from epoch {start_epoch}, step {global_step}") # ---- Training ---- os.makedirs(args.ckpt_dir, exist_ok=True) total_steps = args.epochs * len(train_loader) try: import wandb if args.wandb: wandb.init(project="the-well-jepa", config=vars(args)) except ImportError: args.wandb = False logger.info(f"Starting training: {args.epochs} epochs, ~{total_steps} steps") for epoch in range(start_epoch, args.epochs): model.train() epoch_loss = 0.0 epoch_metrics = {"sim": 0, "var": 0, "cov": 0} n_batches = 0 t0 = time.time() pbar = tqdm(train_loader, desc=f"Epoch {epoch}", leave=False) for batch in pbar: try: x_input, x_target = prepare_batch(batch, device) except Exception as e: logger.warning(f"Batch error: {e}, skipping") continue # LR schedule lr = cosine_lr(global_step, args.warmup, total_steps, args.lr) for pg in optimizer.param_groups: pg["lr"] = lr # EMA schedule ema = cosine_ema(global_step, total_steps, args.ema_start, args.ema_end) model.set_ema_decay(ema) optimizer.zero_grad(set_to_none=True) with autocast(device_type="cuda", dtype=torch.bfloat16, enabled=args.amp): loss, metrics = model.compute_loss(x_input, x_target) scaler.scale(loss).backward() scaler.unscale_(optimizer) nn.utils.clip_grad_norm_(trainable, args.grad_clip) scaler.step(optimizer) scaler.update() # EMA update model.update_target() epoch_loss += loss.item() for k in epoch_metrics: epoch_metrics[k] += metrics[k] n_batches += 1 global_step += 1 pbar.set_postfix( loss=f"{loss.item():.4f}", sim=f"{metrics['sim']:.4f}", ema=f"{ema:.4f}", ) if args.wandb: wandb.log( {"train/loss": loss.item(), "train/lr": lr, "train/ema": ema, **{f"train/{k}": v for k, v in metrics.items()}}, step=global_step, ) avg_loss = epoch_loss / max(n_batches, 1) avg_m = {k: v / max(n_batches, 1) for k, v in epoch_metrics.items()} elapsed = time.time() - t0 logger.info( f"Epoch {epoch}: loss={avg_loss:.4f}, sim={avg_m['sim']:.4f}, " f"var={avg_m['var']:.4f}, cov={avg_m['cov']:.4f}, " f"time={elapsed:.1f}s" ) # Checkpoint if (epoch + 1) % args.save_every == 0 or epoch == args.epochs - 1: ckpt_path = os.path.join(args.ckpt_dir, f"jepa_ep{epoch:04d}.pt") torch.save( { "epoch": epoch, "global_step": global_step, "model": model.state_dict(), "optimizer": optimizer.state_dict(), "scaler": scaler.state_dict(), "args": vars(args), "ch_info": ch_info, }, ckpt_path, ) logger.info(f"Saved {ckpt_path}") logger.info("Training complete.") def main(): p = argparse.ArgumentParser(description="Train Spatial JEPA on The Well") # Data p.add_argument("--dataset", default="turbulent_radiative_layer_2D") p.add_argument("--streaming", action="store_true", default=True) p.add_argument("--no-streaming", dest="streaming", action="store_false") p.add_argument("--local_path", default=None) p.add_argument("--batch_size", type=int, default=16) p.add_argument("--workers", type=int, default=0) p.add_argument("--n_input", type=int, default=1) p.add_argument("--n_output", type=int, default=1) # Model p.add_argument("--latent_ch", type=int, default=128) p.add_argument("--base_ch", type=int, default=32) p.add_argument("--pred_hidden", type=int, default=256) # Optimization p.add_argument("--lr", type=float, default=3e-4) p.add_argument("--wd", type=float, default=0.05) p.add_argument("--warmup", type=int, default=500) p.add_argument("--grad_clip", type=float, default=1.0) p.add_argument("--amp", action="store_true", default=True) p.add_argument("--no-amp", dest="amp", action="store_false") p.add_argument("--epochs", type=int, default=100) p.add_argument("--ema_start", type=float, default=0.996) p.add_argument("--ema_end", type=float, default=1.0) # Checkpointing p.add_argument("--ckpt_dir", default="checkpoints/jepa") p.add_argument("--save_every", type=int, default=5) p.add_argument("--resume", default=None) # Logging p.add_argument("--wandb", action="store_true", default=False) args = p.parse_args() train(args) if __name__ == "__main__": main()