Added code

code/README.md

@@ -0,0 +1,91 @@
+# EdgeNeXt × Optuna × Naruto Sign — код HPO-эксперимента
+
+Учебно-исследовательский код к методичке [`../PROTOCOL_HPO_EdgeNeXt_NarutoSign.md`](../PROTOCOL_HPO_EdgeNeXt_NarutoSign.md)
+(научный руководитель — мнс Павленко Б.В.). Задача: подбор гиперпараметров дообучения
+компактного энкодера **EdgeNeXt** на наборе **Naruto Sign** (классификация ручных печатей).
+
+## Установка
+
+```bash
+# Python 3.10–3.12 (PyTorch ещё не собран под 3.14)
+python -m venv .venv && . .venv/Scripts/activate   # Windows: .venv\Scripts\activate
+pip install -r requirements.txt
+```
+
+## Данные
+
+```python
+import kagglehub
+path = kagglehub.dataset_download("vikranthkanumuru/naruto-hand-sign-dataset")
+print(path)   # внутри — папки-классы (ImageFolder)
+```
+
+Сразу посчитать реальную статистику (заполнить таблицу в STATS.md):
+
+```bash
+python -m src.dataset_stats --data-root <PATH> --out results/dataset_stats.json
+```
+
+## Запуск экспериментов
+
+```bash
+# B0 — baseline (один честный прогон, test + confusion matrix)
+python -m src.run_baseline --data-root <PATH> --regime partial --n-unfrozen 1 --epochs 30
+
+# E1 — ablation режимов (повторить для full / partial(0..4) / mona)
+python -m src.run_baseline --data-root <PATH> --regime full   --epochs 30 --out results/e1_full
+python -m src.run_baseline --data-root <PATH> --regime partial --n-unfrozen 0 --out results/e1_lp
+python -m src.run_baseline --data-root <PATH> --regime mona   --epochs 30 --out results/e1_mona
+
+# E2 — Optuna single-objective (TPE + median pruning), persist to sqlite
+python -m src.optuna_search --data-root <PATH> --sampler tpe --pruner median \
+    --n-trials 60 --epochs 25 --study-name nss_v1 --storage
+
+# E3 — Optuna multi-objective (macro-F1 ↑ vs trainable params ↓)
+python -m src.optuna_search --data-root <PATH> --multi-objective --sampler nsga \
+    --n-trials 80 --study-name nss_mo
+
+# E4 — UMAP анализ признаков (до дообучения и с чекпойнтом — после)
+python -m src.umap_analysis --data-root <PATH> --split val --out results/umap_pretrained
+python -m src.umap_analysis --data-root <PATH> --checkpoint <best.pt> --out results/umap_finetuned
+
+# E5 (опц.) — контроль метода: random vs TPE при равном бюджете
+python -m src.optuna_search --data-root <PATH> --sampler random --n-trials 60 --study-name nss_random
+```
+
+Все скрипты — модули пакета `src`, запускать из папки `code/` через `python -m src.<name>`.
+
+## Визуализация Optuna
+
+```python
+import optuna
+from optuna.visualization import (plot_optimization_history,
+                                  plot_param_importances, plot_pareto_front)
+study = optuna.load_study(study_name="nss_v1", storage="sqlite:///nss_v1.db")
+plot_optimization_history(study).show()
+plot_param_importances(study).show()          # fANOVA — какие гиперпараметры важны
+```
+
+## Структура
+
+| Модуль | Назначение |
+|:--|:--|
+| `src/data.py` | ImageFolder, стратиф. split, transforms (hflip OFF), sampler, class weights |
+| `src/model.py` | EdgeNeXt (timm), режимы full/partial/mona, freeze, param-groups, фичи для UMAP |
+| `src/mona.py` | Conv-MONA адаптер (по `Leiyi-Hu/mona`), вставка hook'ом |
+| `src/losses.py` | CE / label-smoothing / weighted / Focal / effective-number weights |
+| `src/metrics.py` | macro-F1, balanced acc, top-k, MCC, κ, confusion (sklearn) |
+| `src/train.py` | train/eval, mixup, early-stop, Optuna pruning, VRAM-hygiene |
+| `src/optuna_search.py` | пространство поиска, single/multi-objective, sampler+pruner |
+| `src/umap_analysis.py` | эмбеддинги → UMAP 2D + кластеризация + ARI/NMI |
+| `src/run_baseline.py` | одиночный прогон + test + confusion-heatmap |
+| `src/dataset_stats.py` | статистика датасета |
+
+## Заметки по воспроизводимости
+
+- seed=42 по умолчанию (`src/train.py::set_seed`); финальные сравнения — 3 seed (42/123/456).
+- test трогать **один раз** в конце; HPO — только по val (macro-F1).
+- `lr`/`weight_decay` ищутся в **log**-шкале.
+- `RandomHorizontalFlip` **отключён** (печати чувствительны к лево/право).
+- Не включать `weighted_ce/focal` одновременно с `--weighted-sampler` на полную силу (двойная компенсация).
+- Главный риск Naruto Sign — **утечка кадров из одного видео** в train и test: если в именах файлов есть id видео, использовать `StratifiedGroupKFold` (см. методичку §3). Базовый `data.py` делает per-frame split — отметить риск в REPORT.md.

code/requirements.txt

@@ -0,0 +1,16 @@
+# EdgeNeXt + Optuna HPO on Naruto Sign — pinned-ish minimums.
+# Use a Python 3.10–3.12 env (PyTorch wheels are not yet built for 3.14).
+torch>=2.2
+torchvision>=0.17
+timm>=1.0.7            # EdgeNeXt weights + Mixup/SoftTargetCrossEntropy
+optuna>=4.5           # AutoSampler / multi-objective / pruners
+optunahub>=0.2        # AutoSampler lives here (load_module("samplers/auto_sampler"))
+umap-learn>=0.5.5
+scikit-learn>=1.3
+hdbscan>=0.8.33       # optional, density clustering on UMAP mid-dim
+matplotlib>=3.7
+numpy>=1.24
+pillow>=10.0
+pandas>=2.0
+kagglehub>=0.3        # optional, dataset download
+plotly>=5.18          # optional, optuna.visualization interactive plots

code/src/__init__.py

@@ -0,0 +1 @@
+"""EdgeNeXt + Optuna HPO on Naruto Sign — student research package."""

code/src/data.py

