fuse_proj: Initial operational package for 3 researchers (Pavlenko/Blizno/Moroz)

Multimodal fusion research on StripNet+GTA-UAV proxy: - 3 independent fusion tracks: condition-aware (A), token/bottleneck (B), role-aware (C) - Shared interfaces, protocol, dataset audit, baseline benchmarks - Canonical version-chain references to vault (SPEC, ANALYSIS, TRIAGE) - Personalized task plans and decision tables for each researcher - 3 generated DOCX task assignment files with milestones and DoD checklist - Full modality dropout diagnostics and missing-modality robustness requirements - Data contract, benchmark registry, experiment tracking infrastructure Operational documents: - docs/00_project/: MERIDIAN context, protocol, repository reuse guide, experiment specification - docs/01_tasks/: Master assignment + 3 individual researcher tracks + joint integration - docs/02_references/: Core literature, version-chain bases, code maps - docs/03_codebase_guides/: Existing code snapshots from vault - scripts/: gen_task_plans.js (DOCX generation), placeholder infrastructure - vendor_reference/: Snapshots of caption_test, depth_edges_annotate, existing SOFIA/SegModel code - reports/, results/, experiments/: Shared output structure for all 3 researchers 3 DOCX files generated from gen_task_plans.js (Times New Roman 14pt, GOST format): - План_заданий_Павленко_БВ.docx (Condition-Aware track, fusion API owner) - План_заданий_Близно_МВ.docx (Token/Bottleneck track, benchmark owner) - План_заданий_Мороз_ЕС.docx (Role-Aware track, data contract owner) Co-Authored-By: Claude Haiku 4.5 <noreply@anthropic.com>
2026-06-11 17:16:57 +03:00
commit 2c6a00a4ca
155 changed files with 39765 additions and 0 deletions
--- a/vendor_reference/caption_test/src/eval/evaluate.py
+++ b/vendor_reference/caption_test/src/eval/evaluate.py
@@ -0,0 +1,153 @@
+from __future__ import annotations
+
+"""Evaluation for caption quality test.
+
+Recall@K for query(drone+text) -> gallery(satellite).
+delta_r_at_1 compares caption-aware vs baseline runs.
+"""
+
+import json
+import logging
+from pathlib import Path
+
+import gin
+import torch
+from torch.utils.data import DataLoader
+
+from src.models.dual_encoder import DualEncoderCaptionTest
+
+LOGGER = logging.getLogger("caption_test.eval")
+
+
+def _recall_at_k(
+    similarity: torch.Tensor,
+    k_values: tuple[int, ...] = (1, 5, 10),
+) -> dict[int, float]:
+    """Recall@K assuming positives on the diagonal."""
+    n_query = similarity.size(0)
+    targets = torch.arange(n_query, device=similarity.device)
+    sorted_idx = similarity.argsort(dim=1, descending=True)
+    result: dict[int, float] = {}
+    for k in k_values:
+        top_k = sorted_idx[:, :k]
+        hit = (top_k == targets.unsqueeze(1)).any(dim=1).float()
+        result[k] = float(hit.mean().item())
+    return result
+
+
+@torch.no_grad()
+def _encode_dataset(
+    model: DualEncoderCaptionTest,
+    loader: DataLoader,
+    device: str,
+) -> dict[str, torch.Tensor]:
+    """Encode all samples into query and gallery embeddings."""
+    model.eval()
+    all_query: list[torch.Tensor] = []
+    all_gallery: list[torch.Tensor] = []
+
+    for batch in loader:
+        drone_img = batch["drone_img"].to(device, non_blocking=True)
+        sat_img = batch["sat_img"].to(device, non_blocking=True)
+        caption_drone = batch["caption_drone"]
+
+        embeddings = model(
+            drone_img=drone_img,
+            sat_img=sat_img,
+            caption_drone=caption_drone,
+        )
+        all_query.append(embeddings["query"].cpu())
+        all_gallery.append(embeddings["gallery"].cpu())
+
+    return {
+        "query": torch.cat(all_query, dim=0),
+        "gallery": torch.cat(all_gallery, dim=0),
+    }
+
+
+def evaluate_retrieval(
+    model: DualEncoderCaptionTest,
+    loader: DataLoader,
+    device: str,
+    k_values: tuple[int, ...] = (1, 5, 10),
+) -> dict[str, float]:
+    """Compute R@K for query->gallery and gallery->query.
+
+    Returns:
+        Flat dict: r@1_query_to_gallery, r@5_query_to_gallery, etc.
+    """
+    feats = _encode_dataset(model=model, loader=loader, device=device)
+
+    metrics: dict[str, float] = {}
+
+    sim_q2g = feats["query"] @ feats["gallery"].t()
+
+    for k, val in _recall_at_k(sim_q2g, k_values).items():
+        metrics[f"r@{k}_query_to_gallery"] = val
+    for k, val in _recall_at_k(sim_q2g.t(), k_values).items():
+        metrics[f"r@{k}_gallery_to_query"] = val
+
+    # Gate value for diagnostics.
+    metrics["gate"] = model.fusion.gate_value
+
+    return metrics
+
+
+def delta_r_at_1(
+    full_metrics: dict[str, float],
+    baseline_metrics: dict[str, float],
+) -> float:
+    """R@1 gain from adding captions: full - baseline."""
+    key = "r@1_query_to_gallery"
+    return full_metrics[key] - baseline_metrics[key]
+
+
+@gin.configurable
+def run_evaluation_from_checkpoint(
+    checkpoint_path: str,
+    test_query_file: str,
+    data_root: str,
+    output_path: str = "eval_report.json",
+    batch_size: int = 128,
+    num_workers: int = 4,
+    device: str = "cuda",
+) -> dict[str, float]:
+    """Standalone evaluation from checkpoint."""
+    from src.datasets.visloc_with_captions import (
+        GeoLocCaptionDataset,
+        collate_caption_batch,
+    )
+
+    model = DualEncoderCaptionTest().to(device)
+    ckpt = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    model.load_state_dict(ckpt["model_state"])
+    model.eval()
+
+    test_ds = GeoLocCaptionDataset(
+        query_file=test_query_file,
+        data_root=data_root,
+        image_transform=model.preprocess,
+    )
+    test_loader = DataLoader(
+        test_ds,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        collate_fn=collate_caption_batch,
+        pin_memory=True,
+    )
+
+    metrics = evaluate_retrieval(model=model, loader=test_loader, device=device)
+
+    report = {
+        "checkpoint": checkpoint_path,
+        "test_query_file": test_query_file,
+        "metrics": metrics,
+    }
+    out = Path(output_path)
+    out.parent.mkdir(parents=True, exist_ok=True)
+    with out.open("w", encoding="utf-8") as f:
+        json.dump(report, f, indent=2)
+
+    LOGGER.info("evaluation report saved to %s", out)
+    return metrics
--- a/vendor_reference/caption_test/src/losses/hard_negatives.py
+++ b/vendor_reference/caption_test/src/losses/hard_negatives.py
@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+"""Hard negative memory bank for contrastive learning.
+
+MoCo-style FIFO queue of recent gallery embeddings. Each batch gets
+B in-batch negatives + Q queue negatives, significantly increasing
+the effective number of negatives without extra VRAM for forward pass.
+
+Usage:
+    bank = NegativeMemoryBank(size=4096, dim=1024)
+    # In training loop:
+    sim = bank.compute_similarity(query, gallery)  # [B, B + Q]
+    bank.enqueue(gallery.detach())
+"""
+
+import torch
+import torch.nn as nn
+
+
+class NegativeMemoryBank(nn.Module):
+    """FIFO queue of detached gallery embeddings for hard negatives.
+
+    Args:
+        size: Queue capacity (number of stored embeddings).
+        dim: Embedding dimension.
+    """
+
+    def __init__(self, size: int = 4096, dim: int = 1024) -> None:
+        super().__init__()
+        self.size = size
+        self.dim = dim
+        # Queue stored as buffer (not a parameter, moves with .to(device)).
+        self.register_buffer("queue", torch.randn(size, dim))
+        self.queue = nn.functional.normalize(self.queue, dim=-1)
+        self.register_buffer("ptr", torch.zeros(1, dtype=torch.long))
+        self.register_buffer("full", torch.zeros(1, dtype=torch.bool))
+
+    @torch.no_grad()
+    def enqueue(self, embeddings: torch.Tensor) -> None:
+        """Add embeddings to the queue (FIFO). Oldest are overwritten."""
+        batch_size = embeddings.shape[0]
+        ptr = int(self.ptr.item())
+
+        if ptr + batch_size <= self.size:
+            self.queue[ptr:ptr + batch_size] = embeddings.detach()
+        else:
+            # Wrap around.
+            overflow = (ptr + batch_size) - self.size
+            self.queue[ptr:] = embeddings[:batch_size - overflow].detach()
+            self.queue[:overflow] = embeddings[batch_size - overflow:].detach()
+
+        new_ptr = (ptr + batch_size) % self.size
+        self.ptr[0] = new_ptr
+        if not self.full.item() and (new_ptr < ptr or new_ptr == 0):
+            self.full[0] = True
+
+    def get_queue(self) -> torch.Tensor:
+        """Return valid queue entries [Q, dim]."""
+        if self.full.item():
+            return self.queue
+        ptr = int(self.ptr.item())
+        if ptr == 0:
+            return self.queue[:0]  # empty
+        return self.queue[:ptr]
+
+    @property
+    def current_size(self) -> int:
+        if self.full.item():
+            return self.size
+        return int(self.ptr.item())
--- a/vendor_reference/caption_test/src/losses/multi_infonce.py
+++ b/vendor_reference/caption_test/src/losses/multi_infonce.py
@@ -0,0 +1,203 @@
+from __future__ import annotations
+
+"""InfoNCE loss for cross-view geo-localization with optional text fusion.
+
+Single symmetric InfoNCE between query (drone+text fused) and gallery (satellite).
+Asymmetric weighting: query->gallery weighted higher (real use-case direction).
+
+Supports both learnable temperature (CLIP-style logit_scale) and fixed/scheduled.
+"""
+
+import math
+
+import gin
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def _symmetric_info_nce(
+    emb_a: torch.Tensor,
+    emb_b: torch.Tensor,
+    temperature: float | torch.Tensor,
+    label_smoothing: float,
+    weight_a2b: float = 0.5,
+    weight_b2a: float = 0.5,
+    queue_negatives: torch.Tensor | None = None,
+    hard_mining_k: int = 0,
+) -> torch.Tensor:
+    """Weighted symmetric InfoNCE with optional hard negative queue.
+
+    Args:
+        emb_a: Query embeddings [B, D].
+        emb_b: Gallery embeddings [B, D]. Positives on diagonal.
+        queue_negatives: Extra gallery negatives [Q, D] from memory bank.
+        hard_mining_k: If > 0 and queue is non-empty, use only the top-K
+            hardest (highest-similarity) queue entries per query instead
+            of the full queue. Per-query selection — each row gets its
+            own K negatives gathered via `topk`.
+    """
+    batch_size = emb_a.size(0)
+    emb_a_f = emb_a.float()
+    emb_b_f = emb_b.float()
+
+    if queue_negatives is not None and queue_negatives.shape[0] > 0:
+        queue_f = queue_negatives.float()
+        sim_inbatch = emb_a_f @ emb_b_f.t() / temperature  # [B, B]
+        sim_queue = emb_a_f @ queue_f.t() / temperature    # [B, Q]
+
+        if hard_mining_k > 0 and hard_mining_k < queue_f.shape[0]:
+            # Per-row top-K — each query gets its own hardest negatives.
+            sim_queue, _ = sim_queue.topk(k=hard_mining_k, dim=1)  # [B, K]
+
+        # a→b: [B, B + (Q or K)]. Positive at column `i` for row `i`.
+        logits_a2b = torch.cat([sim_inbatch, sim_queue], dim=1)
+        targets_a = torch.arange(batch_size, device=emb_a.device)
+        loss_a2b = F.cross_entropy(logits_a2b, targets_a, label_smoothing=label_smoothing)
+
+        # b→a: gallery sees B in-batch queries (queue is gallery-side, irrelevant here).
+        logits_b2a = sim_inbatch.t()  # [B, B]
+        targets_b = torch.arange(batch_size, device=emb_a.device)
+        loss_b2a = F.cross_entropy(logits_b2a, targets_b, label_smoothing=label_smoothing)
+    else:
+        logits = emb_a_f @ emb_b_f.t() / temperature
+        targets = torch.arange(batch_size, device=emb_a.device)
+        loss_a2b = F.cross_entropy(logits, targets, label_smoothing=label_smoothing)
+        loss_b2a = F.cross_entropy(logits.t(), targets, label_smoothing=label_smoothing)
+
+    return weight_a2b * loss_a2b + weight_b2a * loss_b2a
+
+
+def cosine_temperature(
+    epoch: int,
+    total_epochs: int,
+    tau_init: float = 0.1,
+    tau_final: float = 0.01,
+) -> float:
+    """Cosine-decay schedule for InfoNCE temperature."""
+    total_epochs = max(total_epochs, 1)
+    progress = min(max(epoch / total_epochs, 0.0), 1.0)
+    cosine = 0.5 * (1.0 + math.cos(math.pi * progress))
+    return tau_final + (tau_init - tau_final) * cosine
+
+
+@gin.configurable
+class InfoNCELoss(nn.Module):
+    """Symmetric InfoNCE with learnable or scheduled temperature.
+
+    Args:
+        temperature_init: Initial temperature value.
+        temperature_final: Final temperature (only used if learnable=False).
+        label_smoothing: Cross-entropy label smoothing.
+        weight_q2g: Weight for query->gallery direction.
+        weight_g2q: Weight for gallery->query direction.
+        learnable_temperature: If True, temperature is a learnable parameter
+            (CLIP-style logit_scale). If False, uses cosine schedule.
+        tau_min: Minimum clamp for learnable temperature.
+        tau_max: Maximum clamp for learnable temperature.
+        hard_mining_k: If > 0, mine top-K hardest negatives per query from
+            the memory bank queue instead of using the full queue. 0 disables
+            mining (queue used whole). Typical values: 256-1024 for queue=4096.
+    """
+
+    def __init__(
+        self,
+        temperature_init: float = 0.07,
+        temperature_final: float = 0.01,
+        label_smoothing: float = 0.1,
+        weight_q2g: float = 0.6,
+        weight_g2q: float = 0.4,
+        learnable_temperature: bool = True,
+        tau_min: float = 0.01,
+        tau_max: float = 0.1,
+        hard_mining_k: int = 0,
+    ) -> None:
+        super().__init__()
+        self.temperature_init = temperature_init
+        self.temperature_final = temperature_final
+        self.label_smoothing = label_smoothing
+        self.weight_q2g = weight_q2g
+        self.weight_g2q = weight_g2q
+        self.learnable_temperature = learnable_temperature
+        self.tau_min = tau_min
+        self.tau_max = tau_max
+        self.hard_mining_k = hard_mining_k
+
+        if learnable_temperature:
+            # Store as log(1/tau) like CLIP's logit_scale.
+            init_logit_scale = math.log(1.0 / temperature_init)
+            self.logit_scale = nn.Parameter(torch.tensor(init_logit_scale))
+        else:
+            self.logit_scale = None
+
+    @property
+    def current_temperature(self) -> float:
+        """Current temperature value (for logging)."""
+        if self.logit_scale is not None:
+            tau = 1.0 / self.logit_scale.exp().clamp(
+                min=1.0 / self.tau_max, max=1.0 / self.tau_min,
+            ).item()
+            return tau
+        return self.temperature_init
+
+    def forward(
+        self,
+        embeddings: dict[str, torch.Tensor],
+        epoch: int,
+        total_epochs: int,
+        queue_negatives: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Compute InfoNCE loss with optional hard negative queue.
+
+        Args:
+            embeddings: Dict with 'query' and 'gallery' [B, D] L2-normalized.
+            epoch: Current epoch (0-indexed).
+            total_epochs: Total epochs for temperature schedule.
+            queue_negatives: Extra gallery negatives [Q, D] from memory bank.
+
+        Returns:
+            Dict with 'total', 'temperature', 'gate_q', 'gate_g'.
+        """
+        if self.learnable_temperature:
+            # Clamp logit_scale in logit space first to prevent exp() overflow in fp16.
+            # tau_min=0.01 -> max logit_scale=ln(1/0.01)=4.6
+            # tau_max=0.1  -> min logit_scale=ln(1/0.1)=2.30
+            clamped = self.logit_scale.float().clamp(
+                min=math.log(1.0 / self.tau_max),
+                max=math.log(1.0 / self.tau_min),
+            )
+            logit_scale = clamped.exp()
+            tau = 1.0 / logit_scale
+        else:
+            tau = cosine_temperature(
+                epoch=epoch,
+                total_epochs=total_epochs,
+                tau_init=self.temperature_init,
+                tau_final=self.temperature_final,
+            )
+
+        loss = _symmetric_info_nce(
+            emb_a=embeddings["query"],
+            emb_b=embeddings["gallery"],
+            temperature=tau,
+            label_smoothing=self.label_smoothing,
+            weight_a2b=self.weight_q2g,
+            weight_b2a=self.weight_g2q,
+            queue_negatives=queue_negatives,
+            hard_mining_k=self.hard_mining_k,
+        )
+
+        gate_q = embeddings.get("gate_q", embeddings.get("gate", 1.0))
+        gate_g = embeddings.get("gate_g", 1.0)
+
+        if isinstance(tau, float):
+            tau_out = torch.tensor(tau, device=loss.device)
+        else:
+            tau_out = tau.detach().clone()
+
+        return {
+            "total": loss,
+            "temperature": tau_out,
+            "gate_q": torch.tensor(gate_q, device=loss.device),
+            "gate_g": torch.tensor(gate_g, device=loss.device),
+        }
--- a/vendor_reference/caption_test/src/losses/weighted_infonce.py
+++ b/vendor_reference/caption_test/src/losses/weighted_infonce.py
@@ -0,0 +1,148 @@
+from __future__ import annotations
+
+"""Weighted InfoNCE loss for GTA-UAV cross-view geo-localization.
+
+Adapted from Game4Loc (https://github.com/Yux1angJi/GTA-UAV).
+Uses per-sample label smoothing based on positive_weights (IoU/distance)
+to handle partial overlap between drone and satellite crops.
+
+Standard InfoNCE assumes strict 1-to-1 pairs and treats all non-diagonal
+entries as negatives. In GTA-UAV, multiple satellite crops can validly
+match one drone image (partial IoU overlap), causing false negatives.
+WeightedInfoNCE softens this with adaptive label smoothing per sample.
+"""
+
+import math
+
+import gin
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+@gin.configurable
+class WeightedInfoNCELoss(nn.Module):
+    """Weighted InfoNCE with adaptive per-sample label smoothing.
+
+    For each sample i, eps_i = 1 - (1 - base_smoothing) / (1 + exp(-k * w_i))
+    where w_i is the positive weight (e.g. IoU with matched satellite crop).
+    Higher weight → lower eps → sharper target (strong positive).
+    Lower weight → higher eps → softer target (weak/semi-positive).
+
+    Args:
+        temperature_init: Initial temperature (or learnable logit_scale).
+        learnable_temperature: If True, temperature is learnable (CLIP-style).
+        label_smoothing: Base label smoothing (used when no weights provided).
+        k: Sigmoid steepness for weight → eps mapping.
+        tau_min: Min clamp for learnable temperature.
+        tau_max: Max clamp for learnable temperature.
+    """
+
+    def __init__(
+        self,
+        temperature_init: float = 0.07,
+        learnable_temperature: bool = True,
+        label_smoothing: float = 0.1,
+        k: float = 5.0,
+        tau_min: float = 0.01,
+        tau_max: float = 0.1,
+    ) -> None:
+        super().__init__()
+        self.label_smoothing = label_smoothing
+        self.k = k
+        self.tau_min = tau_min
+        self.tau_max = tau_max
+        self.learnable_temperature = learnable_temperature
+
+        if learnable_temperature:
+            self.logit_scale = nn.Parameter(
+                torch.tensor(math.log(1.0 / temperature_init))
+            )
+        else:
+            self.logit_scale = None
+            self.temperature = temperature_init
+
+    @property
+    def current_temperature(self) -> float:
+        if self.logit_scale is not None:
+            tau = 1.0 / self.logit_scale.exp().clamp(
+                min=1.0 / self.tau_max, max=1.0 / self.tau_min,
+            ).item()
+            return tau
+        return self.temperature
+
+    def _compute_eps(self, positive_weights: torch.Tensor | None, n: int) -> torch.Tensor | list[float]:
+        """Compute per-sample label smoothing from positive weights."""
+        if positive_weights is not None:
+            # Higher weight → lower eps (sharper, stronger positive).
+            return 1.0 - (1.0 - self.label_smoothing) / (1.0 + torch.exp(-self.k * positive_weights))
+        return [self.label_smoothing] * n
+
+    def _weighted_loss(
+        self,
+        sim_matrix: torch.Tensor,
+        eps_all: torch.Tensor | list[float],
+    ) -> torch.Tensor:
+        """Weighted InfoNCE: per-sample interpolation between hard and uniform targets.
+
+        For each row i:
+          L_i = (1-eps_i) * [-sim[i,i] + logsumexp(sim[i,:])]
+              + eps_i * [-mean(sim[i,:]) + logsumexp(sim[i,:])]
+        """
+        n = sim_matrix.shape[0]
+        total_loss = torch.tensor(0.0, device=sim_matrix.device)
+        for i in range(n):
+            eps = eps_all[i] if isinstance(eps_all, list) else eps_all[i]
+            logsumexp = torch.logsumexp(sim_matrix[i, :], dim=0)
+            total_loss += (1 - eps) * (-sim_matrix[i, i] + logsumexp)
+            total_loss += eps * (-sim_matrix[i, :].mean() + logsumexp)
+        return total_loss / n
+
+    def forward(
+        self,
+        embeddings: dict[str, torch.Tensor],
+        epoch: int = 0,
+        total_epochs: int = 1,
+        positive_weights: torch.Tensor | None = None,
+        queue_negatives: torch.Tensor | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Compute weighted InfoNCE loss.
+
+        Args:
+            embeddings: Dict with 'query' [B,D], 'gallery' [B,D], 'gate_q', 'gate_g'.
+            positive_weights: Per-sample weight [B] (e.g. IoU with matched sat crop).
+            queue_negatives: Extra negatives [Q,D] from memory bank (not used with weighted loss).
+        """
+        query = embeddings["query"].float()
+        gallery = embeddings["gallery"].float()
+
+        # Temperature.
+        if self.learnable_temperature:
+            clamped = self.logit_scale.float().clamp(
+                min=math.log(1.0 / self.tau_max),
+                max=math.log(1.0 / self.tau_min),
+            )
+            logit_scale = clamped.exp()
+            tau = 1.0 / logit_scale
+        else:
+            logit_scale = 1.0 / self.temperature
+            tau = self.temperature
+
+        sim_q2g = logit_scale * query @ gallery.t()
+        sim_g2q = sim_q2g.t()
+
+        eps = self._compute_eps(positive_weights, query.shape[0])
+
+        loss_q2g = self._weighted_loss(sim_q2g, eps)
+        loss_g2q = self._weighted_loss(sim_g2q, eps)
+        total = (loss_q2g + loss_g2q) / 2
+
+        gate_q = embeddings.get("gate_q", 1.0)
+        gate_g = embeddings.get("gate_g", 1.0)
+
+        return {
+            "total": total,
+            "temperature": tau if isinstance(tau, torch.Tensor) else torch.tensor(tau, device=total.device),
+            "gate_q": torch.tensor(gate_q, device=total.device) if not isinstance(gate_q, torch.Tensor) else gate_q.detach(),
+            "gate_g": torch.tensor(gate_g, device=total.device) if not isinstance(gate_g, torch.Tensor) else gate_g.detach(),
+        }
--- a/vendor_reference/caption_test/src/models/asymmetric_encoder.py
+++ b/vendor_reference/caption_test/src/models/asymmetric_encoder.py
@@ -0,0 +1,656 @@
+from __future__ import annotations
+
+"""Asymmetric dual encoder for CVGL caption test on GTA-UAV.
+
+Architecture:
+  Query:   DINOv3 ViT-L/16 (LVD, frozen) + LRSCLIP text (L1/L2/L3) -> GatedFusion -> query
+  Gallery: DINOv3 ViT-L/16 (SAT, frozen) -> gallery
+  Loss:    InfoNCE(query, gallery)
+
+DINOv3 checkpoints use a custom key layout (not HuggingFace transformers).
+LRSCLIP (DGTRS-CLIP ViT-L-14) uses open_clip layout with KPS positional embeddings.
+"""
+
+import logging
+import math
+import warnings
+from pathlib import Path
+
+import coloredlogs
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+LOGGER = logging.getLogger("caption_test.model")
+coloredlogs.install(level="INFO", logger=LOGGER, fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
+from safetensors.torch import load_file as load_safetensors
+
+from src.models.adapters import inject_lora_into_dgtrs, inject_mona_into_dinov3
+from src.models.dgtrs.model import DGTRSTextEncoder, load_dgtrs_text_encoder, tokenize_dgtrs
+from src.models.dual_encoder import GatedFusion, ProjectionHead
+from src.models.stripnet import inject_conv_mona_into_stripnet
+from src.models.stripnet_encoder import StripNetEncoder
+
+
+# ---------------------------------------------------------------------------
+# DINOv3 ViT-L/16 — minimal implementation matching checkpoint key layout
+# ---------------------------------------------------------------------------
+
+class DINOv3Attention(nn.Module):
+    """Multi-head self-attention with separate Q/K/V projections."""
+
+    def __init__(self, dim: int = 1024, num_heads: int = 16) -> None:
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        self.scale = self.head_dim ** -0.5
+        self.q_proj = nn.Linear(dim, dim)
+        self.k_proj = nn.Linear(dim, dim, bias=False)
+        self.v_proj = nn.Linear(dim, dim)
+        self.o_proj = nn.Linear(dim, dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B, N, C = x.shape
+        q = self.q_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        k = self.k_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        v = self.v_proj(x).reshape(B, N, self.num_heads, self.head_dim).permute(0, 2, 1, 3)
+        attn = F.scaled_dot_product_attention(q, k, v)
+        x = attn.permute(0, 2, 1, 3).reshape(B, N, C)
+        return self.o_proj(x)
+
+
+class DINOv3LayerScale(nn.Module):
+    """Per-channel learnable scale (lambda)."""
+
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.lambda1 = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * self.lambda1
+
+
+class DINOv3MLP(nn.Module):
+    """SwiGLU-like MLP: up_proj + GELU + down_proj."""
+
+    def __init__(self, dim: int = 1024, mlp_dim: int = 4096) -> None:
+        super().__init__()
+        self.up_proj = nn.Linear(dim, mlp_dim)
+        self.down_proj = nn.Linear(mlp_dim, dim)
+        self.act = nn.GELU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.down_proj(self.act(self.up_proj(x)))
+
+
+class DINOv3Block(nn.Module):
+    """Single DINOv3 transformer block."""
+
+    def __init__(self, dim: int = 1024, num_heads: int = 16, mlp_dim: int = 4096) -> None:
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim)
+        self.attention = DINOv3Attention(dim, num_heads)
+        self.layer_scale1 = DINOv3LayerScale(dim)
+        self.norm2 = nn.LayerNorm(dim)
+        self.mlp = DINOv3MLP(dim, mlp_dim)
+        self.layer_scale2 = DINOv3LayerScale(dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.layer_scale1(self.attention(self.norm1(x)))
+        x = x + self.layer_scale2(self.mlp(self.norm2(x)))
+        return x
+
+
+class DINOv3Embeddings(nn.Module):
+    """Patch embedding + CLS token + register tokens."""
+
+    def __init__(
+        self,
+        dim: int = 1024,
+        patch_size: int = 16,
+        num_registers: int = 4,
+    ) -> None:
+        super().__init__()
+        self.patch_embeddings = nn.Conv2d(3, dim, patch_size, patch_size)
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, dim))
+        self.register_tokens = nn.Parameter(torch.zeros(1, num_registers, dim))
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        B = x.shape[0]
+        patches = self.patch_embeddings(x).flatten(2).transpose(1, 2)  # [B, N, D]
+        N = patches.shape[1]
+
+        cls = self.cls_token.expand(B, -1, -1)
+        reg = self.register_tokens.expand(B, -1, -1)
+
+        # DINOv3: [CLS, registers, patches]
+        x = torch.cat([cls, reg, patches], dim=1)
+
+        # Positional embedding: interpolated sincos (RoPE applied in attention
+        # in original, but pretrained checkpoints bake it into weights).
+        # We use a simple learned-style pos embed computed on the fly.
+        pos = self._get_pos_embed(N, x.device, x.dtype)
+        # pos covers patches only, skip CLS + registers
+        x[:, 1 + reg.shape[1]:] = x[:, 1 + reg.shape[1]:] + pos
+
+        return x
+
+    def _get_pos_embed(self, n_patches: int, device: torch.device, dtype: torch.dtype) -> torch.Tensor:
+        # DINOv3 uses RoPE internally — no additive pos embed needed.
+        # Return zeros as placeholder (weights handle positioning via RoPE).
+        return torch.zeros(1, n_patches, self.cls_token.shape[-1], device=device, dtype=dtype)
+
+
+class DINOv3ViT(nn.Module):
+    """DINOv3 ViT-L/16 matching the checkpoint key layout.
+
+    Checkpoint keys:
+        embeddings.cls_token, embeddings.patch_embeddings.{weight,bias},
+        embeddings.register_tokens, embeddings.mask_token,
+        layer.{i}.attention.{q,k,v,o}_proj.{weight,bias},
+        layer.{i}.layer_scale{1,2}.lambda1,
+        layer.{i}.mlp.{up,down}_proj.{weight,bias},
+        layer.{i}.norm{1,2}.{weight,bias},
+        norm.{weight,bias}
+    """
+
+    def __init__(
+        self,
+        dim: int = 1024,
+        num_heads: int = 16,
+        mlp_dim: int = 4096,
+        num_layers: int = 24,
+        patch_size: int = 16,
+        num_registers: int = 4,
+    ) -> None:
+        super().__init__()
+        self.embeddings = DINOv3Embeddings(dim, patch_size, num_registers)
+        self.layer = nn.ModuleList([
+            DINOv3Block(dim, num_heads, mlp_dim) for _ in range(num_layers)
+        ])
+        self.norm = nn.LayerNorm(dim)
+        self.embed_dim = dim
+        self.gradient_checkpointing = False
+
+    def set_gradient_checkpointing(self, enable: bool = True) -> None:
+        """Enable/disable gradient checkpointing to trade compute for VRAM."""
+        self.gradient_checkpointing = enable
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass. Returns CLS token embedding [B, dim]."""
+        x = self.embeddings(x)
+        for block in self.layer:
+            if self.gradient_checkpointing and self.training:
+                x = torch.utils.checkpoint.checkpoint(
+                    block, x, use_reentrant=False,
+                )
+            else:
+                x = block(x)
+        x = self.norm(x)
+        return x[:, 0]  # CLS token
+
+    @classmethod
+    def from_pretrained(cls, path: str | Path) -> DINOv3ViT:
+        """Load from .pth or .safetensors checkpoint."""
+        model = cls()
+        path = Path(path)
+        LOGGER.info("🧊 Loading DINOv3 from %s", path.name)
+        if path.suffix == ".safetensors":
+            state = load_safetensors(str(path))
+        else:
+            state = torch.load(str(path), map_location="cpu", weights_only=False)
+            if "model" in state:
+                state = state["model"]
+            elif "state_dict" in state:
+                state = state["state_dict"]
+        model.load_state_dict(state, strict=False)
+        n_params = sum(p.numel() for p in model.parameters())
+        LOGGER.info("🧊 DINOv3 loaded: %s params", f"{n_params:,}")
+        return model
+
+
+# LRSCLIPTextEncoder removed — replaced by official DGTRS architecture
+# in src/models/dgtrs/model.py (DGTRSTextEncoder)
+
+
+# ---------------------------------------------------------------------------
+# Text fusion MLP: concat L1/L2/L3 -> project to D
+# ---------------------------------------------------------------------------
+
+class TextFusionMLP(nn.Module):
+    """Fuse L1/L2/L3 text embeddings via concat + MLP.
+
+    [B, 3*text_dim] -> [B, out_dim]
+    """
+
+    def __init__(
+        self,
+        text_dim: int = 768,
+        out_dim: int = 1024,
+    ) -> None:
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(3 * text_dim, out_dim),
+            nn.GELU(),
+            nn.Linear(out_dim, out_dim),
+        )
+
+    def forward(
+        self,
+        z_l1: torch.Tensor,
+        z_l2: torch.Tensor,
+        z_l3: torch.Tensor,
+    ) -> torch.Tensor:
+        """Fuse three text embeddings.
+
+        Args:
+            z_l1: L1 overview [B, text_dim].
+            z_l2: L2 full description [B, text_dim].
+            z_l3: L3 fingerprint [B, text_dim].
+
+        Returns:
+            Fused text embedding [B, out_dim].
+        """
+        cat = torch.cat([z_l1, z_l2, z_l3], dim=-1)
+        return self.mlp(cat)
+
+
+# ---------------------------------------------------------------------------
+# Main model: AsymmetricEncoder
+# ---------------------------------------------------------------------------
+
+# ResidualGateFusin experiment
+from .residual_fusions import ResidualGateType, GatedFusionResidual
+
+class AsymmetricEncoder(nn.Module):
+    """Dual encoder for CVGL with text fusion on both branches.
+
+    Supports two modes:
+      - **shared** (default): single DINOv3 WEB encoder for both drone and satellite,
+        one set of MONA adapters. Saves ~4-5 GB VRAM and halves adapter params.
+      - **asymmetric**: separate DINOv3 encoders (LVD for drone, SAT for satellite),
+        each with their own MONA adapters (legacy mode).
+
+    Query branch:   DINOv3 (drone) + text(L1/L2/L3) -> GatedFusion_q -> query [1024]
+    Gallery branch: DINOv3 (sat)   + text(L1/L2/L3) -> GatedFusion_g -> gallery [1024]
+
+    No projection layers — retrieval space is DINOv3 native 1024-dim.
+    Text fusion MLP is shared between branches (same caption format).
+    Two separate GatedFusion gates (drone/sat may weight text differently).
+
+    For satellite images without captions, GatedFusion passes image features through
+    (text_feat=None → gate acts as identity).
+
+    Args:
+        dino_web_path: Path to DINOv3 LVD checkpoint (used for both branches in shared mode).
+        dino_sat_path: Path to DINOv3 SAT checkpoint (only used in asymmetric mode).
+        lrsclip_path: Path to DGTRS-CLIP checkpoint (text encoder).
+        init_gate: Initial fusion gate (image weight).
+        baseline_mode: If True, gate = 1.0 (text ignored), DGTRS not loaded.
+        shared_encoder: If True, use single DINOv3 WEB for both branches.
+        device: Torch device string.
+    """
+
+    DINO_DIM = 1024
+    TEXT_DIM = 768
+
+    def __init__(
+        self,
+        # !!! ---------------------------------------------------------------
+        gate_type: ResidualGateType = ResidualGateType.simple_residual_one_gate,
+        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",
+        init_gate: float = 0.7,
+        baseline_mode: bool = False,
+        shared_encoder: bool = False,
+        mona_bottleneck: int = 64,
+        mona_last_n_blocks: int = 24,
+        lora_rank: int = 4,
+        device: str = "cuda",
+        backbone: str = "dinov3",
+        stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth",
+        stripnet_mona_last_n_stages: int = 2,
+        stripnet_freeze: bool = True,
+    ) -> None:
+        super().__init__()
+        self.embed_dim = self.DINO_DIM  # native 1024 (StripNet projects 512 -> 1024)
+        self.baseline_mode = baseline_mode
+        self.shared_encoder = shared_encoder
+        self.backbone = backbone
+        self.device = device
+
+        # Image encoder(s) (frozen + MONA adapters).
+        if backbone == "stripnet":
+            # StripNet always operates as shared encoder (one CNN for both branches).
+            self.shared_encoder = True
+            self.image_encoder = StripNetEncoder(checkpoint_path=stripnet_path, out_dim=self.DINO_DIM)
+            if stripnet_freeze:
+                self._freeze(self.image_encoder.backbone)
+                LOGGER.info("StripNet backbone: frozen (Conv-MONA + projection trainable)")
+            else:
+                LOGGER.info("StripNet backbone: UNFROZEN — full fine-tune (use lower lr factor)")
+            if stripnet_mona_last_n_stages > 0:
+                inject_conv_mona_into_stripnet(
+                    self.image_encoder.backbone,
+                    bottleneck=mona_bottleneck,
+                    last_n_stages=stripnet_mona_last_n_stages,
+                )
+            else:
+                LOGGER.info("Conv-MONA disabled (stripnet_mona_last_n_stages=0)")
+            LOGGER.info("StripNet backbone: shared encoder, projection 512 -> %d", self.DINO_DIM)
+        elif shared_encoder:
+            self.image_encoder = DINOv3ViT.from_pretrained(dino_web_path)
+            self._freeze(self.image_encoder)
+            inject_mona_into_dinov3(self.image_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
+            LOGGER.info("Shared encoder mode: single DINOv3 WEB for drone + satellite")
+        else:
+            self.drone_encoder = DINOv3ViT.from_pretrained(dino_web_path)
+            self.sat_encoder = DINOv3ViT.from_pretrained(dino_sat_path)
+            self._freeze(self.drone_encoder)
+            self._freeze(self.sat_encoder)
+            inject_mona_into_dinov3(self.drone_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
+            inject_mona_into_dinov3(self.sat_encoder, bottleneck=mona_bottleneck, last_n_blocks=mona_last_n_blocks)
+            LOGGER.info("Asymmetric encoder mode: DINOv3 WEB (drone) + DINOv3 SAT (satellite)")
+
+        # Text encoder — official DGTRS architecture (frozen + LoRA).
+        if not baseline_mode:
+            self.text_encoder = load_dgtrs_text_encoder(lrsclip_path)
+            self._freeze(self.text_encoder)
+            inject_lora_into_dgtrs(self.text_encoder, rank=lora_rank)
+        else:
+            self.text_encoder = None
+
+        # Shared text fusion MLP: 3×768 -> 1024 (native DINOv3 dim).
+        if not baseline_mode:
+            self.text_fusion = TextFusionMLP(
+                text_dim=self.TEXT_DIM,
+                out_dim=self.DINO_DIM,
+            )
+
+        # Separate gated fusion for query and gallery branches.
+        #! Experimental Gated fusion on query branch.
+        self.fusion_query = GatedFusionResidual(gate_type=gate_type,
+            init_gate=init_gate, baseline_mode=baseline_mode)
+        self.fusion_gallery = GatedFusionResidual(gate_type=gate_type,
+            init_gate=init_gate, baseline_mode=baseline_mode)
+
+
+    @staticmethod
+    def _freeze(module: nn.Module) -> None:
+        for p in module.parameters():
+            p.requires_grad = False
+        module.eval()
+
+    def encode_drone(self, images: torch.Tensor) -> torch.Tensor:
+        """Encode drone images with MONA adapters. Returns [B, 1024]."""
+        if self.shared_encoder:
+            return self.image_encoder(images)
+        return self.drone_encoder(images)
+
+    def encode_satellite(self, images: torch.Tensor) -> torch.Tensor:
+        """Encode satellite images with MONA adapters. Returns [B, 1024]."""
+        if self.shared_encoder:
+            return self.image_encoder(images)
+        return self.sat_encoder(images)
+
+    def encode_text_levels(
+        self,
+        l1_texts: list[str],
+        l2_texts: list[str],
+        l3_texts: list[str],
+    ) -> torch.Tensor | None:
+        """Encode L1/L2/L3 captions and fuse. Returns [B, 1024] or None.
+
+        Returns None if all captions are empty (no text available).
+        For mixed batches (some have captions, some don't), encodes all
+        texts (empty strings tokenize to pad+EOS — their outputs must be
+        masked downstream, see `_fuse_with_mask`).
+        """
+        # Check if any caption is non-empty.
+        if all(t == "" for t in l1_texts):
+            return None
+
+        z_l1 = self._encode_single_text(l1_texts)
+        z_l2 = self._encode_single_text(l2_texts)
+        z_l3 = self._encode_single_text(l3_texts)
+        fused = self.text_fusion(z_l1, z_l2, z_l3)
+        return F.normalize(fused, dim=-1)
+
+    def _encode_single_text(self, texts: list[str]) -> torch.Tensor:
+        """Tokenize and encode a list of strings using DGTRS tokenizer."""
+        tokens = tokenize_dgtrs(list(texts)).to(self.device)
+        return self.text_encoder(tokens)
+
+    def _fuse_with_mask(
+        self,
+        img_feat: torch.Tensor,
+        l1_texts: list[str] | None,
+        l2_texts: list[str] | None,
+        l3_texts: list[str] | None,
+        fusion: GatedFusionResidual,
+    ) -> torch.Tensor:
+        """Fuse image features with optional text, respecting per-sample presence.
+
+        For samples where caption is an empty string, output falls back to
+        pure image features (avoiding noise contamination from empty-string
+        text embeddings). For samples with captions, applies the standard
+        gated fusion `σ(α)·img + (1-σ(α))·text`.
+
+        Returns L2-normalized [B, D] embedding.
+        """
+        if (
+            self.baseline_mode
+            or l1_texts is None
+            or l2_texts is None
+            or l3_texts is None
+        ):
+            return F.normalize(fusion(img_feat, None), dim=-1)
+
+        has_text = torch.tensor(
+            [t != "" for t in l1_texts], dtype=torch.bool, device=img_feat.device,
+        )
+        if not has_text.any():
+            return F.normalize(fusion(img_feat, None), dim=-1)
+
+        z_text = self.encode_text_levels(l1_texts, l2_texts, l3_texts)
+        if z_text is None:
+            return F.normalize(fusion(img_feat, None), dim=-1)
+
+        # Per-sample fusion: text-present samples use full gated fusion,
+        # empty-caption samples pass through pure image features.
+        fused_with_text = fusion(img_feat, z_text)
+        out = torch.where(has_text.unsqueeze(-1), fused_with_text, img_feat)
+        return F.normalize(out, dim=-1)
+
+    def encode_query(
+        self,
+        drone_img: torch.Tensor,
+        caption_l1: list[str] | None = None,
+        caption_l2: list[str] | None = None,
+        caption_l3: list[str] | None = None,
+    ) -> torch.Tensor:
+        """Encode drone → normalized query embedding with per-sample text mask."""
+        drone_feat = self.encode_drone(drone_img)
+        return self._fuse_with_mask(
+            drone_feat, caption_l1, caption_l2, caption_l3, self.fusion_query,
+        )
+
+    def encode_gallery(
+        self,
+        sat_img: torch.Tensor,
+        sat_caption_l1: list[str] | None = None,
+        sat_caption_l2: list[str] | None = None,
+        sat_caption_l3: list[str] | None = None,
+    ) -> torch.Tensor:
+        """Encode satellite → normalized gallery embedding with per-sample text mask."""
+        sat_feat = self.encode_satellite(sat_img)
+        return self._fuse_with_mask(
+            sat_feat, sat_caption_l1, sat_caption_l2, sat_caption_l3, self.fusion_gallery,
+        )
+
+    def forward(
+        self,
+        drone_img: torch.Tensor,
+        sat_img: torch.Tensor,
+        caption_l1: list[str] | None = None,
+        caption_l2: list[str] | None = None,
+        caption_l3: list[str] | None = None,
+        sat_caption_l1: list[str] | None = None,
+        sat_caption_l2: list[str] | None = None,
+        sat_caption_l3: list[str] | None = None,
+    ) -> dict[str, torch.Tensor]:
+        """Forward pass.
+
+        Both branches use per-sample caption masking: samples with an empty
+        caption string fall back to pure image features instead of being
+        fused with noise from empty-string text embeddings.
+
+        Args:
+            drone_img: Drone images [B, 3, 256, 256].
+            sat_img: Satellite images [B, 3, 256, 256].
+            caption_l1/l2/l3: Drone L1/L2/L3 captions.
+            sat_caption_l1/l2/l3: Satellite L1/L2/L3 captions.
+
+        Returns:
+            Dict with 'query' [B, embed_dim], 'gallery' [B, embed_dim],
+            'gate_q', 'gate_g'.
+        """
+        query = self.encode_query(drone_img, caption_l1, caption_l2, caption_l3)
+        gallery = self.encode_gallery(sat_img, sat_caption_l1, sat_caption_l2, sat_caption_l3)
+        return {
+            "query": query,
+            "gallery": gallery,
+            "gate_q": self.fusion_query.gate_value,
+            "gate_g": self.fusion_gallery.gate_value,
+        }
+
+    def trainable_parameters(self) -> list[nn.Parameter]:
+        """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(),
+            "baseline_mode": self.baseline_mode,
+            "shared_encoder": self.shared_encoder,
+            **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,
+            baseline_mode=ckpt.get("baseline_mode", False),
+            shared_encoder=ckpt.get("shared_encoder", False),
+            mona_bottleneck=ckpt.get("mona_bottleneck", 64),
+            mona_last_n_blocks=ckpt.get("mona_last_n_blocks", 24),
+            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)
+        if self.shared_encoder:
+            self.image_encoder.eval()
+        else:
+            self.drone_encoder.eval()
+            self.sat_encoder.eval()
+        if self.text_encoder is not None:
+            self.text_encoder.train(mode)
+        return self
+
+
+# ---------------------------------------------------------------------------
+# Image preprocessing (DINOv3: 256x256, ImageNet normalization)
+# ---------------------------------------------------------------------------
+
+_IMAGENET_MEAN = [0.485, 0.456, 0.406]
+_IMAGENET_STD = [0.229, 0.224, 0.225]
+
+
+def get_dino_transform(image_size: int = 256) -> torch.nn.Module:
+    """Build eval/inference image transform for DINOv3 input."""
+    from torchvision import transforms
+    return transforms.Compose([
+        transforms.Resize(image_size, interpolation=transforms.InterpolationMode.BICUBIC),
+        transforms.CenterCrop(image_size),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
+    ])
+
+
+def get_drone_train_transform(image_size: int = 256) -> torch.nn.Module:
+    """Build training augmentation for drone images.
+
+    Includes: RandomResizedCrop, HFlip, rotation, color jitter,
+    grayscale, Gaussian blur.
+    """
+    from torchvision import transforms
+    return transforms.Compose([
+        transforms.RandomResizedCrop(
+            image_size, scale=(0.7, 1.0),
+            interpolation=transforms.InterpolationMode.BICUBIC,
+        ),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.RandomRotation(degrees=15),
+        transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1),
+        transforms.RandomGrayscale(p=0.05),
+        transforms.GaussianBlur(kernel_size=3, sigma=(0.1, 2.0)),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
+    ])
+
+
+def get_satellite_train_transform(image_size: int = 256) -> torch.nn.Module:
+    """Build training augmentation for satellite images.
+
+    Lighter than drone: no rotation or blur (satellite is nadir/consistent).
+    Includes: RandomResizedCrop, HFlip, color jitter, grayscale.
+    """
+    from torchvision import transforms
+    return transforms.Compose([
+        transforms.RandomResizedCrop(
+            image_size, scale=(0.7, 1.0),
+            interpolation=transforms.InterpolationMode.BICUBIC,
+        ),
+        transforms.RandomHorizontalFlip(p=0.5),
+        transforms.ColorJitter(brightness=0.3, contrast=0.3, saturation=0.2, hue=0.1),
+        transforms.RandomGrayscale(p=0.05),
+        transforms.ToTensor(),
+        transforms.Normalize(mean=_IMAGENET_MEAN, std=_IMAGENET_STD),
+    ])
--- a/vendor_reference/caption_test/src/models/residual_fusions.py
+++ b/vendor_reference/caption_test/src/models/residual_fusions.py
@@ -0,0 +1,154 @@
+import torch
+import torch.nn as nn
+import gin
+
+#! GATE-FUSIONS MODIFICATIONS ---------------------------------
+#! in_dim = 1024
+
+from enum import Enum
+import math
+
+class ResidualGateType(Enum):
+    simple_residual_one_gate = 0,
+    cross_gate = 1,
+    gate_sum = 2,
+    alpha_res_cat = 3,
+    alpha_res_sum = 4
+
+# TODO: add GatedFusionresidual class to gin
+
+def init_bias_for_sigmoid(linear: nn.Linear, value: float) -> None:
+    nn.init.zeros_(linear.weight)
+    nn.init.constant_(linear.bias, math.log(value / (1.0 - value)))
+
+def init_residual_projs(linear: nn.Linear, scale: float) -> None:
+    nn.init.xavier_uniform_(linear.weight, gain=scale)
+    nn.init.zeros_(linear.bias)
+
+RESIDUAL_GATES = {
+    ResidualGateType.alpha_res_sum,
+    ResidualGateType.alpha_res_cat
+}
+
+@gin.configurable
+class GatedFusionResidual(nn.Module):
+    """Learnable gated fusion of image and text embeddings.
+
+    V1 - Simple residual gating with 1 common gate: 
+    V2 - Cross residual gating with 2 cross-gates
+    V3 - Gate + Simple Sum of feats x & y
+    V4 - Alpha-weighted residual sum (per sample)
+    V5 - Alpha-weighted residual concat (per sample)
+    """
+
+    def __init__(self, gate_type: ResidualGateType, 
+                 init_gate: float = 0.7, in_dim = 1024, 
+                 init_res_weight: float = 0.1, residual_proj_scale: float = 0.1,  
+                 baseline_mode: bool = False, 
+                 ) -> None:
+        super().__init__()
+        # alpha is in logit space: sigmoid(alpha) = init_gate
+        init_alpha = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
+        self.alpha = nn.Parameter(init_alpha)
+        
+        # alphas for separated cases
+        if gate_type == ResidualGateType.cross_gate:
+            init_alpha_img_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
+            init_alpha_text_cross_gate = torch.log(torch.tensor(init_gate / (1.0 - init_gate)))
+            self.alpha_img = nn.Parameter(init_alpha_img_cross_gate)
+            self.alpha_text = nn.Parameter(init_alpha_text_cross_gate)
+        
+        # weight for sum and cat residual 
+        if gate_type in RESIDUAL_GATES:
+            self.final_cat_residual_proj = nn.Linear(in_dim * 2, in_dim)
+            self.weight_net_for_sum = nn.Linear(in_dim, 1)
+            self.weight_net_for_cat = nn.Linear(in_dim * 2, 1)
+            init_bias_for_sigmoid(self.weight_net_for_sum, value=init_res_weight)
+            init_bias_for_sigmoid(self.weight_net_for_cat, value=init_res_weight)
+            init_residual_projs(self.final_cat_residual_proj, scale=residual_proj_scale)
+
+        self.gate_type = gate_type
+        self.baseline_mode = baseline_mode
+
+    def FuseSRGF(self,
+        img_feat: torch.Tensor,
+        text_feat: torch.Tensor | None,
+    ) -> torch.Tensor:
+        gate = torch.sigmoid(self.alpha)
+        img_res = img_feat * gate + img_feat
+        text_res = text_feat * (1 - gate) + text_feat
+        fused_vec = img_res + text_res
+        return fused_vec
+    
+    def FuseRCGF(self,
+        img_feat: torch.Tensor,
+        text_feat: torch.Tensor | None,
+    ) -> torch.Tensor:
+        gate_img = torch.sigmoid(self.alpha_img)
+        gate_text = torch.sigmoid(self.alpha_text)
+        z_img = img_feat + gate_text * img_feat
+        z_text = text_feat + gate_img * text_feat
+        fused_vec = z_img + z_text
+        return fused_vec
+    
+    def FuseGSUM(self,
+        img_feat: torch.Tensor,
+        text_feat: torch.Tensor | None,
+    ) -> torch.Tensor:
+        gate = torch.sigmoid(self.alpha)
+        fuzed_vec = img_feat + text_feat + gate * img_feat + (1.0 - gate) * text_feat
+        return fuzed_vec
+    
+    def FuseARGFSum(
+        self,
+        img_feat: torch.Tensor,
+        text_feat: torch.Tensor | None,
+    ) -> torch.Tensor:
+        gate = torch.sigmoid(self.alpha)
+        residual = img_feat + text_feat
+        res_weight = torch.sigmoid(self.weight_net_for_sum(residual))
+        fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual
+        return fuzed_vec
+
+    def FuseARGFCat(
+        self,
+        img_feat: torch.Tensor,
+        text_feat: torch.Tensor | None,
+    ) -> torch.Tensor:
+        gate = torch.sigmoid(self.alpha)
+        cat_vec = torch.cat([img_feat, text_feat], dim=-1)
+        residual = self.final_cat_residual_proj(cat_vec)
+        res_weight = torch.sigmoid(self.weight_net_for_cat(cat_vec))
+        fuzed_vec = gate * img_feat + (1.0 - gate) * text_feat + res_weight * residual
+        return fuzed_vec
+
+    def forward(
+        self,
+        img_feat: torch.Tensor,
+        text_feat: torch.Tensor | None,
+    ) -> torch.Tensor:
+        if text_feat is None or self.baseline_mode:
+            return img_feat
+
+        if self.gate_type == ResidualGateType.simple_residual_one_gate:
+            fused_vec = self.FuseSRGF(img_feat=img_feat, text_feat=text_feat)
+        if self.gate_type == ResidualGateType.cross_gate:
+            fused_vec = self.FuseRCGF(img_feat=img_feat, text_feat=text_feat)
+        if self.gate_type == ResidualGateType.gate_sum:
+            fused_vec = self.FuseGSUM(img_feat=img_feat, text_feat=text_feat)
+        if self.gate_type == ResidualGateType.alpha_res_sum:
+            fused_vec = self.FuseARGFSum(img_feat=img_feat, text_feat=text_feat)
+        if self.gate_type == ResidualGateType.alpha_res_cat:
+            fused_vec = self.FuseARGFCat(img_feat=img_feat, text_feat=text_feat)
+
+        return fused_vec
+    
+    @property
+    def gate_value(self) -> float:
+        """Current gate value (image weight). 1.0 = text ignored."""
+        if self.baseline_mode:
+            return 1.0
+        return torch.sigmoid(self.alpha).item()
+    
+
+    
--- a/vendor_reference/caption_test/src/models/stripnet/conv_mona.py
+++ b/vendor_reference/caption_test/src/models/stripnet/conv_mona.py
@@ -0,0 +1,106 @@
+"""Conv-MONA: 2D adaptation of MONA (CVPR 2025) for hierarchical CNN backbones.
+
+MONA paper applies sequence-form adapters after MSA / MLP in ViT blocks. Here we
+mirror that idea in [B, C, H, W] form: BN → 1×1 Down(C→bn) → multi-scale DWConv
+{3,5,7} mean → +residual → GELU → 1×1 Up(bn→C). Layer scale (γ) channel-wise,
+init 1e-6 for near-identity start. Two adapters per StripNet Block: post-attn
+and post-mlp.
+"""
+from __future__ import annotations
+
+import logging
+import torch
+import torch.nn as nn
+
+from src.models.stripnet.model import StripNet, Block
+
+LOGGER = logging.getLogger("caption_test.stripnet.adapters")
+
+
+class ConvMona(nn.Module):
+    """Single Conv-MONA adapter.
+
+    Args:
+        dim: input channel dim.
+        bottleneck: bottleneck channel dim (e.g. 64).
+        gamma_init: layer-scale init value (1e-6 for near-identity at start).
+    """
+
+    def __init__(self, dim: int, bottleneck: int = 64, gamma_init: float = 1e-6) -> None:
+        super().__init__()
+        self.norm = nn.BatchNorm2d(dim)
+        self.down = nn.Conv2d(dim, bottleneck, kernel_size=1, bias=True)
+        self.dw3 = nn.Conv2d(bottleneck, bottleneck, kernel_size=3, padding=1, groups=bottleneck, bias=True)
+        self.dw5 = nn.Conv2d(bottleneck, bottleneck, kernel_size=5, padding=2, groups=bottleneck, bias=True)
+        self.dw7 = nn.Conv2d(bottleneck, bottleneck, kernel_size=7, padding=3, groups=bottleneck, bias=True)
+        self.act = nn.GELU()
+        self.up = nn.Conv2d(bottleneck, dim, kernel_size=1, bias=True)
+        # Channel-wise layer scale (γ), broadcast across H, W.
+        self.gamma = nn.Parameter(gamma_init * torch.ones(dim), requires_grad=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        h = self.norm(x)
+        h = self.down(h)
+        h = (self.dw3(h) + self.dw5(h) + self.dw7(h)) / 3.0 + h
+        h = self.act(h)
+        h = self.up(h)
+        return self.gamma.view(1, -1, 1, 1) * h
+
+
+def _patched_block_forward(block: Block, mona_attn: ConvMona, mona_mlp: ConvMona):
+    """Closure that wraps a Block.forward with two Conv-MONA residuals."""
+    orig_attn = block.attn
+    orig_mlp = block.mlp
+    orig_norm1 = block.norm1
+    orig_norm2 = block.norm2
+    orig_drop = block.drop_path
+    ls1 = block.layer_scale_1
+    ls2 = block.layer_scale_2
+
+    def forward(x: torch.Tensor) -> torch.Tensor:
+        x = x + orig_drop(ls1.unsqueeze(-1).unsqueeze(-1) * orig_attn(orig_norm1(x))) + mona_attn(x)
+        x = x + orig_drop(ls2.unsqueeze(-1).unsqueeze(-1) * orig_mlp(orig_norm2(x))) + mona_mlp(x)
+        return x
+
+    return forward
+
+
+def inject_conv_mona_into_stripnet(
+    model: StripNet,
+    bottleneck: int = 64,
+    last_n_stages: int = 2,
+    use_bf16: bool = False,
+) -> int:
+    """Inject Conv-MONA adapters into the deepest `last_n_stages` of StripNet.
+
+    Each Block in the targeted stages gets two adapters (post-attn, post-mlp).
+    Returns the number of adapters injected.
+
+    Stages are 1-indexed in StripNet (block1..block4). With `last_n_stages=2`
+    we adapt block3 and block4 — the deepest, semantically richest features.
+    """
+    n_stages = model.num_stages
+    target_stages = list(range(max(1, n_stages - last_n_stages + 1), n_stages + 1))
+    n_added = 0
+
+    for stage_idx in target_stages:
+        blocks: nn.ModuleList = getattr(model, f"block{stage_idx}")
+        dim = model.embed_dims[stage_idx - 1]
+        for blk_idx, block in enumerate(blocks):
+            mona_a = ConvMona(dim=dim, bottleneck=bottleneck)
+            mona_m = ConvMona(dim=dim, bottleneck=bottleneck)
+            if use_bf16:
+                mona_a.to(dtype=torch.bfloat16)
+                mona_m.to(dtype=torch.bfloat16)
+            # Register as submodules so they get moved by .to(device) / .train() etc.
+            block.add_module(f"mona_attn", mona_a)
+            block.add_module(f"mona_mlp", mona_m)
+            block.forward = _patched_block_forward(block, mona_a, mona_m)
+            n_added += 2
+
+    n_trainable = sum(p.numel() for p in model.parameters() if p.requires_grad)
+    LOGGER.info(
+        "🔧 Conv-MONA injected: %d adapters in stages %s, %d trainable params (bottleneck=%d)",
+        n_added, target_stages, n_trainable, bottleneck,
+    )
+    return n_added
--- a/vendor_reference/caption_test/src/models/stripnet/model.py
+++ b/vendor_reference/caption_test/src/models/stripnet/model.py
@@ -0,0 +1,253 @@
+"""StripNet (small) backbone — adapted from Strip-R-CNN (HVision-NKU).
+
+Self-contained: no external utils. State-dict naming follows the upstream
+ImageNet-pretrained checkpoint (`conv_spatial1/2` for the strip kernels).
+"""
+from __future__ import annotations
+
+import logging
+import math
+from functools import partial
+from pathlib import Path
+from typing import List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+LOGGER = logging.getLogger("caption_test.stripnet")
+
+
+def _to_2tuple(x):
+    if isinstance(x, (tuple, list)):
+        return tuple(x)
+    return (x, x)
+
+
+def drop_path(x: torch.Tensor, drop_prob: float = 0.0, training: bool = False) -> torch.Tensor:
+    if drop_prob == 0.0 or not training:
+        return x
+    keep = 1.0 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+    rand = x.new_empty(shape).bernoulli_(keep)
+    if keep > 0:
+        rand.div_(keep)
+    return x * rand
+
+
+class DropPath(nn.Module):
+    def __init__(self, p: float = 0.0) -> None:
+        super().__init__()
+        self.p = p
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return drop_path(x, self.p, self.training)
+
+
+class DWConv(nn.Module):
+    def __init__(self, dim: int) -> None:
+        super().__init__()
+        self.dwconv = nn.Conv2d(dim, dim, 3, 1, 1, bias=True, groups=dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.dwconv(x)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features: int, hidden_features: int, drop: float = 0.0) -> None:
+        super().__init__()
+        self.fc1 = nn.Conv2d(in_features, hidden_features, 1)
+        self.dwconv = DWConv(hidden_features)
+        self.act = nn.GELU()
+        self.fc2 = nn.Conv2d(hidden_features, in_features, 1)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.fc1(x)
+        x = self.dwconv(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class StripGatingUnit(nn.Module):
+    """Strip spatial gating: 5x5 DWConv -> (1, k2) -> (k2, 1) -> 1x1 -> gate."""
+
+    def __init__(self, dim: int, k1: int, k2: int) -> None:
+        super().__init__()
+        self.conv0 = nn.Conv2d(dim, dim, 5, padding=2, groups=dim)
+        # Names match upstream pretrained checkpoint: conv_spatial1 / conv_spatial2.
+        self.conv_spatial1 = nn.Conv2d(dim, dim, kernel_size=(k1, k2), stride=1,
+                                       padding=(k1 // 2, k2 // 2), groups=dim)
+        self.conv_spatial2 = nn.Conv2d(dim, dim, kernel_size=(k2, k1), stride=1,
+                                       padding=(k2 // 2, k1 // 2), groups=dim)
+        self.conv1 = nn.Conv2d(dim, dim, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        attn = self.conv0(x)
+        attn = self.conv_spatial1(attn)
+        attn = self.conv_spatial2(attn)
+        attn = self.conv1(attn)
+        return x * attn
+
+
+class StripAttention(nn.Module):
+    def __init__(self, dim: int, k1: int, k2: int) -> None:
+        super().__init__()
+        self.proj_1 = nn.Conv2d(dim, dim, 1)
+        self.activation = nn.GELU()
+        self.spatial_gating_unit = StripGatingUnit(dim, k1, k2)
+        self.proj_2 = nn.Conv2d(dim, dim, 1)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        residual = x
+        x = self.proj_1(x)
+        x = self.activation(x)
+        x = self.spatial_gating_unit(x)
+        x = self.proj_2(x)
+        return x + residual
+
+
+class Block(nn.Module):
+    def __init__(self, dim: int, mlp_ratio: float, k1: int, k2: int, drop: float, drop_path: float) -> None:
+        super().__init__()
+        self.norm1 = nn.BatchNorm2d(dim)
+        self.norm2 = nn.BatchNorm2d(dim)
+        self.attn = StripAttention(dim, k1, k2)
+        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.mlp = Mlp(dim, int(dim * mlp_ratio), drop=drop)
+        ls_init = 1e-2
+        self.layer_scale_1 = nn.Parameter(ls_init * torch.ones(dim), requires_grad=True)
+        self.layer_scale_2 = nn.Parameter(ls_init * torch.ones(dim), requires_grad=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x + self.drop_path(
+            self.layer_scale_1.unsqueeze(-1).unsqueeze(-1) * self.attn(self.norm1(x))
+        )
+        x = x + self.drop_path(
+            self.layer_scale_2.unsqueeze(-1).unsqueeze(-1) * self.mlp(self.norm2(x))
+        )
+        return x
+
+
+class OverlapPatchEmbed(nn.Module):
+    def __init__(self, patch_size: int, stride: int, in_chans: int, embed_dim: int) -> None:
+        super().__init__()
+        ph, pw = _to_2tuple(patch_size)
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=(ph, pw), stride=stride,
+                              padding=(ph // 2, pw // 2))
+        self.norm = nn.BatchNorm2d(embed_dim)
+
+    def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, int, int]:
+        x = self.proj(x)
+        _, _, H, W = x.shape
+        x = self.norm(x)
+        return x, H, W
+
+
+class StripNet(nn.Module):
+    """Strip-R-CNN backbone: 4-stage hierarchical CNN with strip-shaped DWConv attention.
+
+    Output: list of [B, C_i, H/s_i, W/s_i] per stage. Use `forward_last_features` for
+    the deepest stage only.
+    """
+
+    def __init__(
+        self,
+        embed_dims: List[int] = [64, 128, 320, 512],
+        mlp_ratios: List[int] = [8, 8, 4, 4],
+        k1s: List[int] = [1, 1, 1, 1],
+        k2s: List[int] = [19, 19, 19, 19],
+        depths: List[int] = [2, 2, 4, 2],
+        drop_rate: float = 0.1,
+        drop_path_rate: float = 0.15,
+        in_chans: int = 3,
+    ) -> None:
+        super().__init__()
+        self.depths = depths
+        self.num_stages = len(depths)
+        self.embed_dims = embed_dims
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
+        cur = 0
+        for i in range(self.num_stages):
+            patch_embed = OverlapPatchEmbed(
+                patch_size=7 if i == 0 else 3,
+                stride=4 if i == 0 else 2,
+                in_chans=in_chans if i == 0 else embed_dims[i - 1],
+                embed_dim=embed_dims[i],
+            )
+            block = nn.ModuleList([
+                Block(dim=embed_dims[i], mlp_ratio=mlp_ratios[i], k1=k1s[i], k2=k2s[i],
+                      drop=drop_rate, drop_path=dpr[cur + j])
+                for j in range(depths[i])
+            ])
+            norm = nn.LayerNorm(embed_dims[i], eps=1e-6)
+            cur += depths[i]
+            setattr(self, f"patch_embed{i + 1}", patch_embed)
+            setattr(self, f"block{i + 1}", block)
+            setattr(self, f"norm{i + 1}", norm)
+
+    def forward_features(self, x: torch.Tensor) -> List[torch.Tensor]:
+        B = x.shape[0]
+        outs: List[torch.Tensor] = []
+        for i in range(self.num_stages):
+            patch_embed = getattr(self, f"patch_embed{i + 1}")
+            block = getattr(self, f"block{i + 1}")
+            norm = getattr(self, f"norm{i + 1}")
+            x, H, W = patch_embed(x)
+            for blk in block:
+                x = blk(x)
+            x = x.flatten(2).transpose(1, 2)
+            x = norm(x)
+            x = x.reshape(B, H, W, -1).permute(0, 3, 1, 2).contiguous()
+            outs.append(x)
+        return outs
+
+    def forward_last_features(self, x: torch.Tensor) -> torch.Tensor:
+        return self.forward_features(x)[-1]
+
+
+def get_stripnet_small() -> StripNet:
+    return StripNet(
+        embed_dims=[64, 128, 320, 512],
+        mlp_ratios=[8, 8, 4, 4],
+        k1s=[1, 1, 1, 1],
+        k2s=[19, 19, 19, 19],
+        depths=[2, 2, 4, 2],
+        drop_rate=0.1,
+        drop_path_rate=0.15,
+    )
+
+
+def load_stripnet_small_pretrained(checkpoint_path: str | Path) -> StripNet:
+    """Build StripNet-small and load ImageNet-pretrained weights.
+
+    Strips the classification `head.*` keys. Tolerates missing/extra keys
+    (norm{N}.* are LayerNorm here vs BatchNorm in some forks — we keep LN).
+    """
+    LOGGER.info("📐 Loading StripNet-small from %s", checkpoint_path)
+    model = get_stripnet_small()
+    raw = torch.load(checkpoint_path, map_location="cpu", weights_only=False)
+    state = raw.get("state_dict", raw) if isinstance(raw, dict) else raw
+
+    # Drop classification head + the BN-form norm{N} keys if present (we use LN here).
+    drop_prefixes = ("head.",)
+    cleaned = {k: v for k, v in state.items() if not any(k.startswith(p) for p in drop_prefixes)}
+
+    # The pretrained checkpoint stores norm{N} as BatchNorm2d (running_mean/var/num_batches_tracked).
+    # Our code uses LayerNorm at this position. Strip BN running stats if found; copy weight/bias.
+    for n in (1, 2, 3, 4):
+        for suffix in ("running_mean", "running_var", "num_batches_tracked"):
+            cleaned.pop(f"norm{n}.{suffix}", None)
+
+    missing, unexpected = model.load_state_dict(cleaned, strict=False)
+    if missing:
+        LOGGER.info("StripNet missing keys (expected for newly-init layers): %d", len(missing))
+    if unexpected:
+        LOGGER.info("StripNet unexpected keys (ignored): %d", len(unexpected))
+    LOGGER.info("📐 StripNet-small loaded: %d params", sum(p.numel() for p in model.parameters()))
+    return model
--- a/vendor_reference/caption_test/src/models/stripnet_encoder.py
+++ b/vendor_reference/caption_test/src/models/stripnet_encoder.py
@@ -0,0 +1,48 @@
+"""StripNet image encoder wrapper for the caption-test pipeline.
+
+Exposes the same interface as DINOv3ViT: `forward(images) -> [B, embed_dim]`.
+StripNet's deepest stage produces [B, 512, H/32, W/32]; we apply global average
+pooling (GAP) and project to the target retrieval dimension via Linear(512→1024)
+to match DINOv3 native dim and keep TextFusionMLP unchanged.
+"""
+from __future__ import annotations
+
+import logging
+import torch
+import torch.nn as nn
+
+from src.models.stripnet import StripNet, load_stripnet_small_pretrained
+
+LOGGER = logging.getLogger("caption_test.stripnet_encoder")
+
+
+class StripNetEncoder(nn.Module):
+    """StripNet-small + GAP + projection to `out_dim`.
+
+    Frozen backbone (BatchNorm in eval mode); only the projection head and
+    any injected Conv-MONA adapters are trainable.
+    """
+
+    LAST_STAGE_DIM = 512  # StripNet-small last stage embed dim
+
+    def __init__(self, checkpoint_path: str, out_dim: int = 1024) -> None:
+        super().__init__()
+        self.out_dim = out_dim
+        self.backbone: StripNet = load_stripnet_small_pretrained(checkpoint_path)
+        self.pool = nn.AdaptiveAvgPool2d(1)
+        self.projection = nn.Linear(self.LAST_STAGE_DIM, out_dim)
+        nn.init.trunc_normal_(self.projection.weight, std=0.02)
+        nn.init.zeros_(self.projection.bias)
+
+    def train(self, mode: bool = True):
+        """Override: keep frozen backbone in eval mode (BN running stats stable)."""
+        super().train(mode)
+        # Frozen backbone always in eval; trainable adapters/projection follow `mode`.
+        if not any(p.requires_grad for p in self.backbone.parameters()):
+            self.backbone.eval()
+        return self
+
+    def forward(self, images: torch.Tensor) -> torch.Tensor:
+        feat = self.backbone.forward_last_features(images)  # [B, 512, H/32, W/32]
+        pooled = self.pool(feat).flatten(1)                  # [B, 512]
+        return self.projection(pooled)                        # [B, out_dim]
--- a/vendor_reference/caption_test/src/training/profiling.py
+++ b/vendor_reference/caption_test/src/training/profiling.py
@@ -0,0 +1,166 @@
+from __future__ import annotations
+
+"""PyTorch Profiler wrapper for training performance analysis.
+
+Profiles the first N batches of training to identify bottlenecks
+in CUDA/CPU execution, memory allocation, and data loading.
+
+Exports:
+  - Chrome trace (viewable in chrome://tracing)
+  - TensorBoard plugin data (if TB available)
+  - Summary table to console
+
+Usage:
+    profiler = TrainingProfiler(output_dir, n_warmup=3, n_active=5)
+    for batch_idx, batch in enumerate(loader):
+        with profiler.step_context(batch_idx):
+            # ... training step ...
+        if profiler.is_done(batch_idx):
+            break
+    profiler.export()
+"""
+
+import logging
+from pathlib import Path
+
+import torch
+from torch.profiler import ProfilerActivity, profile, schedule, tensorboard_trace_handler
+
+LOGGER = logging.getLogger("caption_test.profiler")
+
+
+class TrainingProfiler:
+    """PyTorch profiler for first N training batches.
+
+    Args:
+        output_dir: Directory for profiler output.
+        n_warmup: Number of warmup steps (not profiled).
+        n_active: Number of steps to actively profile.
+        n_repeat: Number of profiling cycles.
+        record_shapes: Record tensor shapes.
+        profile_memory: Track memory allocation.
+        with_stack: Record Python call stacks.
+    """
+
+    def __init__(
+        self,
+        output_dir: str | Path,
+        n_warmup: int = 3,
+        n_active: int = 5,
+        n_repeat: int = 1,
+        record_shapes: bool = True,
+        profile_memory: bool = True,
+        with_stack: bool = False,
+    ) -> None:
+        self.output_dir = Path(output_dir) / "profiler"
+        self.output_dir.mkdir(parents=True, exist_ok=True)
+        self.n_warmup = n_warmup
+        self.n_active = n_active
+        self.n_repeat = n_repeat
+        self.total_steps = (n_warmup + n_active) * n_repeat
+
+        self._profiler = profile(
+            activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
+            schedule=schedule(
+                wait=0,
+                warmup=n_warmup,
+                active=n_active,
+                repeat=n_repeat,
+            ),
+            on_trace_ready=tensorboard_trace_handler(str(self.output_dir)),
+            record_shapes=record_shapes,
+            profile_memory=profile_memory,
+            with_stack=with_stack,
+        )
+        self._started = False
+
+    def start(self) -> None:
+        """Start the profiler."""
+        self._profiler.__enter__()
+        self._started = True
+        LOGGER.info(
+            "Profiler started: %d warmup + %d active steps, output: %s",
+            self.n_warmup, self.n_active, self.output_dir,
+        )
+
+    def step(self) -> None:
+        """Signal end of a profiling step."""
+        if self._started:
+            self._profiler.step()
+
+    def is_done(self, batch_idx: int) -> bool:
+        """Check if profiling is complete."""
+        return batch_idx >= self.total_steps
+
+    def export(self) -> None:
+        """Export profiling results and print summary."""
+        if not self._started:
+            return
+
+        self._profiler.__exit__(None, None, None)
+        self._started = False
+
+        # Print key averages summary.
+        summary = self._profiler.key_averages().table(
+            sort_by="cuda_time_total", row_limit=20,
+        )
+        LOGGER.info("Profiler summary (top 20 by CUDA time):\n%s", summary)
+
+        # Export Chrome trace.
+        trace_path = self.output_dir / "chrome_trace.json"
+        self._profiler.export_chrome_trace(str(trace_path))
+        LOGGER.info("Chrome trace exported: %s", trace_path)
+
+        # Memory summary if available.
+        if torch.cuda.is_available():
+            mem_summary = torch.cuda.memory_summary(abbreviated=True)
+            summary_path = self.output_dir / "memory_summary.txt"
+            summary_path.write_text(mem_summary)
+            LOGGER.info("CUDA memory summary: %s", summary_path)
+
+
+def print_model_summary(model: torch.nn.Module, device: str = "cuda") -> str:
+    """Print model summary using torchinfo (if available).
+
+    Falls back to a simple parameter count if torchinfo is not installed.
+
+    Returns:
+        Summary string.
+    """
+    try:
+        from torchinfo import summary as torchinfo_summary
+        info = torchinfo_summary(
+            model,
+            input_data={
+                "drone_img": torch.randn(1, 3, 256, 256, device=device),
+                "sat_img": torch.randn(1, 3, 256, 256, device=device),
+            },
+            col_names=["input_size", "output_size", "num_params", "trainable"],
+            verbose=0,
+            depth=3,
+        )
+        summary_str = str(info)
+        LOGGER.info("Model summary (torchinfo):\n%s", summary_str)
+        return summary_str
+    except ImportError:
+        LOGGER.info("torchinfo not installed, using basic parameter count")
+    except Exception as e:
+        LOGGER.warning("torchinfo failed (%s), using basic parameter count", e)
+
+    # Fallback: simple param count.
+    lines = []
+    total = 0
+    trainable = 0
+    for name, param in model.named_parameters():
+        total += param.numel()
+        if param.requires_grad:
+            trainable += param.numel()
+            lines.append(f"  [trainable] {name}: {list(param.shape)} ({param.numel():,})")
+
+    summary_str = (
+        f"Total parameters: {total:,}\n"
+        f"Trainable parameters: {trainable:,} ({100*trainable/max(total,1):.2f}%)\n"
+        + "\n".join(lines[:30])
+    )
+    LOGGER.info("Model summary:\n%s", summary_str)
+    return summary_str
--- a/vendor_reference/caption_test/src/training/trackers.py
+++ b/vendor_reference/caption_test/src/training/trackers.py
@@ -0,0 +1,206 @@
+from __future__ import annotations
+
+"""Unified experiment tracking: W&B + TensorBoard + CSV.
+
+Auto-detects available backends. Falls back gracefully if wandb/tensorboard
+are not installed.
+
+Usage:
+    tracker = ExperimentTracker(output_dir, config_dict, use_wandb=True, use_tb=True)
+    tracker.log_train(epoch, {"loss": 0.5, "lr": 1e-4})
+    tracker.log_val(epoch, {"r@1_q2g": 0.3})
+    tracker.log_gradients(epoch, grad_norms_dict)
+    tracker.log_image(epoch, "gradcam/drone", image_tensor)
+    tracker.close()
+"""
+
+import logging
+from pathlib import Path
+from typing import Any
+
+import torch
+
+LOGGER = logging.getLogger("caption_test.trackers")
+
+
+def _try_import_wandb():
+    try:
+        import wandb
+        return wandb
+    except ImportError:
+        return None
+
+
+def _try_import_tb():
+    try:
+        from torch.utils.tensorboard import SummaryWriter
+        return SummaryWriter
+    except ImportError:
+        return None
+
+
+class ExperimentTracker:
+    """Unified tracker dispatching to W&B, TensorBoard, and CSV.
+
+    Args:
+        output_dir: Base output directory.
+        config: Dict of hyperparameters to log.
+        use_wandb: Enable Weights & Biases tracking.
+        use_tb: Enable TensorBoard tracking.
+        wandb_project: W&B project name.
+        wandb_run_name: W&B run name (auto-generated if None).
+        wandb_entity: W&B entity (team/user).
+    """
+
+    def __init__(
+        self,
+        output_dir: str | Path,
+        config: dict[str, Any] | None = None,
+        use_wandb: bool = False,
+        use_tb: bool = True,
+        wandb_project: str = "caption-test-gtauav",
+        wandb_run_name: str | None = None,
+        wandb_entity: str | None = None,
+    ) -> None:
+        self.output_dir = Path(output_dir)
+        self._wandb_run = None
+        self._tb_writer = None
+
+        # W&B init.
+        if use_wandb:
+            wandb = _try_import_wandb()
+            if wandb is not None:
+                self._wandb_run = wandb.init(
+                    project=wandb_project,
+                    name=wandb_run_name,
+                    entity=wandb_entity,
+                    config=config or {},
+                    dir=str(self.output_dir),
+                    reinit=True,
+                )
+                LOGGER.info("W&B initialized: %s", self._wandb_run.url)
+            else:
+                LOGGER.warning("wandb not installed, skipping W&B tracking")
+
+        # TensorBoard init.
+        if use_tb:
+            SummaryWriter = _try_import_tb()
+            if SummaryWriter is not None:
+                tb_dir = self.output_dir / "tb_logs"
+                tb_dir.mkdir(parents=True, exist_ok=True)
+                self._tb_writer = SummaryWriter(log_dir=str(tb_dir))
+                LOGGER.info("TensorBoard initialized: %s", tb_dir)
+            else:
+                LOGGER.warning("tensorboard not installed, skipping TB tracking")
+
+    @property
+    def has_wandb(self) -> bool:
+        return self._wandb_run is not None
+
+    @property
+    def has_tb(self) -> bool:
+        return self._tb_writer is not None
+
+    def log_train(self, epoch: int, metrics: dict[str, float], step: int | None = None) -> None:
+        """Log training metrics for an epoch."""
+        if self._wandb_run is not None:
+            self._wandb_run.log(
+                {f"train/{k}": v for k, v in metrics.items()},
+                step=step or epoch,
+            )
+        if self._tb_writer is not None:
+            for k, v in metrics.items():
+                self._tb_writer.add_scalar(f"train/{k}", v, global_step=step or epoch)
+
+    def log_val(self, epoch: int, metrics: dict[str, float], step: int | None = None) -> None:
+        """Log validation metrics."""
+        if self._wandb_run is not None:
+            self._wandb_run.log(
+                {f"val/{k}": v for k, v in metrics.items()},
+                step=step or epoch,
+            )
+        if self._tb_writer is not None:
+            for k, v in metrics.items():
+                self._tb_writer.add_scalar(f"val/{k}", v, global_step=step or epoch)
+
+    def log_gradients(self, epoch: int, grad_norms: dict[str, float], step: int | None = None) -> None:
+        """Log gradient norms per parameter group."""
+        if self._wandb_run is not None:
+            self._wandb_run.log(
+                {f"gradients/{k}": v for k, v in grad_norms.items()},
+                step=step or epoch,
+            )
+        if self._tb_writer is not None:
+            for k, v in grad_norms.items():
+                self._tb_writer.add_scalar(f"gradients/{k}", v, global_step=step or epoch)
+
+    def log_scalar(self, tag: str, value: float, step: int) -> None:
+        """Log a single scalar."""
+        if self._wandb_run is not None:
+            self._wandb_run.log({tag: value}, step=step)
+        if self._tb_writer is not None:
+            self._tb_writer.add_scalar(tag, value, global_step=step)
+
+    def log_image(self, tag: str, image: Any, step: int, caption: str | None = None) -> None:
+        """Log an image (numpy HWC or torch CHW).
+
+        Args:
+            tag: Image tag/name.
+            image: numpy array [H,W,C] or torch tensor [C,H,W].
+            step: Global step.
+            caption: Optional caption for W&B.
+        """
+        if self._wandb_run is not None:
+            wandb = _try_import_wandb()
+            if isinstance(image, torch.Tensor):
+                image_np = image.detach().cpu().permute(1, 2, 0).numpy()
+            else:
+                image_np = image
+            self._wandb_run.log(
+                {tag: wandb.Image(image_np, caption=caption)},
+                step=step,
+            )
+        if self._tb_writer is not None:
+            if isinstance(image, torch.Tensor):
+                self._tb_writer.add_image(tag, image.detach().cpu(), global_step=step)
+            else:
+                self._tb_writer.add_image(tag, image, global_step=step, dataformats="HWC")
+
+    def log_histogram(self, tag: str, values: torch.Tensor, step: int) -> None:
+        """Log a histogram of values (weights, activations, etc.)."""
+        if self._wandb_run is not None:
+            wandb = _try_import_wandb()
+            self._wandb_run.log(
+                {tag: wandb.Histogram(values.detach().cpu().numpy())},
+                step=step,
+            )
+        if self._tb_writer is not None:
+            self._tb_writer.add_histogram(tag, values.detach().cpu(), global_step=step)
+
+    def log_model_graph(self, model: torch.nn.Module, input_example: Any = None) -> None:
+        """Log model graph to TensorBoard (if available)."""
+        if self._tb_writer is not None and input_example is not None:
+            try:
+                self._tb_writer.add_graph(model, input_example)
+            except Exception as e:
+                LOGGER.warning("Failed to log model graph: %s", e)
+
+    def watch_model(self, model: torch.nn.Module, log_freq: int = 100) -> None:
+        """Enable W&B gradient/weight watching."""
+        if self._wandb_run is not None:
+            wandb = _try_import_wandb()
+            wandb.watch(model, log="all", log_freq=log_freq)
+
+    def log_summary(self, summary: dict[str, Any]) -> None:
+        """Log final summary metrics (best R@1, etc.)."""
+        if self._wandb_run is not None:
+            for k, v in summary.items():
+                self._wandb_run.summary[k] = v
+
+    def close(self) -> None:
+        """Flush and close all backends."""
+        if self._tb_writer is not None:
+            self._tb_writer.flush()
+            self._tb_writer.close()
+        if self._wandb_run is not None:
+            self._wandb_run.finish()
--- a/vendor_reference/caption_test/src/training/train_gtauav.py
+++ b/vendor_reference/caption_test/src/training/train_gtauav.py