diff --git a/docs/source/lightly.loss.rst b/docs/source/lightly.loss.rst index 2d2ce0ff1..69272696c 100644 --- a/docs/source/lightly.loss.rst +++ b/docs/source/lightly.loss.rst @@ -56,6 +56,9 @@ lightly.loss .. autoclass:: lightly.loss.regularizer.co2.CO2Regularizer :members: +.. autoclass:: lightly.loss.sparse_spark.SparKPatchReconLoss + :members: + .. autoclass:: lightly.loss.swav_loss.SwaVLoss :members: diff --git a/examples/notebooks/pytorch_lightning/spark.ipynb b/examples/notebooks/pytorch_lightning/spark.ipynb new file mode 100644 index 000000000..d8df49907 --- /dev/null +++ b/examples/notebooks/pytorch_lightning/spark.ipynb @@ -0,0 +1,302 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "0", + "metadata": {}, + "source": [ + "This example requires the following dependencies to be installed:\n", + "pip install \"lightly[timm]\"" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "1", + "metadata": {}, + "outputs": [], + "source": [ + "!pip install \"lightly[timm]\"" + ] + }, + { + "cell_type": "markdown", + "id": "2", + "metadata": {}, + "source": [ + "Note: The model and training settings do not follow the reference settings\n", + "from the paper. The settings are chosen such that the example can easily be\n", + "run on a small dataset with a single GPU." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "3", + "metadata": {}, + "outputs": [], + "source": [ + "import pytorch_lightning as pl\n", + "import timm\n", + "import torch\n", + "import torchvision\n", + "from pytorch_lightning import LightningModule\n", + "from torch import Tensor\n", + "from torch.nn import Module\n", + "from torchvision.transforms import v2" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4", + "metadata": {}, + "outputs": [], + "source": [ + "import lightly.models.utils as model_utils\n", + "from lightly.loss.sparse_spark import SparKPatchReconLoss\n", + "from lightly.models.modules import sparse_spark" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5", + "metadata": {}, + "outputs": [], + "source": [ + "## The global projection head is the same as the Barlow Twins one\n", + "from lightly.models.modules.sparse_spark import (\n", + " SparKDensifier,\n", + " SparKMasker,\n", + " SparKMaskingOutput,\n", + " SparKOutputDecoder,\n", + " UNetDecoder,\n", + " sparse_layer_context,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "6", + "metadata": {}, + "outputs": [], + "source": [ + "def _get_downsample_ratio_from_timm_model(model: Module) -> int:\n", + " if not hasattr(model, \"feature_info\"):\n", + " raise ValueError(\n", + " \"The provided model does not have the required 'feature_info' attribute.\"\n", + " )\n", + " return model.feature_info[-1][\"reduction\"]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "7", + "metadata": {}, + "outputs": [], + "source": [ + "def _get_enc_feat_map_chs_from_timm_model(model: Module) -> list[int]:\n", + " if not hasattr(model, \"feature_info\"):\n", + " raise ValueError(\n", + " \"The provided model does not have the required 'feature_info' attribute.\"\n", + " )\n", + " return [fi[\"num_chs\"] for fi in model.feature_info]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "8", + "metadata": {}, + "outputs": [], + "source": [ + "class SparseSparK(LightningModule):\n", + " def __init__(\n", + " self,\n", + " input_size: int = 416,\n", + " mask_ratio: float = 0.6,\n", + " densify_norm: str = \"bn\",\n", + " sbn=False,\n", + " ):\n", + " super().__init__()\n", + " backbone = timm.create_model(\n", + " model_name=\"resnet18\", drop_path_rate=0.05, features_only=True\n", + " )\n", + " downsample_ratio = _get_downsample_ratio_from_timm_model(backbone)\n", + " enc_feat_map_chs = _get_enc_feat_map_chs_from_timm_model(backbone)\n", + " self.sparse_encoder = sparse_spark.dense_model_to_sparse(\n", + " m=backbone, sbn=sbn, verbose=True\n", + " )\n", + " self.fmap_h = input_size // downsample_ratio\n", + " self.fmap_w = input_size // downsample_ratio\n", + " self.dense_decoder = UNetDecoder(\n", + " up_sample_ratio=downsample_ratio,\n", + " width=enc_feat_map_chs[-1],\n", + " )\n", + " self.masker = SparKMasker(\n", + " feature_map_size=(self.fmap_h, self.fmap_w),\n", + " downsample_ratio=downsample_ratio,\n", + " mask_ratio=mask_ratio,\n", + " )\n", + " self.densifier = SparKDensifier(\n", + " encoder_in_channels=enc_feat_map_chs,\n", + " decoder_in_channel=self.dense_decoder.width,\n", + " densify_norm_str=densify_norm.lower(),\n", + " sbn=sbn,\n", + " )\n", + " self.downsample_ratio = downsample_ratio\n", + " # loss module for patch reconstruction\n", + " self.recon_loss_fn = SparKPatchReconLoss()\n", + " # output decoder for visualization (pass minimal spatial props)\n", + " self.output_decoder = SparKOutputDecoder(\n", + " fmap_h=self.fmap_h,\n", + " fmap_w=self.fmap_w,\n", + " downsample_ratio=downsample_ratio,\n", + " )\n", + "\n", + " def forward(\n", + " self,\n", + " inp_bchw: Tensor,\n", + " vis=False,\n", + " ):\n", + " # step1. Mask\n", + " mask_out: SparKMaskingOutput = self.masker(inp_bchw)\n", + " masked_bchw, per_level_mask = mask_out\n", + " active_b1fHfW = per_level_mask[0]\n", + " active_b1hw = per_level_mask[-1]\n", + " # step2. Encode: get hierarchical encoded sparse features (a list containing 4 feature maps at 4 scales)\n", + " # Use sparse_layer_context to provide the mask to the sparse encoder and densifier.\n", + " with sparse_layer_context(active_mask=active_b1fHfW):\n", + " fea_bcffs: list[Tensor] = self.sparse_encoder(masked_bchw)\n", + " # step3. Densify: get hierarchical dense features for decoding\n", + " to_dec = self.densifier(fea_bcffs)\n", + " # step4. Decode and reconstruct\n", + " rec_bchw = self.dense_decoder(to_dec)\n", + " inp, rec = (\n", + " model_utils.patchify(inp_bchw, self.downsample_ratio),\n", + " model_utils.patchify(rec_bchw, self.downsample_ratio),\n", + " ) # inp and rec: (B, L = f*f, N = C*downsample_ratio**2)\n", + "\n", + " recon_loss, mean, var = self.recon_loss_fn(\n", + " inp_patches=inp, rec_patches=rec, active_mask=active_b1fHfW\n", + " )\n", + "\n", + " if vis:\n", + " return self.output_decoder(\n", + " rec_patches=rec,\n", + " mean=mean,\n", + " var=var,\n", + " inp_bchw=inp_bchw,\n", + " active_mask_full=active_b1hw,\n", + " )\n", + " else:\n", + " return recon_loss\n", + "\n", + " def training_step(self, batch: tuple[Tensor, Tensor], batch_index: int) -> Tensor:\n", + " img, _ = batch\n", + " recon_loss = self.forward(img)\n", + " # Log the training loss to logger and progress bar (per-step and per-epoch)\n", + " self.log(\n", + " \"train_loss\",\n", + " recon_loss,\n", + " on_step=True,\n", + " on_epoch=True,\n", + " prog_bar=True,\n", + " logger=True,\n", + " )\n", + " return recon_loss\n", + "\n", + " def configure_optimizers(self):\n", + " return torch.optim.SGD(\n", + " self.parameters(), lr=0.03, momentum=0.9, weight_decay=1e-4\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "9", + "metadata": {}, + "outputs": [], + "source": [ + "model = SparseSparK(input_size=416)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "10", + "metadata": {}, + "outputs": [], + "source": [ + "dataset = torchvision.datasets.Caltech101(\n", + " \"datasets/caltech101\",\n", + " download=True,\n", + " transform=v2.Compose(\n", + " [\n", + " v2.Resize((416, 416)),\n", + " v2.RGB(),\n", + " v2.ToTensor(),\n", + " v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n", + " ]\n", + " ),\n", + ")\n", + "# or create a dataset from a folder containing images or videos:\n", + "# dataset = LightlyDataset(\"path/to/folder\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "11", + "metadata": {}, + "outputs": [], + "source": [ + "dataloader = torch.utils.data.DataLoader(\n", + " dataset,\n", + " batch_size=4,\n", + " shuffle=True,\n", + " drop_last=True,\n", + " num_workers=8,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "12", + "metadata": {}, + "outputs": [], + "source": [ + "trainer = pl.Trainer(\n", + " max_epochs=30,\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "13", + "metadata": {}, + "outputs": [], + "source": [ + "trainer.fit(\n", + " model=model,\n", + " train_dataloaders=dataloader,\n", + ")" + ] + } + ], + "metadata": { + "jupytext": { + "cell_metadata_filter": "-all", + "main_language": "python", + "notebook_metadata_filter": "-all" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/examples/pytorch_lightning/spark.py b/examples/pytorch_lightning/spark.py new file mode 100644 index 000000000..62101dd5f --- /dev/null +++ b/examples/pytorch_lightning/spark.py @@ -0,0 +1,182 @@ +# This example requires the following dependencies to be installed: +# pip install "lightly[timm]" + +# Note: The model and training settings do not follow the reference settings +# from the paper. The settings are chosen such that the example can easily be +# run on a small dataset with a single GPU. + +import pytorch_lightning as pl +import timm +import torch +import torchvision +from pytorch_lightning import LightningModule +from torch import Tensor +from torch.nn import Module +from torchvision.transforms import v2 + +import lightly.models.utils as model_utils +from lightly.loss.sparse_spark import SparKPatchReconLoss +from lightly.models.modules import sparse_spark + +## The global projection head is the same as the Barlow Twins one +from lightly.models.modules.sparse_spark import ( + SparKDensifier, + SparKMasker, + SparKMaskingOutput, + SparKOutputDecoder, + UNetDecoder, + sparse_layer_context, +) + + +def _get_downsample_ratio_from_timm_model(model: Module) -> int: + if not hasattr(model, "feature_info"): + raise ValueError( + "The provided model does not have the required 'feature_info' attribute." + ) + return model.feature_info[-1]["reduction"] + + +def _get_enc_feat_map_chs_from_timm_model(model: Module) -> list[int]: + if not hasattr(model, "feature_info"): + raise ValueError( + "The provided model does not have the required 'feature_info' attribute." + ) + return [fi["num_chs"] for fi in model.feature_info] + + +class SparseSparK(LightningModule): + def __init__( + self, + input_size: int = 416, + mask_ratio: float = 0.6, + densify_norm: str = "bn", + sbn=False, + ): + super().__init__() + backbone = timm.create_model( + model_name="resnet18", drop_path_rate=0.05, features_only=True + ) + downsample_ratio = _get_downsample_ratio_from_timm_model(backbone) + enc_feat_map_chs = _get_enc_feat_map_chs_from_timm_model(backbone) + self.sparse_encoder = sparse_spark.dense_model_to_sparse( + m=backbone, sbn=sbn, verbose=True + ) + self.fmap_h = input_size // downsample_ratio + self.fmap_w = input_size // downsample_ratio + self.dense_decoder = UNetDecoder( + up_sample_ratio=downsample_ratio, + width=enc_feat_map_chs[-1], + ) + self.masker = SparKMasker( + feature_map_size=(self.fmap_h, self.fmap_w), + downsample_ratio=downsample_ratio, + mask_ratio=mask_ratio, + ) + self.densifier = SparKDensifier( + encoder_in_channels=enc_feat_map_chs, + decoder_in_channel=self.dense_decoder.width, + densify_norm_str=densify_norm.lower(), + sbn=sbn, + ) + self.downsample_ratio = downsample_ratio + # loss module for patch reconstruction + self.recon_loss_fn = SparKPatchReconLoss() + # output decoder for visualization (pass minimal spatial props) + self.output_decoder = SparKOutputDecoder( + fmap_h=self.fmap_h, + fmap_w=self.fmap_w, + downsample_ratio=downsample_ratio, + ) + + def forward( + self, + inp_bchw: Tensor, + vis=False, + ): + # step1. Mask + mask_out: SparKMaskingOutput = self.masker(inp_bchw) + masked_bchw, per_level_mask = mask_out + active_b1fHfW = per_level_mask[0] + active_b1hw = per_level_mask[-1] + # step2. Encode: get hierarchical encoded sparse features (a list containing 4 feature maps at 4 scales) + # Use sparse_layer_context to provide the mask to the sparse encoder and densifier. + with sparse_layer_context(active_mask=active_b1fHfW): + fea_bcffs: list[Tensor] = self.sparse_encoder(masked_bchw) + # step3. Densify: get hierarchical dense features for decoding + to_dec = self.densifier(fea_bcffs) + # step4. Decode and reconstruct + rec_bchw = self.dense_decoder(to_dec) + inp, rec = ( + model_utils.patchify(inp_bchw, self.downsample_ratio), + model_utils.patchify(rec_bchw, self.downsample_ratio), + ) # inp and rec: (B, L = f*f, N = C*downsample_ratio**2) + + recon_loss, mean, var = self.recon_loss_fn( + inp_patches=inp, rec_patches=rec, active_mask=active_b1fHfW + ) + + if vis: + return self.output_decoder( + rec_patches=rec, + mean=mean, + var=var, + inp_bchw=inp_bchw, + active_mask_full=active_b1hw, + ) + else: + return recon_loss + + def training_step(self, batch: tuple[Tensor, Tensor], batch_index: int) -> Tensor: + img, _ = batch + recon_loss = self.forward(img) + # Log the training loss to logger and progress bar (per-step and per-epoch) + self.log( + "train_loss", + recon_loss, + on_step=True, + on_epoch=True, + prog_bar=True, + logger=True, + ) + return recon_loss + + def configure_optimizers(self): + return torch.optim.SGD( + self.parameters(), lr=0.03, momentum=0.9, weight_decay=1e-4 + ) + + +model = SparseSparK(input_size=416) + +dataset = torchvision.datasets.Caltech101( + "datasets/caltech101", + download=True, + transform=v2.Compose( + [ + v2.Resize((416, 416)), + v2.RGB(), + v2.ToTensor(), + v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + ] + ), +) +# or create a dataset from a folder containing images or videos: +# dataset = LightlyDataset("path/to/folder") + +dataloader = torch.utils.data.DataLoader( + dataset, + batch_size=4, + shuffle=True, + drop_last=True, + num_workers=8, +) + +trainer = pl.Trainer( + max_epochs=30, +) + +trainer.fit( + model=model, + train_dataloaders=dataloader, +) diff --git a/lightly/loss/__init__.py b/lightly/loss/__init__.py index e0e2291cf..1bba2eef0 100644 --- a/lightly/loss/__init__.py +++ b/lightly/loss/__init__.py @@ -1,4 +1,4 @@ -"""The lightly.loss package provides loss functions for self-supervised learning. """ +"""The lightly.loss package provides loss functions for self-supervised learning.""" # Copyright (c) 2020. Lightly AG and its affiliates. # All Rights Reserved @@ -16,6 +16,7 @@ from lightly.loss.negative_cosine_similarity import NegativeCosineSimilarity from lightly.loss.ntx_ent_loss import NTXentLoss from lightly.loss.pmsn_loss import PMSNCustomLoss, PMSNLoss +from lightly.loss.sparse_spark import SparKPatchReconLoss from lightly.loss.swav_loss import SwaVLoss from lightly.loss.sym_neg_cos_sim_loss import SymNegCosineSimilarityLoss from lightly.loss.tico_loss import TiCoLoss diff --git a/lightly/loss/sparse_spark.py b/lightly/loss/sparse_spark.py new file mode 100644 index 000000000..664147a57 --- /dev/null +++ b/lightly/loss/sparse_spark.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +import torch +import torch.distributed as dist +from torch import Tensor, nn + +from lightly.utils.dist import gather + + +class SparKPatchReconLoss(nn.Module): + """Computes per-patch normalized reconstruction loss for masked regions. + + Original paper: https://github.com/keyu-tian/SparK + + Calculates L2 loss between reconstructed and original patches, normalized per-patch + to account for varying feature statistics. Loss is computed only on masked (inactive) regions. + + Args: + eps: Small value for numerical stability. Default: 1e-6. + """ + + def __init__(self, eps: float = 1e-6, gather_distributed: bool = False) -> None: + super().__init__() + if gather_distributed and not dist.is_available(): + raise ValueError( + "gather_distributed is True but torch.distributed is not available. " + "Please set gather_distributed=False or install a torch version with " + "distributed support." + ) + self.eps = eps + self.gather_distributed = gather_distributed + + def forward( + self, + inp_patches: Tensor, + rec_patches: Tensor, + active_mask: Tensor, + ) -> tuple[Tensor, Tensor, Tensor]: + """Compute reconstruction loss and per-patch statistics. + + Normalizes original patches based on per-patch mean and variance, then computes + L2 loss between normalized original and reconstructed patches. Averages loss + only over masked (active_mask=False) patches. + + Args: + inp_patches: Original patches of shape (B, L, N) where B=batch, L=length, N=patch_dim. + rec_patches: Reconstructed patches of shape (B, L, N). + active_mask: Boolean mask of shape (B, 1, f, f) indicating active regions. + Must have 4 dimensions (2D spatial mask). + + Returns: + Tuple of: + - recon_loss: Scalar tensor with averaged reconstruction loss on masked regions. + - mean: Per-patch mean of shape (B, L, 1). + - var: Per-patch standard deviation of shape (B, L, 1). + + Raises: + ValueError: If active_mask does not have 4 dimensions. + """ + if active_mask.ndim != 4: + raise ValueError( + "active_mask must be non-flattened with shape (B, 1, f, f)" + ) + + mean = inp_patches.mean(dim=-1, keepdim=True) + var = (inp_patches.var(dim=-1, keepdim=True) + self.eps) ** 0.5 + + inp_normalized = (inp_patches - mean) / var + + l2_loss = ((rec_patches - inp_normalized) ** 2).mean(dim=2) + + non_active = active_mask.logical_not().int().view(active_mask.shape[0], -1) + + local_numerator = (l2_loss * non_active).sum() + local_denominator = non_active.sum() + + if self.gather_distributed and dist.is_available() and dist.is_initialized(): + global_numerator = torch.cat( + gather(local_numerator.unsqueeze(0)), dim=0 + ).sum() + global_denominator = torch.cat( + gather(local_denominator.unsqueeze(0)), dim=0 + ).sum() + else: + global_numerator = local_numerator + global_denominator = local_denominator + + recon_loss = global_numerator / (global_denominator + self.eps) + return recon_loss, mean, var diff --git a/lightly/models/modules/__init__.py b/lightly/models/modules/__init__.py index f5c7c39f4..d4be20e0d 100644 --- a/lightly/models/modules/__init__.py +++ b/lightly/models/modules/__init__.py @@ -8,7 +8,6 @@ # Copyright (c) 2021. Lightly AG and its affiliates. # All Rights Reserved - from lightly.models.modules.heads import ( BarlowTwinsProjectionHead, BYOLPredictionHead, @@ -30,6 +29,12 @@ SwaVPrototypes, ) from lightly.models.modules.nn_memory_bank import NNMemoryBankModule +from lightly.models.modules.sparse_spark import ( + SparKDensifier, + SparKMasker, + SparKOutputDecoder, + dense_model_to_sparse, +) from lightly.utils import dependency as _dependency if _dependency.torchvision_vit_available(): diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py new file mode 100644 index 000000000..89f44afa7 --- /dev/null +++ b/lightly/models/modules/sparse_spark.py @@ -0,0 +1,889 @@ +from __future__ import annotations + +import math +from contextlib import contextmanager +from contextvars import ContextVar +from typing import Generator, NamedTuple, cast + +import torch +import torch.nn as nn +from torch import Tensor + +from lightly.models.utils import random_token_mask + + +def is_pow2n(x: int) -> bool: + """Check if an integer is a power of 2. + + Args: + x: Integer to check. + + Returns: + True if x is a power of 2, False otherwise. + """ + return x > 0 and (x & (x - 1) == 0) + + +def coalesce_to_size_2_t(t: tuple[int, ...]) -> tuple[int, int]: + """Convert a 1-tuple or 2-tuple to standard (H, W) format. + + Args: + t: Tuple of length 1 or 2 containing integers. + + Returns: + A 2-tuple (h, w). If input is 1-tuple, both dimensions are equal. + + Raises: + ValueError: If tuple length is not 1 or 2. + """ + if len(t) == 2: + return (t[0], t[1]) + elif len(t) == 1: + return (t[0], t[0]) + else: + raise ValueError(f"Invalid tuple length: {len(t)}; expected 1 or 2.") + + +_cur_active: ContextVar[Tensor | None] = ContextVar("_cur_active", default=None) + + +@contextmanager +def sparse_layer_context( + active_mask: Tensor, +) -> Generator[None, None, None]: + """Context manager to set the active mask for sparse layers. + + Args: + active_mask: Boolean mask of shape (B, 1, f, f) indicating active regions. + """ + token = _cur_active.set(active_mask) + try: + yield + finally: + _cur_active.reset(token) + + +def _get_active_ex_or_ii( + H: int, W: int, returning_active_ex: bool = True +) -> Tensor | tuple[Tensor, ...]: + """Get active indices or expanded active mask from global _cur_active. + + Converts the global _cur_active mask (shape B, 1, f, f) to a given spatial resolution (H, W). + Uses repeat_interleave to expand the mask to match the target spatial dimensions. + + Args: + H: Target height dimension. + W: Target width dimension. + returning_active_ex: If True, return expanded binary mask (B, 1, H, W). + If False, return tuple of nonzero indices (bi, hi, wi). + + Returns: + If returning_active_ex=True: Tensor of shape (B, 1, H, W) with binary active mask. + If returning_active_ex=False: Tuple of 3 tensors (batch_idx, height_idx, width_idx) for active positions. + + Note: + Optimization opportunity: Consider using gather() for better performance (see TODO). + """ + active_mask = _cur_active.get() + if active_mask is None: + raise RuntimeError( + "_cur_active must be set before calling this function. " + "Use sparse_layer_context to set the mask." + ) + h_repeat, w_repeat = H // active_mask.shape[-2], W // active_mask.shape[-1] + assert ( + h_repeat > 0 + ), f"Target height {H} must be >= mask height {active_mask.shape[-2]}" + assert ( + w_repeat > 0 + ), f"Target width {W} must be >= mask width {active_mask.shape[-1]}" + active_ex = active_mask.repeat_interleave(h_repeat, dim=2).repeat_interleave( + w_repeat, dim=3 + ) + return ( + active_ex + if returning_active_ex + else active_ex.squeeze(1).nonzero(as_tuple=True) + ) + + +def sp_conv_forward( + module: nn.AvgPool2d | nn.MaxPool2d | nn.Conv2d, input: Tensor +) -> Tensor: + """Forward pass for sparse convolution/pooling layers. + + Applies the parent class forward operation and masks the output using the global + active mask to zero out inactive spatial positions. + + Args: + self: ConvTranspose2d, MaxPool2d, or AvgPool2d instance. + input: Input tensor of shape (B, C, H, W). + + Returns: + Masked output tensor of same shape as input, with inactive spatial positions zeroed. + """ + x: Tensor = super(type(module), module).forward(input) # type: ignore[arg-type] + x *= cast( + Tensor, + _get_active_ex_or_ii( + H=input.shape[2], W=input.shape[3], returning_active_ex=True + ), + ) + return x + + +def sp_bn_forward(module: nn.BatchNorm1d | nn.SyncBatchNorm, input: Tensor) -> Tensor: + """Forward pass for sparse batch normalization layers. + + Applies batch norm only to active (unmasked) spatial positions, efficiently handling + sparse feature maps by extracting active features, normalizing them, and reconstructing. + Uses 1D batch norm on flattened active features rather than standard 2D batch norm. + + Args: + self: BatchNorm1d or SyncBatchNorm instance. + input: Input tensor of shape (B, C, H, W) in channels_first format. + + Returns: + Output tensor of same shape as input, with batch norm applied only to active positions. + """ + ii = _get_active_ex_or_ii( + H=input.shape[2], W=input.shape[3], returning_active_ex=False + ) + + bhwc = input.permute(0, 2, 3, 1) + nc = bhwc[ii] + nc = super(type(module), module).forward(nc) # type: ignore[arg-type] + + bchw = torch.zeros_like(bhwc) + bchw[ii] = nc + bchw = bchw.permute(0, 3, 1, 2) + return bchw + + +class SparseConv2d(nn.Conv2d): + """Sparse 2D convolution layer that respects active/mask regions. + + Overrides forward pass to apply global active mask to output, zeroing inactive regions. + Uses function override pattern for efficiency rather than standard subclassing override. + """ + + forward = sp_conv_forward + + +class SparseMaxPooling(nn.MaxPool2d): + """Sparse max pooling layer that respects active/mask regions. + + Overrides forward pass to apply global active mask to output, zeroing inactive regions. + Uses function override pattern for efficiency rather than standard subclassing override. + """ + + forward = sp_conv_forward + + +class SparseAvgPooling(nn.AvgPool2d): + """Sparse average pooling layer that respects active/mask regions. + + Overrides forward pass to apply global active mask to output, zeroing inactive regions. + Uses function override pattern for efficiency rather than standard subclassing override. + """ + + forward = sp_conv_forward + + +class SparseBatchNorm2d(nn.BatchNorm1d): + """Sparse batch normalization for 2D feature maps. + + Overrides forward pass to apply batch norm only to active (unmasked) positions. + Internally converts to 1D batch norm for efficient processing of sparse features. + Uses function override pattern for efficiency rather than standard subclassing override. + """ + + forward = sp_bn_forward + + +class SparseSyncBatchNorm2d(nn.SyncBatchNorm): + """Sparse synchronized batch normalization for 2D feature maps. + + Overrides forward pass to apply synchronized batch norm only to active (unmasked) positions. + Internally converts to 1D batch norm for efficient processing of sparse features. + Uses function override pattern for efficiency rather than standard subclassing override. + Recommended for distributed training scenarios. + """ + + forward = sp_bn_forward + + +class SparseConvNeXtLayerNorm(nn.LayerNorm): + r"""LayerNorm that supports two data formats: channels_last (default) or channels_first. + + The ordering of the dimensions in the inputs. channels_last corresponds to inputs with + shape (batch_size, height, width, channels) while channels_first corresponds to inputs + with shape (batch_size, channels, height, width). + + This sparse implementation only applies normalization to active (unmasked) spatial + positions using indices from the global _cur_active mask. + """ + + def __init__( + self, + normalized_shape: int, + eps: float = 1e-6, + data_format: str = "channels_last", + ) -> None: + if data_format not in ["channels_last", "channels_first"]: + raise NotImplementedError + super().__init__(normalized_shape, eps, elementwise_affine=True) + self.data_format = data_format + + def forward(self, input: Tensor) -> Tensor: + if input.ndim == 4: # BHWC or BCHW + if self.data_format == "channels_last": # BHWC + ii = _get_active_ex_or_ii( + H=input.shape[1], W=input.shape[2], returning_active_ex=False + ) + nc = input[ii] + nc = super().forward(nc) + + input = torch.zeros_like(input) + input[ii] = nc + return input + else: # channels_first, BCHW + ii = _get_active_ex_or_ii( + H=input.shape[2], W=input.shape[3], returning_active_ex=False + ) + bhwc = input.permute(0, 2, 3, 1) + nc = bhwc[ii] + nc = super().forward(nc) + + input = torch.zeros_like(bhwc) + input[ii] = nc + return input.permute(0, 3, 1, 2) + else: # BLC or BC + raise NotImplementedError + + +class SparseConvNeXtBlock(nn.Module): + r"""ConvNeXt Block with sparse computation support. + + There are two equivalent implementations: + (1) DwConv -> LayerNorm (channels_first) -> 1x1 Conv -> GELU -> 1x1 Conv; all in (N, C, H, W) + (2) DwConv -> Permute to (N, H, W, C); LayerNorm (channels_last) -> Linear -> GELU -> Linear; Permute back + + We use (2) as we find it slightly faster in PyTorch. This sparse implementation always + applies masking to the output. + + Args: + dim: Number of input channels. + drop_path: Stochastic depth rate. Default: 0.0 + layer_scale_init_value: Init value for Layer Scale. Default: 1e-6. + ks: Kernel size for depthwise convolution. Default: 7. + """ + + def __init__( + self, + dim: int, + drop_path: float = 0.0, + layer_scale_init_value: float = 1e-6, + ks: int = 7, + ) -> None: + from timm.models.layers import DropPath + + super().__init__() + self.dwconv = nn.Conv2d(dim, dim, kernel_size=ks, padding=ks // 2, groups=dim) + self.norm = SparseConvNeXtLayerNorm(dim, eps=1e-6) + self.pwconv1 = nn.Linear(dim, 4 * dim) + self.act = nn.GELU() + self.pwconv2 = nn.Linear(4 * dim, dim) + self.gamma = ( + nn.Parameter(layer_scale_init_value * torch.ones((dim)), requires_grad=True) + if layer_scale_init_value > 0 + else None + ) + self.drop_path: nn.Module = ( + DropPath(drop_path) if drop_path > 0.0 else nn.Identity() + ) + + def forward(self, x: Tensor) -> Tensor: + input = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act(x) + x = self.pwconv2(x) + if self.gamma is not None: + x = self.gamma * x + x = x.permute(0, 3, 1, 2) + + x *= cast( + Tensor, + _get_active_ex_or_ii(H=x.shape[2], W=x.shape[3], returning_active_ex=True), + ) + + x = input + self.drop_path(x) + return x + + +def dense_model_to_sparse( + m: nn.Module, verbose: bool = False, sbn: bool = False +) -> nn.Module: + """Recursively convert a dense model to sparse by replacing layer types. + + Original paper: https://github.com/keyu-tian/SparK + + Handles Conv2d, MaxPool2d, AvgPool2d, BatchNorm2d, SyncBatchNorm, and LayerNorm layers. + Copies weight and state tensors to maintain the original model's parameters. + + Args: + m: Original dense model or module. + verbose: Whether to print conversion details. Default: False. + sbn: Whether to use SyncBatchNorm instead of BatchNorm2d. Default: False. + + Returns: + Sparse version of the input model with converted layers. + + Raises: + NotImplementedError: If Conv1d layers are encountered. + """ + oup: nn.Module = m + if isinstance(m, nn.Conv2d): + bias = m.bias is not None + + sparse_conv: SparseConv2d = SparseConv2d( + m.in_channels, + m.out_channels, + kernel_size=coalesce_to_size_2_t(m.kernel_size), + stride=coalesce_to_size_2_t(m.stride), + padding=m.padding + if isinstance(m.padding, str) + else coalesce_to_size_2_t(m.padding), + dilation=coalesce_to_size_2_t(m.dilation), + groups=m.groups, + bias=bias, + padding_mode=m.padding_mode, + ) + sparse_conv.weight.data.copy_(m.weight.data) + + if m.bias is not None and sparse_conv.bias is not None: + sparse_conv.bias.data.copy_(m.bias.data) + oup = sparse_conv + + elif isinstance(m, nn.MaxPool2d): + oup = SparseMaxPooling( + m.kernel_size, + stride=m.stride, + padding=m.padding, + dilation=m.dilation, + return_indices=m.return_indices, + ceil_mode=m.ceil_mode, + ) + elif isinstance(m, nn.AvgPool2d): + oup = SparseAvgPooling( + m.kernel_size, + m.stride, + m.padding, + ceil_mode=m.ceil_mode, + count_include_pad=m.count_include_pad, + divisor_override=m.divisor_override, + ) + elif isinstance(m, (nn.BatchNorm2d, nn.SyncBatchNorm)): + if sbn: + sparse_bn: SparseSyncBatchNorm2d = SparseSyncBatchNorm2d( + m.weight.shape[0], + eps=m.eps, + momentum=m.momentum, + affine=m.affine, + track_running_stats=m.track_running_stats, + ) + sparse_bn.weight.data.copy_(m.weight.data) + sparse_bn.bias.data.copy_(m.bias.data) # type: ignore[union-attr] + sparse_bn.running_mean.data.copy_(m.running_mean.data) # type: ignore[union-attr] + sparse_bn.running_var.data.copy_(m.running_var.data) # type: ignore[union-attr] + sparse_bn.num_batches_tracked.data.copy_(m.num_batches_tracked.data) # type: ignore[union-attr] + if hasattr(m, "qconfig"): + sparse_bn.qconfig = m.qconfig # type: ignore[assignment] + oup = sparse_bn + else: + sparse_bn2: SparseBatchNorm2d = SparseBatchNorm2d( + m.weight.shape[0], + eps=m.eps, + momentum=m.momentum, + affine=m.affine, + track_running_stats=m.track_running_stats, + ) + sparse_bn2.weight.data.copy_(m.weight.data) + sparse_bn2.bias.data.copy_(m.bias.data) # type: ignore[union-attr] + sparse_bn2.running_mean.data.copy_(m.running_mean.data) # type: ignore[union-attr] + sparse_bn2.running_var.data.copy_(m.running_var.data) # type: ignore[union-attr] + sparse_bn2.num_batches_tracked.data.copy_(m.num_batches_tracked.data) # type: ignore[union-attr] + if hasattr(m, "qconfig"): + sparse_bn2.qconfig = m.qconfig # type: ignore[assignment] + oup = sparse_bn2 + elif isinstance(m, nn.LayerNorm) and not isinstance(m, SparseConvNeXtLayerNorm): + sparse_ln: SparseConvNeXtLayerNorm = SparseConvNeXtLayerNorm( + m.weight.shape[0], eps=m.eps + ) + sparse_ln.weight.data.copy_(m.weight.data) + sparse_ln.bias.data.copy_(m.bias.data) + oup = sparse_ln + elif isinstance(m, (nn.Conv1d,)): + raise NotImplementedError + + for name, child in m.named_children(): + oup.add_module(name, dense_model_to_sparse(child, verbose=verbose, sbn=sbn)) # type: ignore[union-attr] + del m + return oup + + +class UNetBlock(nn.Module): + """U-Net upsampling block with 2x spatial upsampling and conv refinement. + + Original paper: https://github.com/keyu-tian/SparK + + Combines transposed convolution for 2x upsampling followed by residual convolutions + with batch normalization and ReLU activation. + + Args: + cin: Number of input channels. + cout: Number of output channels. + bn2d: Batch normalization layer class (e.g., nn.BatchNorm2d or nn.SyncBatchNorm). + """ + + def __init__(self, cin: int, cout: int, bn2d: type) -> None: + super().__init__() + self.up_sample = nn.ConvTranspose2d( + cin, cin, kernel_size=4, stride=2, padding=1, bias=True + ) + self.conv = nn.Sequential( + nn.Conv2d(cin, cin, kernel_size=3, stride=1, padding=1, bias=False), + bn2d(cin), + nn.ReLU6(inplace=True), + nn.Conv2d(cin, cout, kernel_size=3, stride=1, padding=1, bias=False), + bn2d(cout), + ) + + def forward(self, x: Tensor) -> Tensor: + """Apply upsampling and convolution refinement. + + Args: + x: Input feature tensor. + + Returns: + Upsampled and refined feature tensor. + """ + x = self.up_sample(x) + return self.conv(x) # type: ignore[no-any-return] + + +class UNetDecoder(nn.Module): + """Lightweight hierarchical decoder for feature map reconstruction. + + Original paper: https://github.com/keyu-tian/SparK + + Applies a series of UNetBlocks to progressively upsample feature maps from deep to shallow, + halving channels at each level according to a simple rule (width //= 2). + Final projection outputs 3 channels for visualization/reconstruction. + + Args: + up_sample_ratio: Total spatial upsampling ratio (must be power of 2). + width: Base channel width at deepest level. Default: 768. + sbn: Whether to use SyncBatchNorm (True) or BatchNorm2d (False). Default: True. + """ + + def __init__( + self, up_sample_ratio: int, width: int = 768, sbn: bool = True + ) -> None: + super().__init__() + self.width = width + assert is_pow2n(up_sample_ratio) + n = round(math.log2(up_sample_ratio)) + channels = [self.width // 2**i for i in range(n + 1)] + bn2d = nn.SyncBatchNorm if sbn else nn.BatchNorm2d + self.dec = nn.ModuleList( + [ + UNetBlock(cin, cout, bn2d) + for (cin, cout) in zip(channels[:-1], channels[1:]) + ] + ) + self.proj = nn.Conv2d(channels[-1], 3, kernel_size=1, stride=1, bias=True) + + self.initialize() + + def forward(self, to_dec: list[Tensor]) -> Tensor: + """Progressively upsample and combine feature maps. + + Args: + to_dec: List of feature tensors from different scales (shallow to deep). + + Returns: + Upsampled feature tensor with 3 output channels. + """ + x = torch.tensor(0.0, device=to_dec[0].device) + for i in range(len(self.dec)): + if i < len(to_dec) and to_dec[i] is not None: + x = x + to_dec[i] + x = self.dec[i](x) + return self.proj(x) # type: ignore[no-any-return] + + def extra_repr(self) -> str: + return f"width={self.width}" + + def initialize(self) -> None: + """Initialize weights and biases using appropriate initialization schemes. + + Uses: + - trunc_normal_ for Linear and Conv2d layers (std=0.02) + - kaiming_normal_ for ConvTranspose2d layers + - constant initialization for batch norm and layer norm (weight=1.0, bias=0.0) + """ + from timm.models.layers import trunc_normal_ + + for m in self.modules(): + if isinstance(m, nn.Linear): + trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Conv2d): + trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)): + nn.init.kaiming_normal_(m.weight, mode="fan_out", nonlinearity="relu") + if m.bias is not None: + nn.init.constant_(m.bias, 0.0) + elif isinstance( + m, (nn.LayerNorm, nn.BatchNorm1d, nn.BatchNorm2d, nn.SyncBatchNorm) + ): + nn.init.constant_(m.bias, 0) + nn.init.constant_(m.weight, 1.0) + + +class SparKDensifiyBlock(nn.Module): + """Block for densifying sparse features by filling masked regions with learned tokens. + + Original paper: https://github.com/keyu-tian/SparK + + Applies normalization to sparse features, then uses learned mask tokens to fill inactive + regions, finally projecting to target channel dimension. + + Args: + e_width: Number of input channels (encoder width). + d_width: Number of output channels (decoder width). + densify_norm_str: Type of normalization ('bn', 'ln', or 'none'). Default: 'bn'. + sbn: Whether to use SyncBatchNorm if densify_norm_str='bn'. Default: False. + use_identity_proj: If True, use identity projection (assumes e_width == d_width). Default: False. + kernel_size: Kernel size for projection convolution. Default: 3. + """ + + def __init__( + self, + e_width: int, + d_width: int, + densify_norm_str: str = "bn", + sbn: bool = False, + use_identity_proj: bool = False, + kernel_size: int = 3, + ) -> None: + from timm.models.layers import trunc_normal_ + + super().__init__() + # mask token + p = nn.Parameter(torch.zeros(1, e_width, 1, 1)) + trunc_normal_(p, mean=0, std=0.02, a=-0.02, b=0.02) + self.mask_token = p + + # densify norm + if densify_norm_str == "bn": + self.densify_norm = (SparseSyncBatchNorm2d if sbn else SparseBatchNorm2d)( + e_width + ) + elif densify_norm_str == "ln": + self.densify_norm = SparseConvNeXtLayerNorm( + e_width, data_format="channels_first" + ) + else: + self.densify_norm = nn.Identity() + + # densify proj + self.densify_proj: nn.Identity | nn.Conv2d + if use_identity_proj: + self.densify_proj = nn.Identity() + print( + f"[SparKDensfiyBlock.__init__]: use nn.Identity() as densify_proj (e_width==d_width)" + ) + else: + densify_proj = nn.Conv2d( + e_width, + d_width, + kernel_size=kernel_size, + stride=1, + padding=kernel_size // 2, + bias=True, + ) + print( + f"[SparKDensfiyBlock.__init__]: densify_proj(ksz={kernel_size}, #para={sum(x.numel() for x in densify_proj.parameters()) / 1e6:.2f}M)" + ) + self.densify_proj = densify_proj + + def _fill_with_mask_tokens(self, features: Tensor, active_mask: Tensor) -> Tensor: + """Fill masked regions with learned mask tokens. + + Args: + features: Input tensor of shape (B, C, H, W). + active_mask: Boolean mask tensor of shape (B, 1, H, W) indicating active regions. + + Returns: + Tensor of shape (B, C, H, W) with masked regions filled with learnable tokens. + """ + mask_tokens: Tensor = self.mask_token.expand(features.size()) + active_expanded: Tensor = active_mask.expand(features.size()) + return torch.where(active_expanded, features, mask_tokens) + + def forward(self, bcff: Tensor, cur_active: Tensor) -> Tensor: + """Densify sparse features by filling masked regions with learned tokens. + + Args: + bcff: Sparse feature tensor of shape (B, C, H, W). + cur_active: Boolean mask tensor of shape (B, 1, H, W) indicating active regions. + + Returns: + Densified and projected feature tensor of shape (B, C', H, W). + """ + bcff = self.densify_norm(bcff) + bcff = self._fill_with_mask_tokens(bcff, cur_active) + bcff = self.densify_proj(bcff) + return bcff + + +class SparKDensifier(nn.Module): + """Stack of densify blocks to convert sparse hierarchical features to dense features. + + Original paper: https://github.com/keyu-tian/SparK + + Processes encoder feature maps from deepest to shallowest scale, applying SparKDensfiyBlock + to each level. Handles the global _cur_active mask, dilating it at each upsampling level. + + Args: + encoder_in_channels: List of channel numbers from encoder, shallow to deep. + decoder_in_channel: Base channel number for decoder (at deepest scale). + densify_norm_str: Type of normalization ('bn', 'ln', or 'none'). Default: 'bn'. + sbn: Whether to use SyncBatchNorm if densify_norm_str='bn'. Default: False. + """ + + def __init__( + self, + encoder_in_channels: list[int], + decoder_in_channel: int, + densify_norm_str: str = "bn", + sbn: bool = False, + ) -> None: + super().__init__() + self.blocks = nn.ModuleList() + self.encoder_in_channels = encoder_in_channels + self.decoder_in_channel = decoder_in_channel + d_width = decoder_in_channel + for i, e_width in enumerate(encoder_in_channels[::-1]): + use_identity = i == 0 and e_width == d_width + kernel_size = 1 if i <= 0 else 3 + block = SparKDensifiyBlock( + e_width=e_width, + d_width=d_width, + densify_norm_str=densify_norm_str, + sbn=sbn, + use_identity_proj=use_identity, + kernel_size=kernel_size, + ) + self.blocks.append(block) + d_width //= 2 + + def forward(self, fea_bcffs: list[Tensor]) -> list[Tensor]: + """Convert sparse features to dense by filling masked regions. + + Args: + fea_bcffs: List of feature tensors from encoder at different scales (shallow + to deep). Each tensor has shape (B, C, f, f) where f varies per scale. + + Returns: + List of densified feature tensors for decoder processing (order reversed and + dilated). + """ + to_dec = [] + fea_bcffs = fea_bcffs[::-1] + active_mask = _cur_active.get() + if active_mask is None: + raise RuntimeError( + "_cur_active must be set before calling SparKDensifier.forward(). " + "Ensure you are within a sparse_layer_context." + ) + active_fmap_current = active_mask + for i, bcff in enumerate(fea_bcffs): + bcff = self.blocks[i](bcff, active_fmap_current) + to_dec.append(bcff) + active_fmap_current = active_fmap_current.repeat_interleave( + 2, dim=2 + ).repeat_interleave(2, dim=3) + return to_dec + + +class SparKMaskingOutput(NamedTuple): + """Output container for SparKMasker forward pass. + + Original paper: https://github.com/keyu-tian/SparK + + Attributes: + masked_bchw: Input image with masks applied at full resolution. + per_level_mask: List of binary masks at each hierarchical level. + """ + + masked_bchw: Tensor + per_level_mask: list[Tensor] + + +class SparKMasker(nn.Module): + """Generates hierarchical random token masks for sparse feature processing. + + Original paper: https://github.com/keyu-tian/SparK + + Creates a binary mask at patch level and expands it hierarchically to match + feature maps at different spatial resolutions in the encoder. + + Args: + feature_map_size: Tuple of (height, width) at patch/feature map level. + downsample_ratio: Total downsampling ratio from input to feature map level. + mask_ratio: Fraction of tokens to mask (make inactive). Default: 0.6. + """ + + def __init__( + self, + feature_map_size: tuple[int, int], + downsample_ratio: int, + mask_ratio: float = 0.6, + ) -> None: + super().__init__() + self.fmap_h, self.fmap_w = feature_map_size + self.downsample_ratio = downsample_ratio + self.mask_ratio = mask_ratio + + def mask(self, B: int, device: torch.device) -> Tensor: + """Generate a random binary mask for features. + + Args: + B: Batch size. + device: Device to create the mask on. + + Returns: + Boolean mask tensor of shape (B, 1, fmap_h, fmap_w) indicating active tokens. + """ + h, w = self.fmap_h, self.fmap_w + index_keep, _ = random_token_mask( + size=(B, h * w), mask_ratio=self.mask_ratio, device=device + ) + return ( + torch.zeros(B, 1, h * w, dtype=torch.bool, device=device) + .scatter_(dim=2, index=index_keep.unsqueeze(1), value=True) + .view(B, 1, h, w) + .bool() + ) + + def forward(self, inp_bchw: Tensor) -> SparKMaskingOutput: + """Generate hierarchical masks for the input. + + Creates masks at multiple scales from the patch level up to full input resolution. + + Args: + inp_bchw: Input image tensor of shape (B, C, H, W). + + Returns: + SparKMaskingOutput containing: + - masked_bchw: Input image masked at full resolution. + - per_level_mask: List of masks at each hierarchical level. + """ + active_b1ff = self.mask(inp_bchw.shape[0], inp_bchw.device) + downsample_ratio = self.downsample_ratio + per_level_mask = [active_b1ff] + for _ in range(int(math.log2(downsample_ratio))): + previous_mask = per_level_mask[-1] + active_b1cHcW = previous_mask.repeat_interleave(2, dim=2).repeat_interleave( + 2, dim=3 + ) + per_level_mask.append(active_b1cHcW) + + active_b1hw = per_level_mask[-1] + masked_bchw = inp_bchw * active_b1hw + return SparKMaskingOutput( + masked_bchw=masked_bchw, per_level_mask=per_level_mask + ) + + +class SparKOutputDecoder(nn.Module): + """Decodes reconstructed patches back to image space and combines with original. + + Original paper: https://github.com/keyu-tian/SparK + + Handles denormalization and unpatchifying of reconstructed patches, then performs + per-pixel blending: uses original pixels where visible (active), reconstructed pixels + where masked (inactive). + + Args: + fmap_h: Height of feature map at patch level. + fmap_w: Width of feature map at patch level. + downsample_ratio: Ratio of input image size to feature map size. + """ + + def __init__(self, fmap_h: int, fmap_w: int, downsample_ratio: int) -> None: + super().__init__() + self.fmap_h = fmap_h + self.fmap_w = fmap_w + self.downsample_ratio = downsample_ratio + + def unpatchify(self, bln: Tensor) -> Tensor: + """Convert flattened patches back to spatial feature map format. + + Reverses the patchify operation: reshapes from (B, L*p*p, N) to (B, N, H, W) + where p is the patch size (downsample_ratio). + + Args: + bln: Flattened patches of shape (B, L, N). + + Returns: + Spatial feature map of shape (B, C, H, W). + """ + p = self.downsample_ratio + h, w = self.fmap_h, self.fmap_w + B, C = bln.shape[0], bln.shape[-1] // p**2 + bln = bln.reshape(shape=(B, h, w, p, p, C)) + bln = torch.einsum("bhwpqc->bchpwq", bln) + bchw = bln.reshape(shape=(B, C, h * p, w * p)) + return bchw + + def forward( + self, + rec_patches: Tensor, + mean: Tensor, + var: Tensor, + inp_bchw: Tensor, + active_mask_full: Tensor, + ) -> tuple[Tensor, Tensor, Tensor]: + """Reconstruct image by blending original and reconstructed regions. + + Denormalizes reconstructed patches, unpatchifies them, and performs pixel-wise + blending based on the active mask: original where visible, reconstructed where masked. + + Args: + rec_patches: Reconstructed patches of shape (B, L, N). + mean: Per-patch mean of shape (B, L, 1). + var: Per-patch standard deviation of shape (B, L, 1). + inp_bchw: Original input image of shape (B, C, H, W). + active_mask_full: Boolean mask of shape (B, 1, H, W) at full resolution. + + Returns: + Tuple of: + - inp_bchw: Original input image (unchanged). + - masked_bchw: Input with inactive regions zeroed. + - rec_or_inp: Blended result using original where visible, reconstructed where masked. + """ + rec_bchw = self.unpatchify(rec_patches * var + mean) + masked_bchw = inp_bchw * active_mask_full + rec_or_inp = torch.where(active_mask_full, inp_bchw, rec_bchw) + + return inp_bchw, masked_bchw, rec_or_inp diff --git a/tests/loss/test_sparse_spark.py b/tests/loss/test_sparse_spark.py new file mode 100644 index 000000000..766878455 --- /dev/null +++ b/tests/loss/test_sparse_spark.py @@ -0,0 +1,83 @@ +import pytest +import torch +import torch.distributed as dist +from pytest_mock import MockerFixture + +pytest.importorskip("timm.models.layers") + +from lightly.loss import SparKPatchReconLoss + + +class TestSparKPatchReconLoss: + def test__gather_distributed(self, mocker: MockerFixture) -> None: + mock_is_available = mocker.patch.object(dist, "is_available", return_value=True) + SparKPatchReconLoss(gather_distributed=True) + mock_is_available.assert_called_once() + + def test__gather_distributed_dist_not_available( + self, mocker: MockerFixture + ) -> None: + mock_is_available = mocker.patch.object( + dist, "is_available", return_value=False + ) + with pytest.raises(ValueError): + SparKPatchReconLoss(gather_distributed=True) + mock_is_available.assert_called_once() + + def test_forward_pass(self) -> None: + loss = SparKPatchReconLoss() + inp_patches = torch.randn((2, 4, 3)) + recon_patches = torch.randn((2, 4, 3)) + mask = torch.randn((2, 1, 2, 2)) > 0 + + loss(inp_patches, recon_patches, mask) + + def test_forward_pass_deterministic(self) -> None: + loss = SparKPatchReconLoss() + inp_patches = torch.randn((2, 4, 3)) + recon_patches = torch.randn((2, 4, 3)) + mask = torch.randn((2, 1, 2, 2)) > 0 + + loss1, _, _ = loss(inp_patches, recon_patches, mask) + loss2, _, _ = loss(inp_patches, recon_patches, mask) + + assert loss1.item() == pytest.approx(loss2.item()) + + def test_forward__compare(self) -> None: + loss = SparKPatchReconLoss() + inp_patches = torch.randn((2, 4, 3)) + recon_patches = torch.randn((2, 4, 3)) + + mask = torch.randn((2, 1, 2, 2)) > 0 + + loss1 = loss(inp_patches, recon_patches, mask)[0] + loss2 = _reference_loss_implementation(inp_patches, recon_patches, mask) + + assert loss1.item() == pytest.approx(loss2.item(), rel=1e-5) + + +def _reference_loss_implementation( + inp: torch.Tensor, + rec: torch.Tensor, + active_b1ff: torch.Tensor, + eps: float = 1e-6, +) -> torch.Tensor: + """ + This loss implementation was taken from the official implementation of the + SparK paper, and serves as a reference to ensure that the + SparKPatchReconLoss implementation in lightly produces the same results. + + https://github.com/keyu-tian/SparK/blob/a63e386f8e5186bc07ad7fce86e06b08f48a61ea/pretrain/spark.py#L112-L120 + """ + mean = inp.mean(dim=-1, keepdim=True) + var = (inp.var(dim=-1, keepdim=True) + eps) ** 0.5 + inp = (inp - mean) / var + l2_loss = ((rec - inp) ** 2).mean( + dim=2, keepdim=False + ) # (B, L, C) ==mean==> (B, L) + + non_active = ( + active_b1ff.logical_not().int().view(active_b1ff.shape[0], -1) + ) # (B, 1, f, f) => (B, L) + recon_loss = l2_loss.mul_(non_active).sum() / (non_active.sum() + eps) + return recon_loss # type: ignore[no-any-return] diff --git a/tests/models/modules/test_sparse_spark.py b/tests/models/modules/test_sparse_spark.py new file mode 100644 index 000000000..8119e4fc7 --- /dev/null +++ b/tests/models/modules/test_sparse_spark.py @@ -0,0 +1,415 @@ +from typing import List + +import pytest +import torch +from torch import Tensor + +pytest.importorskip("timm.models.layers") + +import lightly.models.modules.sparse_spark as sparse_spark +from lightly.models.modules.sparse_spark import ( + SparKDensifiyBlock, + SparKMasker, + SparKMaskingOutput, +) + + +def _create_mask() -> Tensor: + return torch.tensor( + [ + [ + [ + [1, 0], + [0, 1], + ] + ] + ], + dtype=torch.bool, + ) + + +def test__get_active_ex_or_ii_expands_mask() -> None: + H, W = 32, 32 + with sparse_spark.sparse_layer_context(_create_mask()): + active = sparse_spark._get_active_ex_or_ii(H=H, W=W, returning_active_ex=True) + assert not isinstance(active, tuple) + + assert active.shape == (1, 1, H, W) + assert active[:, :, :16, :16].all() + assert active[:, :, :16, 16:].logical_not().all() + assert active[:, :, 16:, :16].logical_not().all() + assert active[:, :, 16:, 16:].all() + + +def test__get_active_ex_or_ii_dont_shrink_mask() -> None: + H, W = 4, 4 + with sparse_spark.sparse_layer_context(torch.ones(1, 1, 32, 32)): + with pytest.raises(AssertionError): + sparse_spark._get_active_ex_or_ii(H=H, W=W, returning_active_ex=False) + + +def test__get_active_ex_or_ii_raise_on_non_active_mask() -> None: + H, W = 32, 32 + with pytest.raises(RuntimeError): + sparse_spark._get_active_ex_or_ii(H=H, W=W, returning_active_ex=False) + + +def test__get_active_ex_or_ii_returning_ex_false_correct_values() -> None: + H, W = 32, 32 + with sparse_spark.sparse_layer_context(_create_mask()): + active_b, active_h, active_w = sparse_spark._get_active_ex_or_ii( + H=H, W=W, returning_active_ex=False + ) + active_ex = sparse_spark._get_active_ex_or_ii( + H=H, W=W, returning_active_ex=True + ) + assert not isinstance(active_ex, tuple) + + active_ex_scattered = torch.zeros_like(active_ex) + active_ex_scattered[active_b, :, active_h, active_w] = 1 + + assert torch.equal(active_ex, active_ex_scattered) + + +def test__sp_conv_forward() -> None: + H, W = 32, 32 + with sparse_spark.sparse_layer_context(_create_mask()): + conv = sparse_spark.SparseConv2d( + in_channels=1, + out_channels=1, + kernel_size=3, + padding=1, + ) + conv.weight.data.fill_(1) + assert conv.bias is not None + conv.bias.data.fill_(0) + + x = torch.ones(1, 1, H, W) + out = conv(x) + + assert out.shape == (1, 1, H, W) + assert out[:, :, :16, :16].all() + assert out[:, :, :16, 16:].logical_not().all() + assert out[:, :, 16:, :16].logical_not().all() + assert out[:, :, 16:, 16:].all() + + +class TestSparKMasker: + def test_forward(self) -> None: + masker = SparKMasker( + feature_map_size=(4, 4), + downsample_ratio=8, + ) + x = torch.ones(1, 1, 32, 32) + mask: SparKMaskingOutput = masker(x) + + for i in range(len(mask.per_level_mask)): + mask_current = mask.per_level_mask[i] + + assert mask_current.shape[0] == 1 + assert mask_current.shape[1] == 1 + assert mask_current.shape[2] == 4 * (2**i) + assert mask_current.shape[3] == 4 * (2**i) + + for i in range(len(mask.per_level_mask) - 1): + mask_current = mask.per_level_mask[i] + mask_next = mask.per_level_mask[i + 1] + assert mask_current.shape[2] * 2 == mask_next.shape[2] + assert mask_current.shape[3] * 2 == mask_next.shape[3] + assert ( + mask_next + == mask_current.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3) + ).all() + + def test_masked_bchw_applies_mask(self) -> None: + masker = SparKMasker( + feature_map_size=(4, 4), + downsample_ratio=8, + ) + x = torch.arange(1, 1025, dtype=torch.float32).view(1, 1, 32, 32) + mask: SparKMaskingOutput = masker(x) + + # masked_bchw should be x * active_mask at full resolution + active_mask_full = mask.per_level_mask[-1] + expected = x * active_mask_full + assert torch.equal(mask.masked_bchw, expected) + + def test_mask_ratio_zero_all_active(self) -> None: + masker = SparKMasker( + feature_map_size=(4, 4), + downsample_ratio=8, + mask_ratio=0.0, + ) + x = torch.ones(1, 1, 32, 32) + mask: SparKMaskingOutput = masker(x) + + # All tokens should be active (True) + assert mask.per_level_mask[0].all() + + def test_batch_independence(self) -> None: + torch.manual_seed(42) # Deterministic masks + masker = SparKMasker( + feature_map_size=(4, 4), + downsample_ratio=8, + ) + x = torch.ones(4, 1, 32, 32) + mask: SparKMaskingOutput = masker(x) + + # Each batch sample should have independent mask + assert mask.per_level_mask[0].shape[0] == 4 + # Masks should differ across batch (with seed, guaranteed different) + assert not torch.equal(mask.per_level_mask[0][0], mask.per_level_mask[0][1]) + + +class TestSparKDensifiyBlock: + def test_fill_with_mask_tokens_preserves_active_regions(self) -> None: + block: SparKDensifiyBlock = SparKDensifiyBlock( + e_width=4, d_width=4, densify_norm_str="none" + ) + block.mask_token.data.fill_(99.0) + + features: Tensor = torch.arange(16, dtype=torch.float32).view(1, 4, 2, 2) + active_mask: Tensor = torch.tensor([[[[True, False], [False, True]]]]) + + result: Tensor = block._fill_with_mask_tokens(features, active_mask) + + assert result[0, :, 0, 0].equal(features[0, :, 0, 0]) + assert result[0, :, 1, 1].equal(features[0, :, 1, 1]) + + def test_fill_with_mask_tokens_fills_inactive_regions(self) -> None: + block: SparKDensifiyBlock = SparKDensifiyBlock( + e_width=4, d_width=4, densify_norm_str="none" + ) + block.mask_token.data.fill_(99.0) + + features: Tensor = torch.arange(16, dtype=torch.float32).view(1, 4, 2, 2) + active_mask: Tensor = torch.tensor([[[[True, False], [False, True]]]]) + + result: Tensor = block._fill_with_mask_tokens(features, active_mask) + + expected_token: Tensor = torch.full((4,), 99.0, dtype=torch.float32) + assert result[0, :, 0, 1].equal(expected_token) + assert result[0, :, 1, 0].equal(expected_token) + + def test_fill_with_mask_tokens_all_active(self) -> None: + block: SparKDensifiyBlock = SparKDensifiyBlock( + e_width=4, d_width=4, densify_norm_str="none" + ) + block.mask_token.data.fill_(99.0) + + features: Tensor = torch.arange(16, dtype=torch.float32).view(1, 4, 2, 2) + active_mask: Tensor = torch.ones(1, 1, 2, 2, dtype=torch.bool) + + result: Tensor = block._fill_with_mask_tokens(features, active_mask) + + assert result.equal(features) + + def test_fill_with_mask_tokens_all_inactive(self) -> None: + block: SparKDensifiyBlock = SparKDensifiyBlock( + e_width=4, d_width=4, densify_norm_str="none" + ) + block.mask_token.data.fill_(99.0) + + features: Tensor = torch.arange(16, dtype=torch.float32).view(1, 4, 2, 2) + active_mask: Tensor = torch.zeros(1, 1, 2, 2, dtype=torch.bool) + + result: Tensor = block._fill_with_mask_tokens(features, active_mask) + + expected: Tensor = torch.full_like(features, 99.0) + assert result.equal(expected) + + def test_fill_with_mask_tokens_batch_independence(self) -> None: + block: SparKDensifiyBlock = SparKDensifiyBlock( + e_width=4, d_width=4, densify_norm_str="none" + ) + block.mask_token.data.fill_(99.0) + + features: Tensor = torch.arange(32, dtype=torch.float32).view(2, 4, 2, 2) + active_mask: Tensor = torch.tensor( + [ + [[[True, False], [False, True]]], + [[[False, True], [True, False]]], + ] + ) + + result: Tensor = block._fill_with_mask_tokens(features, active_mask) + + expected_token: Tensor = torch.full((4,), 99.0, dtype=torch.float32) + # Batch 0: (0,0) active, (0,1) inactive + assert result[0, :, 0, 0].equal(features[0, :, 0, 0]) + assert result[0, :, 0, 1].equal(expected_token) + # Batch 1: (0,0) inactive, (0,1) active + assert result[1, :, 0, 0].equal(expected_token) + assert result[1, :, 0, 1].equal(features[1, :, 0, 1]) + + +class TestSparKDensifier: + def test_forward_raises_without_active_mask(self) -> None: + """Test that forward raises RuntimeError when not in sparse_layer_context.""" + densifier: sparse_spark.SparKDensifier = sparse_spark.SparKDensifier( + encoder_in_channels=[64, 128, 256], + decoder_in_channel=256, + densify_norm_str="bn", + sbn=False, + ) + # Create dummy feature maps (shallow to deep) + fea_bcffs: List[Tensor] = [ + torch.randn(1, 64, 8, 8), + torch.randn(1, 128, 4, 4), + torch.randn(1, 256, 2, 2), + ] + + # Should raise because _cur_active is not set + with pytest.raises(RuntimeError, match="_cur_active must be set"): + densifier(fea_bcffs) + + def test_forward_upsamples_active_mask(self) -> None: + """Test that active_fmap_current is upsampled by 2x at each iteration.""" + densifier: sparse_spark.SparKDensifier = sparse_spark.SparKDensifier( + encoder_in_channels=[64, 128, 256], + decoder_in_channel=256, + densify_norm_str="bn", + sbn=False, + ) + # Create dummy feature maps at different scales (shallow to deep) + fea_bcffs: List[Tensor] = [ + torch.randn(1, 64, 8, 8), + torch.randn(1, 128, 4, 4), + torch.randn(1, 256, 2, 2), + ] + # Initial active mask at deepest scale (2x2) + initial_mask: Tensor = torch.tensor( + [[[[True, False], [False, True]]]], dtype=torch.bool + ) + + with sparse_spark.sparse_layer_context(initial_mask): + to_dec: List[Tensor] = densifier(fea_bcffs) + + # Check that we have 3 output feature maps (one per block) + assert len(to_dec) == 3 + + # Verify upsampling: each level should double in spatial size + # Level 0 (deepest): 2x2 -> input to first block + # Level 1: 2x2 -> 4x4 + # Level 2: 4x4 -> 8x8 + assert to_dec[0].shape == (1, 256, 2, 2) # First block output + assert to_dec[1].shape == (1, 128, 4, 4) # Second block output + assert to_dec[2].shape == (1, 64, 8, 8) # Third block output + + def test_forward_to_dec_contains_all_feature_maps(self) -> None: + """Test that to_dec list contains all densified feature maps.""" + encoder_in_channels: List[int] = [64, 128, 256] + decoder_in_channel: int = 256 + densifier: sparse_spark.SparKDensifier = sparse_spark.SparKDensifier( + encoder_in_channels=encoder_in_channels, + decoder_in_channel=decoder_in_channel, + densify_norm_str="bn", + sbn=False, + ) + fea_bcffs: List[Tensor] = [ + torch.randn(1, 64, 8, 8), + torch.randn(1, 128, 4, 4), + torch.randn(1, 256, 2, 2), + ] + initial_mask: Tensor = torch.ones(1, 1, 2, 2, dtype=torch.bool) + + with sparse_spark.sparse_layer_context(initial_mask): + to_dec: List[Tensor] = densifier(fea_bcffs) + + # Should have one output per block + assert len(to_dec) == len(densifier.blocks) + # All outputs should be Tensors + for i, td in enumerate(to_dec): + assert isinstance(td, Tensor) + + # Verify channel progression: [256, 128, 64] (halved at each level after first) + # Block 0: e_width=256 -> d_width=256 (identity) + # Block 1: e_width=128 -> d_width=128 (256//2) + # Block 2: e_width=64 -> d_width=64 (128//2) + expected_channels: List[int] = [256, 128, 64] + for i, td in enumerate(to_dec): + assert ( + td.shape[1] == expected_channels[i] + ), f"Block {i}: expected {expected_channels[i]} channels, got {td.shape[1]}" + + def test_forward_reverses_fea_bcffs_order(self) -> None: + """Test that fea_bcffs is reversed before processing (deepest first).""" + encoder_in_channels: List[int] = [64, 128, 256] + densifier: sparse_spark.SparKDensifier = sparse_spark.SparKDensifier( + encoder_in_channels=encoder_in_channels, + decoder_in_channel=256, + densify_norm_str="bn", + sbn=False, + ) + # Use distinct shapes to verify order + fea_bcffs: List[Tensor] = [ + torch.randn(1, 64, 8, 8), # shallowest + torch.randn(1, 128, 4, 4), + torch.randn(1, 256, 2, 2), # deepest + ] + initial_mask: Tensor = torch.ones(1, 1, 2, 2, dtype=torch.bool) + + with sparse_spark.sparse_layer_context(initial_mask): + to_dec: List[Tensor] = densifier(fea_bcffs) + + # First block should process deepest feature (2x2, 256 channels) + assert to_dec[0].shape == (1, 256, 2, 2) + # Second block should process middle feature (4x4, 128 channels) + assert to_dec[1].shape == (1, 128, 4, 4) + # Third block should process shallowest feature (8x8, 64 channels) + assert to_dec[2].shape == (1, 64, 8, 8) + + def test_forward_complete_integration(self) -> None: + """Integration test for complete forward pass with proper context.""" + encoder_in_channels: List[int] = [64, 128, 256] + decoder_in_channel: int = 256 + densifier: sparse_spark.SparKDensifier = sparse_spark.SparKDensifier( + encoder_in_channels=encoder_in_channels, + decoder_in_channel=decoder_in_channel, + densify_norm_str="bn", + sbn=False, + ) + batch_size: int = 2 + fea_bcffs: List[Tensor] = [ + torch.randn(batch_size, 64, 8, 8), + torch.randn(batch_size, 128, 4, 4), + torch.randn(batch_size, 256, 2, 2), + ] + # Active mask at deepest level (matches feature map size) + initial_mask: Tensor = torch.randint( + 0, 2, (batch_size, 1, 2, 2), dtype=torch.bool + ) + + with sparse_spark.sparse_layer_context(initial_mask): + to_dec: List[Tensor] = densifier(fea_bcffs) + + # Verify output structure + assert len(to_dec) == 3 + assert all(isinstance(t, Tensor) for t in to_dec) + # Verify batch size preserved + assert all(t.shape[0] == batch_size for t in to_dec) + # Verify spatial dimensions double at each level + assert to_dec[0].shape[2] * 2 == to_dec[1].shape[2] + assert to_dec[1].shape[2] * 2 == to_dec[2].shape[2] + + def test_forward_with_identity_proj_first_block(self) -> None: + """Test that first block uses identity projection when channels match.""" + encoder_in_channels: List[int] = [256, 128, 64] # Same as decoder + densifier: sparse_spark.SparKDensifier = sparse_spark.SparKDensifier( + encoder_in_channels=encoder_in_channels, + decoder_in_channel=256, + densify_norm_str="bn", + sbn=False, + ) + fea_bcffs: List[Tensor] = [ + torch.randn(1, 256, 8, 8), + torch.randn(1, 128, 4, 4), + torch.randn(1, 64, 2, 2), + ] + initial_mask: Tensor = torch.ones(1, 1, 2, 2, dtype=torch.bool) + + with sparse_spark.sparse_layer_context(initial_mask): + to_dec: List[Tensor] = densifier(fea_bcffs) + + # First block should preserve channels (identity projection) + assert to_dec[0].shape[1] == 256