From 44bce3096c45f9cb80b922925d1f44582d0e77d0 Mon Sep 17 00:00:00 2001 From: pikaliov Date: Tue, 21 Apr 2026 18:14:54 +0300 Subject: [PATCH] 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) --- src/models/asymmetric_encoder.py | 50 +++++++++++++++++++++++++++ src/training/train_gtauav.py | 59 +++++++++++++++++++++++++------- 2 files changed, 96 insertions(+), 13 deletions(-) diff --git a/src/models/asymmetric_encoder.py b/src/models/asymmetric_encoder.py index 29091d8..8b0077e 100644 --- a/src/models/asymmetric_encoder.py +++ b/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) diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 6c7665b..4ae7af4 100644 --- a/src/training/train_gtauav.py +++ b/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,17 +212,31 @@ def train(cfg: TrainConfigGTAUAV) -> None: json.dump(vars(cfg), f, indent=2) # Model. - mode_str = "🚫 baseline (no text)" if cfg.baseline_mode else "📝 with text (L1/L2/L3)" - LOGGER.info("🏗️ Building model — %s", mode_str) - model = AsymmetricEncoder( - dino_web_path=cfg.dino_web_path, - dino_sat_path=cfg.dino_sat_path, - lrsclip_path=cfg.lrsclip_path, - proj_dim=cfg.proj_dim, - init_gate=cfg.init_gate, - baseline_mode=cfg.baseline_mode, - device=cfg.device, - ).to(cfg.device) + 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( + dino_web_path=cfg.dino_web_path, + dino_sat_path=cfg.dino_sat_path, + lrsclip_path=cfg.lrsclip_path, + proj_dim=cfg.proj_dim, + init_gate=cfg.init_gate, + baseline_mode=cfg.baseline_mode, + device=cfg.device, + ).to(cfg.device) n_trainable = sum(p.numel() for p in model.trainable_parameters()) n_total = sum(p.numel() for p in model.parameters()) @@ -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