Upload train_jit.py with huggingface_hub

train_jit.py

@@ -0,0 +1,334 @@
+#!/usr/bin/env python3
+"""
+CFM training with unconditional JiT (jit_model_unconditional.JiT).
+Mirrors train_cfm_unet.py (data, TensorBoard, checkpoints); model + YAML differ.
+
+JiT expects forward(x, t); torchdyn NeuralODE calls f(t, x) — use CFMFlowWrapper.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import torch
+import torch.nn as nn
+import torchvision
+from torch.utils.data import DataLoader
+from torch.utils.tensorboard import SummaryWriter
+from torchvision.transforms import v2
+from torchdyn.core import NeuralODE
+
+from torchcfm.conditional_flow_matching import ConditionalFlowMatcher
+
+from jit import JiT
+
+try:
+    import yaml  # type: ignore[import-untyped]
+except ImportError as e:  # pragma: no cover
+    raise ImportError("Please `pip install pyyaml` to use --config.") from e
+
+# Reuse dataset helpers from UNet trainer (same CLI for data)
+from train_unet import load_training_dataset
+
+
+def parse_args() -> argparse.Namespace:
+    p = argparse.ArgumentParser(description="Train unconditional JiT with Conditional Flow Matching")
+
+    p.add_argument(
+        "--dataset",
+        type=str,
+        default="imagenette",
+        choices=["cifar10", "imagenette"],
+        help="Training dataset",
+    )
+    p.add_argument("--data-root", type=str, default=".", help="Root for dataset download/cache")
+    p.add_argument("--cifar-split", type=str, default="train", choices=["train", "test"])
+    p.add_argument("--imagenette-split", type=str, default="train", choices=["train", "val"])
+    p.add_argument("--imagenette-size", type=str, default="160px", choices=["160px", "320px", "full"])
+    p.add_argument("--single-class", action="store_true")
+    p.add_argument("--class-id", type=int, default=0)
+    p.add_argument("--batch-size", type=int, default=64)
+    p.add_argument("--num-workers", type=int, default=4)
+
+    p.add_argument("--epochs", type=int, default=30)
+    p.add_argument("--device", type=str, default=None, help="cuda | cpu (default: auto)")
+    p.add_argument("--log-interval", type=int, default=100)
+    p.add_argument("--seed", type=int, default=0)
+
+    p.add_argument("--save-dir", type=str, default="./runs/cfm_jit/checkpoints")
+    p.add_argument(
+        "--log-dir",
+        type=str,
+        default="./runs/cfm_jit/tensorboard",
+    )
+    p.add_argument("--run-name", type=str, default=None)
+
+    p.add_argument(
+        "--config",
+        type=str,
+        default=None,
+        help="YAML with JiT + CFM hyperparameters (default: jit_config.yaml next to this script)",
+    )
+
+    return p.parse_args()
+
+
+def _dim_from_yaml(value: Any) -> tuple[int, int, int]:
+    if isinstance(value, (list, tuple)) and len(value) == 3:
+        return (int(value[0]), int(value[1]), int(value[2]))
+    raise ValueError("YAML 'dim' must be [C, H, W]")
+
+
+@dataclass
+class JiTTrainConfig:
+    sigma: float
+    dim: tuple[int, int, int]
+    lr: float
+    weight_decay: float
+    inference_steps: int
+    vis_batch_size: int
+    input_size: int
+    patch_size: int
+    hidden_size: int
+    depth: int
+    num_heads: int
+    mlp_ratio: float
+    attn_drop: float
+    proj_drop: float
+    bottleneck_dim: int
+    in_context_len: int
+    in_context_start: int
+
+
+REQUIRED_JIT_YAML_KEYS = (
+    "sigma",
+    "dim",
+    "lr",
+    "weight_decay",
+    "inference_steps",
+    "vis_batch_size",
+    "input_size",
+    "patch_size",
+    "hidden_size",
+    "depth",
+    "num_heads",
+    "mlp_ratio",
+    "attn_drop",
+    "proj_drop",
+    "bottleneck_dim",
+    "in_context_len",
+    "in_context_start",
+)
+
+
+def load_jit_config_yaml(path: str | os.PathLike[str]) -> JiTTrainConfig:
+    path = Path(path)
+    if not path.is_file():
+        raise FileNotFoundError(f"Config file not found: {path.resolve()}")
+
+    with open(path, encoding="utf-8") as f:
+        raw = yaml.safe_load(f)
+    if raw is None or not isinstance(raw, dict):
+        raise ValueError(f"Config must be a YAML mapping: {path}")
+
+    missing = [k for k in REQUIRED_JIT_YAML_KEYS if k not in raw]
+    if missing:
+        raise ValueError(f"Missing keys in {path}: {missing}")
+
+    dim = _dim_from_yaml(raw["dim"])
+    input_size = int(raw["input_size"])
+    if dim[1] != input_size or dim[2] != input_size:
+        raise ValueError(f"dim {dim} must match input_size×input_size ({input_size})")
+
+    return JiTTrainConfig(
+        sigma=float(raw["sigma"]),
+        dim=dim,
+        lr=float(raw["lr"]),
+        weight_decay=float(raw["weight_decay"]),
+        inference_steps=int(raw["inference_steps"]),
+        vis_batch_size=int(raw["vis_batch_size"]),
+        input_size=input_size,
+        patch_size=int(raw["patch_size"]),
+        hidden_size=int(raw["hidden_size"]),
+        depth=int(raw["depth"]),
+        num_heads=int(raw["num_heads"]),
+        mlp_ratio=float(raw["mlp_ratio"]),
+        attn_drop=float(raw["attn_drop"]),
+        proj_drop=float(raw["proj_drop"]),
+        bottleneck_dim=int(raw["bottleneck_dim"]),
+        in_context_len=int(raw["in_context_len"]),
+        in_context_start=int(raw["in_context_start"]),
+    )
+
+
+def build_jit(cfg: JiTTrainConfig) -> JiT:
+    c = cfg.dim[0]
+    return JiT(
+        input_size=cfg.input_size,
+        patch_size=cfg.patch_size,
+        in_channels=c,
+        hidden_size=cfg.hidden_size,
+        depth=cfg.depth,
+        num_heads=cfg.num_heads,
+        mlp_ratio=cfg.mlp_ratio,
+        attn_drop=cfg.attn_drop,
+        proj_drop=cfg.proj_drop,
+        bottleneck_dim=cfg.bottleneck_dim,
+        in_context_len=cfg.in_context_len,
+        in_context_start=cfg.in_context_start,
+    )
+
+
+class CFMFlowWrapper(nn.Module):
+    """
+    torchdyn NeuralODE expects f(t, x) with dx/dt returned.
+    - velocity mode: JiT predicts v directly → return model(x, t).
+    - x_pred mode (JiT denoiser): JiT predicts x1 (clean); v = (x_pred - x) / (1-t).
+    """
+
+    def __init__(self, model: JiT, prediction_mode: str = "velocity", t_eps: float = 1e-5):
+        super().__init__()
+        self.model = model
+        self.prediction_mode = "x_pred"
+        self.t_eps = t_eps
+
+    def forward(self, t: torch.Tensor, x: torch.Tensor, y=None, *args, **kwargs) -> torch.Tensor:
+        batch = x.shape[0]
+        t_flat = torch.as_tensor(t, device=x.device, dtype=torch.float32).reshape(-1)
+        if t_flat.numel() == 1:
+            t_flat = t_flat.expand(batch)
+        elif t_flat.shape[0] != batch:
+            t_flat = t_flat[:batch]
+
+        if self.prediction_mode == "x_pred":
+            x_pred = self.model(x, t_flat)
+            one_minus_t = (1.0 - t_flat).clamp(min=self.t_eps)
+            t_bc = one_minus_t.reshape(-1, *([1] * (x.dim() - 1)))
+            return (x_pred - x) / t_bc
+        return self.model(x, t_flat)
+
+
+def main() -> None:
+    args = parse_args()
+    default_cfg = Path(__file__).resolve().parent / "jit_config.yaml"
+    config_path = Path(args.config).resolve() if args.config else default_cfg
+    cfg = load_jit_config_yaml(config_path)
+    print(f"Loaded JiT config from: {config_path}")
+
+    torch.manual_seed(args.seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(args.seed)
+
+    device = torch.device(args.device or ("cuda" if torch.cuda.is_available() else "cpu"))
+    print(f"Using device: {device}")
+
+    os.makedirs(args.save_dir, exist_ok=True)
+
+    tb_dir = os.path.join(args.log_dir, args.run_name) if args.run_name else args.log_dir
+    os.makedirs(tb_dir, exist_ok=True)
+    writer = SummaryWriter(log_dir=tb_dir)
+    writer.add_text("config/args", str(vars(args)), 0)
+    writer.add_text("config/jit_yaml", config_path.read_text(encoding="utf-8"), 0)
+
+    transforms = v2.Compose(
+        [
+            v2.ToTensor(),
+            v2.ToDtype(torch.float32, scale=True),
+            v2.Resize((cfg.input_size, cfg.input_size)),
+            v2.Normalize(mean=[0.0, 0.0, 0.0], std=[1.0, 1.0, 1.0]),
+        ]
+    )
+    train_dataset = load_training_dataset(args, transforms)
+    print(f"Dataset: {args.dataset}, size={len(train_dataset)}")
+
+    dummy_dataloader = DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=args.num_workers,
+        pin_memory=device.type == "cuda",
+    )
+
+    total_optimizer_steps = len(dummy_dataloader) * args.epochs
+
+    fm = ConditionalFlowMatcher(sigma=cfg.sigma)
+    net_model = build_jit(cfg).to(device)
+    ode_net = CFMFlowWrapper(net_model)
+
+    optim = torch.optim.AdamW(net_model.parameters(), lr=cfg.lr, weight_decay=cfg.weight_decay)
+    scheduler = torch.optim.lr_scheduler.LinearLR(optim, total_iters=max(total_optimizer_steps, 1))
+    t_span = torch.linspace(0, 1, cfg.inference_steps + 1, device=device)
+
+    c, h, w = cfg.dim
+    global_step = 0
+    best_loss = float("inf")
+
+    for ep in range(args.epochs):
+        net_model.train()
+        epoch_loss = 0.0
+        num_batches = 0
+
+        for data in dummy_dataloader:
+            x1 = data[0].to(device, non_blocking=True)
+            x0 = torch.randn_like(x1)
+            t, xt, ut = fm.sample_location_and_conditional_flow(x0, x1)
+            t_b = t.reshape(-1).float()
+            # vt = net_model(xt, t_b)
+            # loss = torch.mean((vt - ut) ** 2)
+            x_pred = net_model(xt, t_b)
+            one_minus_t = (1.0 - t_b).clamp(min=1.0e-5)
+            t_bc = one_minus_t.reshape(-1, *([1] * (xt.dim() - 1)))
+            v_pred = (x_pred - xt) / t_bc
+            v_target = (x1 - xt) / t_bc
+            loss = torch.mean((v_target - v_pred) ** 2)
+
+            optim.zero_grad(set_to_none=True)
+            loss.backward()
+            optim.step()
+            scheduler.step()
+
+            epoch_loss += loss.item()
+            num_batches += 1
+
+            writer.add_scalar("train/loss_step", loss.item(), global_step)
+            writer.add_scalar("train/lr", scheduler.get_last_lr()[0], global_step)
+
+            if global_step % args.log_interval == 0:
+                print(f"[step {global_step}] loss = {loss.item():.6f}")
+
+            global_step += 1
+
+        avg_epoch_loss = epoch_loss / max(num_batches, 1)
+        writer.add_scalar("train/loss_epoch", avg_epoch_loss, ep)
+        print(f"[epoch {ep}] avg loss = {avg_epoch_loss:.6f}")
+
+        net_model.eval()
+        node = NeuralODE(ode_net, solver="euler")
+        with torch.no_grad():
+            x_vis = torch.randn(cfg.vis_batch_size, c, h, w, device=device)
+            traj = node.trajectory(x_vis, t_span=t_span)
+            x_final = traj[-1]
+            x_final = x_final.clamp(0.0, 1.0).cpu()
+            grid = torchvision.utils.make_grid(x_final, nrow=4, padding=2, normalize=False)
+            writer.add_image("samples/neural_ode_final", grid, ep)
+        
+        
+        if ep % 30 == 0:
+            ckpt_path = os.path.join(args.save_dir, f"model_epoch_{ep}.pt")
+            torch.save(net_model.state_dict(), ckpt_path)
+
+        if ep == 0 or avg_epoch_loss < best_loss:
+            best_loss = avg_epoch_loss
+            torch.save(net_model.state_dict(), os.path.join(args.save_dir, "model_best.pt"))
+
+    writer.close()
+    print(f"Done. Checkpoints: {args.save_dir}")
+    print(f"TensorBoard: tensorboard --logdir {tb_dir}")
+
+
+if __name__ == "__main__":
+    main()