@@ -0,0 +1,247 @@
+from __future__ import annotations
+
+"""Naruto Sign data pipeline: stratified split, transforms, imbalance-aware sampling.
+
+The dataset is expected as an ``ImageFolder`` tree::
+
+    data_root/
+        Bird/   img001.png ...
+        Boar/   ...
+        ...
+
+If the dataset already ships a ``train/`` and ``test/`` split, point ``data_root``
+at the parent and set ``use_predefined_split=True``; otherwise a single folder is
+split here in a *stratified*, *seeded* way so the protocol stays reproducible.
+
+Design notes (see PROTOCOL §6, §10):
+* ``hflip`` is **off by default** — hand seals are left/right sensitive, a mirror
+  flip can turn a valid seal into an invalid / different one.
+* Normalization mean/std come from the timm pretrained config so the frozen
+  EdgeNeXt encoder sees inputs in the distribution it was trained on.
+"""
+
+import os
+from collections import Counter
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Subset, WeightedRandomSampler
+from torchvision import transforms
+from torchvision.datasets import ImageFolder
+
+
+@dataclass
+class AugConfig:
+    """Augmentation hyper-parameters (a subset is searched by Optuna)."""
+
+    img_size: int = 256
+    use_hflip: bool = False  # OFF by default: seals are chirality-sensitive.
+    rotation_deg: float = 15.0
+    color_jitter: float = 0.2
+    rrc_scale_min: float = 0.7  # RandomResizedCrop lower scale bound.
+    randaug_magnitude: int = 0  # 0 disables RandAugment.
+    mean: tuple[float, float, float] = (0.485, 0.456, 0.406)
+    std: tuple[float, float, float] = (0.229, 0.224, 0.225)
+
+
+@dataclass
+class DataConfig:
+    """Everything needed to materialize the dataloaders."""
+
+    data_root: str
+    use_predefined_split: bool = False
+    val_frac: float = 0.15
+    test_frac: float = 0.15
+    batch_size: int = 32
+    num_workers: int = 4
+    seed: int = 42
+    use_weighted_sampler: bool = False
+    aug: AugConfig = field(default_factory=AugConfig)
+
+
+def build_transforms(aug: AugConfig, *, train: bool) -> transforms.Compose:
+    """Build train/eval transforms for EdgeNeXt-style input.
+
+    Args:
+        aug: Augmentation configuration.
+        train: If ``True`` return the stochastic train pipeline, else the
+            deterministic eval pipeline (resize -> center-crop).
+
+    Returns:
+        A composed torchvision transform.
+    """
+    if train:
+        ops: list = [
+            transforms.RandomResizedCrop(
+                aug.img_size, scale=(aug.rrc_scale_min, 1.0)
+            )
+        ]
+        if aug.randaug_magnitude > 0:
+            ops.append(transforms.RandAugment(magnitude=aug.randaug_magnitude))
+        if aug.rotation_deg > 0:
+            ops.append(transforms.RandomRotation(aug.rotation_deg))
+        if aug.color_jitter > 0:
+            ops.append(
+                transforms.ColorJitter(
+                    brightness=aug.color_jitter,
+                    contrast=aug.color_jitter,
+                    saturation=aug.color_jitter,
+                )
+            )
+        if aug.use_hflip:
+            ops.append(transforms.RandomHorizontalFlip())
+        ops += [transforms.ToTensor(), transforms.Normalize(aug.mean, aug.std)]
+        return transforms.Compose(ops)
+
+    resize = int(round(aug.img_size * 1.14))  # standard resize->center-crop ratio
+    return transforms.Compose(
+        [
+            transforms.Resize(resize),
+            transforms.CenterCrop(aug.img_size),
+            transforms.ToTensor(),
+            transforms.Normalize(aug.mean, aug.std),
+        ]
+    )
+
+
+def _stratified_indices(
+    targets: list[int], val_frac: float, test_frac: float, seed: int
+) -> tuple[list[int], list[int], list[int]]:
+    """Return stratified train/val/test index lists (per-class shuffle+slice)."""
+    rng = np.random.default_rng(seed)
+    by_class: dict[int, list[int]] = {}
+    for idx, y in enumerate(targets):
+        by_class.setdefault(y, []).append(idx)
+
+    train_idx, val_idx, test_idx = [], [], []
+    for _, idxs in sorted(by_class.items()):
+        idxs = np.array(idxs)
+        rng.shuffle(idxs)
+        n = len(idxs)
+        # Force >=1 only when the fraction is positive (frac=0 -> exactly 0).
+        n_test = max(1, int(round(n * test_frac))) if test_frac > 0 else 0
+        n_val = max(1, int(round(n * val_frac))) if val_frac > 0 else 0
+        test_idx += idxs[:n_test].tolist()
+        val_idx += idxs[n_test : n_test + n_val].tolist()
+        train_idx += idxs[n_test + n_val :].tolist()
+    return train_idx, val_idx, test_idx
+
+
+def _make_sampler(targets: list[int], num_classes: int) -> WeightedRandomSampler:
+    """Inverse-frequency WeightedRandomSampler (one sample weight per element)."""
+    counts = Counter(targets)
+    class_w = {c: 1.0 / max(1, counts.get(c, 0)) for c in range(num_classes)}
+    weights = torch.tensor([class_w[t] for t in targets], dtype=torch.double)
+    return WeightedRandomSampler(weights, num_samples=len(weights), replacement=True)
+
+
+def compute_class_weights(targets: list[int], num_classes: int) -> torch.Tensor:
+    """Normalized inverse-frequency weights for weighted cross-entropy.
+
+    Args:
+        targets: Per-sample integer labels of the *train* split.
+        num_classes: Total number of classes.
+
+    Returns:
+        Float tensor of shape ``[num_classes]`` averaging to ~1.0.
+    """
+    counts = Counter(targets)
+    freq = np.array([counts.get(c, 0) for c in range(num_classes)], dtype=np.float64)
+    freq = np.clip(freq, 1.0, None)
+    w = freq.sum() / (num_classes * freq)
+    return torch.tensor(w, dtype=torch.float32)
+
+
+@dataclass
+class DataBundle:
+    """Container returned by :func:`build_dataloaders`."""
+
+    train_loader: DataLoader
+    val_loader: DataLoader
+    test_loader: DataLoader
+    class_names: list[str]
+    train_targets: list[int]
+    num_classes: int
+
+
+def build_dataloaders(cfg: DataConfig) -> DataBundle:
+    """Build train/val/test dataloaders with reproducible stratified split.
+
+    Args:
+        cfg: Data configuration (paths, split fractions, batch size, sampler).
+
+    Returns:
+        A :class:`DataBundle` with loaders, class names and train targets.
+    """
+    train_tf = build_transforms(cfg.aug, train=True)
+    eval_tf = build_transforms(cfg.aug, train=False)
+
+    if cfg.use_predefined_split:
+        # The Kaggle Naruto Sign download ships train/ and test/ but NO val/.
+        # Never alias val = test (that leaks test into model selection): if val/ is
+        # absent, carve a stratified val OUT of train (seeded).
+        test_ds = ImageFolder(os.path.join(cfg.data_root, "test"), eval_tf)
+        class_names = test_ds.classes
+        val_dir = os.path.join(cfg.data_root, "val")
+        if os.path.isdir(val_dir):
+            train_ds = ImageFolder(os.path.join(cfg.data_root, "train"), train_tf)
+            val_ds = ImageFolder(val_dir, eval_tf)
+            train_targets = list(train_ds.targets)
+        else:
+            base_train = ImageFolder(os.path.join(cfg.data_root, "train"), train_tf)
+            base_eval = ImageFolder(os.path.join(cfg.data_root, "train"), eval_tf)
+            val_frac = cfg.val_frac / (1.0 - cfg.test_frac)  # val share of train
+            tr, va, _ = _stratified_indices(base_train.targets, val_frac, 0.0, cfg.seed)
+            train_ds = Subset(base_train, tr)
+            val_ds = Subset(base_eval, va)
+            train_targets = [base_train.targets[i] for i in tr]
+    else:
+        base_train = ImageFolder(cfg.data_root, train_tf)
+        base_eval = ImageFolder(cfg.data_root, eval_tf)
+        class_names = base_train.classes
+        tr, va, te = _stratified_indices(
+            base_train.targets, cfg.val_frac, cfg.test_frac, cfg.seed
+        )
+        train_ds = Subset(base_train, tr)
+        val_ds = Subset(base_eval, va)
+        test_ds = Subset(base_eval, te)
+        train_targets = [base_train.targets[i] for i in tr]
+
+    num_classes = len(class_names)
+
+    if cfg.use_weighted_sampler:
+        sampler = _make_sampler(train_targets, num_classes)
+        train_loader = DataLoader(
+            train_ds,
+            batch_size=cfg.batch_size,
+            sampler=sampler,
+            num_workers=cfg.num_workers,
+            pin_memory=True,
+            drop_last=False,
+        )
+    else:
+        train_loader = DataLoader(
+            train_ds,
+            batch_size=cfg.batch_size,
+            shuffle=True,
+            num_workers=cfg.num_workers,
+            pin_memory=True,
+            drop_last=False,
+        )
+
+    val_loader = DataLoader(
+        val_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=cfg.num_workers
+    )
+    test_loader = DataLoader(
+        test_ds, batch_size=cfg.batch_size, shuffle=False, num_workers=cfg.num_workers
+    )
+
+    return DataBundle(
+        train_loader=train_loader,
+        val_loader=val_loader,
+        test_loader=test_loader,
+        class_names=class_names,
+        train_targets=train_targets,
+        num_classes=num_classes,
+    )

