File size: 2,768 Bytes
69928d6 | 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 | """Reusable image encoder; AS2 itself does not require an image modality."""
import torch
from torch import nn
class CellCNN(nn.Module):
"""Shared-weight CNN that encodes a single 28×28 grayscale cell image.
Args:
d_model: Output embedding dimension.
backbone: Architecture variant.
``"default"`` — lightweight 3-conv CNN (original SoftASP design).
``"resnet18"`` — ResNet-18 adapted for 1-channel 28×28 input;
weights are loaded externally via ``pretrained_from`` or trained
from scratch. Requires ``torchvision``.
"""
def __init__(self, d_model: int = 128, backbone: str = "default") -> None:
super().__init__()
self.backbone_name = backbone
if backbone == "default":
self.features = nn.Sequential(
nn.Conv2d(1, 32, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(2), # → 32 × 14 × 14
nn.Conv2d(32, 64, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.MaxPool2d(2), # → 64 × 7 × 7
nn.Conv2d(64, 128, kernel_size=3, padding=1),
nn.ReLU(inplace=True),
nn.AdaptiveAvgPool2d(1), # → 128 × 1 × 1
nn.Flatten(), # → 128
)
self.proj = nn.Linear(128, d_model)
elif backbone == "resnet18":
import torchvision.models as _tv # lazy import — optional dep
_r = _tv.resnet18(weights=None)
# Adapt for 1-channel 28×28 input:
# • smaller first conv (3×3, stride 1) preserves 28→28 spatial size
# • remove the 3×3 max-pool so small images keep enough resolution
_r.conv1 = nn.Conv2d(1, 64, kernel_size=3, stride=1, padding=1, bias=False)
# Use all residual stages + avgpool → 512-d feature vector
self.features = nn.Sequential(
_r.conv1,
_r.bn1,
_r.relu,
nn.Identity(),
_r.layer1,
_r.layer2,
_r.layer3,
_r.layer4,
_r.avgpool,
nn.Flatten(), # → 512
)
self.proj = nn.Linear(512, d_model)
else:
raise ValueError(
f"Unknown CellCNN backbone: {backbone!r}. Choose 'default' or 'resnet18'."
)
self.norm = nn.LayerNorm(d_model)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
Args:
x: (N, 1, 28, 28) float32
Returns:
(N, d_model)
"""
h = self.features(x)
return self.norm(self.proj(h))
|