Release AS2 framework with Hydra-selected instantiations

69928d6 verified 9 days ago

2.77 kB

"""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))

	"""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))