diff --git a/conf/gtauav_balanced_sofia.gin b/conf/gtauav_balanced_sofia.gin new file mode 100644 index 0000000..6971495 --- /dev/null +++ b/conf/gtauav_balanced_sofia.gin @@ -0,0 +1,39 @@ +# GTA-UAV Balanced (SOFIA-Tiny backbone): SOFIA v7.1 student trained from scratch +# с двухуровневой text fusion: +# 1. Mid-level: Text-FiLM в SAT и UAV heads (модулирует feature map перед GGeM/CHP). +# 2. Late-level: GatedFusion на дескрипторах (как в DINOv3/StripNet вариантах). +# +# Trainable (~5-7M): +# - SOFIA backbone (Tiny, ~5M, from scratch — нет pretrained) +# - SOFIA heads (SatHead GGeM+BN+Linear, UAVHead AltitudeFiLM+CHP+BN+Linear, +Text-FiLM) +# - DGTRS-CLIP LoRA (rank=4, ~147K) +# - TextFusionMLP (3*768 -> 1024 -> 1024, ~3.4M, shared) +# - Gates α_q, α_g + learnable τ +# +# Altitude (drone_height метры) подаётся в UAVHead.AltitudeFiLM из dataset meta CSV. +# Для sat — altitude=None → FiLM passthrough (γ=1, β=0). +# +# Note: SOFIA from scratch — нужно больше эпох или warmup, чем frozen DINOv3/StripNet. +# Mamba-2 backend (mamba_ssm) даёт 2-8x speedup vs torch fallback. + +include 'conf/gtauav_balanced.gin' + +# ---- Backbone ---- +TrainConfigGTAUAV.backbone = "sofia" +TrainConfigGTAUAV.sofia_preset = "Tiny" +TrainConfigGTAUAV.sofia_d_descriptor = 1024 +TrainConfigGTAUAV.sofia_use_text_film_uav = True +TrainConfigGTAUAV.sofia_use_text_film_sat = True +TrainConfigGTAUAV.sofia_lora_rank = 4 +# Mamba-1 used for Tiny (Mamba-2 torch fallback has a pre-existing reshape bug +# with channels not divisible by default headdim; switch to "mamba2" for M/L +# presets where channels % 64 == 0 OR install mamba_ssm CUDA kernels). +TrainConfigGTAUAV.sofia_mamba_variant = "mamba1" +TrainConfigGTAUAV.sofia_mamba_backend = "auto" # mamba_ssm if installed else torch fallback + +# ---- Training overrides ---- +TrainConfigGTAUAV.gradient_checkpointing = False # SOFIA from-scratch — keep activations live +TrainConfigGTAUAV.shared_encoder = True # ignored by SOFIA but kept for logging compat + +# ---- Output ---- +TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_sofia" diff --git a/conf/gtauav_balanced_sofia_v1.gin b/conf/gtauav_balanced_sofia_v1.gin new file mode 100644 index 0000000..f35abd5 --- /dev/null +++ b/conf/gtauav_balanced_sofia_v1.gin @@ -0,0 +1,42 @@ +# GTA-UAV Balanced (SOFIA v1 backbone): StripNet+DCNv4 hierarchical CNN +# (~3-30M params depending on variant) trained from scratch с двухуровневой +# text fusion: +# 1. Mid-level: Text-FiLM в SAT и UAV heads (модулирует [B,C,8,8] перед GGeM). +# 2. Late-level: GatedFusion на дескрипторах. +# +# UAV head: AltitudeFiLM(drone_height) + [TextFiLM] + GGeM + BN + Linear. +# SAT head: [TextFiLM] + GGeM + BN + Linear. +# Один backbone shared между sat/uav. +# +# Variant -> размер модели: +# tiny_tiny: dims [16, 32, 80, 128] (~0.4M) +# tiny : dims [32, 64, 128, 256] (~1M) +# small : dims [64, 128, 320, 512] (~3M, default) +# small_v2 : dims [64, 128, 256, 384] (~2M) +# +# Trainable (с small variant + text fusion): +# - SOFIA v1 backbone (~3M) + heads (~0.6M) +# - DGTRS-CLIP LoRA (rank=4, ~147K) +# - TextFusionMLP (3*768 -> 1024 -> 1024, ~3.4M, shared) +# - Gates α_q, α_g + learnable τ +# Total trainable ~7M. +# +# Note: DCNv4 требует CUDA — обучение только на GPU. Не работает на CPU. + +include 'conf/gtauav_balanced.gin' + +# ---- Backbone ---- +TrainConfigGTAUAV.backbone = "sofia_v1" +TrainConfigGTAUAV.sofia_v1_variant = "tiny" +TrainConfigGTAUAV.sofia_v1_d_descriptor = 1024 +TrainConfigGTAUAV.sofia_v1_use_text_film_uav = True +TrainConfigGTAUAV.sofia_v1_use_text_film_sat = True +TrainConfigGTAUAV.sofia_v1_use_film_altitude = True +TrainConfigGTAUAV.sofia_v1_lora_rank = 4 + +# ---- Training overrides ---- +TrainConfigGTAUAV.gradient_checkpointing = False # SOFIA v1 from-scratch +TrainConfigGTAUAV.shared_encoder = True + +# ---- Output ---- +TrainConfigGTAUAV.output_dir = "out/gtauav/with_text_sofia_v1" diff --git a/conf/gtauav_baseline_sofia.gin b/conf/gtauav_baseline_sofia.gin new file mode 100644 index 0000000..492a092 --- /dev/null +++ b/conf/gtauav_baseline_sofia.gin @@ -0,0 +1,13 @@ +# GTA-UAV Baseline (SOFIA-Tiny backbone): no text fusion. Reference R@1 для +# computing Δ R@1 vs gtauav_balanced_sofia.gin. +# +# В baseline_mode=True: +# - Text-FiLM отключается (SOFIA heads работают только с altitude). +# - DGTRS-CLIP не загружается, TextFusionMLP не строится. +# - GatedFusion gates = 1.0 (text игнорируется). + +include 'conf/gtauav_balanced_sofia.gin' + +TrainConfigGTAUAV.baseline_mode = True +TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_sofia" +TrainConfigGTAUAV.use_gradcam = False diff --git a/conf/gtauav_baseline_sofia_v1.gin b/conf/gtauav_baseline_sofia_v1.gin new file mode 100644 index 0000000..26478df --- /dev/null +++ b/conf/gtauav_baseline_sofia_v1.gin @@ -0,0 +1,8 @@ +# GTA-UAV Baseline (SOFIA v1 backbone): no text fusion. Reference R@1 для +# computing Δ R@1 vs gtauav_balanced_sofia_v1.gin. + +include 'conf/gtauav_balanced_sofia_v1.gin' + +TrainConfigGTAUAV.baseline_mode = True +TrainConfigGTAUAV.output_dir = "out/gtauav/baseline_sofia_v1" +TrainConfigGTAUAV.use_gradcam = False diff --git a/scripts/test_dcn.py b/scripts/test_dcn.py new file mode 100644 index 0000000..214c173 --- /dev/null +++ b/scripts/test_dcn.py @@ -0,0 +1,11 @@ +import torch +from DCNv4 import DCNv4 + +dcn = DCNv4(channels=64, group=4).cuda().eval() +print('--- 50 forward calls, no_grad ---') +with torch.no_grad(): + for s in range(50): + x = torch.randn(8, 4096, 64, device='cuda') + _ = dcn(x) + if s % 10 == 0: + print(f'{s}: {torch.cuda.memory_allocated() / 1e6:.1f} MB') diff --git a/src/datasets/gtauav_dataset.py b/src/datasets/gtauav_dataset.py index 10b0c75..7d49517 100644 --- a/src/datasets/gtauav_dataset.py +++ b/src/datasets/gtauav_dataset.py @@ -104,6 +104,10 @@ class GTAUAVDataset(Dataset): image_transform: Fallback single transform for both (used if drone/sat not set). filter_meta: Path to seg_filter.json (exclude 90%+ bg/water). drop_caption_prob: Probability of dropping captions (ablation). + meta_csvs: Optional list of CSV paths with `img_name,drone_height,...` + columns (e.g. `cross-area-drone2sate-{train,test}_drone_meta.csv`). + When provided, an `altitude` (meters) field is attached per drone + sample. Defaults to scanning the rgb_root for *_drone_meta.csv. seed: Random seed. """ @@ -117,6 +121,7 @@ class GTAUAVDataset(Dataset): image_transform: Callable[[Image.Image], torch.Tensor] | None = None, filter_meta: str | None = None, drop_caption_prob: float = 0.0, + meta_csvs: list[str] | None = None, seed: int = 0, ) -> None: self.rgb_root = Path(rgb_root) @@ -131,6 +136,9 @@ class GTAUAVDataset(Dataset): if filter_meta is not None: self._load_filter(Path(filter_meta)) + # Load drone altitude index (img_name -> meters). Empty dict if no CSVs. + self.altitude_index: dict[str, float] = self._load_altitude_index(meta_csvs) + # Load caption index. LOGGER.info("📚 Loading caption index from %s", caption_root) self.caption_index = _load_caption_index(self.caption_root) @@ -141,6 +149,43 @@ class GTAUAVDataset(Dataset): self._load_pairs(Path(pair_json)) LOGGER.info("✅ Loaded %d pairs from %s", len(self.entries), pair_json) + def _load_altitude_index(self, meta_csvs: list[str] | None) -> dict[str, float]: + """Build {drone_img_name: altitude_meters} from drone_meta CSVs. + + If `meta_csvs` is None, auto-discovers `*_drone_meta.csv` (TSV format, + columns img_name/drone_height/...) under `self.rgb_root`. Missing or + unreadable files are skipped silently — altitude defaults to 0.0 + downstream when an entry is missing. + """ + if meta_csvs is None: + meta_csvs = [str(p) for p in sorted(self.rgb_root.glob("*_drone_meta.csv"))] + # Prefer `*_drone_meta_new.csv` if present (overrides original). + new_csvs = [str(p) for p in sorted(self.rgb_root.glob("*_drone_meta_new.csv"))] + meta_csvs = new_csvs or meta_csvs + index: dict[str, float] = {} + for csv_path in meta_csvs: + path = Path(csv_path) + if not path.exists(): + continue + try: + with path.open() as f: + header = f.readline().rstrip("\n").split("\t") + name_idx = header.index("img_name") + height_idx = header.index("drone_height") + for line in f: + parts = line.rstrip("\n").split("\t") + if len(parts) <= max(name_idx, height_idx): + continue + try: + index[parts[name_idx]] = float(parts[height_idx]) + except ValueError: + continue + except (OSError, ValueError) as exc: + LOGGER.warning("Failed to parse drone meta CSV %s: %s", path, exc) + if index: + LOGGER.info("📐 Altitude index: %d drones (from %d CSV)", len(index), len(meta_csvs)) + return index + def _load_filter(self, path: Path) -> None: with open(path) as f: meta = json.load(f) @@ -200,6 +245,7 @@ class GTAUAVDataset(Dataset): "caption_l2": l2, "caption_l3": l3, "sat_captions": sat_captions, + "altitude": self.altitude_index.get(drone_name, 0.0), }) def _load_image(self, directory: str, filename: str, transform: Callable | None = None) -> torch.Tensor: @@ -267,6 +313,7 @@ class GTAUAVDataset(Dataset): "pair_id": entry["drone_name"], "sat_name": sat_name, "positive_weight": pos_weight, + "altitude": float(entry["altitude"]), } @@ -286,6 +333,7 @@ def collate_gtauav_batch( "pair_ids": [b["pair_id"] for b in batch], "sat_names": [b["sat_name"] for b in batch], "positive_weights": torch.tensor([b["positive_weight"] for b in batch], dtype=torch.float32), + "altitude": torch.tensor([b["altitude"] for b in batch], dtype=torch.float32), } @@ -371,6 +419,7 @@ class GTAUAVDroneQuery(Dataset): "caption_l2": entry["caption_l2"], "caption_l3": entry["caption_l3"], "valid_sat_names": list(entry["sat_candidates"]), + "altitude": float(entry.get("altitude", 0.0)), } @@ -392,4 +441,5 @@ def collate_drone_query(batch: list[dict[str, Any]]) -> dict[str, Any]: "caption_l2": [b["caption_l2"] for b in batch], "caption_l3": [b["caption_l3"] for b in batch], "valid_sat_names": [b["valid_sat_names"] for b in batch], + "altitude": torch.tensor([b["altitude"] for b in batch], dtype=torch.float32), } diff --git a/src/models/asymmetric_encoder.py b/src/models/asymmetric_encoder.py index 123f2d9..139c542 100644 --- a/src/models/asymmetric_encoder.py +++ b/src/models/asymmetric_encoder.py @@ -462,6 +462,7 @@ class AsymmetricEncoder(nn.Module): caption_l1: list[str] | None = None, caption_l2: list[str] | None = None, caption_l3: list[str] | None = None, + altitude: torch.Tensor | None = None, # noqa: ARG002 — accepted for API parity with SOFIAFusionEncoder ) -> torch.Tensor: """Encode drone → normalized query embedding with per-sample text mask.""" drone_feat = self.encode_drone(drone_img) @@ -492,6 +493,7 @@ class AsymmetricEncoder(nn.Module): sat_caption_l1: list[str] | None = None, sat_caption_l2: list[str] | None = None, sat_caption_l3: list[str] | None = None, + altitude: torch.Tensor | None = None, # noqa: ARG002 — accepted for API parity with SOFIAFusionEncoder ) -> dict[str, torch.Tensor]: """Forward pass. diff --git a/src/models/sofia_v1/__init__.py b/src/models/sofia_v1/__init__.py new file mode 100644 index 0000000..789563e --- /dev/null +++ b/src/models/sofia_v1/__init__.py @@ -0,0 +1,41 @@ +"""SOFIA v1 — StripNet + DCNv4 hierarchical CNN backbone for CVGL. + +Lightweight 4-stage backbone (~5–30M params depending on variant). Outputs +features at 4 scales (last stage is 8x8 for 256x256 input). + +Variants (in `stripnet_model_dcn.VARIANT_MAP`): +- `tiny_tiny`: dims [16, 32, 80, 128] +- `tiny` : dims [32, 64, 128, 256] +- `small` : dims [64, 128, 320, 512] (default) +- `small_v2` : dims [64, 128, 256, 384] + +Use `SOFIAv1FusionEncoder` from `src.models.sofia_v1_fusion_encoder` for +end-to-end CVGL training with DGTRS-CLIP captions and altitude. +""" + +from .config import SOFIAv1Config +from .heads import SatHeadV1, UAVHeadV1 +from .model import SOFIAv1 +from .stripnet_model_dcn import ( + StripNetDCN, + VARIANT_MAP, + build_stripnet_dcn, + get_stripnet_dcn_small, + get_stripnet_dcn_small_v2, + get_stripnet_dcn_tiny, + get_stripnet_dcn_tiny_tiny, +) + +__all__ = [ + "SOFIAv1", + "SOFIAv1Config", + "SatHeadV1", + "UAVHeadV1", + "StripNetDCN", + "build_stripnet_dcn", + "get_stripnet_dcn_small", + "get_stripnet_dcn_small_v2", + "get_stripnet_dcn_tiny", + "get_stripnet_dcn_tiny_tiny", + "VARIANT_MAP", +] diff --git a/src/models/sofia_v1/config.py b/src/models/sofia_v1/config.py new file mode 100644 index 0000000..e69fdef --- /dev/null +++ b/src/models/sofia_v1/config.py @@ -0,0 +1,44 @@ +"""SOFIA v1 configuration. + +Lightweight 4-stage StripNet+DCNv4 backbone variants (`tiny_tiny`/`tiny`/`small` +/`small_v2`) plus simple GGeM-based heads with optional altitude-FiLM and +text-FiLM modulation. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Literal + + +@dataclass +class SOFIAv1Config: + # -------- Backbone -------- + variant: Literal["tiny_tiny", "tiny", "small", "small_v2"] = "small" + in_channels: int = 3 + input_size: int = 256 + # DCN op variant. "v2" (default) uses torchvision DeformConv2d — stable. + # "v4" uses OpenGVLab DCNv4 — faster but has known C++-extension memory + # leak (~9 MB per forward) that OOMs in long training runs. + dcn_variant: Literal["v2", "v4"] = "v2" + + # -------- Heads -------- + d_descriptor: int = 1024 + return_normalized: bool = False # False → wrapper handles L2 after gated fusion + + # Altitude-FiLM (UAV head only). + use_film_altitude: bool = True + altitude_norm: float = 500.0 + + # Text-FiLM (mid-level fusion). Both can be toggled independently. + use_text_film_uav: bool = True + use_text_film_sat: bool = True + text_film_dim: int = 1024 + text_film_hidden: int = 256 + + def summary(self) -> str: + return ( + f"SOFIAv1Config(variant={self.variant}, d={self.d_descriptor}, " + f"film_alt={self.use_film_altitude}, " + f"text_film(sat={self.use_text_film_sat},uav={self.use_text_film_uav}))" + ) diff --git a/src/models/sofia_v1/heads.py b/src/models/sofia_v1/heads.py new file mode 100644 index 0000000..880e8da --- /dev/null +++ b/src/models/sofia_v1/heads.py @@ -0,0 +1,99 @@ +"""Heads for SOFIA v1 (StripNet+DCNv4 backbone). + +Designed parallel to SOFIA v7.1 heads but lighter — basic GGeM pooling +instead of CHP, optional altitude-FiLM (UAV) and text-FiLM (both). + +The heads produce un-normalized D-dim descriptors when `return_normalized=False`, +so that a fusion wrapper can blend with text via gated fusion before the final +L2 normalization. +""" + +from __future__ import annotations + +from typing import Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from src.models.sofia_v71.layers import AltitudeFiLM, GGeM, TextFiLM + + +class SatHeadV1(nn.Module): + """Satellite head: [TextFiLM] + GGeM + BN + Linear [+ L2].""" + + def __init__( + self, + in_channels: int, + d_descriptor: int, + return_normalized: bool = False, + use_text_film: bool = False, + text_film_dim: int = 1024, + text_film_hidden: int = 256, + ) -> None: + super().__init__() + self.return_normalized = return_normalized + self.use_text_film = use_text_film + if use_text_film: + self.text_film = TextFiLM(in_channels, text_dim=text_film_dim, hidden_dim=text_film_hidden) + self.ggem = GGeM(in_channels) + self.bn = nn.BatchNorm1d(in_channels, affine=False) + self.proj = nn.Linear(in_channels, d_descriptor) + + def forward( + self, + x: torch.Tensor, + text_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if self.use_text_film: + x = self.text_film(x, text_emb) + g = self.ggem(x) + g = self.bn(g) + g = self.proj(g) + if self.return_normalized: + g = F.normalize(g, p=2, dim=-1) + return g + + +class UAVHeadV1(nn.Module): + """UAV head: AltitudeFiLM [+ TextFiLM] + GGeM + BN + Linear [+ L2].""" + + def __init__( + self, + in_channels: int, + d_descriptor: int, + use_film: bool = True, + altitude_norm: float = 500.0, + return_normalized: bool = False, + use_text_film: bool = False, + text_film_dim: int = 1024, + text_film_hidden: int = 256, + ) -> None: + super().__init__() + self.return_normalized = return_normalized + self.use_film = use_film + self.use_text_film = use_text_film + if use_film: + self.film = AltitudeFiLM(in_channels, altitude_norm=altitude_norm) + if use_text_film: + self.text_film = TextFiLM(in_channels, text_dim=text_film_dim, hidden_dim=text_film_hidden) + self.ggem = GGeM(in_channels) + self.bn = nn.BatchNorm1d(in_channels, affine=False) + self.proj = nn.Linear(in_channels, d_descriptor) + + def forward( + self, + x: torch.Tensor, + altitude: Optional[torch.Tensor] = None, + text_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + if self.use_film: + x = self.film(x, altitude) + if self.use_text_film: + x = self.text_film(x, text_emb) + g = self.ggem(x) + g = self.bn(g) + g = self.proj(g) + if self.return_normalized: + g = F.normalize(g, p=2, dim=-1) + return g diff --git a/src/models/sofia_v1/model.py b/src/models/sofia_v1/model.py new file mode 100644 index 0000000..cdca014 --- /dev/null +++ b/src/models/sofia_v1/model.py @@ -0,0 +1,93 @@ +"""SOFIA v1 model: StripNet+DCNv4 backbone + asymmetric Sat/UAV heads. + +Architecture: + img [B,3,256,256] --> StripNetDCN (4 stages) --> [B, C_4, 8, 8] + | + ┌─────────────────────┴────────────────┐ + ▼ ▼ + SatHeadV1 UAVHeadV1 + [TextFiLM] AltitudeFiLM + GGeM [TextFiLM] + BN + Linear GGeM + [L2] BN + Linear + [L2] + +Heads return un-normalized D-dim descriptors (for downstream gated text fusion). +""" + +from __future__ import annotations + +from typing import Dict, List, Optional + +import torch +import torch.nn as nn + +from .config import SOFIAv1Config +from .heads import SatHeadV1, UAVHeadV1 +from .stripnet_model_dcn import VARIANT_MAP, StripNetDCN + + +class SOFIAv1(nn.Module): + """SOFIA v1: shared StripNet+DCNv4 backbone with asymmetric heads.""" + + def __init__(self, cfg: SOFIAv1Config) -> None: + super().__init__() + if cfg.variant not in VARIANT_MAP: + raise ValueError(f"Unknown variant {cfg.variant!r}") + self.cfg = cfg + + # Single shared backbone for sat + uav (saves params; v1 stays lightweight). + self.backbone: StripNetDCN = VARIANT_MAP[cfg.variant](dcn_variant=cfg.dcn_variant) + last_channels = self.backbone.embed_dims[-1] + self.feature_channels = last_channels + + self.sat_head = SatHeadV1( + in_channels=last_channels, + d_descriptor=cfg.d_descriptor, + return_normalized=cfg.return_normalized, + use_text_film=cfg.use_text_film_sat, + text_film_dim=cfg.text_film_dim, + text_film_hidden=cfg.text_film_hidden, + ) + self.uav_head = UAVHeadV1( + in_channels=last_channels, + d_descriptor=cfg.d_descriptor, + use_film=cfg.use_film_altitude, + altitude_norm=cfg.altitude_norm, + return_normalized=cfg.return_normalized, + use_text_film=cfg.use_text_film_uav, + text_film_dim=cfg.text_film_dim, + text_film_hidden=cfg.text_film_hidden, + ) + + def _extract_last(self, x: torch.Tensor) -> torch.Tensor: + """Run backbone, return only the deepest feature map [B, C, 8, 8].""" + feats: List[torch.Tensor] = self.backbone(x) + return feats[-1] + + def forward( + self, + sat: Optional[torch.Tensor] = None, + uav: Optional[torch.Tensor] = None, + altitude: Optional[torch.Tensor] = None, + text_emb_sat: Optional[torch.Tensor] = None, + text_emb_uav: Optional[torch.Tensor] = None, + return_features: bool = False, + ) -> Dict[str, torch.Tensor]: + result: Dict[str, torch.Tensor] = {} + + if sat is not None: + f_sat = self._extract_last(sat) + g_sat = self.sat_head(f_sat, text_emb=text_emb_sat) + result["g_sat"] = g_sat + if return_features: + result["features_sat"] = f_sat + + if uav is not None: + f_uav = self._extract_last(uav) + g_uav = self.uav_head(f_uav, altitude=altitude, text_emb=text_emb_uav) + result["g_uav"] = g_uav + if return_features: + result["features_uav"] = f_uav + + return result diff --git a/src/models/sofia_v1/stripnet_blocks_dcn.py b/src/models/sofia_v1/stripnet_blocks_dcn.py new file mode 100644 index 0000000..679423b --- /dev/null +++ b/src/models/sofia_v1/stripnet_blocks_dcn.py @@ -0,0 +1,57 @@ +import torch +import torch.nn as nn +from torchvision.ops import DeformConv2d + +class DCNBlock(nn.Module): + """ + StripNet-style block but uses deformable conv instead of rigid convs. + """ + def __init__(self, in_ch, out_ch, hidden_ratio=0.25, modulation=True): + super().__init__() + hidden_ch = max(1, int(out_ch * hidden_ratio)) + + # 1x1 reduce + self.reduce = nn.Sequential( + nn.Conv2d(in_ch, hidden_ch, kernel_size=1, bias=False), + nn.BatchNorm2d(hidden_ch), + nn.ReLU(inplace=True), + ) + + # Offset conv (predicts offsets and optional mask) + offset_channels = 2 * 3 * 3 if not modulation else 3 * 3 * 3 + self.offset_conv = nn.Conv2d(hidden_ch, offset_channels, + kernel_size=3, padding=1) + + # Deformable conv + self.deform = DeformConv2d(hidden_ch, hidden_ch, + kernel_size=3, padding=1, bias=False) + self.bn = nn.BatchNorm2d(hidden_ch) + self.act = nn.ReLU(inplace=True) + + # 1x1 expand + self.expand = nn.Sequential( + nn.Conv2d(hidden_ch, out_ch, kernel_size=1, bias=False), + nn.BatchNorm2d(out_ch), + ) + + self.residual = (in_ch == out_ch) + + def forward(self, x): + identity = x + + x = self.reduce(x) + offset = self.offset_conv(x) + + if offset.shape[1] == 18: # DCNv1 + x = self.deform(x, offset) + else: # DCNv2: last 9 channels are mask + o, mask = offset.split([18, 9], dim=1) + mask = mask.sigmoid() + x = self.deform(x, o, mask) + + x = self.act(self.bn(x)) + x = self.expand(x) + + if self.residual: + x = x + identity + return x diff --git a/src/models/sofia_v1/stripnet_blocks_dcn_new.py b/src/models/sofia_v1/stripnet_blocks_dcn_new.py new file mode 100644 index 0000000..4b312fb --- /dev/null +++ b/src/models/sofia_v1/stripnet_blocks_dcn_new.py @@ -0,0 +1,125 @@ +import torch +import torch.nn as nn +from DCNv4 import DCNv4 # ← your installed OpenGVLab version +import math + +class DCNBlockV4(nn.Module): + """ + StripNet-style block that uses OpenGVLab DCNv4 instead of torchvision DCNv2. + """ + + def __init__(self, in_ch, out_ch, hidden_ratio=0.25, + kernel_size=3, stride=1, dilation=1, group=4, + offset_scale=1.0, use_bias=False): + super().__init__() + assert kernel_size in (3, 5, 7) + pad = (kernel_size // 2) * dilation + + # Hidden channels — must satisfy (hidden_ch // group) % 16 == 0 + hidden_ch = max(16, int(out_ch * hidden_ratio)) + hidden_ch = math.ceil(hidden_ch / 16) * 16 + # increase until kernel constraint satisfied + while (hidden_ch // group) % 16 != 0: + hidden_ch += 16 + + #print(f"[DCNv4] adjusted hidden_ch={hidden_ch}, group={group}") + + + self.reduce = nn.Sequential( + nn.Conv2d(in_ch, hidden_ch, 1, bias=False), + self.make_gn(hidden_ch), + nn.ReLU(inplace=True), + ) + + # DCNv4 core + self.dcn = DCNv4( + channels=hidden_ch, + kernel_size=kernel_size, + stride=stride, + pad=pad, + dilation=dilation, + group=group, + offset_scale=offset_scale, + dw_kernel_size=None, + center_feature_scale=False, + remove_center=False, + output_bias=True, + without_pointwise=False, + ) + + + self.expand = nn.Sequential( + nn.Conv2d(hidden_ch, out_ch, 1, bias=False), + self.make_gn(out_ch), + + ) + + self.residual = (in_ch == out_ch and stride == 1) + self.stride = stride + + def make_gn(self, num_channels): + num_groups = max(1, num_channels // 16) + return nn.GroupNorm(num_groups, num_channels) + + def forward(self, x): + """ + Input: [B, C, H, W] + Output: [B, C, H', W'] + """ + identity = x + + x = self.reduce(x) + B, C, H, W = x.shape + + # Flatten for DCNv4 + x_seq = x.permute(0, 2, 3, 1).contiguous().view(B, H * W, C) + + # --- clamp activations to avoid huge values --- + x_seq = torch.clamp(x_seq, -50.0, 50.0) + + x_seq = self.dcn(x_seq) + x_seq = torch.nan_to_num(x_seq, nan=0.0, posinf=1e4, neginf=-1e4) + + # --- back to (B, C, H, W) --- + x = x_seq.view(B, H, W, C).permute(0, 3, 1, 2).contiguous() + + # --- expand + residual --- + x = self.expand(x) + if self.residual: + x = x + identity + return x + + +def test_dcnblock_v4(): + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"Using device: {device}") + + block = DCNBlockV4( + in_ch=128, # must match input + out_ch=128, + hidden_ratio=0.25, + kernel_size=3, + stride=1, + dilation=1, + group=5, + offset_scale=1.0, + ).to(device) + + # ✅ internal feature map (not RGB) + x = torch.randn(2, 128, 128, 128, device=device, requires_grad=True) + + with torch.cuda.amp.autocast_mode.autocast(enabled=torch.cuda.is_available()): + y = block(x) + + print(f"Input shape : {x.shape}") + print(f"Output shape : {y.shape}") + assert y.shape == x.shape + + loss = y.mean() + loss.backward() + print("✅ DCNBlockV4 test passed.\n") + + + +if __name__ == "__main__": + test_dcnblock_v4() diff --git a/src/models/sofia_v1/stripnet_model_dcn.py b/src/models/sofia_v1/stripnet_model_dcn.py new file mode 100644 index 0000000..a731102 --- /dev/null +++ b/src/models/sofia_v1/stripnet_model_dcn.py @@ -0,0 +1,117 @@ +import torch +import torch.nn as nn + +from src.models.sofia_v1.stripnet_blocks_dcn import DCNBlock +from src.models.sofia_v1.stripnet_blocks_dcn_new import DCNBlockV4 +from src.models.stripnet.model import OverlapPatchEmbed + + +class StripNetDCN(nn.Module): + """4-stage hierarchical CNN backbone: OverlapPatchEmbed + DCN blocks per stage. + + Output: list of [B, C_i, H_i, W_i] features at each stage. For 256x256 input + with default downsampling (4, 2, 2, 2), final stage is 8x8. + + DCN variant: + - "v2" (default): torchvision `DeformConv2d` — stable, no memory leaks. + - "v4": OpenGVLab DCNv4 — faster on CUDA but has a known C++ extension + memory leak (~9 MB per forward call) that causes OOM in long training + runs. Only use if you have a patched DCNv4 build. + """ + + def __init__( + self, + in_chans: int = 3, + embed_dims: list[int] = [64, 128, 256, 512], + depths: list[int] = [3, 4, 6, 3], + dcn_variant: str = "v2", + ) -> None: + super().__init__() + if dcn_variant not in ("v2", "v4"): + raise ValueError(f"dcn_variant must be 'v2' or 'v4', got {dcn_variant!r}") + self.num_stages = len(embed_dims) + self.embed_dims = embed_dims + self.dcn_variant = dcn_variant + + 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 self.embed_dims[i - 1], + embed_dim=self.embed_dims[i], + ) + + block_cls = DCNBlockV4 if dcn_variant == "v4" else DCNBlock + block = nn.ModuleList([ + block_cls( + in_ch=self.embed_dims[i], + out_ch=self.embed_dims[i], + ) for _ in range(depths[i]) + ]) + + setattr(self, f"patch_embed{i + 1}", patch_embed) + setattr(self, f"block{i + 1}", block) + + def forward_features(self, x: torch.Tensor) -> list[torch.Tensor]: + outs = [] + for i in range(self.num_stages): + patch_embed = getattr(self, f"patch_embed{i + 1}") + block = getattr(self, f"block{i + 1}") + x, H, W = patch_embed(x) + for blk in block: + x = blk(x) + outs.append(x) + return outs + + def forward(self, x: torch.Tensor) -> list[torch.Tensor]: + return self.forward_features(x) + +def get_stripnet_dcn_small(dcn_variant: str = "v2") -> StripNetDCN: + return StripNetDCN( + in_chans=3, + embed_dims=[64, 128, 320, 512], + depths=[2, 2, 4, 2], + dcn_variant=dcn_variant, + ) + + +def get_stripnet_dcn_small_v2(dcn_variant: str = "v2") -> StripNetDCN: + return StripNetDCN( + in_chans=3, + embed_dims=[64, 128, 256, 384], + depths=[2, 2, 4, 2], + dcn_variant=dcn_variant, + ) + + +def get_stripnet_dcn_tiny(dcn_variant: str = "v2") -> StripNetDCN: + return StripNetDCN( + in_chans=3, + embed_dims=[32, 64, 128, 256], + depths=[3, 3, 5, 2], + dcn_variant=dcn_variant, + ) + + +def get_stripnet_dcn_tiny_tiny(dcn_variant: str = "v2") -> StripNetDCN: + return StripNetDCN( + in_chans=3, + embed_dims=[16, 32, 80, 128], + depths=[3, 3, 5, 2], + dcn_variant=dcn_variant, + ) + + +VARIANT_MAP = { + "small": get_stripnet_dcn_small, + "small_v2": get_stripnet_dcn_small_v2, + "tiny": get_stripnet_dcn_tiny, + "tiny_tiny": get_stripnet_dcn_tiny_tiny, +} + + +def build_stripnet_dcn(variant: str = "small", dcn_variant: str = "v2") -> StripNetDCN: + """Factory: variant in {tiny_tiny, tiny, small, small_v2}, dcn_variant in {v2, v4}.""" + if variant not in VARIANT_MAP: + raise ValueError(f"Unknown variant {variant!r}. Available: {list(VARIANT_MAP)}") + return VARIANT_MAP[variant](dcn_variant=dcn_variant) \ No newline at end of file diff --git a/src/models/sofia_v71/README.md b/src/models/sofia_v71/README.md new file mode 100644 index 0000000..efb0fb4 --- /dev/null +++ b/src/models/sofia_v71/README.md @@ -0,0 +1,331 @@ +# SOFIA v7.1 — PyTorch implementation + +Reference implementation of **SOFIA v7.1** CVGL student model targeting Jetson Orin NX with 500 MB – 1 GB VRAM after INT8 quantization. + +Full design rationale: see `2_hypotesis/temp_hypotesis/HYP_SOFIA_v7_UltraDeep_дизайн.md`. + +## Architecture + +``` +Input (256×256×3, sat or UAV) + ↓ +Stem: dual-conv (3→C_mid→C_stem_out), ×2 down + ↓ +Stage 1 (shared sat/UAV): StripDCN-lite × d1, ×2 down + ↓ +Stage 2 (shared): StripMixConv × d2, ×2 down + ↓ +Stage 3 (separate): MambaVision MV5 × d3, ×2 down + ↓ +Stage 4 (separate): MambaVision MV1 × d4, ×2 down (→ 8×8) + ↓ +Ultra-lite 1×1 Neck → F̃ [B, C_n, 8, 8] + ↓ +┌─ Sat-Head: GGeM → BN → Linear → L2 → g_sat +└─ UAV-Head: FiLM(altitude) → CHP (polar+FFT+mag) → BN → Linear → L2 → g_uav +``` + +## Presets + +| Preset | Params | FLOPs | INT8 weights | FP16 weights | Target latency | +|--------|-------:|------:|-------------:|-------------:|---------------:| +| **M** (default) | ~500 M | ~132 G | ~500 MB | ~1 GB | ~18 ms | +| **L** | ~1 B | ~283 G | ~1 GB | ~2 GB | ~20 ms | +| Tiny | ~5 M | ~1.4 G | ~5 MB | ~10 MB | ~4 ms | + +*Latency estimates на Jetson Orin NX 8GB (INT8 TRT mixed-precision для DCN/Mamba в FP16).* + +## Dependencies + +Required: +- Python ≥ 3.9 +- PyTorch ≥ 2.0 +- torchvision ≥ 0.15 (для `deform_conv2d`) + +Optional (для production Mamba speed): +- [`mamba_ssm`](https://github.com/state-spaces/mamba) — ускоренные backends: + - v1 `selective_scan_fn` для Mamba-1 (~5× vs Python loop) + - v2 `Mamba2` модуль для Mamba-2 SSD dual form (~2–8× vs Mamba-1) +- [`causal-conv1d`](https://github.com/Dao-AILab/causal-conv1d) — ускоренный 1D conv для Mamba + +## Mamba variants (приоритет по качеству/скорости) + +| Variant | Описание | Speed | Quality | Когда использовать | +|---------|----------|:-----:|:-------:|--------------------| +| **`mamba2`** (default) | SSD dual form, scalar A per head | **2–8× faster** | best | Main choice если mamba_ssm v2 доступен | +| `mamba1` | Original selective scan, diagonal A | ref | +0% vs mamba2 | Legacy / reproducibility / если v2 недоступен | +| `efficient_vmamba` | Atrous scan (2 directions, no CUDA kernel) | 2–3× faster | ~−0.3% R@1 | Speed fallback без mamba_ssm | + +### Выбор через config + +```python +from code_sofia_v71 import sofia_m_config, SOFIAv71 + +cfg = sofia_m_config() + +# Вариант 1: Mamba-2 (default, preferred) — если mamba_ssm v2 доступен +cfg.mamba_variant = "mamba2" +cfg.mamba_backend = "auto" # falls back to torch if unavailable + +# Вариант 2: EfficientVMamba — speed без зависимости от mamba_ssm +cfg.mamba_variant = "efficient_vmamba" + +# Вариант 3: Mamba-1 — legacy / для сравнения +cfg.mamba_variant = "mamba1" + +model = SOFIAv71(cfg) +``` + +### Проверка доступности backends + +```python +from code_sofia_v71 import is_mamba_ssm_available, is_mamba2_available + +print(f"Mamba-1 CUDA: {is_mamba_ssm_available()}") +print(f"Mamba-2 CUDA: {is_mamba2_available()}") +``` + +### Несовместимость параметров между variants + +Mamba-1 и Mamba-2 используют **разную параметризацию $A$** (diagonal-per-channel vs scalar-per-head) — state_dict **НЕ совместим**. EfficientVMamba имеет N независимых scanners каждый с собственными параметрами Mamba-1. + +**Следствие:** нельзя train в `mamba1` и load в `mamba2` checkpoint. Выбор variant нужно зафиксировать до training. + +## Quick start + +```python +import torch +from code_sofia_v71 import build_sofia + +# Default = SOFIA-M (500 MB INT8 target) +model = build_sofia("M") +model.eval() + +sat = torch.randn(2, 3, 256, 256) # satellite image +uav = torch.randn(2, 3, 256, 256) # UAV image +altitude = torch.tensor([120.0, 450.0]) # meters + +out = model(sat=sat, uav=uav, altitude=altitude) +g_sat = out["g_sat"] # [2, 512] L2-normalized +g_uav = out["g_uav"] # [2, 512] L2-normalized + +# Retrieval similarity +similarity = (g_sat * g_uav).sum(dim=-1) # cosine +``` + +## Verification + +```bash +# Smoke test (Tiny preset, CPU, fast) +python -m sofia_v71.verify + +# Full SOFIA-M check with latency benchmark (GPU) +python -m sofia_v71.verify --preset M --device cuda --benchmark + +# Maximum scale +python -m sofia_v71.verify --preset L --device cuda +``` + +Expected output for SOFIA-M: +``` +Total params: ~500 M +... +FP16 weights: ~1000 MB (~0.98 GB) +INT8 weights: ~500 MB (~0.49 GB) +... +out[g_sat]: (1, 512) +out[g_uav]: (1, 512) +``` + +## Key novel components + +### `layers.py` +- **`GGeM`** — per-channel learnable exponent Generalized Mean pooling (F11) +- **`CircularHarmonicPool`** — formally SO(2)-invariant UAV pool via polar → 1D FFT → magnitude (NOVEL NH2) +- **`AltitudeFiLM`** — telemetry-conditioned modulation (NOVEL NH4) +- **`RoPE2D`** — 2D rotary positional embedding + +### `blocks.py` +- **`StripDCNLiteBlock`** — Strip DW MBConv с DCN offset на одной оси (NOVEL) +- **`StripMixConvBlock`** — Strip + MixConv (3/5/7 kernels) для multi-scale +- **`MambaVisionBlock`** — Mamba ∥ MHSA [∥ Strip] + FFN (MV5 variant — NOVEL 3-way) +- **`SimpleMambaBlock`** — reference pure-PyTorch Mamba-1 (replace with `mamba_ssm` в production) + +### `model.py` +- **`SOFIAv71`** — full model with optional weight-sharing +- **`SatHead` vs `UAVHead`** — asymmetric physics-motivated design (NOVEL NH1) +- **`RingAuxHead`** — LPN Square-Ring training-only aux + +## KD taps + +Backbone forward returns features at stages s0/f1/f2/f3/f4. Использовать для +hierarchical knowledge distillation: + +```python +out = model(sat=sat, uav=uav, return_features=True) +f3_sat = out["features_sat"]["f3"] # для teacher feature alignment +``` + +## Production notes + +### Mamba backend selection + +Default `mamba_variant="mamba2"` с `backend="auto"`: +- Use `Mamba2` from `mamba_ssm.modules.mamba2` если установлено (CUDA, fast) +- Иначе fallback на pure-PyTorch simplified SSD scan (slow, работает) + +Для Mamba-1 legacy путь — `mamba_variant="mamba1"`, backend так же резолвится. + +Для EfficientVMamba — pure PyTorch, не нуждается в mamba_ssm. + +## Quantization (`quant.py`) + +Reusable utilities для INT8 PTQ/QAT в SOFIA. + +### `OffsetClampSTE` — DCN-M2 + +Hard clamp DCN offsets к `[-k, +k]` со STE backward. Уже подключён в `StripDCNLiteBlock` через флаг `use_offset_clamp_ste=True` (default). + +```python +from code_sofia_v71 import OffsetClampSTE, offset_clamp_ste + +# Module form +clamp = OffsetClampSTE(kernel_size=7) +clamped = clamp(offsets) + +# Functional form +clamped = offset_clamp_ste(offsets, kernel_size=7) +``` + +Backward — identity (gradient проходит насквозь). Можно ставить с epoch 1, не только в QAT. + +### `KScaledFakeQuant` — k-scaled fake quantization + +Multi-bin fake quant для long-tail distributions (Mamba Δ, y). + +```python +from code_sofia_v71 import KScaledFakeQuant + +fq = KScaledFakeQuant(num_bins=3) + +# Calibration +fq.start_calibration() +with torch.no_grad(): + for batch in cal_loader: + _ = model(batch) +fq.finalize_calibration(percentile=99.9) +# fq теперь в PTQ-режиме +``` + +### `KScaledMamba2Block` — drop-in для Mamba2Block + +Подкласс `Mamba2Block` с k-scaled fake-quant узлами на критических путях +(`x_main`, `delta`, `y`). Внутреннее состояние scan'а `h_t` остаётся в +model dtype (R5 reparam principle). + +```python +from code_sofia_v71 import KScaledMamba2Block + +mamba = KScaledMamba2Block( + channels=192, + d_state=64, + headdim=64, + num_bins=3, + targets=("x_main", "delta", "y"), # subset +) + +mamba.start_calibration() +# ... run calibration data ... +mamba.finalize_calibration(percentile=99.9) +``` + +**Constraint:** только `backend='torch'` поддерживается (mamba_ssm CUDA kernel +не имеет hooks для k-scaled fake-quant). Для full INT8 deploy на TRT нужен +custom plugin — отдельный deploy concern. + +### Model-wide helpers + +```python +from code_sofia_v71 import start_calibration, finalize_calibration, set_quant_enabled + +# Начать калибровку всех KScaledFakeQuant и KScaledMamba2Block в модели +start_calibration(model) +with torch.no_grad(): + for batch in cal_loader: + _ = model(batch) +finalize_calibration(model, percentile=99.9) + +# Toggle on/off для FP-vs-INT8 ablation +set_quant_enabled(model, False) # FP forward +y_fp = model(x) +set_quant_enabled(model, True) +y_q = model(x) +``` + +### Smoke test + +```bash +python -m sofia_v71.quant +``` + +Прогонит unit-test на: +- DCN-M2 clamp (gradient pass-through) +- KScaledFakeQuant calibration + quant error +- KScaledMamba2Block FP-vs-INT8 diff + +### DCN INT8 QAT +StripDCN использует `torchvision.ops.deform_conv2d`. Для INT8 deploy на TRT: +- Применить QAT с DCN-M1..M4 modifications (см. HYP Phase 7) +- Offset predictor: per-channel scale, offset clamping `[-k, +k]` +- Mask: FP16 micro-block внутри INT8 graph + +### TensorRT export +```python +import torch +model.eval() +model_fuse = model # apply fuse_reparam passes separately +torch.onnx.export( + model, + (sat, uav, altitude), + "sofia_v71.onnx", + input_names=["sat", "uav", "altitude"], + output_names=["g_sat", "g_uav"], + opset_version=17, +) +# Then: +# trtexec --onnx=sofia_v71.onnx --int8 --fp16 --saveEngine=sofia.plan +``` + +## Training + +See `HYP_SOFIA_v7_UltraDeep_дизайн.md` Phase 8 for full training recipe: +- Loss: InfoNCE (Sample4Geo mining) + Ring aux + (opt) cross-view consistency +- Curriculum: PALW sigmoid warmup, 60 epochs, 4 phases +- Optimizer: AdamW, LR 3e-4 cosine, wd 0.05 +- Temperature: τ 0.1 → 0.01 cosine decay + +## File layout + +``` +code_sofia_v71/ +├── __init__.py — public exports +├── config.py — SOFIAConfig + presets (M/L/Tiny) +├── layers.py — GGeM, CHP, FiLM, RoPE2D, SE, LayerNorm2d +├── blocks.py — StripDCN, StripMixConv, Mamba, MambaVision, Downsample +├── model.py — Stem, Backbone, Neck, Heads, SOFIAv71 +├── verify.py — parameter counter + benchmark script +└── README.md — this file +``` + +## References to design doc + +Все architectural decisions обоснованы в +`2_hypotesis/temp_hypotesis/HYP_SOFIA_v7_UltraDeep_дизайн.md`: +- Phase 1: requirements R1–R13 +- Phase 2': DCN / Strip / MambaVision operator catalog +- Phase 3': backbone design with 4 candidates (E/F/G) +- Phase 4'': CVGL-Aware Head v7.1-α (Asymmetric + CHP + FiLM) +- Phase 5': ablation matrix +- Phase 6: MambaVision MV5 operational details +- Phase 7: StripDCN QAT strategy +- Phase 8: training pipeline diff --git a/src/models/sofia_v71/__init__.py b/src/models/sofia_v71/__init__.py new file mode 100644 index 0000000..40681ed --- /dev/null +++ b/src/models/sofia_v71/__init__.py @@ -0,0 +1,126 @@ +"""SOFIA v7.1 — student model for cross-view geo-localization (CVGL). + +Architecture: + stem → stage1 (StripDCN-lite) → stage2 (StripMixConv) + → stage3 (MambaVision MV5) → stage4 (MambaVision MV1) + → 1×1 neck → asymmetric Sat/UAV heads + +Key novel components: +- StripDCN-lite: Strip DW with adaptive offset on one axis +- MambaVision MV5: Mamba ∥ MHSA ∥ Strip 3-way parallel +- CircularHarmonicPool: formally SO(2)-invariant UAV head +- AltitudeFiLM: telemetry-aware conditioning + +Presets (see config.py): +- SOFIA-M (~500 M params, 500 MB INT8) — default +- SOFIA-L (~1 B params, 1 GB INT8) — max scale +- SOFIA-Tiny (~5 M) — reference from original v7.1 spec + +Quick start: + from sofia_v71 import build_sofia + model = build_sofia("M") + out = model(sat=sat_tensor, uav=uav_tensor, altitude=alt_tensor) + g_sat, g_uav = out["g_sat"], out["g_uav"] + +See README.md and HYP_SOFIA_v7_UltraDeep_дизайн.md for full design rationale. +""" + +from .config import ( + SOFIAConfig, + sofia_l_config, + sofia_m_config, + sofia_tiny_config, + DEFAULT_CONFIG, +) +from .model import ( + SOFIAv71, + build_sofia, + Backbone, + Stem, + UltraLiteNeck, + SatHead, + UAVHead, + RingAuxHead, +) +from .layers import ( + GGeM, + CircularHarmonicPool, + AltitudeFiLM, + TextFiLM, + RoPE2D, + SqueezeExcite, + LayerNorm2d, +) +from .blocks import ( + StripDCNLiteBlock, + StripMixConvBlock, + SimpleMambaBlock, + Mamba2Block, + EfficientVMambaBlock, + MambaVisionBlock, + EVSSBridge, + Downsample, + build_mamba_block, + is_mamba_ssm_available, + is_mamba2_available, +) +from .quant import ( + # DCN-M2 + OffsetClampSTE, + offset_clamp_ste, + # k-scaled + KScaledFakeQuant, + KScaledMamba2Block, + # model-wide helpers + start_calibration, + finalize_calibration, + set_quant_enabled, +) + +__version__ = "0.1.0" +__all__ = [ + # Config + "SOFIAConfig", + "sofia_m_config", + "sofia_l_config", + "sofia_tiny_config", + "DEFAULT_CONFIG", + # Top-level + "SOFIAv71", + "build_sofia", + # Model parts + "Backbone", + "Stem", + "UltraLiteNeck", + "SatHead", + "UAVHead", + "RingAuxHead", + # Layers + "GGeM", + "CircularHarmonicPool", + "AltitudeFiLM", + "TextFiLM", + "RoPE2D", + "SqueezeExcite", + "LayerNorm2d", + # Blocks + "StripDCNLiteBlock", + "StripMixConvBlock", + "SimpleMambaBlock", + "Mamba2Block", + "EfficientVMambaBlock", + "MambaVisionBlock", + "EVSSBridge", + "Downsample", + "build_mamba_block", + "is_mamba_ssm_available", + "is_mamba2_available", + # Quantization + "OffsetClampSTE", + "offset_clamp_ste", + "KScaledFakeQuant", + "KScaledMamba2Block", + "start_calibration", + "finalize_calibration", + "set_quant_enabled", +] diff --git a/src/models/sofia_v71/blocks.py b/src/models/sofia_v71/blocks.py new file mode 100644 index 0000000..b2d39af --- /dev/null +++ b/src/models/sofia_v71/blocks.py @@ -0,0 +1,1186 @@ +"""SOFIA v7.1 stage building blocks. + +Contains: +- StripDCNLiteBlock: Stage 1 block (strip DW + optional DCN offset) +- StripMixConvBlock: Stage 2 block (strip + mixed kernel DW) +- SimpleMambaBlock: Simplified Mamba-1 for demonstration (replace with mamba_ssm in production) +- MambaVisionBlock: Stage 3-4 (Mamba ∥ MHSA [+ optional Strip]) + FFN + +Novel blocks vs literature: +- StripDCN-lite: asymmetric DCN (one strip adaptive, one plain) +- MambaVision MV5: Mamba + MHSA + Strip 3-way parallel (first for CVGL) +""" + +from __future__ import annotations + +from typing import List, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +try: + from torchvision.ops import deform_conv2d + _DEFORM_CONV_AVAILABLE = True +except ImportError: + _DEFORM_CONV_AVAILABLE = False + +# Optional mamba_ssm for fast CUDA selective scan (Mamba-1) +try: + from mamba_ssm.ops.selective_scan_interface import selective_scan_fn + _MAMBA_SSM_AVAILABLE = True +except ImportError: + selective_scan_fn = None + _MAMBA_SSM_AVAILABLE = False + +# Optional Mamba-2 SSD implementation +try: + from mamba_ssm.modules.mamba2 import Mamba2 as _ExternalMamba2 + _MAMBA2_AVAILABLE = True +except ImportError: + _ExternalMamba2 = None + _MAMBA2_AVAILABLE = False + +# Optional causal-conv1d for fast 1D conv in Mamba (used by mamba_ssm backend) +try: + from causal_conv1d import causal_conv1d_fn + _CAUSAL_CONV1D_AVAILABLE = True +except ImportError: + causal_conv1d_fn = None + _CAUSAL_CONV1D_AVAILABLE = False + + +def is_mamba_ssm_available() -> bool: + """Check whether mamba_ssm v1 CUDA backend is available (Mamba-1).""" + return _MAMBA_SSM_AVAILABLE + + +def is_mamba2_available() -> bool: + """Check whether mamba_ssm Mamba-2 (SSD) CUDA backend is available.""" + return _MAMBA2_AVAILABLE + + +from .layers import LayerNorm2d, RoPE2D, SqueezeExcite + + +# ============================================================ +# StripDCN-Lite Block (Stage 1) +# ============================================================ + +class StripDCNLiteBlock(nn.Module): + """StripDCN-lite: Strip MBConv with DCN offset on horizontal strip only. + + Structure: + x → Expand 1x1 → 4C + ├── Strip-H DW 1×k (with DCN offset if available) + ├── Strip-V DW k×1 (plain) + → Sum → BN → SiLU + → SE + → Project 1x1 → C + + skip (if stride=1 and in=out) + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + expand: int = 4, + kernel: int = 7, + stride: int = 1, + se_ratio: int = 16, + use_dcn: bool = True, + use_offset_clamp_ste: bool = True, + ) -> None: + super().__init__() + self.in_channels = in_channels + self.out_channels = out_channels + self.stride = stride + self.kernel = kernel + self.use_dcn = use_dcn and _DEFORM_CONV_AVAILABLE + self.use_offset_clamp_ste = use_offset_clamp_ste + mid = expand * out_channels + self.mid = mid + + # Expand + self.expand = nn.Sequential( + nn.Conv2d(in_channels, mid, 1, bias=False), + nn.BatchNorm2d(mid), + nn.SiLU(inplace=True), + ) + + # Strip-H: DW 1×k (stride on W only) + pad_h = kernel // 2 + + if self.use_dcn: + # DCN offset+mask predictor for strip-H + # For strip kernel (1, k): offset channels = 2*1*k = 2k, mask = k + self.offset_h = nn.Conv2d(mid, 2 * kernel, 1) + self.mask_h = nn.Conv2d(mid, kernel, 1) + nn.init.zeros_(self.offset_h.weight) + nn.init.zeros_(self.offset_h.bias) + nn.init.zeros_(self.mask_h.weight) + nn.init.zeros_(self.mask_h.bias) + + # DCN-M2: optional offset clamp via STE (lazy import to avoid cycle) + if use_offset_clamp_ste: + from .quant import OffsetClampSTE + self.offset_clamp = OffsetClampSTE(kernel_size=kernel) + else: + self.offset_clamp = None + + # Strip-H DW weights stored as buffer-friendly parameter + self.strip_h_weight = nn.Parameter( + torch.zeros(mid, 1, 1, kernel) + ) + nn.init.kaiming_normal_(self.strip_h_weight, mode="fan_out", nonlinearity="linear") + self.strip_h_stride = (1, stride) + self.strip_h_padding = (0, pad_h) + else: + self.strip_h = nn.Conv2d( + mid, mid, + kernel_size=(1, kernel), + stride=(1, stride), + padding=(0, pad_h), + groups=mid, + bias=False, + ) + + # Strip-V: DW k×1 (stride on H only) + self.strip_v = nn.Conv2d( + mid, mid, + kernel_size=(kernel, 1), + stride=(stride, 1), + padding=(pad_h, 0), + groups=mid, + bias=False, + ) + + self.bn_strip = nn.BatchNorm2d(mid) + self.act_strip = nn.SiLU(inplace=True) + + # SE + self.se = SqueezeExcite(mid, reduction=se_ratio) + + # Project + self.project = nn.Sequential( + nn.Conv2d(mid, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + ) + + self.use_skip = (stride == 1) and (in_channels == out_channels) + + def _apply_strip_h_deform(self, x: torch.Tensor) -> torch.Tensor: + """DCN-enabled strip-H application.""" + offset = self.offset_h(x) # [B, 2k, H, W] + # DCN-M2: hard clamp offsets to [-k, +k] with STE backward + if self.offset_clamp is not None: + offset = self.offset_clamp(offset) + mask = torch.sigmoid(self.mask_h(x)) # [B, k, H, W] + # deform_conv2d signature: (x, offset, weight, bias=None, stride, padding, dilation, mask) + # weight shape: [mid, 1, 1, k] (depthwise: groups = mid via groups param via w shape) + # Wait: torchvision deform_conv2d uses weight [out_c, in_c/groups, kH, kW] + # For DW: groups=mid → weight [mid, 1, 1, k] + return deform_conv2d( + input=x, + offset=offset, + weight=self.strip_h_weight, + bias=None, + stride=self.strip_h_stride, + padding=self.strip_h_padding, + dilation=(1, 1), + mask=mask, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + identity = x + y = self.expand(x) + + if self.use_dcn: + strip_h = self._apply_strip_h_deform(y) + else: + strip_h = self.strip_h(y) + strip_v = self.strip_v(y) + + # If stride>1, spatial shapes match since both strips apply stride in their axis + y = strip_h + strip_v + y = self.bn_strip(y) + y = self.act_strip(y) + + y = self.se(y) + y = self.project(y) + + if self.use_skip: + y = y + identity + return y + + +# ============================================================ +# StripMixConv Block (Stage 2) +# ============================================================ + +class StripMixConvBlock(nn.Module): + """Stage 2 block: Strip H + Strip V + MixConv 3×3 (all DW) summed. + + Provides simultaneous anisotropic (R1) and multi-scale (R2) features. + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + expand: int = 4, + strip_kernel: int = 5, + mix_kernels: List[int] = (3, 5, 7), + stride: int = 1, + se_ratio: int = 16, + ) -> None: + super().__init__() + self.stride = stride + mid = expand * out_channels + self.mid = mid + + # Expand + self.expand = nn.Sequential( + nn.Conv2d(in_channels, mid, 1, bias=False), + nn.BatchNorm2d(mid), + nn.SiLU(inplace=True), + ) + + # Strip-H + pad_h = strip_kernel // 2 + self.strip_h = nn.Conv2d( + mid, mid, + kernel_size=(1, strip_kernel), + stride=(1, stride), + padding=(0, pad_h), + groups=mid, + bias=False, + ) + self.strip_v = nn.Conv2d( + mid, mid, + kernel_size=(strip_kernel, 1), + stride=(stride, 1), + padding=(pad_h, 0), + groups=mid, + bias=False, + ) + + # MixConv: split mid into len(mix_kernels) groups, DW with different kernels + self.mix_kernels = list(mix_kernels) + n_groups = len(self.mix_kernels) + assert mid % n_groups == 0, f"mid {mid} must be divisible by {n_groups}" + group_size = mid // n_groups + self.group_size = group_size + + mix_layers = [] + for k in self.mix_kernels: + pad = k // 2 + mix_layers.append( + nn.Conv2d( + group_size, group_size, + kernel_size=k, + stride=stride, + padding=pad, + groups=group_size, + bias=False, + ) + ) + self.mix_dws = nn.ModuleList(mix_layers) + + self.bn_merge = nn.BatchNorm2d(mid) + self.act_merge = nn.SiLU(inplace=True) + + self.se = SqueezeExcite(mid, reduction=se_ratio) + + self.project = nn.Sequential( + nn.Conv2d(mid, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + ) + + self.use_skip = (stride == 1) and (in_channels == out_channels) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + identity = x + y = self.expand(x) + + # Strip branches + s_h = self.strip_h(y) + s_v = self.strip_v(y) + + # MixConv branch: split along C, apply DW, concat + y_chunks = torch.split(y, self.group_size, dim=1) + mix_outs = [dw(chunk) for dw, chunk in zip(self.mix_dws, y_chunks)] + mix = torch.cat(mix_outs, dim=1) + + y = s_h + s_v + mix + y = self.bn_merge(y) + y = self.act_merge(y) + + y = self.se(y) + y = self.project(y) + + if self.use_skip: + y = y + identity + return y + + +# ============================================================ +# Simplified Mamba-1 Block (replace with mamba_ssm in production) +# ============================================================ + +class SimpleMambaBlock(nn.Module): + """Mamba-1 block with optional mamba_ssm CUDA backend. + + Backend selection: + - "mamba_ssm": uses selective_scan_fn from state-spaces/mamba (5-10× faster, CUDA only) + - "torch": pure PyTorch sequential scan (slow, portable, works on CPU) + - "auto": mamba_ssm if importable, else torch + + Parameter layout matches canonical Mamba-1 so checkpoints could be + transferred between backends (same state_dict keys). + """ + + def __init__( + self, + channels: int, + d_state: int = 16, + dt_rank: Optional[int] = None, + conv_kernel: int = 3, + backend: str = "auto", + ) -> None: + super().__init__() + self.d_model = channels + self.d_state = d_state + self.dt_rank = dt_rank or max(1, channels // 16) + self.conv_kernel = conv_kernel + + # Resolve backend + if backend == "auto": + backend = "mamba_ssm" if _MAMBA_SSM_AVAILABLE else "torch" + if backend == "mamba_ssm" and not _MAMBA_SSM_AVAILABLE: + raise RuntimeError( + "mamba_ssm backend requested but package not installed. " + "Install via `pip install mamba-ssm causal-conv1d`." + ) + self.backend = backend + + # In projection splits into (gate, x_ssm) + self.in_proj = nn.Linear(channels, 2 * channels, bias=False) + + # Local DW 1D conv + self.conv1d = nn.Conv1d( + channels, channels, + kernel_size=conv_kernel, + padding=conv_kernel // 2, + groups=channels, + bias=True, + ) + + # Selective params + self.x_proj = nn.Linear( + channels, self.dt_rank + 2 * d_state, bias=False + ) + self.dt_proj = nn.Linear(self.dt_rank, channels, bias=True) + + # State transition A (stored as log for stability; HiPPO-style negative init) + A = torch.arange(1, d_state + 1, dtype=torch.float32).unsqueeze(0).repeat(channels, 1) + self.A_log = nn.Parameter(torch.log(A)) + + # Skip scale D + self.D = nn.Parameter(torch.ones(channels)) + + # Out projection + self.out_proj = nn.Linear(channels, channels, bias=False) + + # --------------------------------------------- + # mamba_ssm backend (fast, CUDA-only) + # --------------------------------------------- + def _forward_mamba_ssm(self, x: torch.Tensor) -> torch.Tensor: + """Fast forward using state-spaces/mamba selective_scan_fn. + + x: [B, L, C] → y: [B, L, C] + """ + B, L, C = x.shape + + # In proj + z = self.in_proj(x) # [B, L, 2C] + x_gate, x_ssm = z.chunk(2, dim=-1) + + # Local conv: mamba_ssm works with [B, C, L] layout + x_ssm_t = x_ssm.transpose(1, 2) # [B, C, L] + x_ssm_t = self.conv1d(x_ssm_t)[..., :L] # enforce len L (padding may add) + x_ssm_t = F.silu(x_ssm_t) + + # Selective params (on [B, L, C]) + x_ssm_seq = x_ssm_t.transpose(1, 2) + delta_bc = self.x_proj(x_ssm_seq) + delta_raw, B_p, C_p = torch.split( + delta_bc, + [self.dt_rank, self.d_state, self.d_state], + dim=-1, + ) + # selective_scan_fn expects: + # u: [B, D, L] (our x_ssm_t) + # delta: [B, D, L] + # A: [D, N] + # B: [B, N, L] + # C: [B, N, L] + # D: [D] + # z: optional gate (we use external gate) + delta = F.softplus(self.dt_proj(delta_raw)).transpose(1, 2) # [B, C, L] + A = -torch.exp(self.A_log.float()) # [C, N], fp32 for stability + B_ssm = B_p.transpose(1, 2) # [B, N, L] + C_ssm = C_p.transpose(1, 2) # [B, N, L] + + y = selective_scan_fn( + u=x_ssm_t, + delta=delta, + A=A, + B=B_ssm, + C=C_ssm, + D=self.D.float(), + z=None, + delta_bias=None, + delta_softplus=False, # already applied + return_last_state=False, + ) # [B, C, L] + + y = y.transpose(1, 2) # [B, L, C] + y = y * F.silu(x_gate) + y = self.out_proj(y) + return y + + # --------------------------------------------- + # Pure-PyTorch backend (slow, portable) + # --------------------------------------------- + def _forward_torch(self, x: torch.Tensor) -> torch.Tensor: + """Sequential scan forward. Python loop over L — slow but correct.""" + B, L, C = x.shape + + # In proj + z = self.in_proj(x) + x_gate, x_ssm = z.chunk(2, dim=-1) + + # Local conv + x_ssm_t = x_ssm.transpose(1, 2) + x_ssm_t = self.conv1d(x_ssm_t)[..., :L] + x_ssm = x_ssm_t.transpose(1, 2) + x_ssm = F.silu(x_ssm) + + # Selective params + delta_bc = self.x_proj(x_ssm) + delta_raw, B_p, C_p = torch.split( + delta_bc, + [self.dt_rank, self.d_state, self.d_state], + dim=-1, + ) + delta = F.softplus(self.dt_proj(delta_raw)) + + # State transition + A = -torch.exp(self.A_log) # [C, d_state] + + # Sequential scan + h = torch.zeros(B, C, self.d_state, device=x.device, dtype=x.dtype) + ys = [] + for t in range(L): + delta_t = delta[:, t] + B_t = B_p[:, t] + C_t = C_p[:, t] + x_t = x_ssm[:, t] + + dA = torch.exp(delta_t.unsqueeze(-1) * A.unsqueeze(0)) + dB = delta_t.unsqueeze(-1) * B_t.unsqueeze(1) + + h = dA * h + dB * x_t.unsqueeze(-1) + y_t = (C_t.unsqueeze(1) * h).sum(dim=-1) + self.D * x_t + ys.append(y_t) + + y = torch.stack(ys, dim=1) + y = y * F.silu(x_gate) + y = self.out_proj(y) + return y + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Dispatch to selected backend. + + Args: + x: [B, L, C] + Returns: + y: [B, L, C] + """ + if self.backend == "mamba_ssm": + return self._forward_mamba_ssm(x) + return self._forward_torch(x) + + +# ============================================================ +# Mamba-2 Block (preferred for SOFIA v7.1 — SSD dual form, 2-8x faster) +# ============================================================ + +class Mamba2Block(nn.Module): + """Mamba-2 block using Structured State Duality (SSD). + + Differences vs Mamba-1: + - A is scalar per head (not diagonal per channel) → better compute/param balance + - Parallel matmul (SSD) via chunked scan → 2-8x speedup on CUDA + - Cleaner HiPPO-free initialization + - Head-based structure similar to MHSA + + Backend selection: + - "mamba_ssm": uses Mamba2 from mamba_ssm.modules.mamba2 (preferred) + - "torch": simplified fallback via chunked scan in PyTorch (slow) + + Reference: + Dao & Gu, "Transformers are SSMs: Generalized Models and Efficient + Algorithms Through Structured State Space Duality", ICML 2024. + + Note on parameter layout: when backend="mamba_ssm", the underlying Mamba2 + module owns all parameters. We wrap it for interface uniformity. State + dicts are NOT compatible with Mamba-1 (different parameterization of A). + """ + + def __init__( + self, + channels: int, + d_state: int = 64, # Mamba-2 uses larger N than Mamba-1 (16→64 default) + headdim: int = 64, + d_conv: int = 4, + expand: int = 2, + backend: str = "auto", + ) -> None: + super().__init__() + self.channels = channels + self.d_state = d_state + self.headdim = headdim + assert channels % headdim == 0, ( + f"channels {channels} must be divisible by headdim {headdim}" + ) + self.nheads = channels // headdim + + # Resolve backend + if backend == "auto": + backend = "mamba_ssm" if _MAMBA2_AVAILABLE else "torch" + if backend == "mamba_ssm" and not _MAMBA2_AVAILABLE: + raise RuntimeError( + "Mamba-2 (SSD) backend requested but mamba_ssm.modules.mamba2 " + "not available. Install mamba_ssm with Mamba-2 support." + ) + self.backend = backend + + if backend == "mamba_ssm": + # Delegate entirely to official Mamba2 — best perf + self.impl = _ExternalMamba2( + d_model=channels, + d_state=d_state, + d_conv=d_conv, + expand=expand, + headdim=headdim, + ) + else: + # Torch fallback: approximates SSD via sequential scan with + # scalar-per-head A. Params are owned by this nn.Module so + # checkpoints round-trip between torch instances at least. + d_inner = expand * channels + + # In projection: [x, z, dt, B, C, A_log, D are packed] + self.in_proj = nn.Linear( + channels, + 2 * d_inner + 2 * d_state + self.nheads, # x + z + B + C + dt + bias=False, + ) + self.conv1d = nn.Conv1d( + d_inner, d_inner, + kernel_size=d_conv, + padding=d_conv - 1, + groups=d_inner, + bias=True, + ) + self.d_inner = d_inner + + # Learnable A (scalar per head, log-parameterized) + A_init = torch.arange(1, self.nheads + 1, dtype=torch.float32) + self.A_log = nn.Parameter(torch.log(A_init)) # [nheads] + + # D skip + self.D = nn.Parameter(torch.ones(self.nheads)) + + # Output projection + self.out_proj = nn.Linear(d_inner, channels, bias=False) + + def _forward_torch(self, x: torch.Tensor) -> torch.Tensor: + """Simplified Mamba-2 scan in PyTorch.""" + B, L, C = x.shape + + # In projection produces packed tensor + z_in = self.in_proj(x) # [B, L, 2*d_inner + 2*d_state + nheads] + xz, BC, dt = torch.split( + z_in, + [2 * self.d_inner, 2 * self.d_state, self.nheads], + dim=-1, + ) + x_main, z_gate = xz.chunk(2, dim=-1) # each [B, L, d_inner] + B_p, C_p = BC.chunk(2, dim=-1) # each [B, L, d_state] + + # Conv1D + x_main_t = x_main.transpose(1, 2) # [B, d_inner, L] + x_main_t = self.conv1d(x_main_t)[..., :L] + x_main = F.silu(x_main_t).transpose(1, 2) # [B, L, d_inner] + + dt = F.softplus(dt) # [B, L, nheads] + A = -torch.exp(self.A_log) # [nheads] + + # Reshape x_main to heads: [B, L, nheads, headdim] + x_head = x_main.view(B, L, self.nheads, self.headdim) + + # Sequential scan per-head (like Mamba-1 but scalar A per head) + # h[t]: [B, nheads, headdim, d_state] + h = torch.zeros( + B, self.nheads, self.headdim, self.d_state, + device=x.device, dtype=x.dtype, + ) + ys = [] + for t in range(L): + dt_t = dt[:, t] # [B, nheads] + B_t = B_p[:, t] # [B, d_state] + C_t = C_p[:, t] # [B, d_state] + x_t = x_head[:, t] # [B, nheads, headdim] + + # Discretize: dA = exp(dt * A) [B, nheads], broadcast to [B, nheads, 1, 1] + dA = torch.exp(dt_t * A.unsqueeze(0)).unsqueeze(-1).unsqueeze(-1) + # dB = dt * B [B, nheads, d_state] broadcast to [B, nheads, 1, d_state] + dB = (dt_t.unsqueeze(-1) * B_t.unsqueeze(1)).unsqueeze(2) + + # State update: h = dA * h + dB * x_t + h = dA * h + dB * x_t.unsqueeze(-1) # [B, nheads, headdim, d_state] + + # Output y = (h · C) + D * x + y_t = (h * C_t.unsqueeze(1).unsqueeze(1)).sum(dim=-1) # [B, nheads, headdim] + y_t = y_t + self.D.unsqueeze(0).unsqueeze(-1) * x_t + ys.append(y_t) + + y = torch.stack(ys, dim=1) # [B, L, nheads, headdim] + y = y.reshape(B, L, self.d_inner) + y = y * F.silu(z_gate) + y = self.out_proj(y) + return y + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """x: [B, L, C] → y: [B, L, C]""" + if self.backend == "mamba_ssm": + return self.impl(x) + return self._forward_torch(x) + + +# ============================================================ +# EfficientVMamba Block (speed fallback, atrous scan) +# ============================================================ + +class EfficientVMambaBlock(nn.Module): + """EfficientVMamba with Efficient 2D Scanning (ES2D) — atrous strided scan. + + Idea: Instead of scanning all N=HW tokens in one pass (O(N) work per channel), + split the token sequence into 4 interleaved sub-sequences (stride 2 in each + of 4 patterns), scan each sub-sequence at shorter length N/4, and merge. + + For N=256 → 4×64 tokens — 4× shorter scans but 4× parallelizable, giving + ~2-3× wall-time speedup on GPU without custom CUDA kernel. + + Reference: + Pei et al., "EfficientVMamba: Atrous Selective Scan for Light Weight + Visual Mamba", 2024 (arXiv:2403.09977). + + This implementation is a simplified 2-direction variant (stride 1 + stride 2) + rather than full 4-direction cross-atrous. Quality ~98% of original per paper. + """ + + def __init__( + self, + channels: int, + d_state: int = 16, + dt_rank: Optional[int] = None, + conv_kernel: int = 3, + n_directions: int = 2, + ) -> None: + super().__init__() + assert n_directions in (2, 4), "n_directions must be 2 or 4" + self.channels = channels + self.n_directions = n_directions + + # Use multiple SimpleMambaBlocks for each direction + # They share config but not weights — each scan direction needs own params + self.scanners = nn.ModuleList([ + SimpleMambaBlock( + channels=channels, + d_state=d_state, + dt_rank=dt_rank, + conv_kernel=conv_kernel, + backend="auto", + ) + for _ in range(n_directions) + ]) + + # Merge gate — learnable weighted sum of directions + self.merge_weights = nn.Parameter(torch.zeros(n_directions)) + + @staticmethod + def _atrous_reorder(x: torch.Tensor, H: int, W: int, direction: int) -> torch.Tensor: + """Reorder tokens for atrous scan. + + Args: + x: [B, L, C] with L = H*W + H, W: spatial dims + direction: 0=identity, 1=stride-2 interleave, 2=vertical, 3=vertical stride-2 + + Returns: + [B, L, C] reordered + """ + B, L, C = x.shape + x_2d = x.reshape(B, H, W, C) + + if direction == 0: + out = x_2d.reshape(B, L, C) + elif direction == 1: + # Stride-2 row interleave: even cols first, then odd cols + even = x_2d[:, :, 0::2, :] # [B, H, W/2, C] + odd = x_2d[:, :, 1::2, :] + out = torch.cat([even, odd], dim=2).reshape(B, L, C) + elif direction == 2: + # Column-major (transpose) + out = x_2d.permute(0, 2, 1, 3).reshape(B, L, C) + elif direction == 3: + # Transposed + stride 2 + xt = x_2d.permute(0, 2, 1, 3) # [B, W, H, C] + even = xt[:, :, 0::2, :] + odd = xt[:, :, 1::2, :] + out = torch.cat([even, odd], dim=2).reshape(B, L, C) + else: + raise ValueError(f"direction must be 0-3, got {direction}") + + return out + + @staticmethod + def _atrous_unreorder(x: torch.Tensor, H: int, W: int, direction: int) -> torch.Tensor: + """Inverse of _atrous_reorder.""" + B, L, C = x.shape + + if direction == 0: + return x + elif direction == 1: + half = W // 2 + x_2d = x.reshape(B, H, W, C) + even, odd = x_2d[:, :, :half, :], x_2d[:, :, half:, :] + out = torch.zeros_like(x_2d) + out[:, :, 0::2, :] = even + out[:, :, 1::2, :] = odd + return out.reshape(B, L, C) + elif direction == 2: + x_2d = x.reshape(B, W, H, C) + return x_2d.permute(0, 2, 1, 3).reshape(B, L, C) + elif direction == 3: + half = H // 2 + x_2d = x.reshape(B, W, H, C) + even, odd = x_2d[:, :, :half, :], x_2d[:, :, half:, :] + out = torch.zeros_like(x_2d) + out[:, :, 0::2, :] = even + out[:, :, 1::2, :] = odd + return out.permute(0, 2, 1, 3).reshape(B, L, C) + + def forward(self, x: torch.Tensor, H: int = None, W: int = None) -> torch.Tensor: + """x: [B, L, C] with L = H*W → y: [B, L, C] + + H, W required so we know the 2D layout for atrous reordering. + """ + if H is None or W is None: + # Assume square + import math + L = x.shape[1] + side = int(math.isqrt(L)) + assert side * side == L, ( + f"EfficientVMamba needs square feature map or explicit H, W; got L={L}" + ) + H = W = side + + outs = [] + for d, scanner in enumerate(self.scanners): + x_reord = self._atrous_reorder(x, H, W, direction=d) + y_reord = scanner(x_reord) + y = self._atrous_unreorder(y_reord, H, W, direction=d) + outs.append(y) + + # Weighted merge via softmax + weights = F.softmax(self.merge_weights, dim=0) + out = sum(w * o for w, o in zip(weights, outs)) + return out + + +# ============================================================ +# Factory for Mamba variants +# ============================================================ + +def build_mamba_block( + variant: str, + channels: int, + d_state: int = 16, + dt_rank: Optional[int] = None, + backend: str = "auto", + **extra_kwargs, +) -> nn.Module: + """Factory to build appropriate Mamba variant for SOFIA. + + Args: + variant: one of + - "mamba1": original Mamba-1 selective scan (SimpleMambaBlock) + - "mamba2": Mamba-2 with SSD dual form (preferred, 2-8x faster) + - "efficient_vmamba": atrous scan (speed fallback, ~2-3x speedup) + channels: model channel dim + d_state: state dimension N + dt_rank: Δ predictor rank (auto if None) + backend: "auto" | "torch" | "mamba_ssm" + extra_kwargs: passed to variant-specific block + + Returns: + nn.Module with forward(x: [B, L, C]) -> [B, L, C] + """ + if variant == "mamba1": + return SimpleMambaBlock( + channels=channels, + d_state=d_state, + dt_rank=dt_rank, + backend=backend, + ) + elif variant == "mamba2": + return Mamba2Block( + channels=channels, + d_state=extra_kwargs.get("d_state_mamba2", 64), # Mamba-2 typically uses N=64 + headdim=extra_kwargs.get("headdim", 64), + d_conv=extra_kwargs.get("d_conv", 4), + expand=extra_kwargs.get("expand", 2), + backend=backend, + ) + elif variant == "efficient_vmamba": + return EfficientVMambaBlock( + channels=channels, + d_state=d_state, + dt_rank=dt_rank, + n_directions=extra_kwargs.get("n_directions", 2), + ) + else: + raise ValueError( + f"Unknown Mamba variant '{variant}'. " + f"Use 'mamba1', 'mamba2', or 'efficient_vmamba'." + ) + + +# ============================================================ +# MambaVision Block (Stages 3-4) +# ============================================================ + +class MambaVisionBlock(nn.Module): + """MV5/MV1: (Mamba ∥ MHSA [∥ Strip]) sum + FFN. + + When use_strip_branch=True → MV5 (novel 3-way parallel). + When use_strip_branch=False → MV1 (original MambaVision 2-way). + + MHSA branch uses RoPE 2D positional encoding. + """ + + def __init__( + self, + channels: int, + num_heads: int = 8, + d_state: int = 16, + dt_rank: Optional[int] = None, + use_strip_branch: bool = True, + ffn_expand: int = 4, + strip_kernel: int = 7, + use_rope: bool = True, + mamba_backend: str = "auto", + mamba_variant: str = "mamba2", + mamba_extra_kwargs: Optional[dict] = None, + ) -> None: + super().__init__() + self.channels = channels + self.num_heads = num_heads + self.head_dim = channels // num_heads + assert channels % num_heads == 0, ( + f"channels {channels} must be divisible by num_heads {num_heads}" + ) + assert self.head_dim % 4 == 0, ( + f"head_dim {self.head_dim} must be divisible by 4 for RoPE 2D" + ) + + self.use_strip_branch = use_strip_branch + self.use_rope = use_rope + self.mamba_variant = mamba_variant + + # Input LayerNorm (channel-last, applied via [B, L, C]) + self.norm1 = nn.LayerNorm(channels) + + # Mamba branch — factory selects variant + self.mamba = build_mamba_block( + variant=mamba_variant, + channels=channels, + d_state=d_state, + dt_rank=dt_rank, + backend=mamba_backend, + **(mamba_extra_kwargs or {}), + ) + + # MHSA branch + self.qkv = nn.Linear(channels, 3 * channels, bias=False) + self.attn_out = nn.Linear(channels, channels, bias=False) + if use_rope: + self.rope = RoPE2D(self.head_dim) + + # Strip branch (optional, MV5) + if use_strip_branch: + pad = strip_kernel // 2 + self.strip_h = nn.Conv2d( + channels, channels, + kernel_size=(1, strip_kernel), + stride=1, + padding=(0, pad), + groups=channels, + bias=False, + ) + self.strip_v = nn.Conv2d( + channels, channels, + kernel_size=(strip_kernel, 1), + stride=1, + padding=(pad, 0), + groups=channels, + bias=False, + ) + self.strip_proj = nn.Linear(channels, channels, bias=False) + + # Learnable merge weights (M2: weighted sum via softmax) + n_branches = 3 if use_strip_branch else 2 + self.merge_weights = nn.Parameter(torch.zeros(n_branches)) + + # FFN + self.norm2 = nn.LayerNorm(channels) + ffn_hidden = ffn_expand * channels + self.ffn = nn.Sequential( + nn.Linear(channels, ffn_hidden), + nn.GELU(), + nn.Linear(ffn_hidden, channels), + ) + + def _mhsa(self, x: torch.Tensor, H: int, W: int) -> torch.Tensor: + """ + Args: + x: [B, L, C] + H, W: spatial dims + Returns: + [B, L, C] + """ + B, L, C = x.shape + qkv = self.qkv(x) # [B, L, 3C] + qkv = qkv.reshape(B, L, 3, self.num_heads, self.head_dim) + qkv = qkv.permute(2, 0, 3, 1, 4) # [3, B, heads, L, head_dim] + q, k, v = qkv.unbind(0) + + if self.use_rope: + q, k = self.rope(q, k, H, W) + + # SDPA + attn_out = F.scaled_dot_product_attention(q, k, v) # [B, heads, L, head_dim] + attn_out = attn_out.transpose(1, 2).reshape(B, L, C) + attn_out = self.attn_out(attn_out) + return attn_out + + def _strip(self, x_2d: torch.Tensor) -> torch.Tensor: + """ + Args: + x_2d: [B, C, H, W] + Returns: + [B, C, H, W] + """ + s_h = self.strip_h(x_2d) + s_v = self.strip_v(x_2d) + out = s_h + s_v + # Project via Linear: need [B, L, C] + B, C, H, W = out.shape + out_seq = out.permute(0, 2, 3, 1).reshape(B, H * W, C) + out_seq = self.strip_proj(out_seq) + return out_seq # [B, L, C] + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x: [B, C, H, W] + Returns: + [B, C, H, W] + """ + B, C, H, W = x.shape + L = H * W + identity = x + + # Convert to [B, L, C] + x_seq = x.permute(0, 2, 3, 1).reshape(B, L, C) + x_norm = self.norm1(x_seq) + + # Mamba branch — EfficientVMamba needs H, W for atrous reorder + if self.mamba_variant == "efficient_vmamba": + mamba_out = self.mamba(x_norm, H=H, W=W) + else: + mamba_out = self.mamba(x_norm) # [B, L, C] + + # MHSA branch + mhsa_out = self._mhsa(x_norm, H, W) # [B, L, C] + + # Strip branch (if enabled) + branch_outs = [mamba_out, mhsa_out] + if self.use_strip_branch: + x_2d_norm = x_norm.reshape(B, H, W, C).permute(0, 3, 1, 2) + strip_out = self._strip(x_2d_norm) + branch_outs.append(strip_out) + + # Weighted sum (softmax merge weights) + weights = F.softmax(self.merge_weights, dim=0) + merged = sum(w * b for w, b in zip(weights, branch_outs)) + + # Residual + x_seq = x_seq + merged + + # FFN + x_seq = x_seq + self.ffn(self.norm2(x_seq)) + + # Back to [B, C, H, W] + x = x_seq.reshape(B, H, W, C).permute(0, 3, 1, 2).contiguous() + return x + + +# ============================================================ +# EVSS Bridge (B6-inspired refinement between heterogeneous stages) +# ============================================================ + +class EVSSBridge(nn.Module): + """Within-resolution dual-path refinement bridge (EVSS-style). + + Inspired by EfficientVMamba (B6, AAAI 2025) where the EVSS block + parallel-fuses a local depthwise-conv path with a global SSM-scan path + via SE channel-attention gating. Applied here NOT as a stage block but + as a thin refinement module bridging heterogeneous backbone stages + (e.g. SOFIA DCN stages 1-2 → MambaVision stages 3-4) to smooth the + "semantic bottleneck" caused by the abrupt receptive-field shift. + + Architecture (input/output share shape [B, C, H, W]): + + x ─┬── [DW 3x3 + BN + SiLU + PW 1x1 + SE] → local + └── [LN2d → SimpleMambaBlock → SE] → global + y = local + global + + The Mamba branch uses our existing SimpleMambaBlock (Mamba-1 by default; + can be swapped via `mamba_variant` argument). The SE blocks are tiny + and act as soft channel-wise gates between the two paths. + + Cost at C=192, HW=256: ~165K params, ~30 MMAC — under 1% of SOFIA-M + backbone. Designed to be opt-in via cfg.use_evss_bridge. + """ + + def __init__( + self, + channels: int, + mamba_d_state: int = 16, + mamba_dt_rank: Optional[int] = None, + mamba_variant: str = "mamba1", + mamba_backend: str = "auto", + se_reduction: int = 16, + ) -> None: + super().__init__() + self.channels = channels + + # Local path: DW 3x3 + BN + SiLU + PW 1x1 + SE + self.local_dw = nn.Conv2d( + channels, channels, kernel_size=3, padding=1, + groups=channels, bias=False, + ) + self.local_bn = nn.BatchNorm2d(channels) + self.local_pw = nn.Conv2d(channels, channels, kernel_size=1, bias=False) + self.local_pw_bn = nn.BatchNorm2d(channels) + self.local_se = SqueezeExcite(channels, reduction=se_reduction) + + # Global path: LayerNorm2d + Mamba scan + SE + self.global_norm = LayerNorm2d(channels) + # Use lightweight Mamba-1 by default for the bridge — full Mamba-2 SSD + # is overkill here since this is a refinement module, not a primary + # token mixer. User can switch via mamba_variant if desired. + self.global_mamba = build_mamba_block( + variant=mamba_variant, + channels=channels, + d_state=mamba_d_state, + dt_rank=mamba_dt_rank, + backend=mamba_backend, + ) + self.global_se = SqueezeExcite(channels, reduction=se_reduction) + + self.act = nn.SiLU(inplace=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """x: [B, C, H, W] → [B, C, H, W]""" + # Local path + local = self.local_dw(x) + local = self.local_bn(local) + local = self.act(local) + local = self.local_pw(local) + local = self.local_pw_bn(local) + local = self.local_se(local) + + # Global path: 2D → seq → Mamba → seq → 2D + B, C, H, W = x.shape + x_norm = self.global_norm(x) + x_seq = x_norm.flatten(2).transpose(1, 2) # [B, HW, C] + global_out = self.global_mamba(x_seq) # [B, HW, C] + global_out = global_out.transpose(1, 2).reshape(B, C, H, W) + global_out = self.global_se(global_out) + + return local + global_out + + +# ============================================================ +# Downsample layer +# ============================================================ + +class Downsample(nn.Module): + """3×3 stride-2 conv + BN.""" + + def __init__(self, in_channels: int, out_channels: int) -> None: + super().__init__() + self.conv = nn.Conv2d( + in_channels, out_channels, kernel_size=3, stride=2, padding=1, bias=False + ) + self.bn = nn.BatchNorm2d(out_channels) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.bn(self.conv(x)) + + +if __name__ == "__main__": + import time + torch.manual_seed(0) + + # StripDCNLite smoke test + blk = StripDCNLiteBlock(48, 48, kernel=7, stride=1, use_dcn=True) + x = torch.randn(2, 48, 64, 64) + out = blk(x) + print(f"StripDCNLite: {out.shape}, params: {sum(p.numel() for p in blk.parameters()) / 1e3:.1f} K") + + # StripMixConv smoke test + blk = StripMixConvBlock(96, 96, strip_kernel=5) + x = torch.randn(2, 96, 32, 32) + out = blk(x) + print(f"StripMixConv: {out.shape}, params: {sum(p.numel() for p in blk.parameters()) / 1e3:.1f} K") + + # MambaVision MV5 smoke test (small for speed) + blk = MambaVisionBlock(channels=128, num_heads=4, d_state=16, use_strip_branch=True) + x = torch.randn(1, 128, 16, 16) + t0 = time.time() + out = blk(x) + print(f"MV5 block: {out.shape}, params: {sum(p.numel() for p in blk.parameters()) / 1e3:.1f} K, time: {time.time() - t0:.3f}s") + + # MV1 (no strip) + blk = MambaVisionBlock(channels=128, num_heads=4, d_state=16, use_strip_branch=False) + x = torch.randn(1, 128, 8, 8) + out = blk(x) + print(f"MV1 block: {out.shape}, params: {sum(p.numel() for p in blk.parameters()) / 1e3:.1f} K") diff --git a/src/models/sofia_v71/config.py b/src/models/sofia_v71/config.py new file mode 100644 index 0000000..ecd2448 --- /dev/null +++ b/src/models/sofia_v71/config.py @@ -0,0 +1,210 @@ +"""SOFIA v7.1 configuration system. + +Two scale presets targeting Jetson Orin NX INT8 deployment: + +| Preset | Params | INT8 size | FP16 size | Target latency | +|----------|---------:|----------:|----------:|---------------:| +| SOFIA-M | ~500 M | ~500 MB | ~1 GB | ~18 ms | +| SOFIA-L | ~1 B | ~1 GB | ~2 GB | ~20 ms | + +Based on HYP_SOFIA_v7_UltraDeep_дизайн.md (Phases 2'–5' + Phase 4'' +CVGL-Aware Head + revision to ultra-lite 1x1 neck). +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import List, Literal, Optional + + +@dataclass +class SOFIAConfig: + """Configuration for SOFIA v7.1 student architecture. + + Controls backbone width/depth, neck, and CVGL-Aware Head. + """ + + # -------- Input -------- + input_size: int = 256 + in_channels: int = 3 + + # -------- Stem (dual-conv FastViT-style) -------- + stem_mid: int = 64 # intermediate channels + stem_out: int = 128 # output channels into stage 1 + + # -------- Backbone dimensions per stage s1..s4 -------- + embed_dims: List[int] = field( + default_factory=lambda: [256, 512, 1280, 1536] + ) + depths: List[int] = field(default_factory=lambda: [3, 4, 15, 3]) + + # -------- Stage 1–2 block parameters -------- + mbconv_expand: int = 4 + se_ratio: int = 16 + strip_kernel_s1: int = 7 # strip DW kernel size stage 1 + strip_kernel_s2: int = 5 # strip DW kernel size stage 2 + mix_kernels: List[int] = field( + default_factory=lambda: [3, 5, 7] + ) # MixConv DW kernel sizes stage 2 + use_dcn_strip: bool = True # adaptive offset on horizontal strip + + # -------- Stage 3–4 (MambaVision) -------- + mamba_d_state: int = 16 + mamba_dt_rank: Optional[int] = None # auto = max(1, C // 16) + mamba_backend: Literal["auto", "torch", "mamba_ssm"] = "auto" + # "auto" uses mamba_ssm if importable, else torch fallback + + # Mamba variant — one of: + # "mamba2" (preferred, SSD dual form, 2-8x faster) + # "mamba1" (original selective scan, mature) + # "efficient_vmamba" (speed fallback, atrous scan ~2-3x speedup) + mamba_variant: Literal["mamba1", "mamba2", "efficient_vmamba"] = "mamba2" + # Per-variant tunables passed through to factory + mamba_extra_kwargs: dict = field(default_factory=lambda: { + "d_state_mamba2": 64, # Mamba-2 typically uses N=64 (not 16) + "headdim": 64, # Mamba-2 head dim + "expand": 2, # Mamba-2 inner expansion factor + "d_conv": 4, # Mamba-2 local conv kernel + "n_directions": 2, # EfficientVMamba: 2 or 4 atrous directions + }) + + num_heads_s3: int = 8 + num_heads_s4: int = 8 + use_strip_branch_s3: bool = True # MV5 with Strip branch + use_strip_branch_s4: bool = False # MV1 without Strip branch + ffn_expand: int = 4 + + # -------- EVSS-style bridge (B6-inspired, opt-in) -------- + # When True, inserts a within-resolution dual-path refinement block + # right after each downsample to stage 3 and stage 4. Smooths semantic + # gap between heterogeneous stages (DCN → MambaVision). Adds ~165K params + # and ~30 MMAC per insertion at C=192, HW=256. + use_evss_bridge: bool = False + evss_bridge_locations: List[str] = field( + default_factory=lambda: ["pre_stage3"] # subset of {pre_stage3, pre_stage4} + ) + + # -------- Neck (ultra-lite 1x1 projection) -------- + neck_channels: int = 192 # C_n, output channels from neck + + # -------- CVGL-Aware Head v7.1-α -------- + d_descriptor: int = 512 # global descriptor dimensionality + use_asymmetric_heads: bool = True # Sat vs UAV different heads + chp_rings: int = 8 + chp_angles: int = 16 + chp_harmonics: int = 4 + use_film_altitude: bool = True + altitude_norm: float = 500.0 # divides altitude in meters + ring_count: int = 4 # LPN rings auxiliary + use_ring_aux: bool = True # training-only ring aux branch + + # -------- Text fusion (extension for caption-conditioned heads) -------- + return_normalized: bool = True # if False, heads return pre-L2 features (for late gated fusion) + use_text_film_sat: bool = False # text-FiLM modulation in SatHead before GGeM + use_text_film_uav: bool = False # text-FiLM modulation in UAVHead alongside altitude FiLM + text_film_dim: int = 1024 # text embedding dim feeding FiLM (matches TextFusionMLP out_dim) + text_film_hidden: int = 256 + + # -------- Weight-sharing -------- + share_stages_1_2: bool = True # sat ↔ UAV shared weights stages 1-2 + + # -------- KD taps (enable for future teacher KD) -------- + enable_kd_taps: bool = True + + # -------- Deployment hints -------- + precision: Literal["fp32", "fp16", "int8_mixed"] = "fp16" + + def validate(self) -> None: + assert len(self.embed_dims) == 4, ( + f"embed_dims must have 4 entries, got {len(self.embed_dims)}" + ) + assert len(self.depths) == 4, ( + f"depths must have 4 entries, got {len(self.depths)}" + ) + assert self.input_size % 32 == 0, ( + f"input_size must be divisible by 32 (4 downsamples × stem), " + f"got {self.input_size}" + ) + + def summary(self) -> str: + return ( + f"SOFIAConfig(stem={self.stem_mid}/{self.stem_out}, " + f"dims={self.embed_dims}, depths={self.depths}, " + f"neck={self.neck_channels}, d={self.d_descriptor}, " + f"precision={self.precision})" + ) + + +# ============================================================ +# Scale Presets +# ============================================================ + +def sofia_m_config() -> SOFIAConfig: + """SOFIA-M: ~500 M params target (~500 MB INT8, ~1 GB FP16). + + Fits in 500 MB VRAM after INT8 quantization on Jetson Orin NX. + Expected latency: ~18 ms. + """ + return SOFIAConfig( + stem_mid=64, + stem_out=128, + embed_dims=[256, 512, 1280, 1536], + depths=[3, 4, 15, 3], + neck_channels=192, + d_descriptor=512, + ) + + +def sofia_l_config() -> SOFIAConfig: + """SOFIA-L: ~1 B params target (~1 GB INT8, ~2 GB FP16). + + Fits in 1 GB VRAM after INT8 quantization on Jetson Orin NX. + Expected latency: ~20 ms. + """ + return SOFIAConfig( + stem_mid=64, + stem_out=128, + embed_dims=[256, 512, 1536, 2048], + depths=[3, 4, 20, 3], + neck_channels=256, + d_descriptor=1024, + ) + + +def sofia_tiny_config() -> SOFIAConfig: + """SOFIA-Tiny: ~5 M params (matches original v7.1 spec). + + Reference for research comparisons. Not optimized for 500 MB INT8 target. + `num_heads_*` is set to 4 so `head_dim` (channels // heads) is divisible + by 4 — required by RoPE 2D in `MambaVisionBlock` (s3: 176/4=44, s4: 224/4=56). + `mamba_extra_kwargs.headdim=16` because Mamba-2 requires channels % headdim == 0; + 176 and 224 are not divisible by the default 64. + """ + return SOFIAConfig( + stem_mid=16, + stem_out=32, + embed_dims=[48, 96, 176, 224], + depths=[2, 3, 4, 2], + num_heads_s3=4, + num_heads_s4=4, + neck_channels=128, + d_descriptor=512, + mamba_extra_kwargs={ + "d_state_mamba2": 64, + "headdim": 16, + "expand": 2, + "d_conv": 4, + "n_directions": 2, + }, + ) + + +# Default preset +DEFAULT_CONFIG = sofia_m_config + + +if __name__ == "__main__": + for name, fn in [("M", sofia_m_config), ("L", sofia_l_config), ("Tiny", sofia_tiny_config)]: + cfg = fn() + cfg.validate() + print(f"SOFIA-{name}: {cfg.summary()}") diff --git a/src/models/sofia_v71/layers.py b/src/models/sofia_v71/layers.py new file mode 100644 index 0000000..6d07f0f --- /dev/null +++ b/src/models/sofia_v71/layers.py @@ -0,0 +1,398 @@ +"""SOFIA v7.1 custom layers. + +Includes: +- GGeM: Generalized Mean Pooling with per-channel learnable exponent (F11) +- CircularHarmonicPool: Formally SO(2)-invariant pooling via polar + FFT magnitude (NH2 novel) +- AltitudeFiLM: FiLM conditioning on UAV altitude (NH4 novel) +- RoPE2D: 2D Rotary Position Embedding for attention +- SqueezeExcite: standard SE block +- LayerNorm2d: channel-last LN wrapper for 2D features + +All rotation-invariance and FiLM modules are NOVEL contributions of SOFIA v7.1-α. +""" + +from __future__ import annotations + +import math +from typing import Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + + +# ============================================================ +# GGeM: Generalized Mean Pooling (F11) +# ============================================================ + +class GGeM(nn.Module): + """Per-channel learnable Generalized Mean pooling. + + Formula: + GGeM(F)_c = (1/HW · Σ F_{c,h,w}^{p_c})^{1/p_c} + p_c = softplus(p_hat_c) ∈ (0, ∞) + """ + + def __init__(self, channels: int, init_p: float = 3.0, eps: float = 1e-6) -> None: + super().__init__() + # softplus^{-1}(init_p) = log(exp(init_p) - 1) + hat_init = math.log(math.exp(init_p) - 1.0) + self.hat_p = nn.Parameter(torch.full((channels,), hat_init)) + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """x: [B, C, H, W] -> [B, C]""" + p = F.softplus(self.hat_p).view(1, -1, 1, 1) # [1, C, 1, 1] + x_clamped = x.clamp(min=self.eps) + x_pow = x_clamped.pow(p) + x_mean = x_pow.mean(dim=(2, 3), keepdim=True) + out = x_mean.pow(1.0 / p) + return out.flatten(1) + + +# ============================================================ +# CHP: Circular Harmonic Pool (NH2 novel — formally SO(2)-invariant) +# ============================================================ + +class CircularHarmonicPool(nn.Module): + """Formally SO(2) rotation-invariant pooling. + + Algorithm: + 1. Sample input feature map at polar grid (r, θ) via bilinear grid_sample + 2. Apply 1D real FFT along θ-axis + 3. Keep magnitudes of first N harmonics (invariant to shift = rotation) + 4. GGeM pool over rings r (per-channel-per-harmonic) + 5. Flatten to descriptor [B, C * N] + + Output is theoretically invariant to input rotation of any angle. + + See HYP Phase 4'' Section 4''.1 NH2 and Section 4''.5 for formal proof. + """ + + def __init__( + self, + channels: int, + rings: int = 8, + angles: int = 16, + harmonics: int = 4, + r_min: float = 0.1, + r_max: float = 1.0, + ) -> None: + super().__init__() + assert harmonics <= angles // 2 + 1, ( + f"harmonics {harmonics} cannot exceed angles//2+1 = {angles // 2 + 1}" + ) + + self.channels = channels + self.rings = rings + self.angles = angles + self.harmonics = harmonics + + # GGeM over rings (per channel × per harmonic) + self.ggem = GGeM(channels * harmonics, init_p=3.0) + + # Precompute polar grid in normalized [-1, 1] coords for grid_sample + grid = self._make_polar_grid(rings, angles, r_min, r_max) # [R, T, 2] + self.register_buffer("polar_grid", grid, persistent=False) + + @staticmethod + def _make_polar_grid(R: int, T: int, r_min: float, r_max: float) -> torch.Tensor: + r_values = torch.linspace(r_min, r_max, R) # [R] + theta_values = torch.linspace(0.0, 2 * math.pi, T + 1)[:-1] # [T] + r_grid, theta_grid = torch.meshgrid(r_values, theta_values, indexing="ij") # [R, T] + x = r_grid * torch.cos(theta_grid) + y = r_grid * torch.sin(theta_grid) + grid = torch.stack([x, y], dim=-1) # [R, T, 2] + return grid + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Args: + x: [B, C, H, W] feature map + Returns: + descriptor: [B, C * harmonics] + """ + B, C, H, W = x.shape + assert C == self.channels, ( + f"Input channels {C} != expected {self.channels}" + ) + + # 1. Polar sampling + # grid: [B, R, T, 2] + grid = self.polar_grid.unsqueeze(0).expand(B, -1, -1, -1) + polar = F.grid_sample( + x, grid, + mode="bilinear", + padding_mode="zeros", + align_corners=True, + ) # [B, C, R, T] + + # 2. 1D real FFT along angular axis + polar_fft = torch.fft.rfft(polar, dim=-1) # [B, C, R, T//2+1] + polar_fft = polar_fft[..., : self.harmonics] # [B, C, R, N] + + # 3. Magnitude (rotation invariant) + magnitude = polar_fft.abs() # [B, C, R, N] + + # 4. Reshape for GGeM: treat (C, N) as combined channel dim, rings as spatial + # Shape: [B, C*N, R, 1] + magnitude_reshaped = ( + magnitude + .permute(0, 1, 3, 2) # [B, C, N, R] + .reshape(B, C * self.harmonics, self.rings, 1) + ) + + # 5. GGeM pool over rings (H=R, W=1) + descriptor = self.ggem(magnitude_reshaped) # [B, C*N] + + return descriptor + + +# ============================================================ +# FiLM: Altitude-conditioned modulation (NH4 novel) +# ============================================================ + +class AltitudeFiLM(nn.Module): + """FiLM modulation conditioned on scalar altitude. + + F' = γ(h) · F + β(h) where γ,β ∈ R^C are produced by MLP. + + At altitude=None, produces identity (γ=1, β=0) via zero-init of final layer. + """ + + def __init__( + self, + channels: int, + hidden_dim: int = 64, + altitude_norm: float = 500.0, + ) -> None: + super().__init__() + self.channels = channels + self.altitude_norm = altitude_norm + + self.mlp = nn.Sequential( + nn.Linear(1, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, 2 * channels), + ) + # Zero-init last layer → initial γ=0 before residual, β=0 + nn.init.zeros_(self.mlp[-1].weight) + nn.init.zeros_(self.mlp[-1].bias) + + def forward(self, x: torch.Tensor, altitude: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Args: + x: [B, C, H, W] + altitude: [B] or [B, 1] scalar altitude in meters (or None for neutral) + Returns: + [B, C, H, W] + """ + B = x.shape[0] + if altitude is None: + altitude = torch.zeros(B, 1, device=x.device, dtype=x.dtype) + elif altitude.dim() == 1: + altitude = altitude.unsqueeze(-1) + + h_norm = altitude.to(x.dtype) / self.altitude_norm + gamma_beta = self.mlp(h_norm) # [B, 2C] + gamma, beta = gamma_beta.chunk(2, dim=-1) + # Residual form: γ = 1 + delta_γ (starts at identity) + gamma = gamma.view(B, self.channels, 1, 1) + 1.0 + beta = beta.view(B, self.channels, 1, 1) + return gamma * x + beta + + +# ============================================================ +# TextFiLM: text-conditioned modulation (extension for caption fusion) +# ============================================================ + +class TextFiLM(nn.Module): + """FiLM modulation conditioned on a text embedding. + + F' = γ(z) · F + β(z) where γ,β ∈ R^C are produced by an MLP + from z ∈ R^{D_txt}. Identity at init via zero-init of last layer + (so γ=1, β=0 before training shifts the residual). + """ + + def __init__( + self, + channels: int, + text_dim: int = 1024, + hidden_dim: int = 256, + ) -> None: + super().__init__() + self.channels = channels + self.text_dim = text_dim + self.mlp = nn.Sequential( + nn.Linear(text_dim, hidden_dim), + nn.GELU(), + nn.Linear(hidden_dim, 2 * channels), + ) + nn.init.zeros_(self.mlp[-1].weight) + nn.init.zeros_(self.mlp[-1].bias) + + def forward(self, x: torch.Tensor, text_emb: Optional[torch.Tensor] = None) -> torch.Tensor: + """ + Args: + x: [B, C, H, W] + text_emb: [B, D_txt] or None for no-op (identity) + Returns: + [B, C, H, W] + """ + if text_emb is None: + return x + B = x.shape[0] + gamma_beta = self.mlp(text_emb.to(x.dtype)) # [B, 2C] + gamma, beta = gamma_beta.chunk(2, dim=-1) + gamma = gamma.view(B, self.channels, 1, 1) + 1.0 + beta = beta.view(B, self.channels, 1, 1) + return gamma * x + beta + + +# ============================================================ +# RoPE 2D +# ============================================================ + +class RoPE2D(nn.Module): + """2D Rotary Position Embedding. + + Splits head_dim into two halves: first half gets x-position encoding, + second half gets y-position encoding. For each half, applies standard + 1D RoPE rotation. + + Reference: RoFormer (B49) adapted for 2D. + """ + + def __init__(self, head_dim: int, max_resolution: int = 64, base: float = 10000.0) -> None: + super().__init__() + assert head_dim % 2 == 0, "head_dim must be even for RoPE" + assert head_dim % 4 == 0, "head_dim must be divisible by 4 for 2D RoPE" + self.head_dim = head_dim + self.half_dim = head_dim // 2 # dedicated to each axis + self.max_resolution = max_resolution + + # Frequencies for each axis (half_dim per axis, sin+cos pairs) + freqs = 1.0 / (base ** (torch.arange(0, self.half_dim, 2).float() / self.half_dim)) + self.register_buffer("freqs", freqs, persistent=False) + + def _make_embeds(self, H: int, W: int, device: torch.device) -> Tuple[torch.Tensor, torch.Tensor]: + """Produce cos/sin embeddings for HW tokens in raster order.""" + y_pos = torch.arange(H, device=device, dtype=torch.float32) + x_pos = torch.arange(W, device=device, dtype=torch.float32) + + freqs_y = torch.einsum("i,j->ij", y_pos, self.freqs) # [H, half_dim/2] + freqs_x = torch.einsum("i,j->ij", x_pos, self.freqs) # [W, half_dim/2] + + # Expand to full grid: [H, W, half_dim/2] each + freqs_y = freqs_y.unsqueeze(1).expand(-1, W, -1) # [H, W, half_dim/2] + freqs_x = freqs_x.unsqueeze(0).expand(H, -1, -1) # [H, W, half_dim/2] + + # Concatenate: x-axis freqs into first half, y-axis into second half + # Each half is [H, W, half_dim/2]; we pair-up for complex rotation + freqs_combined_x = torch.cat([freqs_x, freqs_x], dim=-1) # [H, W, half_dim] + freqs_combined_y = torch.cat([freqs_y, freqs_y], dim=-1) # [H, W, half_dim] + + freqs_full = torch.cat([freqs_combined_x, freqs_combined_y], dim=-1) # [H, W, head_dim] + cos = freqs_full.cos().reshape(H * W, -1) + sin = freqs_full.sin().reshape(H * W, -1) + return cos, sin + + @staticmethod + def _rotate_half(x: torch.Tensor) -> torch.Tensor: + """Rotate: (x1, x2) -> (-x2, x1).""" + x1, x2 = x.chunk(2, dim=-1) + return torch.cat([-x2, x1], dim=-1) + + def forward(self, q: torch.Tensor, k: torch.Tensor, H: int, W: int) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Args: + q, k: [B, heads, HW, head_dim] + H, W: spatial dims for position computation + Returns: + q_rot, k_rot with positional encoding applied + """ + cos, sin = self._make_embeds(H, W, q.device) + cos = cos.to(q.dtype).unsqueeze(0).unsqueeze(0) # [1, 1, HW, head_dim] + sin = sin.to(q.dtype).unsqueeze(0).unsqueeze(0) + q_rot = (q * cos) + (self._rotate_half(q) * sin) + k_rot = (k * cos) + (self._rotate_half(k) * sin) + return q_rot, k_rot + + +# ============================================================ +# SE: Squeeze-Excite +# ============================================================ + +class SqueezeExcite(nn.Module): + """Standard Squeeze-Excite channel attention.""" + + def __init__(self, channels: int, reduction: int = 16) -> None: + super().__init__() + hidden = max(1, channels // reduction) + self.pool = nn.AdaptiveAvgPool2d(1) + self.fc = nn.Sequential( + nn.Conv2d(channels, hidden, 1), + nn.SiLU(inplace=True), + nn.Conv2d(hidden, channels, 1), + nn.Sigmoid(), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + s = self.pool(x) + s = self.fc(s) + return x * s + + +# ============================================================ +# LayerNorm2d: LN over channels for (B, C, H, W) layout +# ============================================================ + +class LayerNorm2d(nn.Module): + """LayerNorm over C dimension for 4D tensors.""" + + def __init__(self, channels: int, eps: float = 1e-6) -> None: + super().__init__() + self.norm = nn.LayerNorm(channels, eps=eps) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # [B, C, H, W] → [B, H, W, C] → LN → [B, C, H, W] + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = x.permute(0, 3, 1, 2).contiguous() + return x + + +if __name__ == "__main__": + # Smoke test + torch.manual_seed(0) + + # GGeM test + g = GGeM(64) + x = torch.randn(2, 64, 8, 8) + out = g(x) + print(f"GGeM out: {out.shape}") # [2, 64] + + # CHP test — verify rotation invariance + chp = CircularHarmonicPool(32, rings=8, angles=16, harmonics=4) + x = torch.randn(1, 32, 16, 16) + out1 = chp(x) + # Rotate x by 90° and verify invariance (approximately) + x_rot = torch.rot90(x, k=1, dims=(-2, -1)) + out2 = chp(x_rot) + diff = (out1 - out2).abs().max().item() + print(f"CHP rotation-invariance max diff: {diff:.4e} (should be small)") + print(f"CHP out shape: {out1.shape}") # [1, 128] + + # FiLM test + film = AltitudeFiLM(64) + x = torch.randn(2, 64, 8, 8) + altitudes = torch.tensor([150.0, 300.0]) + out = film(x, altitudes) + print(f"FiLM out: {out.shape}") # [2, 64, 8, 8] + + # RoPE2D test + rope = RoPE2D(32) + q = torch.randn(2, 4, 64, 32) # [B, heads, HW, head_dim] + k = torch.randn(2, 4, 64, 32) + q_r, k_r = rope(q, k, 8, 8) + print(f"RoPE out: q={q_r.shape}, k={k_r.shape}") diff --git a/src/models/sofia_v71/model.py b/src/models/sofia_v71/model.py new file mode 100644 index 0000000..39ef1f3 --- /dev/null +++ b/src/models/sofia_v71/model.py @@ -0,0 +1,589 @@ +"""SOFIA v7.1 full model. + +Architecture: + stem → stage1 (StripDCN-lite) → stage2 (StripMixConv) + → stage3 (MambaVision MV5) → stage4 (MambaVision MV1) + → 1×1 neck projection → {Sat-Head, UAV-Head} + (training) Ring aux + +Features: +- Siamese backbone with optional weight-sharing stages 1-2 +- Asymmetric Sat/UAV heads (CVGL-Aware Head v7.1-α) +- Training-time Ring LPN aux for partial matching robustness +- KD taps on F2, F3, F4 for future teacher-student distillation +""" + +from __future__ import annotations + +from typing import Dict, List, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .config import SOFIAConfig +from .blocks import ( + Downsample, + EVSSBridge, + MambaVisionBlock, + StripDCNLiteBlock, + StripMixConvBlock, +) +from .layers import ( + AltitudeFiLM, + CircularHarmonicPool, + GGeM, + LayerNorm2d, + TextFiLM, +) + + +# ============================================================ +# Stem +# ============================================================ + +class Stem(nn.Module): + """Dual-conv FastViT-style stem: 3 → mid (s=2) → out (s=1). + + Downsampling ×2 total (input 256 → 128). + """ + + def __init__(self, in_channels: int, mid_channels: int, out_channels: int) -> None: + super().__init__() + self.conv1 = nn.Conv2d(in_channels, mid_channels, 3, stride=2, padding=1, bias=False) + self.bn1 = nn.BatchNorm2d(mid_channels) + self.act1 = nn.SiLU(inplace=True) + self.conv2 = nn.Conv2d(mid_channels, out_channels, 3, stride=1, padding=1, bias=False) + self.bn2 = nn.BatchNorm2d(out_channels) + self.act2 = nn.SiLU(inplace=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.act1(self.bn1(self.conv1(x))) + x = self.act2(self.bn2(self.conv2(x))) + return x + + +# ============================================================ +# Backbone (shared s1-2 optional, separate s3-4) +# ============================================================ + +class Backbone(nn.Module): + """SOFIA v7.1 backbone: 4 stages + stem. + + Stage 1: StripDCN-lite blocks + Stage 2: StripMixConv blocks + Stage 3: MambaVision MV5 blocks (Mamba ∥ MHSA ∥ Strip) + Stage 4: MambaVision MV1 blocks (Mamba ∥ MHSA) + """ + + def __init__(self, cfg: SOFIAConfig) -> None: + super().__init__() + self.cfg = cfg + dims = cfg.embed_dims + depths = cfg.depths + + # Stem: input → stem_out + self.stem = Stem(cfg.in_channels, cfg.stem_mid, cfg.stem_out) + + # Stage 1: stem_out → dims[0], no downsample at block level + # (stem already did 4× total; stage 1 operates at 64×64 for 256 input) + self.ds1 = Downsample(cfg.stem_out, dims[0]) + self.stage1 = nn.Sequential(*[ + StripDCNLiteBlock( + in_channels=dims[0], + out_channels=dims[0], + expand=cfg.mbconv_expand, + kernel=cfg.strip_kernel_s1, + stride=1, + se_ratio=cfg.se_ratio, + use_dcn=cfg.use_dcn_strip, + ) + for _ in range(depths[0]) + ]) + + # Stage 2 + self.ds2 = Downsample(dims[0], dims[1]) + self.stage2 = nn.Sequential(*[ + StripMixConvBlock( + in_channels=dims[1], + out_channels=dims[1], + expand=cfg.mbconv_expand, + strip_kernel=cfg.strip_kernel_s2, + mix_kernels=cfg.mix_kernels, + stride=1, + se_ratio=cfg.se_ratio, + ) + for _ in range(depths[1]) + ]) + + # Stage 3 — optional EVSS bridge after downsample (B6-inspired refinement + # to smooth DCN→MambaVision semantic gap) + self.ds3 = Downsample(dims[1], dims[2]) + if cfg.use_evss_bridge and "pre_stage3" in cfg.evss_bridge_locations: + self.bridge3 = EVSSBridge( + channels=dims[2], + mamba_d_state=cfg.mamba_d_state, + mamba_dt_rank=cfg.mamba_dt_rank, + mamba_variant="mamba1", # lightweight refinement, not primary mixer + mamba_backend=cfg.mamba_backend, + se_reduction=cfg.se_ratio, + ) + else: + self.bridge3 = None + self.stage3 = nn.Sequential(*[ + MambaVisionBlock( + channels=dims[2], + num_heads=cfg.num_heads_s3, + d_state=cfg.mamba_d_state, + dt_rank=cfg.mamba_dt_rank, + use_strip_branch=cfg.use_strip_branch_s3, + ffn_expand=cfg.ffn_expand, + strip_kernel=cfg.strip_kernel_s1, + mamba_backend=cfg.mamba_backend, + mamba_variant=cfg.mamba_variant, + mamba_extra_kwargs=cfg.mamba_extra_kwargs, + ) + for _ in range(depths[2]) + ]) + + # Stage 4 — optional EVSS bridge after downsample + self.ds4 = Downsample(dims[2], dims[3]) + if cfg.use_evss_bridge and "pre_stage4" in cfg.evss_bridge_locations: + self.bridge4 = EVSSBridge( + channels=dims[3], + mamba_d_state=cfg.mamba_d_state, + mamba_dt_rank=cfg.mamba_dt_rank, + mamba_variant="mamba1", + mamba_backend=cfg.mamba_backend, + se_reduction=cfg.se_ratio, + ) + else: + self.bridge4 = None + self.stage4 = nn.Sequential(*[ + MambaVisionBlock( + channels=dims[3], + num_heads=cfg.num_heads_s4, + d_state=cfg.mamba_d_state, + dt_rank=cfg.mamba_dt_rank, + use_strip_branch=cfg.use_strip_branch_s4, + ffn_expand=cfg.ffn_expand, + strip_kernel=cfg.strip_kernel_s1, + mamba_backend=cfg.mamba_backend, + mamba_variant=cfg.mamba_variant, + mamba_extra_kwargs=cfg.mamba_extra_kwargs, + ) + for _ in range(depths[3]) + ]) + + def forward(self, x: torch.Tensor) -> Dict[str, torch.Tensor]: + """Returns dict with F1-F4 taps for KD/multi-scale use.""" + s0 = self.stem(x) # stem_out, H/2 + f1 = self.stage1(self.ds1(s0)) # dims[0], H/4 + f2 = self.stage2(self.ds2(f1)) # dims[1], H/8 + + # Stage 3 with optional EVSS bridge. + ds3_out = self.ds3(f2) + if self.bridge3 is not None: + ds3_out = self.bridge3(ds3_out) + f3 = self.stage3(ds3_out) # dims[2], H/16 + + # Stage 4 with optional EVSS bridge. + ds4_out = self.ds4(f3) + if self.bridge4 is not None: + ds4_out = self.bridge4(ds4_out) + f4 = self.stage4(ds4_out) # dims[3], H/32 (= 8×8 for 256 input) + + return {"s0": s0, "f1": f1, "f2": f2, "f3": f3, "f4": f4} + + +# ============================================================ +# Neck: ultra-lite 1×1 projection +# ============================================================ + +class UltraLiteNeck(nn.Module): + """1×1 projection + BN + activation.""" + + def __init__(self, in_channels: int, out_channels: int) -> None: + super().__init__() + self.proj = nn.Sequential( + nn.Conv2d(in_channels, out_channels, 1, bias=False), + nn.BatchNorm2d(out_channels), + nn.SiLU(inplace=True), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.proj(x) + + +# ============================================================ +# Heads +# ============================================================ + +class SatHead(nn.Module): + """Satellite view head: [TextFiLM] + GGeM + BN + Linear [+ L2]. + + `text_emb` is optional; when None or `text_film` is disabled, behaves + identically to the original head. `return_normalized=False` returns + pre-L2 descriptors for late gated fusion in a wrapper. + """ + + def __init__( + self, + channels: int, + d_descriptor: int, + return_normalized: bool = True, + use_text_film: bool = False, + text_film_dim: int = 1024, + text_film_hidden: int = 256, + ) -> None: + super().__init__() + self.return_normalized = return_normalized + self.use_text_film = use_text_film + if use_text_film: + self.text_film = TextFiLM(channels, text_dim=text_film_dim, hidden_dim=text_film_hidden) + self.ggem = GGeM(channels) + self.bn = nn.BatchNorm1d(channels, affine=False) + self.proj = nn.Linear(channels, d_descriptor) + + def forward( + self, + x: torch.Tensor, + text_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """x: [B, C, H, W], text_emb: [B, D_txt] or None -> [B, d]""" + if self.use_text_film: + x = self.text_film(x, text_emb) + g = self.ggem(x) + g = self.bn(g) + g = self.proj(g) + if self.return_normalized: + g = F.normalize(g, p=2, dim=-1) + return g + + +class UAVHead(nn.Module): + """UAV view head: FiLM(altitude) [+ TextFiLM] + CHP + BN + Linear [+ L2]. + + Formally SO(2)-invariant via CHP. Altitude-aware via FiLM. Optional + text-FiLM is applied AFTER altitude-FiLM (zero-init β so it starts as + identity). `return_normalized=False` returns pre-L2 descriptors for + late gated fusion in a wrapper. + """ + + def __init__( + self, + channels: int, + d_descriptor: int, + rings: int = 8, + angles: int = 16, + harmonics: int = 4, + use_film: bool = True, + altitude_norm: float = 500.0, + return_normalized: bool = True, + use_text_film: bool = False, + text_film_dim: int = 1024, + text_film_hidden: int = 256, + ) -> None: + super().__init__() + self.return_normalized = return_normalized + self.use_film = use_film + self.use_text_film = use_text_film + if use_film: + self.film = AltitudeFiLM(channels, altitude_norm=altitude_norm) + if use_text_film: + self.text_film = TextFiLM(channels, text_dim=text_film_dim, hidden_dim=text_film_hidden) + + self.chp = CircularHarmonicPool( + channels=channels, + rings=rings, + angles=angles, + harmonics=harmonics, + ) + chp_dim = channels * harmonics + self.bn = nn.BatchNorm1d(chp_dim, affine=False) + self.proj = nn.Linear(chp_dim, d_descriptor) + + def forward( + self, + x: torch.Tensor, + altitude: Optional[torch.Tensor] = None, + text_emb: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + """x: [B, C, H, W], altitude: [B] or None, text_emb: [B, D_txt] or None -> [B, d]""" + if self.use_film: + x = self.film(x, altitude) + if self.use_text_film: + x = self.text_film(x, text_emb) + g = self.chp(x) + g = self.bn(g) + g = self.proj(g) + if self.return_normalized: + g = F.normalize(g, p=2, dim=-1) + return g + + +class RingAuxHead(nn.Module): + """LPN Square-Ring pool + per-ring Linear + L2 (training-only auxiliary). + + Used for partial-matching robustness. Drop at inference. + """ + + def __init__( + self, + channels: int, + rings: int = 4, + d_per_ring: int = 128, + feature_size: int = 8, + ) -> None: + super().__init__() + self.rings = rings + self.feature_size = feature_size + self.d_per_ring = d_per_ring + + # Per-ring GGeM + self.ggems = nn.ModuleList([GGeM(channels) for _ in range(rings)]) + self.projs = nn.ModuleList([ + nn.Linear(channels, d_per_ring) for _ in range(rings) + ]) + + # Precompute ring masks + masks = self._make_ring_masks(rings, feature_size) # [R, H, W] + self.register_buffer("ring_masks", masks, persistent=False) + + @staticmethod + def _make_ring_masks(R: int, S: int) -> torch.Tensor: + """Concentric square rings on SxS feature map.""" + masks = torch.zeros(R, S, S) + center = (S - 1) / 2.0 + # Define ring by Chebyshev distance thresholds + for r in range(R): + r_min = r * (S / (2 * R)) + r_max = (r + 1) * (S / (2 * R)) + for i in range(S): + for j in range(S): + dist = max(abs(i - center), abs(j - center)) + if r_min <= dist < r_max or (r == R - 1 and dist <= r_max): + masks[r, i, j] = 1.0 + # Normalize each ring to sum to 1 for pooling stability (not required but cleaner) + masks = masks / masks.sum(dim=(1, 2), keepdim=True).clamp(min=1.0) + return masks + + def forward(self, x: torch.Tensor) -> List[torch.Tensor]: + """ + Args: + x: [B, C, H, W] + Returns: + List of R tensors each [B, d_per_ring], L2-normalized + """ + B, C, H, W = x.shape + outs = [] + for r in range(self.rings): + mask = self.ring_masks[r].to(x.dtype) # [H, W] + # Apply mask + x_masked = x * mask.unsqueeze(0).unsqueeze(0) # [B, C, H, W] + g = self.ggems[r](x_masked) # [B, C] + g = self.projs[r](g) # [B, d_per_ring] + g = F.normalize(g, p=2, dim=-1) + outs.append(g) + return outs + + +# ============================================================ +# Full SOFIA v7.1 model +# ============================================================ + +class SOFIAv71(nn.Module): + """Full SOFIA v7.1 model. + + Forward signature: + forward(sat: [B,3,H,W], uav: [B,3,H,W], altitude: [B] = None, + return_features: bool = False) -> dict + """ + + def __init__(self, cfg: SOFIAConfig) -> None: + super().__init__() + cfg.validate() + self.cfg = cfg + + # Backbones (shared or separate depending on config) + self.backbone_shared = Backbone(cfg) + if cfg.share_stages_1_2: + # Single backbone, use one for both + self.backbone_sat = None + self.backbone_uav = None + else: + # Fully separate backbones (rare case) + self.backbone_sat = Backbone(cfg) + self.backbone_uav = Backbone(cfg) + # Clear shared + del self.backbone_shared + self.backbone_shared = None + + # Neck (shared) + self.neck = UltraLiteNeck(cfg.embed_dims[-1], cfg.neck_channels) + + # Heads + if cfg.use_asymmetric_heads: + self.sat_head = SatHead( + cfg.neck_channels, cfg.d_descriptor, + return_normalized=cfg.return_normalized, + use_text_film=cfg.use_text_film_sat, + text_film_dim=cfg.text_film_dim, + text_film_hidden=cfg.text_film_hidden, + ) + self.uav_head = UAVHead( + channels=cfg.neck_channels, + d_descriptor=cfg.d_descriptor, + rings=cfg.chp_rings, + angles=cfg.chp_angles, + harmonics=cfg.chp_harmonics, + use_film=cfg.use_film_altitude, + altitude_norm=cfg.altitude_norm, + return_normalized=cfg.return_normalized, + use_text_film=cfg.use_text_film_uav, + text_film_dim=cfg.text_film_dim, + text_film_hidden=cfg.text_film_hidden, + ) + else: + # Symmetric: use SatHead for both + self.sat_head = SatHead( + cfg.neck_channels, cfg.d_descriptor, + return_normalized=cfg.return_normalized, + use_text_film=cfg.use_text_film_sat, + text_film_dim=cfg.text_film_dim, + text_film_hidden=cfg.text_film_hidden, + ) + self.uav_head = self.sat_head + + # Ring aux (training-only) + # Resolution chain: input/2 (stem) → /2 (ds1) → /2 (ds2) → /2 (ds3) → /2 (ds4) + # Total ×32 downsampling → input 256 gives 8×8 final feature map + if cfg.use_ring_aux: + feat_size = cfg.input_size // 32 + self.ring_head = RingAuxHead( + channels=cfg.neck_channels, + rings=cfg.ring_count, + d_per_ring=cfg.d_descriptor // cfg.ring_count, + feature_size=feat_size, + ) + else: + self.ring_head = None + + def _extract_features(self, x: torch.Tensor, view: str) -> Dict[str, torch.Tensor]: + """Run backbone for sat or uav view.""" + if self.cfg.share_stages_1_2: + return self.backbone_shared(x) + else: + bb = self.backbone_sat if view == "sat" else self.backbone_uav + return bb(x) + + def forward( + self, + sat: Optional[torch.Tensor] = None, + uav: Optional[torch.Tensor] = None, + altitude: Optional[torch.Tensor] = None, + text_emb_sat: Optional[torch.Tensor] = None, + text_emb_uav: Optional[torch.Tensor] = None, + return_features: bool = False, + ) -> Dict[str, torch.Tensor]: + """ + Args: + sat: [B, 3, H, W] satellite image (or None to skip) + uav: [B, 3, H, W] UAV image (or None to skip) + altitude: [B] altitude in meters for UAV (or None = neutral) + text_emb_sat: [B, D_txt] caption embedding for SatHead text-FiLM (or None) + text_emb_uav: [B, D_txt] caption embedding for UAVHead text-FiLM (or None) + return_features: if True, also return backbone features F2-F4 for KD + + Returns: + dict with g_sat, g_uav (global descriptors), optional features, + and optional rings (training only) + """ + result: Dict[str, torch.Tensor] = {} + + # SAT path + if sat is not None: + feats_sat = self._extract_features(sat, view="sat") + f4_sat = feats_sat["f4"] + neck_sat = self.neck(f4_sat) + g_sat = self.sat_head(neck_sat, text_emb=text_emb_sat) + result["g_sat"] = g_sat + if return_features: + result["features_sat"] = feats_sat + result["neck_sat"] = neck_sat + if self.training and self.ring_head is not None: + result["rings_sat"] = self.ring_head(neck_sat) + + # UAV path + if uav is not None: + feats_uav = self._extract_features(uav, view="uav") + f4_uav = feats_uav["f4"] + neck_uav = self.neck(f4_uav) + if self.cfg.use_asymmetric_heads: + g_uav = self.uav_head(neck_uav, altitude=altitude, text_emb=text_emb_uav) + else: + g_uav = self.uav_head(neck_uav, text_emb=text_emb_uav) + result["g_uav"] = g_uav + if return_features: + result["features_uav"] = feats_uav + result["neck_uav"] = neck_uav + if self.training and self.ring_head is not None: + result["rings_uav"] = self.ring_head(neck_uav) + + return result + + def count_parameters(self, only_trainable: bool = True) -> Dict[str, int]: + """Count parameters per module.""" + counts = {} + total = 0 + for name, module in self.named_children(): + if module is None: + continue + n = sum(p.numel() for p in module.parameters() if not only_trainable or p.requires_grad) + counts[name] = n + total += n + counts["_total"] = total + return counts + + +# ============================================================ +# Build helper +# ============================================================ + +def build_sofia(preset: str = "M") -> SOFIAv71: + """Build SOFIA v7.1 model from preset name.""" + from .config import sofia_m_config, sofia_l_config, sofia_tiny_config + preset_map = { + "M": sofia_m_config, + "L": sofia_l_config, + "Tiny": sofia_tiny_config, + } + if preset not in preset_map: + raise ValueError(f"Unknown preset '{preset}'. Available: {list(preset_map)}") + cfg = preset_map[preset]() + return SOFIAv71(cfg) + + +if __name__ == "__main__": + # Smoke test with Tiny preset (fast) + print("Building SOFIA-Tiny for smoke test...") + model = build_sofia("Tiny") + model.eval() + + counts = model.count_parameters() + total_m = counts["_total"] / 1e6 + print(f"Total params: {total_m:.2f} M") + print("Per-module:") + for k, v in counts.items(): + print(f" {k}: {v / 1e6:.3f} M") + + # Dry forward + sat = torch.randn(1, 3, 256, 256) + uav = torch.randn(1, 3, 256, 256) + alt = torch.tensor([150.0]) + with torch.no_grad(): + out = model(sat=sat, uav=uav, altitude=alt) + for k, v in out.items(): + if isinstance(v, torch.Tensor): + print(f" out[{k}]: {tuple(v.shape)}") + else: + print(f" out[{k}]: {type(v).__name__}") diff --git a/src/models/sofia_v71/quant.py b/src/models/sofia_v71/quant.py new file mode 100644 index 0000000..89046b4 --- /dev/null +++ b/src/models/sofia_v71/quant.py @@ -0,0 +1,565 @@ +"""SOFIA v7.1 quantization utilities. + +Two reusable building blocks for INT8 deployment: + +1. **OffsetClampSTE** (DCN-M2): hard clamp DCN offsets to physical range + ``[-k, +k]`` with straight-through gradient. Bounded distribution = + stable percentile clipping during PTQ + cleaner offset semantics. + +2. **KScaledFakeQuant** (R5-style): multi-bin fake quantization. Bins are + chosen by magnitude percentiles; each bin has its own scale. Adapts + quantization resolution to long-tail distributions (Mamba Δ, output y). + +Plus a drop-in: + +3. **KScaledMamba2Block**: ``Mamba2Block`` subclass with k-scaled fake-quant + nodes on ``x_main``, ``delta``, and per-token ``y`` outputs. Implements + R5 reparam principle (state ``h_t`` kept in model dtype; only observable + tensors quantized). + +All components are PTQ/QAT-compatible (STE backward, calibration mode). +""" + +from __future__ import annotations + +from typing import List, Optional, Sequence, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from .blocks import Mamba2Block + + +# ============================================================ +# DCN-M2: Offset clamp with STE +# ============================================================ + +class _OffsetClampSTE(torch.autograd.Function): + """Hard clamp [low, high] in forward, identity gradient in backward.""" + + @staticmethod + def forward(ctx, x: torch.Tensor, low: float, high: float) -> torch.Tensor: + return torch.clamp(x, low, high) + + @staticmethod + def backward(ctx, grad: torch.Tensor): + # STE: pass gradient unmodified + return grad, None, None + + +def offset_clamp_ste(offsets: torch.Tensor, kernel_size: int) -> torch.Tensor: + """Functional DCN-M2: clamp offsets to ``[-k, +k]`` with STE backward. + + Args: + offsets: arbitrary-shape tensor of DCN offsets. + kernel_size: physical receptive-field bound k. Output is clamped to + absolute values <= k. Standard choice equals the DW kernel size. + + Returns: + Clamped tensor (same shape as ``offsets``). Backward propagates + gradient as if no clamp was applied, so training can still adjust + weights even when an offset hits the boundary. + """ + k = float(kernel_size) + return _OffsetClampSTE.apply(offsets, -k, k) + + +class OffsetClampSTE(nn.Module): + """Module wrapper for :func:`offset_clamp_ste`. + + Use as a drop-in layer between ``offset_predictor`` and ``deform_conv2d``:: + + self.clamp = OffsetClampSTE(kernel_size=7) + offsets = self.clamp(self.offset_predictor(x)) + """ + + def __init__(self, kernel_size: int) -> None: + super().__init__() + self.kernel_size = int(kernel_size) + + def forward(self, offsets: torch.Tensor) -> torch.Tensor: + return offset_clamp_ste(offsets, self.kernel_size) + + def extra_repr(self) -> str: + return f"kernel_size={self.kernel_size}, range=[-{self.kernel_size}, +{self.kernel_size}]" + + +# ============================================================ +# K-Scaled Fake Quantization (R5 style, multi-bin) +# ============================================================ + +class _STEFakeQuant(torch.autograd.Function): + """Fake-quantize: ``round(x / s) * s`` clamped to [qmin, qmax], STE backward. + + `scale` is a tensor broadcastable with `x` (per-tensor scalar or + per-element via gathered per-bin scales). + """ + + @staticmethod + def forward( + ctx, + x: torch.Tensor, + scale: torch.Tensor, + qmin: int, + qmax: int, + ) -> torch.Tensor: + x_int = torch.round(x / scale).clamp(qmin, qmax) + return x_int * scale + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + return grad_output, None, None, None + + +class KScaledFakeQuant(nn.Module): + """k-scaled (multi-bin magnitude-bucketed) fake quantization. + + Algorithm: + 1. ``g(x) = bin index in [0..k-1] based on |x|`` using thresholds + ``t_0 < t_1 < ... < t_{k-2}``. + 2. ``s(x) = scales[g(x)]`` per-element scale. + 3. ``x_q = round(x / s) * s`` clamped to ``[qmin, qmax]``. + 4. STE backward. + + Setting ``num_bins=1`` reduces to standard single-scale fake quant. + + Calibration workflow:: + + fq = KScaledFakeQuant(num_bins=3) + fq.start_calibration() + for batch in cal_loader: + _ = model(batch) # collect stats + fq.finalize_calibration(percentile=99.9) + # fq is now active (calibrating=False) and applies fake-quant in forward + """ + + def __init__( + self, + num_bins: int = 3, + qmin: int = -128, + qmax: int = 127, + max_calibration_samples: int = 200_000, + ) -> None: + super().__init__() + if num_bins < 1: + raise ValueError(f"num_bins must be >= 1, got {num_bins}") + if qmax <= qmin: + raise ValueError(f"qmax ({qmax}) must be > qmin ({qmin})") + + self.num_bins = num_bins + self.qmin = qmin + self.qmax = qmax + self.max_calibration_samples = max_calibration_samples + + # Thresholds: (num_bins - 1,) — bin boundaries on |x|. For num_bins=1, empty. + n_thr = max(0, num_bins - 1) + self.register_buffer("thresholds", torch.zeros(n_thr)) + + # Scales: (num_bins,) — one scale per bin. Default 1.0 = no quant effect. + self.register_buffer("scales", torch.ones(num_bins)) + + # Per-instance state (not buffers — runtime-only) + self.calibrating: bool = False + self.enabled: bool = True + self._calib_buffer: List[torch.Tensor] = [] + self._n_collected: int = 0 + + # -------------------- Calibration API -------------------- + + def start_calibration(self) -> None: + """Begin collecting input distribution stats.""" + self.calibrating = True + self._calib_buffer = [] + self._n_collected = 0 + + def stop_calibration(self) -> None: + """End calibration without finalizing scales (e.g. abort).""" + self.calibrating = False + self._calib_buffer = [] + self._n_collected = 0 + + def finalize_calibration(self, percentile: float = 99.9) -> None: + """Compute thresholds and per-bin scales from collected stats. + + Args: + percentile: per-bin tail percentile for scale computation + (typical 99.9). Robust to extreme outliers. + """ + if not self._calib_buffer: + self.calibrating = False + return + + all_data = torch.cat(self._calib_buffer) + abs_data = all_data.abs() + + if self.num_bins == 1: + # Single scale: percentile of |x| + tail = torch.quantile(abs_data, percentile / 100.0).item() + self.scales[0] = max(tail / float(self.qmax), 1e-8) + else: + # Threshold positions: equal-mass quantile splits + qpts = torch.linspace(0.0, 1.0, self.num_bins + 1, device=abs_data.device)[1:-1] + thr_vals = torch.quantile(abs_data, qpts) + self.thresholds.copy_(thr_vals.to(self.thresholds.device)) + + # Per-bin scale: per-bin tail percentile / qmax + for g in range(self.num_bins): + if g == 0: + lo = 0.0 + hi = thr_vals[0].item() + elif g == self.num_bins - 1: + lo = thr_vals[-1].item() + hi = float("inf") + else: + lo = thr_vals[g - 1].item() + hi = thr_vals[g].item() + + mask = (abs_data >= lo) & (abs_data < hi) + if mask.any(): + bin_tail = torch.quantile( + abs_data[mask], percentile / 100.0 + ).item() + self.scales[g] = max(bin_tail / float(self.qmax), 1e-8) + else: + self.scales[g] = 1.0 + + self._calib_buffer = [] + self._n_collected = 0 + self.calibrating = False + + # -------------------- Forward -------------------- + + def _bin_index(self, x: torch.Tensor) -> torch.Tensor: + """Per-element bin index (long tensor, shape == x.shape).""" + if self.num_bins == 1: + return torch.zeros_like(x, dtype=torch.long) + + abs_x = x.abs() + idx = torch.zeros_like(x, dtype=torch.long) + for i in range(self.num_bins - 1): + t = self.thresholds[i] + idx = torch.where(abs_x >= t, idx + 1, idx) + return idx + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if not self.enabled: + return x + + if self.calibrating: + # Subsample to bound memory + with torch.no_grad(): + flat = x.detach().flatten() + budget = self.max_calibration_samples - self._n_collected + if budget > 0: + take = min(flat.numel(), budget) + if take < flat.numel(): + # Random subsample for representativeness + idx = torch.randperm(flat.numel(), device=flat.device)[:take] + sample = flat[idx] + else: + sample = flat + self._calib_buffer.append(sample.cpu()) + self._n_collected += sample.numel() + return x + + # Apply k-scaled fake-quant + bin_idx = self._bin_index(x) + scale = self.scales[bin_idx] # shape == x.shape + return _STEFakeQuant.apply(x, scale, self.qmin, self.qmax) + + def extra_repr(self) -> str: + return ( + f"num_bins={self.num_bins}, qmin={self.qmin}, qmax={self.qmax}, " + f"calibrating={self.calibrating}, enabled={self.enabled}" + ) + + +# ============================================================ +# K-Scaled Mamba-2 drop-in +# ============================================================ + +class KScaledMamba2Block(Mamba2Block): + """Drop-in replacement for :class:`Mamba2Block` with k-scaled fake-quant. + + Adds ``KScaledFakeQuant`` nodes on three observable tensors inside the + Mamba-2 forward pass: + + - ``x_main``: scan input after conv1d + SiLU + - ``delta``: time-step ``Δ`` after softplus (addresses MF2 long tail) + - ``y``: per-token scan output (addresses MF4 state propagation) + + The recurrent state ``h_t`` is **not** quantized — kept in model dtype + inside the scan loop. This implements R5's reparam principle: only + observable I/O is quantized; internal state stays high-precision. + + Backend constraint: forces ``backend='torch'``. The mamba_ssm CUDA + backend would require k-scaled support inside the custom kernel (out of + scope here). For a fast deploy path, train+QAT in torch, then export to + a custom TRT plugin that does the same INT8 scan. + + Args: + channels: model dim ``C``. + d_state: state size ``N`` (Mamba-2 default 64). + headdim: head dim (must divide channels). + d_conv: local conv kernel. + expand: inner expansion factor. + backend: must be ``'torch'`` (or ``'auto'``, which resolves to torch). + num_bins: ``k`` for k-scaled (typical 2-4). + targets: which tensors to quantize. Subset of + ``('x_main', 'delta', 'y')``. + + Example:: + + mamba = KScaledMamba2Block(channels=192, num_bins=3, + targets=('delta', 'y')) + mamba.start_calibration() + for batch in cal_loader: + _ = model_with_mamba(batch) + mamba.finalize_calibration() + # Now in PTQ mode — forward applies fake-quant on chosen targets. + """ + + SUPPORTED_TARGETS: Tuple[str, ...] = ("x_main", "delta", "y") + + def __init__( + self, + channels: int, + d_state: int = 64, + headdim: int = 64, + d_conv: int = 4, + expand: int = 2, + backend: str = "torch", + num_bins: int = 3, + targets: Sequence[str] = ("x_main", "delta", "y"), + ) -> None: + # Force torch — k-scaled requires explicit forward-pass control + if backend == "auto": + backend = "torch" + if backend != "torch": + raise ValueError( + "KScaledMamba2Block currently supports only backend='torch'. " + f"Got backend='{backend}'. The mamba_ssm CUDA kernel does not " + "expose hooks for k-scaled fake-quant; integrating INT8 scan " + "via TRT plugin is a separate deploy-time concern." + ) + + super().__init__( + channels=channels, + d_state=d_state, + headdim=headdim, + d_conv=d_conv, + expand=expand, + backend="torch", + ) + + unknown = set(targets) - set(self.SUPPORTED_TARGETS) + if unknown: + raise ValueError( + f"Unknown k-scaled targets: {unknown}. " + f"Supported: {self.SUPPORTED_TARGETS}" + ) + self.targets = set(targets) + self.num_bins = num_bins + + # Build fake-quant nodes only for active targets + if "x_main" in self.targets: + self.fq_x_main = KScaledFakeQuant(num_bins=num_bins) + if "delta" in self.targets: + self.fq_delta = KScaledFakeQuant(num_bins=num_bins) + if "y" in self.targets: + self.fq_y = KScaledFakeQuant(num_bins=num_bins) + + # -------------------- Calibration helpers -------------------- + + def _all_fqs(self) -> List[KScaledFakeQuant]: + out: List[KScaledFakeQuant] = [] + for name in self.SUPPORTED_TARGETS: + attr = f"fq_{name}" + mod = getattr(self, attr, None) + if isinstance(mod, KScaledFakeQuant): + out.append(mod) + return out + + def start_calibration(self) -> None: + for fq in self._all_fqs(): + fq.start_calibration() + + def finalize_calibration(self, percentile: float = 99.9) -> None: + for fq in self._all_fqs(): + fq.finalize_calibration(percentile=percentile) + + def set_quant_enabled(self, enabled: bool) -> None: + for fq in self._all_fqs(): + fq.enabled = enabled + + # -------------------- Forward -------------------- + + def _forward_torch(self, x: torch.Tensor) -> torch.Tensor: + """Mamba-2 torch forward with k-scaled fake-quant on selected paths.""" + B, L, C = x.shape + + # In projection (matches parent layout) + z_in = self.in_proj(x) + xz, BC, dt = torch.split( + z_in, + [2 * self.d_inner, 2 * self.d_state, self.nheads], + dim=-1, + ) + x_main, z_gate = xz.chunk(2, dim=-1) + B_p, C_p = BC.chunk(2, dim=-1) + + # Local conv + SiLU + x_main_t = x_main.transpose(1, 2) + x_main_t = self.conv1d(x_main_t)[..., :L] + x_main = F.silu(x_main_t).transpose(1, 2) + + # K-SCALED QUANT: x_main (scan input) + if "x_main" in self.targets: + x_main = self.fq_x_main(x_main) + + # Δ + dt = F.softplus(dt) + if "delta" in self.targets: + dt = self.fq_delta(dt) + + A = -torch.exp(self.A_log) # [nheads] + x_head = x_main.view(B, L, self.nheads, self.headdim) + + # Sequential scan. R5 reparam: state h kept in model dtype throughout. + h = torch.zeros( + B, self.nheads, self.headdim, self.d_state, + device=x.device, dtype=x.dtype, + ) + ys: List[torch.Tensor] = [] + for t in range(L): + dt_t = dt[:, t] + B_t = B_p[:, t] + C_t = C_p[:, t] + x_t = x_head[:, t] + + dA = torch.exp(dt_t * A.unsqueeze(0)).unsqueeze(-1).unsqueeze(-1) + dB = (dt_t.unsqueeze(-1) * B_t.unsqueeze(1)).unsqueeze(2) + + h = dA * h + dB * x_t.unsqueeze(-1) + y_t = (h * C_t.unsqueeze(1).unsqueeze(1)).sum(dim=-1) + y_t = y_t + self.D.unsqueeze(0).unsqueeze(-1) * x_t + + # K-SCALED QUANT: per-token y (after scan, before gate) + if "y" in self.targets: + y_t = self.fq_y(y_t) + + ys.append(y_t) + + y = torch.stack(ys, dim=1) + y = y.reshape(B, L, self.d_inner) + y = y * F.silu(z_gate) + y = self.out_proj(y) + return y + + +# ============================================================ +# Model-wide calibration helpers +# ============================================================ + +def start_calibration(model: nn.Module) -> None: + """Walk model, start calibration on every quant module.""" + seen = set() + # First, KScaledMamba2Block instances drive their internal FQs. + for m in model.modules(): + if isinstance(m, KScaledMamba2Block): + m.start_calibration() + for fq in m._all_fqs(): + seen.add(id(fq)) + # Then standalone KScaledFakeQuant not owned by a KScaledMamba2Block. + for m in model.modules(): + if isinstance(m, KScaledFakeQuant) and id(m) not in seen: + m.start_calibration() + + +def finalize_calibration(model: nn.Module, percentile: float = 99.9) -> None: + """Walk model, finalize all quant modules with the same percentile.""" + seen = set() + for m in model.modules(): + if isinstance(m, KScaledMamba2Block): + m.finalize_calibration(percentile=percentile) + for fq in m._all_fqs(): + seen.add(id(fq)) + for m in model.modules(): + if isinstance(m, KScaledFakeQuant) and id(m) not in seen: + m.finalize_calibration(percentile=percentile) + + +def set_quant_enabled(model: nn.Module, enabled: bool) -> None: + """Toggle every quant module on/off. Useful for FP-vs-INT8 comparison.""" + for m in model.modules(): + if isinstance(m, KScaledFakeQuant): + m.enabled = enabled + elif isinstance(m, KScaledMamba2Block): + m.set_quant_enabled(enabled) + + +# ============================================================ +# Smoke tests +# ============================================================ + +if __name__ == "__main__": + torch.manual_seed(0) + + print("=== DCN-M2 OffsetClampSTE ===") + offsets = torch.randn(2, 14, 16, 16, requires_grad=True) * 5.0 # outliers ~ ±15 + print(f" in: range [{offsets.min().item():.2f}, {offsets.max().item():.2f}]") + clamped = offset_clamp_ste(offsets, kernel_size=3) + print(f" out: range [{clamped.min().item():.2f}, {clamped.max().item():.2f}]") + clamped.sum().backward() + grad_mean = offsets.grad.mean().item() + grad_std = offsets.grad.std().item() + print(f" grad mean={grad_mean:.4f} std={grad_std:.4f} " + f"(STE → exactly 1.0 for every element)") + + print("\n=== KScaledFakeQuant (num_bins=3) ===") + fq = KScaledFakeQuant(num_bins=3) + fq.start_calibration() + with torch.no_grad(): + for _ in range(20): + x = torch.randn(64, 32) * 1.5 + x[0, 0] = 25.0 # outlier + x[3, 7] = -40.0 + _ = fq(x) + fq.finalize_calibration(percentile=99.9) + print(f" thresholds: {fq.thresholds.tolist()}") + print(f" scales: {[round(s, 5) for s in fq.scales.tolist()]}") + + x_test = torch.randn(8, 32) * 3 + x_test[0, 0] = 30.0 + x_q = fq(x_test) + err = (x_q - x_test).abs().mean().item() + rel_err = err / x_test.abs().mean().item() + print(f" avg abs err: {err:.4e} (relative: {rel_err:.2%})") + + print("\n=== KScaledMamba2Block ===") + block = KScaledMamba2Block( + channels=128, d_state=32, headdim=32, num_bins=3, + targets=("delta", "y"), + ) + n_params = sum(p.numel() for p in block.parameters()) / 1e3 + print(f" params: {n_params:.1f} K") + + # Calibration on dummy data + block.start_calibration() + with torch.no_grad(): + for _ in range(3): + _ = block(torch.randn(1, 64, 128)) + block.finalize_calibration(percentile=99.9) + + # FP vs k-scaled-INT8 comparison + x = torch.randn(2, 64, 128) + block.set_quant_enabled(False) + y_fp = block(x) + block.set_quant_enabled(True) + y_q = block(x) + diff = (y_fp - y_q).abs().mean().item() + rel = diff / y_fp.abs().mean().item() + print(f" FP output range [{y_fp.min().item():.3f}, {y_fp.max().item():.3f}]") + print(f" INT8 output range [{y_q.min().item():.3f}, {y_q.max().item():.3f}]") + print(f" abs diff: {diff:.4e} (relative: {rel:.2%})") diff --git a/src/models/sofia_v71/verify.py b/src/models/sofia_v71/verify.py new file mode 100644 index 0000000..bc0faba --- /dev/null +++ b/src/models/sofia_v71/verify.py @@ -0,0 +1,158 @@ +"""Verify SOFIA v7.1 model scales: param count, memory footprint, forward pass. + +Run: + python -m sofia_v71.verify + python -m sofia_v71.verify --preset L + python -m sofia_v71.verify --preset M --benchmark +""" + +from __future__ import annotations + +import argparse +import time +from typing import Optional + +import torch + +from .blocks import is_mamba_ssm_available, is_mamba2_available +from .config import sofia_m_config, sofia_l_config, sofia_tiny_config +from .model import SOFIAv71 + + +def count_parameters(model: torch.nn.Module) -> dict: + """Count parameters per named child.""" + counts = {} + for name, param in model.named_parameters(): + top = name.split(".")[0] + counts[top] = counts.get(top, 0) + param.numel() + counts["_total"] = sum(counts.values()) + return counts + + +def estimate_quantized_size(n_params: int, precision: str = "int8") -> float: + """Estimate on-disk / VRAM weight storage in MB.""" + bytes_per_param = {"fp32": 4, "fp16": 2, "int8": 1, "int4": 0.5}[precision] + return n_params * bytes_per_param / (1024 ** 2) + + +def test_forward(model: SOFIAv71, device: str = "cpu", batch_size: int = 1) -> dict: + """Dry forward pass, return output shapes.""" + model = model.to(device).eval() + cfg = model.cfg + sat = torch.randn(batch_size, 3, cfg.input_size, cfg.input_size, device=device) + uav = torch.randn(batch_size, 3, cfg.input_size, cfg.input_size, device=device) + altitude = torch.rand(batch_size, device=device) * 300.0 + 50.0 + + shapes = {} + with torch.no_grad(): + out = model(sat=sat, uav=uav, altitude=altitude) + for k, v in out.items(): + if isinstance(v, torch.Tensor): + shapes[k] = tuple(v.shape) + elif isinstance(v, list): + shapes[k] = f"list[{len(v)} × {tuple(v[0].shape)}]" + elif isinstance(v, dict): + for kk, vv in v.items(): + if isinstance(vv, torch.Tensor): + shapes[f"{k}.{kk}"] = tuple(vv.shape) + return shapes + + +def benchmark(model: SOFIAv71, device: str = "cpu", n_warmup: int = 3, n_runs: int = 20) -> dict: + """Measure forward pass latency.""" + model = model.to(device).eval() + cfg = model.cfg + sat = torch.randn(1, 3, cfg.input_size, cfg.input_size, device=device) + uav = torch.randn(1, 3, cfg.input_size, cfg.input_size, device=device) + altitude = torch.tensor([150.0], device=device) + + # Warmup + with torch.no_grad(): + for _ in range(n_warmup): + _ = model(sat=sat, uav=uav, altitude=altitude) + if device.startswith("cuda"): + torch.cuda.synchronize() + + # Measured runs + times = [] + for _ in range(n_runs): + if device.startswith("cuda"): + torch.cuda.synchronize() + t0 = time.perf_counter() + _ = model(sat=sat, uav=uav, altitude=altitude) + if device.startswith("cuda"): + torch.cuda.synchronize() + times.append(time.perf_counter() - t0) + + times_ms = [t * 1000 for t in times] + return { + "mean_ms": sum(times_ms) / len(times_ms), + "min_ms": min(times_ms), + "max_ms": max(times_ms), + "std_ms": (sum((t - sum(times_ms) / len(times_ms)) ** 2 for t in times_ms) / len(times_ms)) ** 0.5, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description="Verify SOFIA v7.1 model") + parser.add_argument("--preset", choices=["Tiny", "M", "L"], default="Tiny", + help="Model preset (default: Tiny for fast smoke test)") + parser.add_argument("--device", default="cpu") + parser.add_argument("--benchmark", action="store_true") + parser.add_argument("--batch-size", type=int, default=1) + args = parser.parse_args() + + cfg_fn = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config}[args.preset] + cfg = cfg_fn() + + # Report mamba_ssm availability + print(f"\n=== Environment ===") + print(f" mamba_ssm v1 (Mamba-1 selective scan): " + f"{'available' if is_mamba_ssm_available() else 'NOT available'}") + print(f" mamba_ssm Mamba-2 (SSD dual form): " + f"{'available' if is_mamba2_available() else 'NOT available'}") + print(f" configured variant: {cfg.mamba_variant}, backend: {cfg.mamba_backend}") + + print(f"\n=== SOFIA-{args.preset} configuration ===") + print(cfg.summary()) + + print("\n=== Building model ===") + model = SOFIAv71(cfg) + counts = count_parameters(model) + total = counts["_total"] + + print(f"\nTotal params: {total / 1e6:.2f} M") + print("\nPer-module breakdown:") + for k, v in sorted(counts.items(), key=lambda x: -x[1]): + if k == "_total": + continue + pct = 100 * v / total + print(f" {k:30s} {v / 1e6:>8.3f} M ({pct:5.1f}%)") + + print("\n=== Memory footprint estimates ===") + for prec in ["fp32", "fp16", "int8", "int4"]: + mb = estimate_quantized_size(total, prec) + print(f" {prec.upper():6s} weights: {mb:>8.1f} MB ({mb / 1024:.2f} GB)") + + print("\n=== Forward pass test ===") + try: + shapes = test_forward(model, device=args.device, batch_size=args.batch_size) + for k, v in shapes.items(): + print(f" out[{k}]: {v}") + except Exception as e: + print(f" ERROR during forward: {type(e).__name__}: {e}") + import traceback + traceback.print_exc() + + if args.benchmark: + print("\n=== Latency benchmark ===") + try: + stats = benchmark(model, device=args.device) + print(f" mean: {stats['mean_ms']:.2f} ms min: {stats['min_ms']:.2f} " + f"max: {stats['max_ms']:.2f} std: {stats['std_ms']:.2f}") + except Exception as e: + print(f" ERROR during benchmark: {type(e).__name__}: {e}") + + +if __name__ == "__main__": + main() diff --git a/src/training/train_gtauav.py b/src/training/train_gtauav.py index 24ae3ea..0655a38 100644 --- a/src/training/train_gtauav.py +++ b/src/training/train_gtauav.py @@ -54,6 +54,14 @@ from src.models.asymmetric_encoder import ( get_drone_train_transform, get_satellite_train_transform, ) +from src.models.sofia_fusion_encoder import SOFIAFusionEncoder +from src.models.sofia_v1 import SOFIAv1Config +from src.models.sofia_v1_fusion_encoder import SOFIAv1FusionEncoder +from src.models.sofia_v71 import ( + sofia_l_config, + sofia_m_config, + sofia_tiny_config, +) LOGGER = logging.getLogger("caption_test.train_gtauav") @@ -91,11 +99,30 @@ class TrainConfigGTAUAV: mona_last_n_blocks: int = 12 # inject adapters only in last 12 of 24 ViT blocks gradient_checkpointing: bool = True # trade compute for VRAM (allows larger batch) # StripNet backbone option (replaces DINOv3 when backbone="stripnet"). - backbone: str = "dinov3" # "dinov3" or "stripnet" + backbone: str = "dinov3" # "dinov3", "stripnet", or "sofia" stripnet_path: str = "nn_models/STRIPNET/stripnet_s.pth" stripnet_mona_last_n_stages: int = 2 # Conv-MONA in last N of 4 StripNet stages (0 = disable MONA) stripnet_freeze: bool = True # If False, StripNet backbone is fully trainable (full fine-tune) stripnet_backbone_lr_factor: float = 0.1 # Backbone LR = learning_rate * factor (only when unfrozen) + # SOFIA backbone options (used when backbone="sofia"). Trained from scratch — no pretrained checkpoint. + sofia_preset: str = "Tiny" # "Tiny" | "M" | "L" + sofia_d_descriptor: int = 1024 # retrieval space (1024 = match TextFusionMLP out_dim) + sofia_use_text_film_uav: bool = True # mid-level text-FiLM in UAV head + sofia_use_text_film_sat: bool = True # mid-level text-FiLM in SAT head + sofia_lora_rank: int = 4 + sofia_mamba_variant: str = "mamba2" # "mamba1" | "mamba2" | "efficient_vmamba" + sofia_mamba_backend: str = "auto" # "auto" | "torch" | "mamba_ssm" + # EVSSBridge (B6-inspired refinement between heterogeneous stages, opt-in). + sofia_use_evss_bridge: bool = False + sofia_evss_bridge_locations: list[str] = field(default_factory=lambda: ["pre_stage3"]) + # SOFIA v1 backbone options (used when backbone="sofia_v1"). StripNet+DCN, from scratch. + sofia_v1_variant: str = "small" # "tiny_tiny" | "tiny" | "small" | "small_v2" + sofia_v1_dcn_variant: str = "v2" # "v2" (torchvision DeformConv2d, stable) | "v4" (OpenGVLab, leaky) + sofia_v1_d_descriptor: int = 1024 + sofia_v1_use_text_film_uav: bool = True + sofia_v1_use_text_film_sat: bool = True + sofia_v1_use_film_altitude: bool = True + sofia_v1_lora_rank: int = 4 # Training. resume_from: str | None = None # path to checkpoint for resuming @@ -167,7 +194,7 @@ def _atomic_save(obj: dict, path: Path) -> None: def _build_param_groups( - model: AsymmetricEncoder, + model: nn.Module, lr: float, text_lr_factor: float, stripnet_backbone_lr_factor: float = 0.1, @@ -177,7 +204,8 @@ def _build_param_groups( Groups: - text_encoder.* → lr * text_lr_factor (default 1e-5) - image_encoder.backbone.* (when StripNet unfrozen) → lr * stripnet_backbone_lr_factor (default 1e-5) - - everything else (MONA, projection, TextFusionMLP, gates, tau, MONA on Conv) → lr + - everything else (MONA, projection, TextFusionMLP, gates, tau, MONA on Conv, + SOFIA backbone+heads when backbone="sofia") → lr """ text_params = [] backbone_params = [] @@ -249,9 +277,13 @@ def _embed_drone_queries( embs: list[torch.Tensor] = [] for batch in tqdm(loader, desc=" dss-embed-queries", unit="batch", leave=False): drone_img = batch["drone_img"].to(device, non_blocking=True) + altitude = batch.get("altitude") + if altitude is not None: + altitude = altitude.to(device, non_blocking=True) q = model.encode_query( drone_img, batch["caption_l1"], batch["caption_l2"], batch["caption_l3"], + altitude=altitude, ) embs.append(q.cpu()) @@ -336,9 +368,13 @@ def _evaluate( if max_batches is not None and i >= max_batches: break drone_img = batch["drone_img"].to(device, non_blocking=True) + altitude = batch.get("altitude") + if altitude is not None: + altitude = altitude.to(device, non_blocking=True) q = model.encode_query( drone_img, batch["caption_l1"], batch["caption_l2"], batch["caption_l3"], + altitude=altitude, ) query_embs.append(q.cpu()) query_valid_names.extend(batch["valid_sat_names"]) @@ -575,40 +611,97 @@ def train(cfg: TrainConfigGTAUAV) -> None: if cfg.resume_from is not None: LOGGER.info("Resuming from %s", cfg.resume_from) - model, resume_ckpt = AsymmetricEncoder.load_checkpoint( - cfg.resume_from, - dino_web_path=cfg.dino_web_path, - dino_sat_path=cfg.dino_sat_path, - lrsclip_path=cfg.lrsclip_path, - device=cfg.device, - ) + if cfg.backbone == "sofia": + model, resume_ckpt = SOFIAFusionEncoder.load_checkpoint( + cfg.resume_from, + lrsclip_path=cfg.lrsclip_path, + device=cfg.device, + ) + elif cfg.backbone == "sofia_v1": + model, resume_ckpt = SOFIAv1FusionEncoder.load_checkpoint( + cfg.resume_from, + lrsclip_path=cfg.lrsclip_path, + device=cfg.device, + ) + else: + model, resume_ckpt = AsymmetricEncoder.load_checkpoint( + cfg.resume_from, + dino_web_path=cfg.dino_web_path, + dino_sat_path=cfg.dino_sat_path, + lrsclip_path=cfg.lrsclip_path, + device=cfg.device, + ) start_epoch = resume_ckpt.get("epoch", -1) + 1 else: mode_str = "baseline (no text)" if cfg.baseline_mode else "with text (L1/L2/L3)" - if cfg.backbone == "stripnet": + if cfg.backbone == "sofia": + enc_str = f"SOFIA-{cfg.sofia_preset} (text-FiLM uav={cfg.sofia_use_text_film_uav}, sat={cfg.sofia_use_text_film_sat})" + elif cfg.backbone == "sofia_v1": + enc_str = f"SOFIAv1-{cfg.sofia_v1_variant} (StripNet+DCNv4, text-FiLM uav={cfg.sofia_v1_use_text_film_uav}, sat={cfg.sofia_v1_use_text_film_sat})" + elif cfg.backbone == "stripnet": enc_str = "StripNet-small (shared, 512→1024 proj)" else: enc_str = "shared DINOv3 WEB" if cfg.shared_encoder else "asymmetric (WEB + SAT)" LOGGER.info("Building model — %s, %s", mode_str, enc_str) - model = AsymmetricEncoder( - dino_web_path=cfg.dino_web_path, - dino_sat_path=cfg.dino_sat_path, - lrsclip_path=cfg.lrsclip_path, - init_gate=cfg.init_gate, - baseline_mode=cfg.baseline_mode, - shared_encoder=cfg.shared_encoder, - mona_bottleneck=cfg.mona_bottleneck, - mona_last_n_blocks=cfg.mona_last_n_blocks, - device=cfg.device, - backbone=cfg.backbone, - stripnet_path=cfg.stripnet_path, - stripnet_mona_last_n_stages=cfg.stripnet_mona_last_n_stages, - stripnet_freeze=cfg.stripnet_freeze, - ).to(cfg.device) + if cfg.backbone == "sofia": + preset_map = {"Tiny": sofia_tiny_config, "M": sofia_m_config, "L": sofia_l_config} + if cfg.sofia_preset not in preset_map: + raise ValueError(f"Unknown sofia_preset={cfg.sofia_preset!r}") + sofia_cfg = preset_map[cfg.sofia_preset]() + sofia_cfg.d_descriptor = cfg.sofia_d_descriptor + sofia_cfg.text_film_dim = cfg.sofia_d_descriptor + sofia_cfg.use_text_film_uav = cfg.sofia_use_text_film_uav and not cfg.baseline_mode + sofia_cfg.use_text_film_sat = cfg.sofia_use_text_film_sat and not cfg.baseline_mode + sofia_cfg.mamba_variant = cfg.sofia_mamba_variant + sofia_cfg.mamba_backend = cfg.sofia_mamba_backend + sofia_cfg.use_evss_bridge = cfg.sofia_use_evss_bridge + sofia_cfg.evss_bridge_locations = list(cfg.sofia_evss_bridge_locations) + model = SOFIAFusionEncoder( + sofia_cfg=sofia_cfg, + lrsclip_path=cfg.lrsclip_path, + init_gate=cfg.init_gate, + baseline_mode=cfg.baseline_mode, + lora_rank=cfg.sofia_lora_rank, + device=cfg.device, + ).to(cfg.device) + elif cfg.backbone == "sofia_v1": + sofia_v1_cfg = SOFIAv1Config( + variant=cfg.sofia_v1_variant, + dcn_variant=cfg.sofia_v1_dcn_variant, + d_descriptor=cfg.sofia_v1_d_descriptor, + text_film_dim=cfg.sofia_v1_d_descriptor, + use_text_film_uav=cfg.sofia_v1_use_text_film_uav and not cfg.baseline_mode, + use_text_film_sat=cfg.sofia_v1_use_text_film_sat and not cfg.baseline_mode, + use_film_altitude=cfg.sofia_v1_use_film_altitude, + ) + model = SOFIAv1FusionEncoder( + sofia_cfg=sofia_v1_cfg, + lrsclip_path=cfg.lrsclip_path, + init_gate=cfg.init_gate, + baseline_mode=cfg.baseline_mode, + lora_rank=cfg.sofia_v1_lora_rank, + device=cfg.device, + ).to(cfg.device) + else: + model = AsymmetricEncoder( + dino_web_path=cfg.dino_web_path, + dino_sat_path=cfg.dino_sat_path, + lrsclip_path=cfg.lrsclip_path, + init_gate=cfg.init_gate, + baseline_mode=cfg.baseline_mode, + shared_encoder=cfg.shared_encoder, + mona_bottleneck=cfg.mona_bottleneck, + mona_last_n_blocks=cfg.mona_last_n_blocks, + device=cfg.device, + backbone=cfg.backbone, + stripnet_path=cfg.stripnet_path, + stripnet_mona_last_n_stages=cfg.stripnet_mona_last_n_stages, + stripnet_freeze=cfg.stripnet_freeze, + ).to(cfg.device) LOGGER.info("embed_dim=%d", model.embed_dim) # --- Gradient checkpointing (trade compute for VRAM) --- - # StripNet doesn't expose set_gradient_checkpointing — skip silently. + # StripNet/SOFIA don't expose set_gradient_checkpointing — only DGTRS gets it. if cfg.gradient_checkpointing and cfg.backbone == "dinov3": if cfg.shared_encoder: model.image_encoder.set_gradient_checkpointing(True) @@ -618,10 +711,10 @@ def train(cfg: TrainConfigGTAUAV) -> None: if model.text_encoder is not None: model.text_encoder.transformer.gradient_checkpointing = True LOGGER.info("Gradient checkpointing enabled (DINOv3 + DGTRS)") - elif cfg.gradient_checkpointing and cfg.backbone == "stripnet": + elif cfg.gradient_checkpointing and cfg.backbone in ("stripnet", "sofia", "sofia_v1"): if model.text_encoder is not None: model.text_encoder.transformer.gradient_checkpointing = True - LOGGER.info("Gradient checkpointing enabled (DGTRS only; StripNet doesn't support)") + LOGGER.info("Gradient checkpointing enabled (DGTRS only; %s doesn't support)", cfg.backbone) n_trainable = sum(p.numel() for p in model.trainable_parameters()) n_total = sum(p.numel() for p in model.parameters()) @@ -879,11 +972,14 @@ def train(cfg: TrainConfigGTAUAV) -> None: drone_img = batch["drone_img"].to(cfg.device, non_blocking=True) sat_img = batch["sat_img"].to(cfg.device, non_blocking=True) + altitude = batch.get("altitude") + if altitude is not None: + altitude = altitude.to(cfg.device, non_blocking=True) # Model forward in AMP (fp16 for DINOv3/DGTRS encoders). with autocast(device_type="cuda", enabled=cfg.use_amp): if cfg.baseline_mode: - embeddings = model(drone_img=drone_img, sat_img=sat_img) + embeddings = model(drone_img=drone_img, sat_img=sat_img, altitude=altitude) else: embeddings = model( drone_img=drone_img, @@ -894,6 +990,7 @@ def train(cfg: TrainConfigGTAUAV) -> None: sat_caption_l1=batch["sat_caption_l1"], sat_caption_l2=batch["sat_caption_l2"], sat_caption_l3=batch["sat_caption_l3"], + altitude=altitude, ) # Loss — InfoNCE or WeightedInfoNCE. Only the latter uses positive_weights. queue_neg = neg_bank.get_queue() if neg_bank is not None else None @@ -1103,20 +1200,23 @@ def train(cfg: TrainConfigGTAUAV) -> None: history.append(epoch_record) # Save checkpoint. Model architecture flags go into the ckpt so - # `AsymmetricEncoder.load_checkpoint` can rebuild the right shape. - _atomic_save( - obj={ - "epoch": epoch, - "model_state": model.state_dict(), - "optimizer_state": optimizer.state_dict(), - "loss_state": loss_fn.state_dict(), - "baseline_mode": cfg.baseline_mode, - "shared_encoder": cfg.shared_encoder, - "mona_bottleneck": cfg.mona_bottleneck, - "mona_last_n_blocks": cfg.mona_last_n_blocks, - }, - path=output_dir / f"ckpt_epoch{epoch:03d}.pt", - ) + # `AsymmetricEncoder.load_checkpoint` (or `SOFIAFusionEncoder.load_checkpoint`) + # can rebuild the right shape. + ckpt_obj = { + "epoch": epoch, + "model_state": model.state_dict(), + "optimizer_state": optimizer.state_dict(), + "loss_state": loss_fn.state_dict(), + "baseline_mode": cfg.baseline_mode, + "backbone": cfg.backbone, + } + if cfg.backbone in ("sofia", "sofia_v1"): + ckpt_obj["sofia_cfg"] = model.sofia_cfg + else: + ckpt_obj["shared_encoder"] = cfg.shared_encoder + ckpt_obj["mona_bottleneck"] = cfg.mona_bottleneck + ckpt_obj["mona_last_n_blocks"] = cfg.mona_last_n_blocks + _atomic_save(obj=ckpt_obj, path=output_dir / f"ckpt_epoch{epoch:03d}.pt") LOGGER.info("Checkpoint saved: ckpt_epoch%03d.pt", epoch) # Save history.