File size: 4,135 Bytes
12db9dc | 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 | """MRI-CORE stem: MRI-pretrained ViT-B feature extractor.
Loads the MRI-CORE checkpoint (DINOv2-style ViT-B, SAM-init, pretrained on
6.1M MRI slices) and exposes a timm-compatible surface so the existing
MoR/MoE cores can use it as their frozen stem:
stem.embed_dim -> 768
stem.num_prefix_tokens -> 1 (CLS only, no registers)
stem.pos_embed -> [1, 197, 768]
stem.forward_features(x)-> [S, 197, 768] (frozen, eval mode)
Key layout in the checkpoint: teacher.backbone.{cls_token, pos_embed,
patch_embed.proj, blocks.<g>.<i>, norm} with blocks nested 4 x 3 = 12.
"""
from __future__ import annotations
import os
import torch
import torch.nn as nn
import torch.nn.functional as F
_CKPT = "MRI_CORE_vitb.pth"
class Block(nn.Module):
def __init__(self, dim: int = 768, n_heads: int = 12, mlp_ratio: float = 4.0):
super().__init__()
self.norm1 = nn.LayerNorm(dim)
self.attn = nn.MultiheadAttention(dim, n_heads, batch_first=True)
self.norm2 = nn.LayerNorm(dim)
self.mlp = nn.Sequential(
nn.Linear(dim, int(dim * mlp_ratio)),
nn.GELU(),
nn.Linear(int(dim * mlp_ratio), dim),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = x + self.attn(self.norm1(x), self.norm1(x), self.norm1(x))[0]
x = x + self.mlp(self.norm2(x))
return x
class MRICoreStem(nn.Module):
"""Frozen MRI-CORE ViT-B feature extractor (embed_dim=768, 197 tokens)."""
def __init__(self, ckpt_path: str | None = None, freeze: bool = True):
super().__init__()
if ckpt_path is None:
for base in (
r"G:\RSNA-Knee\cache",
os.path.join(os.path.dirname(__file__), "..", "..", "weights"),
):
cand = os.path.join(base, _CKPT)
if os.path.isfile(cand):
ckpt_path = cand
break
if ckpt_path is None or not os.path.isfile(ckpt_path):
raise FileNotFoundError(f"MRI-CORE checkpoint not found ({_CKPT})")
self.embed_dim = 768
self.num_prefix_tokens = 1
self.patch_embed = nn.Conv2d(3, 768, kernel_size=16, stride=16, bias=True)
self.cls_token = nn.Parameter(torch.zeros(1, 1, 768))
self.pos_embed = nn.Parameter(torch.zeros(1, 197, 768))
self.mask_token = nn.Parameter(torch.zeros(1, 768))
self.blocks = nn.ModuleList(
[nn.ModuleList([Block() for _ in range(3)]) for _ in range(4)]
)
self.norm = nn.LayerNorm(768)
st = torch.load(ckpt_path, map_location="cpu", weights_only=True)
inner = st["teacher"]
if isinstance(inner, dict) and "state_dict" in inner:
inner = inner["state_dict"]
own = self.state_dict()
prefix = "backbone."
missing = []
for k in list(inner.keys()):
if not k.startswith(prefix):
inner.pop(k)
continue
name = k[len(prefix):]
if name not in own:
continue
if tuple(inner[k].shape) == tuple(own[name].shape):
own[name] = inner[k]
else:
missing.append((name, tuple(inner[k].shape), tuple(own[name].shape)))
self.load_state_dict(own)
if missing:
print(f"mricore: {len(missing)} shape mismatches skipped "
f"(e.g. {missing[0]})", flush=True)
if freeze:
for p in self.parameters():
p.requires_grad_(False)
self.eval()
def forward_features(self, x: torch.Tensor) -> torch.Tensor:
"""[S, 3, H, W] -> [S, 197, 768] (CLS + 196 patch tokens)."""
x = self.patch_embed(x).flatten(2).transpose(1, 2) # [S, 196, 768]
x = torch.cat([self.cls_token.expand(x.size(0), -1, -1), x], dim=1)
x = x + self.pos_embed
for group in self.blocks:
for blk in group:
x = blk(x)
return self.norm(x)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.forward_features(x)
|