Added code

code/src/model.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+"""EdgeNeXt builder + transfer-learning regimes (freeze / partial / MONA).
+
+Three fine-tuning regimes are exposed (see PROTOCOL §5):
+* ``full``    — train the whole backbone + head (one or two LR groups).
+* ``partial`` — freeze the first ``4 - n_unfrozen`` stages, train the rest + head;
+                ``n_unfrozen=0`` is pure linear probing (feature extraction).
+* ``mona``    — freeze the whole backbone, train only MONA adapters + head (PEFT).
+
+``feature_extract_features`` returns L2-normalized pooled embeddings for the
+UMAP analysis (see ``umap_analysis.py``).
+"""
+
+from dataclasses import dataclass
+
+import timm
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from .mona import attach_mona_adapters
+
+
+@dataclass
+class ModelConfig:
+    """Backbone + regime configuration (a subset is searched by Optuna)."""
+
+    model_name: str = "edgenext_small.usi_in1k"
+    num_classes: int = 13  # Naruto Sign: 12 seals + 'zero' (verify after download)
+    pretrained: bool = True
+    drop_rate: float = 0.0  # classifier dropout
+    drop_path_rate: float = 0.0  # stochastic depth
+    regime: str = "partial"  # {"full", "partial", "mona"}
+    n_unfrozen_stages: int = 1  # used by "partial": 0..4
+    mona_stages: tuple[int, ...] = (2, 3)  # used by "mona"
+    mona_bottleneck: int = 64
+    mona_kernels: tuple[int, ...] = (3, 5, 7)
+
+
+def _set_norm_eval(module: nn.Module) -> None:
+    """Put frozen BatchNorm/LayerNorm/GroupNorm in eval mode (freeze stats)."""
+    for m in module.modules():
+        if isinstance(m, (nn.BatchNorm2d, nn.LayerNorm, nn.GroupNorm)):
+            m.eval()
+
+
+def build_model(cfg: ModelConfig) -> nn.Module:
+    """Create an EdgeNeXt and apply the requested fine-tuning regime.
+
+    Args:
+        cfg: Model configuration.
+
+    Returns:
+        The configured ``nn.Module``. Trainable parameters depend on ``regime``.
+    """
+    model = timm.create_model(
+        cfg.model_name,
+        pretrained=cfg.pretrained,
+        num_classes=cfg.num_classes,
+        drop_rate=cfg.drop_rate,
+        drop_path_rate=cfg.drop_path_rate,
+    )
+
+    if cfg.regime == "full":
+        return model
+
+    # Freeze the whole backbone first; head stays trainable in every regime.
+    for name, p in model.named_parameters():
+        if not _is_head(name):
+            p.requires_grad_(False)
+
+    if cfg.regime == "partial":
+        n = max(0, min(cfg.n_unfrozen_stages, len(model.stages)))
+        if n > 0:
+            for idx in range(len(model.stages) - n, len(model.stages)):
+                for p in model.stages[idx].parameters():
+                    p.requires_grad_(True)
+        # final norm before the head (if present) follows the head.
+        for name, p in model.named_parameters():
+            if "norm_pre" in name or name.startswith("norm"):
+                p.requires_grad_(True)
+    elif cfg.regime == "mona":
+        attach_mona_adapters(
+            model,
+            list(cfg.mona_stages),
+            bottleneck=cfg.mona_bottleneck,
+            kernels=cfg.mona_kernels,
+        )  # adapters are created trainable by default
+    else:
+        raise ValueError(f"unknown regime: {cfg.regime!r}")
+
+    return model
+
+
+def _is_head(param_name: str) -> bool:
+    """Heuristic: classifier-head parameters in timm models contain 'head'/'fc'."""
+    return "head" in param_name or param_name.startswith("fc")
+
+
+def freeze_consistency(model: nn.Module) -> None:
+    """Set frozen norm layers to eval so their running stats are not updated.
+
+    Call this **after** ``model.train()`` in every training step for the
+    ``partial``/``mona`` regimes (see PROTOCOL §5.2 — the classic BN pitfall).
+    """
+    for module in model.modules():
+        if isinstance(module, (nn.BatchNorm2d, nn.LayerNorm, nn.GroupNorm)):
+            if not any(p.requires_grad for p in module.parameters(recurse=False)):
+                module.eval()
+
+
+def build_param_groups(
+    model: nn.Module, base_lr: float, *, backbone_lr_mult: float = 0.1
+) -> list[dict]:
+    """Two LR groups: a smaller LR for backbone, the base LR for head/adapters.
+
+    Discriminative learning rates (lower layers -> lower LR) are a standard
+    transfer-learning trick (see PROTOCOL §5.3). Only parameters with
+    ``requires_grad=True`` are included.
+
+    Args:
+        model: The configured model.
+        base_lr: LR for the head and MONA adapters.
+        backbone_lr_mult: Multiplier applied to ``base_lr`` for backbone params.
+
+    Returns:
+        A list of optimizer parameter-group dicts.
+    """
+    head, backbone = [], []
+    for name, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        if _is_head(name) or "mona_adapters" in name:
+            head.append(p)
+        else:
+            backbone.append(p)
+
+    groups: list[dict] = []
+    if head:
+        groups.append({"params": head, "lr": base_lr})
+    if backbone:
+        groups.append({"params": backbone, "lr": base_lr * backbone_lr_mult})
+    return groups
+
+
+def count_trainable(model: nn.Module) -> tuple[int, int]:
+    """Return ``(trainable, total)`` parameter counts."""
+    total = sum(p.numel() for p in model.parameters())
+    trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    return trainable, total
+
+
+@torch.inference_mode()
+def extract_features(
+    model: nn.Module, x: torch.Tensor, *, l2_normalize: bool = True
+) -> torch.Tensor:
+    """Return pooled pre-logit embeddings ``[B, num_features]`` for UMAP.
+
+    Uses timm's ``forward_features`` + ``forward_head(..., pre_logits=True)`` so
+    the classifier is bypassed.
+
+    Args:
+        model: A timm model in eval mode.
+        x: Input batch ``[B, 3, H, W]`` already normalized.
+        l2_normalize: If ``True`` L2-normalize embeddings (cosine geometry).
+
+    Returns:
+        Float tensor on CPU of shape ``[B, num_features]``.
+    """
+    feats = model.forward_features(x)
+    emb = model.forward_head(feats, pre_logits=True)
+    if l2_normalize:
+        emb = F.normalize(emb, dim=1)
+    return emb.detach().cpu()