feat: add persistent Trackio Space for image logging (space_id + sync)

train_phase2.py

@@ -19,6 +19,8 @@ Visual diagnostics logged to Trackio (state-of-the-art for JEPA):
   9. Cross-Attention Weights in Perceiver — which evidence each query attends to
   10. Eigenspectrum Plot — singular value distribution of latent space
 
+All images are persisted to HF Space JorgeAV/MR-JEPA-Trackio via space_id parameter.
+
 Usage:
     python train_phase2.py --checkpoint checkpoints/hybrid_main_best.pt
     python train_phase2.py --epochs 10 --backbone_lr 1e-5
@@ -712,6 +714,8 @@ def main():
     parser.add_argument("--max_eval_samples", type=int, default=500)
     parser.add_argument("--vis_interval", type=int, default=100)
     parser.add_argument("--output_dir", default="./outputs/mrjepa_phase2")
+    parser.add_argument("--trackio_space", default="JorgeAV/MR-JEPA-Trackio",
+                        help="HF Space ID for persistent Trackio dashboard")
     args = parser.parse_args()
 
     log.info("Downloading Phase 1 training script...")
@@ -753,9 +757,12 @@ def main():
     log.info(f"Device: {device}")
     os.makedirs(cfg.output_dir, exist_ok=True)
 
+    # ── Initialize Trackio with persistent HF Space ──
     import trackio
     trackio.init(
-        name=args.run_name, project="MR-JEPA",
+        name=args.run_name,
+        project="MR-JEPA",
+        space_id=args.trackio_space,
         config={
             "phase": 2, "epochs": args.epochs,
             "core_lr": args.core_lr, "backbone_lr": args.backbone_lr, "text_lr": args.text_lr,
@@ -767,7 +774,7 @@ def main():
             "backbone": cfg.backbone, "K": cfg.K, "use_jepa": cfg.use_jepa, "loss_fn": cfg.loss_fn,
         }
     )
-    log.info("Trackio initialized with visual diagnostics")
+    log.info(f"Trackio initialized → Space: https://huggingface.co/spaces/{args.trackio_space}")
 
     log.info("Building model...")
     model = p1.MRJEPAModel(cfg)
@@ -839,102 +846,115 @@ def main():
     amp_dtype = torch.bfloat16 if cfg.bf16 else torch.float32
     trainable = [p for p in model.parameters() if p.requires_grad]
 
