Initial commit

spectre/models/eomt.py

@@ -1,230 +0,0 @@
-# Adapted from https://github.com/tue-mps/eomt/
-from __future__ import annotations
-
-import math
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from spectre.models.layers import LayerNorm3d
-
-
-class ScaleBlock(nn.Module):
-    def __init__(
-        self, 
-        embed_dim: int, 
-        scale_factors: Union[int, Tuple[int, int, int]] = (2, 2, 2),
-        conv1_layer: nn.Module = nn.ConvTranspose3d,
-    ):
-        super().__init__()
-
-        self.conv1 = conv1_layer(
-            embed_dim,
-            embed_dim,
-            kernel_size=scale_factors,
-            stride=scale_factors,
-        )
-        self.act = nn.GELU()
-        self.conv2 = nn.Conv3d(
-            embed_dim,
-            embed_dim,
-            kernel_size=3,
-            padding=1,
-            groups=embed_dim,
-            bias=False,
-        )
-        self.norm = LayerNorm3d(embed_dim)
-
-    def forward(self, x):
-        x = self.conv1(x)
-        x = self.act(x)
-        x = self.conv2(x)
-        x = self.norm(x)
-
-        return x
-
-
-def compute_upscale_stages(patch_size, min_size=4):
-    # Compute how many times to upscale per dimension
-    num_stages = []
-    for size in patch_size:
-        stages = max(0, int(math.log2(size)) - int(math.log2(min_size)))
-        num_stages.append(stages)
-    return num_stages
-
-
-class EoMT(nn.Module):
-    def __init__(
-        self,
-        backbone: "VisionTransformer",
-        num_classes: int,
-        num_q: int,
-        num_blocks=4,
-        masked_attn_enabled=True,
-    ):
-        super().__init__()
-        self.backbone = backbone
-        self.num_q = num_q
-        self.num_blocks = num_blocks
-        self.masked_attn_enabled = masked_attn_enabled
-
-        self.register_buffer("attn_mask_probs", torch.ones(num_blocks))
-
-        self.q = nn.Embedding(num_q, self.backbone.embed_dim)
-
-        self.class_head = nn.Linear(self.backbone.embed_dim, num_classes + 1)
-
-        self.mask_head = nn.Sequential(
-            nn.Linear(self.backbone.embed_dim, self.backbone.embed_dim),
-            nn.GELU(),
-            nn.Linear(self.backbone.embed_dim, self.backbone.embed_dim),
-            nn.GELU(),
-            nn.Linear(self.backbone.embed_dim, self.backbone.embed_dim),
-        )
-
-        patch_size = self.backbone.patch_embed.patch_size
-        num_upscale_stages = compute_upscale_stages(patch_size, min_size=4)
-
-        # Build per-stage scale factors list
-        max_stages = max(num_upscale_stages)
-        upscale_blocks = []
-        for stage_idx in range(max_stages):
-            # for each dimension, upscale by 2 only if this dimension still has
-            # remaining upscales at this stage
-            scale_factors = tuple(
-                2 if stage_idx < num_upscale_stages[dim] else 1
-                for dim in range(len(patch_size))
-            )
-            upscale_blocks.append(ScaleBlock(self.backbone.embed_dim,
-                                            scale_factors=scale_factors))
-
-        self.upscale = nn.Sequential(*upscale_blocks)
-
-    def _predict(self, x: torch.Tensor):
-        q = x[:, : self.num_q, :]
-
-        class_logits = self.class_head(q)
-
-        x = x[:, self.num_q + self.backbone.num_prefix_tokens :, :]
-        x = x.transpose(1, 2).reshape(
-            x.shape[0], -1, *self.backbone.patch_embed.grid_size
-        )
-
-        mask_logits = torch.einsum(
-            "bqc, bchwd -> bqhwd", self.mask_head(q), self.upscale(x)
-        )
-
-        return mask_logits, class_logits
-
-    @torch.compiler.disable
-    def _disable_attn_mask(self, attn_mask, prob):
-        if prob < 1:
-            random_queries = (
-                torch.rand(attn_mask.shape[0], self.num_q, device=attn_mask.device)
-                > prob
-            )
-            attn_mask[
-                :, : self.num_q, self.num_q + self.backbone.num_prefix_tokens :
-            ][random_queries] = True
-
-        return attn_mask
-
-    def _attn(self, module: 'Attention', x: torch.Tensor, mask: Optional[torch.Tensor], rope=None):
-        B, N, C = x.shape
-
-        q = module.q(x).reshape(B, N, module.num_heads, module.head_dim).permute(0, 2, 1, 3)
-        kv = module.kv(x).reshape(B, N, 2, module.num_heads, module.head_dim)
-        k, v = kv.permute(2, 0, 3, 1, 4).unbind(0)
-        q, k = module.q_norm(q), module.k_norm(k)
-
-        if mask is not None:
-            mask = mask[:, None, ...].expand(-1, module.num_heads, -1, -1)
-
-        dropout_p = module.attn_drop.p if self.training else 0.0
-
-        if rope is not None:
-            if isinstance(rope, list):
-                rope = tuple(torch.stack([r[i] for r in rope], dim=0) for i in range(2))
-            q, k = module.apply_rotary_pos_emb(q, k, rope)
-
-        if module.fused_attn:
-            x = F.scaled_dot_product_attention(q, k, v, mask, dropout_p)
-        else:
-            attn = (q @ k.transpose(-2, -1)) * module.scale
-            if mask is not None:
-                attn = attn.masked_fill(~mask, float("-inf"))
-            attn = F.softmax(attn, dim=-1)
-            attn = module.attn_drop(attn)
-            x = attn @ v
-
-        x = module.proj_drop(module.proj(x.transpose(1, 2).reshape(B, N, C)))
-
-        return x
-
-    def forward(self, x: torch.Tensor):
-        x = self.backbone.patch_embed(x)
-        x, rope = self.backbone._pos_embed(x)
-        x = self.backbone.patch_drop(x)
-        x = self.backbone.norm_pre(x)
-
-        attn_mask = None
-        mask_logits_per_layer, class_logits_per_layer = [], []
-
-        for i, block in enumerate(self.backbone.blocks):
-            if i == len(self.backbone.blocks) - self.num_blocks:
-                x = torch.cat(
-                    (self.q.weight[None, :, :].expand(x.shape[0], -1, -1), x), dim=1
-                )
-
-            if (
-                self.masked_attn_enabled
-                and i >= len(self.backbone.blocks) - self.num_blocks
-            ):
-                mask_logits, class_logits = self._predict(self.backbone.norm(x))
-                mask_logits_per_layer.append(mask_logits)
-                class_logits_per_layer.append(class_logits)
-
-                attn_mask = torch.ones(
-                    x.shape[0],
-                    x.shape[1],
-                    x.shape[1],
-                    dtype=torch.bool,
-                    device=x.device,
-                )
-                interpolated = F.interpolate(
-                    mask_logits,
-                    self.backbone.patch_embed.grid_size,
-                    mode="trilinear",
-                )
-                interpolated = interpolated.view(
-                    interpolated.size(0), interpolated.size(1), -1
-                )
-                attn_mask[
-                    :,
-                    : self.num_q,
-                    self.num_q + self.backbone.num_prefix_tokens :,
-                ] = (
-                    interpolated > 0
-                )
-                attn_mask = self._disable_attn_mask(
-                    attn_mask,
-                    self.attn_mask_probs[
-                        i - len(self.backbone.blocks) + self.num_blocks
-                    ],
-                )
-
-            x = x + block.drop_path1(
-                block.ls1(self._attn(block.attn, block.norm1(x), attn_mask, rope))
-            )
-            x = x + block.drop_path2(block.ls2(block.mlp(block.norm2(x))))
-
-        mask_logits, class_logits = self._predict(self.backbone.norm(x))
-        mask_logits_per_layer.append(mask_logits)
-        class_logits_per_layer.append(class_logits)
-
-        return (
-            mask_logits_per_layer,
-            class_logits_per_layer,
-        )

spectre/models/resnet.py