code/src/dataset_stats.py

@@ -0,0 +1,86 @@
+from __future__ import annotations
+
+"""Compute and dump Naruto Sign dataset statistics (PROTOCOL §3, step 1).
+
+Run this right after downloading the data to fill the real numbers (counts per
+class, image-size distribution, imbalance ratio) into the protocol's statistics
+table — those numbers must come from the actual download, not from assumptions.
+"""
+
+import argparse
+import json
+import os
+from collections import Counter
+
+from PIL import Image
+
+
+def scan_imagefolder(root: str) -> dict:
+    """Scan an ImageFolder-style tree and return per-class statistics.
+
+    Args:
+        root: Dataset root containing one sub-directory per class.
+
+    Returns:
+        Dict with class counts, totals, imbalance ratio and image-size stats.
+    """
+    exts = {".png", ".jpg", ".jpeg", ".bmp", ".webp"}
+    classes = sorted(
+        d for d in os.listdir(root) if os.path.isdir(os.path.join(root, d))
+    )
+    counts: Counter = Counter()
+    widths, heights = [], []
+    for cls in classes:
+        cls_dir = os.path.join(root, cls)
+        for fn in os.listdir(cls_dir):
+            if os.path.splitext(fn)[1].lower() not in exts:
+                continue
+            counts[cls] += 1
+            try:
+                with Image.open(os.path.join(cls_dir, fn)) as im:
+                    widths.append(im.width)
+                    heights.append(im.height)
+            except OSError:
+                pass
+
+    total = sum(counts.values())
+    per_class = {c: counts.get(c, 0) for c in classes}
+    nonzero = [v for v in per_class.values() if v > 0]
+    imbalance = (max(nonzero) / min(nonzero)) if nonzero else float("nan")
+
+    def _stats(xs: list[int]) -> dict:
+        if not xs:
+            return {}
+        xs_sorted = sorted(xs)
+        return {
+            "min": xs_sorted[0],
+            "median": xs_sorted[len(xs_sorted) // 2],
+            "max": xs_sorted[-1],
+        }
+
+    return {
+        "num_classes": len(classes),
+        "total_images": total,
+        "per_class": per_class,
+        "imbalance_ratio_max_over_min": imbalance,
+        "width": _stats(widths),
+        "height": _stats(heights),
+    }
+
+
+def main() -> None:
+    """CLI entry point."""
+    p = argparse.ArgumentParser(description="Naruto Sign dataset statistics")
+    p.add_argument("--data-root", required=True)
+    p.add_argument("--out", default="results/dataset_stats.json")
+    args = p.parse_args()
+
+    stats = scan_imagefolder(args.data_root)
+    os.makedirs(os.path.dirname(args.out) or ".", exist_ok=True)
+    with open(args.out, "w", encoding="utf-8") as f:
+        json.dump(stats, f, ensure_ascii=False, indent=2)
+    print(json.dumps(stats, ensure_ascii=False, indent=2))
+
+
+if __name__ == "__main__":
+    main()

code/src/losses.py

@@ -0,0 +1,101 @@
+from __future__ import annotations
+
+"""Loss functions for (possibly imbalanced) multi-class classification.
+
+See PROTOCOL §6: cross-entropy, label-smoothed CE, weighted CE, focal loss and
+class-balanced (effective-number) weighting. The choice of loss is itself a
+searchable hyper-parameter.
+"""
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class FocalLoss(nn.Module):
+    r"""Multi-class focal loss (Lin et al., 2017).
+
+    :math:`\mathrm{FL}(p_t) = -\alpha_t (1 - p_t)^\gamma \log(p_t)`.
+
+    Args:
+        gamma: Focusing parameter; ``0`` recovers cross-entropy.
+        weight: Optional per-class ``alpha`` weights ``[C]``.
+        label_smoothing: Optional smoothing applied to the CE term.
+    """
+
+    def __init__(
+        self,
+        gamma: float = 2.0,
+        weight: torch.Tensor | None = None,
+        label_smoothing: float = 0.0,
+    ) -> None:
+        super().__init__()
+        self.gamma = gamma
+        self.register_buffer("weight", weight if weight is not None else None)
+        self.label_smoothing = label_smoothing
+
+    def forward(self, logits: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        ce = F.cross_entropy(
+            logits,
+            target,
+            weight=self.weight,
+            label_smoothing=self.label_smoothing,
+            reduction="none",
+        )
+        pt = torch.exp(-ce)
+        return ((1.0 - pt) ** self.gamma * ce).mean()
+
+
+def effective_number_weights(
+    class_counts: list[int], beta: float = 0.999
+) -> torch.Tensor:
+    r"""Class-balanced weights via effective number of samples (Cui et al., 2019).
+
+    :math:`w_c \propto (1 - \beta) / (1 - \beta^{n_c})`, normalized to mean ~1.
+
+    Args:
+        class_counts: Per-class train-sample counts.
+        beta: Re-weighting hyper-parameter in ``[0, 1)``.
+
+    Returns:
+        Float tensor ``[C]``.
+    """
+    counts = np.asarray(class_counts, dtype=np.float64)
+    eff = 1.0 - np.power(beta, np.clip(counts, 1.0, None))
+    w = (1.0 - beta) / eff
+    w = w / w.sum() * len(counts)
+    return torch.tensor(w, dtype=torch.float32)
+
+
+def build_criterion(
+    name: str,
+    *,
+    class_weights: torch.Tensor | None = None,
+    label_smoothing: float = 0.0,
+    focal_gamma: float = 2.0,
+) -> nn.Module:
+    """Factory for the searchable loss choice.
+
+    Args:
+        name: One of ``{"ce", "ce_ls", "weighted_ce", "focal"}``.
+        class_weights: Per-class weights for the weighted variants.
+        label_smoothing: Smoothing for the CE-based losses.
+        focal_gamma: Focusing parameter for focal loss.
+
+    Returns:
+        A loss ``nn.Module``.
+    """
+    if name == "ce":
+        return nn.CrossEntropyLoss(label_smoothing=label_smoothing)
+    if name == "ce_ls":
+        return nn.CrossEntropyLoss(label_smoothing=max(label_smoothing, 0.1))
+    if name == "weighted_ce":
+        return nn.CrossEntropyLoss(
+            weight=class_weights, label_smoothing=label_smoothing
+        )
+    if name == "focal":
+        return FocalLoss(
+            gamma=focal_gamma, weight=class_weights, label_smoothing=label_smoothing
+        )
+    raise ValueError(f"unknown loss: {name!r}")

code/src/metrics.py

@@ -0,0 +1,92 @@
+from __future__ import annotations
+
+"""Classification metrics for an imbalanced multi-class problem (PROTOCOL §6).
+
+The *primary* selection metric for HPO is **macro-F1** (equal weight per class),
+which does not let frequent classes dominate. Accuracy, balanced accuracy, MCP
+and the confusion matrix are reported alongside it.
+"""
+
+from dataclasses import dataclass
+
+import numpy as np
+from sklearn.metrics import (
+    accuracy_score,
+    balanced_accuracy_score,
+    cohen_kappa_score,
+    confusion_matrix,
+    f1_score,
+    matthews_corrcoef,
+    precision_recall_fscore_support,
+    top_k_accuracy_score,
+)
+
+
+@dataclass
+class ClassificationReport:
+    """Aggregated metrics for one evaluation pass."""
+
+    accuracy: float
+    balanced_accuracy: float
+    macro_f1: float
+    weighted_f1: float
+    top5_accuracy: float
+    mcc: float
+    kappa: float
+    per_class_f1: list[float]
+    confusion: list[list[int]]
+
+    def as_dict(self) -> dict:
+        """JSON-serializable view (used when persisting eval reports)."""
+        return {
+            "accuracy": self.accuracy,
+            "balanced_accuracy": self.balanced_accuracy,
+            "macro_f1": self.macro_f1,
+            "weighted_f1": self.weighted_f1,
+            "top5_accuracy": self.top5_accuracy,
+            "mcc": self.mcc,
+            "kappa": self.kappa,
+            "per_class_f1": self.per_class_f1,
+            "confusion": self.confusion,
+        }
+
+
+def compute_report(
+    y_true: np.ndarray, y_pred: np.ndarray, y_prob: np.ndarray, num_classes: int
+) -> ClassificationReport:
+    """Compute the full metric bundle.
+
+    Args:
+        y_true: Ground-truth labels ``[N]``.
+        y_pred: Predicted labels ``[N]`` (argmax of probabilities).
+        y_prob: Class probabilities ``[N, C]`` (for top-k).
+        num_classes: Number of classes ``C``.
+
+    Returns:
+        A populated :class:`ClassificationReport`.
+    """
+    labels = list(range(num_classes))
+    per_class = precision_recall_fscore_support(
+        y_true, y_pred, labels=labels, average=None, zero_division=0
+    )
+    k = min(5, num_classes)
+    try:
+        top5 = float(
+            top_k_accuracy_score(y_true, y_prob, k=k, labels=labels)
+        )
+    except ValueError:
+        top5 = float("nan")
+
+    return ClassificationReport(
+        accuracy=float(accuracy_score(y_true, y_pred)),
+        balanced_accuracy=float(balanced_accuracy_score(y_true, y_pred)),
+        macro_f1=float(f1_score(y_true, y_pred, average="macro", zero_division=0)),
+        weighted_f1=float(
+            f1_score(y_true, y_pred, average="weighted", zero_division=0)
+        ),
+        top5_accuracy=top5,
+        mcc=float(matthews_corrcoef(y_true, y_pred)),
+        kappa=float(cohen_kappa_score(y_true, y_pred)),
+        per_class_f1=[float(v) for v in per_class[2]],
+        confusion=confusion_matrix(y_true, y_pred, labels=labels).tolist(),
+    )

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()

code/src/mona.py

@@ -0,0 +1,185 @@
+from __future__ import annotations
+
+"""MONA-style parameter-efficient adapter (conv variant for NCHW feature maps).
+
+Reference: Yin et al., "5%>100%: Breaking Performance Shackles of Full
+Fine-Tuning on Visual Recognition Tasks" (Mona), arXiv:2408.08345, CVPR 2025;
+code: https://github.com/Leiyi-Hu/mona (classes ``Mona`` / ``MonaOp``).
+
+The original Mona block targets transformer token sequences ``[B, N, C]``.
+EdgeNeXt produces 4-D feature maps ``[B, C, H, W]``, so this is the *conv* form
+(``Conv-MONA``): a frozen backbone receives a tiny trainable residual adapter.
+The structure mirrors the verified repo code:
+
+    x_hat = LayerNorm(x) * gamma + x * gamma_x      # scaled LayerNorm (2 scales)
+    u     = down_proj(x_hat)                         # 1x1 conv  C -> r
+    u     = MonaOp(u)                                # (DW3+DW5+DW7)/k + u, then +1x1
+    u     = dropout(gelu(u))
+    out   = x + up_proj(u)                           # outer residual (up zero-init)
+
+Only the adapters (+ classifier head) train; the backbone stays frozen. This is
+the *middle path* between full fine-tuning and pure linear probing (PROTOCOL §5).
+"""
+
+import torch
+from torch import nn
+
+
+class _ScaledChannelLayerNorm(nn.Module):
+    """Mona "scaled LayerNorm" for NCHW: ``LN(x) * gamma + x * gamma_x``.
+
+    ``gamma`` starts near zero (norm branch begins almost off) and ``gamma_x``
+    starts at one (identity passthrough), so the adapter is near-identity at init.
+    """
+
+    def __init__(self, num_channels: int, gamma_init: float = 1e-6) -> None:
+        super().__init__()
+        self.norm = nn.GroupNorm(1, num_channels)  # GroupNorm(1, C) == LayerNorm over C
+        self.gamma = nn.Parameter(torch.ones(1, num_channels, 1, 1) * gamma_init)
+        self.gamma_x = nn.Parameter(torch.ones(1, num_channels, 1, 1))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.norm(x) * self.gamma + x * self.gamma_x
+
+
+class MonaOp(nn.Module):
+    """Multi-cognitive conv aggregator: depth-wise multi-kernel + 1x1 projector.
+
+    Two residuals, as in the reference implementation.
+
+    Args:
+        dim: Bottleneck channel count.
+        kernels: Depth-wise kernel sizes (e.g. ``(3, 5, 7)``).
+    """
+
+    def __init__(self, dim: int, kernels: tuple[int, ...] = (3, 5, 7)) -> None:
+        super().__init__()
+        self.convs = nn.ModuleList(
+            [nn.Conv2d(dim, dim, k, padding=k // 2, groups=dim) for k in kernels]
+        )
+        self.projector = nn.Conv2d(dim, dim, kernel_size=1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        identity = x
+        x = sum(c(x) for c in self.convs) / len(self.convs) + identity
+        return x + self.projector(x)
+
+
+class MonaAdapter(nn.Module):
+    """Conv-MONA adapter operating on ``[B, C, H, W]`` feature maps.
+
+    Args:
+        channels: Input/output channels ``C`` of the wrapped block.
+        bottleneck: Bottleneck width ``r`` (searched by Optuna).
+        kernels: Multi-cognitive depth-wise kernel sizes.
+        gamma_init: Init for the scaled-LayerNorm ``gamma`` branch.
+        dropout: Dropout inside the bottleneck.
+    """
+
+    def __init__(
+        self,
+        channels: int,
+        bottleneck: int = 64,
+        kernels: tuple[int, ...] = (3, 5, 7),
+        gamma_init: float = 1e-6,
+        dropout: float = 0.1,
+    ) -> None:
+        super().__init__()
+        r = max(1, min(bottleneck, channels))
+        self.pre_norm = _ScaledChannelLayerNorm(channels, gamma_init=gamma_init)
+        self.down = nn.Conv2d(channels, r, kernel_size=1)
+        self.op = MonaOp(r, kernels=kernels)
+        self.act = nn.GELU()
+        self.dropout = nn.Dropout(dropout)
+        self.up = nn.Conv2d(r, channels, kernel_size=1)
+
+        nn.init.zeros_(self.up.weight)  # outer residual starts as identity
+        if self.up.bias is not None:
+            nn.init.zeros_(self.up.bias)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply the residual adapter branch and add it to ``x``."""
+        u = self.down(self.pre_norm(x))
+        u = self.dropout(self.act(self.op(u)))
+        return x + self.up(u)
+
+
+class StageMonaHook:
+    """Forward hook adding a :class:`MonaAdapter` to a stage's output tensor."""
+
+    def __init__(self, adapter: MonaAdapter) -> None:
+        self.adapter = adapter
+
+    def __call__(
+        self, module: nn.Module, inputs: tuple, output: torch.Tensor
+    ) -> torch.Tensor:
+        if not isinstance(output, torch.Tensor):  # defensive: some stages return tuples
+            return output
+        return self.adapter(output)
+
+
+@torch.no_grad()
+def _detect_stage_channels(
+    model: nn.Module, stage_indices: list[int], img_size: int = 64
+) -> dict[int, int]:
+    """Dry-run the backbone to read each target stage's output channel count.
+
+    Robust across timm versions (does not rely on ``feature_info`` being a
+    ``FeatureInfo`` object on a non-``features_only`` model).
+    """
+    recorded: dict[int, int] = {}
+    handles = []
+
+    def make_recorder(idx: int):
+        def _rec(_m, _i, out):
+            t = out[0] if isinstance(out, (tuple, list)) else out
+            recorded[idx] = t.shape[1]
+
+        return _rec
+
+    for idx in stage_indices:
+        handles.append(model.stages[idx].register_forward_hook(make_recorder(idx)))
+
+    was_training = model.training
+    model.eval()
+    device = next(model.parameters()).device
+    model(torch.zeros(1, 3, img_size, img_size, device=device))
+    model.train(was_training)
+    for h in handles:
+        h.remove()
+    return recorded
+
+
+def attach_mona_adapters(
+    model: nn.Module,
+    stage_indices: list[int],
+    *,
+    bottleneck: int = 64,
+    kernels: tuple[int, ...] = (3, 5, 7),
+    gamma_init: float = 1e-6,
+) -> nn.ModuleList:
+    """Insert MONA adapters after the chosen EdgeNeXt stages.
+
+    Args:
+        model: A timm EdgeNeXt exposing ``model.stages`` (indexable).
+        stage_indices: 0-based stage indices to adapt (e.g. ``[2, 3]``).
+        bottleneck: Adapter bottleneck width.
+        kernels: Multi-cognitive depth-wise kernel sizes.
+        gamma_init: Init for the scaled-norm branch.
+
+    Returns:
+        The ``ModuleList`` of created adapters (also attached to ``model``).
+    """
+    channels = _detect_stage_channels(model, stage_indices)
+    adapters = nn.ModuleList()
+    for idx in stage_indices:
+        adapter = MonaAdapter(
+            channels[idx],
+            bottleneck=bottleneck,
+            kernels=kernels,
+            gamma_init=gamma_init,
+        )
+        model.stages[idx].register_forward_hook(StageMonaHook(adapter))
+        adapters.append(adapter)
+    model.add_module("mona_adapters", adapters)  # part of the module tree + .to()
+    return adapters

code/src/optuna_search.py

@@ -0,0 +1,232 @@
+from __future__ import annotations
+
+"""Optuna hyper-parameter search for EdgeNeXt on Naruto Sign (PROTOCOL §7-§9).
+
+Single-objective (maximize val macro-F1) or multi-objective
+(maximize macro-F1 + minimize trainable params) search with configurable
+sampler and pruner. Define-by-run: the search space is described inside
+``suggest_configs`` using ``trial.suggest_*`` calls, so conditional spaces
+(e.g. focal-only ``focal_gamma``) are expressed naturally.
+
+Usage::
+
+    python -m src.optuna_search --data-root /path/Naruto_Sign \
+        --n-trials 60 --sampler tpe --pruner median --study-name nss_v1
+
+    # Multi-objective (accuracy vs trainable params):
+    python -m src.optuna_search --data-root /path/Naruto_Sign \
+        --n-trials 80 --multi-objective --sampler nsga
+"""
+
+import argparse
+import json
+import os
+
+import optuna
+import torch
+
+from .data import AugConfig, DataConfig
+from .model import ModelConfig
+from .train import TrainConfig, train_model
+
+EDGENEXT_VARIANTS = [
+    "edgenext_xx_small.in1k",
+    "edgenext_x_small.in1k",
+    "edgenext_small.usi_in1k",
+]
+
+
+def suggest_configs(
+    trial: optuna.Trial, data_root: str, num_classes: int, epochs: int
+) -> tuple[ModelConfig, DataConfig, TrainConfig]:
+    """Sample one configuration from the search space (define-by-run).
+
+    Args:
+        trial: The Optuna trial.
+        data_root: Path to the ImageFolder dataset root.
+        num_classes: Number of classes.
+        epochs: Max epochs per trial (training budget).
+
+    Returns:
+        ``(model_cfg, data_cfg, train_cfg)`` for :func:`train_model`.
+    """
+    # --- backbone + transfer-learning regime (H2/H3) ----------------------
+    model_name = trial.suggest_categorical("model_name", EDGENEXT_VARIANTS)
+    regime = trial.suggest_categorical("regime", ["full", "partial", "mona"])
+    drop_path = trial.suggest_float("drop_path_rate", 0.0, 0.3)
+    drop_rate = trial.suggest_float("drop_rate", 0.0, 0.4)
+
+    n_unfrozen, mona_stages, mona_bottleneck, mona_kernels = 1, (2, 3), 64, (3, 5, 7)
+    if regime == "partial":
+        n_unfrozen = trial.suggest_int("n_unfrozen_stages", 0, 4)
+    elif regime == "mona":
+        mona_bottleneck = trial.suggest_categorical("mona_bottleneck", [16, 32, 64, 96])
+        kernel_set = trial.suggest_categorical(
+            "mona_kernels", ["3", "3-5", "3-5-7", "3-5-7-9"]
+        )
+        mona_kernels = tuple(int(k) for k in kernel_set.split("-"))
+        last_only = trial.suggest_categorical("mona_last_only", [True, False])
+        mona_stages = (3,) if last_only else (2, 3)
+
+    model_cfg = ModelConfig(
+        model_name=model_name,
+        num_classes=num_classes,
+        pretrained=True,
+        drop_rate=drop_rate,
+        drop_path_rate=drop_path,
+        regime=regime,
+        n_unfrozen_stages=n_unfrozen,
+        mona_stages=mona_stages,
+        mona_bottleneck=mona_bottleneck,
+        mona_kernels=mona_kernels,
+    )
+
+    # --- optimization (H1) -----------------------------------------------
+    lr = trial.suggest_float("lr", 1e-5, 5e-3, log=True)
+    weight_decay = trial.suggest_float("weight_decay", 1e-6, 1e-2, log=True)
+    optimizer = trial.suggest_categorical("optimizer", ["adamw", "sgd"])
+    backbone_lr_mult = trial.suggest_float("backbone_lr_mult", 0.02, 1.0, log=True)
+
+    # --- loss + imbalance handling (H4) ----------------------------------
+    loss_name = trial.suggest_categorical(
+        "loss_name", ["ce", "ce_ls", "weighted_ce", "focal"]
+    )
+    label_smoothing = trial.suggest_float("label_smoothing", 0.0, 0.2)
+    focal_gamma = trial.suggest_float("focal_gamma", 0.5, 4.0) if loss_name == "focal" else 2.0
+    use_weighted_sampler = trial.suggest_categorical("weighted_sampler", [True, False])
+
+    # --- augmentation (H6) -----------------------------------------------
+    img_size = trial.suggest_categorical("img_size", [224, 256])
+    mixup_alpha = trial.suggest_float("mixup_alpha", 0.0, 0.4)
+    rrc_scale_min = trial.suggest_float("rrc_scale_min", 0.5, 0.9)
+    randaug = trial.suggest_int("randaug_magnitude", 0, 9)
+
+    aug = AugConfig(
+        img_size=img_size,
+        use_hflip=False,  # fixed OFF: chirality-sensitive seals (PROTOCOL §3)
+        rrc_scale_min=rrc_scale_min,
+        randaug_magnitude=randaug,
+    )
+    data_cfg = DataConfig(
+        data_root=data_root,
+        batch_size=trial.suggest_categorical("batch_size", [16, 32, 64]),
+        use_weighted_sampler=use_weighted_sampler,
+        aug=aug,
+    )
+    train_cfg = TrainConfig(
+        epochs=epochs,
+        lr=lr,
+        weight_decay=weight_decay,
+        optimizer=optimizer,
+        backbone_lr_mult=backbone_lr_mult,
+        loss_name=loss_name,
+        label_smoothing=label_smoothing,
+        focal_gamma=focal_gamma,
+        mixup_alpha=mixup_alpha,
+    )
+    return model_cfg, data_cfg, train_cfg
+
+
+def make_study(args: argparse.Namespace) -> optuna.Study:
+    """Construct a study with the requested sampler + pruner."""
+    samplers = {
+        "tpe": optuna.samplers.TPESampler(multivariate=True, seed=args.seed),
+        "random": optuna.samplers.RandomSampler(seed=args.seed),
+        "cmaes": optuna.samplers.CmaEsSampler(seed=args.seed),
+        "nsga": optuna.samplers.NSGAIISampler(seed=args.seed),
+    }
+    # AutoSampler lives in OptunaHub; load it when requested (PROTOCOL §2.5).
+    if args.sampler == "auto":
+        import optunahub
+
+        sampler = optunahub.load_module("samplers/auto_sampler").AutoSampler()
+    else:
+        sampler = samplers[args.sampler]
+
+    pruners = {
+        "median": optuna.pruners.MedianPruner(n_warmup_steps=5),
+        "asha": optuna.pruners.SuccessiveHalvingPruner(),
+        "hyperband": optuna.pruners.HyperbandPruner(),
+        "none": optuna.pruners.NopPruner(),
+    }
+    storage = f"sqlite:///{args.study_name}.db" if args.storage else None
+
+    if args.multi_objective:
+        return optuna.create_study(
+            study_name=args.study_name,
+            storage=storage,
+            load_if_exists=bool(storage),
+            sampler=sampler,
+            directions=["maximize", "minimize"],  # macro-F1 up, params down
+        )
+    return optuna.create_study(
+        study_name=args.study_name,
+        storage=storage,
+        load_if_exists=bool(storage),
+        sampler=sampler,
+        pruner=pruners[args.pruner],
+        direction="maximize",
+    )
+
+
+def main() -> None:
+    """CLI entry point."""
+    p = argparse.ArgumentParser(description="Optuna HPO for EdgeNeXt / Naruto Sign")
+    p.add_argument("--data-root", required=True)
+    p.add_argument("--num-classes", type=int, default=13)  # 12 seals + 'zero'
+    p.add_argument("--n-trials", type=int, default=60)
+    p.add_argument("--epochs", type=int, default=30)
+    p.add_argument("--sampler", choices=["tpe", "random", "cmaes", "nsga", "auto"], default="tpe")
+    p.add_argument("--pruner", choices=["median", "asha", "hyperband", "none"], default="median")
+    p.add_argument("--multi-objective", action="store_true")
+    p.add_argument("--study-name", default="nss_hpo")
+    p.add_argument("--storage", action="store_true", help="persist to sqlite db")
+    p.add_argument("--out", default="results")
+    p.add_argument("--seed", type=int, default=42)
+    args = p.parse_args()
+
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    os.makedirs(args.out, exist_ok=True)
+
+    def objective(trial: optuna.Trial):
+        model_cfg, data_cfg, train_cfg = suggest_configs(
+            trial, args.data_root, args.num_classes, args.epochs
+        )
+        # Multi-objective studies do not support pruning -> do not pass the trial
+        # (trial.report/should_prune raise on multi-objective trials).
+        prune_trial = None if args.multi_objective else trial
+        result = train_model(model_cfg, data_cfg, train_cfg, device, trial=prune_trial)
+        trial.set_user_attr("trainable_params", result["trainable_params"])
+        if args.multi_objective:
+            return result["best_val_macro_f1"], float(result["trainable_params"])
+        return result["best_val_macro_f1"]
+
+    study = make_study(args)
+    study.optimize(objective, n_trials=args.n_trials, gc_after_trial=True)
+
+    _persist(study, args)
+
+
+def _persist(study: optuna.Study, args: argparse.Namespace) -> None:
+    """Write best params / trials dataframe atomically."""
+    df = study.trials_dataframe()
+    df.to_csv(os.path.join(args.out, f"{args.study_name}_trials.csv"), index=False)
+
+    if args.multi_objective:
+        best = [
+            {"values": t.values, "params": t.params} for t in study.best_trials
+        ]
+        payload = {"pareto_front": best}
+    else:
+        payload = {"best_value": study.best_value, "best_params": study.best_params}
+
+    tmp = os.path.join(args.out, f"{args.study_name}_best.json.tmp")
+    final = os.path.join(args.out, f"{args.study_name}_best.json")
+    with open(tmp, "w", encoding="utf-8") as f:
+        json.dump(payload, f, ensure_ascii=False, indent=2)
+    os.replace(tmp, final)
+    print(f"[optuna] wrote {final}")
+
+
+if __name__ == "__main__":
+    main()

code/src/run_baseline.py

@@ -0,0 +1,90 @@
+from __future__ import annotations
+
+"""Train a single configuration and dump test metrics + plots (PROTOCOL §9, B0).
+
+Used for the H2/H3 regime ablation: run the same recipe with ``--regime`` in
+{full, partial, mona} at equal budget and compare. Also produces the baseline
+that the Optuna search must beat.
+"""
+
+import argparse
+import json
+import os
+
+import torch
+
+from .data import AugConfig, DataConfig
+from .model import ModelConfig
+from .train import TrainConfig, train_model
+
+
+def plot_confusion(confusion: list[list[int]], class_names: list[str], out_png: str) -> None:
+    """Save a confusion-matrix heatmap."""
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    cm = np.array(confusion, dtype=float)
+    cm_norm = cm / np.clip(cm.sum(axis=1, keepdims=True), 1, None)
+    fig, ax = plt.subplots(figsize=(8, 7))
+    im = ax.imshow(cm_norm, cmap="viridis", vmin=0, vmax=1)
+    ax.set_xticks(range(len(class_names)))
+    ax.set_yticks(range(len(class_names)))
+    ax.set_xticklabels(class_names, rotation=90, fontsize=7)
+    ax.set_yticklabels(class_names, fontsize=7)
+    ax.set_xlabel("predicted")
+    ax.set_ylabel("true")
+    ax.set_title("Confusion matrix (row-normalized)")
+    fig.colorbar(im)
+    fig.tight_layout()
+    fig.savefig(out_png, dpi=150)
+    plt.close(fig)
+
+
+def main() -> None:
+    """CLI entry point."""
+    p = argparse.ArgumentParser(description="Single EdgeNeXt training run")
+    p.add_argument("--data-root", required=True)
+    p.add_argument("--num-classes", type=int, default=13)  # 12 seals + 'zero'
+    p.add_argument("--model-name", default="edgenext_small.usi_in1k")
+    p.add_argument("--regime", choices=["full", "partial", "mona"], default="partial")
+    p.add_argument("--n-unfrozen", type=int, default=1)
+    p.add_argument("--epochs", type=int, default=30)
+    p.add_argument("--lr", type=float, default=1e-3)
+    p.add_argument("--loss", default="ce_ls")
+    p.add_argument("--weighted-sampler", action="store_true")
+    p.add_argument("--out", default="results/baseline")
+    p.add_argument("--seed", type=int, default=42)
+    args = p.parse_args()
+
+    os.makedirs(args.out, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model_cfg = ModelConfig(
+        model_name=args.model_name,
+        num_classes=args.num_classes,
+        regime=args.regime,
+        n_unfrozen_stages=args.n_unfrozen,
+    )
+    data_cfg = DataConfig(
+        data_root=args.data_root,
+        use_weighted_sampler=args.weighted_sampler,
+        seed=args.seed,
+        aug=AugConfig(),
+    )
+    train_cfg = TrainConfig(epochs=args.epochs, lr=args.lr, loss_name=args.loss, seed=args.seed)
+
+    result = train_model(model_cfg, data_cfg, train_cfg, device, eval_test=True)
+
+    with open(os.path.join(args.out, "report.json"), "w", encoding="utf-8") as f:
+        json.dump(result, f, ensure_ascii=False, indent=2)
+    if "test" in result:
+        plot_confusion(
+            result["test"]["confusion"],
+            result["class_names"],
+            os.path.join(args.out, "confusion.png"),
+        )
+    print(json.dumps({k: v for k, v in result.items() if k != "history"}, indent=2))
+
+
+if __name__ == "__main__":
+    main()

code/src/train.py

@@ -0,0 +1,225 @@
+from __future__ import annotations
+
+"""Training / evaluation loop with optional Optuna pruning.
+
+``train_model`` is regime-agnostic: it builds the model via :mod:`model`, the
+loss via :mod:`losses`, and reports macro-F1 each epoch. When an Optuna ``trial``
+is passed, it reports the intermediate validation macro-F1 and honors
+``trial.should_prune()`` (median / ASHA / Hyperband pruning — see PROTOCOL §2.3).
+"""
+
+import gc
+import random
+from dataclasses import dataclass, field
+
+import numpy as np
+import torch
+from torch import nn
+
+from .data import DataConfig, build_dataloaders, compute_class_weights
+from .losses import build_criterion, effective_number_weights
+from .metrics import ClassificationReport, compute_report
+from .model import (
+    ModelConfig,
+    build_model,
+    build_param_groups,
+    count_trainable,
+    freeze_consistency,
+)
+
+
+def set_seed(seed: int = 42) -> None:
+    """Seed Python/NumPy/PyTorch for reproducibility (PROTOCOL §10)."""
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+
+@dataclass
+class TrainConfig:
+    """Optimization configuration (a subset is searched by Optuna)."""
+
+    epochs: int = 30
+    lr: float = 1e-3
+    weight_decay: float = 1e-4
+    optimizer: str = "adamw"  # {"adamw", "sgd"}
+    backbone_lr_mult: float = 0.1
+    warmup_epochs: int = 2
+    loss_name: str = "ce_ls"  # {"ce", "ce_ls", "weighted_ce", "focal"}
+    label_smoothing: float = 0.1
+    focal_gamma: float = 2.0
+    cb_beta: float = 0.999  # effective-number beta for weighted/focal weights
+    mixup_alpha: float = 0.0  # 0 disables mixup
+    grad_clip: float = 0.0
+    amp: bool = True
+    early_stop_patience: int = 8
+    seed: int = 42
+
+
+def _build_optimizer(model: nn.Module, cfg: TrainConfig) -> torch.optim.Optimizer:
+    groups = build_param_groups(
+        model, cfg.lr, backbone_lr_mult=cfg.backbone_lr_mult
+    )
+    if cfg.optimizer == "adamw":
+        return torch.optim.AdamW(groups, weight_decay=cfg.weight_decay)
+    if cfg.optimizer == "sgd":
+        return torch.optim.SGD(
+            groups, momentum=0.9, weight_decay=cfg.weight_decay, nesterov=True
+        )
+    raise ValueError(f"unknown optimizer: {cfg.optimizer!r}")
+
+
+def _build_scheduler(
+    optimizer: torch.optim.Optimizer, cfg: TrainConfig, steps_per_epoch: int
+):
+    total = cfg.epochs * max(1, steps_per_epoch)
+    warmup = cfg.warmup_epochs * max(1, steps_per_epoch)
+
+    def lr_lambda(step: int) -> float:
+        if step < warmup:
+            return (step + 1) / max(1, warmup)
+        progress = (step - warmup) / max(1, total - warmup)
+        return 0.5 * (1.0 + np.cos(np.pi * min(1.0, progress)))
+
+    return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+
+@torch.inference_mode()
+def evaluate(
+    model: nn.Module, loader, device: torch.device, num_classes: int
+) -> ClassificationReport:
+    """Run a full evaluation pass and return the metric bundle."""
+    model.eval()
+    probs, trues = [], []
+    for x, y in loader:
+        x = x.to(device, non_blocking=True)
+        logits = model(x)
+        probs.append(torch.softmax(logits, dim=1).cpu().numpy())
+        trues.append(y.numpy())
+    y_prob = np.concatenate(probs)
+    y_true = np.concatenate(trues)
+    y_pred = y_prob.argmax(axis=1)
+    return compute_report(y_true, y_pred, y_prob, num_classes)
+
+
+def train_model(
+    model_cfg: ModelConfig,
+    data_cfg: DataConfig,
+    train_cfg: TrainConfig,
+    device: torch.device,
+    *,
+    trial=None,
+    eval_test: bool = False,
+) -> dict:
+    """Train one configuration end-to-end.
+
+    Args:
+        model_cfg: Backbone + regime configuration.
+        data_cfg: Data / split / sampler configuration.
+        train_cfg: Optimization configuration.
+        device: Target device.
+        trial: Optional Optuna ``Trial`` for intermediate reporting + pruning.
+        eval_test: If ``True`` also evaluate the held-out test split at the end.
+
+    Returns:
+        Dict with ``best_val_macro_f1``, ``history``, ``trainable_params`` and
+        (optionally) ``test`` metrics.
+
+    Raises:
+        optuna.TrialPruned: If the trial is pruned by the configured pruner.
+    """
+    set_seed(train_cfg.seed)
+    bundle = build_dataloaders(data_cfg)
+    model = build_model(model_cfg).to(device)
+    trainable, total = count_trainable(model)
+
+    # Class-imbalance weights for the weighted/focal losses.
+    counts = [bundle.train_targets.count(c) for c in range(bundle.num_classes)]
+    if train_cfg.loss_name in {"weighted_ce", "focal"}:
+        weights = effective_number_weights(counts, beta=train_cfg.cb_beta).to(device)
+    else:
+        weights = None
+    criterion = build_criterion(
+        train_cfg.loss_name,
+        class_weights=weights,
+        label_smoothing=train_cfg.label_smoothing,
+        focal_gamma=train_cfg.focal_gamma,
+    )
+
+    mixup_fn = None
+    if train_cfg.mixup_alpha > 0:
+        from timm.data import Mixup
+        from timm.loss import SoftTargetCrossEntropy
+
+        mixup_fn = Mixup(
+            mixup_alpha=train_cfg.mixup_alpha,
+            cutmix_alpha=0.0,  # pure mixup; cutmix is a separate (cautious) knob
+            num_classes=bundle.num_classes,
+        )
+        train_criterion: nn.Module = SoftTargetCrossEntropy()
+    else:
+        train_criterion = criterion
+
+    optimizer = _build_optimizer(model, train_cfg)
+    scheduler = _build_scheduler(optimizer, train_cfg, len(bundle.train_loader))
+    scaler = torch.cuda.amp.GradScaler(enabled=train_cfg.amp)
+
+    best_f1, best_epoch, history = -1.0, -1, []
+    for epoch in range(train_cfg.epochs):
+        model.train()
+        freeze_consistency(model)  # keep frozen BN/LN in eval (PROTOCOL §5.2)
+        for x, y in bundle.train_loader:
+            x = x.to(device, non_blocking=True)
+            y = y.to(device, non_blocking=True)
+            if mixup_fn is not None:
+                x, y = mixup_fn(x, y)
+            optimizer.zero_grad(set_to_none=True)
+            with torch.cuda.amp.autocast(enabled=train_cfg.amp):
+                loss = train_criterion(model(x), y)
+            scaler.scale(loss).backward()
+            if train_cfg.grad_clip > 0:
+                scaler.unscale_(optimizer)
+                nn.utils.clip_grad_norm_(model.parameters(), train_cfg.grad_clip)
+            scaler.step(optimizer)
+            scaler.update()
+            scheduler.step()
+
+        report = evaluate(model, bundle.val_loader, device, bundle.num_classes)
+        history.append({"epoch": epoch, "val_macro_f1": report.macro_f1})
+
+        if report.macro_f1 > best_f1:
+            best_f1, best_epoch = report.macro_f1, epoch
+
+        if trial is not None:
+            trial.report(report.macro_f1, epoch)
+            if trial.should_prune():
+                import optuna
+
+                _cleanup(model)
+                raise optuna.TrialPruned()
+
+        if epoch - best_epoch >= train_cfg.early_stop_patience:
+            break  # early stop on validation macro-F1
+
+    out: dict = {
+        "best_val_macro_f1": best_f1,
+        "best_epoch": best_epoch,
+        "history": history,
+        "trainable_params": trainable,
+        "total_params": total,
+    }
+    if eval_test:
+        test_report = evaluate(model, bundle.test_loader, device, bundle.num_classes)
+        out["test"] = test_report.as_dict()
+        out["class_names"] = bundle.class_names
+    _cleanup(model)
+    return out
+
+
+def _cleanup(model: nn.Module) -> None:
+    """Free GPU memory between trials (PROTOCOL §10, VRAM hygiene)."""
+    del model
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()

code/src/umap_analysis.py

@@ -0,0 +1,147 @@
+from __future__ import annotations
+
+"""Feature extraction + UMAP visualization + clustering quality (PROTOCOL §8).
+
+Pipeline:
+1. Run a (pretrained or fine-tuned) EdgeNeXt over a split and collect L2-normalized
+   pooled embeddings ``[N, D]`` with their true labels.
+2. UMAP -> 2D for plotting (colored by class).
+3. UMAP -> ``cluster_dim`` (10-50) for *clustering* (authors recommend clustering on
+   a mid-dim embedding, NOT on the 2D picture), then KMeans / HDBSCAN.
+4. Report silhouette + agreement with labels (ARI / NMI).
+
+WARNING: do not over-interpret absolute distances or cluster sizes on the 2D UMAP
+plot; fix ``random_state`` for reproducibility.
+"""
+
+import argparse
+import json
+import os
+
+import numpy as np
+import torch
+
+from .data import AugConfig, DataConfig, build_dataloaders, build_transforms
+from .model import ModelConfig, build_model, extract_features
+
+
+@torch.inference_mode()
+def collect_embeddings(
+    model, loader, device: torch.device
+) -> tuple[np.ndarray, np.ndarray]:
+    """Collect L2-normalized embeddings and labels over a dataloader."""
+    model.eval()
+    embs, labels = [], []
+    for x, y in loader:
+        x = x.to(device, non_blocking=True)
+        embs.append(extract_features(model, x, l2_normalize=True).numpy())
+        labels.append(y.numpy())
+    return np.concatenate(embs), np.concatenate(labels)
+
+
+def run_umap(
+    emb: np.ndarray, n_components: int, n_neighbors: int, min_dist: float, seed: int
+) -> np.ndarray:
+    """UMAP embedding with cosine metric (suits L2-normalized features)."""
+    import umap
+
+    reducer = umap.UMAP(
+        n_components=n_components,
+        n_neighbors=n_neighbors,
+        min_dist=min_dist,
+        metric="cosine",
+        random_state=seed,
+    )
+    return reducer.fit_transform(emb)
+
+
+def cluster_and_score(
+    emb_mid: np.ndarray, labels: np.ndarray, num_classes: int
+) -> dict:
+    """KMeans on a mid-dim UMAP embedding; score vs ground-truth labels.
+
+    Returns:
+        Dict with silhouette, ARI and NMI.
+    """
+    from sklearn.cluster import KMeans
+    from sklearn.metrics import (
+        adjusted_rand_score,
+        normalized_mutual_info_score,
+        silhouette_score,
+    )
+
+    km = KMeans(n_clusters=num_classes, n_init=10, random_state=42)
+    assign = km.fit_predict(emb_mid)
+    return {
+        "silhouette": float(silhouette_score(emb_mid, assign)),
+        "ari_vs_labels": float(adjusted_rand_score(labels, assign)),
+        "nmi_vs_labels": float(normalized_mutual_info_score(labels, assign)),
+    }
+
+
+def plot_2d(
+    emb2d: np.ndarray, labels: np.ndarray, class_names: list[str], out_png: str
+) -> None:
+    """Scatter the 2D UMAP colored by true class."""
+    import matplotlib.pyplot as plt
+
+    fig, ax = plt.subplots(figsize=(9, 7))
+    for c, name in enumerate(class_names):
+        m = labels == c
+        ax.scatter(emb2d[m, 0], emb2d[m, 1], s=10, label=name, alpha=0.7)
+    ax.set_title("UMAP of EdgeNeXt features (Naruto Sign)")
+    ax.legend(markerscale=2, bbox_to_anchor=(1.02, 1), loc="upper left", fontsize=8)
+    fig.tight_layout()
+    fig.savefig(out_png, dpi=150)
+    plt.close(fig)
+
+
+def main() -> None:
+    """CLI: extract -> UMAP 2D plot + mid-dim clustering report."""
+    p = argparse.ArgumentParser(description="UMAP feature analysis")
+    p.add_argument("--data-root", required=True)
+    p.add_argument("--model-name", default="edgenext_small.usi_in1k")
+    p.add_argument("--num-classes", type=int, default=13)  # 12 seals + 'zero'
+    p.add_argument("--checkpoint", default="", help="optional fine-tuned state_dict")
+    p.add_argument("--split", choices=["train", "val", "test"], default="val")
+    p.add_argument("--n-neighbors", type=int, default=15)
+    p.add_argument("--min-dist", type=float, default=0.1)
+    p.add_argument("--cluster-dim", type=int, default=20)
+    p.add_argument("--out", default="results/umap")
+    p.add_argument("--seed", type=int, default=42)
+    args = p.parse_args()
+
+    os.makedirs(args.out, exist_ok=True)
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    model = build_model(
+        ModelConfig(model_name=args.model_name, num_classes=args.num_classes, regime="full")
+    ).to(device)
+    if args.checkpoint:
+        model.load_state_dict(torch.load(args.checkpoint, map_location=device))
+
+    bundle = build_dataloaders(
+        DataConfig(data_root=args.data_root, aug=AugConfig(), seed=args.seed)
+    )
+    loader = {
+        "train": bundle.train_loader,
+        "val": bundle.val_loader,
+        "test": bundle.test_loader,
+    }[args.split]
+
+    emb, labels = collect_embeddings(model, loader, device)
+    np.save(os.path.join(args.out, "embeddings.npy"), emb)
+    np.save(os.path.join(args.out, "labels.npy"), labels)
+
+    emb2d = run_umap(emb, 2, args.n_neighbors, args.min_dist, args.seed)
+    plot_2d(emb2d, labels, bundle.class_names, os.path.join(args.out, "umap_2d.png"))
+
+    emb_mid = run_umap(emb, args.cluster_dim, args.n_neighbors, 0.0, args.seed)
+    report = cluster_and_score(emb_mid, labels, bundle.num_classes)
+    with open(os.path.join(args.out, "cluster_report.json"), "w", encoding="utf-8") as f:
+        json.dump(report, f, ensure_ascii=False, indent=2)
+    print(f"[umap] {report}")
+
+
+if __name__ == "__main__":
+    main()