Remove src/src/models/plant_vit_v2.py

src/src/models/plant_vit_v2.py

@@ -1,169 +0,0 @@
-"""
-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).
-3. Pure 16 SwiGLU Mixture-of-Experts (MoESwiGLUv2 with Gating Residuals & Context Projection).
-4. Sub-Center ArcFace Metric Head (K=3 Phenology Centers per Class).
-5. Hierarchical Multi-Task Taxonomic Classification Heads (Species + Genus + Family).
-"""
-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):
-        super().__init__()
-        self.dim = dim
-        self.out_dim = out_dim
-
-        self.norm_mid = nn.LayerNorm(dim)
-        self.norm_deep = nn.LayerNorm(dim)
-
-        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 forward(self, mid_tokens: torch.Tensor, deep_tokens: torch.Tensor) -> torch.Tensor:
-        # Pool spatial tokens over sequence dimension
-        mid_pooled = self.norm_mid(mid_tokens.mean(dim=1))
-        deep_pooled = self.norm_deep(deep_tokens.mean(dim=1))
-
-        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
-    ):
-        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
-
-        if freeze_stem:
-            for p in self.stem.parameters():
-                p.requires_grad = False
-
-        # 2. Pure 16-SwiGLU MoE Layer
-        if self.use_moe:
-            self.moe_layer = MoESwiGLUv2(
-                d_model=self.dim,
-                d_ff=d_ff,
-                num_ffn=num_ffn,
-                top_k=top_k,
-                gating_residual=True,
-                gate_ctx=True
-            )
-            self.moe_norm = nn.LayerNorm(self.dim)
-        else:
-            self.moe_layer = 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):
-        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)
-        gate_logits = None
-        if self.use_moe and self.moe_layer is not None:
-            ctx = deep_tokens[:, 0]
-            moe_out, aux_loss, gate_logits = self.moe_layer(self.moe_norm(deep_tokens), ctx=ctx)
-            deep_tokens = deep_tokens + moe_out
-
-        return mid_tokens, deep_tokens, aux_loss, gate_logits
-
-    def forward(self, x: torch.Tensor, y_species: torch.Tensor = None):
-        mid_tokens, deep_tokens, moe_aux, _ = self.forward_features(x)
-
-        # 1. Multi-Layer Pyramid Fusion (Block 8 Texture + Block 11 Macro)
-        emb = self.fusion(mid_tokens, deep_tokens)
-
-        # 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
-        arcface_logits = self.subcenter_arcface(emb, y_species)
-
-        aux_dict = {"moe_aux": moe_aux}
-
-        return (
-            logits_species,
-            logits_genus,
-            logits_family,
-            arcface_logits,
-            emb,
-            aux_dict
-        )