train_diffusion.py

#!/usr/bin/env python3
"""
Training script for conditional DDPM on The Well datasets.
Includes periodic evaluation with WandB video logging.

Usage:
    python train_diffusion.py --dataset turbulent_radiative_layer_2D --wandb
    python train_diffusion.py --dataset active_matter --batch_size 4 --wandb
"""
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 unet import UNet
from diffusion import GaussianDiffusion

# --- logging setup (suppress noisy library logs) ---
logging.basicConfig(level=logging.WARNING)
logger = logging.getLogger("train_diffusion")
logger.setLevel(logging.INFO)
_h = logging.StreamHandler()
_h.setFormatter(logging.Formatter("%(asctime)s [%(levelname)s] %(message)s", datefmt="%H:%M:%S"))
logger.addHandler(_h)
logger.propagate = False

# Also let eval_utils log through us
logging.getLogger("eval_utils").setLevel(logging.INFO)
logging.getLogger("eval_utils").addHandler(_h)
logging.getLogger("eval_utils").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 train(args):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    logger.info(f"Device: {device}")

    # ---- WandB ----
    wandb_run = None
    if args.wandb:
        import wandb

        wandb_run = wandb.init(
            project="the-well-diffusion",
            name=f"{args.dataset}_bs{args.batch_size}_lr{args.lr}",
            config=vars(args),
        )
        logger.info(f"WandB run: {wandb_run.url}")

    # ---- Data: train ----
    logger.info(f"Loading training data: {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"]

    # ---- Data: validation (single-step) ----
    logger.info("Loading validation data...")
    val_loader, _ = create_dataloader(
        dataset_name=args.dataset,
        split="valid",
        batch_size=args.batch_size,
        n_steps_input=args.n_input,
        n_steps_output=args.n_output,
        num_workers=0,
        streaming=args.streaming,
        local_path=args.local_path,
    )

    # ---- Data: rollout validation (multi-step output for GT comparison) ----
    logger.info(f"Loading rollout data (n_steps_output={args.n_rollout})...")
    rollout_loader, _ = create_dataloader(
        dataset_name=args.dataset,
        split="valid",
        batch_size=1,
        n_steps_input=args.n_input,
        n_steps_output=args.n_rollout,
        num_workers=0,
        streaming=args.streaming,
        local_path=args.local_path,
    )

    # ---- Model ----
    unet = UNet(
        in_channels=c_out + c_in,
        out_channels=c_out,
        base_ch=args.base_ch,
        ch_mults=tuple(args.ch_mults),
        n_res=args.n_res,
        attn_levels=tuple(args.attn_levels),
        dropout=args.dropout,
    )
    diffusion = GaussianDiffusion(unet, timesteps=args.timesteps).to(device)

    n_params = sum(p.numel() for p in diffusion.parameters() if p.requires_grad)
    logger.info(f"Model parameters: {n_params:,}")

    if wandb_run:
        wandb_run.summary["n_params"] = n_params

    # ---- Optimizer ----
    optimizer = torch.optim.AdamW(diffusion.parameters(), lr=args.lr, weight_decay=args.wd)
    scaler = GradScaler("cuda", enabled=args.amp)

    # ---- Checkpoint 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)
        diffusion.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 loop ----
    os.makedirs(args.ckpt_dir, exist_ok=True)
    total_steps = args.epochs * len(train_loader)

    logger.info(f"Starting training: {args.epochs} epochs, ~{total_steps} steps")
    logger.info(f"Eval every {args.eval_every} epochs, rollout {args.n_rollout} steps")

    for epoch in range(start_epoch, args.epochs):
        diffusion.train()
        epoch_loss = 0.0
        n_batches = 0
        t0 = time.time()

        pbar = tqdm(train_loader, desc=f"Epoch {epoch}", leave=False)
        for batch in pbar:
            try:
                x_cond, x_target = prepare_batch(batch, device)
            except Exception as e:
                logger.warning(f"Batch error: {e}, skipping")
                continue

            lr = cosine_lr(global_step, args.warmup, total_steps, args.lr)
            for pg in optimizer.param_groups:
                pg["lr"] = lr

            optimizer.zero_grad(set_to_none=True)

            with autocast(device_type="cuda", dtype=torch.bfloat16, enabled=args.amp):
                loss = diffusion.training_loss(x_target, x_cond)

            scaler.scale(loss).backward()
            scaler.unscale_(optimizer)
            nn.utils.clip_grad_norm_(diffusion.parameters(), args.grad_clip)
            scaler.step(optimizer)
            scaler.update()

            epoch_loss += loss.item()
            n_batches += 1
            global_step += 1

            pbar.set_postfix(loss=f"{loss.item():.4f}", lr=f"{lr:.2e}")

            if wandb_run and global_step % 20 == 0:
                wandb_run.log({"train/loss": loss.item(), "train/lr": lr}, step=global_step)

        avg_loss = epoch_loss / max(n_batches, 1)
        elapsed = time.time() - t0
        logger.info(
            f"Epoch {epoch}: loss={avg_loss:.4f}, batches={n_batches}, "
            f"time={elapsed:.1f}s, lr={lr:.2e}"
        )
        if wandb_run:
            wandb_run.log({"train/epoch_loss": avg_loss, "epoch": epoch}, step=global_step)

        # ---- Evaluation with video logging ----
        if (epoch + 1) % args.eval_every == 0:
            from eval_utils import run_evaluation

            logger.info("=" * 40)
            logger.info(f"EVALUATION at epoch {epoch}")
            logger.info("=" * 40)

            eval_metrics = run_evaluation(
                model=diffusion,
                val_loader=val_loader,
                rollout_loader=rollout_loader,
                device=device,
                global_step=global_step,
                wandb_run=wandb_run,
                n_val_batches=args.eval_batches,
                n_rollout=args.n_rollout,
                ddim_steps=args.ddim_steps,
            )

            logger.info(
                f"  val/mse={eval_metrics['val/mse']:.6f}, "
                f"rollout_mse_mean={eval_metrics['val/rollout_mse_mean']:.6f}"
            )
            logger.info("=" * 40)

        # ---- Checkpoint ----
        if (epoch + 1) % args.save_every == 0 or epoch == args.epochs - 1:
            ckpt_path = os.path.join(args.ckpt_dir, f"diffusion_ep{epoch:04d}.pt")
            torch.save(
                {
                    "epoch": epoch,
                    "global_step": global_step,
                    "model": diffusion.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}")

    if wandb_run:
        wandb_run.finish()
    logger.info("Training complete.")


def main():
    p = argparse.ArgumentParser(description="Train conditional DDPM 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=8)
    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("--base_ch", type=int, default=64)
    p.add_argument("--ch_mults", type=int, nargs="+", default=[1, 2, 4, 8])
    p.add_argument("--n_res", type=int, default=2)
    p.add_argument("--attn_levels", type=int, nargs="+", default=[3])
    p.add_argument("--dropout", type=float, default=0.1)
    p.add_argument("--timesteps", type=int, default=1000)
    # Optimization
    p.add_argument("--lr", type=float, default=1e-4)
    p.add_argument("--wd", type=float, default=0.01)
    p.add_argument("--warmup", type=int, default=1000)
    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)
    # Evaluation
    p.add_argument("--eval_every", type=int, default=5, help="Eval every N epochs")
    p.add_argument("--eval_batches", type=int, default=4, help="Val batches for MSE")
    p.add_argument("--n_rollout", type=int, default=20, help="Rollout steps for video")
    p.add_argument("--ddim_steps", type=int, default=50, help="DDIM steps for eval sampling")
    # Checkpointing
    p.add_argument("--ckpt_dir", default="checkpoints/diffusion")
    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()