geospatiallabamsob/mantis-vision-api

1"""Train the schema-driven multi-head EfficientNet-B0 seaweed model.2 3Two-phase schedule:4  1. Frozen backbone, train only the heads (fast, stabilizes the new heads5     before touching pretrained weights).6  2. Unfreeze the backbone and fine-tune end-to-end at a lower LR.7 8The loss is multi-task (see src/losses.py) and driven entirely by the active9Schema: one term per measurement (classification/regression/segmentation),10each masked to the samples where it applies and has a value. Adding a new11measurement to the schema needs no change here — it just becomes one more12term in the sum, at 0 loss until labeled data exists for it.13 14Label-noise robustness: label smoothing on the primary (background-carrying)15classification measurement plus heavy augmentation (blur/brightness/noise,16see src/data/transforms.py) is the first-line defense. For heavier noise,17swap that measurement's criterion in src/losses.py for a Generalized18Cross-Entropy / symmetric loss — the rest of the pipeline is unaffected.19 20Usage:21    python -m src.train22"""23from __future__ import annotations24 25import sys26import time27from pathlib import Path28 29import torch30from torch.optim import AdamW31from tqdm import tqdm32 33sys.path.insert(0, str(Path(__file__).resolve().parents[1]))34from config import SCHEMA as _default_schema  # noqa: E40235from config import Config, Schema, config as _default_config  # noqa: E40236from src.data.dataset import get_dataloaders  # noqa: E40237from src.losses import build_criterions, compute_losses  # noqa: E40238from src.models.efficientnet import build_model, save_checkpoint, unfreeze_backbone  # noqa: E40239from src.utils.logger import get_logger  # noqa: E40240from src.utils.seed import get_device, set_seed  # noqa: E40241 42 43def _to_device(targets: dict, device) -> dict:44    return {key: value.to(device) for key, value in targets.items()}45 46 47def run_epoch(model, loader, schema: Schema, criterions, cfg: Config, optimizer, device, train: bool) -> tuple[float, float]:48    model.train() if train else model.eval()49 50    primary = schema.primary_classification()51    total_loss, correct, total = 0.0, 0, 052    context = torch.enable_grad() if train else torch.no_grad()53 54    with context:55        for images, targets in tqdm(loader, leave=False):56            images = images.to(device)57            targets = _to_device(targets, device)58 59            if train:60                optimizer.zero_grad()61 62            outputs = model(images)63            loss, _ = compute_losses(outputs, targets, schema, criterions, cfg)64 65            if train:66                loss.backward()67                optimizer.step()68 69            batch_size = images.size(0)70            total_loss += loss.item() * batch_size71            # Track the primary classification's accuracy as the72            # human-readable progress signal (skipped if the schema declares73            # no such measurement).74            if primary is not None:75                correct += (outputs[primary.key].argmax(1) == targets[f"{primary.key}_id"]).sum().item()76            total += batch_size77 78    accuracy = correct / total if primary is not None else 0.079    return total_loss / total, accuracy80 81 82def train(cfg: Config | None = None, schema: Schema | None = None) -> None:83    """cfg/schema default to the process-wide config.config / config.SCHEMA84    (what every real invocation uses); both are overridable so tests and85    tooling can point at a synthetic Config/Schema without touching global86    state."""87    cfg = cfg if cfg is not None else _default_config88    schema = schema if schema is not None else _default_schema89 90    set_seed(cfg.seed)91    device = get_device(cfg.device)92    logger = get_logger("train", cfg.logs_dir)93    logger.info("Using device: %s", device)94    logger.info("Measurements: %s", [m.key for m in schema.measurements])95 96    data = get_dataloaders(cfg, schema)97 98    model = build_model(schema, freeze_backbone=True)99    model.to(device)100 101    criterions = build_criterions(schema, cfg)102    cfg.checkpoints_dir.mkdir(parents=True, exist_ok=True)103 104    best_val_loss = float("inf")105    epochs_without_improvement = 0106    best_path = cfg.checkpoints_dir / "best_model.pt"107 108    def maybe_early_stop(val_loss: float, epoch_label: str) -> bool:109        nonlocal best_val_loss, epochs_without_improvement110        if best_val_loss - val_loss > cfg.early_stopping_min_delta:111            best_val_loss = val_loss112            epochs_without_improvement = 0113            save_checkpoint(model, schema, best_path, extra={"val_loss": val_loss, "epoch": epoch_label})114            logger.info("New best model saved (val_loss=%.4f) -> %s", val_loss, best_path)115        else:116            epochs_without_improvement += 1117            logger.info("No improvement (%d/%d)", epochs_without_improvement, cfg.early_stopping_patience)118        return epochs_without_improvement >= cfg.early_stopping_patience119 120    # --- Phase 1: frozen backbone ---121    optimizer = AdamW(122        filter(lambda p: p.requires_grad, model.parameters()),123        lr=cfg.frozen_lr,124        weight_decay=cfg.weight_decay,125    )126    logger.info("Phase 1: training heads (backbone frozen)")127    for epoch in range(1, cfg.frozen_epochs + 1):128        start = time.time()129        train_loss, train_acc = run_epoch(model, data.train, schema, criterions, cfg, optimizer, device, train=True)130        val_loss, val_acc = run_epoch(model, data.val, schema, criterions, cfg, optimizer, device, train=False)131        logger.info(132            "[frozen %02d/%02d] train_loss=%.4f acc=%.4f val_loss=%.4f val_acc=%.4f (%.1fs)",133            epoch, cfg.frozen_epochs, train_loss, train_acc, val_loss, val_acc, time.time() - start,134        )135        if maybe_early_stop(val_loss, f"frozen-{epoch}"):136            logger.info("Early stopping during frozen phase.")137            return138 139    # --- Phase 2: fine-tune the whole network ---140    unfreeze_backbone(model)141    optimizer = AdamW(model.parameters(), lr=cfg.finetune_lr, weight_decay=cfg.weight_decay)142    epochs_without_improvement = 0  # reset patience for the new phase143    logger.info("Phase 2: fine-tuning full network (backbone unfrozen)")144    for epoch in range(1, cfg.finetune_epochs + 1):145        start = time.time()146        train_loss, train_acc = run_epoch(model, data.train, schema, criterions, cfg, optimizer, device, train=True)147        val_loss, val_acc = run_epoch(model, data.val, schema, criterions, cfg, optimizer, device, train=False)148        logger.info(149            "[finetune %02d/%02d] train_loss=%.4f acc=%.4f val_loss=%.4f val_acc=%.4f (%.1fs)",150            epoch, cfg.finetune_epochs, train_loss, train_acc, val_loss, val_acc, time.time() - start,151        )152        if maybe_early_stop(val_loss, f"finetune-{epoch}"):153            logger.info("Early stopping during fine-tune phase.")154            return155 156    logger.info("Training complete. Best val_loss=%.4f at %s", best_val_loss, best_path)157 158 159if __name__ == "__main__":160    train()161