Upload train_fast.py

train_fast.py

@@ -0,0 +1,106 @@
+import os, random, json, sys
+from pathlib import Path
+import numpy as np, torch
+from torch.utils.data import Dataset, DataLoader
+
+HUB_DIR = Path(__file__).parent.resolve()
+sys.path.insert(0, str(HUB_DIR / "src"))
+sys.path.insert(0, str(HUB_DIR / "src" / "models"))
+
+from models.grid_jepa import GridJEPA
+import trackio
+
+MAX_GRID = 30
+
+def load_tasks(d):
+    tasks = []
+    for p in sorted(Path(d).glob("*.json")):
+        with open(p) as f: tasks.append(json.load(f))
+    return tasks
+
+def to_tensor(grid):
+    arr = np.array(grid, dtype=np.int64)
+    H, W = arr.shape
+    t = torch.zeros(MAX_GRID, MAX_GRID, dtype=torch.long)
+    t[:H, :W] = torch.from_numpy(arr)
+    return t
+
+class ARCDataset(Dataset):
+    def __init__(self, d):
+        self.samples = []
+        for task in load_tasks(d):
+            for pair in task.get("train", []):
+                self.samples.append({"input": to_tensor(pair["input"]), "output": to_tensor(pair["output"])})
+    def __len__(self): return len(self.samples)
+    def __getitem__(self, idx):
+        s = self.samples[idx]
+        return {"context_grid": s["input"], "target_grid": s["output"]}
+
+def collate(batch):
+    return {"context_grid": torch.stack([b["context_grid"] for b in batch]), "target_grid": torch.stack([b["target_grid"] for b in batch])}
+
+def sample_masks(B, H, W, ratio=0.4, device="cpu"):
+    N = H * W; nt = max(1, int(N * ratio))
+    ctx = torch.zeros(B, N, dtype=torch.bool, device=device)
+    tgt = torch.zeros(B, N, dtype=torch.bool, device=device)
+    for b in range(B):
+        idx = list(range(N)); random.shuffle(idx)
+        tgt[b, idx[:nt]] = True; ctx[b, idx[nt:]] = True
+    return ctx, tgt
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"[TRAIN] Device: {device}", flush=True)
+run = trackio.init(project="arc-agi-3", name="grid-jepa-fast", group="pretrain")
+print("[TRAIN] Trackio OK", flush=True)
+
+model = GridJEPA(num_colors=10, embed_dim=192, encoder_depth=6, predictor_depth=6, num_heads=6, max_grid_size=MAX_GRID, ema_decay=0.996).to(device)
+print(f"[TRAIN] Params: {sum(p.numel() for p in model.parameters()):,}", flush=True)
+
+data_dir = HUB_DIR / "data" / "training"
+if not data_dir.exists(): data_dir = Path("/app/arc_data_source/data/training")
+ds = ARCDataset(str(data_dir))
+print(f"[TRAIN] Samples: {len(ds)}", flush=True)
+
+loader = DataLoader(ds, batch_size=8, shuffle=True, collate_fn=collate, num_workers=0)
+opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=0.05)
+
+save_dir = Path("/app/arc-jepa-hf/checkpoints")
+save_dir.mkdir(exist_ok=True)
+
+for epoch in range(1, 51):
+    model.train(); epoch_loss = 0.0; n = 0
+    for batch in loader:
+        ctx_g = batch["context_grid"].to(device)
+        tgt_g = batch["target_grid"].to(device)
+        B = ctx_g.shape[0]
+        ctx_mask, target_mask = sample_masks(B, MAX_GRID, MAX_GRID, ratio=0.5, device=device)
+        a_key = torch.zeros(B, dtype=torch.long, device=device)
+        a_pos = torch.zeros(B, dtype=torch.long, device=device)
+        opt.zero_grad()
+        loss, _ = model(tgt_g, ctx_mask, target_mask, a_key, a_pos)
+        if torch.isnan(loss) or loss.item() == 0: continue
+        loss.backward()
+        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+        opt.step(); model.update_ema()
+        epoch_loss += loss.item(); n += 1
+    avg = epoch_loss / max(n, 1)
+    print(f"[TRAIN] Epoch {epoch}/50: loss={avg:.4f}", flush=True)
+    run.log({"epoch_loss": avg, "epoch": epoch})
+    if epoch % 10 == 0 or epoch == 50:
+        ck = save_dir / f"ckpt_{epoch}.pt"
+        torch.save({"epoch": epoch, "model": model.state_dict()}, ck)
+        print(f"[TRAIN] Saved {ck}", flush=True)
+
+final = save_dir / "final.pt"
+torch.save({"model": model.state_dict(), "epoch": 50}, final)
+print(f"[TRAIN] Final: {final}", flush=True)
+
+from huggingface_hub import HfApi
+try:
+    api = HfApi()
+    api.upload_file(path_or_fileobj=str(final), path_in_repo="checkpoints/final.pt", repo_id="guychuk/arc-agi-3-grid-jepa", repo_type="model")
+    print("[TRAIN] Pushed to hub", flush=True)
+except Exception as e:
+    print(f"[TRAIN] Push failed: {e}", flush=True)
+run.finish()
+print("[TRAIN] Done", flush=True)