Upload test_pipeline.py with huggingface_hub

test_pipeline.py

@@ -0,0 +1,321 @@
+#!/usr/bin/env python3
+"""
+End-to-end test: data loading → model forward → backward.
+Verifies that the full pipeline works before committing to long training.
+
+Usage:
+    python test_pipeline.py
+    python test_pipeline.py --dataset active_matter --no-streaming --local_path /data/well
+"""
+import argparse
+import sys
+import time
+import traceback
+
+import torch
+import torch.nn as nn
+
+
+def fmt_mem():
+    if torch.cuda.is_available():
+        alloc = torch.cuda.memory_allocated() / 1e9
+        res = torch.cuda.memory_reserved() / 1e9
+        total = torch.cuda.get_device_properties(0).total_memory / 1e9
+        return f"alloc={alloc:.2f}GB, reserved={res:.2f}GB, total={total:.1f}GB"
+    return "CPU only"
+
+
+def test_data_loading(args):
+    """Test 1: Load data and print shapes."""
+    print("\n" + "=" * 60)
+    print("TEST 1: Data Loading")
+    print("=" * 60)
+
+    from data_pipeline import create_dataloader, prepare_batch, get_channel_info, get_data_info
+
+    t0 = time.time()
+    loader, dataset = create_dataloader(
+        dataset_name=args.dataset,
+        split="train",
+        batch_size=args.batch_size,
+        streaming=args.streaming,
+        local_path=args.local_path,
+    )
+    print(f"  Dataset created in {time.time() - t0:.1f}s")
+    print(f"  Dataset length: {len(dataset)}")
+
+    # Probe shapes
+    info = get_data_info(dataset)
+    print(f"  Sample fields:")
+    for k, v in info.items():
+        print(f"    {k}: {v}")
+
+    ch = get_channel_info(dataset)
+    print(f"  Channel info: {ch}")
+
+    # Load one batch
+    t0 = time.time()
+    batch = next(iter(loader))
+    print(f"  First batch loaded in {time.time() - t0:.1f}s")
+    print(f"  Batch keys: {list(batch.keys())}")
+    for k, v in batch.items():
+        if isinstance(v, torch.Tensor):
+            print(f"    {k}: {v.shape} ({v.dtype})")
+
+    # Prepare for model
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    x_in, x_out = prepare_batch(batch, device)
+    print(f"  Model input:  {x_in.shape} ({x_in.dtype})")
+    print(f"  Model target: {x_out.shape} ({x_out.dtype})")
+    print(f"  GPU memory: {fmt_mem()}")
+
+    return ch, x_in, x_out
+
+
+def test_diffusion(ch, x_in, x_out):
+    """Test 2: Diffusion model forward + backward."""
+    print("\n" + "=" * 60)
+    print("TEST 2: Diffusion Model")
+    print("=" * 60)
+
+    from unet import UNet
+    from diffusion import GaussianDiffusion
+
+    c_in = ch["input_channels"]
+    c_out = ch["output_channels"]
+
+    unet = UNet(
+        in_channels=c_out + c_in,
+        out_channels=c_out,
+        base_ch=64,
+        ch_mults=(1, 2, 4, 8),
+        n_res=2,
+        attn_levels=(3,),
+    )
+    model = GaussianDiffusion(unet, timesteps=1000)
+    device = x_in.device
+    model = model.to(device)
+
+    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    print(f"  Parameters: {n_params:,}")
+    print(f"  GPU memory after model: {fmt_mem()}")
+
+    # Forward
+    t0 = time.time()
+    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        loss = model.training_loss(x_out, x_in)
+    print(f"  Forward pass: loss={loss.item():.4f} ({time.time() - t0:.3f}s)")
+    print(f"  GPU memory after forward: {fmt_mem()}")
+
+    # Backward
+    t0 = time.time()
+    loss.backward()
+    print(f"  Backward pass: ({time.time() - t0:.3f}s)")
+    print(f"  GPU memory after backward: {fmt_mem()}")
+
+    # Quick sampling test (just 5 steps for speed)
+    model.eval()
+    model.T = 5  # temporarily reduce for testing
+    model.betas = model.betas[:5]
+    model.alphas = model.alphas[:5]
+    model.alpha_bar = model.alpha_bar[:5]
+    model.sqrt_alpha_bar = model.sqrt_alpha_bar[:5]
+    model.sqrt_one_minus_alpha_bar = model.sqrt_one_minus_alpha_bar[:5]
+    model.sqrt_recip_alpha = model.sqrt_recip_alpha[:5]
+    model.posterior_variance = model.posterior_variance[:5]
+
+    t0 = time.time()
+    with torch.no_grad():
+        sample = model.sample(x_in[:2], shape=(2, c_out, x_in.shape[2], x_in.shape[3]))
+    print(f"  Sampling (5 steps, B=2): shape={sample.shape} ({time.time() - t0:.3f}s)")
+
+    del model
+    torch.cuda.empty_cache()
+    print(f"  DIFFUSION OK")
+
+
+def test_jepa(ch, x_in, x_out):
+    """Test 3: JEPA forward + backward."""
+    print("\n" + "=" * 60)
+    print("TEST 3: JEPA Model")
+    print("=" * 60)
+
+    from jepa import JEPA
+
+    c_in = ch["input_channels"]
+    device = x_in.device
+
+    model = JEPA(
+        in_channels=c_in,
+        latent_channels=128,
+        base_ch=32,
+        pred_hidden=256,
+    ).to(device)
+
+    n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    total_params = sum(p.numel() for p in model.parameters())
+    print(f"  Trainable parameters: {n_params:,}")
+    print(f"  Total parameters (incl EMA target): {total_params:,}")
+    print(f"  GPU memory after model: {fmt_mem()}")
+
+    # Forward
+    t0 = time.time()
+    with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+        loss, metrics = model.compute_loss(x_in, x_out)
+    print(f"  Forward: loss={loss.item():.4f}, metrics={metrics} ({time.time() - t0:.3f}s)")
+    print(f"  GPU memory after forward: {fmt_mem()}")
+
+    # Backward
+    t0 = time.time()
+    loss.backward()
+    print(f"  Backward: ({time.time() - t0:.3f}s)")
+    print(f"  GPU memory after backward: {fmt_mem()}")
+
+    # EMA update
+    model.update_target()
+    print(f"  EMA update: OK")
+
+    # Check latent shapes
+    model.eval()
+    with torch.no_grad():
+        z_pred, z_target = model(x_in[:2], x_out[:2])
+    print(f"  Latent shapes: pred={z_pred.shape}, target={z_target.shape}")
+
+    del model
+    torch.cuda.empty_cache()
+    print(f"  JEPA OK")
+
+
+def test_training_step(ch, loader):
+    """Test 4: Full training step with optimizer and grad scaling."""
+    print("\n" + "=" * 60)
+    print("TEST 4: Full Training Step")
+    print("=" * 60)
+
+    from data_pipeline import prepare_batch
+    from unet import UNet
+    from diffusion import GaussianDiffusion
+
+    c_in = ch["input_channels"]
+    c_out = ch["output_channels"]
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    unet = UNet(in_channels=c_out + c_in, out_channels=c_out, base_ch=64)
+    model = GaussianDiffusion(unet, timesteps=1000).to(device)
+    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
+    scaler = torch.amp.GradScaler("cuda")
+
+    model.train()
+    losses = []
+
+    for i, batch in enumerate(loader):
+        if i >= 3:
+            break
+
+        x_in, x_out = prepare_batch(batch, device)
+        optimizer.zero_grad(set_to_none=True)
+
+        with torch.amp.autocast("cuda", dtype=torch.bfloat16):
+            loss = model.training_loss(x_out, x_in)
+
+        scaler.scale(loss).backward()
+        scaler.unscale_(optimizer)
+        nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        scaler.step(optimizer)
+        scaler.update()
+
+        losses.append(loss.item())
+        print(f"  Step {i}: loss={loss.item():.4f}, mem={fmt_mem()}")
+
+    print(f"  3 training steps completed. Losses: {[f'{l:.4f}' for l in losses]}")
+    del model, optimizer, scaler
+    torch.cuda.empty_cache()
+    print(f"  TRAINING STEP OK")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--dataset", default="turbulent_radiative_layer_2D")
+    parser.add_argument("--streaming", action="store_true", default=True)
+    parser.add_argument("--no-streaming", dest="streaming", action="store_false")
+    parser.add_argument("--local_path", default=None)
+    parser.add_argument("--batch_size", type=int, default=4)
+    args = parser.parse_args()
+
+    print("=" * 60)
+    print("THE WELL - Pipeline End-to-End Test")
+    print("=" * 60)
+    print(f"Dataset:   {args.dataset}")
+    print(f"Streaming: {args.streaming}")
+    print(f"Batch:     {args.batch_size}")
+    print(f"Device:    {'cuda' if torch.cuda.is_available() else 'cpu'}")
+    if torch.cuda.is_available():
+        print(f"GPU:       {torch.cuda.get_device_name(0)}")
+    print(f"Memory:    {fmt_mem()}")
+
+    results = {}
+
+    # Test 1: Data
+    try:
+        ch, x_in, x_out = test_data_loading(args)
+        results["data"] = "PASS"
+    except Exception as e:
+        print(f"  FAIL: {e}")
+        traceback.print_exc()
+        results["data"] = f"FAIL: {e}"
+        sys.exit(1)
+
+    # Test 2: Diffusion
+    try:
+        test_diffusion(ch, x_in, x_out)
+        results["diffusion"] = "PASS"
+    except Exception as e:
+        print(f"  FAIL: {e}")
+        traceback.print_exc()
+        results["diffusion"] = f"FAIL: {e}"
+
+    # Test 3: JEPA
+    try:
+        test_jepa(ch, x_in, x_out)
+        results["jepa"] = "PASS"
+    except Exception as e:
+        print(f"  FAIL: {e}")
+        traceback.print_exc()
+        results["jepa"] = f"FAIL: {e}"
+
+    # Test 4: Training step
+    try:
+        loader, _ = __import__("data_pipeline").create_dataloader(
+            dataset_name=args.dataset,
+            split="train",
+            batch_size=args.batch_size,
+            streaming=args.streaming,
+            local_path=args.local_path,
+        )
+        test_training_step(ch, loader)
+        results["training_step"] = "PASS"
+    except Exception as e:
+        print(f"  FAIL: {e}")
+        traceback.print_exc()
+        results["training_step"] = f"FAIL: {e}"
+
+    # Summary
+    print("\n" + "=" * 60)
+    print("SUMMARY")
+    print("=" * 60)
+    all_pass = True
+    for name, status in results.items():
+        icon = "PASS" if status == "PASS" else "FAIL"
+        print(f"  [{icon}] {name}")
+        if status != "PASS":
+            all_pass = False
+
+    if all_pass:
+        print("\nAll tests passed! Pipeline is ready for training.")
+    else:
+        print("\nSome tests failed. Check output above.")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()