Add model save/load and --resume for training continuation

- AsymmetricEncoder.save_checkpoint(): saves model_state + metadata
- AsymmetricEncoder.load_checkpoint(): rebuilds model with frozen backbones,
  then loads trainable weights from checkpoint
- --resume flag restores optimizer, loss (learnable tau), and scheduler state
- Training continues from the saved epoch

Usage:
  python -m src.training.train_gtauav --resume out/gtauav/with_text/ckpt_epoch004.pt

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

src/models/asymmetric_encoder.py

@@ -527,6 +527,56 @@ class AsymmetricEncoder(nn.Module):
         """Return list of parameters that require grad."""
         return [p for p in self.parameters() if p.requires_grad]
 
+    def save_checkpoint(self, path: str | Path, **extra) -> None:
+        """Save model checkpoint with metadata."""
+        path = Path(path)
+        path.parent.mkdir(parents=True, exist_ok=True)
+        ckpt = {
+            "model_state": self.state_dict(),
+            "proj_dim": self.proj_dim,
+            "baseline_mode": self.baseline_mode,
+            **extra,
+        }
+        tmp = path.with_suffix(path.suffix + ".tmp")
+        torch.save(ckpt, tmp)
+        tmp.replace(path)
+        LOGGER.info("💾 Model saved to %s", path)
+
+    @classmethod
+    def load_checkpoint(
+        cls,
+        path: str | Path,
+        dino_web_path: str = "nn_models/DINO_WEB/dinov3-vitl16-pretrain-lvd1689m.pth",
+        dino_sat_path: str = "nn_models/DINO_SAT/model.safetensors",
+        lrsclip_path: str = "nn_models/LRSCLIP/DGTRS-CLIP-ViT-L-14.pt",
+        device: str = "cuda",
+    ) -> tuple[AsymmetricEncoder, dict]:
+        """Load model from checkpoint.
+
+        First builds the model (loading frozen backbones), then loads
+        the saved trainable weights on top.
+
+        Returns:
+            (model, checkpoint_dict) — model ready for eval/resume,
+            checkpoint_dict has optimizer_state, epoch, etc.
+        """
+        path = Path(path)
+        LOGGER.info("📂 Loading checkpoint from %s", path)
+        ckpt = torch.load(str(path), map_location="cpu", weights_only=False)
+
+        model = cls(
+            dino_web_path=dino_web_path,
+            dino_sat_path=dino_sat_path,
+            lrsclip_path=lrsclip_path,
+            proj_dim=ckpt.get("proj_dim", 512),
+            baseline_mode=ckpt.get("baseline_mode", False),
+            device=device,
+        )
+        model.load_state_dict(ckpt["model_state"], strict=False)
+        model = model.to(device)
+        LOGGER.info("✅ Checkpoint loaded (epoch=%s)", ckpt.get("epoch", "?"))
+        return model, ckpt
+
     def train(self, mode: bool = True) -> AsymmetricEncoder:
         """Override to keep frozen encoders in eval mode."""
         super().train(mode)

src/training/train_gtauav.py

@@ -65,6 +65,7 @@ class TrainConfigGTAUAV:
     baseline_mode: bool = False
 
     # Training.
+    resume_from: str | None = None  # path to checkpoint for resuming
     output_dir: str = "out/gtauav/with_text"
     epochs: int = 10
     batch_size: int = 64
@@ -211,6 +212,20 @@ def train(cfg: TrainConfigGTAUAV) -> None:
         json.dump(vars(cfg), f, indent=2)
 
     # Model.
+    start_epoch = 0
+    resume_ckpt = None
+
+    if cfg.resume_from is not None:
+        LOGGER.info("🔄 Resuming from %s", cfg.resume_from)
+        model, resume_ckpt = AsymmetricEncoder.load_checkpoint(
+            cfg.resume_from,
+            dino_web_path=cfg.dino_web_path,
+            dino_sat_path=cfg.dino_sat_path,
+            lrsclip_path=cfg.lrsclip_path,
+            device=cfg.device,
+        )
+        start_epoch = resume_ckpt.get("epoch", -1) + 1
+    else:
         mode_str = "🚫 baseline (no text)" if cfg.baseline_mode else "📝 with text (L1/L2/L3)"
         LOGGER.info("🏗️  Building model — %s", mode_str)
         model = AsymmetricEncoder(
@@ -308,11 +323,24 @@ def train(cfg: TrainConfigGTAUAV) -> None:
     )
     scaler = GradScaler(enabled=cfg.use_amp)
 
+    # Restore optimizer/scheduler/loss state on resume.
+    if resume_ckpt is not None:
+        if "optimizer_state" in resume_ckpt:
+            optimizer.load_state_dict(resume_ckpt["optimizer_state"])
+            LOGGER.info("🔄 Optimizer state restored")
+        if "loss_state" in resume_ckpt:
+            loss_fn.load_state_dict(resume_ckpt["loss_state"])
+            LOGGER.info("🔄 Loss state restored (tau=%.4f)", loss_fn.current_temperature)
+        # Advance scheduler to the correct step.
+        for _ in range(start_epoch * steps_per_epoch):
+            scheduler.step()
+        LOGGER.info("🔄 Resuming from epoch %d", start_epoch)
+
     history: list[dict] = []
 
-    LOGGER.info("🚀 Starting training for %d epochs", cfg.epochs)
+    LOGGER.info("🚀 Starting training for %d epochs (from epoch %d)", cfg.epochs, start_epoch)
 
-    for epoch in range(cfg.epochs):
+    for epoch in range(start_epoch, cfg.epochs):
         model.train()
         epoch_start = time.time()
         agg: dict[str, float] = {}
@@ -448,6 +476,10 @@ def main() -> None:
         "--baseline", action="store_true",
         help="Run baseline mode (no text).",
     )
+    parser.add_argument(
+        "--resume", type=str, default=None,
+        help="Path to checkpoint to resume training from.",
+    )
     parser.add_argument(
         "--output-dir", type=str, default=None,
         help="Override output directory.",
@@ -484,6 +516,7 @@ def main() -> None:
 
     cfg = TrainConfigGTAUAV()
     cfg.baseline_mode = args.baseline
+    cfg.resume_from = args.resume
     cfg.batch_size = args.batch_size
     cfg.epochs = args.epochs
     cfg.learning_rate = args.lr