Spaces:
Sleeping
Sleeping
| """EpisodicTrainer — trains AppearanceScorer and ChangeScorer jointly. | |
| The training loop: | |
| • Loads prebuilt corpus prototypes from ``checkpoints/prototypes.pt``. | |
| • Feeds (tile_2022, tile_2024, mask_2022, mask_2024) batches to the detector. | |
| • Computes combined BCE + Dice loss on ``prob_map`` vs ``mask_2024``. | |
| • Validates every epoch; saves best checkpoint by validation IoU. | |
| """ | |
| from __future__ import annotations | |
| import logging | |
| from pathlib import Path | |
| from typing import Optional | |
| import torch | |
| import torch.nn as nn | |
| from omegaconf import DictConfig | |
| from torch.optim import AdamW | |
| from torch.optim.lr_scheduler import CosineAnnealingLR, LinearLR, SequentialLR | |
| from torch.utils.data import DataLoader | |
| from torch.utils.tensorboard import SummaryWriter | |
| from data.dataset import ConstructionChangeDataset | |
| from evaluation.metrics import SegmentationEvaluator | |
| from losses.segmentation_losses import CombinedSegmentationLoss | |
| from models.detector import ConstructionChangeDetector | |
| logger = logging.getLogger(__name__) | |
| class EpisodicTrainer: | |
| def __init__(self, cfg: DictConfig) -> None: | |
| self.cfg = cfg | |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| logger.info("Device: %s", self.device) | |
| # --- Model --- | |
| self.model = ConstructionChangeDetector( | |
| k_prototypes=cfg.model.k_prototypes | |
| ).to(self.device) | |
| self.model.load_sam2(cfg.model.sam2_hf_id) | |
| # --- Prototypes --- | |
| proto_path = Path(cfg.training.checkpoint_dir) / "prototypes.pt" | |
| if not proto_path.exists(): | |
| raise FileNotFoundError( | |
| f"Prototypes not found at {proto_path}. " | |
| "Run scripts/build_prototypes.py first." | |
| ) | |
| ckpt = torch.load(proto_path, map_location="cpu") | |
| self.prototypes: torch.Tensor = ckpt["prototypes"].to(self.device) | |
| logger.info("Loaded %d prototypes.", self.prototypes.shape[0]) | |
| # --- Loss --- | |
| self.criterion = CombinedSegmentationLoss( | |
| bce_weight=cfg.training.loss_bce_weight, | |
| dice_weight=cfg.training.loss_dice_weight, | |
| ) | |
| # --- Optimizer (only trainable heads) --- | |
| trainable = list(self.model.appearance_scorer.parameters()) + \ | |
| list(self.model.change_scorer.parameters()) | |
| self.optimizer = AdamW( | |
| trainable, | |
| lr=cfg.training.lr, | |
| weight_decay=cfg.training.weight_decay, | |
| ) | |
| # --- LR schedule: linear warmup → cosine annealing --- | |
| warmup_steps = cfg.training.warmup_epochs | |
| total_steps = cfg.training.epochs | |
| warmup_sched = LinearLR( | |
| self.optimizer, start_factor=0.1, total_iters=warmup_steps | |
| ) | |
| cosine_sched = CosineAnnealingLR( | |
| self.optimizer, T_max=total_steps - warmup_steps, eta_min=1e-6 | |
| ) | |
| self.scheduler = SequentialLR( | |
| self.optimizer, | |
| schedulers=[warmup_sched, cosine_sched], | |
| milestones=[warmup_steps], | |
| ) | |
| # --- AMP --- | |
| self.scaler = torch.cuda.amp.GradScaler(enabled=cfg.training.amp) | |
| # --- Data --- | |
| self.train_loader, self.val_loader = self._build_dataloaders() | |
| # --- Logging --- | |
| log_dir = Path(cfg.training.log_dir) | |
| log_dir.mkdir(parents=True, exist_ok=True) | |
| self.writer = SummaryWriter(str(log_dir)) | |
| self.ckpt_dir = Path(cfg.training.checkpoint_dir) | |
| self.ckpt_dir.mkdir(parents=True, exist_ok=True) | |
| self.best_iou = 0.0 | |
| self.log_every = cfg.training.log_every | |
| # ------------------------------------------------------------------ | |
| def _build_dataloaders(self): | |
| common = dict( | |
| tile_dir=self.cfg.data.tile_dir, | |
| image_size=self.cfg.data.image_size, | |
| min_fg_pixels=self.cfg.data.min_fg_pixels, | |
| val_split=self.cfg.training.val_split, | |
| seed=self.cfg.training.seed, | |
| ) | |
| train_ds = ConstructionChangeDataset(augment=True, split="train", **common) | |
| val_ds = ConstructionChangeDataset(augment=False, split="val", **common) | |
| train_loader = DataLoader( | |
| train_ds, | |
| batch_size=self.cfg.training.batch_size, | |
| shuffle=True, | |
| num_workers=self.cfg.data.num_workers, | |
| pin_memory=True, | |
| drop_last=True, | |
| ) | |
| val_loader = DataLoader( | |
| val_ds, | |
| batch_size=self.cfg.training.batch_size, | |
| shuffle=False, | |
| num_workers=self.cfg.data.num_workers, | |
| pin_memory=True, | |
| ) | |
| return train_loader, val_loader | |
| # ------------------------------------------------------------------ | |
| def train_epoch(self, epoch: int) -> float: | |
| self.model.train() | |
| # Keep encoder in eval mode regardless | |
| if self.model.image_encoder is not None: | |
| self.model.image_encoder.eval() | |
| total_loss = 0.0 | |
| global_step = epoch * len(self.train_loader) | |
| for step, batch in enumerate(self.train_loader): | |
| tile_2022 = batch["tile_2022"].to(self.device) | |
| tile_2024 = batch["tile_2024"].to(self.device) | |
| mask_2022 = batch["mask_2022"].to(self.device) | |
| mask_2024 = batch["mask_2024"].to(self.device) | |
| self.optimizer.zero_grad(set_to_none=True) | |
| with torch.cuda.amp.autocast(enabled=self.cfg.training.amp): | |
| out = self.model(tile_2022, tile_2024, mask_2022, self.prototypes) | |
| loss_dict = self.criterion(out["prob_map"], mask_2024) | |
| self.scaler.scale(loss_dict["loss"]).backward() | |
| self.scaler.unscale_(self.optimizer) | |
| nn.utils.clip_grad_norm_( | |
| list(self.model.appearance_scorer.parameters()) + | |
| list(self.model.change_scorer.parameters()), | |
| max_norm=1.0, | |
| ) | |
| self.scaler.step(self.optimizer) | |
| self.scaler.update() | |
| total_loss += loss_dict["loss"].item() | |
| if step % self.log_every == 0: | |
| gstep = global_step + step | |
| self.writer.add_scalar("train/loss", loss_dict["loss"].item(), gstep) | |
| self.writer.add_scalar("train/bce_loss", loss_dict["bce_loss"].item(), gstep) | |
| self.writer.add_scalar("train/dice_loss", loss_dict["dice_loss"].item(), gstep) | |
| return total_loss / max(len(self.train_loader), 1) | |
| def validate(self, epoch: int) -> dict: | |
| self.model.eval() | |
| evaluator = SegmentationEvaluator(threshold=0.5) | |
| for batch in self.val_loader: | |
| tile_2022 = batch["tile_2022"].to(self.device) | |
| tile_2024 = batch["tile_2024"].to(self.device) | |
| mask_2022 = batch["mask_2022"].to(self.device) | |
| mask_2024 = batch["mask_2024"].to(self.device) | |
| with torch.cuda.amp.autocast(enabled=self.cfg.training.amp): | |
| out = self.model(tile_2022, tile_2024, mask_2022, self.prototypes) | |
| evaluator.update(out["prob_map"].cpu(), mask_2024.cpu()) | |
| metrics = evaluator.compute() | |
| for k, v in metrics.items(): | |
| self.writer.add_scalar(f"val/{k}", v, epoch) | |
| return metrics | |
| def save_checkpoint(self, epoch: int, is_best: bool = False) -> None: | |
| state = { | |
| "epoch": epoch, | |
| "best_iou": self.best_iou, | |
| **self.model.state_dict_trainable(), | |
| "optimizer": self.optimizer.state_dict(), | |
| } | |
| name = "best.pt" if is_best else f"epoch_{epoch:03d}.pt" | |
| torch.save(state, self.ckpt_dir / name) | |
| def load_checkpoint(self, path: str) -> int: | |
| state = torch.load(path, map_location=self.device) | |
| self.model.load_trainable_state_dict(state) | |
| self.best_iou = state.get("best_iou", 0.0) | |
| return state.get("epoch", 0) | |
| # ------------------------------------------------------------------ | |
| def train(self) -> None: | |
| logger.info("Starting training for %d epochs …", self.cfg.training.epochs) | |
| for epoch in range(self.cfg.training.epochs): | |
| train_loss = self.train_epoch(epoch) | |
| metrics = self.validate(epoch) | |
| self.scheduler.step() | |
| logger.info( | |
| "Epoch %03d | loss=%.4f | iou=%.4f | f1=%.4f | " | |
| "precision=%.4f | recall=%.4f", | |
| epoch, train_loss, | |
| metrics["iou"], metrics["f1"], | |
| metrics["precision"], metrics["recall"], | |
| ) | |
| if metrics["iou"] > self.best_iou: | |
| self.best_iou = metrics["iou"] | |
| self.save_checkpoint(epoch, is_best=True) | |
| logger.info(" ↑ New best IoU=%.4f — saved best.pt", self.best_iou) | |
| if (epoch + 1) % self.cfg.training.checkpoint_every == 0: | |
| self.save_checkpoint(epoch, is_best=False) | |
| self.writer.close() | |
| logger.info("Training complete. Best IoU=%.4f", self.best_iou) | |