| """
|
| Vbai-2.6AD Model
|
| ================
|
| Multimodal Alzheimer's classifier with REAL MRI<->biomarker pairing.
|
|
|
| Streams:
|
| * MRI encoder : 3D ResNet (CBAM/SE) + ASPP → 512-d
|
| * Tabular encoder: MLP on (values + missing-mask) → 256-d
|
| * Fusion : bidirectional cross-attention + gated combine → 512-d
|
|
|
| Heads:
|
| * mri_logits : Stage-1 MRI-only prediction
|
| * tab_logits : Tabular-only prediction (used as auxiliary)
|
| * fused_logits : Final 3-way classification (CN/MCI/AD)
|
| * progression : will_progress (sigmoid), time_to_conversion (months),
|
| time_distribution (24 bins, 5-month resolution)
|
|
|
| Training-time tricks (in dataset/loss, not here):
|
| * modality dropout
|
| * per-feature random masking
|
| * cross-modal contrastive loss
|
| """
|
| from __future__ import annotations
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
|
|
| import config as C
|
|
|
|
|
|
|
|
|
|
|
| class ChannelAttention3D(nn.Module):
|
| def __init__(self, ch, r=16):
|
| super().__init__()
|
| m = max(ch // r, 8)
|
| self.mlp = nn.Sequential(nn.Linear(ch, m), nn.ReLU(inplace=True), nn.Linear(m, ch))
|
|
|
| def forward(self, x):
|
| a = x.mean(dim=[2, 3, 4]); b = x.amax(dim=[2, 3, 4])
|
| attn = torch.sigmoid(self.mlp(a) + self.mlp(b))
|
| return x * attn[..., None, None, None]
|
|
|
|
|
| class SpatialAttention3D(nn.Module):
|
| def __init__(self, k=7):
|
| super().__init__()
|
| self.conv = nn.Conv3d(2, 1, k, padding=k // 2, bias=False)
|
|
|
| def forward(self, x):
|
| avg = x.mean(dim=1, keepdim=True); mx = x.amax(dim=1, keepdim=True)
|
| attn = torch.sigmoid(self.conv(torch.cat([avg, mx], dim=1)))
|
| return x * attn
|
|
|
|
|
| class CBAM3D(nn.Module):
|
| def __init__(self, ch, r=16):
|
| super().__init__()
|
| self.c = ChannelAttention3D(ch, r); self.s = SpatialAttention3D()
|
| def forward(self, x): return self.s(self.c(x))
|
|
|
|
|
| class SEBlock3D(nn.Module):
|
| def __init__(self, ch, r=16):
|
| super().__init__()
|
| m = max(ch // r, 8)
|
| self.fc = nn.Sequential(nn.Linear(ch, m), nn.ReLU(True), nn.Linear(m, ch), nn.Sigmoid())
|
| def forward(self, x):
|
| s = x.mean(dim=[2, 3, 4]); s = self.fc(s)[..., None, None, None]
|
| return x * s
|
|
|
|
|
|
|
|
|
|
|
| class ResBlock3D(nn.Module):
|
| def __init__(self, in_ch, out_ch, stride=1, use_cbam=True, use_se=True, drop_path=0.0):
|
| super().__init__()
|
| self.conv1 = nn.Conv3d(in_ch, out_ch, 3, stride, 1, bias=False)
|
| self.bn1 = nn.BatchNorm3d(out_ch)
|
| self.conv2 = nn.Conv3d(out_ch, out_ch, 3, 1, 1, bias=False)
|
| self.bn2 = nn.BatchNorm3d(out_ch)
|
| self.act = nn.GELU()
|
| self.cbam = CBAM3D(out_ch) if use_cbam else nn.Identity()
|
| self.se = SEBlock3D(out_ch) if use_se else nn.Identity()
|
| self.drop_path = drop_path
|
| self.skip = nn.Identity() if (in_ch == out_ch and stride == 1) else nn.Sequential(
|
| nn.Conv3d(in_ch, out_ch, 1, stride, bias=False), nn.BatchNorm3d(out_ch))
|
|
|
| def _stochastic(self, x):
|
| if not self.training or self.drop_path == 0.0:
|
| return x
|
| keep = 1.0 - self.drop_path
|
| mask = torch.empty(x.shape[0], 1, 1, 1, 1, device=x.device).bernoulli_(keep)
|
| return x * mask / keep
|
|
|
| def forward(self, x):
|
| identity = self.skip(x)
|
| out = self.act(self.bn1(self.conv1(x)))
|
| out = self.bn2(self.conv2(out))
|
| out = self.cbam(out); out = self.se(out)
|
| out = self._stochastic(out)
|
| return self.act(out + identity)
|
|
|
|
|
| class ASPP3D(nn.Module):
|
| def __init__(self, in_ch, out_ch, dilations=(1, 6, 12, 18)):
|
| super().__init__()
|
| per = out_ch // len(dilations)
|
| self.branches = nn.ModuleList([
|
| nn.Sequential(nn.Conv3d(in_ch, per, 3, padding=d, dilation=d, bias=False),
|
| nn.BatchNorm3d(per), nn.GELU())
|
| for d in dilations
|
| ])
|
| self.gp = nn.Sequential(
|
| nn.AdaptiveAvgPool3d(1),
|
| nn.Conv3d(in_ch, per, 1, bias=False),
|
| nn.BatchNorm3d(per), nn.GELU())
|
| self.fuse = nn.Sequential(nn.Conv3d(per * (len(dilations) + 1), out_ch, 1, bias=False),
|
| nn.BatchNorm3d(out_ch), nn.GELU())
|
|
|
| def forward(self, x):
|
| feats = [b(x) for b in self.branches]
|
| g = self.gp(x)
|
| g = F.interpolate(g, size=x.shape[2:], mode="trilinear", align_corners=False)
|
| feats.append(g)
|
| return self.fuse(torch.cat(feats, dim=1))
|
|
|
|
|
|
|
|
|
|
|
| class MRIEncoder3D(nn.Module):
|
| def __init__(self, mcfg: C.ModelConfig):
|
| super().__init__()
|
| ch = mcfg.mri_encoder_channels
|
| self.stem = nn.Sequential(
|
| nn.Conv3d(1, ch[0], 7, 2, 3, bias=False), nn.BatchNorm3d(ch[0]), nn.GELU(),
|
| nn.MaxPool3d(3, 2, 1))
|
| depths = [2, 2, 2, 2]
|
| dp = [0.0, 0.05, 0.1, 0.15]
|
| self.stage1 = self._make(ch[0], ch[0], depths[0], 1, mcfg, dp[0])
|
| self.stage2 = self._make(ch[0], ch[1], depths[1], 2, mcfg, dp[1])
|
| self.stage3 = self._make(ch[1], ch[2], depths[2], 2, mcfg, dp[2])
|
| self.stage4 = self._make(ch[2], ch[3], depths[3], 2, mcfg, dp[3])
|
| self.aspp = ASPP3D(ch[3], mcfg.mri_bottleneck_channels)
|
| self.pool = nn.AdaptiveAvgPool3d(1)
|
| self.proj = nn.Sequential(
|
| nn.Linear(mcfg.mri_bottleneck_channels, mcfg.mri_feature_dim),
|
| nn.GELU(), nn.Dropout(mcfg.mri_dropout))
|
|
|
| def _make(self, in_ch, out_ch, n, stride, mcfg, dp):
|
| layers = [ResBlock3D(in_ch, out_ch, stride, mcfg.use_cbam, mcfg.use_se_block, dp)]
|
| for _ in range(1, n):
|
| layers.append(ResBlock3D(out_ch, out_ch, 1, mcfg.use_cbam, mcfg.use_se_block, dp))
|
| return nn.Sequential(*layers)
|
|
|
| def forward(self, x):
|
| x = self.stem(x)
|
| x = self.stage1(x); x = self.stage2(x); x = self.stage3(x); x = self.stage4(x)
|
| x = self.aspp(x); x = self.pool(x).flatten(1)
|
| return self.proj(x)
|
|
|
|
|
|
|
|
|
|
|
| class TabularEncoder(nn.Module):
|
| def __init__(self, mcfg: C.ModelConfig):
|
| super().__init__()
|
| prev = mcfg.num_tabular_inputs
|
| layers = []
|
| for h in mcfg.tabular_hidden_dims:
|
| layers += [nn.Linear(prev, h), nn.LayerNorm(h), nn.GELU(), nn.Dropout(mcfg.tabular_dropout)]
|
| prev = h
|
| layers += [nn.Linear(prev, mcfg.tabular_feature_dim)]
|
| self.net = nn.Sequential(*layers)
|
|
|
| def forward(self, x):
|
| return self.net(x)
|
|
|
|
|
|
|
|
|
|
|
| class CrossModalFusion(nn.Module):
|
| def __init__(self, mri_dim, tab_dim, fdim, heads=8, dropout=0.1):
|
| super().__init__()
|
| self.pm = nn.Linear(mri_dim, fdim); self.pt = nn.Linear(tab_dim, fdim)
|
| self.a_mt = nn.MultiheadAttention(fdim, heads, dropout=dropout, batch_first=True)
|
| self.a_tm = nn.MultiheadAttention(fdim, heads, dropout=dropout, batch_first=True)
|
| self.lnm = nn.LayerNorm(fdim); self.lnt = nn.LayerNorm(fdim)
|
| self.gate = nn.Sequential(nn.Linear(fdim * 2, fdim), nn.Sigmoid())
|
| self.out = nn.Sequential(nn.Linear(fdim * 2, fdim), nn.GELU(), nn.Dropout(dropout))
|
|
|
| def forward(self, m, t):
|
| m1 = self.pm(m).unsqueeze(1); t1 = self.pt(t).unsqueeze(1)
|
| ma, _ = self.a_mt(m1, t1, t1); ta, _ = self.a_tm(t1, m1, m1)
|
| m2 = self.lnm(m1 + ma).squeeze(1); t2 = self.lnt(t1 + ta).squeeze(1)
|
| cat = torch.cat([m2, t2], dim=-1)
|
| g = self.gate(cat); o = self.out(cat)
|
| return g * m2 + (1 - g) * t2 + o
|
|
|
|
|
|
|
|
|
|
|
| class ClsHead(nn.Module):
|
| def __init__(self, in_dim, num_classes, dropout=0.3):
|
| super().__init__()
|
| self.h = nn.Sequential(
|
| nn.Linear(in_dim, 256), nn.GELU(), nn.Dropout(dropout),
|
| nn.Linear(256, 128), nn.GELU(), nn.Dropout(dropout),
|
| nn.Linear(128, num_classes))
|
| def forward(self, x): return self.h(x)
|
|
|
|
|
| class ProgressionHead(nn.Module):
|
| def __init__(self, in_dim, hidden=256, max_months=120, n_bins=24):
|
| super().__init__()
|
| self.max_months = float(max_months); self.n_bins = n_bins
|
| self.shared = nn.Sequential(nn.Linear(in_dim, hidden), nn.GELU(), nn.Dropout(0.3))
|
| self.binary = nn.Linear(hidden, 1)
|
| self.time = nn.Sequential(nn.Linear(hidden, 64), nn.GELU(), nn.Linear(64, 1))
|
| self.dist = nn.Linear(hidden, n_bins)
|
|
|
| def forward(self, x):
|
| h = self.shared(x)
|
| logits = self.binary(h).squeeze(-1)
|
| return {
|
| "will_progress_logits": logits,
|
| "will_progress": torch.sigmoid(logits),
|
| "time_to_conversion": torch.clamp(F.softplus(self.time(h)).squeeze(-1),
|
| min=0.0, max=self.max_months),
|
| "time_distribution": F.softmax(self.dist(h), dim=-1),
|
| }
|
|
|
|
|
|
|
|
|
|
|
| class Vbai26ADModel(nn.Module):
|
| def __init__(self, mcfg: C.ModelConfig | None = None):
|
| super().__init__()
|
| self.cfg = mcfg or C.ModelConfig()
|
| self.mri_encoder = MRIEncoder3D(self.cfg)
|
| self.tab_encoder = TabularEncoder(self.cfg)
|
|
|
| self.mri_classifier = ClsHead(self.cfg.mri_feature_dim, self.cfg.num_classes, self.cfg.mri_dropout)
|
| self.tab_classifier = ClsHead(self.cfg.tabular_feature_dim, self.cfg.num_classes, self.cfg.tabular_dropout)
|
|
|
| self.fusion = CrossModalFusion(
|
| self.cfg.mri_feature_dim, self.cfg.tabular_feature_dim,
|
| self.cfg.fusion_dim, self.cfg.fusion_num_heads, self.cfg.fusion_dropout)
|
| self.fused_classifier = ClsHead(self.cfg.fusion_dim, self.cfg.num_classes, self.cfg.fusion_dropout)
|
|
|
| self.progression_head = ProgressionHead(
|
| self.cfg.fusion_dim, self.cfg.progression_hidden_dim,
|
| self.cfg.max_progression_months, self.cfg.num_time_bins)
|
|
|
|
|
| self.contrast_mri = nn.Sequential(nn.Linear(self.cfg.mri_feature_dim, 128))
|
| self.contrast_tab = nn.Sequential(nn.Linear(self.cfg.tabular_feature_dim, 128))
|
|
|
| def forward(self, mri=None, tab=None):
|
| out = {}
|
| m_feat = t_feat = None
|
| if mri is not None:
|
| m_feat = self.mri_encoder(mri)
|
| out["mri_features"] = m_feat
|
| out["mri_logits"] = self.mri_classifier(m_feat)
|
| if tab is not None:
|
| t_feat = self.tab_encoder(tab)
|
| out["tab_features"] = t_feat
|
| out["tab_logits"] = self.tab_classifier(t_feat)
|
| if m_feat is not None and t_feat is not None:
|
| f = self.fusion(m_feat, t_feat)
|
| out["fused_features"] = f
|
| out["fused_logits"] = self.fused_classifier(f)
|
| out["progression"] = self.progression_head(f)
|
|
|
| out["zm"] = F.normalize(self.contrast_mri(m_feat), dim=-1)
|
| out["zt"] = F.normalize(self.contrast_tab(t_feat), dim=-1)
|
| elif m_feat is not None:
|
| out["fused_logits"] = out["mri_logits"]
|
| elif t_feat is not None:
|
| out["fused_logits"] = out["tab_logits"]
|
| return out
|
|
|
| def get_param_groups(self, lr_backbone, lr_fusion):
|
| backbone = list(self.mri_encoder.parameters()) + list(self.tab_encoder.parameters())
|
| fusion = (list(self.fusion.parameters()) + list(self.fused_classifier.parameters())
|
| + list(self.progression_head.parameters())
|
| + list(self.mri_classifier.parameters())
|
| + list(self.tab_classifier.parameters())
|
| + list(self.contrast_mri.parameters()) + list(self.contrast_tab.parameters()))
|
| return [{"params": backbone, "lr": lr_backbone},
|
| {"params": fusion, "lr": lr_fusion}]
|
|
|
| @torch.no_grad()
|
| def predict(self, mri=None, tab=None):
|
| self.eval()
|
| out = self.forward(mri=mri, tab=tab)
|
| probs = F.softmax(out["fused_logits"], dim=-1)
|
| pred = probs.argmax(dim=-1)
|
| result = {"pred_class": pred, "class_probs": probs,
|
| "class_names": [C.CLASS_NAMES[c] for c in pred.cpu().tolist()]}
|
| if "progression" in out:
|
| p = out["progression"]
|
| result["will_progress"] = p["will_progress"]
|
| result["time_to_conversion_months"] = p["time_to_conversion"]
|
| result["time_distribution"] = p["time_distribution"]
|
| return result
|
|
|
|
|
| def count_params(model):
|
| return sum(p.numel() for p in model.parameters() if p.requires_grad)
|
|
|
