| """3D ResNet over the 8-channel atlas volume (4 MR types + 4 presence channels).""" |
| from __future__ import annotations |
|
|
| import math |
|
|
| import torch.nn as nn |
|
|
|
|
| def _norm(c: int) -> nn.Module: |
| |
| return nn.GroupNorm(math.gcd(c, 32), c) |
|
|
|
|
| class BasicBlock(nn.Module): |
| def __init__(self, cin: int, cout: int, stride: int = 1): |
| super().__init__() |
| self.conv1 = nn.Conv3d(cin, cout, 3, stride, 1, bias=False) |
| self.n1 = _norm(cout) |
| self.conv2 = nn.Conv3d(cout, cout, 3, 1, 1, bias=False) |
| self.n2 = _norm(cout) |
| self.act = nn.SiLU(inplace=True) |
| self.down = None |
| if stride != 1 or cin != cout: |
| self.down = nn.Sequential(nn.Conv3d(cin, cout, 1, stride, bias=False), _norm(cout)) |
|
|
| def forward(self, x): |
| idt = x if self.down is None else self.down(x) |
| x = self.act(self.n1(self.conv1(x))) |
| x = self.n2(self.conv2(x)) |
| return self.act(x + idt) |
|
|
|
|
| class ResNet3D(nn.Module): |
| """(B, 8, D, H, W) -> (B, num_classes) logits. Total stride 16 (stride-2 stem + three stride-2 stages).""" |
|
|
| def __init__(self, in_ch: int = 8, num_classes: int = 74, width: int = 48, |
| blocks: tuple[int, ...] = (2, 2, 2, 2), dropout: float = 0.2): |
| super().__init__() |
| w = width |
| self.stem = nn.Sequential( |
| nn.Conv3d(in_ch, w, 3, 2, 1, bias=False), _norm(w), nn.SiLU(inplace=True), |
| nn.Conv3d(w, w, 3, 1, 1, bias=False), _norm(w), nn.SiLU(inplace=True), |
| ) |
| stages, cin = [], w |
| for i, (cout, n) in enumerate(zip([w, w * 2, w * 4, w * 8], blocks)): |
| layers = [BasicBlock(cin, cout, stride=2 if i > 0 else 1)] |
| layers += [BasicBlock(cout, cout) for _ in range(n - 1)] |
| stages.append(nn.Sequential(*layers)) |
| cin = cout |
| self.stages = nn.Sequential(*stages) |
| self.head = nn.Sequential(nn.AdaptiveAvgPool3d(1), nn.Flatten(), nn.Dropout(dropout), |
| nn.Linear(cin, num_classes)) |
| for m in self.modules(): |
| if isinstance(m, nn.Conv3d): |
| nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") |
|
|
| def forward(self, x): |
| return self.head(self.stages(self.stem(x))) |
|
|
|
|
| ARCHS = { |
| "r18": dict(width=48, blocks=(2, 2, 2, 2)), |
| "r18w": dict(width=64, blocks=(2, 2, 2, 2)), |
| "r26w": dict(width=64, blocks=(2, 3, 4, 2)), |
| } |
|
|
|
|
| def build_model(arch: str = "r18", **kw) -> ResNet3D: |
| return ResNet3D(**ARCHS[arch], **kw) |
|
|
