Update src\models\plant_vit_v2.py

src/models/plant_vit_v2.py

@@ -0,0 +1,295 @@
+"""
+PlantViT v2 Architecture (PlantDetect v2.0):
+1. DINOv3 ViT-Base Backbone @ 448x448 High-Res (784 patch tokens).
+2. Multi-Layer Feature Pyramid Fusion (Block 8 Micro-Texture + Block 11 Macro-Semantics) — CLS excluded.
+3. Pure SwiGLU Mixture-of-Experts (MoESwiGLUv2 with Gating Residuals & Context Projection).
+   - v1 heterogeneous: 16 FFN + 1 zero + 1 copy + 2 const = 20 total (moe_vit.py)
+   - v2 pure: 16 SwiGLU experts (d_ff=512) configurable via --num-ffn 8/16/20, stacked 1-2 layers
+4. Sub-Center ArcFace Metric Head (K=3 Phenology Centers per Class) with diversity regularizer.
+5. Hierarchical Multi-Task Taxonomic Classification Heads (Species + Genus + Family) — Unknown_Family fallback.
+"""
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import timm
+
+from .moe_swiglu_v2 import MoESwiGLUv2
+from .subcenter_arcface import SubCenterArcFace
+
+class MultiLayerFeatureFusion(nn.Module):
+    def __init__(self, dim: int = 768, out_dim: int = 512, use_attn_pool: bool = True):
+        super().__init__()
+        self.dim = dim
+        self.out_dim = out_dim
+        self.use_attn_pool = use_attn_pool
+
+        self.norm_mid = nn.LayerNorm(dim)
+        self.norm_deep = nn.LayerNorm(dim)
+
+        # Attention pooling (learned query) retains spatial signal vs pure mean — P1 fix
+        if use_attn_pool:
+            self.pool_q_mid = nn.Parameter(torch.randn(dim) * (dim ** -0.5))
+            self.pool_q_deep = nn.Parameter(torch.randn(dim) * (dim ** -0.5))
+            self.scale = dim ** -0.5
+        else:
+            self.pool_q_mid = None
+            self.pool_q_deep = None
+
+        self.fusion_proj = nn.Sequential(
+            nn.Linear(dim * 2, out_dim),
+            nn.GELU(),
+            nn.LayerNorm(out_dim),
+            nn.Linear(out_dim, out_dim),
+            nn.LayerNorm(out_dim)
+        )
+
+    def _pool(self, tokens: torch.Tensor, q: torch.Tensor | None) -> torch.Tensor:
+        # tokens [B,T,D], q [D]
+        if q is not None and self.use_attn_pool:
+            # Attentive pooling: w = softmax(tokens @ q / sqrt(D))
+            attn = (tokens @ q) * self.scale  # [B,T]
+            w = torch.softmax(attn, dim=1)  # [B,T]
+            pooled = (tokens * w.unsqueeze(-1)).sum(dim=1)  # [B,D]
+            return pooled
+        return tokens.mean(dim=1)
+
+    def forward(self, mid_tokens: torch.Tensor, deep_tokens: torch.Tensor) -> torch.Tensor:
+        # mid_tokens / deep_tokens are PATCH tokens only (CLS excluded by caller).
+        mid_pooled = self.norm_mid(self._pool(mid_tokens, self.pool_q_mid))
+        deep_pooled = self.norm_deep(self._pool(deep_tokens, self.pool_q_deep))
+
+        fused = torch.cat([mid_pooled, deep_pooled], dim=-1)
+        emb = self.fusion_proj(fused)
+        return F.normalize(emb, dim=-1)
+
+class PlantViTv2(nn.Module):
+    def __init__(
+        self,
+        stem_name: str = "vit_base_patch16_dinov3.lvd1689m",
+        img_size: int = 448,
+        n_species: int = 1000,
+        n_genera: int = 358,
+        n_families: int = 80,
+        embed_dim: int = 512,
+        use_moe: bool = True,
+        num_ffn: int = 16,
+        d_ff: int = 512,
+        top_k: int = 2,
+        k_subcenters: int = 3,
+        pretrained: bool = True,
+        freeze_stem: bool = False,
+        num_moe_layers: int = 1,
+        enable_grad_checkpoint: bool = False,
+        shared_experts: int = 0,
+        granularity_factor: int = 1,
+    ):
+        super().__init__()
+        self.stem_name = stem_name
+        self.img_size = img_size
+        self.n_species = n_species
+        self.n_genera = n_genera
+        self.n_families = n_families
+        self.embed_dim = embed_dim
+        self.use_moe = use_moe
+
+        # 1. DINOv3 Backbone
+        stem_model_name = stem_name if stem_name.startswith("vit_") else f"{stem_name}.lvd1689m"
+        self.stem = timm.create_model(
+            stem_model_name,
+            pretrained=pretrained,
+            img_size=img_size,
+            num_classes=0,
+            dynamic_img_size=True
+        )
+        self.dim = self.stem.embed_dim
+        # Number of prefix tokens (CLS + registers) to exclude from patch averaging.
+        # v1 correctly excluded via n_prefix; v2 previously averaged CLS into fusion.
+        self.n_prefix = int(getattr(self.stem, "num_prefix_tokens", 1))
+
+        if freeze_stem:
+            for p in self.stem.parameters():
+                p.requires_grad = False
+
+        # Real grad checkpointing for 448px (784+5 tok vs 441 at 336). timm DINOv3 has no set_grad_checkpointing,
+        # so we checkpoint forward_intermediates via torch.utils.checkpoint when training.
+        self.enable_grad_checkpoint = bool(enable_grad_checkpoint)
+        # Try timm native if available (harmless), but our manual checkpoint below is the real saving.
+        if self.enable_grad_checkpoint and hasattr(self.stem, "set_grad_checkpointing"):
+            try:
+                self.stem.set_grad_checkpointing(True)
+            except Exception:
+                pass
+
+        # Granularity: DeepSeekMoE-style fine split — keep FLOPs ~const (num*d_ff)
+        # factor=2: 16×512 → 32×256, factor=4: 64×128 — cheap sweep for 3k species
+        self.shared_experts = int(shared_experts)
+        self.granularity_factor = int(granularity_factor) if granularity_factor and granularity_factor >= 1 else 1
+        eff_num_ffn = num_ffn * self.granularity_factor
+        eff_d_ff = max(64, d_ff // self.granularity_factor)  # clamp to keep matrix sane
+        if self.granularity_factor > 1:
+            print(
+                f"[PlantViTv2] granularity x{self.granularity_factor}: {num_ffn}x{d_ff} -> {eff_num_ffn}x{eff_d_ff} (FLOPs ~const)",
+                flush=True,
+            )
+
+        # 2. Pure SwiGLU MoE Layer(s) — configurable stacking restores depth lost from v1's 2 blocks
+        self.num_moe_layers = int(num_moe_layers) if use_moe else 0
+        self.eff_num_ffn = eff_num_ffn
+        self.eff_d_ff = eff_d_ff
+        if self.use_moe and self.num_moe_layers > 0:
+            self.moe_layers = nn.ModuleList(
+                [
+                    MoESwiGLUv2(
+                        d_model=self.dim,
+                        d_ff=eff_d_ff,
+                        num_ffn=eff_num_ffn,
+                        top_k=top_k,
+                        gating_residual=True,
+                        gate_ctx=True,
+                        shared_experts=self.shared_experts,
+                    )
+                    for _ in range(self.num_moe_layers)
+                ]
+            )
+            self.moe_norms = nn.ModuleList([nn.LayerNorm(self.dim) for _ in range(self.num_moe_layers)])
+            # Backward compat: single-layer alias used by older checkpoints
+            self.moe_layer = self.moe_layers[0] if self.num_moe_layers == 1 else None
+            self.moe_norm = self.moe_norms[0] if self.num_moe_layers == 1 else None
+        else:
+            self.moe_layers = None
+            self.moe_norms = None
+            self.moe_layer = None
+            self.moe_norm = None
+
+        # 3. Multi-Layer Feature Fusion (Block 8 + Block 11)
+        self.fusion = MultiLayerFeatureFusion(dim=self.dim, out_dim=embed_dim)
+
+        # 4. Multi-Task Taxonomic Heads
+        self.head_species = nn.Linear(embed_dim, n_species, bias=False)
+        self.head_genus = nn.Linear(embed_dim, n_genera, bias=False)
+        self.head_family = nn.Linear(embed_dim, n_families, bias=False)
+
+        # 5. Sub-Center ArcFace Metric Head (K=3 Sub-Centers)
+        self.subcenter_arcface = SubCenterArcFace(
+            embed_dim=embed_dim,
+            n_classes=n_species,
+            k_subcenters=k_subcenters,
+            m=0.5,
+            s=30.0
+        )
+
+    def _to_tokens(self, feat: torch.Tensor) -> torch.Tensor:
+        """Converts [B, C, H, W] or [B, T, C] to standard [B, T, C]."""
+        if feat.dim() == 4:
+            return feat.flatten(2).transpose(1, 2)
+        return feat
+
+    def forward_features(self, x: torch.Tensor):
+        # Manual checkpointing for DINOv3 (timm has no set_grad_checkpointing)
+        if self.enable_grad_checkpoint and self.training:
+            # checkpoint expects pure function; wrap forward_intermediates
+            def _run_stem(inp):
+                return self.stem.forward_intermediates(inp, indices=[7, 11])
+
+            intermediates, final_feat = torch.utils.checkpoint.checkpoint(
+                _run_stem, x, use_reentrant=False
+            )
+        else:
+            intermediates, final_feat = self.stem.forward_intermediates(x, indices=[7, 11])
+
+        if isinstance(intermediates, (list, tuple)):
+            mid_raw = intermediates[0]
+            deep_raw = intermediates[1] if len(intermediates) > 1 else final_feat
+        else:
+            mid_raw = intermediates
+            deep_raw = final_feat
+
+        if isinstance(mid_raw, (list, tuple)): mid_raw = mid_raw[0]
+        if isinstance(deep_raw, (list, tuple)): deep_raw = deep_raw[0]
+
+        mid_tokens = self._to_tokens(mid_raw)
+        deep_tokens = self._to_tokens(deep_raw)
+
+        aux_loss = torch.tensor(0.0, device=x.device)
+        intra_loss = torch.tensor(0.0, device=x.device)
+        cross_loss = torch.tensor(0.0, device=x.device)
+        gate_logits = None
+        gate_weights_list: list[torch.Tensor] = []
+        if self.use_moe and self.moe_layers is not None and len(self.moe_layers) > 0:
+            prev_gate = None
+            total_aux = 0.0
+            total_intra = 0.0
+            for norm, moe in zip(self.moe_norms, self.moe_layers):
+                ctx = deep_tokens[:, 0]  # CLS as global context
+                ret = moe(norm(deep_tokens), prev_gate=prev_gate, ctx=ctx)
+                # Backward compat: MoESwiGLUv2 now returns 5 (with intra + gate_weights), old 3
+                if len(ret) == 5:
+                    moe_out, aux, gate_logits, intra, gate_weights = ret
+                elif len(ret) == 4:
+                    moe_out, aux, gate_logits, intra = ret
+                    gate_weights = torch.softmax(gate_logits, dim=-1)
+                else:
+                    moe_out, aux, gate_logits = ret
+                    intra = torch.tensor(0.0, device=x.device)
+                    gate_weights = torch.softmax(gate_logits, dim=-1)
+                deep_tokens = deep_tokens + moe_out
+                total_aux = total_aux + aux
+                total_intra = total_intra + intra
+                gate_weights_list.append(gate_weights)
+                prev_gate = gate_logits
+            aux_loss = total_aux / max(1, len(self.moe_layers))
+            intra_loss = total_intra / max(1, len(self.moe_layers))
+            # Cross-layer joint routing coherence (Feb 2026): maximize joint Top-k confidence across depth
+            # For tokens, joint = max_prob(L) * max_prob(L+1); loss = -log(joint) encourages coherent pathways
+            if len(gate_weights_list) >= 2:
+                cross_terms = []
+                for a, b in zip(gate_weights_list[:-1], gate_weights_list[1:]):
+                    # a,b [B,T,E] — per-token max prob
+                    max_a = a.max(dim=-1).values.clamp(min=1e-6)  # [B,T]
+                    max_b = b.max(dim=-1).values.clamp(min=1e-6)
+                    joint = (max_a * max_b).clamp(min=1e-6)
+                    # Negative log likelihood of joint
+                    cross_terms.append(-torch.log(joint).mean())
+                cross_loss = torch.stack(cross_terms).mean() if cross_terms else torch.tensor(0.0, device=x.device)
+            # Backward compat alias
+            if self.num_moe_layers == 1:
+                self.moe_layer = self.moe_layers[0]
+                self.moe_norm = self.moe_norms[0]
+
+        # Return intra + cross for trainer (add to aux_dict)
+        return mid_tokens, deep_tokens, aux_loss, gate_logits, intra_loss, cross_loss
+
+    def forward(self, x: torch.Tensor, y_species: torch.Tensor = None):
+        mid_tokens, deep_tokens, moe_aux, _, intra_loss, cross_loss = self.forward_features(x)
+
+        # 1. Multi-Layer Pyramid Fusion (Block 8 Texture + Block 11 Macro)
+        # Exclude CLS/registers from spatial averaging (fixes v2 bug where CLS diluted patch stats).
+        mid_patch = mid_tokens[:, self.n_prefix :] if mid_tokens.size(1) > self.n_prefix else mid_tokens
+        deep_patch = deep_tokens[:, self.n_prefix :] if deep_tokens.size(1) > self.n_prefix else deep_tokens
+        emb = self.fusion(mid_patch, deep_patch)
+
+        # 2. Multi-Task Linear Logits
+        logits_species = self.head_species(emb)
+        logits_genus = self.head_genus(emb)
+        logits_family = self.head_family(emb)
+
+        # 3. Sub-Center ArcFace Margin Logits — BF16 autocast (s=30) overflows in bf16, force FP32
+        # Keep BF16 for backbone/fusion (fast), drop to FP32 only here.
+        if emb.is_cuda:
+            with torch.amp.autocast("cuda", enabled=False):
+                arcface_logits = self.subcenter_arcface(emb.float(), y_species)
+                div_loss = self.subcenter_arcface.diversity_loss()
+        else:
+            arcface_logits = self.subcenter_arcface(emb.float(), y_species)
+            div_loss = self.subcenter_arcface.diversity_loss()
+
+        aux_dict = {"moe_aux": moe_aux, "div_loss": div_loss, "intra_loss": intra_loss, "cross_loss": cross_loss}
+
+        return (
+            logits_species,
+            logits_genus,
+            logits_family,
+            arcface_logits,
+            emb,
+            aux_dict
+        )