@@ -1,726 +0,0 @@
-import os
-import math
-from urllib.parse import urlparse
-from typing import Type, Any, Tuple, List, Optional, Union, Dict
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from spectre.utils import to_ntuple
-
-
-def get_padding(kernel_size: int, stride: int, dilation: int = 1) -> int:
-    padding = ((stride - 1) + dilation * (kernel_size - 1)) // 2
-    return padding
-
-
-class BasicBlock(nn.Module):
-    expansion = 1
-
-    def __init__(
-            self,
-            inplanes: int,
-            planes: int,
-            stride: int = 1,
-            downsample: Optional[nn.Module] = None,
-            cardinality: int = 1,
-            base_width: int = 64,
-            reduce_first: int = 1,
-            dilation: int = 1,
-            first_dilation: Optional[int] = None,
-            act_layer: Type[nn.Module] = nn.ReLU,
-            norm_layer: Type[nn.Module] = nn.BatchNorm3d,
-    ):
-        """
-        Args:
-            inplanes: Input channel dimensionality.
-            planes: Used to determine output channel dimensionalities.
-            stride: Stride used in convolution layers.
-            downsample: Optional downsample layer for residual path.
-            cardinality: Number of convolution groups.
-            base_width: Base width used to determine output channel dimensionality.
-            reduce_first: Reduction factor for first convolution output width of residual blocks.
-            dilation: Dilation rate for convolution layers.
-            first_dilation: Dilation rate for first convolution layer.
-            act_layer: Activation layer.
-            norm_layer: Normalization layer.
-        """
-        super(BasicBlock, self).__init__()
-
-        assert cardinality == 1, 'BasicBlock only supports cardinality of 1'
-        assert base_width == 64, 'BasicBlock does not support changing base width'
-        first_planes = planes // reduce_first
-        outplanes = planes * self.expansion
-        first_dilation = first_dilation or dilation
-
-        self.conv1 = nn.Conv3d(
-            inplanes, first_planes, kernel_size=3, stride=stride, padding=first_dilation,
-            dilation=first_dilation, bias=False)
-        self.bn1 = norm_layer(first_planes)
-        self.act1 = act_layer()
-
-        self.conv2 = nn.Conv3d(
-            first_planes, outplanes, kernel_size=3, padding=dilation, dilation=dilation, bias=False)
-        self.bn2 = norm_layer(outplanes)
-
-        self.act2 = act_layer()
-        self.downsample = downsample
-        self.stride = stride
-        self.dilation = dilation
-
-    def zero_init_last(self):
-        if getattr(self.bn2, 'weight', None) is not None:
-            nn.init.zeros_(self.bn2.weight)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        shortcut = x
-
-        x = self.conv1(x)
-        x = self.bn1(x)
-        x = self.act1(x)
-
-        x = self.conv2(x)
-        x = self.bn2(x)
-
-        if self.downsample is not None:
-            shortcut = self.downsample(shortcut)
-        x = x + shortcut
-        x = self.act2(x)
-
-        return x
-
-
-class Bottleneck(nn.Module):
-    expansion = 4
-
-    def __init__(
-            self,
-            inplanes: int,
-            planes: int,
-            stride: int = 1,
-            downsample: Optional[nn.Module] = None,
-            cardinality: int = 1,
-            base_width: int = 64,
-            reduce_first: int = 1,
-            dilation: int = 1,
-            first_dilation: Optional[int] = None,
-            act_layer: Type[nn.Module] = nn.ReLU,
-            norm_layer: Type[nn.Module] = nn.BatchNorm3d,
-    ):
-        """
-        Args:
-            inplanes: Input channel dimensionality.
-            planes: Used to determine output channel dimensionalities.
-            stride: Stride used in convolution layers.
-            downsample: Optional downsample layer for residual path.
-            cardinality: Number of convolution groups.
-            base_width: Base width used to determine output channel dimensionality.
-            reduce_first: Reduction factor for first convolution output width of residual blocks.
-            dilation: Dilation rate for convolution layers.
-            first_dilation: Dilation rate for first convolution layer.
-            act_layer: Activation layer.
-            norm_layer: Normalization layer.
-        """
-        super(Bottleneck, self).__init__()
-
-        width = int(math.floor(planes * (base_width / 64)) * cardinality)
-        first_planes = width // reduce_first
-        outplanes = planes * self.expansion
-        first_dilation = first_dilation or dilation
-
-        self.conv1 = nn.Conv3d(inplanes, first_planes, kernel_size=1, bias=False)
-        self.bn1 = norm_layer(first_planes)
-        self.act1 = act_layer()
-
-        self.conv2 = nn.Conv3d(
-            first_planes, width, kernel_size=3, stride=stride,
-            padding=first_dilation, dilation=first_dilation, groups=cardinality, bias=False)
-        self.bn2 = norm_layer(width)
-        self.act2 = act_layer()
-
-        self.conv3 = nn.Conv3d(width, outplanes, kernel_size=1, bias=False)
-        self.bn3 = norm_layer(outplanes)
-
-        self.act3 = act_layer()
-        self.downsample = downsample
-        self.stride = stride
-        self.dilation = dilation
-
-    def zero_init_last(self):
-        if getattr(self.bn3, 'weight', None) is not None:
-            nn.init.zeros_(self.bn3.weight)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        shortcut = x
-
-        x = self.conv1(x)
-        x = self.bn1(x)
-        x = self.act1(x)
-
-        x = self.conv2(x)
-        x = self.bn2(x)
-        x = self.act2(x)
-
-        x = self.conv3(x)
-        x = self.bn3(x)
-
-        if self.downsample is not None:
-            shortcut = self.downsample(shortcut)
-        x = x + shortcut
-        x = self.act3(x)
-
-        return x
-
-
-def downsample_conv(
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int,
-        stride: int = 1,
-        dilation: int = 1,
-        first_dilation: Optional[int] = None,
-        norm_layer: Optional[Type[nn.Module]] = None,
-) -> nn.Module:
-    norm_layer = norm_layer or nn.BatchNorm3d
-    kernel_size = 1 if stride == 1 and dilation == 1 else kernel_size
-    first_dilation = (first_dilation or dilation) if kernel_size > 1 else 1
-    p = get_padding(kernel_size, stride, first_dilation)
-
-    return nn.Sequential(*[
-        nn.Conv3d(
-            in_channels, out_channels, kernel_size, stride=stride, padding=p, dilation=first_dilation, bias=False),
-        norm_layer(out_channels)
-    ])
-
-
-def downsample_avg(
-        in_channels: int,
-        out_channels: int,
-        kernel_size: int,
-        stride: int = 1,
-        dilation: int = 1,
-        first_dilation: Optional[int] = None,
-        norm_layer: Optional[Type[nn.Module]] = None,
-) -> nn.Module:
-    norm_layer = norm_layer or nn.BatchNorm3d
-    avg_stride = stride if dilation == 1 else 1
-    if stride == 1 and dilation == 1:
-        pool = nn.Identity()
-    else:
-        pool = nn.AvgPool3d(2, avg_stride, ceil_mode=True, count_include_pad=False)
-
-    return nn.Sequential(*[
-        pool,
-        nn.Conv3d(in_channels, out_channels, 1, stride=1, padding=0, bias=False),
-        norm_layer(out_channels)
-    ])
-
-
-def make_blocks(
-        block_fns: Tuple[Union[BasicBlock, Bottleneck]],
-        channels: Tuple[int, ...],
-        block_repeats: Tuple[int, ...],
-        inplanes: int,
-        reduce_first: int = 1,
-        output_stride: int = 32,
-        down_kernel_size: int = 1,
-        avg_down: bool = False,
-        **kwargs,
-) -> Tuple[List[Tuple[str, nn.Module]], List[Dict[str, Any]]]:
-    stages = []
-    feature_info = []
-    net_num_blocks = sum(block_repeats)
-    net_block_idx = 0
-    net_stride = 4
-    dilation = prev_dilation = 1
-    for stage_idx, (block_fn, planes, num_blocks) in enumerate(zip(block_fns, channels, block_repeats)):
-        stage_name = f'layer{stage_idx + 1}'  # never liked this name, but weight compat requires it
-        stride = 1 if stage_idx == 0 else 2
-        if net_stride >= output_stride:
-            dilation *= stride
-            stride = 1
-        else:
-            net_stride *= stride
-
-        downsample = None
-        if stride != 1 or inplanes != planes * block_fn.expansion:
-            down_kwargs = dict(
-                in_channels=inplanes,
-                out_channels=planes * block_fn.expansion,
-                kernel_size=down_kernel_size,
-                stride=stride,
-                dilation=dilation,
-                first_dilation=prev_dilation,
-                norm_layer=kwargs.get('norm_layer'),
-            )
-            downsample = downsample_avg(**down_kwargs) if avg_down else downsample_conv(**down_kwargs)
-
-        block_kwargs = dict(reduce_first=reduce_first, dilation=dilation, **kwargs)
-        blocks = []
-        for block_idx in range(num_blocks):
-            downsample = downsample if block_idx == 0 else None
-            stride = stride if block_idx == 0 else 1
-            blocks.append(block_fn(
-                inplanes,
-                planes,
-                stride,
-                downsample,
-                first_dilation=prev_dilation,
-                **block_kwargs,
-            ))
-            prev_dilation = dilation
-            inplanes = planes * block_fn.expansion
-            net_block_idx += 1
-
-        stages.append((stage_name, nn.Sequential(*blocks)))
-        feature_info.append(dict(num_chs=inplanes, reduction=net_stride, module=stage_name))
-
-    return stages, feature_info
-
-
-def feature_take_indices(
-        num_features: int,
-        indices: Optional[Union[int, List[int]]] = None,
-        as_set: bool = False,
-) -> Tuple[List[int], int]:
-    """ Determine the absolute feature indices to 'take' from.
-
-    Note: This function can be called in forward() so must be torchscript compatible,
-    which requires some incomplete typing and workaround hacks.
-
-    Args:
-        num_features: total number of features to select from
-        indices: indices to select,
-          None -> select all
-          int -> select last n
-          list/tuple of int -> return specified (-ve indices specify from end)
-        as_set: return as a set
-
-    Returns:
-        List (or set) of absolute (from beginning) indices, Maximum index
-    """
-    if indices is None:
-        indices = num_features  # all features if None
-
-    if isinstance(indices, int):
-        # convert int -> last n indices
-        assert 0 < indices <= num_features, f'last-n ({indices}) is out of range (1 to {num_features})'
-        take_indices = [num_features - indices + i for i in range(indices)]
-    else:
-        take_indices: List[int] = []
-        for i in indices:
-            idx = num_features + i if i < 0 else i
-            assert 0 <= idx < num_features, f'feature index {idx} is out of range (0 to {num_features - 1})'
-            take_indices.append(idx)
-
-    if not torch.jit.is_scripting() and as_set:
-        return set(take_indices), max(take_indices)
-
-    return take_indices, max(take_indices)
-
-
-class ResNet(nn.Module):
-    """ResNet / ResNeXt
-
-    This class implements all variants of ResNet, ResNeXt that
-      * have > 1 stride in the 3x3 conv layer of bottleneck
-      * have conv-bn-act ordering
-
-    This ResNet impl supports a number of stem and downsample options based on the v1c, v1d, v1e, and v1s
-    variants included in the MXNet Gluon ResNetV1b model. The C and D variants are also discussed in the
-    'Bag of Tricks' paper: https://arxiv.org/pdf/1812.01187. The B variant is equivalent to torchvision default.
-
-    ResNet variants (the same modifications can be used in SE/ResNeXt models as well):
-      * normal, b - 7x7 stem, stem_width = 64, same as torchvision ResNet, NVIDIA ResNet 'v1.5', Gluon v1b
-      * c - 3 layer deep 3x3 stem, stem_width = 32 (32, 32, 64)
-      * d - 3 layer deep 3x3 stem, stem_width = 32 (32, 32, 64), average pool in downsample
-      * e - 3 layer deep 3x3 stem, stem_width = 64 (64, 64, 128), average pool in downsample
-      * s - 3 layer deep 3x3 stem, stem_width = 64 (64, 64, 128)
-      * t - 3 layer deep 3x3 stem, stem width = 32 (24, 48, 64), average pool in downsample
-      * tn - 3 layer deep 3x3 stem, stem width = 32 (24, 32, 64), average pool in downsample
-
-    ResNeXt
-      * normal - 7x7 stem, stem_width = 64, standard cardinality and base widths
-      * same c,d, e, s variants as ResNet can be enabled
-    """
-
-    def __init__(
-            self,
-            block: Union[BasicBlock, Bottleneck],
-            layers: Tuple[int, ...],
-            num_classes: int = 1000,
-            in_chans: int = 1,
-            output_stride: int = 32,
-            global_pool: str = 'avg',
-            cardinality: int = 1,
-            base_width: int = 64,
-            stem_width: int = 64,
-            stem_type: str = '',
-            replace_stem_pool: bool = False,
-            block_reduce_first: int = 1,
-            down_kernel_size: int = 1,
-            avg_down: bool = False,
-            channels: Optional[Tuple[int, ...]] = (64, 128, 256, 512),
-            act_layer: Type[nn.Module] = nn.ReLU,
-            norm_layer: Type[nn.Module] = nn.BatchNorm3d,
-            drop_rate: float = 0.0,
-            zero_init_last: bool = True,
-            block_args: Optional[Dict[str, Any]] = None,
-    ):
-        """
-        Args:
-            block (nn.Module): class for the residual block. Options are BasicBlock, Bottleneck.
-            layers (List[int]) : number of layers in each block
-            num_classes (int): number of classification classes (default 1000)
-            in_chans (int): number of input (color) channels. (default 3)
-            output_stride (int): output stride of the network, 32, 16, or 8. (default 32)
-            global_pool (str): Global pooling type. One of 'avg', 'max' (default 'avg')
-            cardinality (int): number of convolution groups for 3x3 conv in Bottleneck. (default 1)
-            base_width (int): bottleneck channels factor. `planes * base_width / 64 * cardinality` (default 64)
-            stem_width (int): number of channels in stem convolutions (default 64)
-            stem_type (str): The type of stem (default ''):
-                * '', default - a single 7x7 conv with a width of stem_width
-                * 'deep' - three 3x3 convolution layers of widths stem_width, stem_width, stem_width * 2
-                * 'deep_tiered' - three 3x3 conv layers of widths stem_width//4 * 3, stem_width, stem_width * 2
-            replace_stem_pool (bool): replace stem max-pooling layer with a 3x3 stride-2 convolution
-            block_reduce_first (int): Reduction factor for first convolution output width of residual blocks,
-                1 for all archs except senets, where 2 (default 1)
-            down_kernel_size (int): kernel size of residual block downsample path,
-                1x1 for most, 3x3 for senets (default: 1)
-            avg_down (bool): use avg pooling for projection skip connection between stages/downsample (default False)
-            act_layer (str, nn.Module): activation layer
-            norm_layer (str, nn.Module): normalization layer
-            drop_rate (float): Dropout probability before classifier, for training (default 0.)
-            zero_init_last (bool): zero-init the last weight in residual path (usually last BN affine weight)
-            block_args (dict): Extra kwargs to pass through to block module
-        """
-        super(ResNet, self).__init__()
-        block_args = block_args or dict()
-        assert output_stride in (8, 16, 32)
-        self.num_classes = num_classes
-        self.drop_rate = drop_rate
-
-        # Stem
-        deep_stem = 'deep' in stem_type
-        inplanes = stem_width * 2 if deep_stem else 64
-        if deep_stem:
-            stem_chs = (stem_width, stem_width)
-            if 'tiered' in stem_type:
-                stem_chs = (3 * (stem_width // 4), stem_width)
-            self.conv1 = nn.Sequential(*[
-                nn.Conv3d(in_chans, stem_chs[0], 3, stride=2, padding=1, bias=False),
-                norm_layer(stem_chs[0]),
-                act_layer(),
-                nn.Conv3d(stem_chs[0], stem_chs[1], 3, stride=1, padding=1, bias=False),
-                norm_layer(stem_chs[1]),
-                act_layer(),
-                nn.Conv3d(stem_chs[1], inplanes, 3, stride=1, padding=1, bias=False)])
-        else:
-            self.conv1 = nn.Conv3d(in_chans, inplanes, kernel_size=7, stride=2, padding=3, bias=False)
-        self.bn1 = norm_layer(inplanes)
-        self.act1 = act_layer()
-        self.feature_info = [dict(num_chs=inplanes, reduction=2, module='act1')]
-
-        # Stem pooling. The name 'maxpool' remains for weight compatibility.
-        if replace_stem_pool:
-            self.maxpool = nn.Sequential(*filter(None, [
-                nn.Conv3d(inplanes, inplanes, 3, stride=2, padding=1, bias=False),
-                norm_layer(inplanes),
-                act_layer(),
-            ]))
-        else:
-            self.maxpool = nn.MaxPool3d(kernel_size=3, stride=2, padding=1)
-
-        # Feature Blocks
-        block_fns = to_ntuple(len(channels))(block)
-        stage_modules, stage_feature_info = make_blocks(
-            block_fns,
-            channels,
-            layers,
-            inplanes,
-            cardinality=cardinality,
-            base_width=base_width,
-            output_stride=output_stride,
-            reduce_first=block_reduce_first,
-            avg_down=avg_down,
-            down_kernel_size=down_kernel_size,
-            act_layer=act_layer,
-            norm_layer=norm_layer,
-            **block_args,
-        )
-        for stage in stage_modules:
-            self.add_module(*stage)  # layer1, layer2, etc
-        self.feature_info.extend(stage_feature_info)
-
-        # Head (Pooling and Classifier)
-        self.num_features = self.head_hidden_size = channels[-1] * block_fns[-1].expansion
-        if global_pool == 'avg':
-            self.global_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
-        elif global_pool == 'max':
-            self.global_pool = nn.AdaptiveMaxPool3d((1, 1, 1))
-        else:
-            raise NotImplementedError('Global pooling type not supported: {}'.format(global_pool))
-        
-        self.fc = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
-
-        self.init_weights(zero_init_last=zero_init_last)
-
-    @torch.jit.ignore
-    def init_weights(self, zero_init_last: bool = True):
-        for n, m in self.named_modules():
-            if isinstance(m, nn.Conv3d):
-                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
-        if zero_init_last:
-            for m in self.modules():
-                if hasattr(m, 'zero_init_last'):
-                    m.zero_init_last()
-
-    @torch.jit.ignore
-    def group_matcher(self, coarse: bool = False):
-        matcher = dict(stem=r'^conv1|bn1|maxpool', blocks=r'^layer(\d+)' if coarse else r'^layer(\d+)\.(\d+)')
-        return matcher
-
-    @torch.jit.ignore
-    def get_classifier(self, name_only: bool = False):
-        return 'fc' if name_only else self.fc
-
-    def reset_classifier(self, num_classes: int, global_pool: str = 'avg'):
-        self.num_classes = num_classes
-        if global_pool == 'avg':
-            self.global_pool = nn.AdaptiveAvgPool3d((1, 1, 1))
-        elif global_pool == 'max':
-            self.global_pool = nn.AdaptiveMaxPool3d((1, 1, 1))
-        else:
-            raise NotImplementedError('Global pooling type not supported: {}'.format(global_pool))
-        
-        self.fc = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity()
-
-    def forward_intermediates(
-            self,
-            x: torch.Tensor,
-            indices: Optional[Union[int, List[int]]] = None,
-            norm: bool = False,
-            stop_early: bool = False,
-            output_fmt: str = 'NCHWD',
-            intermediates_only: bool = False,
-    ) -> Union[List[torch.Tensor], Tuple[torch.Tensor, List[torch.Tensor]]]:
-        """ Forward features that returns intermediates.
-
-        Args:
-            x: Input image tensor
-            indices: Take last n blocks if int, all if None, select matching indices if sequence
-            norm: Apply norm layer to compatible intermediates
-            stop_early: Stop iterating over blocks when last desired intermediate hit
-            output_fmt: Shape of intermediate feature outputs
-            intermediates_only: Only return intermediate features
-        Returns:
-
-        """
-        assert output_fmt in ('NCHWD',), 'Output shape must be NCHWD.'
-        intermediates = []
-        take_indices, max_index = feature_take_indices(5, indices)
-
-        # forward pass
-        feat_idx = 0
-        x = self.conv1(x)
-        x = self.bn1(x)
-        x = self.act1(x)
-        if feat_idx in take_indices:
-            intermediates.append(x)
-        x = self.maxpool(x)
-
-        layer_names = ('layer1', 'layer2', 'layer3', 'layer4')
-        if stop_early:
-            layer_names = layer_names[:max_index]
-        for n in layer_names:
-            feat_idx += 1
-            x = getattr(self, n)(x)  # won't work with torchscript, but keeps code reasonable, FML
-            if feat_idx in take_indices:
-                intermediates.append(x)
-
-        if intermediates_only:
-            return intermediates
-
-        return x, intermediates
-
-    def prune_intermediate_layers(
-            self,
-            indices: Union[int, List[int]] = 1,
-            prune_norm: bool = False,
-            prune_head: bool = True,
-    ):
-        """ Prune layers not required for specified intermediates.
-        """
-        take_indices, max_index = feature_take_indices(5, indices)
-        layer_names = ('layer1', 'layer2', 'layer3', 'layer4')
-        layer_names = layer_names[max_index:]
-        for n in layer_names:
-            setattr(self, n, nn.Identity())
-        if prune_head:
-            self.reset_classifier(0, '')
-        return take_indices
-
-    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.conv1(x)
-        x = self.bn1(x)
-        x = self.act1(x)
-        x = self.maxpool(x)
-
-        x = self.layer1(x)
-        x = self.layer2(x)
-        x = self.layer3(x)
-        x = self.layer4(x)
-        return x
-
-    def forward_head(self, x: torch.Tensor, pre_logits: bool = False) -> torch.Tensor:
-        x = self.global_pool(x)
-        x = x.flatten(1)
-        if self.drop_rate:
-            x = F.dropout(x, p=float(self.drop_rate), training=self.training)
-        return x if pre_logits else self.fc(x)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.forward_features(x)
-        x = self.forward_head(x)
-        return x
-    
-    @classmethod
-    def from_pretrained(
-            cls,
-            checkpoint_path_or_url: Union[str, os.PathLike],
-            verbose: bool = True,
-            **kwargs
-    ) -> 'ResNet':
-        """Load pretrained model weights from a local path or a URL."""
-        model = cls(**kwargs)
-
-        def _is_url(path: str) -> bool:
-            try:
-                parsed = urlparse(str(path))
-                return parsed.scheme in ('http', 'https')
-            except Exception:
-                return False
-
-        if _is_url(checkpoint_path_or_url):
-            if verbose:
-                print(f"Downloading pretrained weights from URL: {checkpoint_path_or_url}")
-            state_dict = torch.hub.load_state_dict_from_url(
-                checkpoint_path_or_url, map_location='cpu', weights_only=False, progress=verbose)
-        else:
-            local_path = os.fspath(checkpoint_path_or_url)
-            if not os.path.exists(local_path):
-                raise FileNotFoundError(f"Checkpoint file not found: {local_path}")
-        if verbose:
-            print(f"Loading checkpoint from local path: {local_path}")
-            state_dict = torch.load(local_path, map_location='cpu', weights_only=False)
-
-        msg = model.load_state_dict(state_dict, strict=False)
-        if verbose:
-            print(f"Loaded pretrained weights with msg: {msg}")
-        return model
-
-    
-
-def resnet18(
-    checkpoint_path_or_url: Optional[str] = None,
-    **kwargs
-) -> ResNet:
-    """ResNet-18 model with 3D operations.
-    """
-    kwargs = dict(
-        block=BasicBlock,
-        layers=[2, 2, 2, 2],
-        cardinality=1,
-        **kwargs,
-    )
-    if checkpoint_path_or_url:
-        return ResNet.from_pretrained(checkpoint_path_or_url, **kwargs)
-    return ResNet(**kwargs)
-
-
-def resnet34(
-    checkpoint_path_or_url: Optional[str] = None,
-    **kwargs
-) -> ResNet:
-    """ResNet-34 model with 3D operations.
-    """
-    kwargs = dict(
-        block=BasicBlock,
-        layers=[3, 4, 6, 3],
-        cardinality=1,
-        **kwargs,
-    )
-    if checkpoint_path_or_url:
-        return ResNet.from_pretrained(checkpoint_path_or_url, **kwargs)
-    return ResNet(**kwargs)
-
-
-def resnet50(
-    checkpoint_path_or_url: Optional[str] = None,
-    **kwargs
-) -> ResNet:
-    """ResNet-50 model with 3D operations.
-    """
-    kwargs = dict(
-        block=Bottleneck,
-        layers=[3, 4, 6, 3],
-        cardinality=1,
-        **kwargs,
-    )
-    if checkpoint_path_or_url:
-        return ResNet.from_pretrained(checkpoint_path_or_url, **kwargs)
-    return ResNet(**kwargs)
-
-
-def resnet101(
-    checkpoint_path_or_url: Optional[str] = None,
-    **kwargs
-) -> ResNet:
-    """ResNet-101 model with 3D operations.
-    """
-    kwargs = dict(
-        block=Bottleneck,
-        layers=[3, 4, 23, 3],
-        cardinality=1,
-        **kwargs,
-    )
-    if checkpoint_path_or_url:
-        return ResNet.from_pretrained(checkpoint_path_or_url, **kwargs)
-    return ResNet(**kwargs)
-
-
-def resnext50(
-    checkpoint_path_or_url: Optional[str] = None,
-    **kwargs
-) -> ResNet:
-    """ResNeXt-50 model with 3D operations.
-    """
-    kwargs = dict(
-        block=Bottleneck,
-        layers=[3, 4, 6, 3],
-        cardinality=32,
-        base_width=4,
-        **kwargs,
-    )
-    if checkpoint_path_or_url:
-        return ResNet.from_pretrained(checkpoint_path_or_url, **kwargs)
-    return ResNet(**kwargs)
-
-
-def resnext101(
-    checkpoint_path_or_url: Optional[str] = None,
-    **kwargs
-) -> ResNet:
-    """ResNeXt-101 model with 3D operations.
-    """
-    kwargs = dict(
-        block=Bottleneck,
-        layers=[3, 4, 23, 3],
-        cardinality=32,
-        base_width=8,
-        **kwargs,
-    )
-    if checkpoint_path_or_url:
-        return ResNet.from_pretrained(checkpoint_path_or_url, **kwargs)
-    return ResNet(**kwargs)

spectre/models/seomt.py

@@ -1,394 +0,0 @@
-# Adapted from https://github.com/tue-mps/eomt/
-
-import math
-from typing import Optional, Tuple, Union
-
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-
-from spectre.models import VisionTransformer
-from spectre.models.layers import LayerNorm3d
-
-
-class ScaleBlock(nn.Module):
-    def __init__(
-        self, 
-        embed_dim: int, 
-        scale_factors: Union[int, Tuple[int, int, int]] = (2, 2, 2),
-        conv1_layer: nn.Module = nn.ConvTranspose3d,
-    ):
-        super().__init__()
-
-        self.conv1 = conv1_layer(
-            embed_dim,
-            embed_dim,
-            kernel_size=scale_factors,
-            stride=scale_factors,
-        )
-        self.act = nn.GELU()
-        self.conv2 = nn.Conv3d(
-            embed_dim,
-            embed_dim,
-            kernel_size=3,
-            padding=1,
-            groups=embed_dim,
-            bias=False,
-        )
-        self.norm = LayerNorm3d(embed_dim)
-
-    def forward(self, x):
-        # print(x.shape)
-        x = self.conv1(x)
-        x = self.act(x)
-        x = self.conv2(x)
-        x = self.norm(x)
-
-        return x
-
-
-def compute_upscale_stages(patch_size, min_size=4):
-    # Compute how many times to upscale per dimension
-    num_stages = []
-    for size in patch_size:
-        stages = max(0, int(math.log2(size)) - int(math.log2(min_size)))
-        num_stages.append(stages)
-    return num_stages
-
-def voxel_shuffle_3d(x: torch.Tensor, r: int = 2) -> torch.Tensor:
-    """
-    Rearranges channels of a 5D tensor (N, C*r^3, D, H, W) to
-    (N, C, D*r, H*r, W*r).
-    """
-    n, c, d, h, w = x.size()
-    assert c % (r ** 3) == 0, f"Channels {c} not divisible by r^3={r**3}"
-    out_c = c // (r ** 3)
-    x = x.view(n, out_c, r, r, r, d, h, w)          # (N, C, r, r, r, D, H, W)
-    x = x.permute(0, 1, 5, 2, 6, 3, 7, 4).contiguous()  # (N, C, D, r, H, r, W, r)
-    x = x.view(n, out_c, d * r, h * r, w * r)       # (N, C, D*r, H*r, W*r)
-    return x
-
-
-class MLPUpBlock3D(nn.Module):
-    """
-    2x upsampling via:
-      - 1x1x1 Conv (per-voxel MLP) expanding channels by 2^3
-      - 3D voxel shuffle to double D/H/W
-      - optional norm + activation
-    """
-    def __init__(self, channels: int, norm=nn.Identity, activation=nn.GELU):
-        super().__init__()
-        self.proj = nn.Conv3d(channels, channels * 8, kernel_size=1, bias=True)
-        self.norm = norm(channels) if norm is not None else nn.Identity()
-        self.act = activation() if activation is not None else nn.Identity()
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.proj(x)           # (N, C*8, D, H, W)
-        x = voxel_shuffle_3d(x, 2) # (N, C, 2D, 2H, 2W)
-        x = self.norm(x)
-        x = self.act(x)
-        return x
-
-
-class SmolMLPDecoder3D(nn.Module):
-    """
-    Simple decoder with two MLP upsampling layers (2x each) for total 4x upsampling.
-    """
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        norm=LayerNorm3d,
-        activation=nn.GELU,
-    ):
-        super().__init__()
-        self.up1 = MLPUpBlock3D(in_channels, norm=norm, activation=activation)
-        self.up2 = MLPUpBlock3D(in_channels, norm=norm, activation=activation)
-        self.head = nn.Conv3d(in_channels, out_channels, kernel_size=1, bias=True)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        x = self.up1(x)   # 2x
-        x = self.up2(x)   # another 2x => 4x total
-        x = self.head(x)  # map to desired output channels
-        return x
-    
-class SpatialMLPDecoder3D(nn.Module):
-    """
-    Per-voxel MLP that expands logits from (B,Q,H,W,D) to (B,Q,H*s,W*s,D*s)
-    by predicting a learned s^3 block per voxel, then rearranging spatially.
-    """
-    def __init__(self, num_classes: int, upscale_factor: int = 4, hidden_mul: int = 4):
-        super().__init__()
-        self.num_classes = num_classes
-        self.s = upscale_factor
-        hidden = hidden_mul * num_classes
-        self.mlp = nn.Sequential(
-            nn.Linear(num_classes, hidden),
-            nn.GELU(),
-            nn.Linear(hidden, num_classes * (upscale_factor ** 3)),
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # x: (B, Q, H, W, D)
-        B, Q, H, W, D = x.shape
-        s = self.s
-        x = x.permute(0, 2, 3, 4, 1).contiguous().view(B * H * W * D, Q)  # (BHW D, Q)
-        x = self.mlp(x)                                                   # (BHW D, Q*s^3)
-        x = x.view(B, H, W, D, Q, s, s, s)                                # (B,H,W,D,Q,s,s,s)
-        x = x.permute(0, 4, 1, 5, 2, 6, 3, 7).contiguous()                # (B,Q,H,s,W,s,D,s)
-        x = x.view(B, Q, H * s, W * s, D * s)
-        return x
-
-
-class SimpleConvDecoder3D(nn.Module):
-    """
-    Depthwise ConvTranspose3d upsampler:
-    - Assumes input & output channels == num_classes
-    - Uses groups=num_classes for class-wise deconvolution
-    """
-    def __init__(self, num_classes: int, upscale_factor: int = 4):
-        super().__init__()
-        s = upscale_factor
-        self.deconv = nn.ConvTranspose3d(
-            in_channels=num_classes,
-            out_channels=num_classes,
-            kernel_size=s,
-            stride=s,
-            padding=0,
-            output_padding=0,
-            groups=num_classes,
-            bias=True,
-        )
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        # x: (B, Q, H, W, D)
-        return self.deconv(x)
-
-
-class SEoMT(nn.Module):
-    def __init__(
-        self,
-        backbone: VisionTransformer,
-        num_classes: int,
-        # num_q: int,
-        num_blocks=4,
-        masked_attn_enabled=True,
-        return_only_final_layer=False,
-        upscale_output=True,
-        for_nnunet=False,
-        decoder=False,
-    ):
-        super().__init__()
-        self.backbone = backbone
-        self.num_q = num_classes
-        self.num_blocks = num_blocks
-        self.masked_attn_enabled = masked_attn_enabled
-        self.return_only_final_layer = return_only_final_layer
-        self.upscale_output = upscale_output
-        self.register_buffer("attn_mask_probs", torch.ones(num_blocks))
-        self.for_nnunet = for_nnunet
-        self.q = nn.Embedding(num_classes, self.backbone.embed_dim)
-
-
-        self.mask_head = nn.Sequential(
-            nn.Linear(self.backbone.embed_dim, self.backbone.embed_dim),
-            nn.GELU(),
-            nn.Linear(self.backbone.embed_dim, self.backbone.embed_dim),
-            nn.GELU(),
-            nn.Linear(self.backbone.embed_dim, self.backbone.embed_dim),
-        )
-
-        patch_size = self.backbone.patch_embed.patch_size
-        num_upscale_stages = compute_upscale_stages(patch_size, min_size=4)
-
-        # Build per-stage scale factors list
-        max_stages = max(num_upscale_stages)
-        upscale_blocks = []
-        for stage_idx in range(max_stages):
-            # for each dimension, upscale by 2 only if this dimension still has
-            # remaining upscales at this stage
-            scale_factors = tuple(
-                2 if stage_idx < num_upscale_stages[dim] else 1
-                for dim in range(len(patch_size))
-            )
-            upscale_blocks.append(ScaleBlock(self.backbone.embed_dim,
-                                            scale_factors=scale_factors))
-
-        self.upscale = nn.Sequential(*upscale_blocks)
-
-    def _predict(self, x: torch.Tensor, stage: int = None):
-        q = x[:, : self.num_q, :]
-        # print(stage)
-        # class_logits = self.class_head(q)
-        x = x[:, self.num_q + self.backbone.num_prefix_tokens :, :]
-        x = x.transpose(1, 2).reshape(
-            x.shape[0], -1, *self.backbone.patch_embed.grid_size
-        )
-        mask_logits = torch.einsum(
-            "bqc, bchwd -> bqhwd", self.mask_head(q), self.upscale(x)
-        )
-
-
-        return mask_logits
-
-    @torch.compiler.disable
-    def _disable_attn_mask(self, attn_mask, prob):
-        if prob < 1:
-            random_queries = (
-                torch.rand(attn_mask.shape[0], self.num_q, device=attn_mask.device)
-                > prob
-            )
-            attn_mask[
-                :, : self.num_q, self.num_q + self.backbone.num_prefix_tokens :
-            ][random_queries] = True
-
-        return attn_mask
-
-    def _attn(self, module: 'Attention', x: torch.Tensor, mask: Optional[torch.Tensor], rope = None):
-        B, N, C = x.shape
-
-        q = module.q(x).reshape(B, N, module.num_heads, module.head_dim).permute(0, 2, 1, 3)
-        kv = module.kv(x).reshape(B, N, 2, module.num_heads, module.head_dim)
-        k, v = kv.permute(2, 0, 3, 1, 4).unbind(0)
-        q, k = module.q_norm(q), module.k_norm(k)
-
-        if mask is not None:
-            mask = mask[:, None, ...].expand(-1, module.num_heads, -1, -1)
-
-        dropout_p = module.attn_drop.p if self.training else 0.0
-
-        if rope is not None:
-            if isinstance(rope, list):
-                rope = tuple(torch.stack([r[i] for r in rope], dim=0) for i in range(2))
-            q, k = module.apply_rotary_pos_emb(q, k, rope)
-
-        if module.fused_attn:
-            x = F.scaled_dot_product_attention(q, k, v, mask, dropout_p)
-        else:
-            attn = (q @ k.transpose(-2, -1)) * module.scale
-            if mask is not None:
-                attn = attn.masked_fill(~mask, float("-inf"))
-            attn = F.softmax(attn, dim=-1)
-            attn = module.attn_drop(attn)
-            x = attn @ v
-
-        x = module.proj_drop(module.proj(x.transpose(1, 2).reshape(B, N, C)))
-
-        return x
-
-    def forward(self, x: torch.Tensor):
-
-        if self.for_nnunet: # swap data order, will be incoming at czyx - cxyz
-            x = x.permute(0, 1, 3, 4, 2).contiguous()
-
-        self.backbone.patch_embed.set_input_size(x.shape[2:])
-        x = self.backbone.patch_embed(x)
-        x, rope = self.backbone._pos_embed(x)
-        x = self.backbone.patch_drop(x)
-        x = self.backbone.norm_pre(x)
-        attn_mask = None
-        mask_logits_per_layer = []
-
-        for i, block in enumerate(self.backbone.blocks):
-            if i == len(self.backbone.blocks) - self.num_blocks:
-                x = torch.cat(
-                    (self.q.weight[None, :, :].expand(x.shape[0], -1, -1), x), dim=1
-                )
-
-            if (
-                self.masked_attn_enabled
-                and i >= len(self.backbone.blocks) - self.num_blocks
-            ):
-                mask_logits = self._predict(self.backbone.norm(x))
-
-                if self.for_nnunet:
-                     # swap back to czyx
-                    
-                    if self.upscale_output:
-                        # Upscale to original input size / stage
-
-                        stage = len(self.backbone.blocks) -i 
-                        if stage is not None:
-                            input_size = tuple(
-                                int(self.backbone.patch_embed.patch_size[dim] * self.backbone.patch_embed.grid_size[dim] / 2**(stage))
-                                for dim in range(len(self.backbone.patch_embed.patch_size))
-                            )
-
-                        mask_logits_per_layer.append(F.interpolate(mask_logits, input_size, mode="trilinear").permute(0, 1, 4, 2, 3).contiguous())
-                    else:
-                        mask_logits_per_layer.append(mask_logits.permute(0, 1, 4, 2, 3).contiguous())
-                else:
-                    mask_logits_per_layer.append(mask_logits)
-
-                attn_mask = torch.ones(
-                    x.shape[0],
-                    x.shape[1],
-                    x.shape[1],
-                    dtype=torch.bool,
-                    device=x.device,
-                )
-                interpolated = F.interpolate(
-                    mask_logits,
-                    self.backbone.patch_embed.grid_size,
-                    mode="trilinear",
-                )
-                interpolated = interpolated.view(
-                    interpolated.size(0), interpolated.size(1), -1
-                )
-                attn_mask[
-                    :,
-                    : self.num_q,
-                    self.num_q + self.backbone.num_prefix_tokens :,
-                ] = (
-                    interpolated > 0
-                )
-                attn_mask = self._disable_attn_mask(
-                    attn_mask,
-                    self.attn_mask_probs[
-                        i - len(self.backbone.blocks) + self.num_blocks
-                    ],
-                )
-            x = x + block.drop_path1(
-                block.ls1(self._attn(block.attn, block.norm1(x), attn_mask, rope=rope))
-            )
-            x = x + block.drop_path2(block.ls2(block.mlp(block.norm2(x))))
-
-        mask_logits = self._predict(self.backbone.norm(x))
-        if self.for_nnunet:
-            input_size = tuple(
-                                int(self.backbone.patch_embed.patch_size[dim] * self.backbone.patch_embed.grid_size[dim] / 2**0)
-                                for dim in range(len(self.backbone.patch_embed.patch_size))
-                            )
-            mask_logits_per_layer.append(F.interpolate(mask_logits, input_size, mode="trilinear").permute(0, 1, 4, 2, 3).contiguous())
-        else:
-            mask_logits_per_layer.append(mask_logits)
-
-        if self.for_nnunet:
-            # return in reversed order for deep supervision
-            mask_logits_per_layer = mask_logits_per_layer[::-1]
-        return mask_logits_per_layer if not self.return_only_final_layer else mask_logits_per_layer[0]
-    
-
-if __name__ == "__main__":
-    from spectre.models import vit_large_patch16_128
-
-    model = SEoMT(
-        backbone=vit_large_patch16_128(pos_embed='rope',
-            rope_kwargs={
-                "base": 1000.0,  # works for most 3D models
-            },),
-        num_classes=4,
-        num_blocks=4,
-        masked_attn_enabled=True,
-        return_only_final_layer=True,
-        for_nnunet=True,
-        upscale_output=True,
-        decoder=False,
-    )
-    # print number of parameters
-    print(f"Number of parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad)}")
-
-    x = torch.randn(2, 1, 64, 128, 128)
-    out = model(x)
-    for o in out:
-        print(o.shape)

spectre/models/upsample_anything.py

@@ -1,319 +0,0 @@
-import math
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from torch.optim.lr_scheduler import LambdaLR
-
-
-def UPA(hr_image, lr_volume, device="cuda", use_amp=True):
-    """
-    hr_image: numpy or torch [C,Hh,Wh,Dh]
-    lr_volume: torch [1,C,Hl,Wl,Dl]
-    """
-
-    hr = torch.as_tensor(hr_image).unsqueeze(0).float().to(device)
-
-    _, _, Hh, Wh, Dh = hr.shape
-    _, _, Hl, Wl, Dl = lr_volume.shape
-    scale = Hh // Hl
-    assert Wh // Wl == scale and Dh // Dl == scale, "Inconsistent scale factors"
-
-    lr_volume = lr_volume.to(device).float()
-    lr = F.interpolate(hr, scale_factor=1/scale, mode="trilinear", align_corners=False)
-
-    model = LearnablePixelwiseAnisoJBU3D(
-        Hl, Wl, Dl, scale=scale
-    ).to(device)
-
-    model.train()
-    opt = torch.optim.Adam(model.parameters(), lr=1e-1)
-    max_steps = 350
-    gamma = (1e-9 / 1e-1) ** (1.0 / max_steps)
-    scheduler = LambdaLR(opt, lr_lambda=lambda step: gamma ** step)
-    scaler = torch.amp.GradScaler(device=device, enabled=use_amp)
-
-    for step in range(max_steps):
-        opt.zero_grad(set_to_none=True)
-        with torch.amp.autocast(device_type=device, enabled=use_amp):
-            pred = model(lr, hr)
-            loss = F.l1_loss(pred, hr)
-
-        scaler.scale(loss).backward()
-        scaler.step(opt)
-        scaler.update()
-        scheduler.step()
-
-        if step % 50 == 0:
-            print(f"step {step}: loss={loss.item():.5f}")
-
-    model.eval()
-    with torch.inference_mode(), \
-        torch.amp.autocast(device_type=device, enabled=use_amp, dtype=torch.float16):
-        out = model(lr_volume, hr)
- 
-    return out
-
-
-@torch.no_grad()
-def _build_offsets_3d(R_max: int, device):
-    offs = torch.arange(-R_max, R_max + 1, device=device)
-    dX, dY, dZ = torch.meshgrid(offs, offs, offs, indexing="ij")
-    return (
-        dX.reshape(-1),
-        dY.reshape(-1),
-        dZ.reshape(-1),
-    )  # [K]
-
-
-def gather_lr_scalar_3d(map_lr, Ui, Vi, Wi):
-    """
-    map_lr: [1,1,Hl,Wl,Dl] or [Hl,Wl,Dl]
-    Ui,Vi,Wi: [Bn,Hh,Wh,Dh]
-    """
-    Hl, Wl, Dl = map_lr.shape[-3:]
-    flat = Hl * Wl * Dl
-    idx = (Ui * Wl * Dl + Vi * Dl + Wi).reshape(-1)
-    t = map_lr.view(flat)
-    vals = t.index_select(0, idx)
-    return vals.view(Ui.shape)
-
-
-def gs_jbu_aniso_noparent_3d(
-    feat_lr,          # [1,C,Hl,Wl,Dl]
-    guide_hr,         # [1,G,Hh,Wh,Dh]
-    scale,
-    sigma_x_map,
-    sigma_y_map,
-    sigma_z_map,
-    sigma_r_map,
-    R_max=3,
-    alpha_dyn=2.0,
-    C_chunk=64,
-    Nn_chunk=125,
-):
-    _, C, Hl, Wl, Dl = feat_lr.shape
-    _, _, Hh, Wh, Dh = guide_hr.shape
-    device = feat_lr.device
-    dtype_feat = feat_lr.dtype
-
-    # HR grid
-    x = torch.arange(Hh, device=device, dtype=torch.float32)
-    y = torch.arange(Wh, device=device, dtype=torch.float32)
-    z = torch.arange(Dh, device=device, dtype=torch.float32)
-    X, Y, Z = torch.meshgrid(x, y, z, indexing="ij")
-
-    u = (X + 0.5) / scale - 0.5
-    v = (Y + 0.5) / scale - 0.5
-    w = (Z + 0.5) / scale - 0.5
-
-    uc = torch.round(u).clamp(0, Hl - 1).long()
-    vc = torch.round(v).clamp(0, Wl - 1).long()
-    wc = torch.round(w).clamp(0, Dl - 1).long()
-
-    # Dynamic radius
-    sigma_eff = torch.maximum(
-        sigma_x_map,
-        torch.maximum(sigma_y_map, sigma_z_map),
-    )
-    sigma_eff_hr = F.interpolate(
-        sigma_eff, (Hh, Wh, Dh), mode="trilinear", align_corners=False
-    )
-    # sigma_eff_hr = sigma_eff_hr.squeeze(0).squeeze(0)
-    R_map = torch.ceil(alpha_dyn * sigma_eff_hr).clamp(1, R_max).long()
-
-    dX_all, dY_all, dZ_all = _build_offsets_3d(R_max, device)
-
-    num = torch.zeros(C, Hh, Wh, Dh, device=device, dtype=torch.float32)
-    den = torch.zeros(Hh, Wh, Dh, device=device, dtype=torch.float32)
-    m = torch.full((Hh, Wh, Dh), -1e9, device=device, dtype=torch.float32)
-
-    feat_flat = feat_lr[0].permute(1, 2, 3, 0).reshape(-1, C)
-    guide_lr = F.interpolate(
-        guide_hr, (Hl, Wl, Dl), mode="trilinear", align_corners=False
-    )
-
-    for n0 in range(0, len(dX_all), Nn_chunk):
-        dX = dX_all[n0:n0+Nn_chunk][:, None, None, None]
-        dY = dY_all[n0:n0+Nn_chunk][:, None, None, None]
-        dZ = dZ_all[n0:n0+Nn_chunk][:, None, None, None]
-
-        Ui = torch.clamp(uc.unsqueeze(0) + dX, 0, Hl - 1)
-        Vi = torch.clamp(vc.unsqueeze(0) + dY, 0, Wl - 1)
-        Wi = torch.clamp(wc.unsqueeze(0) + dZ, 0, Dl - 1)
-
-        # mask = (dX**2 + dY**2 + dZ**2 <= R_map[None, ...] ** 2)
-        mask = (dX**2 + dY**2 + dZ**2 <= R_map**2).squeeze(0).squeeze(0)
-
-        cx = (Ui.float() + 0.5) * scale - 0.5
-        cy = (Vi.float() + 0.5) * scale - 0.5
-        cz = (Wi.float() + 0.5) * scale - 0.5
-
-        dx = X.unsqueeze(0) - cx
-        dy = Y.unsqueeze(0) - cy
-        dz = Z.unsqueeze(0) - cz
-
-        sx = gather_lr_scalar_3d(sigma_x_map, Ui, Vi, Wi).clamp_min(1e-6)
-        sy = gather_lr_scalar_3d(sigma_y_map, Ui, Vi, Wi).clamp_min(1e-6)
-        sz = gather_lr_scalar_3d(sigma_z_map, Ui, Vi, Wi).clamp_min(1e-6)
-        sr = gather_lr_scalar_3d(sigma_r_map, Ui, Vi, Wi).clamp_min(1e-6)
-
-        log_ws = (
-            -(dx**2)/(2*sx**2)
-            -(dy**2)/(2*sy**2)
-            -(dz**2)/(2*sz**2)
-        )
-
-        diff2 = 0.0
-        for g in range(guide_hr.shape[1]):
-            g0 = gather_lr_scalar_3d(guide_lr[0, g], Ui, Vi, Wi)
-            diff2 += (guide_hr[0, g] - g0) ** 2
-
-        log_wr = -diff2 / (2 * sr**2 + 1e-8)
-        log_w = torch.where(mask, log_ws + log_wr, -1e9)
-
-        m_chunk = log_w.max(dim=0).values
-        m_new = torch.maximum(m, m_chunk)
-
-        scale_old = torch.exp(m - m_new)
-        num *= scale_old
-        den *= scale_old
-
-        w = torch.exp(log_w - m_new)
-        den += w.sum(0)
-
-        idx_flat = (Ui * Wl * Dl + Vi * Dl + Wi).reshape(-1)
-
-        for c0 in range(0, C, C_chunk):
-            c1 = min(c0 + C_chunk, C)
-            f = feat_flat.index_select(0, idx_flat)[:, c0:c1]
-            f = f.view(w.shape + (c1 - c0,))
-            num[c0:c1] += (f * w[..., None]).sum(0).permute(3, 0, 1, 2)
-
-        m = m_new
-
-    out = (num / den.clamp_min(1e-8)).unsqueeze(0)
-    return out.to(dtype_feat)
-
-
-class LearnablePixelwiseAnisoJBU3D(nn.Module):
-    def __init__(
-        self,
-        Hl,
-        Wl,
-        Dl,
-        scale,
-        init_sigma=1.5,
-        init_sigma_r=0.1,
-        R_max=3,
-        alpha_dyn=2.0,
-    ):
-        super().__init__()
-        self.scale = scale
-        self.R_max = R_max
-        self.alpha_dyn = alpha_dyn
-
-        self.sx_raw = nn.Parameter(torch.full((1,1,Hl,Wl,Dl), math.log(init_sigma)))
-        self.sy_raw = nn.Parameter(torch.full((1,1,Hl,Wl,Dl), math.log(init_sigma)))
-        self.sz_raw = nn.Parameter(torch.full((1,1,Hl,Wl,Dl), math.log(init_sigma)))
-        self.sr_raw = nn.Parameter(torch.full((1,1,Hl,Wl,Dl), math.log(init_sigma_r)))
-
-    def forward(self, feat_lr, guide_hr):
-        return gs_jbu_aniso_noparent_3d(
-            feat_lr,
-            guide_hr,
-            self.scale,
-            torch.exp(self.sx_raw),
-            torch.exp(self.sy_raw),
-            torch.exp(self.sz_raw),
-            torch.exp(self.sr_raw),
-            R_max=self.R_max,
-            alpha_dyn=self.alpha_dyn,
-        )
-
-
-if __name__ == "__main__":
-    import argparse
-
-    import numpy as np
-    import nibabel as nib
-    import monai.transforms as transforms
-
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--image_path", type=str, required=True)
-    parser.add_argument("--mask_path", type=str, required=True)
-    parser.add_argument("--device", type=str, default="cuda")
-    parser.add_argument("--use_amp", action="store_true")
-    args = parser.parse_args()
-
-    transform = transforms.Compose([
-        transforms.LoadImaged(keys=("image", "mask")),
-        transforms.EnsureChannelFirstd(keys=("image", "mask"), channel_dim="no_channel"),
-        transforms.ScaleIntensityRanged(
-            keys=("image",),
-            a_min=-150,
-            a_max=250,
-            b_min=0.0,
-            b_max=1.0,
-            clip=True,
-        ),
-        transforms.Orientationd(keys=("image", "mask"), axcodes="RAS"),
-        transforms.RandWeightedCropd(
-            keys=("image", "mask"), 
-            w_key="mask", 
-            spatial_size=(128, 128, 64), 
-            num_samples=1,
-        ),
-        transforms.CopyItemsd(keys=("mask"), times=1, names=("mask_low_res")),
-        transforms.Resized(keys=("mask_low_res"), spatial_size=(16, 16, 8), mode="nearest", align_corners=False)
-    ])
-    sample = transform({
-        "image": args.image_path,
-        "mask": args.mask_path,
-    })[0]
-
-    nib.save(
-        nib.Nifti1Image(
-            (F.interpolate(sample["mask_low_res"].unsqueeze(0), size=(128, 128, 64), mode="nearest").squeeze(0).squeeze(0).cpu().numpy().astype(np.uint8)),
-            affine=np.eye(4),
-        ),
-        "mask_low_res_upscaled.nii.gz",
-    )
-
-    sample["mask_low_res"] = F.one_hot(
-        sample["mask_low_res"].long().squeeze(0), num_classes=4,
-    ).permute(3, 0, 1, 2).unsqueeze(0).float()
-
-    print(sample["mask_low_res"].shape)
-
-    mask_out = UPA(
-        sample["image"],
-        sample["mask_low_res"],
-        device=args.device,
-        use_amp=args.use_amp,
-    )
-
-    mask_out = mask_out.argmax(dim=1, keepdim=True)
-
-    nib.save(
-        nib.Nifti1Image(
-            (sample["image"] * 255).squeeze(0).cpu().numpy().astype(np.uint8),
-            affine=np.eye(4),
-        ),
-        "image.nii.gz",
-    )
-    nib.save(
-        nib.Nifti1Image(
-            sample["mask"].squeeze(0).cpu().numpy().astype(np.uint8),
-            affine=np.eye(4),
-        ),
-        "mask.nii.gz",
-    )
-    torch.save(mask_out.squeeze(0).squeeze(0).cpu(), "upsampled_mask.pt")
-    nib.save(
-        nib.Nifti1Image(
-            mask_out.squeeze(0).squeeze(0).cpu().numpy().astype(np.uint8),
-            affine=np.eye(4),
-        ),
-        "upsampled_mask.nii.gz",
-    )
-    

spectre/utils/checkpointing.py

@@ -1,238 +0,0 @@
-import os
-import random
-import warnings
-from typing import Optional, Any
-
-import torch
-import numpy as np
-import torch.distributed as dist
-
-
-def _get_local_rng_state() -> dict:
-    """Return a picklable dict with local RNG states (cpu & cuda, numpy, random)."""
-    state = {
-        "torch": torch.get_rng_state().cpu(),
-        "numpy": np.random.get_state(),
-        "random": random.getstate(),
-    }
-
-    if torch.cuda.is_available():
-        # make sure CUDA states are stored on CPU so they are picklable
-        cuda_states = [s.cpu() for s in torch.cuda.get_rng_state_all()]
-        state["cuda"] = cuda_states
-    else:
-        state["cuda"] = None
-
-    return state
-
-
-def _set_local_rng_state(state: dict) -> None:
-    """Set local RNG states from the dict produced by _get_local_rng_state()."""
-    if state is None:
-        return
-
-    if "torch" in state and state["torch"] is not None:
-        torch.set_rng_state(state["torch"])
-    if "cuda" in state and state["cuda"] is not None and torch.cuda.is_available():
-        try:
-            # move back to CUDA tensors for this process and set them
-            cuda_states = [s.cuda() for s in state["cuda"]]
-            torch.cuda.set_rng_state_all(cuda_states)
-        except Exception:
-            # fallback: try setting per-device RNG if set_rng_state_all fails
-            for i, s in enumerate(state["cuda"]):
-                try:
-                    torch.cuda.set_rng_state(s.cuda(), device=i)
-                except Exception:
-                    # ignore if device mismatch
-                    pass
-
-    if "numpy" in state and state["numpy"] is not None:
-        np.random.set_state(state["numpy"])
-    if "random" in state and state["random"] is not None:
-        random.setstate(state["random"])
-
-
-def save_state(ckpt_path: str, epoch: Optional[int] = None, **named_objects: Any) -> None:
-    """
-    Save a checkpoint that includes:
-      - epoch (optional)
-      - state_dicts for provided named_objects
-      - rng_states: list of per-rank RNG dictionaries (one entry per world rank)
-
-    If torch.distributed is initialized the RNG states from all ranks are gathered and
-    stored in checkpoint["rng_states"] (list indexed by rank). Only rank 0 writes the file.
-    In single-process mode the checkpoint contains a single-item rng_states list.
-    """
-    os.makedirs(os.path.dirname(ckpt_path), exist_ok=True)
-
-    # prepare local RNG state
-    local_rng = _get_local_rng_state()
-
-    # distributed path: gather RNG states from all ranks
-    if dist.is_available() and dist.is_initialized():
-        rank = dist.get_rank()
-        world_size = dist.get_world_size()
-        all_states = [None] * world_size
-        # gather python objects (picklable)
-        dist.all_gather_object(all_states, local_rng)
-
-        # only rank 0 writes the checkpoint file
-        if rank == 0:
-            checkpoint = {}
-            if epoch is not None:
-                checkpoint["epoch"] = epoch
-            checkpoint["rng_states"] = all_states
-
-            # save provided objects' state_dicts (rank 0's local state_dicts)
-            for name, obj in named_objects.items():
-                checkpoint[name] = obj.state_dict()
-
-            torch.save(checkpoint, ckpt_path)
-
-        # ensure everyone waits until rank 0 finished writing
-        dist.barrier()
-
-    else:
-        # single-process fallback
-        checkpoint = {}
-        if epoch is not None:
-            checkpoint["epoch"] = epoch
-        checkpoint["rng_states"] = [local_rng]
-        for name, obj in named_objects.items():
-            checkpoint[name] = obj.state_dict()
-        torch.save(checkpoint, ckpt_path)
-
-
-def load_state(ckpt_path: str, **named_objects: Any) -> int:
-    """
-    Load checkpoint saved by save_state.
-
-    - Each process loads the same file and restores its own RNG state (checkpoint['rng_states'][rank]).
-    - Named objects that exist in the checkpoint will have their state_dict loaded.
-    - Returns epoch (int) if present, otherwise 0.
-    """
-    if not os.path.isfile(ckpt_path):
-        warnings.warn(f"Checkpoint file not found: {ckpt_path}")
-        return 0
-
-    # load on all ranks (shared FS assumed)
-    checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=False)
-    epoch = checkpoint.get("epoch", 0)
-
-    # load state_dicts into provided objects
-    for name, obj in named_objects.items():
-        if name in checkpoint:
-            try:
-                obj.load_state_dict(checkpoint[name])
-            except Exception as e:
-                warnings.warn(f"Failed to load state_dict for '{name}': {e}")
-        else:
-            warnings.warn(f"No state_dict found for '{name}' in checkpoint.")
-
-    # restore this rank's RNG state
-    rng_states = checkpoint.get("rng_states", None)
-    if rng_states is not None:
-        if dist.is_available() and dist.is_initialized():
-            rank = dist.get_rank()
-            if rank < len(rng_states):
-                my_state = rng_states[rank]
-            else:
-                my_state = None
-        else:
-            # single-process file: first element
-            my_state = rng_states[0] if len(rng_states) > 0 else None
-
-        try:
-            _set_local_rng_state(my_state)
-        except Exception as e:
-            warnings.warn(f"Failed to restore RNG state: {e}")
-
-    else:
-        warnings.warn("No 'rng_states' found in checkpoint; RNGs not restored.")
-
-    return epoch
-
-
-def extract_model_from_checkpoint_dinov2(checkpoint_path: str):
-    # Load the checkpoint
-    checkpoint = torch.load(
-        checkpoint_path, 
-        weights_only=False, 
-        map_location="cpu"
-    )
-    
-    # Get model state dict
-    model_state = checkpoint.get("model", checkpoint)
-
-    # Create output folder
-    output_dir = str(checkpoint_path).replace(".pt", "")
-    os.makedirs(output_dir, exist_ok=True)
-
-    # Quick check: compare the parameters of head_teacher_ibot vs head_teacher_dino
-    teacher_dino_keys = [k for k in model_state.keys() if k.startswith("head_teacher_dino.")]
-    teacher_ibot_keys = [k for k in model_state.keys() if k.startswith("head_teacher_ibot.")]
-
-    ibot_separate = True
-    if teacher_dino_keys and teacher_ibot_keys:
-        if all(torch.equal(model_state[dino_key], model_state[ibot_key]) \
-               for dino_key, ibot_key in zip(teacher_dino_keys, teacher_ibot_keys)):
-            ibot_separate = False  # Same weights → no separate ibot head
-    
-    # Define the components to save
-    components = {
-        "backbone_teacher.pt": "backbone_teacher.vit",
-        "backbone_student.pt": "backbone_student.vit",
-        "head_student_dino.pt": "head_student_dino",
-        "head_teacher_dino.pt": "head_teacher_dino"
-    }
-
-    # Add ibot heads only if separate
-    if ibot_separate:
-        components["head_student_ibot.pt"] = "head_student_ibot"
-        components["head_teacher_ibot.pt"] = "head_teacher_ibot"
-
-    # Extract and save each component
-    for filename, key in components.items():
-        sub_state_dict = {k.replace(f"{key}.", ""): v for k, v in model_state.items() if k.startswith(key)}
-        if not sub_state_dict:
-            print(f"[WARNING] No parameters found for {key}, skipping...")
-            continue
-        torch.save(sub_state_dict, os.path.join(output_dir, filename))
-    
-    print(f"Components extracted to: {output_dir}")
-
-
-def extract_model_from_checkpoint_siglip(checkpoint_path: str):
-    # Load the checkpoint
-    checkpoint = torch.load(
-        checkpoint_path, 
-        weights_only=False, 
-        map_location="cpu",
-    )
-    
-    # Get model state dict
-    model_state = checkpoint.get("model", checkpoint)
-
-    # Create output folder
-    output_dir = str(checkpoint_path).replace(".pt", "")
-    os.makedirs(output_dir, exist_ok=True)
-    
-    # Define the components to save
-    components = {
-        "backbone_image.pt": "backbone_image",
-        "backbone_text.pt": "backbone_text",
-        "feature_comb_image.pt": "feature_comb_image",
-        "projection_image.pt": "projection_image",
-        "projection_text.pt": "projection_text"
-    }
-
-    # Extract and save each component
-    for filename, key in components.items():
-        sub_state_dict = {k.replace(f"{key}.", ""): v for k, v in model_state.items() if k.startswith(key)}
-        if not sub_state_dict:
-            print(f"[WARNING] No parameters found for {key}, skipping...")
-            continue
-        torch.save(sub_state_dict, os.path.join(output_dir, filename))
-    
-    print(f"Components extracted to: {output_dir}")

spectre/utils/collate.py

@@ -1,120 +0,0 @@
-from typing import List, Callable, Optional
-
-import torch
-
-MONAI_IMPORT_ERROR = None
-try:
-    from monai.data import list_data_collate
-except ImportError as e:
-    list_data_collate = lambda x: x  # type: ignore
-    MONAI_IMPORT_ERROR = e
-
-
-def extended_collate_dino(samples_list: List) -> dict:
-    """
-    Applies MONAI's list_data_collate first and then extends it with DINOv2 masking logic.
-
-    Args:
-        samples_list: List of samples containing 'global_crops' and 'local_crops'.
-        mask_ratio: Tuple defining the range of masking ratios.
-        mask_probability: Probability of applying masking.
-        dtype: Data type to cast the collated tensors.
-        n_tokens: Number of tokens for masking.
-        mask_generator: Function to generate masks.
-
-    Returns:
-        A dictionary with collated global/local crops and corresponding masks.
-    """
-    if MONAI_IMPORT_ERROR is not None:
-        raise ImportError(
-            "MONAI is required to use extended_collate_dino but not installed. "
-            "Please install MONAI to use this collate function."
-        ) from MONAI_IMPORT_ERROR
-    
-    # Apply MONAI's list_data_collate
-    collated_data = list_data_collate(samples_list)
-
-    # Extract crops
-    global_views = torch.cat(collated_data["image_global_views"], dim=0)
-    local_views = torch.cat(collated_data["image_local_views"], dim=0)
-
-    return {
-        "global_views": global_views,
-        "local_views": local_views,
-    }
-    
-
-def extended_collate_siglip(
-    samples_list: List,
-    tokenizer: Optional[Callable] = None,
-    tokenizer_padding: bool = True,
-    tokenizer_truncation: bool = True,
-    tokenizer_max_length: Optional[int] = 1024,
-    return_filenames: bool = False
-) -> dict:
-    """
-    Applies SigLIP collate and then extends it with tokenization logic.
-    
-    Args:
-        samples_list: List of samples containing 'image' and 'report'.
-        tokenizer: Tokenizer function to apply on the reports.
-    
-    Returns:
-        A dictionary with collated images and tokenized text.
-    """
-    if MONAI_IMPORT_ERROR is not None:
-        raise ImportError(
-            "MONAI is required to use extended_collate_siglip but not installed. "
-            "Please install MONAI to use this collate function."
-        ) from MONAI_IMPORT_ERROR
-    
-    collated_data = list_data_collate(samples_list)
-
-    if return_filenames:
-        if "image" in collated_data.keys():
-            if (
-                hasattr(samples_list[0]["image"].data, "meta") 
-                and "filename_or_obj" in samples_list[0]["image"].data.meta
-            ):
-                collated_data["filename"] = [s["image"].data.meta["filename_or_obj"] for s in samples_list]
-
-    if tokenizer is not None and "report" in collated_data.keys():
-        tokenizer_output = tokenizer.batch_encode_plus(
-            collated_data["report"], 
-            add_special_tokens=True,
-            padding=tokenizer_padding,
-            truncation=tokenizer_truncation,
-            max_length=tokenizer_max_length,
-        )
-        
-        collated_data["input_ids"] = torch.tensor(tokenizer_output["input_ids"])
-        collated_data["attention_mask"] = torch.tensor(tokenizer_output["attention_mask"])
-
-    return collated_data
-
-
-def collate_add_filenames(samples_list: List) -> dict:
-    """
-    Applies MONAI's list_data_collate and adds filenames to the collated output.
-
-    Args:
-        samples_list: List of samples containing 'image' with metadata.
-    Returns:
-        A dictionary with collated images and filenames.
-    """
-    if MONAI_IMPORT_ERROR is not None:
-        raise ImportError(
-            "MONAI is required to use collate_add_filenames but not installed. "
-            "Please install MONAI to use this collate function."
-        ) from MONAI_IMPORT_ERROR
-    
-    collated_data = list_data_collate(samples_list)
-
-    if "image" in collated_data.keys():
-        if (
-            hasattr(samples_list[0]["image"].data, "meta") 
-            and "filename_or_obj" in samples_list[0]["image"].data.meta
-        ):
-            collated_data["filename"] = [s["image"].data.meta["filename_or_obj"] for s in samples_list]
-
-    return collated_data

spectre/utils/config.py

@@ -1,91 +0,0 @@
-import os
-import math
-
-from spectre.utils import _utils, distributed
-
-OMEGACONF_IMPORT_ERROR = None
-try:
-    from omegaconf import OmegaConf
-except ImportError as e:
-    OmegaConf = None  # type: ignore
-    OMEGACONF_IMPORT_ERROR = e
-
-
-def apply_scaling_rules_to_cfg(cfg):
-    """
-    Apply learing rate scaling rules to the configuration object.
-    """
-    base_lr = cfg.optim.base_lr
-    cfg.optim.lr = base_lr
-
-    # Apply scaling rules
-    if cfg.optim.scaling_rule == "constant":
-        return cfg
-    
-    try:
-        scaling_type, ref_batch_size = cfg.optim.scaling_rule.split("_wrt_")
-        ref_batch_size = float(ref_batch_size)
-    except ValueError:
-        raise NotImplementedError(f"Unknown scaling rule: {cfg.optim.scaling_rule}")
-    
-    scale_factor = cfg.train.batch_size_per_gpu * distributed.get_global_size()
-    scale_factor /= ref_batch_size
-    scale_factor *= cfg.train.grad_accum_steps
-    
-    if scaling_type == "sqrt":
-        cfg.optim.lr *= math.sqrt(scale_factor)
-    elif scaling_type == "linear":
-        cfg.optim.lr *= scale_factor
-    else:
-        raise NotImplementedError(f"Unsupported scaling type: {scaling_type}")
-    
-    return cfg
-
-
-def write_config(cfg, output_dir, name="config.yaml"):
-    if OMEGACONF_IMPORT_ERROR is not None:
-        raise ImportError(
-            "OmegaConf is required to use write_config but not installed. "
-            "Please install OmegaConf to use this function."
-        ) from OMEGACONF_IMPORT_ERROR
-    
-    saved_cfg_path = os.path.join(output_dir, name)
-    with open(saved_cfg_path, "w") as f:
-        OmegaConf.save(config=cfg, f=f)
-    return saved_cfg_path
-
-
-def get_cfg_from_args(args, default_config):
-    if OMEGACONF_IMPORT_ERROR is not None:
-        raise ImportError(
-            "OmegaConf is required to use get_cfg_from_args but not installed. "
-            "Please install OmegaConf to use this function."
-        ) from OMEGACONF_IMPORT_ERROR
-    
-    args.output_dir = os.path.abspath(args.output_dir)
-    args.opts = [] if args.opts is None else args.opts
-    args.opts += [f"train.output_dir={args.output_dir}"]
-    default_cfg = OmegaConf.create(default_config)
-    cfg = OmegaConf.load(args.config_file)
-    cfg = OmegaConf.merge(default_cfg, cfg, OmegaConf.from_cli(args.opts))
-    return cfg
-
-
-def random_seed(args):
-    seed = getattr(args, "seed", 0)
-    rank = distributed.get_global_rank()
-
-    _utils.fix_random_seeds(seed + rank)
-
-
-def setup(args, default_config):
-    """
-    Create configs and perform basic setups.
-    """
-    cfg = get_cfg_from_args(args, default_config)
-    os.makedirs(args.output_dir, exist_ok=True)
-    random_seed(args)
-    accelerator = distributed.init_distributed(cfg)
-    apply_scaling_rules_to_cfg(cfg)
-    write_config(cfg, args.output_dir)
-    return cfg, accelerator

spectre/utils/dataloader.py

@@ -1,126 +0,0 @@
-from __future__ import annotations
-import os
-from typing import Union, Callable, Optional, List
-
-import torch
-from torch.utils.data import ConcatDataset
-
-MONAI_IMPORT_ERROR = None
-try:
-    import monai.data as data
-except ImportError as e:
-    data = None  # type: ignore
-    MONAI_IMPORT_ERROR = e
-
-
-
-def get_dataloader(
-    datasets: Union[str, List[str]],
-    data_dir: str,
-    include_reports: bool = False,
-    include_labels: bool = False,
-    cache_dataset: bool = False,
-    cache_dir: Optional[str] = None,
-    use_gds: bool = False,
-    transform: Optional[Callable] = None,
-    fraction: float = 1.0,
-    batch_size: int = 64,
-    num_workers: int = 4,
-    pin_memory: bool = True,
-    shuffle: bool = True,
-    collate_fn: Optional[Callable] = None,
-    drop_last: bool = True,
-    persistent_workers: bool = True,
-    use_thread: bool = False,
-) -> "DataLoader":
-    """
-    Get dataloader for training.
-    """
-    if MONAI_IMPORT_ERROR is not None:
-        raise ImportError(
-            "MONAI is required to use get_dataloader but not installed. "
-            "Please install MONAI to use this function."
-        ) from MONAI_IMPORT_ERROR
-
-    if isinstance(datasets, str):
-        datasets = [datasets]
-
-    # Validate constraints
-    if include_reports:
-        assert set(datasets).issubset({"ct_rate", "merlin", "inspect"}), \
-            "When include_reports=True, only 'ct_rate', 'merlin', and 'inspect' are allowed."
-    if include_labels:
-        assert set(datasets).issubset({"abdomen_atlas", "abdomenct_1k"}), \
-            "When include_labels=True, only 'abdomen_atlas' and 'abdomenct_1k' are allowed."
-    if use_gds:
-        assert cache_dataset, "GDS requires cache_dataset=True."
-        assert torch.cuda.is_available(), "GDS requires CUDA to be available."
-
-    # Dataset configurations
-    DATASET_CONFIGS = {
-        "ct_rate": {"folder": "CT-RATE", "base_name": "CTRate",
-            "extra": {"include_reports": include_reports}},
-        "inspect": {"folder": "INSPECT", "base_name": "Inspect",
-            "extra": {"include_reports": include_reports}},
-        "merlin": {"folder": "MERLIN", "base_name": "Merlin",
-            "extra": {"include_reports": include_reports}},
-        "nlst": {"folder": "NLST", "base_name": "Nlst"},
-        "amos": {"folder": "Amos", "base_name": "Amos"},
-        "abdomen_atlas": {"folder": "AbdomenAtlas1.0Mini", "base_name": "AbdomenAtlas",
-            "extra": {"include_labels": include_labels}},
-        "panorama": {"folder": "PANORAMA", "base_name": "Panorama"},
-        "abdomenct_1k": {"folder": "AbdomenCT-1K", "base_name": "AbdomenCT1K",
-            "extra": {"include_labels": include_labels}},
-    }
-    
-    datasets_list = []
-    for ds in datasets:
-        if ds.lower() not in DATASET_CONFIGS:
-            raise NotImplementedError(f"Dataset {ds} not implemented.")
-
-        cfg = DATASET_CONFIGS[ds.lower()]
-        folder = cfg["folder"]
-        extra_args = cfg.get("extra", {})
-
-        kwargs = {
-            "data_dir": os.path.join(data_dir, folder),
-            "transform": transform,
-            "fraction": fraction,
-            **extra_args,
-        }
-
-        base_name = cfg["base_name"]
-        class_suffix = "Dataset"
-        if cache_dataset:
-            class_suffix = "GDSDataset" if use_gds else "PersistentDataset"
-
-        class_name = f"{base_name}{class_suffix}"
-        DatasetClass = getattr(__import__("spectre.data", fromlist=[class_name]), class_name)
-
-        if cache_dataset:
-            kwargs["cache_dir"] = os.path.join(cache_dir, folder)
-            if use_gds:
-                kwargs["device"] = torch.cuda.current_device()
-
-        datasets_list.append(DatasetClass(**kwargs))
-
-    dataset = datasets_list[0] if len(datasets_list) == 1 else ConcatDataset(datasets_list)
-
-    loader_cls = getattr(data, "ThreadDataLoader" if use_thread else "DataLoader")
-    loader_kwargs = {
-        "dataset": dataset,
-        "batch_size": batch_size,
-        "num_workers": num_workers,
-        "shuffle": shuffle,
-        "drop_last": drop_last,
-    }
-
-    if not use_thread:
-        loader_kwargs.update({
-            "pin_memory": pin_memory, 
-            "persistent_workers": persistent_workers
-        })
-    if collate_fn is not None:
-        loader_kwargs["collate_fn"] = collate_fn
-
-    return loader_cls(**loader_kwargs)

spectre/utils/distributed.py

@@ -1,92 +0,0 @@
-import os
-
-import torch.distributed as dist
-
-ACCELERATE_IMPORT_ERROR = None
-try:
-    from accelerate import Accelerator, DataLoaderConfiguration
-except ImportError as e:
-    Accelerator = None  # type: ignore
-    DataLoaderConfiguration = None  # type: ignore
-    ACCELERATE_IMPORT_ERROR = e
-
-
-def is_enabled() -> bool:
-    """
-    Returns:
-        True if distributed training is enabled
-    """
-    return dist.is_available() and dist.is_initialized()
-
-
-def get_global_size() -> int:
-    """
-    Returns:
-        Number of processes in the distributed group
-    """
-    if not is_enabled():
-        return 1
-    return dist.get_world_size()
-
-
-def get_global_rank() -> int:
-    """
-    Returns:
-        The rank of the current process in the distributed group
-    """
-    if not is_enabled():
-        return 0
-    return dist.get_rank()
-
-
-def get_local_size() -> int:
-    """
-    Returns:
-        Number of processes on the current machine
-    """
-    if not is_enabled():
-        return 1
-    return int(os.environ.get("LOCAL_SIZE", 1))
-
-
-def get_local_rank() -> int:
-    """
-    Returns:
-        The rank of the current process on the current machine
-    """
-    if not is_enabled():
-        return 0
-    return int(os.environ.get("LOCAL_RANK", 0))
-
-
-def init_distributed(cfg):
-    """
-    Initialize distributed training.
-    """
-    if ACCELERATE_IMPORT_ERROR is not None:
-        raise ImportError(
-            "Accelerate is required to use init_distributed but not installed. "
-            "Please install Accelerate to use this function."
-        ) from ACCELERATE_IMPORT_ERROR
-
-    # Initialize accelerator
-    dataloader_config = DataLoaderConfiguration(
-        non_blocking=cfg.train.pin_memory,
-    )
-    accelerator = Accelerator(
-        gradient_accumulation_steps=cfg.train.grad_accum_steps,
-        log_with="wandb" if cfg.train.log_wandb else None,
-        dataloader_config=dataloader_config,
-    )
-
-    # Initialize wandb
-    if cfg.train.log_wandb:
-        accelerator.init_trackers(
-            project_name="spectre",
-            config={k: v for d in cfg.values() for k, v in d.items()},
-            init_kwargs={
-                "dir": os.path.join(cfg.train.output_dir, "logs"),
-            },
-        )
-    
-    return accelerator

spectre/utils/lora.py

@@ -1,38 +0,0 @@
-import torch.nn as nn
-import loralib as lora
-
-
-def add_lora_adapters(
-        root_module: nn.Module,
-        r: int = 8,
-        lora_alpha: int = 32,
-        lora_dropout: float = 0.05,
-        target_keywords: tuple[str, ...] = ("q_proj", "k_proj", "v_proj", "o_proj")
-    ) -> None:
-    """
-    Recursively traverses the model and replaces every `nn.Linear`
-    whose name contains one of `target_keywords` with a LoRA-augmented
-    linear layer from loralib.
-    """
-
-    for name, child in list(root_module.named_children()):
-        # If the child is itself a container, recurse first
-        add_lora_adapters(child, r, lora_alpha, lora_dropout, target_keywords)
-
-        # Replace target linear layers
-        if isinstance(child, nn.Linear) and any(k in name for k in target_keywords):
-            lora_layer = lora.Linear(                             # loralib wrapper
-                in_features=child.in_features,
-                out_features=child.out_features,
-                r=r,
-                lora_alpha=lora_alpha,
-                lora_dropout=lora_dropout,
-                bias=child.bias is not None,
-            )
-
-            # copy original weights so that behaviour is identical pre-training
-            lora_layer.weight.data = child.weight.data.clone()
-            if child.bias is not None:
-                lora_layer.bias.data = child.bias.data.clone()
-
-            setattr(root_module, name, lora_layer)               # hot-swap!

spectre/utils/masking.py

@@ -1,196 +0,0 @@
-import math
-from typing import Optional, Tuple, Union
-
-import torch
-
-
-def _random_block_mask(
-    size: Tuple[int, int, int],
-    num_masks: int,
-    min_num_masks_per_block: int = 4,
-    max_num_masks_per_block: Optional[int] = None,
-    max_attempts_per_block: int = 10,
-    generator: Optional[torch.Generator] = None,
-    device: Optional[Union[torch.device, str]] = None,
-) -> torch.Tensor:
-    """3D helper: generate a (H, W, D) boolean mask by placing cuboidal blocks.
-
-    - size: (H, W, D)
-    - num_masks: target total number of masked voxels for this image
-    - min_num_masks_per_block / max_num_masks_per_block: voxel-range per block
-    """
-    H, W, D = size
-    total = H * W * D
-    num_masks = min(max(0, int(num_masks)), total)
-
-    if max_num_masks_per_block is None:
-        max_num_masks_per_block = max(1, num_masks)
-
-    mask = torch.zeros((H, W, D), dtype=torch.bool, device=device)
-    masked_count = 0
-    global_attempts = 0
-
-    orders = [(0, 1, 2), (1, 2, 0), (2, 0, 1)]
-
-    # Try to place blocks until we have enough masked voxels or we exceed attempts
-    while masked_count < num_masks and global_attempts < max_attempts_per_block:
-        global_attempts += 1
-
-        # choose target voxels for this block
-        target_voxels = int(torch.randint(
-            min_num_masks_per_block, max_num_masks_per_block + 1, (1,), generator=generator
-        ).item())
-
-        found = False
-        local_attempts = 0
-        while not found and local_attempts < max_attempts_per_block:
-            local_attempts += 1
-
-            # random pick order for dims to reduce bias
-            order_idx = int(torch.randint(0, 3, (1,), generator=generator).item())
-            order = orders[order_idx]
-
-            # pick first dimension
-            if order[0] == 0:
-                h = int(torch.randint(1, min(H, target_voxels) + 1, (1,), generator=generator).item())
-            elif order[0] == 1:
-                w = int(torch.randint(1, min(W, target_voxels) + 1, (1,), generator=generator).item())
-            else:
-                d = int(torch.randint(1, min(D, target_voxels) + 1, (1,), generator=generator).item())
-
-            # progressively choose remaining dims while ensuring feasibility
-            try:
-                if order[0] == 0:
-                    # h chosen -> pick w then compute d_needed
-                    max_w = max(1, min(W, target_voxels // h))
-                    w = int(torch.randint(1, max_w + 1, (1,), generator=generator).item())
-                    d_needed = math.ceil(target_voxels / (h * w))
-                    if d_needed <= D:
-                        d = max(1, d_needed)
-                        found = True
-                elif order[0] == 1:
-                    # w chosen -> pick d then compute h_needed
-                    max_d = max(1, min(D, target_voxels // w))
-                    d = int(torch.randint(1, max_d + 1, (1,), generator=generator).item())
-                    h_needed = math.ceil(target_voxels / (d * w))
-                    if h_needed <= H:
-                        h = max(1, h_needed)
-                        found = True
-                else:
-                    # d chosen -> pick h then compute w_needed
-                    max_h = max(1, min(H, target_voxels // d))
-                    h = int(torch.randint(1, max_h + 1, (1,), generator=generator).item())
-                    w_needed = math.ceil(target_voxels / (d * h))
-                    if w_needed <= W:
-                        w = max(1, w_needed)
-                        found = True
-
-            except ValueError:
-                # in case of invalid ranges (defensive); just continue trying
-                continue
-
-            # fallback alternative attempt: try simple factorization heuristics
-            if not found:
-                # attempt small-to-large factorization
-                for hh in range(1, min(H, target_voxels) + 1):
-                    for ww in range(1, min(W, target_voxels // hh) + 1):
-                        dd = math.ceil(target_voxels / (hh * ww))
-                        if dd <= D:
-                            h, w, d = hh, ww, dd
-                            found = True
-                            break
-                    if found:
-                        break
-
-        if not found:
-            # couldn't find a fitting block this global attempt; move on
-            continue
-
-        # clamp block dims to volume just in case and ensure at least 1
-        h = min(max(1, int(h)), H)
-        w = min(max(1, int(w)), W)
-        d = min(max(1, int(d)), D)
-
-        # choose random location so block fits
-        x0 = int(torch.randint(0, (H - h) + 1, (1,), generator=generator).item()) if H - h > 0 else 0
-        y0 = int(torch.randint(0, (W - w) + 1, (1,), generator=generator).item()) if W - w > 0 else 0
-        z0 = int(torch.randint(0, (D - d) + 1, (1,), generator=generator).item()) if D - d > 0 else 0
-
-        mask[x0 : x0 + h, y0 : y0 + w, z0 : z0 + d] = True
-        masked_count = int(mask.sum().item())
-
-    # If still short, fill remaining voxels at random positions
-    if masked_count < num_masks:
-        remaining = num_masks - masked_count
-        indices = torch.nonzero(~mask, as_tuple=False)
-        if indices.numel() > 0:
-            perm = torch.randperm(indices.shape[0], generator=generator, device=mask.device)
-            pick = indices[perm[:remaining]]
-            mask[pick[:, 0], pick[:, 1], pick[:, 2]] = True
-
-    return mask
-
-
-def random_block_mask(
-    size: Tuple[int, int, int, int],
-    batch_mask_ratio: float = 0.5,
-    min_image_mask_ratio: float = 0.1,
-    max_image_mask_ratio: float = 0.5,
-    min_num_masks_per_block: int = 4,
-    max_num_masks_per_block: Optional[int] = None,
-    max_attempts_per_block: int = 10,
-    generator: Optional[torch.Generator] = None,
-    device: Optional[Union[torch.device, str]] = None,
-) -> torch.Tensor:
-    """Create random block masks for 3D volumes only.
-
-    Args:
-        size: (B, H, W, D)
-        batch_mask_ratio: fraction of images in the batch to apply masking to
-        min_image_mask_ratio / max_image_mask_ratio: per-image mask fraction range
-        min_num_masks_per_block / max_num_masks_per_block: voxels per block range
-        max_attempts_per_block: attempts to find a fitting block
-        generator: optional torch.Generator for reproducibility.
-        device: device for returned tensor
-
-    Returns:
-        boolean tensor with shape (B, H, W, D)
-    """
-    if len(size) != 4:
-        raise ValueError("size must be (B, H, W, D) for 3D masking.")
-
-    B, H, W, D = size
-
-    if max_image_mask_ratio < min_image_mask_ratio:
-        raise ValueError("max_image_mask_ratio must be >= min_image_mask_ratio.")
-
-    num_images_masked = int(B * batch_mask_ratio)
-    probs = torch.linspace(min_image_mask_ratio, max_image_mask_ratio, num_images_masked + 1).tolist()
-
-    image_masks = []
-    total_voxels = H * W * D
-
-    for prob_min, prob_max in zip(probs[:-1], probs[1:]):
-        # choose number of masked voxels for this image
-        u = float(prob_min + (prob_max - prob_min) * torch.rand(1, generator=generator).item())
-        num_mask = int(total_voxels * u)
-        image_masks.append(
-            _random_block_mask(
-                size=(H, W, D),
-                num_masks=num_mask,
-                min_num_masks_per_block=min_num_masks_per_block,
-                max_num_masks_per_block=max_num_masks_per_block,
-                max_attempts_per_block=max_attempts_per_block,
-                generator=generator,
-                device=device,
-            )
-        )
-
-    # Add non-masked images (all False) to fill the batch
-    for _ in range(num_images_masked, B):
-        image_masks.append(torch.zeros((H, W, D), dtype=torch.bool, device=device))
-
-    perm = torch.randperm(B, generator=generator).tolist()
-    image_masks = [image_masks[i] for i in perm]
-    
-    return torch.stack(image_masks)

spectre/utils/param_groups.py

@@ -1,118 +0,0 @@
-import torch.nn as nn
-
-
-def get_vit_lr_decay_rate(
-    name: str,
-    llrd_factor: float = 1.0,
-    num_layers: int = 12,
-    force_is_backbone: bool = False,
-    shift: int = 0,
-) -> float:
-    """
-    Get the layer-wise learning rate decay (LLRD) rate for a given parameter name.
-
-    Args:
-        name:
-            The name of the parameter.
-        llrd_factor:
-            The decay factor for each layer.
-        num_layers:
-            The total number of layers in the model.
-        force_is_backbone:
-            If True, forces the function to treat the parameter as part of the backbone.
-        shift:
-            An integer to shift the layer ids, useful when combining multiple modules.
-
-    Returns:
-        The learning rate multiplier for the parameter.
-    """
-    layer_id = num_layers + 1
-    if name.startswith("backbone") or force_is_backbone:
-        if (
-            ".pos_embed" in name
-            or ".patch_embed" in name
-            or ".patch_proj" in name
-            or ".mask_token" in name
-            or ".cls_token" in name
-            or ".reg_token" in name
-        ):
-            layer_id = 0
-        elif ".blocks." in name:
-            layer_id = int(name[name.find(".blocks.") :].split(".")[2]) + 1 + shift
-
-    return llrd_factor ** (num_layers + 1 - layer_id)
-
-
-def get_param_groups_with_decay(
-    model: nn.Module,
-    llrd_factor: float = 1.0,
-    patch_embed_lr_mult: float = 1.0,
-    projection_head_wd_mult: float = 1.0,
-    lora_lr_factor: float = 1.0,
-    num_layers: int | None = None,
-):
-    
-    force_is_backbone = False
-    shift = 0
-    if num_layers is not None:
-        num_layers = num_layers
-    elif hasattr(model, "n_blocks"):
-        num_layers = model.n_blocks
-        force_is_backbone = True
-    elif hasattr(model, "blocks"):
-        num_layers = len(model.blocks)
-        force_is_backbone = True
-    elif hasattr(model, "backbone") and hasattr(model.backbone, "blocks"):
-        num_layers = len(model.backbone.blocks)
-    elif hasattr(model, "backbone_student") and hasattr(model.backbone_student, "blocks"):  # DINO specific
-        num_layers = len(model.backbone_student.blocks)
-    elif hasattr(model, "backbone_student") and hasattr(model.backbone_student, "vit") and hasattr(model.backbone_student.vit, "blocks"):  # DINOv2 specific
-        num_layers = len(model.backbone_student.vit.blocks)
-    elif hasattr(model, "backbone_image") and hasattr(model.backbone_image, "blocks"):  # SigLIP specific
-        if not hasattr(model, "feature_comb_image") or model.feature_comb_image is None:
-            num_layers = len(model.backbone_image.blocks)
-        else:
-            num_layers = len(model.backbone_image.blocks) + len(model.feature_comb_image.blocks)
-            shift = len(model.backbone_image.blocks)
-            force_is_backbone = True
-    else:
-        num_layers = 0
-    
-    all_param_groups = []
-    for n, p in model.named_parameters():
-        if not p.requires_grad:
-            continue
-        if not "lora_" in n:
-            s = shift if "feature_comb" in n else 0
-            llrd_rate = get_vit_lr_decay_rate(
-                n, llrd_factor, num_layers, force_is_backbone, s,
-            )
-
-            d = {
-                "name": n,
-                "params": p,
-                "lr_mult": llrd_rate,
-                "wd_mult": 1.0,
-            }
-
-            if "head" in n or "projection" in n:
-                d["wd_mult"] = projection_head_wd_mult
-
-            # No weight-decay on biases, norm parameters, layer scale gamma, learned tokens and embeddings
-            if n.endswith("bias") or "norm" in n or "gamma" in n or "fourrier_w" in n:
-                d["wd_mult"] = 0.0
-
-            if "patch_embed" in n:
-                d["lr_mult"] *= patch_embed_lr_mult
-        
-        else:
-            # LoRA parameters
-            d = {
-                "name": n,
-                "params": p,
-                "lr_mult": lora_lr_factor,
-                "wd_mult": 1.0,
-            }
-        
-        all_param_groups.append(d)
-    return all_param_groups

spectre/utils/scheduler.py

@@ -1,236 +0,0 @@
-import warnings
-from typing import Optional
-
-import numpy as np
-import torch
-
-
-def linear_warmup_schedule(
-    step: int,
-    warmup_steps: int,
-    start_value: float,
-    end_value: float,
-) -> float:
-    if warmup_steps < 0:
-        raise ValueError(f"Warmup steps {warmup_steps} can't be negative.")
-    if step < 0:
-        raise ValueError(f"Current step number {step} can't be negative.")
-    if start_value < 0:
-        raise ValueError(f"Start value {start_value} can't be negative.")
-    if end_value <= 0:
-        raise ValueError(f"End value {end_value} can't be non-positive.")
-    if start_value > end_value:
-        raise ValueError(
-            f"Start value {start_value} must be less than or equal to end value {end_value}."
-        )
-    if step < warmup_steps:
-        return start_value + step / warmup_steps * (end_value - start_value)
-    else:
-        return end_value
-    
-
-def cosine_schedule(
-    step: int,
-    max_steps: int,
-    start_value: float,
-    end_value: float,
-    period: Optional[int] = None,
-) -> float:
-    """Use cosine decay to gradually modify start_value to reach target end_value.
-
-    Args:
-        step:
-            Current step number.
-        max_steps:
-            Total number of steps.
-        start_value:
-            Starting value.
-        end_value:
-            Target value.
-        period:
-            The number of steps over which the cosine function completes a full cycle.
-            Defaults to max_steps.
-
-    Returns:
-        Cosine decay value.
-
-    """
-    if step < 0:
-        raise ValueError(f"Current step number {step} can't be negative.")
-    if max_steps < 1:
-        raise ValueError(f"Total step number {max_steps} must be >= 1.")
-    if period is None and step > max_steps:
-        warnings.warn(
-            f"Current step number {step} exceeds max_steps {max_steps}.",
-            category=RuntimeWarning,
-        )
-    if period is not None and period <= 0:
-        raise ValueError(f"Period {period} must be >= 1")
-
-    decay: float
-    if period is not None:  # "cycle" based on period, if provided
-        decay = (
-            end_value
-            - (end_value - start_value) * (np.cos(2 * np.pi * step / period) + 1) / 2
-        )
-    elif max_steps == 1:
-        # Avoid division by zero
-        decay = end_value
-    elif step == max_steps:
-        # Special case for Pytorch Lightning which updates LR scheduler also for epoch
-        # after last training epoch.
-        decay = end_value
-    else:
-        decay = (
-            end_value
-            - (end_value - start_value)
-            * (np.cos(np.pi * step / (max_steps - 1)) + 1)
-            / 2
-        )
-    return decay
-
-
-def cosine_warmup_schedule(
-    step: int,
-    max_steps: int,
-    start_value: float,
-    end_value: float,
-    warmup_steps: int,
-    warmup_start_value: float,
-    warmup_end_value: Optional[float] = None,
-    period: Optional[int] = None,
-) -> float:
-    """Use cosine decay to gradually modify start_value to reach target end_value.
-
-    Uses linear warmup for the first warmup_steps steps.
-
-    Args:
-        step:
-            Current step number.
-        max_steps:
-            Total number of steps.
-        start_value:
-            Starting value.
-        end_value:
-            Target value.
-        warmup_steps:
-            Number of steps for warmup.
-        warmup_start_value:
-            Starting value for warmup.
-        warmup_end_value:
-            Target value for warmup. Defaults to start_value.
-        period:
-            The number of steps over which the cosine function completes a full cycle.
-            Defaults to max_steps - warmup_steps.
-    Returns:
-        Cosine decay value.
-    """
-    if warmup_steps < 0:
-        raise ValueError(f"Warmup steps {warmup_steps} can't be negative.")
-    if warmup_steps > max_steps:
-        raise ValueError(f"Warmup steps {warmup_steps} must be <= max_steps.")
-    if step > max_steps:
-        warnings.warn(
-            f"Current step number {step} exceeds max_steps {max_steps}.",
-            category=RuntimeWarning,
-        )
-
-    if warmup_end_value is None:
-        warmup_end_value = start_value
-
-    if step < warmup_steps:
-        # Use step + 1 to reach warmup_end_value at end of warmup. This means that the
-        # initial warmup_start_value is skipped which is oftentimes desired when setting
-        # it to 0 as this would result in no parameter updates.
-        return (
-            warmup_start_value
-            + (warmup_end_value - warmup_start_value) * (step + 1) / warmup_steps
-        )
-    else:
-        max_steps = max_steps - warmup_steps if period is None else 1
-        return cosine_schedule(
-            step=step - warmup_steps,
-            max_steps=max_steps,
-            start_value=start_value,
-            end_value=end_value,
-            period=period,
-        )
-
-
-class CosineWarmupScheduler(torch.optim.lr_scheduler.LambdaLR):
-    """Cosine warmup scheduler for learning rate.
-
-    Args:
-        optimizer:
-            Optimizer object to schedule the learning rate.
-        warmup_epochs:
-            Number of warmup epochs or steps.
-        max_epochs:
-            Total number of training epochs or steps.
-        last_epoch:
-            The index of last epoch or step.
-        start_value:
-            Starting learning rate.
-        end_value:
-            Target learning rate.
-        verbose:
-            If True, prints a message to stdout for each update.
-        warmup_start_value:
-            Starting learning rate for warmup.
-        warmup_end_value:
-            Target learning rate for warmup. Defaults to start_value.
-
-    Note: The `epoch` arguments do not necessarily have to be epochs. Any step or index
-    can be used. The naming follows the PyTorch convention to use `epoch` for the steps
-    in the scheduler.
-    """
-
-    def __init__(
-        self,
-        optimizer: torch.optim.Optimizer,
-        warmup_epochs: int,
-        max_epochs: int,
-        last_epoch: int = -1,
-        start_value: float = 1.0,
-        end_value: float = 0.001,
-        period: Optional[int] = None,
-        verbose: bool = False,
-        warmup_start_value: float = 0.0,
-        warmup_end_value: Optional[float] = None,
-    ) -> None:
-        self.warmup_epochs = warmup_epochs
-        self.max_epochs = max_epochs
-        self.start_value = start_value
-        self.end_value = end_value
-        self.period = period
-        self.warmup_start_value = warmup_start_value
-        self.warmup_end_value = warmup_end_value
-
-        super().__init__(
-            optimizer=optimizer,
-            lr_lambda=self.scale_lr,
-            last_epoch=last_epoch,
-            verbose=verbose,
-        )
-
-    def scale_lr(self, epoch: int) -> float:
-        """Scale learning rate according to the current epoch number.
-
-        Args:
-            epoch:
-                Current epoch number.
-
-        Returns:
-            Scaled learning rate.
-
-        """
-        return cosine_warmup_schedule(
-            step=epoch,
-            max_steps=self.max_epochs,
-            start_value=self.start_value,
-            end_value=self.end_value,
-            warmup_steps=self.warmup_epochs,
-            warmup_start_value=self.warmup_start_value,
-            warmup_end_value=self.warmup_end_value,
-            period=self.period,
-        )