-    for epoch in range(cfg.epochs):
-        model.train()
-        epoch_losses = defaultdict(list)
-        epoch_correct = 0
-        epoch_total = 0
-        optimizer.zero_grad()
-
-        for batch_idx, batch in enumerate(train_dl):
-            batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
-            vis_tok = model.vis(batch["pixel_values"]).float()
-            txt_tok = model.txt(batch["input_ids"], batch["attention_mask"]).float()
-
-            with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=cfg.bf16 and device.type == "cuda"):
-                evidence, _, ev_mask = model.evidence(vis_tok, txt_tok, batch["attention_mask"])
-                if model._use_rollout:
-                    traj, z_final, z_proj = model.rollout(evidence)
-                else:
-                    B = batch["batch_size"]
-                    z0 = model.rollout.init_tokens.expand(B, -1, -1) + \
-                         model.rollout.z0_proj(F.adaptive_avg_pool1d(
-                             evidence.permute(0,2,1), model.rollout.num_tokens).permute(0,2,1))
-                    z_final = z0
-                    z_proj = model.rollout.out_proj(z0).unsqueeze(1)
-
-                if model._use_jepa:
-                    target_proj = model.target(vis_tok.detach(), txt_tok.detach(), batch["attention_mask"].detach())
-                else:
-                    target_proj = None
-
-                opt_emb = model.encode_options(batch["opt_input_ids"], batch["opt_attention_mask"])
-                opt_emb = opt_emb.view(batch["batch_size"], cfg.max_options, -1)
-                logits = model.disc(z_final, opt_emb, batch["opt_mask"])
-                task_loss = F.cross_entropy(logits, batch["labels"])
-
-                if model._use_jepa and target_proj is not None:
-                    losses = model.jepa_loss(z_proj, target_proj, task_loss)
-                else:
-                    losses = {"total": task_loss, "jepa": torch.tensor(0.0), "task": task_loss, "reg": torch.tensor(0.0)}
-                loss = losses["total"] / cfg.grad_accum
-
-            loss.backward()
-
-            if (batch_idx + 1) % cfg.grad_accum == 0:
-                nn.utils.clip_grad_norm_(trainable, cfg.max_grad_norm)
-                optimizer.step(); scheduler.step(); optimizer.zero_grad()
-                model.update_target(global_step, total_steps)
-                global_step += 1
-                if global_step % args.vis_interval == 0 and global_step > 0:
-                    log.info(f"Generating visual diagnostics at step {global_step}...")
-                    log_visual_diagnostics(model, batch, device, cfg, global_step, epoch,
-                                          diagnostics_collector=diag_collector, vis_interval=args.vis_interval)
-
-            preds = logits.argmax(dim=-1)
-            for k, v in losses.items():
-                if isinstance(v, torch.Tensor):
-                    epoch_losses[k].append(v.item())
-            epoch_correct += (preds == batch["labels"]).sum().item()
-            epoch_total += batch["batch_size"]
-
-            if batch_idx % 50 == 0:
-                avg = {k: np.mean(v[-50:]) for k, v in epoch_losses.items()}
-                acc = epoch_correct / max(epoch_total, 1) * 100
-                lrs = scheduler.get_last_lr()
-                log.info(f"P2 E{epoch} B{batch_idx}/{len(train_dl)} | "
-                         f"loss={avg.get('total',0):.4f} jepa={avg.get('jepa',0):.4f} "
-                         f"task={avg.get('task',0):.4f} | acc={acc:.1f}%")
-                trackio.log({
-                    "train/loss": avg.get("total", 0), "train/jepa_loss": avg.get("jepa", 0),
-                    "train/task_loss": avg.get("task", 0), "train/reg_loss": avg.get("reg", 0),
-                    "train/accuracy": acc, "train/lr": lrs[0] if lrs else 0,
-                    "train/backbone_lr": lrs[1] if len(lrs) > 1 else 0,
-                    "train/text_lr": lrs[2] if len(lrs) > 2 else 0,
-                    "train/ema_momentum": model.target.mom,
-                    "train/epoch": epoch, "train/step": global_step,
-                })
-
-        eval_acc = p1.evaluate(model, eval_dl, device, cfg)
-        train_acc = epoch_correct / max(epoch_total, 1) * 100
-        log.info(f"=== Phase 2 Epoch {epoch} | Train: {train_acc:.1f}% | Eval: {eval_acc:.1f}% ===")
-        trackio.log({"eval/accuracy": eval_acc, "eval/epoch": epoch,
-                     "eval/train_accuracy": train_acc, "eval/best_accuracy": max(best_acc, eval_acc)})
-
-        log.info(f"Generating epoch-end visual diagnostics...")
-        diag_batch = next(iter(eval_dl))
-        diag_batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in diag_batch.items()}
-        log_visual_diagnostics(model, diag_batch, device, cfg, global_step, epoch,
-                              diagnostics_collector=diag_collector, vis_interval=args.vis_interval)
-
-        if eval_acc > best_acc:
-            best_acc = eval_acc
-            p1.save_checkpoint(model, cfg, epoch, eval_acc, is_best=True)
-            log.info(f"New best accuracy: {best_acc:.1f}%")
-
-    log.info(f"Phase 2 complete. Best eval accuracy: {best_acc:.1f}%")
-    diag_collector.detach()
-    trackio.log({"final/best_accuracy": best_acc, "final/phase": 2, "final/total_steps": global_step})
+    try:
+        for epoch in range(cfg.epochs):
+            model.train()
+            epoch_losses = defaultdict(list)
+            epoch_correct = 0
+            epoch_total = 0
+            optimizer.zero_grad()
+
+            for batch_idx, batch in enumerate(train_dl):
+                batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in batch.items()}
+                vis_tok = model.vis(batch["pixel_values"]).float()
+                txt_tok = model.txt(batch["input_ids"], batch["attention_mask"]).float()
+
+                with torch.autocast(device_type="cuda", dtype=amp_dtype, enabled=cfg.bf16 and device.type == "cuda"):
+                    evidence, _, ev_mask = model.evidence(vis_tok, txt_tok, batch["attention_mask"])
+                    if model._use_rollout:
+                        traj, z_final, z_proj = model.rollout(evidence)
+                    else:
+                        B = batch["batch_size"]
+                        z0 = model.rollout.init_tokens.expand(B, -1, -1) + \
+                             model.rollout.z0_proj(F.adaptive_avg_pool1d(
+                                 evidence.permute(0,2,1), model.rollout.num_tokens).permute(0,2,1))
+                        z_final = z0
+                        z_proj = model.rollout.out_proj(z0).unsqueeze(1)
+
+                    if model._use_jepa:
+                        target_proj = model.target(vis_tok.detach(), txt_tok.detach(), batch["attention_mask"].detach())
+                    else:
+                        target_proj = None
+
+                    opt_emb = model.encode_options(batch["opt_input_ids"], batch["opt_attention_mask"])
+                    opt_emb = opt_emb.view(batch["batch_size"], cfg.max_options, -1)
+                    logits = model.disc(z_final, opt_emb, batch["opt_mask"])
+                    task_loss = F.cross_entropy(logits, batch["labels"])
+
+                    if model._use_jepa and target_proj is not None:
+                        losses = model.jepa_loss(z_proj, target_proj, task_loss)
+                    else:
+                        losses = {"total": task_loss, "jepa": torch.tensor(0.0), "task": task_loss, "reg": torch.tensor(0.0)}
+                    loss = losses["total"] / cfg.grad_accum
+
+                loss.backward()
+
+                if (batch_idx + 1) % cfg.grad_accum == 0:
+                    nn.utils.clip_grad_norm_(trainable, cfg.max_grad_norm)
+                    optimizer.step(); scheduler.step(); optimizer.zero_grad()
+                    model.update_target(global_step, total_steps)
+                    global_step += 1
+                    if global_step % args.vis_interval == 0 and global_step > 0:
+                        log.info(f"Generating visual diagnostics at step {global_step}...")
+                        log_visual_diagnostics(model, batch, device, cfg, global_step, epoch,
+                                              diagnostics_collector=diag_collector, vis_interval=args.vis_interval)
+
+                preds = logits.argmax(dim=-1)
+                for k, v in losses.items():
+                    if isinstance(v, torch.Tensor):
+                        epoch_losses[k].append(v.item())
+                epoch_correct += (preds == batch["labels"]).sum().item()
+                epoch_total += batch["batch_size"]
+
+                if batch_idx % 50 == 0:
+                    avg = {k: np.mean(v[-50:]) for k, v in epoch_losses.items()}
+                    acc = epoch_correct / max(epoch_total, 1) * 100
+                    lrs = scheduler.get_last_lr()
+                    log.info(f"P2 E{epoch} B{batch_idx}/{len(train_dl)} | "
+                             f"loss={avg.get('total',0):.4f} jepa={avg.get('jepa',0):.4f} "
+                             f"task={avg.get('task',0):.4f} | acc={acc:.1f}%")
+                    trackio.log({
+                        "train/loss": avg.get("total", 0), "train/jepa_loss": avg.get("jepa", 0),
+                        "train/task_loss": avg.get("task", 0), "train/reg_loss": avg.get("reg", 0),
+                        "train/accuracy": acc, "train/lr": lrs[0] if lrs else 0,
+                        "train/backbone_lr": lrs[1] if len(lrs) > 1 else 0,
+                        "train/text_lr": lrs[2] if len(lrs) > 2 else 0,
+                        "train/ema_momentum": model.target.mom,
+                        "train/epoch": epoch, "train/step": global_step,
+                    })
+
+            eval_acc = p1.evaluate(model, eval_dl, device, cfg)
+            train_acc = epoch_correct / max(epoch_total, 1) * 100
+            log.info(f"=== Phase 2 Epoch {epoch} | Train: {train_acc:.1f}% | Eval: {eval_acc:.1f}% ===")
+            trackio.log({"eval/accuracy": eval_acc, "eval/epoch": epoch,
+                         "eval/train_accuracy": train_acc, "eval/best_accuracy": max(best_acc, eval_acc)})
+
+            log.info(f"Generating epoch-end visual diagnostics...")
+            diag_batch = next(iter(eval_dl))
+            diag_batch = {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in diag_batch.items()}
+            log_visual_diagnostics(model, diag_batch, device, cfg, global_step, epoch,
+                                  diagnostics_collector=diag_collector, vis_interval=args.vis_interval)
+
+            if eval_acc > best_acc:
+                best_acc = eval_acc
+                p1.save_checkpoint(model, cfg, epoch, eval_acc, is_best=True)
+                log.info(f"New best accuracy: {best_acc:.1f}%")
+
+        log.info(f"Phase 2 complete. Best eval accuracy: {best_acc:.1f}%")
+
+    finally:
+        # ── Ensure Trackio data is persisted even if training crashes ──
+        diag_collector.detach()
+        trackio.log({"final/best_accuracy": best_acc, "final/phase": 2, "final/total_steps": global_step})
+        log.info("Finishing Trackio and syncing to Space...")
+        trackio.finish()
+        # Belt-and-suspenders: explicit sync to ensure all images are uploaded
+        try:
+            trackio.sync(project="MR-JEPA", space_id=args.trackio_space)
+            log.info(f"Trackio synced to https://huggingface.co/spaces/{args.trackio_space}")
+        except Exception as e:
+            log.warning(f"Trackio sync failed (data may still be available via finish): {e}")
+
     if cfg.push_to_hub:
         p1.push_results(cfg, best_acc)