File size: 6,229 Bytes
3217f9a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | """PlantViT - thin adapter over RSNA MoR/MoE for single-image plant ID.
Reuses moe_vit.MoEViT and mor_vit.MoRViT frozen DINOv3 stem + routing.
Single image = single study with S=1 slice, so pool reduces to patchmean.
Adds 512-D embedder head per architecture.md §4.
Usage:
from src.models.plant_vit import PlantViT
model = PlantViT(stem_name="vit_base_patch16_dinov3", n_classes=500, use_moe=True, embed_dim=512)
logits, embed, aux = model(x) # x [B,3,336,336] -> logits [B,500], embed [B,512] L2
"""
from __future__ import annotations
import torch
import torch.nn as nn
import torch.nn.functional as F
from .moe_vit import MoEViT
from .mor_vit import MoRViT
class EmbedHead(nn.Module):
"""Plant-specific projection: 768/384 -> 512 L2 per architecture.md §4."""
def __init__(self, in_dim: int, embed_dim: int = 512):
super().__init__()
self.ln1 = nn.LayerNorm(in_dim)
self.proj = nn.Linear(in_dim, embed_dim, bias=False)
self.ln2 = nn.LayerNorm(embed_dim)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# x [B, D]
x = self.ln1(x)
x = self.proj(x)
x = self.ln2(x)
return F.normalize(x, dim=-1)
class PlantViT(nn.Module):
"""Wrapper that exposes (logits, embed, aux_losses) with interchangeable stem.
Args:
stem_name: timm name, e.g. vit_small_patch16_dinov3, vit_base_patch16_dinov3, vit_large_patch16_dinov3, convnext_tiny
n_classes: 500 for WA top500
embed_dim: 512 (test 256/128 later)
use_moe: True -> MoE++ (MoEViT), False -> MoR (MoRViT), None -> dense
freeze_stem: keep DINOv3 frozen for CE baseline
All MoE/MoR kwargs (num_ffn, top_k, d_ff, n_recursions, capacities etc.) passed through.
"""
def __init__(
self,
*,
stem_name: str = "vit_base_patch16_dinov3",
n_classes: int = 500,
embed_dim: int = 512,
use_moe: bool | None = True,
freeze_stem: bool = True,
# MoE knobs
num_ffn: int = 4,
top_k: int = 2,
d_ff: int = 384,
# MoR knobs
n_recursions: int = 3,
capacities: tuple[float, ...] = (1.0, 0.667, 0.333),
core_blocks: int = 2,
router_hidden: int = 128,
pretrained: bool = True,
**kwargs,
):
super().__init__()
self.use_moe = use_moe
# build core - keep RSNA defaults, just swap n_classes and stem
if use_moe is True:
self.core = MoEViT(
stem_name=stem_name, pretrained=pretrained, freeze_stem=freeze_stem,
core_blocks=core_blocks, d_ff=d_ff, num_ffn=num_ffn, top_k=top_k,
n_classes=n_classes, **{k: v for k, v in kwargs.items() if k in ["gate_ctx","per_target","tau","n_zero","n_copy","n_const","gating_residual"]},
)
feat_dim = self.core.core_dim # 384
elif use_moe is False:
self.core = MoRViT(
stem_name=stem_name, pretrained=pretrained, freeze_stem=freeze_stem,
n_recursions=n_recursions, capacities=capacities, core_blocks=core_blocks,
router_hidden=router_hidden, n_classes=n_classes,
)
feat_dim = self.core.core_dim
else:
# dense baseline - MoR dense=True
self.core = MoRViT(
stem_name=stem_name, pretrained=pretrained, freeze_stem=freeze_stem,
n_recursions=n_recursions, capacities=capacities, core_blocks=core_blocks,
router_hidden=router_hidden, n_classes=n_classes, dense=True,
)
feat_dim = self.core.core_dim
# embedder on pooled slice feat (before head pool, use same dim)
# For MoE/MoR, slice feat dim = core.core_dim (384) ; for base ViT-B 768 -> still 384 via input_proj
# If stem is vit_base (768) we still route through input_proj -> 384, so embed in 384
self.embed = EmbedHead(feat_dim, embed_dim)
self.n_classes = n_classes
self.embed_dim = embed_dim
self.stem_name = stem_name
def forward(self, x: torch.Tensor, step: int = 0):
"""x [B,3,H,W] single image per plant -> logits [B,C], embed [B,512], aux dict"""
# RSNA core expects [S,3,H,W] where S is slices per study; for plant S=B single slice per sample
# Use core.features single-batch path
feats = self.core.features(x) # [B, D] (MoEViT/MoRViT handle bfloat16 internally)
logits = self.core.head(self.core.pool(feats.unsqueeze(1).repeat(1,1,1)) if False else feats) # bypass buggy repeat - just pool slice feats
# Actually MoEViT/MoRViT head expects pooled study vector [D]; we have [B,D] already pooled per image
# Re-derive logits via head directly for single-slice case:
# MoEViT/MoRViT: features -> [S,D] -> pool -> [D] -> head -> [C]
# For B independent images, features is [B,D] already per-image slice feat, so head per row:
# Use core.head on each row individually via batched linear - head is Sequential Linear/GELU/Linear
# For per_target head, need [B, n_classes, D] - not used for plant (per_target=False by default)
# Simpler: call core.head on feats directly (feats is [B,D] -> head expects [D] but Linear broadcasts on last dim)
# Do batched:
if isinstance(self.core.head, nn.Sequential):
logits = self.core.head(feats) # [B,C] broadcasts
else:
# PerTarget not used
logits = self.core.head(feats)
embed = self.embed(feats)
aux = {}
# surface aux losses if present
try:
aux["moe_aux"] = float(self.core.moe_aux_loss().item()) if hasattr(self.core, "moe_aux_loss") else 0.0
except:
aux["moe_aux"] = 0.0
try:
aux["tv"] = float(self.core.routing_tv_loss().item()) if hasattr(self.core, "routing_tv_loss") else float(self.core.routing_smoothness_loss().item() if hasattr(self.core, "routing_smoothness_loss") else 0.0)
except:
aux["tv"] = 0.0
return logits, embed, aux
def features(self, x: torch.Tensor) -> torch.Tensor:
return self.core.features(x)
|