From bd0855cdf6abd5f7da52b34d27d24db8eedb00e3 Mon Sep 17 00:00:00 2001 From: John Sutor Date: Wed, 4 Dec 2024 19:29:02 -0500 Subject: [PATCH 01/85] Initial work on SparK --- lightly/models/modules/sparse/spark.py | 61 ++++++ lightly/models/spark.py | 252 +++++++++++++++++++++++++ lightly/models/sparse/decoder.py | 0 lightly/models/sparse/resnet.py | 223 ++++++++++++++++++++++ 4 files changed, 536 insertions(+) create mode 100644 lightly/models/modules/sparse/spark.py create mode 100644 lightly/models/spark.py create mode 100644 lightly/models/sparse/decoder.py create mode 100644 lightly/models/sparse/resnet.py diff --git a/lightly/models/modules/sparse/spark.py b/lightly/models/modules/sparse/spark.py new file mode 100644 index 000000000..e1f47ff07 --- /dev/null +++ b/lightly/models/modules/sparse/spark.py @@ -0,0 +1,61 @@ +# Code adapted from https://github.com/keyu-tian/SparK/blob/main/pretrain/encoder.py + +import torch +import torch.nn as nn + + +def _get_active_ex_or_ii( + mask: torch.BoolTensor, H: int, W: int, returning_active_ex: bool = True +): + h_repeat, w_repeat = H // mask.shape[-2], W // mask.shape[-1] + active_ex = 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) + ) # ii: bi, hi, wi + + +def sp_conv_forward(self, x: torch.Tensor): + x = super(type(self), self).forward(x) + x *= _get_active_ex_or_ii( + self.sparse_mask.mask, H=x.shape[2], W=x.shape[3], returning_active_ex=True + ) # (BCHW) *= (B1HW), mask the output of conv + return x + + +def sp_bn_forward(self, x: torch.Tensor): + ii = _get_active_ex_or_ii( + self.sparse_mask.mask, H=x.shape[2], W=x.shape[3], returning_active_ex=False + ) + + bhwc = x.permute(0, 2, 3, 1) + nc = bhwc[ii] + nc = super(type(self), self).forward(nc) + + bchw = torch.zeros_like(bhwc) + bchw[ii] = nc + bchw = bchw.permute(0, 3, 1, 2) + return bchw + + +class SparseConv2d(nn.Conv2d): + forward = sp_conv_forward + + +class SparseMaxPooling(nn.MaxPool2d): + forward = sp_conv_forward + + +class SparseAvgPooling(nn.AvgPool2d): + forward = sp_conv_forward + + +class SparseBatchNorm2d(nn.BatchNorm1d): + forward = sp_bn_forward + + +class SparseSyncBatchNorm2d(nn.SyncBatchNorm): + forward = sp_bn_forward diff --git a/lightly/models/spark.py b/lightly/models/spark.py new file mode 100644 index 000000000..51b84c01a --- /dev/null +++ b/lightly/models/spark.py @@ -0,0 +1,252 @@ +# Code adapted from https://github.com/keyu-tian/SparK/blob/main/ + +import math +from typing import List + +import torch +import torch.nn as nn +from timm.models.layers import trunc_normal_ + +def is_pow2n(n: int) -> bool: + return n > 0 and (n & (n - 1)) == 0 + + +class UNetBlock(nn.Module): + def __init__(self, cin, cout, bn2d): + """ + a UNet block with 2x up sampling + """ + 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): + x = self.up_sample(x) + return self.conv(x) + + +class LightDecoder(nn.Module): + def __init__(self, up_sample_ratio, width=768, sync_batch_norm=True): + 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 sync_batch_norm 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[torch.Tensor]): + x = 0 + for i, d in enumerate(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) + + def extra_repr(self) -> str: + return f"width={self.width}" + + def initialize(self): + 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 SparK(nn.Module): + def __init__( + self, sparse_encoder: encoder.SparseEncoder, dense_decoder: LightDecoder, + mask_ratio=0.6, densify_norm='bn', sbn=False, + ): + super().__init__() + input_size, downsample_raito = sparse_encoder.input_size, sparse_encoder.downsample_raito + self.downsample_raito = downsample_raito + self.fmap_h, self.fmap_w = input_size // downsample_raito, input_size // downsample_raito + self.mask_ratio = mask_ratio + self.len_keep = round(self.fmap_h * self.fmap_w * (1 - mask_ratio)) + + self.sparse_encoder = sparse_encoder + self.dense_decoder = dense_decoder + + self.sbn = sbn + self.hierarchy = len(sparse_encoder.enc_feat_map_chs) + self.densify_norm_str = densify_norm.lower() + self.densify_norms = nn.ModuleList() + self.densify_projs = nn.ModuleList() + self.mask_tokens = nn.ParameterList() + + # build the `densify` layers + e_widths, d_width = self.sparse_encoder.enc_feat_map_chs, self.dense_decoder.width + e_widths: List[int] + for i in range(self.hierarchy): # from the smallest feat map to the largest; i=0: the last feat map; i=1: the second last feat map ... + e_width = e_widths.pop() + # create mask token + p = nn.Parameter(torch.zeros(1, e_width, 1, 1)) + trunc_normal_(p, mean=0, std=.02, a=-.02, b=.02) + self.mask_tokens.append(p) + + # create densify norm + if self.densify_norm_str == 'bn': + densify_norm = (encoder.SparseSyncBatchNorm2d if self.sbn else encoder.SparseBatchNorm2d)(e_width) + elif self.densify_norm_str == 'ln': + densify_norm = encoder.SparseConvNeXtLayerNorm(e_width, data_format='channels_first', sparse=True) + else: + densify_norm = nn.Identity() + self.densify_norms.append(densify_norm) + + # create densify proj + if i == 0 and e_width == d_width: + densify_proj = nn.Identity() # todo: NOTE THAT CONVNEXT-S WOULD USE THIS, because it has a width of 768 that equals to the decoder's width 768 + print(f'[SparK.__init__, densify {i+1}/{self.hierarchy}]: use nn.Identity() as densify_proj') + else: + kernel_size = 1 if i <= 0 else 3 + densify_proj = nn.Conv2d(e_width, d_width, kernel_size=kernel_size, stride=1, padding=kernel_size // 2, bias=True) + print(f'[SparK.__init__, densify {i+1}/{self.hierarchy}]: densify_proj(ksz={kernel_size}, #para={sum(x.numel() for x in densify_proj.parameters()) / 1e6:.2f}M)') + self.densify_projs.append(densify_proj) + + # todo: the decoder's width follows a simple halfing rule; you can change it to any other rule + d_width //= 2 + + print(f'[SparK.__init__] dims of mask_tokens={tuple(p.numel() for p in self.mask_tokens)}') + + # these are deprecated and would never be used; can be removed. + self.register_buffer('imn_m', torch.empty(1, 3, 1, 1)) + self.register_buffer('imn_s', torch.empty(1, 3, 1, 1)) + self.register_buffer('norm_black', torch.zeros(1, 3, input_size, input_size)) + self.vis_active = self.vis_active_ex = self.vis_inp = self.vis_inp_mask = ... + + def mask(self, B: int, device, generator=None): + h, w = self.fmap_h, self.fmap_w + idx = torch.rand(B, h * w, generator=generator).argsort(dim=1) + idx = idx[:, :self.len_keep].to(device) # (B, len_keep) + return torch.zeros(B, h * w, dtype=torch.bool, device=device).scatter_(dim=1, index=idx, value=True).view(B, 1, h, w) + + def forward(self, inp_bchw: torch.Tensor, active_b1ff=None, vis=False): + # step1. Mask + if active_b1ff is None: # rand mask + active_b1ff: torch.BoolTensor = self.mask(inp_bchw.shape[0], inp_bchw.device) # (B, 1, f, f) + encoder._cur_active = active_b1ff # (B, 1, f, f) + active_b1hw = active_b1ff.repeat_interleave(self.downsample_raito, 2).repeat_interleave(self.downsample_raito, 3) # (B, 1, H, W) + masked_bchw = inp_bchw * active_b1hw + + # step2. Encode: get hierarchical encoded sparse features (a list containing 4 feature maps at 4 scales) + fea_bcffs: List[torch.Tensor] = self.sparse_encoder(masked_bchw) + fea_bcffs.reverse() # after reversion: from the smallest feature map to the largest + + # step3. Densify: get hierarchical dense features for decoding + cur_active = active_b1ff # (B, 1, f, f) + to_dec = [] + for i, bcff in enumerate(fea_bcffs): # from the smallest feature map to the largest + if bcff is not None: + bcff = self.densify_norms[i](bcff) + mask_tokens = self.mask_tokens[i].expand_as(bcff) + bcff = torch.where(cur_active.expand_as(bcff), bcff, mask_tokens) # fill in empty (non-active) positions with [mask] tokens + bcff: torch.Tensor = self.densify_projs[i](bcff) + to_dec.append(bcff) + cur_active = cur_active.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3) # dilate the mask map, from (B, 1, f, f) to (B, 1, H, W) + + # step4. Decode and reconstruct + rec_bchw = self.dense_decoder(to_dec) + inp, rec = self.patchify(inp_bchw), self.patchify(rec_bchw) # inp and rec: (B, L = f*f, N = C*downsample_raito**2) + mean = inp.mean(dim=-1, keepdim=True) + var = (inp.var(dim=-1, keepdim=True) + 1e-6) ** .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() + 1e-8) # loss only on masked (non-active) patches + + if vis: + masked_bchw = inp_bchw * active_b1hw + rec_bchw = self.unpatchify(rec * var + mean) + rec_or_inp = torch.where(active_b1hw, inp_bchw, rec_bchw) + return inp_bchw, masked_bchw, rec_or_inp + else: + return recon_loss + + def patchify(self, bchw): + p = self.downsample_raito + h, w = self.fmap_h, self.fmap_w + B, C = bchw.shape[:2] + bchw = bchw.reshape(shape=(B, C, h, p, w, p)) + bchw = torch.einsum('bchpwq->bhwpqc', bchw) + bln = bchw.reshape(shape=(B, h * w, C * p ** 2)) # (B, f*f, 3*downsample_raito**2) + return bln + + def unpatchify(self, bln): + p = self.downsample_raito + 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 __repr__(self): + return ( + f'\n' + f'[SparK.config]: {pformat(self.get_config(), indent=2, width=250)}\n' + f'[SparK.structure]: {super(SparK, self).__repr__().replace(SparK.__name__, "")}' + ) + + def get_config(self): + return { + # self + 'mask_ratio': self.mask_ratio, + 'densify_norm_str': self.densify_norm_str, + 'sbn': self.sbn, 'hierarchy': self.hierarchy, + + # enc + 'sparse_encoder.input_size': self.sparse_encoder.input_size, + # dec + 'dense_decoder.width': self.dense_decoder.width, + } + + def state_dict(self, destination=None, prefix='', keep_vars=False, with_config=False): + state = super(SparK, self).state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars) + if with_config: + state['config'] = self.get_config() + return state + + def load_state_dict(self, state_dict, strict=True): + config: dict = state_dict.pop('config', None) + incompatible_keys = super(SparK, self).load_state_dict(state_dict, strict=strict) + if config is not None: + for k, v in self.get_config().items(): + ckpt_v = config.get(k, None) + if ckpt_v != v: + err = f'[SparseMIM.load_state_dict] config mismatch: this.{k}={v} (ckpt.{k}={ckpt_v})' + if strict: + raise AttributeError(err) + else: + print(err, file=sys.stderr) + return incompatible_keys \ No newline at end of file diff --git a/lightly/models/sparse/decoder.py b/lightly/models/sparse/decoder.py new file mode 100644 index 000000000..e69de29bb diff --git a/lightly/models/sparse/resnet.py b/lightly/models/sparse/resnet.py new file mode 100644 index 000000000..d90ecb2c9 --- /dev/null +++ b/lightly/models/sparse/resnet.py @@ -0,0 +1,223 @@ +# Code adapted from https://github.com/keyu-tian/SparK/blob/main/pretrain/models/resnet.py +from typing import Union + +import torch +import torch.nn as nn +from torch import List, Optional, Tensor + +from lightly.models.modules.sparse.spark import ( + SparseAvgPooling, + SparseBatchNorm2d, + SparseConv2d, + SparseMaxPooling, + SparseSyncBatchNorm2d, +) + + +class SparseMask: + """ + Basic class to store the mask for the sparse model. + """ + + def __init__(self): + self.mask: Union[Tensor, None] = None + + +class SparseResnet(nn.Module): + def __init__( + self, + backbone: nn.Module, + input_size: int, + downsample_ratio: Optional[int] = None, + ): + """Sparse ResNet Encoder as used by SparK [0] + + Default params are the ones explained in the original code base. The backbone is assumed + to follow the same API as the ResNet models from torchvision. + [0] Designing BERT for Convolutional Networks: Sparse and Hierarchical Masked Modeling https://arxiv.org/abs/2301.03580 + + Attributes: + backbone: + Backbone model to extract features from images. Should have both + the methods get_downsample_ratio() and get_feature_map_channels() + implemented. + input_size: + Size of the input image. + + """ + super().__init__() + self.input_size = input_size + if downsample_ratio is None: + self.downsample_ratio = self._get_downsample_ratio(backbone) + else: + self.downsample_ratio = downsample_ratio + self.downsample_ratio, self.enc_feat_map_chs = ( + self._get_downsample_ratio(backbone), + self._get_feature_map_channels(backbone), + ) + self.sparse_mask = SparseMask() + self.sparse_backbone = self.dense_model_to_sparse( + m=backbone, sparse_mask=self.sparse_mask + ) + + def forward( + self, x: torch.Tensor, mask: torch.BoolTensor, hierarchical: bool = False + ) -> Union[torch.Tensor, List[torch.Tensor]]: + """ + Forward pass of the sparse encoder. + + Args: + x: The input tensor. + mask: The mask to apply to the input tensor. This should be a binary mask the + size of the smallest resolution feature map [0] of the backbone model. + + Returns: + The output tensor. + + [0] SparK codebase https://github.com/keyu-tian/SparK/tree/main/pretrain#some-details-how-we-mask-images-and-how-to-set-the-patch-size + """ + assert mask.shape == ( + 1, + 1, + self.input_size // self.downsample_ratio, + self.input_size // self.downsample_ratio, + ), "Mask shape must be (1, 1, H // downsample_ratio, W // downsample_ratio)" + self.sparse_mask.mask = mask + + x = self.sparse_backbone.conv1(x) + x = self.sparse_backbone.bn1(x) + x = self.sparse_backbone.act1(x) + x = self.sparse_backbone.maxpool(x) + + if hierarchical: + ls = [] + x = self.sparse_backbone.layer1(x) + ls.append(x) + x = self.sparse_backbone.layer2(x) + ls.append(x) + x = self.sparse_backbone.layer3(x) + ls.append(x) + x = self.sparse_backbone.layer4(x) + ls.append(x) + return ls + else: + x = self.sparse_backbone.global_pool(x) + x = self.sparse_backbone.fc(x) + return x + + def _get_downsample_ratio(self, backbone: nn.Module): + """ + Try to get the downsample ratio of the backbone model. + + Returns: + The downsample ratio of the backbone model. + """ + try: + return backbone.get_downsample_ratio() + except AttributeError: + x = torch.randn(1, 3, self.input_size, self.input_size) + out = nn.Sequential(*list(backbone.children())[:-2])(x) + + return self.input_size // out.shape[-1] + + def _get_feature_map_channels(self, backbone: nn.Module): + """ + Try to get the number of channels in the feature maps of the backbone model. + + Returns: + The number of channels in the feature maps of the backbone model. + """ + try: + return backbone.get_feature_map_channels() + except AttributeError: + x = torch.randn(1, 3, self.input_size, self.input_size) + out = nn.Sequential(*list(backbone.children())[:-2])(x) + + return out.shape[1] + + def dense_model_to_sparse( + self, m: nn.Module, sparse_mask: SparseMask, sync_batch_norm: bool = False + ) -> nn.Module: + """ + Convert a dense model to a sparse model. + + Args: + m: The dense model to convert. + sparse_mask: The sparse mask to use for the conversion. + sync_batch_norm: Whether to convert BatchNorm2d to SyncBatchNorm. + + Returns: + The sparse model. + """ + with torch.no_grad(): + oup = m + if isinstance(m, nn.Conv2d): + m: nn.Conv2d + bias = m.bias is not None + oup = SparseConv2d( + m.in_channels, + m.out_channels, + kernel_size=m.kernel_size, + stride=m.stride, + padding=m.padding, + dilation=m.dilation, + groups=m.groups, + bias=bias, + padding_mode=m.padding_mode, + ) + oup.sparse_mask = sparse_mask + oup.weight.copy_(m.weight) + if bias: + oup.bias.copy_(m.bias) + elif isinstance(m, nn.MaxPool2d): + 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, + ) + oup.sparse_mask = sparse_mask + elif isinstance(m, nn.AvgPool2d): + 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, + ) + oup.sparse_mask = sparse_mask + elif isinstance(m, (nn.BatchNorm2d, nn.SyncBatchNorm)): + m: nn.BatchNorm2d + oup = ( + SparseSyncBatchNorm2d + if isinstance(m, nn.SyncBatchNorm) + else SparseBatchNorm2d + )( + m.weight.shape[0], + eps=m.eps, + momentum=m.momentum, + affine=m.affine, + track_running_stats=m.track_running_stats, + ) + oup.sparse_mask = sparse_mask + oup.weight.copy_(m.weight) + oup.bias.copy_(m.bias) + oup.running_mean.copy_(m.running_mean) + oup.running_var.copy_(m.running_var) + oup.num_batches_tracked.copy_(m.num_batches_tracked) + if hasattr(m, "qconfig"): + oup.qconfig = m.qconfig + elif isinstance(m, (nn.Conv1d,)): + raise NotImplementedError + + for name, child in m.named_children(): + oup.add_module( + name, self.dense_model_to_sparse(child, sparse_mask=sparse_mask) + ) + del m + return oup From 42bd8497b3ac7ca5902f7652e18ea3e53b6f8bd1 Mon Sep 17 00:00:00 2001 From: John Sutor Date: Tue, 17 Dec 2024 13:28:01 -0500 Subject: [PATCH 02/85] feat(spark): More work on spark --- lightly/models/modules/sparse/spark.py | 207 +++++++++++++++++- lightly/models/spark.py | 146 +++++++----- .../models/sparse/{resnet.py => encoder.py} | 98 +++++---- 3 files changed, 346 insertions(+), 105 deletions(-) rename lightly/models/sparse/{resnet.py => encoder.py} (74%) diff --git a/lightly/models/modules/sparse/spark.py b/lightly/models/modules/sparse/spark.py index e1f47ff07..bf5eb5b0f 100644 --- a/lightly/models/modules/sparse/spark.py +++ b/lightly/models/modules/sparse/spark.py @@ -1,9 +1,56 @@ -# Code adapted from https://github.com/keyu-tian/SparK/blob/main/pretrain/encoder.py +# Code adapted from https://github.com/keyu-tian/SparK/blob/main/pretrain/encoder.py and https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/drop.py + +from typing import Literal, Optional import torch import torch.nn as nn +def drop_path(x, drop_prob: float = 0.0, training: bool = False): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, + the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... + See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for + changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use + 'survival rate' as the argument. + + Attributes: + x: + The input tensor. + drop_prob: + The drop probability of the path. Default: 0. + training: + Whether the model is in training mode. Default: False + """ + if drop_prob == 0.0 or not training: + return x + keep_prob = 1 - drop_prob + shape = (x.shape[0],) + (1,) * ( + x.ndim - 1 + ) + random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) + random_tensor.floor_() + output = x.div(keep_prob) * random_tensor + return output + + +class DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + Attributes: + drop_prob: + The drop probability of the path. Default: None. + """ + + def __init__(self, drop_prob: Optional[float] = None): + super(DropPath, self).__init__() + self.drop_prob = drop_prob + + def forward(self, x): + return drop_path(x, self.drop_prob, self.training) + + def _get_active_ex_or_ii( mask: torch.BoolTensor, H: int, W: int, returning_active_ex: bool = True ): @@ -15,17 +62,16 @@ def _get_active_ex_or_ii( active_ex if returning_active_ex else active_ex.squeeze(1).nonzero(as_tuple=True) - ) # ii: bi, hi, wi + ) -def sp_conv_forward(self, x: torch.Tensor): +def sp_conv_forward(self, x: torch.Tensor) -> torch.Tensor: x = super(type(self), self).forward(x) x *= _get_active_ex_or_ii( self.sparse_mask.mask, H=x.shape[2], W=x.shape[3], returning_active_ex=True - ) # (BCHW) *= (B1HW), mask the output of conv + ) return x - def sp_bn_forward(self, x: torch.Tensor): ii = _get_active_ex_or_ii( self.sparse_mask.mask, H=x.shape[2], W=x.shape[3], returning_active_ex=False @@ -59,3 +105,154 @@ class SparseBatchNorm2d(nn.BatchNorm1d): class SparseSyncBatchNorm2d(nn.SyncBatchNorm): 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). + + + Attributes: + normalized_shape: + Input shape from an expected input of size + normalized_shape or a single integer. + eps: + A value added to the denominator for numerical stability. Default: 1e-6. + data_format: + The ordering of the dimensions in the inputs. Default: "channels_last". + sparse: + Whether to use sparse computation. Default: True. + + """ + + def __init__( + self, + normalized_shape: int, + eps: float = 1e-6, + data_format: Literal["channels_last", "channels_first"] = "channels_last", + sparse: bool = True, + ): + if data_format not in ["channels_last", "channels_first"]: + raise NotImplementedError + super().__init__(normalized_shape, eps, elementwise_affine=True) + self.data_format = data_format + self.sparse = sparse + + def forward(self, x): + if x.ndim == 4: + if self.data_format == "channels_last": + if self.sparse: + ii = _get_active_ex_or_ii( + H=x.shape[1], W=x.shape[2], returning_active_ex=False + ) + nc = x[ii] + nc = super(SparseConvNeXtLayerNorm, self).forward(nc) + + x = torch.zeros_like(x) + x[ii] = nc + return x + else: + return super(SparseConvNeXtLayerNorm, self).forward(x) + else: + if self.sparse: + ii = _get_active_ex_or_ii( + H=x.shape[2], W=x.shape[3], returning_active_ex=False + ) + bhwc = x.permute(0, 2, 3, 1) + nc = bhwc[ii] + nc = super(SparseConvNeXtLayerNorm, self).forward(nc) + + x = torch.zeros_like(bhwc) + x[ii] = nc + return x.permute(0, 3, 1, 2) + else: + u = x.mean(1, keepdim=True) + s = (x - u).pow(2).mean(1, keepdim=True) + x = (x - u) / torch.sqrt(s + self.eps) + x = self.weight[:, None, None] * x + self.bias[:, None, None] + return x + else: # BLC or BC + if self.sparse: + raise NotImplementedError + else: + return super(SparseConvNeXtLayerNorm, self).forward(x) + + def __repr__(self): + return ( + super(SparseConvNeXtLayerNorm, self).__repr__()[:-1] + + f', ch={self.data_format.split("_")[-1]}, sp={self.sparse})' + ) + + +class SparseConvNeXtBlock(nn.Module): + r""" + ConvNeXt Block. 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 + + Attributes: + 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. + sparse: + Whether to use sparse computation. Default: True. + kernel_size: + Kernel size of depthwise convolution. Default: 7. + """ + + def __init__( + self, + dim, + drop_path: float = 0.0, + layer_scale_init_value: float = 1e-6, + sparse: bool = True, + kernel_size=7, + ): + super().__init__() + self.dwconv = nn.Conv2d( + dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim + ) + self.norm = SparseConvNeXtLayerNorm(dim, eps=1e-6, sparse=sparse) + 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() + ) + self.sparse = sparse + + def forward(self, x): + input = x + x = self.dwconv(x) + x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) + x = self.norm(x) + x = self.pwconv1(x) + x = self.act( + x + ) # GELU(0) == (0), so there is no need to mask x (no need to `x *= _get_active_ex_or_ii`) + x = self.pwconv2(x) + if self.gamma is not None: + x = self.gamma * x + x = x.permute(0, 3, 1, 2) # (N, H, W, C) -> (N, C, H, W) + + if self.sparse: + x *= _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 __repr__(self): + return super(SparseConvNeXtBlock, self).__repr__()[:-1] + f", sp={self.sparse})" diff --git a/lightly/models/spark.py b/lightly/models/spark.py index 51b84c01a..c344ef3d6 100644 --- a/lightly/models/spark.py +++ b/lightly/models/spark.py @@ -1,11 +1,14 @@ # Code adapted from https://github.com/keyu-tian/SparK/blob/main/ import math -from typing import List +from typing import List, Literal, Tuple, Union +from pprint import pformat +import sys import torch import torch.nn as nn -from timm.models.layers import trunc_normal_ +from lightly.models.sparse.encoder import SparseResnet +from lightly.models.modules.sparse.spark import SparseBatchNorm2d, SparseSyncBatchNorm2d, SparseConvNeXtLayerNorm def is_pow2n(n: int) -> bool: return n > 0 and (n & (n - 1)) == 0 @@ -65,11 +68,11 @@ def extra_repr(self) -> str: def initialize(self): for m in self.modules(): if isinstance(m, nn.Linear): - trunc_normal_(m.weight, std=0.02) + torch.nn.init.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) + torch.nn.init.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)): @@ -83,14 +86,35 @@ def initialize(self): nn.init.constant_(m.weight, 1.0) class SparK(nn.Module): + """ + The SparK model as used by SparK [0] + + Default params are the ones explained in the original code base. The backbone is assumed + to follow the same API as the ResNet models from torchvision or a ConvNext model. + [0] Designing BERT for Convolutional Networks: Sparse and Hierarchical Masked Modeling https://arxiv.org/abs/2301.03580 + + Attributes: + sparse_encoder: + Sparse encoder to extract features from images. Should have both + the methods get_downsample_ratio() and get_feature_map_channels() + implemented. + dense_decoder: + Dense decoder to reconstruct the image from the sparse features. + mask_ratio: + Ratio of the image to mask. Default: 0.6 + densify_norm: + Type of normalization to use for densification. Default: 'bn' + sbn: + Whether to use SyncBatchNorm. Default: False + """ def __init__( - self, sparse_encoder: encoder.SparseEncoder, dense_decoder: LightDecoder, - mask_ratio=0.6, densify_norm='bn', sbn=False, + self, sparse_encoder: SparseResnet, dense_decoder: LightDecoder, + mask_ratio: float=0.6, densify_norm: Literal['batch_norm', 'layer_norm', 'identity']='bn', sbn: bool=False, ): super().__init__() - input_size, downsample_raito = sparse_encoder.input_size, sparse_encoder.downsample_raito - self.downsample_raito = downsample_raito - self.fmap_h, self.fmap_w = input_size // downsample_raito, input_size // downsample_raito + input_size, downsample_ratio = sparse_encoder.input_size, sparse_encoder.downsample_ratio + self.downsample_ratio = downsample_ratio + self.fmap_h, self.fmap_w = input_size // downsample_ratio, input_size // downsample_ratio self.mask_ratio = mask_ratio self.len_keep = round(self.fmap_h * self.fmap_w * (1 - mask_ratio)) @@ -104,106 +128,114 @@ def __init__( self.densify_projs = nn.ModuleList() self.mask_tokens = nn.ParameterList() - # build the `densify` layers e_widths, d_width = self.sparse_encoder.enc_feat_map_chs, self.dense_decoder.width e_widths: List[int] - for i in range(self.hierarchy): # from the smallest feat map to the largest; i=0: the last feat map; i=1: the second last feat map ... + for i in range(self.hierarchy): e_width = e_widths.pop() - # create mask token p = nn.Parameter(torch.zeros(1, e_width, 1, 1)) - trunc_normal_(p, mean=0, std=.02, a=-.02, b=.02) + torch.nn.init.trunc_normal_(p, mean=0, std=.02, a=-.02, b=.02) self.mask_tokens.append(p) - # create densify norm - if self.densify_norm_str == 'bn': - densify_norm = (encoder.SparseSyncBatchNorm2d if self.sbn else encoder.SparseBatchNorm2d)(e_width) - elif self.densify_norm_str == 'ln': - densify_norm = encoder.SparseConvNeXtLayerNorm(e_width, data_format='channels_first', sparse=True) + + if self.densify_norm_str == 'batch_norm': + densify_norm = (SparseSyncBatchNorm2d if self.sbn else SparseBatchNorm2d)(e_width) + elif self.densify_norm_str == 'layer_norm': + densify_norm = SparseConvNeXtLayerNorm(e_width, data_format='channels_first', sparse=True) else: densify_norm = nn.Identity() self.densify_norms.append(densify_norm) - # create densify proj if i == 0 and e_width == d_width: - densify_proj = nn.Identity() # todo: NOTE THAT CONVNEXT-S WOULD USE THIS, because it has a width of 768 that equals to the decoder's width 768 + densify_proj = nn.Identity() print(f'[SparK.__init__, densify {i+1}/{self.hierarchy}]: use nn.Identity() as densify_proj') else: kernel_size = 1 if i <= 0 else 3 densify_proj = nn.Conv2d(e_width, d_width, kernel_size=kernel_size, stride=1, padding=kernel_size // 2, bias=True) print(f'[SparK.__init__, densify {i+1}/{self.hierarchy}]: densify_proj(ksz={kernel_size}, #para={sum(x.numel() for x in densify_proj.parameters()) / 1e6:.2f}M)') self.densify_projs.append(densify_proj) - - # todo: the decoder's width follows a simple halfing rule; you can change it to any other rule d_width //= 2 print(f'[SparK.__init__] dims of mask_tokens={tuple(p.numel() for p in self.mask_tokens)}') - # these are deprecated and would never be used; can be removed. - self.register_buffer('imn_m', torch.empty(1, 3, 1, 1)) - self.register_buffer('imn_s', torch.empty(1, 3, 1, 1)) - self.register_buffer('norm_black', torch.zeros(1, 3, input_size, input_size)) - self.vis_active = self.vis_active_ex = self.vis_inp = self.vis_inp_mask = ... def mask(self, B: int, device, generator=None): + """ + Generate a mask for the input tensor + + Attributes: + B: + Batch size + device: + Device to put the mask on + generator: + Random number generator + """ h, w = self.fmap_h, self.fmap_w idx = torch.rand(B, h * w, generator=generator).argsort(dim=1) idx = idx[:, :self.len_keep].to(device) # (B, len_keep) return torch.zeros(B, h * w, dtype=torch.bool, device=device).scatter_(dim=1, index=idx, value=True).view(B, 1, h, w) - def forward(self, inp_bchw: torch.Tensor, active_b1ff=None, vis=False): - # step1. Mask - if active_b1ff is None: # rand mask - active_b1ff: torch.BoolTensor = self.mask(inp_bchw.shape[0], inp_bchw.device) # (B, 1, f, f) - encoder._cur_active = active_b1ff # (B, 1, f, f) - active_b1hw = active_b1ff.repeat_interleave(self.downsample_raito, 2).repeat_interleave(self.downsample_raito, 3) # (B, 1, H, W) + def forward(self, inp_bchw: torch.Tensor, active_b1ff: torch.BoolTensor=None, return_loss: bool=True) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """ + Forward pass of the SparK model + + Attributes: + inp_bchw: + Input tensor + active_b1ff: + Active mask + return_loss: + Whether to return the loss + """ + if active_b1ff is None: + active_b1ff: torch.BoolTensor = self.mask(inp_bchw.shape[0], inp_bchw.device) + active_b1hw = active_b1ff.repeat_interleave(self.downsample_ratio, 2).repeat_interleave(self.downsample_ratio, 3) masked_bchw = inp_bchw * active_b1hw - # step2. Encode: get hierarchical encoded sparse features (a list containing 4 feature maps at 4 scales) - fea_bcffs: List[torch.Tensor] = self.sparse_encoder(masked_bchw) - fea_bcffs.reverse() # after reversion: from the smallest feature map to the largest + fea_bcffs: List[torch.Tensor] = self.sparse_encoder.forward(masked_bchw, active_b1ff, hierarchical=True) + fea_bcffs.reverse() - # step3. Densify: get hierarchical dense features for decoding - cur_active = active_b1ff # (B, 1, f, f) + cur_active = active_b1ff to_dec = [] - for i, bcff in enumerate(fea_bcffs): # from the smallest feature map to the largest + for i, bcff in enumerate(fea_bcffs): if bcff is not None: bcff = self.densify_norms[i](bcff) mask_tokens = self.mask_tokens[i].expand_as(bcff) - bcff = torch.where(cur_active.expand_as(bcff), bcff, mask_tokens) # fill in empty (non-active) positions with [mask] tokens + bcff = torch.where(cur_active.expand_as(bcff), bcff, mask_tokens) bcff: torch.Tensor = self.densify_projs[i](bcff) to_dec.append(bcff) - cur_active = cur_active.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3) # dilate the mask map, from (B, 1, f, f) to (B, 1, H, W) + cur_active = cur_active.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3) - # step4. Decode and reconstruct rec_bchw = self.dense_decoder(to_dec) - inp, rec = self.patchify(inp_bchw), self.patchify(rec_bchw) # inp and rec: (B, L = f*f, N = C*downsample_raito**2) + inp, rec = self.patchify(inp_bchw), self.patchify(rec_bchw) mean = inp.mean(dim=-1, keepdim=True) var = (inp.var(dim=-1, keepdim=True) + 1e-6) ** .5 inp = (inp - mean) / var - l2_loss = ((rec - inp) ** 2).mean(dim=2, keepdim=False) # (B, L, C) ==mean==> (B, L) + l2_loss = ((rec - inp) ** 2).mean(dim=2, keepdim=False) - 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() + 1e-8) # loss only on masked (non-active) patches + non_active = active_b1ff.logical_not().int().view(active_b1ff.shape[0], -1) + recon_loss = l2_loss.mul_(non_active).sum() / (non_active.sum() + 1e-8) - if vis: - masked_bchw = inp_bchw * active_b1hw - rec_bchw = self.unpatchify(rec * var + mean) - rec_or_inp = torch.where(active_b1hw, inp_bchw, rec_bchw) - return inp_bchw, masked_bchw, rec_or_inp - else: + if return_loss: return recon_loss + + masked_bchw = inp_bchw * active_b1hw + rec_bchw = self.unpatchify(rec * var + mean) + rec_or_inp = torch.where(active_b1hw, inp_bchw, rec_bchw) + return inp_bchw, masked_bchw, rec_or_inp + def patchify(self, bchw): - p = self.downsample_raito + p = self.downsample_ratio h, w = self.fmap_h, self.fmap_w B, C = bchw.shape[:2] bchw = bchw.reshape(shape=(B, C, h, p, w, p)) bchw = torch.einsum('bchpwq->bhwpqc', bchw) - bln = bchw.reshape(shape=(B, h * w, C * p ** 2)) # (B, f*f, 3*downsample_raito**2) + bln = bchw.reshape(shape=(B, h * w, C * p ** 2)) return bln def unpatchify(self, bln): - p = self.downsample_raito + 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)) @@ -220,14 +252,10 @@ def __repr__(self): def get_config(self): return { - # self 'mask_ratio': self.mask_ratio, 'densify_norm_str': self.densify_norm_str, 'sbn': self.sbn, 'hierarchy': self.hierarchy, - - # enc 'sparse_encoder.input_size': self.sparse_encoder.input_size, - # dec 'dense_decoder.width': self.dense_decoder.width, } diff --git a/lightly/models/sparse/resnet.py b/lightly/models/sparse/encoder.py similarity index 74% rename from lightly/models/sparse/resnet.py rename to lightly/models/sparse/encoder.py index d90ecb2c9..bc4f6704a 100644 --- a/lightly/models/sparse/resnet.py +++ b/lightly/models/sparse/encoder.py @@ -11,8 +11,14 @@ SparseConv2d, SparseMaxPooling, SparseSyncBatchNorm2d, + SparseConvNeXtBlock, + SparseConvNeXtLayerNorm ) +from torchvision.models.convnext import CNBlock, LayerNorm2d + +from torchvision.models import ResNet, ConvNeXt + class SparseMask: """ @@ -23,17 +29,17 @@ def __init__(self): self.mask: Union[Tensor, None] = None -class SparseResnet(nn.Module): +class SparseEncoder(nn.Module): def __init__( self, backbone: nn.Module, input_size: int, downsample_ratio: Optional[int] = None, ): - """Sparse ResNet Encoder as used by SparK [0] + """Sparse encoder as used by SparK [0] Default params are the ones explained in the original code base. The backbone is assumed - to follow the same API as the ResNet models from torchvision. + to follow the same API as the ResNet or ConvNext models from torchvision. [0] Designing BERT for Convolutional Networks: Sparse and Hierarchical Masked Modeling https://arxiv.org/abs/2301.03580 Attributes: @@ -43,17 +49,18 @@ def __init__( implemented. input_size: Size of the input image. + downsample_ratio: """ super().__init__() self.input_size = input_size if downsample_ratio is None: - self.downsample_ratio = self._get_downsample_ratio(backbone) + self.downsample_ratio = self.get_downsample_ratio(backbone) else: self.downsample_ratio = downsample_ratio - self.downsample_ratio, self.enc_feat_map_chs = ( - self._get_downsample_ratio(backbone), - self._get_feature_map_channels(backbone), + + self.enc_feat_map_chs = ( + self.get_feature_map_channels(backbone), ) self.sparse_mask = SparseMask() self.sparse_backbone = self.dense_model_to_sparse( @@ -84,28 +91,45 @@ def forward( ), "Mask shape must be (1, 1, H // downsample_ratio, W // downsample_ratio)" self.sparse_mask.mask = mask - x = self.sparse_backbone.conv1(x) - x = self.sparse_backbone.bn1(x) - x = self.sparse_backbone.act1(x) - x = self.sparse_backbone.maxpool(x) - - if hierarchical: - ls = [] - x = self.sparse_backbone.layer1(x) - ls.append(x) - x = self.sparse_backbone.layer2(x) - ls.append(x) - x = self.sparse_backbone.layer3(x) - ls.append(x) - x = self.sparse_backbone.layer4(x) - ls.append(x) - return ls + ls = [] + + if isinstance(self.sparse_backbone, ConvNeXt): + if hierarchical: + for i in range(0,8,2): + x = self.sparse_backbone.features[i](x) + x = self.sparse_backbone.features[i+1](x) + ls.append(x) + else: + x = self.sparse_backbone.avgpool(x) + x = self.sparse_backbone.classifier(x) + return x + + elif isinstance(self.sparse_backbone, ResNet): + x = self.sparse_backbone.conv1(x) + x = self.sparse_backbone.bn1(x) + x = self.sparse_backbone.act1(x) + x = self.sparse_backbone.maxpool(x) + + if hierarchical: + ls = [] + x = self.sparse_backbone.layer1(x) + ls.append(x) + x = self.sparse_backbone.layer2(x) + ls.append(x) + x = self.sparse_backbone.layer3(x) + ls.append(x) + x = self.sparse_backbone.layer4(x) + ls.append(x) + return ls + else: + x = self.sparse_backbone.global_pool(x) + x = self.sparse_backbone.fc(x) + return x + else: - x = self.sparse_backbone.global_pool(x) - x = self.sparse_backbone.fc(x) - return x + raise NotImplementedError("Backbone not supported") - def _get_downsample_ratio(self, backbone: nn.Module): + def get_downsample_ratio(self, backbone: nn.Module) -> int: """ Try to get the downsample ratio of the backbone model. @@ -120,20 +144,6 @@ def _get_downsample_ratio(self, backbone: nn.Module): return self.input_size // out.shape[-1] - def _get_feature_map_channels(self, backbone: nn.Module): - """ - Try to get the number of channels in the feature maps of the backbone model. - - Returns: - The number of channels in the feature maps of the backbone model. - """ - try: - return backbone.get_feature_map_channels() - except AttributeError: - x = torch.randn(1, 3, self.input_size, self.input_size) - out = nn.Sequential(*list(backbone.children())[:-2])(x) - - return out.shape[1] def dense_model_to_sparse( self, m: nn.Module, sparse_mask: SparseMask, sync_batch_norm: bool = False @@ -212,6 +222,12 @@ def dense_model_to_sparse( oup.num_batches_tracked.copy_(m.num_batches_tracked) if hasattr(m, "qconfig"): oup.qconfig = m.qconfig + elif isinstance(m, CNBlock): + m: CNBlock + oup = SparseConvNeXtBlock( + m.weight.shape[0], + m.layer_scale. + ) elif isinstance(m, (nn.Conv1d,)): raise NotImplementedError From 724819c765cbec2f9b66ac2fc01a7c299745af6f Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 5 Feb 2026 19:12:28 -0300 Subject: [PATCH 03/85] refactor: should not nest modules. --- .../modules/{sparse/spark.py => sparse_spark.py} | 15 +++++++-------- 1 file changed, 7 insertions(+), 8 deletions(-) rename lightly/models/modules/{sparse/spark.py => sparse_spark.py} (97%) diff --git a/lightly/models/modules/sparse/spark.py b/lightly/models/modules/sparse_spark.py similarity index 97% rename from lightly/models/modules/sparse/spark.py rename to lightly/models/modules/sparse_spark.py index bf5eb5b0f..481246fc9 100644 --- a/lightly/models/modules/sparse/spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -26,11 +26,9 @@ def drop_path(x, drop_prob: float = 0.0, training: bool = False): if drop_prob == 0.0 or not training: return x keep_prob = 1 - drop_prob - shape = (x.shape[0],) + (1,) * ( - x.ndim - 1 - ) + shape = (x.shape[0],) + (1,) * (x.ndim - 1) random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) - random_tensor.floor_() + random_tensor.floor_() output = x.div(keep_prob) * random_tensor return output @@ -62,16 +60,17 @@ def _get_active_ex_or_ii( active_ex if returning_active_ex else active_ex.squeeze(1).nonzero(as_tuple=True) - ) + ) def sp_conv_forward(self, x: torch.Tensor) -> torch.Tensor: x = super(type(self), self).forward(x) x *= _get_active_ex_or_ii( self.sparse_mask.mask, H=x.shape[2], W=x.shape[3], returning_active_ex=True - ) + ) return x + def sp_bn_forward(self, x: torch.Tensor): ii = _get_active_ex_or_ii( self.sparse_mask.mask, H=x.shape[2], W=x.shape[3], returning_active_ex=False @@ -155,7 +154,7 @@ def forward(self, x): return x else: return super(SparseConvNeXtLayerNorm, self).forward(x) - else: + else: if self.sparse: ii = _get_active_ex_or_ii( H=x.shape[2], W=x.shape[3], returning_active_ex=False @@ -182,7 +181,7 @@ def forward(self, x): def __repr__(self): return ( super(SparseConvNeXtLayerNorm, self).__repr__()[:-1] - + f', ch={self.data_format.split("_")[-1]}, sp={self.sparse})' + + f", ch={self.data_format.split('_')[-1]}, sp={self.sparse})" ) From b0023f0c5b94efc49495ac9828967a2bdb6cdf77 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 5 Feb 2026 19:13:33 -0300 Subject: [PATCH 04/85] refactor: removed empty file --- lightly/models/sparse/decoder.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 lightly/models/sparse/decoder.py diff --git a/lightly/models/sparse/decoder.py b/lightly/models/sparse/decoder.py deleted file mode 100644 index e69de29bb..000000000 From 81209909850155b06760934ab1280d781ce754e6 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 5 Feb 2026 19:15:39 -0300 Subject: [PATCH 05/85] refactor: adhered to correct directory structure. --- lightly/models/modules/sparse_spark.py | 229 +++++++++++++++++++++++ lightly/models/spark.py | 211 ++++++++++++++-------- lightly/models/sparse/encoder.py | 239 ------------------------- 3 files changed, 364 insertions(+), 315 deletions(-) delete mode 100644 lightly/models/sparse/encoder.py diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 481246fc9..bf6c29f16 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -4,6 +4,14 @@ import torch import torch.nn as nn +from typing import Union + +import torch +import torch.nn as nn +from torch import List, Optional, Tensor +from torchvision.models.convnext import CNBlock, LayerNorm2d + +from torchvision.models import ResNet, ConvNeXt def drop_path(x, drop_prob: float = 0.0, training: bool = False): @@ -255,3 +263,224 @@ def forward(self, x): def __repr__(self): return super(SparseConvNeXtBlock, self).__repr__()[:-1] + f", sp={self.sparse})" + +# Code adapted from https://github.com/keyu-tian/SparK/blob/main/pretrain/models/resnet.py + + +class SparseMask: + """ + Basic class to store the mask for the sparse model. + """ + + def __init__(self): + self.mask: Union[Tensor, None] = None + + +class SparseEncoder(nn.Module): + def __init__( + self, + backbone: nn.Module, + input_size: int, + downsample_ratio: Optional[int] = None, + ): + """Sparse encoder as used by SparK [0] + + Default params are the ones explained in the original code base. The backbone is assumed + to follow the same API as the ResNet or ConvNext models from torchvision. + [0] Designing BERT for Convolutional Networks: Sparse and Hierarchical Masked Modeling https://arxiv.org/abs/2301.03580 + + Attributes: + backbone: + Backbone model to extract features from images. Should have both + the methods get_downsample_ratio() and get_feature_map_channels() + implemented. + input_size: + Size of the input image. + downsample_ratio: + + """ + super().__init__() + self.input_size = input_size + if downsample_ratio is None: + self.downsample_ratio = self.get_downsample_ratio(backbone) + else: + self.downsample_ratio = downsample_ratio + + self.enc_feat_map_chs = ( + self.get_feature_map_channels(backbone), + ) + self.sparse_mask = SparseMask() + self.sparse_backbone = self.dense_model_to_sparse( + m=backbone, sparse_mask=self.sparse_mask + ) + + def forward( + self, x: torch.Tensor, mask: torch.BoolTensor, hierarchical: bool = False + ) -> Union[torch.Tensor, List[torch.Tensor]]: + """ + Forward pass of the sparse encoder. + + Args: + x: The input tensor. + mask: The mask to apply to the input tensor. This should be a binary mask the + size of the smallest resolution feature map [0] of the backbone model. + + Returns: + The output tensor. + + [0] SparK codebase https://github.com/keyu-tian/SparK/tree/main/pretrain#some-details-how-we-mask-images-and-how-to-set-the-patch-size + """ + assert mask.shape == ( + 1, + 1, + self.input_size // self.downsample_ratio, + self.input_size // self.downsample_ratio, + ), "Mask shape must be (1, 1, H // downsample_ratio, W // downsample_ratio)" + self.sparse_mask.mask = mask + + ls = [] + + if isinstance(self.sparse_backbone, ConvNeXt): + if hierarchical: + for i in range(0,8,2): + x = self.sparse_backbone.features[i](x) + x = self.sparse_backbone.features[i+1](x) + ls.append(x) + else: + x = self.sparse_backbone.avgpool(x) + x = self.sparse_backbone.classifier(x) + return x + + elif isinstance(self.sparse_backbone, ResNet): + x = self.sparse_backbone.conv1(x) + x = self.sparse_backbone.bn1(x) + x = self.sparse_backbone.act1(x) + x = self.sparse_backbone.maxpool(x) + + if hierarchical: + ls = [] + x = self.sparse_backbone.layer1(x) + ls.append(x) + x = self.sparse_backbone.layer2(x) + ls.append(x) + x = self.sparse_backbone.layer3(x) + ls.append(x) + x = self.sparse_backbone.layer4(x) + ls.append(x) + return ls + else: + x = self.sparse_backbone.global_pool(x) + x = self.sparse_backbone.fc(x) + return x + + else: + raise NotImplementedError("Backbone not supported") + + def get_downsample_ratio(self, backbone: nn.Module) -> int: + """ + Try to get the downsample ratio of the backbone model. + + Returns: + The downsample ratio of the backbone model. + """ + try: + return backbone.get_downsample_ratio() + except AttributeError: + x = torch.randn(1, 3, self.input_size, self.input_size) + out = nn.Sequential(*list(backbone.children())[:-2])(x) + + return self.input_size // out.shape[-1] + + + def dense_model_to_sparse( + self, m: nn.Module, sparse_mask: SparseMask, sync_batch_norm: bool = False + ) -> nn.Module: + """ + Convert a dense model to a sparse model. + + Args: + m: The dense model to convert. + sparse_mask: The sparse mask to use for the conversion. + sync_batch_norm: Whether to convert BatchNorm2d to SyncBatchNorm. + + Returns: + The sparse model. + """ + with torch.no_grad(): + oup = m + if isinstance(m, nn.Conv2d): + m: nn.Conv2d + bias = m.bias is not None + oup = SparseConv2d( + m.in_channels, + m.out_channels, + kernel_size=m.kernel_size, + stride=m.stride, + padding=m.padding, + dilation=m.dilation, + groups=m.groups, + bias=bias, + padding_mode=m.padding_mode, + ) + oup.sparse_mask = sparse_mask + oup.weight.copy_(m.weight) + if bias: + oup.bias.copy_(m.bias) + elif isinstance(m, nn.MaxPool2d): + 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, + ) + oup.sparse_mask = sparse_mask + elif isinstance(m, nn.AvgPool2d): + 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, + ) + oup.sparse_mask = sparse_mask + elif isinstance(m, (nn.BatchNorm2d, nn.SyncBatchNorm)): + m: nn.BatchNorm2d + oup = ( + SparseSyncBatchNorm2d + if isinstance(m, nn.SyncBatchNorm) + else SparseBatchNorm2d + )( + m.weight.shape[0], + eps=m.eps, + momentum=m.momentum, + affine=m.affine, + track_running_stats=m.track_running_stats, + ) + oup.sparse_mask = sparse_mask + oup.weight.copy_(m.weight) + oup.bias.copy_(m.bias) + oup.running_mean.copy_(m.running_mean) + oup.running_var.copy_(m.running_var) + oup.num_batches_tracked.copy_(m.num_batches_tracked) + if hasattr(m, "qconfig"): + oup.qconfig = m.qconfig + elif isinstance(m, CNBlock): + m: CNBlock + oup = SparseConvNeXtBlock( + m.weight.shape[0], + m.layer_scale. + ) + elif isinstance(m, (nn.Conv1d,)): + raise NotImplementedError + + for name, child in m.named_children(): + oup.add_module( + name, self.dense_model_to_sparse(child, sparse_mask=sparse_mask) + ) + del m + return oup \ No newline at end of file diff --git a/lightly/models/spark.py b/lightly/models/spark.py index c344ef3d6..4c1b4ca30 100644 --- a/lightly/models/spark.py +++ b/lightly/models/spark.py @@ -1,14 +1,20 @@ # Code adapted from https://github.com/keyu-tian/SparK/blob/main/ import math -from typing import List, Literal, Tuple, Union -from pprint import pformat import sys +from pprint import pformat +from typing import List, Literal, Tuple, Union import torch import torch.nn as nn + +from lightly.models.modules.spark import ( + SparseBatchNorm2d, + SparseConvNeXtLayerNorm, + SparseSyncBatchNorm2d, +) from lightly.models.sparse.encoder import SparseResnet -from lightly.models.modules.sparse.spark import SparseBatchNorm2d, SparseSyncBatchNorm2d, SparseConvNeXtLayerNorm + def is_pow2n(n: int) -> bool: return n > 0 and (n & (n - 1)) == 0 @@ -85,6 +91,7 @@ def initialize(self): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) + class SparK(nn.Module): """ The SparK model as used by SparK [0] @@ -107,57 +114,86 @@ class SparK(nn.Module): sbn: Whether to use SyncBatchNorm. Default: False """ + def __init__( - self, sparse_encoder: SparseResnet, dense_decoder: LightDecoder, - mask_ratio: float=0.6, densify_norm: Literal['batch_norm', 'layer_norm', 'identity']='bn', sbn: bool=False, + self, + sparse_encoder: SparseResnet, + dense_decoder: LightDecoder, + mask_ratio: float = 0.6, + densify_norm: Literal["batch_norm", "layer_norm", "identity"] = "bn", + sbn: bool = False, ): super().__init__() - input_size, downsample_ratio = sparse_encoder.input_size, sparse_encoder.downsample_ratio + input_size, downsample_ratio = ( + sparse_encoder.input_size, + sparse_encoder.downsample_ratio, + ) self.downsample_ratio = downsample_ratio - self.fmap_h, self.fmap_w = input_size // downsample_ratio, input_size // downsample_ratio + self.fmap_h, self.fmap_w = ( + input_size // downsample_ratio, + input_size // downsample_ratio, + ) self.mask_ratio = mask_ratio self.len_keep = round(self.fmap_h * self.fmap_w * (1 - mask_ratio)) - + self.sparse_encoder = sparse_encoder self.dense_decoder = dense_decoder - + self.sbn = sbn self.hierarchy = len(sparse_encoder.enc_feat_map_chs) self.densify_norm_str = densify_norm.lower() self.densify_norms = nn.ModuleList() self.densify_projs = nn.ModuleList() self.mask_tokens = nn.ParameterList() - - e_widths, d_width = self.sparse_encoder.enc_feat_map_chs, self.dense_decoder.width + + e_widths, d_width = ( + self.sparse_encoder.enc_feat_map_chs, + self.dense_decoder.width, + ) e_widths: List[int] for i in range(self.hierarchy): e_width = e_widths.pop() p = nn.Parameter(torch.zeros(1, e_width, 1, 1)) - torch.nn.init.trunc_normal_(p, mean=0, std=.02, a=-.02, b=.02) + torch.nn.init.trunc_normal_(p, mean=0, std=0.02, a=-0.02, b=0.02) self.mask_tokens.append(p) - - - if self.densify_norm_str == 'batch_norm': - densify_norm = (SparseSyncBatchNorm2d if self.sbn else SparseBatchNorm2d)(e_width) - elif self.densify_norm_str == 'layer_norm': - densify_norm = SparseConvNeXtLayerNorm(e_width, data_format='channels_first', sparse=True) + + if self.densify_norm_str == "batch_norm": + densify_norm = ( + SparseSyncBatchNorm2d if self.sbn else SparseBatchNorm2d + )(e_width) + elif self.densify_norm_str == "layer_norm": + densify_norm = SparseConvNeXtLayerNorm( + e_width, data_format="channels_first", sparse=True + ) else: densify_norm = nn.Identity() self.densify_norms.append(densify_norm) - + if i == 0 and e_width == d_width: - densify_proj = nn.Identity() - print(f'[SparK.__init__, densify {i+1}/{self.hierarchy}]: use nn.Identity() as densify_proj') + densify_proj = nn.Identity() + print( + f"[SparK.__init__, densify {i + 1}/{self.hierarchy}]: use nn.Identity() as densify_proj" + ) else: kernel_size = 1 if i <= 0 else 3 - densify_proj = nn.Conv2d(e_width, d_width, kernel_size=kernel_size, stride=1, padding=kernel_size // 2, bias=True) - print(f'[SparK.__init__, densify {i+1}/{self.hierarchy}]: densify_proj(ksz={kernel_size}, #para={sum(x.numel() for x in densify_proj.parameters()) / 1e6:.2f}M)') + densify_proj = nn.Conv2d( + e_width, + d_width, + kernel_size=kernel_size, + stride=1, + padding=kernel_size // 2, + bias=True, + ) + print( + f"[SparK.__init__, densify {i + 1}/{self.hierarchy}]: densify_proj(ksz={kernel_size}, #para={sum(x.numel() for x in densify_proj.parameters()) / 1e6:.2f}M)" + ) self.densify_projs.append(densify_proj) d_width //= 2 - - print(f'[SparK.__init__] dims of mask_tokens={tuple(p.numel() for p in self.mask_tokens)}') - - + + print( + f"[SparK.__init__] dims of mask_tokens={tuple(p.numel() for p in self.mask_tokens)}" + ) + def mask(self, B: int, device, generator=None): """ Generate a mask for the input tensor @@ -172,10 +208,19 @@ def mask(self, B: int, device, generator=None): """ h, w = self.fmap_h, self.fmap_w idx = torch.rand(B, h * w, generator=generator).argsort(dim=1) - idx = idx[:, :self.len_keep].to(device) # (B, len_keep) - return torch.zeros(B, h * w, dtype=torch.bool, device=device).scatter_(dim=1, index=idx, value=True).view(B, 1, h, w) - - def forward(self, inp_bchw: torch.Tensor, active_b1ff: torch.BoolTensor=None, return_loss: bool=True) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + idx = idx[:, : self.len_keep].to(device) # (B, len_keep) + return ( + torch.zeros(B, h * w, dtype=torch.bool, device=device) + .scatter_(dim=1, index=idx, value=True) + .view(B, 1, h, w) + ) + + def forward( + self, + inp_bchw: torch.Tensor, + active_b1ff: torch.BoolTensor = None, + return_loss: bool = True, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: """ Forward pass of the SparK model @@ -187,94 +232,108 @@ def forward(self, inp_bchw: torch.Tensor, active_b1ff: torch.BoolTensor=None, re return_loss: Whether to return the loss """ - if active_b1ff is None: - active_b1ff: torch.BoolTensor = self.mask(inp_bchw.shape[0], inp_bchw.device) - active_b1hw = active_b1ff.repeat_interleave(self.downsample_ratio, 2).repeat_interleave(self.downsample_ratio, 3) + if active_b1ff is None: + active_b1ff: torch.BoolTensor = self.mask( + inp_bchw.shape[0], inp_bchw.device + ) + active_b1hw = active_b1ff.repeat_interleave( + self.downsample_ratio, 2 + ).repeat_interleave(self.downsample_ratio, 3) masked_bchw = inp_bchw * active_b1hw - - fea_bcffs: List[torch.Tensor] = self.sparse_encoder.forward(masked_bchw, active_b1ff, hierarchical=True) - fea_bcffs.reverse() - - cur_active = active_b1ff + + fea_bcffs: List[torch.Tensor] = self.sparse_encoder.forward( + masked_bchw, active_b1ff, hierarchical=True + ) + fea_bcffs.reverse() + + cur_active = active_b1ff to_dec = [] - for i, bcff in enumerate(fea_bcffs): + for i, bcff in enumerate(fea_bcffs): if bcff is not None: bcff = self.densify_norms[i](bcff) mask_tokens = self.mask_tokens[i].expand_as(bcff) - bcff = torch.where(cur_active.expand_as(bcff), bcff, mask_tokens) + bcff = torch.where(cur_active.expand_as(bcff), bcff, mask_tokens) bcff: torch.Tensor = self.densify_projs[i](bcff) to_dec.append(bcff) - cur_active = cur_active.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3) - + cur_active = cur_active.repeat_interleave(2, dim=2).repeat_interleave( + 2, dim=3 + ) + rec_bchw = self.dense_decoder(to_dec) - inp, rec = self.patchify(inp_bchw), self.patchify(rec_bchw) + inp, rec = self.patchify(inp_bchw), self.patchify(rec_bchw) mean = inp.mean(dim=-1, keepdim=True) - var = (inp.var(dim=-1, keepdim=True) + 1e-6) ** .5 + var = (inp.var(dim=-1, keepdim=True) + 1e-6) ** 0.5 inp = (inp - mean) / var - l2_loss = ((rec - inp) ** 2).mean(dim=2, keepdim=False) - - non_active = active_b1ff.logical_not().int().view(active_b1ff.shape[0], -1) - recon_loss = l2_loss.mul_(non_active).sum() / (non_active.sum() + 1e-8) - + l2_loss = ((rec - inp) ** 2).mean(dim=2, keepdim=False) + + non_active = active_b1ff.logical_not().int().view(active_b1ff.shape[0], -1) + recon_loss = l2_loss.mul_(non_active).sum() / (non_active.sum() + 1e-8) + if return_loss: return recon_loss - + masked_bchw = inp_bchw * active_b1hw rec_bchw = self.unpatchify(rec * var + mean) rec_or_inp = torch.where(active_b1hw, inp_bchw, rec_bchw) return inp_bchw, masked_bchw, rec_or_inp - def patchify(self, bchw): p = self.downsample_ratio h, w = self.fmap_h, self.fmap_w B, C = bchw.shape[:2] bchw = bchw.reshape(shape=(B, C, h, p, w, p)) - bchw = torch.einsum('bchpwq->bhwpqc', bchw) - bln = bchw.reshape(shape=(B, h * w, C * p ** 2)) + bchw = torch.einsum("bchpwq->bhwpqc", bchw) + bln = bchw.reshape(shape=(B, h * w, C * p**2)) return bln - + def unpatchify(self, bln): p = self.downsample_ratio h, w = self.fmap_h, self.fmap_w - B, C = bln.shape[0], bln.shape[-1] // p ** 2 + 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) + bln = torch.einsum("bhwpqc->bchpwq", bln) bchw = bln.reshape(shape=(B, C, h * p, w * p)) return bchw - + def __repr__(self): return ( - f'\n' - f'[SparK.config]: {pformat(self.get_config(), indent=2, width=250)}\n' - f'[SparK.structure]: {super(SparK, self).__repr__().replace(SparK.__name__, "")}' + f"\n" + f"[SparK.config]: {pformat(self.get_config(), indent=2, width=250)}\n" + f"[SparK.structure]: {super(SparK, self).__repr__().replace(SparK.__name__, '')}" ) - + def get_config(self): return { - 'mask_ratio': self.mask_ratio, - 'densify_norm_str': self.densify_norm_str, - 'sbn': self.sbn, 'hierarchy': self.hierarchy, - 'sparse_encoder.input_size': self.sparse_encoder.input_size, - 'dense_decoder.width': self.dense_decoder.width, + "mask_ratio": self.mask_ratio, + "densify_norm_str": self.densify_norm_str, + "sbn": self.sbn, + "hierarchy": self.hierarchy, + "sparse_encoder.input_size": self.sparse_encoder.input_size, + "dense_decoder.width": self.dense_decoder.width, } - - def state_dict(self, destination=None, prefix='', keep_vars=False, with_config=False): - state = super(SparK, self).state_dict(destination=destination, prefix=prefix, keep_vars=keep_vars) + + def state_dict( + self, destination=None, prefix="", keep_vars=False, with_config=False + ): + state = super(SparK, self).state_dict( + destination=destination, prefix=prefix, keep_vars=keep_vars + ) if with_config: - state['config'] = self.get_config() + state["config"] = self.get_config() return state - + def load_state_dict(self, state_dict, strict=True): - config: dict = state_dict.pop('config', None) - incompatible_keys = super(SparK, self).load_state_dict(state_dict, strict=strict) + config: dict = state_dict.pop("config", None) + incompatible_keys = super(SparK, self).load_state_dict( + state_dict, strict=strict + ) if config is not None: for k, v in self.get_config().items(): ckpt_v = config.get(k, None) if ckpt_v != v: - err = f'[SparseMIM.load_state_dict] config mismatch: this.{k}={v} (ckpt.{k}={ckpt_v})' + err = f"[SparseMIM.load_state_dict] config mismatch: this.{k}={v} (ckpt.{k}={ckpt_v})" if strict: raise AttributeError(err) else: print(err, file=sys.stderr) - return incompatible_keys \ No newline at end of file + return incompatible_keys diff --git a/lightly/models/sparse/encoder.py b/lightly/models/sparse/encoder.py deleted file mode 100644 index bc4f6704a..000000000 --- a/lightly/models/sparse/encoder.py +++ /dev/null @@ -1,239 +0,0 @@ -# Code adapted from https://github.com/keyu-tian/SparK/blob/main/pretrain/models/resnet.py -from typing import Union - -import torch -import torch.nn as nn -from torch import List, Optional, Tensor - -from lightly.models.modules.sparse.spark import ( - SparseAvgPooling, - SparseBatchNorm2d, - SparseConv2d, - SparseMaxPooling, - SparseSyncBatchNorm2d, - SparseConvNeXtBlock, - SparseConvNeXtLayerNorm -) - -from torchvision.models.convnext import CNBlock, LayerNorm2d - -from torchvision.models import ResNet, ConvNeXt - - -class SparseMask: - """ - Basic class to store the mask for the sparse model. - """ - - def __init__(self): - self.mask: Union[Tensor, None] = None - - -class SparseEncoder(nn.Module): - def __init__( - self, - backbone: nn.Module, - input_size: int, - downsample_ratio: Optional[int] = None, - ): - """Sparse encoder as used by SparK [0] - - Default params are the ones explained in the original code base. The backbone is assumed - to follow the same API as the ResNet or ConvNext models from torchvision. - [0] Designing BERT for Convolutional Networks: Sparse and Hierarchical Masked Modeling https://arxiv.org/abs/2301.03580 - - Attributes: - backbone: - Backbone model to extract features from images. Should have both - the methods get_downsample_ratio() and get_feature_map_channels() - implemented. - input_size: - Size of the input image. - downsample_ratio: - - """ - super().__init__() - self.input_size = input_size - if downsample_ratio is None: - self.downsample_ratio = self.get_downsample_ratio(backbone) - else: - self.downsample_ratio = downsample_ratio - - self.enc_feat_map_chs = ( - self.get_feature_map_channels(backbone), - ) - self.sparse_mask = SparseMask() - self.sparse_backbone = self.dense_model_to_sparse( - m=backbone, sparse_mask=self.sparse_mask - ) - - def forward( - self, x: torch.Tensor, mask: torch.BoolTensor, hierarchical: bool = False - ) -> Union[torch.Tensor, List[torch.Tensor]]: - """ - Forward pass of the sparse encoder. - - Args: - x: The input tensor. - mask: The mask to apply to the input tensor. This should be a binary mask the - size of the smallest resolution feature map [0] of the backbone model. - - Returns: - The output tensor. - - [0] SparK codebase https://github.com/keyu-tian/SparK/tree/main/pretrain#some-details-how-we-mask-images-and-how-to-set-the-patch-size - """ - assert mask.shape == ( - 1, - 1, - self.input_size // self.downsample_ratio, - self.input_size // self.downsample_ratio, - ), "Mask shape must be (1, 1, H // downsample_ratio, W // downsample_ratio)" - self.sparse_mask.mask = mask - - ls = [] - - if isinstance(self.sparse_backbone, ConvNeXt): - if hierarchical: - for i in range(0,8,2): - x = self.sparse_backbone.features[i](x) - x = self.sparse_backbone.features[i+1](x) - ls.append(x) - else: - x = self.sparse_backbone.avgpool(x) - x = self.sparse_backbone.classifier(x) - return x - - elif isinstance(self.sparse_backbone, ResNet): - x = self.sparse_backbone.conv1(x) - x = self.sparse_backbone.bn1(x) - x = self.sparse_backbone.act1(x) - x = self.sparse_backbone.maxpool(x) - - if hierarchical: - ls = [] - x = self.sparse_backbone.layer1(x) - ls.append(x) - x = self.sparse_backbone.layer2(x) - ls.append(x) - x = self.sparse_backbone.layer3(x) - ls.append(x) - x = self.sparse_backbone.layer4(x) - ls.append(x) - return ls - else: - x = self.sparse_backbone.global_pool(x) - x = self.sparse_backbone.fc(x) - return x - - else: - raise NotImplementedError("Backbone not supported") - - def get_downsample_ratio(self, backbone: nn.Module) -> int: - """ - Try to get the downsample ratio of the backbone model. - - Returns: - The downsample ratio of the backbone model. - """ - try: - return backbone.get_downsample_ratio() - except AttributeError: - x = torch.randn(1, 3, self.input_size, self.input_size) - out = nn.Sequential(*list(backbone.children())[:-2])(x) - - return self.input_size // out.shape[-1] - - - def dense_model_to_sparse( - self, m: nn.Module, sparse_mask: SparseMask, sync_batch_norm: bool = False - ) -> nn.Module: - """ - Convert a dense model to a sparse model. - - Args: - m: The dense model to convert. - sparse_mask: The sparse mask to use for the conversion. - sync_batch_norm: Whether to convert BatchNorm2d to SyncBatchNorm. - - Returns: - The sparse model. - """ - with torch.no_grad(): - oup = m - if isinstance(m, nn.Conv2d): - m: nn.Conv2d - bias = m.bias is not None - oup = SparseConv2d( - m.in_channels, - m.out_channels, - kernel_size=m.kernel_size, - stride=m.stride, - padding=m.padding, - dilation=m.dilation, - groups=m.groups, - bias=bias, - padding_mode=m.padding_mode, - ) - oup.sparse_mask = sparse_mask - oup.weight.copy_(m.weight) - if bias: - oup.bias.copy_(m.bias) - elif isinstance(m, nn.MaxPool2d): - 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, - ) - oup.sparse_mask = sparse_mask - elif isinstance(m, nn.AvgPool2d): - 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, - ) - oup.sparse_mask = sparse_mask - elif isinstance(m, (nn.BatchNorm2d, nn.SyncBatchNorm)): - m: nn.BatchNorm2d - oup = ( - SparseSyncBatchNorm2d - if isinstance(m, nn.SyncBatchNorm) - else SparseBatchNorm2d - )( - m.weight.shape[0], - eps=m.eps, - momentum=m.momentum, - affine=m.affine, - track_running_stats=m.track_running_stats, - ) - oup.sparse_mask = sparse_mask - oup.weight.copy_(m.weight) - oup.bias.copy_(m.bias) - oup.running_mean.copy_(m.running_mean) - oup.running_var.copy_(m.running_var) - oup.num_batches_tracked.copy_(m.num_batches_tracked) - if hasattr(m, "qconfig"): - oup.qconfig = m.qconfig - elif isinstance(m, CNBlock): - m: CNBlock - oup = SparseConvNeXtBlock( - m.weight.shape[0], - m.layer_scale. - ) - elif isinstance(m, (nn.Conv1d,)): - raise NotImplementedError - - for name, child in m.named_children(): - oup.add_module( - name, self.dense_model_to_sparse(child, sparse_mask=sparse_mask) - ) - del m - return oup From a7138d0638f7f885aeb3e58d5fd5027f8664a99e Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 5 Feb 2026 19:17:44 -0300 Subject: [PATCH 06/85] feat: put everything into the sparse spark module. --- lightly/models/modules/sparse_spark.py | 345 ++++++++++++++++++++++++- 1 file changed, 340 insertions(+), 5 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index bf6c29f16..39ec02d3b 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -2,10 +2,6 @@ from typing import Literal, Optional -import torch -import torch.nn as nn -from typing import Union - import torch import torch.nn as nn from torch import List, Optional, Tensor @@ -483,4 +479,343 @@ def dense_model_to_sparse( name, self.dense_model_to_sparse(child, sparse_mask=sparse_mask) ) del m - return oup \ No newline at end of file + return oup + + +import math +import sys +from pprint import pformat +from typing import List, Literal, Tuple, Union + +import torch +import torch.nn as nn + +from lightly.models.modules.spark import ( + SparseBatchNorm2d, + SparseConvNeXtLayerNorm, + SparseSyncBatchNorm2d, +) +from lightly.models.sparse.encoder import SparseResnet + + +def is_pow2n(n: int) -> bool: + return n > 0 and (n & (n - 1)) == 0 + + +class UNetBlock(nn.Module): + def __init__(self, cin, cout, bn2d): + """ + a UNet block with 2x up sampling + """ + 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): + x = self.up_sample(x) + return self.conv(x) + + +class LightDecoder(nn.Module): + def __init__(self, up_sample_ratio, width=768, sync_batch_norm=True): + 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 sync_batch_norm 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[torch.Tensor]): + x = 0 + for i, d in enumerate(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) + + def extra_repr(self) -> str: + return f"width={self.width}" + + def initialize(self): + for m in self.modules(): + if isinstance(m, nn.Linear): + torch.nn.init.trunc_normal_(m.weight, std=0.02) + if m.bias is not None: + nn.init.constant_(m.bias, 0) + elif isinstance(m, nn.Conv2d): + torch.nn.init.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 SparK(nn.Module): + """ + The SparK model as used by SparK [0] + + Default params are the ones explained in the original code base. The backbone is assumed + to follow the same API as the ResNet models from torchvision or a ConvNext model. + [0] Designing BERT for Convolutional Networks: Sparse and Hierarchical Masked Modeling https://arxiv.org/abs/2301.03580 + + Attributes: + sparse_encoder: + Sparse encoder to extract features from images. Should have both + the methods get_downsample_ratio() and get_feature_map_channels() + implemented. + dense_decoder: + Dense decoder to reconstruct the image from the sparse features. + mask_ratio: + Ratio of the image to mask. Default: 0.6 + densify_norm: + Type of normalization to use for densification. Default: 'bn' + sbn: + Whether to use SyncBatchNorm. Default: False + """ + + def __init__( + self, + sparse_encoder: SparseResnet, + dense_decoder: LightDecoder, + mask_ratio: float = 0.6, + densify_norm: Literal["batch_norm", "layer_norm", "identity"] = "bn", + sbn: bool = False, + ): + super().__init__() + input_size, downsample_ratio = ( + sparse_encoder.input_size, + sparse_encoder.downsample_ratio, + ) + self.downsample_ratio = downsample_ratio + self.fmap_h, self.fmap_w = ( + input_size // downsample_ratio, + input_size // downsample_ratio, + ) + self.mask_ratio = mask_ratio + self.len_keep = round(self.fmap_h * self.fmap_w * (1 - mask_ratio)) + + self.sparse_encoder = sparse_encoder + self.dense_decoder = dense_decoder + + self.sbn = sbn + self.hierarchy = len(sparse_encoder.enc_feat_map_chs) + self.densify_norm_str = densify_norm.lower() + self.densify_norms = nn.ModuleList() + self.densify_projs = nn.ModuleList() + self.mask_tokens = nn.ParameterList() + + e_widths, d_width = ( + self.sparse_encoder.enc_feat_map_chs, + self.dense_decoder.width, + ) + e_widths: List[int] + for i in range(self.hierarchy): + e_width = e_widths.pop() + p = nn.Parameter(torch.zeros(1, e_width, 1, 1)) + torch.nn.init.trunc_normal_(p, mean=0, std=0.02, a=-0.02, b=0.02) + self.mask_tokens.append(p) + + if self.densify_norm_str == "batch_norm": + densify_norm = ( + SparseSyncBatchNorm2d if self.sbn else SparseBatchNorm2d + )(e_width) + elif self.densify_norm_str == "layer_norm": + densify_norm = SparseConvNeXtLayerNorm( + e_width, data_format="channels_first", sparse=True + ) + else: + densify_norm = nn.Identity() + self.densify_norms.append(densify_norm) + + if i == 0 and e_width == d_width: + densify_proj = nn.Identity() + print( + f"[SparK.__init__, densify {i + 1}/{self.hierarchy}]: use nn.Identity() as densify_proj" + ) + else: + kernel_size = 1 if i <= 0 else 3 + densify_proj = nn.Conv2d( + e_width, + d_width, + kernel_size=kernel_size, + stride=1, + padding=kernel_size // 2, + bias=True, + ) + print( + f"[SparK.__init__, densify {i + 1}/{self.hierarchy}]: densify_proj(ksz={kernel_size}, #para={sum(x.numel() for x in densify_proj.parameters()) / 1e6:.2f}M)" + ) + self.densify_projs.append(densify_proj) + d_width //= 2 + + print( + f"[SparK.__init__] dims of mask_tokens={tuple(p.numel() for p in self.mask_tokens)}" + ) + + def mask(self, B: int, device, generator=None): + """ + Generate a mask for the input tensor + + Attributes: + B: + Batch size + device: + Device to put the mask on + generator: + Random number generator + """ + h, w = self.fmap_h, self.fmap_w + idx = torch.rand(B, h * w, generator=generator).argsort(dim=1) + idx = idx[:, : self.len_keep].to(device) # (B, len_keep) + return ( + torch.zeros(B, h * w, dtype=torch.bool, device=device) + .scatter_(dim=1, index=idx, value=True) + .view(B, 1, h, w) + ) + + def forward( + self, + inp_bchw: torch.Tensor, + active_b1ff: torch.BoolTensor = None, + return_loss: bool = True, + ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: + """ + Forward pass of the SparK model + + Attributes: + inp_bchw: + Input tensor + active_b1ff: + Active mask + return_loss: + Whether to return the loss + """ + if active_b1ff is None: + active_b1ff: torch.BoolTensor = self.mask( + inp_bchw.shape[0], inp_bchw.device + ) + active_b1hw = active_b1ff.repeat_interleave( + self.downsample_ratio, 2 + ).repeat_interleave(self.downsample_ratio, 3) + masked_bchw = inp_bchw * active_b1hw + + fea_bcffs: List[torch.Tensor] = self.sparse_encoder.forward( + masked_bchw, active_b1ff, hierarchical=True + ) + fea_bcffs.reverse() + + cur_active = active_b1ff + to_dec = [] + for i, bcff in enumerate(fea_bcffs): + if bcff is not None: + bcff = self.densify_norms[i](bcff) + mask_tokens = self.mask_tokens[i].expand_as(bcff) + bcff = torch.where(cur_active.expand_as(bcff), bcff, mask_tokens) + bcff: torch.Tensor = self.densify_projs[i](bcff) + to_dec.append(bcff) + cur_active = cur_active.repeat_interleave(2, dim=2).repeat_interleave( + 2, dim=3 + ) + + rec_bchw = self.dense_decoder(to_dec) + inp, rec = self.patchify(inp_bchw), self.patchify(rec_bchw) + mean = inp.mean(dim=-1, keepdim=True) + var = (inp.var(dim=-1, keepdim=True) + 1e-6) ** 0.5 + inp = (inp - mean) / var + l2_loss = ((rec - inp) ** 2).mean(dim=2, keepdim=False) + + non_active = active_b1ff.logical_not().int().view(active_b1ff.shape[0], -1) + recon_loss = l2_loss.mul_(non_active).sum() / (non_active.sum() + 1e-8) + + if return_loss: + return recon_loss + + masked_bchw = inp_bchw * active_b1hw + rec_bchw = self.unpatchify(rec * var + mean) + rec_or_inp = torch.where(active_b1hw, inp_bchw, rec_bchw) + return inp_bchw, masked_bchw, rec_or_inp + + def patchify(self, bchw): + p = self.downsample_ratio + h, w = self.fmap_h, self.fmap_w + B, C = bchw.shape[:2] + bchw = bchw.reshape(shape=(B, C, h, p, w, p)) + bchw = torch.einsum("bchpwq->bhwpqc", bchw) + bln = bchw.reshape(shape=(B, h * w, C * p**2)) + return bln + + def unpatchify(self, bln): + 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 __repr__(self): + return ( + f"\n" + f"[SparK.config]: {pformat(self.get_config(), indent=2, width=250)}\n" + f"[SparK.structure]: {super(SparK, self).__repr__().replace(SparK.__name__, '')}" + ) + + def get_config(self): + return { + "mask_ratio": self.mask_ratio, + "densify_norm_str": self.densify_norm_str, + "sbn": self.sbn, + "hierarchy": self.hierarchy, + "sparse_encoder.input_size": self.sparse_encoder.input_size, + "dense_decoder.width": self.dense_decoder.width, + } + + def state_dict( + self, destination=None, prefix="", keep_vars=False, with_config=False + ): + state = super(SparK, self).state_dict( + destination=destination, prefix=prefix, keep_vars=keep_vars + ) + if with_config: + state["config"] = self.get_config() + return state + + def load_state_dict(self, state_dict, strict=True): + config: dict = state_dict.pop("config", None) + incompatible_keys = super(SparK, self).load_state_dict( + state_dict, strict=strict + ) + if config is not None: + for k, v in self.get_config().items(): + ckpt_v = config.get(k, None) + if ckpt_v != v: + err = f"[SparseMIM.load_state_dict] config mismatch: this.{k}={v} (ckpt.{k}={ckpt_v})" + if strict: + raise AttributeError(err) + else: + print(err, file=sys.stderr) + return incompatible_keys \ No newline at end of file From 9b869e91dd7d4706d41dfb4a717c2293b6b4f594 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 5 Feb 2026 19:19:36 -0300 Subject: [PATCH 07/85] refactor: removed redundant super calls with class. --- lightly/models/modules/sparse_spark.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 39ec02d3b..b9b283c84 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -151,13 +151,13 @@ def forward(self, x): H=x.shape[1], W=x.shape[2], returning_active_ex=False ) nc = x[ii] - nc = super(SparseConvNeXtLayerNorm, self).forward(nc) + nc = super().forward(nc) x = torch.zeros_like(x) x[ii] = nc return x else: - return super(SparseConvNeXtLayerNorm, self).forward(x) + return super().forward(x) else: if self.sparse: ii = _get_active_ex_or_ii( @@ -165,7 +165,7 @@ def forward(self, x): ) bhwc = x.permute(0, 2, 3, 1) nc = bhwc[ii] - nc = super(SparseConvNeXtLayerNorm, self).forward(nc) + nc = super().forward(nc) x = torch.zeros_like(bhwc) x[ii] = nc @@ -180,11 +180,11 @@ def forward(self, x): if self.sparse: raise NotImplementedError else: - return super(SparseConvNeXtLayerNorm, self).forward(x) + return super().forward(x) def __repr__(self): return ( - super(SparseConvNeXtLayerNorm, self).__repr__()[:-1] + super().__repr__()[:-1] + f", ch={self.data_format.split('_')[-1]}, sp={self.sparse})" ) @@ -258,7 +258,7 @@ def forward(self, x): return x def __repr__(self): - return super(SparseConvNeXtBlock, self).__repr__()[:-1] + f", sp={self.sparse})" + return super().__repr__()[:-1] + f", sp={self.sparse})" # Code adapted from https://github.com/keyu-tian/SparK/blob/main/pretrain/models/resnet.py @@ -781,7 +781,7 @@ def __repr__(self): return ( f"\n" f"[SparK.config]: {pformat(self.get_config(), indent=2, width=250)}\n" - f"[SparK.structure]: {super(SparK, self).__repr__().replace(SparK.__name__, '')}" + f"[SparK.structure]: {super().__repr__().replace(SparK.__name__, '')}" ) def get_config(self): From 36cb94c9cab014ea0f4b92fa05daa708d60f0130 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 5 Feb 2026 19:24:27 -0300 Subject: [PATCH 08/85] refactor: removed spark code. --- lightly/models/spark.py | 339 ---------------------------------------- 1 file changed, 339 deletions(-) diff --git a/lightly/models/spark.py b/lightly/models/spark.py index 4c1b4ca30..e69de29bb 100644 --- a/lightly/models/spark.py +++ b/lightly/models/spark.py @@ -1,339 +0,0 @@ -# Code adapted from https://github.com/keyu-tian/SparK/blob/main/ - -import math -import sys -from pprint import pformat -from typing import List, Literal, Tuple, Union - -import torch -import torch.nn as nn - -from lightly.models.modules.spark import ( - SparseBatchNorm2d, - SparseConvNeXtLayerNorm, - SparseSyncBatchNorm2d, -) -from lightly.models.sparse.encoder import SparseResnet - - -def is_pow2n(n: int) -> bool: - return n > 0 and (n & (n - 1)) == 0 - - -class UNetBlock(nn.Module): - def __init__(self, cin, cout, bn2d): - """ - a UNet block with 2x up sampling - """ - 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): - x = self.up_sample(x) - return self.conv(x) - - -class LightDecoder(nn.Module): - def __init__(self, up_sample_ratio, width=768, sync_batch_norm=True): - 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 sync_batch_norm 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[torch.Tensor]): - x = 0 - for i, d in enumerate(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) - - def extra_repr(self) -> str: - return f"width={self.width}" - - def initialize(self): - for m in self.modules(): - if isinstance(m, nn.Linear): - torch.nn.init.trunc_normal_(m.weight, std=0.02) - if m.bias is not None: - nn.init.constant_(m.bias, 0) - elif isinstance(m, nn.Conv2d): - torch.nn.init.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 SparK(nn.Module): - """ - The SparK model as used by SparK [0] - - Default params are the ones explained in the original code base. The backbone is assumed - to follow the same API as the ResNet models from torchvision or a ConvNext model. - [0] Designing BERT for Convolutional Networks: Sparse and Hierarchical Masked Modeling https://arxiv.org/abs/2301.03580 - - Attributes: - sparse_encoder: - Sparse encoder to extract features from images. Should have both - the methods get_downsample_ratio() and get_feature_map_channels() - implemented. - dense_decoder: - Dense decoder to reconstruct the image from the sparse features. - mask_ratio: - Ratio of the image to mask. Default: 0.6 - densify_norm: - Type of normalization to use for densification. Default: 'bn' - sbn: - Whether to use SyncBatchNorm. Default: False - """ - - def __init__( - self, - sparse_encoder: SparseResnet, - dense_decoder: LightDecoder, - mask_ratio: float = 0.6, - densify_norm: Literal["batch_norm", "layer_norm", "identity"] = "bn", - sbn: bool = False, - ): - super().__init__() - input_size, downsample_ratio = ( - sparse_encoder.input_size, - sparse_encoder.downsample_ratio, - ) - self.downsample_ratio = downsample_ratio - self.fmap_h, self.fmap_w = ( - input_size // downsample_ratio, - input_size // downsample_ratio, - ) - self.mask_ratio = mask_ratio - self.len_keep = round(self.fmap_h * self.fmap_w * (1 - mask_ratio)) - - self.sparse_encoder = sparse_encoder - self.dense_decoder = dense_decoder - - self.sbn = sbn - self.hierarchy = len(sparse_encoder.enc_feat_map_chs) - self.densify_norm_str = densify_norm.lower() - self.densify_norms = nn.ModuleList() - self.densify_projs = nn.ModuleList() - self.mask_tokens = nn.ParameterList() - - e_widths, d_width = ( - self.sparse_encoder.enc_feat_map_chs, - self.dense_decoder.width, - ) - e_widths: List[int] - for i in range(self.hierarchy): - e_width = e_widths.pop() - p = nn.Parameter(torch.zeros(1, e_width, 1, 1)) - torch.nn.init.trunc_normal_(p, mean=0, std=0.02, a=-0.02, b=0.02) - self.mask_tokens.append(p) - - if self.densify_norm_str == "batch_norm": - densify_norm = ( - SparseSyncBatchNorm2d if self.sbn else SparseBatchNorm2d - )(e_width) - elif self.densify_norm_str == "layer_norm": - densify_norm = SparseConvNeXtLayerNorm( - e_width, data_format="channels_first", sparse=True - ) - else: - densify_norm = nn.Identity() - self.densify_norms.append(densify_norm) - - if i == 0 and e_width == d_width: - densify_proj = nn.Identity() - print( - f"[SparK.__init__, densify {i + 1}/{self.hierarchy}]: use nn.Identity() as densify_proj" - ) - else: - kernel_size = 1 if i <= 0 else 3 - densify_proj = nn.Conv2d( - e_width, - d_width, - kernel_size=kernel_size, - stride=1, - padding=kernel_size // 2, - bias=True, - ) - print( - f"[SparK.__init__, densify {i + 1}/{self.hierarchy}]: densify_proj(ksz={kernel_size}, #para={sum(x.numel() for x in densify_proj.parameters()) / 1e6:.2f}M)" - ) - self.densify_projs.append(densify_proj) - d_width //= 2 - - print( - f"[SparK.__init__] dims of mask_tokens={tuple(p.numel() for p in self.mask_tokens)}" - ) - - def mask(self, B: int, device, generator=None): - """ - Generate a mask for the input tensor - - Attributes: - B: - Batch size - device: - Device to put the mask on - generator: - Random number generator - """ - h, w = self.fmap_h, self.fmap_w - idx = torch.rand(B, h * w, generator=generator).argsort(dim=1) - idx = idx[:, : self.len_keep].to(device) # (B, len_keep) - return ( - torch.zeros(B, h * w, dtype=torch.bool, device=device) - .scatter_(dim=1, index=idx, value=True) - .view(B, 1, h, w) - ) - - def forward( - self, - inp_bchw: torch.Tensor, - active_b1ff: torch.BoolTensor = None, - return_loss: bool = True, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """ - Forward pass of the SparK model - - Attributes: - inp_bchw: - Input tensor - active_b1ff: - Active mask - return_loss: - Whether to return the loss - """ - if active_b1ff is None: - active_b1ff: torch.BoolTensor = self.mask( - inp_bchw.shape[0], inp_bchw.device - ) - active_b1hw = active_b1ff.repeat_interleave( - self.downsample_ratio, 2 - ).repeat_interleave(self.downsample_ratio, 3) - masked_bchw = inp_bchw * active_b1hw - - fea_bcffs: List[torch.Tensor] = self.sparse_encoder.forward( - masked_bchw, active_b1ff, hierarchical=True - ) - fea_bcffs.reverse() - - cur_active = active_b1ff - to_dec = [] - for i, bcff in enumerate(fea_bcffs): - if bcff is not None: - bcff = self.densify_norms[i](bcff) - mask_tokens = self.mask_tokens[i].expand_as(bcff) - bcff = torch.where(cur_active.expand_as(bcff), bcff, mask_tokens) - bcff: torch.Tensor = self.densify_projs[i](bcff) - to_dec.append(bcff) - cur_active = cur_active.repeat_interleave(2, dim=2).repeat_interleave( - 2, dim=3 - ) - - rec_bchw = self.dense_decoder(to_dec) - inp, rec = self.patchify(inp_bchw), self.patchify(rec_bchw) - mean = inp.mean(dim=-1, keepdim=True) - var = (inp.var(dim=-1, keepdim=True) + 1e-6) ** 0.5 - inp = (inp - mean) / var - l2_loss = ((rec - inp) ** 2).mean(dim=2, keepdim=False) - - non_active = active_b1ff.logical_not().int().view(active_b1ff.shape[0], -1) - recon_loss = l2_loss.mul_(non_active).sum() / (non_active.sum() + 1e-8) - - if return_loss: - return recon_loss - - masked_bchw = inp_bchw * active_b1hw - rec_bchw = self.unpatchify(rec * var + mean) - rec_or_inp = torch.where(active_b1hw, inp_bchw, rec_bchw) - return inp_bchw, masked_bchw, rec_or_inp - - def patchify(self, bchw): - p = self.downsample_ratio - h, w = self.fmap_h, self.fmap_w - B, C = bchw.shape[:2] - bchw = bchw.reshape(shape=(B, C, h, p, w, p)) - bchw = torch.einsum("bchpwq->bhwpqc", bchw) - bln = bchw.reshape(shape=(B, h * w, C * p**2)) - return bln - - def unpatchify(self, bln): - 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 __repr__(self): - return ( - f"\n" - f"[SparK.config]: {pformat(self.get_config(), indent=2, width=250)}\n" - f"[SparK.structure]: {super(SparK, self).__repr__().replace(SparK.__name__, '')}" - ) - - def get_config(self): - return { - "mask_ratio": self.mask_ratio, - "densify_norm_str": self.densify_norm_str, - "sbn": self.sbn, - "hierarchy": self.hierarchy, - "sparse_encoder.input_size": self.sparse_encoder.input_size, - "dense_decoder.width": self.dense_decoder.width, - } - - def state_dict( - self, destination=None, prefix="", keep_vars=False, with_config=False - ): - state = super(SparK, self).state_dict( - destination=destination, prefix=prefix, keep_vars=keep_vars - ) - if with_config: - state["config"] = self.get_config() - return state - - def load_state_dict(self, state_dict, strict=True): - config: dict = state_dict.pop("config", None) - incompatible_keys = super(SparK, self).load_state_dict( - state_dict, strict=strict - ) - if config is not None: - for k, v in self.get_config().items(): - ckpt_v = config.get(k, None) - if ckpt_v != v: - err = f"[SparseMIM.load_state_dict] config mismatch: this.{k}={v} (ckpt.{k}={ckpt_v})" - if strict: - raise AttributeError(err) - else: - print(err, file=sys.stderr) - return incompatible_keys From 0afc280f603f7f6573e577199b4347f0165bff4b Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 5 Feb 2026 19:25:15 -0300 Subject: [PATCH 09/85] refactor: removing redundant super calls --- lightly/models/modules/sparse_spark.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index b9b283c84..c0b1ad6ad 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -797,7 +797,7 @@ def get_config(self): def state_dict( self, destination=None, prefix="", keep_vars=False, with_config=False ): - state = super(SparK, self).state_dict( + state = super().state_dict( destination=destination, prefix=prefix, keep_vars=keep_vars ) if with_config: @@ -806,7 +806,7 @@ def state_dict( def load_state_dict(self, state_dict, strict=True): config: dict = state_dict.pop("config", None) - incompatible_keys = super(SparK, self).load_state_dict( + incompatible_keys = super().load_state_dict( state_dict, strict=strict ) if config is not None: From 280a479d268b42897eb837e71c2a171206d38b10 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 6 Feb 2026 08:34:02 -0300 Subject: [PATCH 10/85] refactor: remove empty file --- lightly/models/spark.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) delete mode 100644 lightly/models/spark.py diff --git a/lightly/models/spark.py b/lightly/models/spark.py deleted file mode 100644 index e69de29bb..000000000 From 7e0f2bca7db7a9bd904a610435e33a816ba6ec46 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 6 Feb 2026 10:38:54 -0300 Subject: [PATCH 11/85] refactor: porting original code. starting from scratch --- lightly/models/modules/sparse_spark.py | 687 +++++++++---------------- 1 file changed, 244 insertions(+), 443 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index c0b1ad6ad..58cb005c3 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -1,88 +1,61 @@ -# Code adapted from https://github.com/keyu-tian/SparK/blob/main/pretrain/encoder.py and https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/drop.py - -from typing import Literal, Optional +import math +import sys +from pprint import pformat +from typing import List import torch import torch.nn as nn -from torch import List, Optional, Tensor -from torchvision.models.convnext import CNBlock, LayerNorm2d +from timm.models.layers import DropPath, trunc_normal_ +from torch.nn.common_types import _size_2_t -from torchvision.models import ResNet, ConvNeXt +def is_pow2n(x): + return x > 0 and (x & (x - 1) == 0) -def drop_path(x, drop_prob: float = 0.0, training: bool = False): - """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). - This is the same as the DropConnect impl I created for EfficientNet, etc networks, however, - the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... - See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for - changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use - 'survival rate' as the argument. +def coalesce_to_size_2_t(t: tuple[int, ...]) -> _size_2_t: + if len(t) == 2: + return t + elif len(t) == 1: + return t[0], t[0] + else: + raise ValueError(f"Invalid tuple length: {len(t)}; expected 1 or 2.") - Attributes: - x: - The input tensor. - drop_prob: - The drop probability of the path. Default: 0. - training: - Whether the model is in training mode. Default: False - """ - if drop_prob == 0.0 or not training: - return x - keep_prob = 1 - drop_prob - shape = (x.shape[0],) + (1,) * (x.ndim - 1) - random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) - random_tensor.floor_() - output = x.div(keep_prob) * random_tensor - return output +_cur_active: torch.Tensor = None # B1ff -class DropPath(nn.Module): - """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). - Attributes: - drop_prob: - The drop probability of the path. Default: None. - """ - - def __init__(self, drop_prob: Optional[float] = None): - super(DropPath, self).__init__() - self.drop_prob = drop_prob - - def forward(self, x): - return drop_path(x, self.drop_prob, self.training) - - -def _get_active_ex_or_ii( - mask: torch.BoolTensor, H: int, W: int, returning_active_ex: bool = True -): - h_repeat, w_repeat = H // mask.shape[-2], W // mask.shape[-1] - active_ex = mask.repeat_interleave(h_repeat, dim=2).repeat_interleave( +# todo: try to use `gather` for speed? +def _get_active_ex_or_ii(H: int, W: int, returning_active_ex=True): + h_repeat, w_repeat = H // _cur_active.shape[-2], W // _cur_active.shape[-1] + active_ex = _cur_active.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) - ) + ) # ii: bi, hi, wi -def sp_conv_forward(self, x: torch.Tensor) -> torch.Tensor: +def sp_conv_forward(self, x: torch.Tensor): x = super(type(self), self).forward(x) x *= _get_active_ex_or_ii( - self.sparse_mask.mask, H=x.shape[2], W=x.shape[3], returning_active_ex=True - ) + H=x.shape[2], W=x.shape[3], returning_active_ex=True + ) # (BCHW) *= (B1HW), mask the output of conv return x def sp_bn_forward(self, x: torch.Tensor): - ii = _get_active_ex_or_ii( - self.sparse_mask.mask, H=x.shape[2], W=x.shape[3], returning_active_ex=False - ) + ii = _get_active_ex_or_ii(H=x.shape[2], W=x.shape[3], returning_active_ex=False) bhwc = x.permute(0, 2, 3, 1) - nc = bhwc[ii] - nc = super(type(self), self).forward(nc) + nc = bhwc[ + ii + ] # select the features on non-masked positions to form a flatten feature `nc` + nc = super(type(self), self).forward( + nc + ) # use BN1d to normalize this flatten feature `nc` bchw = torch.zeros_like(bhwc) bchw[ii] = nc @@ -91,23 +64,23 @@ def sp_bn_forward(self, x: torch.Tensor): class SparseConv2d(nn.Conv2d): - forward = sp_conv_forward + forward = sp_conv_forward # hack: override the forward function; see `sp_conv_forward` above for more details class SparseMaxPooling(nn.MaxPool2d): - forward = sp_conv_forward + forward = sp_conv_forward # hack: override the forward function; see `sp_conv_forward` above for more details class SparseAvgPooling(nn.AvgPool2d): - forward = sp_conv_forward + forward = sp_conv_forward # hack: override the forward function; see `sp_conv_forward` above for more details class SparseBatchNorm2d(nn.BatchNorm1d): - forward = sp_bn_forward + forward = sp_bn_forward # hack: override the forward function; see `sp_bn_forward` above for more details class SparseSyncBatchNorm2d(nn.SyncBatchNorm): - forward = sp_bn_forward + forward = sp_bn_forward # hack: override the forward function; see `sp_bn_forward` above for more details class SparseConvNeXtLayerNorm(nn.LayerNorm): @@ -115,27 +88,10 @@ class SparseConvNeXtLayerNorm(nn.LayerNorm): 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). - - - Attributes: - normalized_shape: - Input shape from an expected input of size - normalized_shape or a single integer. - eps: - A value added to the denominator for numerical stability. Default: 1e-6. - data_format: - The ordering of the dimensions in the inputs. Default: "channels_last". - sparse: - Whether to use sparse computation. Default: True. - """ def __init__( - self, - normalized_shape: int, - eps: float = 1e-6, - data_format: Literal["channels_last", "channels_first"] = "channels_last", - sparse: bool = True, + self, normalized_shape, eps=1e-6, data_format="channels_last", sparse=True ): if data_format not in ["channels_last", "channels_first"]: raise NotImplementedError @@ -144,28 +100,28 @@ def __init__( self.sparse = sparse def forward(self, x): - if x.ndim == 4: - if self.data_format == "channels_last": + if x.ndim == 4: # BHWC or BCHW + if self.data_format == "channels_last": # BHWC if self.sparse: ii = _get_active_ex_or_ii( H=x.shape[1], W=x.shape[2], returning_active_ex=False ) nc = x[ii] - nc = super().forward(nc) + nc = super(SparseConvNeXtLayerNorm, self).forward(nc) x = torch.zeros_like(x) x[ii] = nc return x else: - return super().forward(x) - else: + return super(SparseConvNeXtLayerNorm, self).forward(x) + else: # channels_first, BCHW if self.sparse: ii = _get_active_ex_or_ii( H=x.shape[2], W=x.shape[3], returning_active_ex=False ) bhwc = x.permute(0, 2, 3, 1) nc = bhwc[ii] - nc = super().forward(nc) + nc = super(SparseConvNeXtLayerNorm, self).forward(nc) x = torch.zeros_like(bhwc) x[ii] = nc @@ -180,49 +136,38 @@ def forward(self, x): if self.sparse: raise NotImplementedError else: - return super().forward(x) + return super(SparseConvNeXtLayerNorm, self).forward(x) def __repr__(self): return ( - super().__repr__()[:-1] + super(SparseConvNeXtLayerNorm, self).__repr__()[:-1] + f", ch={self.data_format.split('_')[-1]}, sp={self.sparse})" ) class SparseConvNeXtBlock(nn.Module): - r""" - ConvNeXt Block. There are two equivalent implementations: + r"""ConvNeXt Block. 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 - Attributes: - 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. - sparse: - Whether to use sparse computation. Default: True. - kernel_size: - Kernel size of depthwise convolution. Default: 7. + Args: + dim (int): Number of input channels. + drop_path (float): Stochastic depth rate. Default: 0.0 + layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. """ def __init__( - self, - dim, - drop_path: float = 0.0, - layer_scale_init_value: float = 1e-6, - sparse: bool = True, - kernel_size=7, + self, dim, drop_path=0.0, layer_scale_init_value=1e-6, sparse=True, ks=7 ): super().__init__() self.dwconv = nn.Conv2d( - dim, dim, kernel_size=kernel_size, padding=kernel_size // 2, groups=dim - ) + dim, dim, kernel_size=ks, padding=ks // 2, groups=dim + ) # depthwise conv self.norm = SparseConvNeXtLayerNorm(dim, eps=1e-6, sparse=sparse) - self.pwconv1 = nn.Linear(dim, 4 * dim) + self.pwconv1 = nn.Linear( + dim, 4 * dim + ) # pointwise/1x1 convs, implemented with linear layers self.act = nn.GELU() self.pwconv2 = nn.Linear(4 * dim, dim) self.gamma = ( @@ -258,248 +203,106 @@ def forward(self, x): return x def __repr__(self): - return super().__repr__()[:-1] + f", sp={self.sparse})" - -# Code adapted from https://github.com/keyu-tian/SparK/blob/main/pretrain/models/resnet.py - - -class SparseMask: - """ - Basic class to store the mask for the sparse model. - """ - - def __init__(self): - self.mask: Union[Tensor, None] = None + return super(SparseConvNeXtBlock, self).__repr__()[:-1] + f", sp={self.sparse})" class SparseEncoder(nn.Module): - def __init__( - self, - backbone: nn.Module, - input_size: int, - downsample_ratio: Optional[int] = None, - ): - """Sparse encoder as used by SparK [0] - - Default params are the ones explained in the original code base. The backbone is assumed - to follow the same API as the ResNet or ConvNext models from torchvision. - [0] Designing BERT for Convolutional Networks: Sparse and Hierarchical Masked Modeling https://arxiv.org/abs/2301.03580 - - Attributes: - backbone: - Backbone model to extract features from images. Should have both - the methods get_downsample_ratio() and get_feature_map_channels() - implemented. - input_size: - Size of the input image. - downsample_ratio: - - """ - super().__init__() - self.input_size = input_size - if downsample_ratio is None: - self.downsample_ratio = self.get_downsample_ratio(backbone) - else: - self.downsample_ratio = downsample_ratio - - self.enc_feat_map_chs = ( - self.get_feature_map_channels(backbone), + def __init__(self, cnn, input_size, sbn=False, verbose=False): + super(SparseEncoder, self).__init__() + self.sp_cnn = SparseEncoder.dense_model_to_sparse( + m=cnn, verbose=verbose, sbn=sbn ) - self.sparse_mask = SparseMask() - self.sparse_backbone = self.dense_model_to_sparse( - m=backbone, sparse_mask=self.sparse_mask + self.input_size, self.downsample_raito, self.enc_feat_map_chs = ( + input_size, + cnn.get_downsample_ratio(), + cnn.get_feature_map_channels(), ) - def forward( - self, x: torch.Tensor, mask: torch.BoolTensor, hierarchical: bool = False - ) -> Union[torch.Tensor, List[torch.Tensor]]: - """ - Forward pass of the sparse encoder. - - Args: - x: The input tensor. - mask: The mask to apply to the input tensor. This should be a binary mask the - size of the smallest resolution feature map [0] of the backbone model. - - Returns: - The output tensor. - - [0] SparK codebase https://github.com/keyu-tian/SparK/tree/main/pretrain#some-details-how-we-mask-images-and-how-to-set-the-patch-size - """ - assert mask.shape == ( - 1, - 1, - self.input_size // self.downsample_ratio, - self.input_size // self.downsample_ratio, - ), "Mask shape must be (1, 1, H // downsample_ratio, W // downsample_ratio)" - self.sparse_mask.mask = mask - - ls = [] - - if isinstance(self.sparse_backbone, ConvNeXt): - if hierarchical: - for i in range(0,8,2): - x = self.sparse_backbone.features[i](x) - x = self.sparse_backbone.features[i+1](x) - ls.append(x) - else: - x = self.sparse_backbone.avgpool(x) - x = self.sparse_backbone.classifier(x) - return x - - elif isinstance(self.sparse_backbone, ResNet): - x = self.sparse_backbone.conv1(x) - x = self.sparse_backbone.bn1(x) - x = self.sparse_backbone.act1(x) - x = self.sparse_backbone.maxpool(x) - - if hierarchical: - ls = [] - x = self.sparse_backbone.layer1(x) - ls.append(x) - x = self.sparse_backbone.layer2(x) - ls.append(x) - x = self.sparse_backbone.layer3(x) - ls.append(x) - x = self.sparse_backbone.layer4(x) - ls.append(x) - return ls - else: - x = self.sparse_backbone.global_pool(x) - x = self.sparse_backbone.fc(x) - return x - - else: - raise NotImplementedError("Backbone not supported") - - def get_downsample_ratio(self, backbone: nn.Module) -> int: - """ - Try to get the downsample ratio of the backbone model. - - Returns: - The downsample ratio of the backbone model. - """ - try: - return backbone.get_downsample_ratio() - except AttributeError: - x = torch.randn(1, 3, self.input_size, self.input_size) - out = nn.Sequential(*list(backbone.children())[:-2])(x) - - return self.input_size // out.shape[-1] - - - def dense_model_to_sparse( - self, m: nn.Module, sparse_mask: SparseMask, sync_batch_norm: bool = False - ) -> nn.Module: - """ - Convert a dense model to a sparse model. - - Args: - m: The dense model to convert. - sparse_mask: The sparse mask to use for the conversion. - sync_batch_norm: Whether to convert BatchNorm2d to SyncBatchNorm. - - Returns: - The sparse model. - """ - with torch.no_grad(): - oup = m - if isinstance(m, nn.Conv2d): - m: nn.Conv2d - bias = m.bias is not None - oup = SparseConv2d( - m.in_channels, - m.out_channels, - kernel_size=m.kernel_size, - stride=m.stride, - padding=m.padding, - dilation=m.dilation, - groups=m.groups, - bias=bias, - padding_mode=m.padding_mode, - ) - oup.sparse_mask = sparse_mask - oup.weight.copy_(m.weight) - if bias: - oup.bias.copy_(m.bias) - elif isinstance(m, nn.MaxPool2d): - 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, - ) - oup.sparse_mask = sparse_mask - elif isinstance(m, nn.AvgPool2d): - 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, - ) - oup.sparse_mask = sparse_mask - elif isinstance(m, (nn.BatchNorm2d, nn.SyncBatchNorm)): - m: nn.BatchNorm2d - oup = ( - SparseSyncBatchNorm2d - if isinstance(m, nn.SyncBatchNorm) - else SparseBatchNorm2d - )( - m.weight.shape[0], - eps=m.eps, - momentum=m.momentum, - affine=m.affine, - track_running_stats=m.track_running_stats, - ) - oup.sparse_mask = sparse_mask - oup.weight.copy_(m.weight) - oup.bias.copy_(m.bias) - oup.running_mean.copy_(m.running_mean) - oup.running_var.copy_(m.running_var) - oup.num_batches_tracked.copy_(m.num_batches_tracked) - if hasattr(m, "qconfig"): - oup.qconfig = m.qconfig - elif isinstance(m, CNBlock): - m: CNBlock - oup = SparseConvNeXtBlock( - m.weight.shape[0], - m.layer_scale. - ) - elif isinstance(m, (nn.Conv1d,)): - raise NotImplementedError - - for name, child in m.named_children(): - oup.add_module( - name, self.dense_model_to_sparse(child, sparse_mask=sparse_mask) - ) - del m - return oup - - -import math -import sys -from pprint import pformat -from typing import List, Literal, Tuple, Union + @staticmethod + def dense_model_to_sparse(m: nn.Module, verbose=False, sbn=False): + oup = m + if isinstance(m, nn.Conv2d): + bias = m.bias is not None + + oup = 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, + ) + oup.weight.data.copy_(m.weight.data) + + if m.bias is not None and oup.bias is not None: + oup.bias.data.copy_(m.bias.data) + + elif isinstance(m, nn.MaxPool2d): + 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): + 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)): + m: nn.BatchNorm2d + oup = (SparseSyncBatchNorm2d if sbn else SparseBatchNorm2d)( + m.weight.shape[0], + eps=m.eps, + momentum=m.momentum, + affine=m.affine, + track_running_stats=m.track_running_stats, + ) + oup.weight.data.copy_(m.weight.data) + oup.bias.data.copy_(m.bias.data) + oup.running_mean.data.copy_(m.running_mean.data) + oup.running_var.data.copy_(m.running_var.data) + oup.num_batches_tracked.data.copy_(m.num_batches_tracked.data) + if hasattr(m, "qconfig"): + oup.qconfig = m.qconfig + elif isinstance(m, nn.LayerNorm) and not isinstance(m, SparseConvNeXtLayerNorm): + m: nn.LayerNorm + oup = SparseConvNeXtLayerNorm(m.weight.shape[0], eps=m.eps) + oup.weight.data.copy_(m.weight.data) + oup.bias.data.copy_(m.bias.data) + elif isinstance(m, (nn.Conv1d,)): + raise NotImplementedError -import torch -import torch.nn as nn + for name, child in m.named_children(): + oup.add_module( + name, + SparseEncoder.dense_model_to_sparse(child, verbose=verbose, sbn=sbn), + ) + del m + return oup -from lightly.models.modules.spark import ( - SparseBatchNorm2d, - SparseConvNeXtLayerNorm, - SparseSyncBatchNorm2d, -) -from lightly.models.sparse.encoder import SparseResnet + def forward(self, x): + return self.sp_cnn(x, hierarchical=True) -def is_pow2n(n: int) -> bool: - return n > 0 and (n & (n - 1)) == 0 +# Copyright (c) ByteDance, Inc. and its affiliates. +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. class UNetBlock(nn.Module): @@ -525,13 +328,17 @@ def forward(self, x): class LightDecoder(nn.Module): - def __init__(self, up_sample_ratio, width=768, sync_batch_norm=True): + def __init__( + self, up_sample_ratio, width=768, sbn=True + ): # todo: the decoder's width follows a simple halfing rule; you can change it to any other rule 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 sync_batch_norm else nn.BatchNorm2d + channels = [ + self.width // 2**i for i in range(n + 1) + ] # todo: the decoder's width follows a simple halfing rule; you can change it to any other rule + bn2d = nn.SyncBatchNorm if sbn else nn.BatchNorm2d self.dec = nn.ModuleList( [ UNetBlock(cin, cout, bn2d) @@ -556,11 +363,11 @@ def extra_repr(self) -> str: def initialize(self): for m in self.modules(): if isinstance(m, nn.Linear): - torch.nn.init.trunc_normal_(m.weight, std=0.02) + trunc_normal_(m.weight, std=0.02) if m.bias is not None: nn.init.constant_(m.bias, 0) elif isinstance(m, nn.Conv2d): - torch.nn.init.trunc_normal_(m.weight, std=0.02) + 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)): @@ -573,47 +380,33 @@ def initialize(self): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) + # Copyright (c) ByteDance, Inc. and its affiliates. -class SparK(nn.Module): - """ - The SparK model as used by SparK [0] - - Default params are the ones explained in the original code base. The backbone is assumed - to follow the same API as the ResNet models from torchvision or a ConvNext model. - [0] Designing BERT for Convolutional Networks: Sparse and Hierarchical Masked Modeling https://arxiv.org/abs/2301.03580 - - Attributes: - sparse_encoder: - Sparse encoder to extract features from images. Should have both - the methods get_downsample_ratio() and get_feature_map_channels() - implemented. - dense_decoder: - Dense decoder to reconstruct the image from the sparse features. - mask_ratio: - Ratio of the image to mask. Default: 0.6 - densify_norm: - Type of normalization to use for densification. Default: 'bn' - sbn: - Whether to use SyncBatchNorm. Default: False - """ +# All rights reserved. +# +# This source code is licensed under the license found in the +# LICENSE file in the root directory of this source tree. + + +class SparK(nn.Module): def __init__( self, - sparse_encoder: SparseResnet, + sparse_encoder: SparseEncoder, dense_decoder: LightDecoder, - mask_ratio: float = 0.6, - densify_norm: Literal["batch_norm", "layer_norm", "identity"] = "bn", - sbn: bool = False, + mask_ratio=0.6, + densify_norm="bn", + sbn=False, ): super().__init__() - input_size, downsample_ratio = ( + input_size, downsample_raito = ( sparse_encoder.input_size, - sparse_encoder.downsample_ratio, + sparse_encoder.downsample_raito, ) - self.downsample_ratio = downsample_ratio + self.downsample_raito = downsample_raito self.fmap_h, self.fmap_w = ( - input_size // downsample_ratio, - input_size // downsample_ratio, + input_size // downsample_raito, + input_size // downsample_raito, ) self.mask_ratio = mask_ratio self.len_keep = round(self.fmap_h * self.fmap_w * (1 - mask_ratio)) @@ -628,22 +421,27 @@ def __init__( self.densify_projs = nn.ModuleList() self.mask_tokens = nn.ParameterList() + # build the `densify` layers e_widths, d_width = ( self.sparse_encoder.enc_feat_map_chs, self.dense_decoder.width, ) e_widths: List[int] - for i in range(self.hierarchy): + for i in range( + self.hierarchy + ): # from the smallest feat map to the largest; i=0: the last feat map; i=1: the second last feat map ... e_width = e_widths.pop() + # create mask token p = nn.Parameter(torch.zeros(1, e_width, 1, 1)) - torch.nn.init.trunc_normal_(p, mean=0, std=0.02, a=-0.02, b=0.02) + trunc_normal_(p, mean=0, std=0.02, a=-0.02, b=0.02) self.mask_tokens.append(p) - if self.densify_norm_str == "batch_norm": + # create densify norm + if self.densify_norm_str == "bn": densify_norm = ( SparseSyncBatchNorm2d if self.sbn else SparseBatchNorm2d )(e_width) - elif self.densify_norm_str == "layer_norm": + elif self.densify_norm_str == "ln": densify_norm = SparseConvNeXtLayerNorm( e_width, data_format="channels_first", sparse=True ) @@ -651,8 +449,9 @@ def __init__( densify_norm = nn.Identity() self.densify_norms.append(densify_norm) + # create densify proj if i == 0 and e_width == d_width: - densify_proj = nn.Identity() + densify_proj = nn.Identity() # todo: NOTE THAT CONVNEXT-S WOULD USE THIS, because it has a width of 768 that equals to the decoder's width 768 print( f"[SparK.__init__, densify {i + 1}/{self.hierarchy}]: use nn.Identity() as densify_proj" ) @@ -670,24 +469,21 @@ def __init__( f"[SparK.__init__, densify {i + 1}/{self.hierarchy}]: densify_proj(ksz={kernel_size}, #para={sum(x.numel() for x in densify_proj.parameters()) / 1e6:.2f}M)" ) self.densify_projs.append(densify_proj) + + # todo: the decoder's width follows a simple halfing rule; you can change it to any other rule d_width //= 2 print( f"[SparK.__init__] dims of mask_tokens={tuple(p.numel() for p in self.mask_tokens)}" ) + # these are deprecated and would never be used; can be removed. + self.register_buffer("imn_m", torch.empty(1, 3, 1, 1)) + self.register_buffer("imn_s", torch.empty(1, 3, 1, 1)) + self.register_buffer("norm_black", torch.zeros(1, 3, input_size, input_size)) + self.vis_active = self.vis_active_ex = self.vis_inp = self.vis_inp_mask = ... + def mask(self, B: int, device, generator=None): - """ - Generate a mask for the input tensor - - Attributes: - B: - Batch size - device: - Device to put the mask on - generator: - Random number generator - """ h, w = self.fmap_h, self.fmap_w idx = torch.rand(B, h * w, generator=generator).argsort(dim=1) idx = idx[:, : self.len_keep].to(device) # (B, len_keep) @@ -697,79 +493,81 @@ def mask(self, B: int, device, generator=None): .view(B, 1, h, w) ) - def forward( - self, - inp_bchw: torch.Tensor, - active_b1ff: torch.BoolTensor = None, - return_loss: bool = True, - ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - """ - Forward pass of the SparK model - - Attributes: - inp_bchw: - Input tensor - active_b1ff: - Active mask - return_loss: - Whether to return the loss - """ - if active_b1ff is None: + def forward(self, inp_bchw: torch.Tensor, active_b1ff=None, vis=False): + # step1. Mask + if active_b1ff is None: # rand mask active_b1ff: torch.BoolTensor = self.mask( inp_bchw.shape[0], inp_bchw.device - ) + ) # (B, 1, f, f) + _cur_active = active_b1ff # (B, 1, f, f) active_b1hw = active_b1ff.repeat_interleave( - self.downsample_ratio, 2 - ).repeat_interleave(self.downsample_ratio, 3) + self.downsample_raito, 2 + ).repeat_interleave(self.downsample_raito, 3) # (B, 1, H, W) masked_bchw = inp_bchw * active_b1hw - fea_bcffs: List[torch.Tensor] = self.sparse_encoder.forward( - masked_bchw, active_b1ff, hierarchical=True - ) - fea_bcffs.reverse() + # step2. Encode: get hierarchical encoded sparse features (a list containing 4 feature maps at 4 scales) + fea_bcffs: List[torch.Tensor] = self.sparse_encoder(masked_bchw) + fea_bcffs.reverse() # after reversion: from the smallest feature map to the largest - cur_active = active_b1ff + # step3. Densify: get hierarchical dense features for decoding + cur_active = active_b1ff # (B, 1, f, f) to_dec = [] - for i, bcff in enumerate(fea_bcffs): + for i, bcff in enumerate( + fea_bcffs + ): # from the smallest feature map to the largest if bcff is not None: bcff = self.densify_norms[i](bcff) mask_tokens = self.mask_tokens[i].expand_as(bcff) - bcff = torch.where(cur_active.expand_as(bcff), bcff, mask_tokens) + bcff = torch.where( + cur_active.expand_as(bcff), bcff, mask_tokens + ) # fill in empty (non-active) positions with [mask] tokens bcff: torch.Tensor = self.densify_projs[i](bcff) to_dec.append(bcff) cur_active = cur_active.repeat_interleave(2, dim=2).repeat_interleave( 2, dim=3 - ) + ) # dilate the mask map, from (B, 1, f, f) to (B, 1, H, W) + # step4. Decode and reconstruct rec_bchw = self.dense_decoder(to_dec) - inp, rec = self.patchify(inp_bchw), self.patchify(rec_bchw) + inp, rec = ( + self.patchify(inp_bchw), + self.patchify(rec_bchw), + ) # inp and rec: (B, L = f*f, N = C*downsample_raito**2) mean = inp.mean(dim=-1, keepdim=True) var = (inp.var(dim=-1, keepdim=True) + 1e-6) ** 0.5 inp = (inp - mean) / var - l2_loss = ((rec - inp) ** 2).mean(dim=2, keepdim=False) - - non_active = active_b1ff.logical_not().int().view(active_b1ff.shape[0], -1) - recon_loss = l2_loss.mul_(non_active).sum() / (non_active.sum() + 1e-8) - - if return_loss: + 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() + 1e-8 + ) # loss only on masked (non-active) patches + + if vis: + masked_bchw = inp_bchw * active_b1hw + rec_bchw = self.unpatchify(rec * var + mean) + rec_or_inp = torch.where(active_b1hw, inp_bchw, rec_bchw) + return inp_bchw, masked_bchw, rec_or_inp + else: return recon_loss - masked_bchw = inp_bchw * active_b1hw - rec_bchw = self.unpatchify(rec * var + mean) - rec_or_inp = torch.where(active_b1hw, inp_bchw, rec_bchw) - return inp_bchw, masked_bchw, rec_or_inp - def patchify(self, bchw): - p = self.downsample_ratio + p = self.downsample_raito h, w = self.fmap_h, self.fmap_w B, C = bchw.shape[:2] bchw = bchw.reshape(shape=(B, C, h, p, w, p)) bchw = torch.einsum("bchpwq->bhwpqc", bchw) - bln = bchw.reshape(shape=(B, h * w, C * p**2)) + bln = bchw.reshape( + shape=(B, h * w, C * p**2) + ) # (B, f*f, 3*downsample_raito**2) return bln def unpatchify(self, bln): - p = self.downsample_ratio + p = self.downsample_raito 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)) @@ -781,23 +579,26 @@ def __repr__(self): return ( f"\n" f"[SparK.config]: {pformat(self.get_config(), indent=2, width=250)}\n" - f"[SparK.structure]: {super().__repr__().replace(SparK.__name__, '')}" + f"[SparK.structure]: {super(SparK, self).__repr__().replace(SparK.__name__, '')}" ) def get_config(self): return { + # self "mask_ratio": self.mask_ratio, "densify_norm_str": self.densify_norm_str, "sbn": self.sbn, "hierarchy": self.hierarchy, + # enc "sparse_encoder.input_size": self.sparse_encoder.input_size, + # dec "dense_decoder.width": self.dense_decoder.width, } def state_dict( self, destination=None, prefix="", keep_vars=False, with_config=False ): - state = super().state_dict( + state = super(SparK, self).state_dict( destination=destination, prefix=prefix, keep_vars=keep_vars ) if with_config: @@ -806,7 +607,7 @@ def state_dict( def load_state_dict(self, state_dict, strict=True): config: dict = state_dict.pop("config", None) - incompatible_keys = super().load_state_dict( + incompatible_keys = super(SparK, self).load_state_dict( state_dict, strict=strict ) if config is not None: @@ -818,4 +619,4 @@ def load_state_dict(self, state_dict, strict=True): raise AttributeError(err) else: print(err, file=sys.stderr) - return incompatible_keys \ No newline at end of file + return incompatible_keys From 9a9dff6f84f1d26e50993cb3129722a091c44ea3 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 6 Feb 2026 10:39:05 -0300 Subject: [PATCH 12/85] refactor: fixing type hint problems. --- lightly/models/modules/sparse_spark.py | 30 ++++++++++++++------------ 1 file changed, 16 insertions(+), 14 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 58cb005c3..b82a323c6 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -243,7 +243,6 @@ def dense_model_to_sparse(m: nn.Module, verbose=False, sbn=False): oup.bias.data.copy_(m.bias.data) elif isinstance(m, nn.MaxPool2d): - m: nn.MaxPool2d oup = SparseMaxPooling( m.kernel_size, stride=m.stride, @@ -253,7 +252,6 @@ def dense_model_to_sparse(m: nn.Module, verbose=False, sbn=False): ceil_mode=m.ceil_mode, ) elif isinstance(m, nn.AvgPool2d): - m: nn.AvgPool2d oup = SparseAvgPooling( m.kernel_size, m.stride, @@ -263,7 +261,6 @@ def dense_model_to_sparse(m: nn.Module, verbose=False, sbn=False): divisor_override=m.divisor_override, ) elif isinstance(m, (nn.BatchNorm2d, nn.SyncBatchNorm)): - m: nn.BatchNorm2d oup = (SparseSyncBatchNorm2d if sbn else SparseBatchNorm2d)( m.weight.shape[0], eps=m.eps, @@ -279,7 +276,6 @@ def dense_model_to_sparse(m: nn.Module, verbose=False, sbn=False): if hasattr(m, "qconfig"): oup.qconfig = m.qconfig elif isinstance(m, nn.LayerNorm) and not isinstance(m, SparseConvNeXtLayerNorm): - m: nn.LayerNorm oup = SparseConvNeXtLayerNorm(m.weight.shape[0], eps=m.eps) oup.weight.data.copy_(m.weight.data) oup.bias.data.copy_(m.bias.data) @@ -360,7 +356,7 @@ def forward(self, to_dec: List[torch.Tensor]): def extra_repr(self) -> str: return f"width={self.width}" - def initialize(self): + def initialize(self) -> None: for m in self.modules(): if isinstance(m, nn.Linear): trunc_normal_(m.weight, std=0.02) @@ -380,7 +376,7 @@ def initialize(self): nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) - # Copyright (c) ByteDance, Inc. and its affiliates. + # Copyright (c) ByteDance, Inc. and its affiliates. # All rights reserved. @@ -483,7 +479,7 @@ def __init__( self.register_buffer("norm_black", torch.zeros(1, 3, input_size, input_size)) self.vis_active = self.vis_active_ex = self.vis_inp = self.vis_inp_mask = ... - def mask(self, B: int, device, generator=None): + def mask(self, B: int, device, generator=None) -> torch.BoolTensor: h, w = self.fmap_h, self.fmap_w idx = torch.rand(B, h * w, generator=generator).argsort(dim=1) idx = idx[:, : self.len_keep].to(device) # (B, len_keep) @@ -493,12 +489,16 @@ def mask(self, B: int, device, generator=None): .view(B, 1, h, w) ) - def forward(self, inp_bchw: torch.Tensor, active_b1ff=None, vis=False): + def forward( + self, + inp_bchw: torch.Tensor, + active_b1ff: None | torch.BoolTensor = None, + vis=False, + ): # step1. Mask - if active_b1ff is None: # rand mask - active_b1ff: torch.BoolTensor = self.mask( - inp_bchw.shape[0], inp_bchw.device - ) # (B, 1, f, f) + active_b1ff = active_b1ff or self.mask( + inp_bchw.shape[0], inp_bchw.device + ) # (B, 1, f, f) _cur_active = active_b1ff # (B, 1, f, f) active_b1hw = active_b1ff.repeat_interleave( self.downsample_raito, 2 @@ -597,7 +597,7 @@ def get_config(self): def state_dict( self, destination=None, prefix="", keep_vars=False, with_config=False - ): + ) -> dict[str, torch.Tensor]: state = super(SparK, self).state_dict( destination=destination, prefix=prefix, keep_vars=keep_vars ) @@ -605,7 +605,9 @@ def state_dict( state["config"] = self.get_config() return state - def load_state_dict(self, state_dict, strict=True): + def load_state_dict( + self, state_dict: dict[str, torch.Tensor], strict=True + ) -> dict[str, torch.Tensor]: config: dict = state_dict.pop("config", None) incompatible_keys = super(SparK, self).load_state_dict( state_dict, strict=strict From 407a4c0397f274d83a6009df654e692564c505b8 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 6 Feb 2026 10:44:43 -0300 Subject: [PATCH 13/85] refactor: removing unecessary redundant super --- lightly/models/modules/sparse_spark.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index b82a323c6..b09dead84 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -107,7 +107,7 @@ def forward(self, x): H=x.shape[1], W=x.shape[2], returning_active_ex=False ) nc = x[ii] - nc = super(SparseConvNeXtLayerNorm, self).forward(nc) + nc = super().forward(nc) x = torch.zeros_like(x) x[ii] = nc @@ -121,7 +121,7 @@ def forward(self, x): ) bhwc = x.permute(0, 2, 3, 1) nc = bhwc[ii] - nc = super(SparseConvNeXtLayerNorm, self).forward(nc) + nc = super().forward(nc) x = torch.zeros_like(bhwc) x[ii] = nc @@ -136,13 +136,13 @@ def forward(self, x): if self.sparse: raise NotImplementedError else: - return super(SparseConvNeXtLayerNorm, self).forward(x) + return super().forward(x) def __repr__(self): - return ( - super(SparseConvNeXtLayerNorm, self).__repr__()[:-1] - + f", ch={self.data_format.split('_')[-1]}, sp={self.sparse})" - ) + return ( + super().__repr__()[:-1] + + f", ch={self.data_format.split('_')[-1]}, sp={self.sparse})" + ) class SparseConvNeXtBlock(nn.Module): @@ -203,12 +203,12 @@ def forward(self, x): return x def __repr__(self): - return super(SparseConvNeXtBlock, self).__repr__()[:-1] + f", sp={self.sparse})" + return super().__repr__()[:-1] + f", sp={self.sparse})" class SparseEncoder(nn.Module): def __init__(self, cnn, input_size, sbn=False, verbose=False): - super(SparseEncoder, self).__init__() + super().__init__() self.sp_cnn = SparseEncoder.dense_model_to_sparse( m=cnn, verbose=verbose, sbn=sbn ) @@ -579,7 +579,7 @@ def __repr__(self): return ( f"\n" f"[SparK.config]: {pformat(self.get_config(), indent=2, width=250)}\n" - f"[SparK.structure]: {super(SparK, self).__repr__().replace(SparK.__name__, '')}" + f"[SparK.structure]: {super().__repr__().replace(SparK.__name__, '')}" ) def get_config(self): @@ -598,7 +598,7 @@ def get_config(self): def state_dict( self, destination=None, prefix="", keep_vars=False, with_config=False ) -> dict[str, torch.Tensor]: - state = super(SparK, self).state_dict( + state = super().state_dict( destination=destination, prefix=prefix, keep_vars=keep_vars ) if with_config: @@ -609,7 +609,7 @@ def load_state_dict( self, state_dict: dict[str, torch.Tensor], strict=True ) -> dict[str, torch.Tensor]: config: dict = state_dict.pop("config", None) - incompatible_keys = super(SparK, self).load_state_dict( + incompatible_keys = super().load_state_dict( state_dict, strict=strict ) if config is not None: From 47eecf38a93f0a10c8fd9f4136dbb563c3756822 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 6 Feb 2026 11:03:43 -0300 Subject: [PATCH 14/85] fix: indentation --- lightly/models/modules/sparse_spark.py | 56 +++++++++++++------------- 1 file changed, 28 insertions(+), 28 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index b09dead84..68610a16a 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -99,50 +99,52 @@ def __init__( self.data_format = data_format self.sparse = sparse - def forward(self, x): - if x.ndim == 4: # BHWC or BCHW + def forward(self, input: torch.Tensor): + if input.ndim == 4: # BHWC or BCHW if self.data_format == "channels_last": # BHWC if self.sparse: ii = _get_active_ex_or_ii( - H=x.shape[1], W=x.shape[2], returning_active_ex=False + H=input.shape[1], W=input.shape[2], returning_active_ex=False ) - nc = x[ii] - nc = super().forward(nc) + nc = input[ii] + nc = super().forward(nc) - x = torch.zeros_like(x) - x[ii] = nc - return x + input = torch.zeros_like(input) + input[ii] = nc + return input else: - return super(SparseConvNeXtLayerNorm, self).forward(x) + return super(SparseConvNeXtLayerNorm, self).forward(input) else: # channels_first, BCHW if self.sparse: ii = _get_active_ex_or_ii( - H=x.shape[2], W=x.shape[3], returning_active_ex=False + H=input.shape[2], W=input.shape[3], returning_active_ex=False ) - bhwc = x.permute(0, 2, 3, 1) + bhwc = input.permute(0, 2, 3, 1) nc = bhwc[ii] - nc = super().forward(nc) + nc = super().forward(nc) - x = torch.zeros_like(bhwc) - x[ii] = nc - return x.permute(0, 3, 1, 2) + input = torch.zeros_like(bhwc) + input[ii] = nc + return input.permute(0, 3, 1, 2) else: - u = x.mean(1, keepdim=True) - s = (x - u).pow(2).mean(1, keepdim=True) - x = (x - u) / torch.sqrt(s + self.eps) - x = self.weight[:, None, None] * x + self.bias[:, None, None] - return x + u = input.mean(1, keepdim=True) + s = (input - u).pow(2).mean(1, keepdim=True) + input = (input - u) / torch.sqrt(s + self.eps) + input = ( + self.weight[:, None, None] * input + self.bias[:, None, None] + ) + return input else: # BLC or BC if self.sparse: raise NotImplementedError else: - return super().forward(x) + return super().forward(input) def __repr__(self): - return ( - super().__repr__()[:-1] - + f", ch={self.data_format.split('_')[-1]}, sp={self.sparse})" - ) + return ( + super().__repr__()[:-1] + + f", ch={self.data_format.split('_')[-1]}, sp={self.sparse})" + ) class SparseConvNeXtBlock(nn.Module): @@ -609,9 +611,7 @@ def load_state_dict( self, state_dict: dict[str, torch.Tensor], strict=True ) -> dict[str, torch.Tensor]: config: dict = state_dict.pop("config", None) - incompatible_keys = super().load_state_dict( - state_dict, strict=strict - ) + incompatible_keys = super().load_state_dict(state_dict, strict=strict) if config is not None: for k, v in self.get_config().items(): ckpt_v = config.get(k, None) From 7a956b324e096eb5c3660bde0469e8bcdb59b428 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 6 Feb 2026 13:51:47 -0300 Subject: [PATCH 15/85] feat: working module --- examples/pytorch_lightning/spark.py | 92 ++++++++++++++++++++++++++ lightly/models/modules/sparse_spark.py | 50 +++++++++++++- 2 files changed, 139 insertions(+), 3 deletions(-) create mode 100644 examples/pytorch_lightning/spark.py diff --git a/examples/pytorch_lightning/spark.py b/examples/pytorch_lightning/spark.py new file mode 100644 index 000000000..cb612cdce --- /dev/null +++ b/examples/pytorch_lightning/spark.py @@ -0,0 +1,92 @@ +# This example requires the following dependencies to be installed: +# pip install lightly + +# 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.callbacks import RichProgressBar +from torchvision.transforms import v2 + +## The global projection head is the same as the Barlow Twins one +from lightly.models.modules.sparse_spark import LightDecoder, SparK, SparseEncoder + + +class SparseSpark(pl.LightningModule): + def __init__(self) -> None: + super().__init__() + backbone = timm.create_model( + "resnet50", drop_path_rate=0.05, features_only=True + ) + self.sparse_encoder = SparseEncoder( + backbone, input_size=416, sbn=False, verbose=True + ) + self.dense_decoder = LightDecoder(self.sparse_encoder.downsample_raito) + self.spark = SparK( + sparse_encoder=self.sparse_encoder, + dense_decoder=self.dense_decoder, + ) + + def forward(self, x): + return self.spark(x) + + def training_step(self, batch, batch_index) -> torch.Tensor: + img, target = 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): + optim = torch.optim.SGD(self.parameters(), lr=0.01, momentum=0.9) + return optim + + +model = SparseSpark() + + +# we ignore object detection annotations by setting target_transform to return 0 +def target_transform(t): + return 0 + + +dataset = torchvision.datasets.Caltech101( + "datasets/caltech101", + download=True, + transform=v2.Compose( + [ + v2.Resize((416, 416)), + v2.RGB(), + v2.ToTensor(), + ] + ), + target_transform=target_transform, +) +# 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=0, +) + +accelerator = "gpu" if torch.cuda.is_available() else "cpu" + +trainer = pl.Trainer( + max_epochs=10, devices=1, accelerator=accelerator, callbacks=[RichProgressBar()] +) +trainer.fit(model=model, train_dataloaders=dataloader) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 68610a16a..9f285c325 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -3,6 +3,7 @@ from pprint import pformat from typing import List +import timm import torch import torch.nn as nn from timm.models.layers import DropPath, trunc_normal_ @@ -208,7 +209,26 @@ def __repr__(self): return super().__repr__()[:-1] + f", sp={self.sparse})" +def get_downsample_ratio_from_timm_model(model: nn.Module) -> int: + return model.feature_info[-1]["reduction"] + + +def get_enc_feat_map_chs_from_timm_model(model: nn.Module) -> List[int]: + return [fi["num_chs"] for fi in model.feature_info] + + class SparseEncoder(nn.Module): + """ + Converts a dense CNN model to a sparse CNN model by replacing standard layers + + Attributes: + enc_feat_map_chs: List[int]: list of channel numbers of feature maps at different scales, in the order from shallow to deep + + + """ + + enc_feat_map_chs: List[int] + def __init__(self, cnn, input_size, sbn=False, verbose=False): super().__init__() self.sp_cnn = SparseEncoder.dense_model_to_sparse( @@ -216,8 +236,8 @@ def __init__(self, cnn, input_size, sbn=False, verbose=False): ) self.input_size, self.downsample_raito, self.enc_feat_map_chs = ( input_size, - cnn.get_downsample_ratio(), - cnn.get_feature_map_channels(), + get_downsample_ratio_from_timm_model(cnn), + get_enc_feat_map_chs_from_timm_model(cnn), ) @staticmethod @@ -293,7 +313,7 @@ def dense_model_to_sparse(m: nn.Module, verbose=False, sbn=False): return oup def forward(self, x): - return self.sp_cnn(x, hierarchical=True) + return self.sp_cnn(x) # Copyright (c) ByteDance, Inc. and its affiliates. @@ -386,6 +406,29 @@ def initialize(self) -> None: # This source code is licensed under the license found in the # LICENSE file in the root directory of this source tree. +pretrain_default_model_kwargs = { + "your_convnet": dict(), + "resnet50": dict(drop_path_rate=0.05), + "resnet101": dict(drop_path_rate=0.08), + "resnet152": dict(drop_path_rate=0.10), + "resnet200": dict(drop_path_rate=0.15), + "convnext_small": dict(sparse=True, drop_path_rate=0.2), + "convnext_base": dict(sparse=True, drop_path_rate=0.3), + "convnext_large": dict(sparse=True, drop_path_rate=0.4), +} + + +def build_sparse_encoder( + name: str, input_size: int, sbn=False, drop_path_rate=0.0, verbose=False +): + kwargs = pretrain_default_model_kwargs[name] + if drop_path_rate != 0: + kwargs["drop_path_rate"] = drop_path_rate + print(f"[build_sparse_encoder] model kwargs={kwargs}") + cnn = timm.create_model(name, **kwargs) + + return SparseEncoder(cnn, input_size=input_size, sbn=sbn, verbose=verbose) + class SparK(nn.Module): def __init__( @@ -501,6 +544,7 @@ def forward( active_b1ff = active_b1ff or self.mask( inp_bchw.shape[0], inp_bchw.device ) # (B, 1, f, f) + global _cur_active _cur_active = active_b1ff # (B, 1, f, f) active_b1hw = active_b1ff.repeat_interleave( self.downsample_raito, 2 From 9d0f5c4a02f17f5b7503f13945b69b4a2f9d32b6 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 08:47:26 -0300 Subject: [PATCH 16/85] refactor: using library already implemente masking --- lightly/models/modules/sparse_spark.py | 143 +++++++++++++++---------- 1 file changed, 86 insertions(+), 57 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 9f285c325..35792d9dc 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -9,6 +9,8 @@ from timm.models.layers import DropPath, trunc_normal_ from torch.nn.common_types import _size_2_t +from lightly.models.utils import random_token_mask + def is_pow2n(x): return x > 0 and (x & (x - 1) == 0) @@ -430,6 +432,64 @@ def build_sparse_encoder( return SparseEncoder(cnn, input_size=input_size, sbn=sbn, verbose=verbose) +class SparKDensfiyBlock(nn.Module): + 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, + ): + 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", sparse=True + ) + else: + self.densify_norm = nn.Identity() + + # densify proj + 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 forward(self, bcff: torch.Tensor, cur_active: torch.BoolTensor): + if bcff is None: + return None + bcff = self.densify_norm(bcff) + mask_tokens = self.mask_token.expand_as(bcff) + bcff = torch.where(cur_active.expand_as(bcff), bcff, mask_tokens) + bcff = self.densify_proj(bcff) + return bcff + + class SparK(nn.Module): def __init__( self, @@ -458,9 +518,7 @@ def __init__( self.sbn = sbn self.hierarchy = len(sparse_encoder.enc_feat_map_chs) self.densify_norm_str = densify_norm.lower() - self.densify_norms = nn.ModuleList() - self.densify_projs = nn.ModuleList() - self.mask_tokens = nn.ParameterList() + self.densify_blocks = nn.ModuleList() # build the `densify` layers e_widths, d_width = ( @@ -468,54 +526,28 @@ def __init__( self.dense_decoder.width, ) e_widths: List[int] - for i in range( - self.hierarchy - ): # from the smallest feat map to the largest; i=0: the last feat map; i=1: the second last feat map ... + for i in range(self.hierarchy): + # from the smallest feat map to the largest; i=0: the last feat map; i=1: the second last feat map ... e_width = e_widths.pop() - # create 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_tokens.append(p) - - # create densify norm - if self.densify_norm_str == "bn": - densify_norm = ( - SparseSyncBatchNorm2d if self.sbn else SparseBatchNorm2d - )(e_width) - elif self.densify_norm_str == "ln": - densify_norm = SparseConvNeXtLayerNorm( - e_width, data_format="channels_first", sparse=True - ) - else: - densify_norm = nn.Identity() - self.densify_norms.append(densify_norm) - - # create densify proj - if i == 0 and e_width == d_width: - densify_proj = nn.Identity() # todo: NOTE THAT CONVNEXT-S WOULD USE THIS, because it has a width of 768 that equals to the decoder's width 768 - print( - f"[SparK.__init__, densify {i + 1}/{self.hierarchy}]: use nn.Identity() as densify_proj" - ) - else: - kernel_size = 1 if i <= 0 else 3 - densify_proj = nn.Conv2d( - e_width, - d_width, - kernel_size=kernel_size, - stride=1, - padding=kernel_size // 2, - bias=True, - ) - print( - f"[SparK.__init__, densify {i + 1}/{self.hierarchy}]: densify_proj(ksz={kernel_size}, #para={sum(x.numel() for x in densify_proj.parameters()) / 1e6:.2f}M)" - ) - self.densify_projs.append(densify_proj) + # fork arguments that depend on the position (previously used idx inside the block) + use_identity = i == 0 and e_width == d_width + kernel_size = 1 if i <= 0 else 3 + # build densify block and append + block = SparKDensfiyBlock( + e_width=e_width, + d_width=d_width, + densify_norm_str=self.densify_norm_str, + sbn=self.sbn, + use_identity_proj=use_identity, + kernel_size=kernel_size, + ) + self.densify_blocks.append(block) # todo: the decoder's width follows a simple halfing rule; you can change it to any other rule d_width //= 2 print( - f"[SparK.__init__] dims of mask_tokens={tuple(p.numel() for p in self.mask_tokens)}" + f"[SparK.__init__] dims of mask_tokens={tuple(b.mask_token.numel() for b in self.densify_blocks)}" ) # these are deprecated and would never be used; can be removed. @@ -524,14 +556,16 @@ def __init__( self.register_buffer("norm_black", torch.zeros(1, 3, input_size, input_size)) self.vis_active = self.vis_active_ex = self.vis_inp = self.vis_inp_mask = ... - def mask(self, B: int, device, generator=None) -> torch.BoolTensor: + def mask(self, B: int, device: torch.device) -> torch.BoolTensor: h, w = self.fmap_h, self.fmap_w - idx = torch.rand(B, h * w, generator=generator).argsort(dim=1) - idx = idx[:, : self.len_keep].to(device) # (B, len_keep) - return ( - torch.zeros(B, h * w, dtype=torch.bool, device=device) - .scatter_(dim=1, index=idx, value=True) + index_keep, _ = random_token_mask( + size=(B, h * w), mask_ratio=self.mask_ratio, device=device + ) + return torch.BoolTensor( + 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( @@ -562,12 +596,7 @@ def forward( fea_bcffs ): # from the smallest feature map to the largest if bcff is not None: - bcff = self.densify_norms[i](bcff) - mask_tokens = self.mask_tokens[i].expand_as(bcff) - bcff = torch.where( - cur_active.expand_as(bcff), bcff, mask_tokens - ) # fill in empty (non-active) positions with [mask] tokens - bcff: torch.Tensor = self.densify_projs[i](bcff) + bcff = self.densify_blocks[i](bcff, cur_active) to_dec.append(bcff) cur_active = cur_active.repeat_interleave(2, dim=2).repeat_interleave( 2, dim=3 From 4aa69f7e368d8458cba8b2058feaf3926fef4d28 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 08:50:56 -0300 Subject: [PATCH 17/85] feat: using patchify --- lightly/models/modules/sparse_spark.py | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 35792d9dc..55ec56b29 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -9,7 +9,7 @@ from timm.models.layers import DropPath, trunc_normal_ from torch.nn.common_types import _size_2_t -from lightly.models.utils import random_token_mask +from lightly.models.utils import patchify, random_token_mask, unpatchify def is_pow2n(x): @@ -634,12 +634,7 @@ def patchify(self, bchw): p = self.downsample_raito h, w = self.fmap_h, self.fmap_w B, C = bchw.shape[:2] - bchw = bchw.reshape(shape=(B, C, h, p, w, p)) - bchw = torch.einsum("bchpwq->bhwpqc", bchw) - bln = bchw.reshape( - shape=(B, h * w, C * p**2) - ) # (B, f*f, 3*downsample_raito**2) - return bln + return patchify(bchw, p) # (B, L=f*f, N=C*p*p) def unpatchify(self, bln): p = self.downsample_raito From 8b95a51fbd448841fcff355cdee8cdd2a12d631c Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 08:54:32 -0300 Subject: [PATCH 18/85] refactor: putting densification into a single module --- lightly/models/modules/sparse_spark.py | 39 ++++++++++++++++---------- 1 file changed, 24 insertions(+), 15 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 55ec56b29..7272656d4 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -9,7 +9,7 @@ from timm.models.layers import DropPath, trunc_normal_ from torch.nn.common_types import _size_2_t -from lightly.models.utils import patchify, random_token_mask, unpatchify +from lightly.models.utils import patchify, random_token_mask def is_pow2n(x): @@ -490,6 +490,25 @@ def forward(self, bcff: torch.Tensor, cur_active: torch.BoolTensor): return bcff +class SparKDensifier(nn.Module): + def __init__(self, blocks: List[SparKDensfiyBlock]): + super().__init__() + self.blocks = nn.ModuleList(blocks) + + def forward(self, fea_bcffs: List[torch.Tensor], cur_active: torch.BoolTensor): + to_dec = [] + for i, bcff in enumerate( + fea_bcffs + ): # from the smallest feature map to the largest + if bcff is not None: + bcff = self.blocks[i](bcff, cur_active) + to_dec.append(bcff) + cur_active = cur_active.repeat_interleave(2, dim=2).repeat_interleave( + 2, dim=3 + ) # dilate the mask map, from (B, 1, f, f) to (B, 1, H, W) + return to_dec + + class SparK(nn.Module): def __init__( self, @@ -518,7 +537,6 @@ def __init__( self.sbn = sbn self.hierarchy = len(sparse_encoder.enc_feat_map_chs) self.densify_norm_str = densify_norm.lower() - self.densify_blocks = nn.ModuleList() # build the `densify` layers e_widths, d_width = ( @@ -526,6 +544,7 @@ def __init__( self.dense_decoder.width, ) e_widths: List[int] + densify_blocks = [] for i in range(self.hierarchy): # from the smallest feat map to the largest; i=0: the last feat map; i=1: the second last feat map ... e_width = e_widths.pop() @@ -541,11 +560,11 @@ def __init__( use_identity_proj=use_identity, kernel_size=kernel_size, ) - self.densify_blocks.append(block) - + densify_blocks.append(block) # todo: the decoder's width follows a simple halfing rule; you can change it to any other rule d_width //= 2 + self.densifier = SparKDensifier(densify_blocks) print( f"[SparK.__init__] dims of mask_tokens={tuple(b.mask_token.numel() for b in self.densify_blocks)}" ) @@ -590,17 +609,7 @@ def forward( fea_bcffs.reverse() # after reversion: from the smallest feature map to the largest # step3. Densify: get hierarchical dense features for decoding - cur_active = active_b1ff # (B, 1, f, f) - to_dec = [] - for i, bcff in enumerate( - fea_bcffs - ): # from the smallest feature map to the largest - if bcff is not None: - bcff = self.densify_blocks[i](bcff, cur_active) - to_dec.append(bcff) - cur_active = cur_active.repeat_interleave(2, dim=2).repeat_interleave( - 2, dim=3 - ) # dilate the mask map, from (B, 1, f, f) to (B, 1, H, W) + to_dec = self.densifier(fea_bcffs, active_b1ff) # step4. Decode and reconstruct rec_bchw = self.dense_decoder(to_dec) From 014f724e218acdf61cf21a94685568730efaca62 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 08:55:45 -0300 Subject: [PATCH 19/85] typo: raito -> ratio --- lightly/models/modules/sparse_spark.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 7272656d4..6fdfbedec 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -236,7 +236,7 @@ def __init__(self, cnn, input_size, sbn=False, verbose=False): self.sp_cnn = SparseEncoder.dense_model_to_sparse( m=cnn, verbose=verbose, sbn=sbn ) - self.input_size, self.downsample_raito, self.enc_feat_map_chs = ( + self.input_size, self.downsample_ratio, self.enc_feat_map_chs = ( input_size, get_downsample_ratio_from_timm_model(cnn), get_enc_feat_map_chs_from_timm_model(cnn), @@ -519,14 +519,14 @@ def __init__( sbn=False, ): super().__init__() - input_size, downsample_raito = ( + input_size, downsample_ratio = ( sparse_encoder.input_size, - sparse_encoder.downsample_raito, + sparse_encoder.downsample_ratio, ) - self.downsample_raito = downsample_raito + self.downsample_ratio = downsample_ratio self.fmap_h, self.fmap_w = ( - input_size // downsample_raito, - input_size // downsample_raito, + input_size // downsample_ratio, + input_size // downsample_ratio, ) self.mask_ratio = mask_ratio self.len_keep = round(self.fmap_h * self.fmap_w * (1 - mask_ratio)) @@ -600,8 +600,8 @@ def forward( global _cur_active _cur_active = active_b1ff # (B, 1, f, f) active_b1hw = active_b1ff.repeat_interleave( - self.downsample_raito, 2 - ).repeat_interleave(self.downsample_raito, 3) # (B, 1, H, W) + self.downsample_ratio, 2 + ).repeat_interleave(self.downsample_ratio, 3) # (B, 1, H, W) masked_bchw = inp_bchw * active_b1hw # step2. Encode: get hierarchical encoded sparse features (a list containing 4 feature maps at 4 scales) @@ -640,13 +640,13 @@ def forward( return recon_loss def patchify(self, bchw): - p = self.downsample_raito + p = self.downsample_ratio h, w = self.fmap_h, self.fmap_w B, C = bchw.shape[:2] return patchify(bchw, p) # (B, L=f*f, N=C*p*p) def unpatchify(self, bln): - p = self.downsample_raito + 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)) From de01b0a70d936f4ff9f96fc4ea263935388e36a1 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 09:01:37 -0300 Subject: [PATCH 20/85] feat: encapsulated logic to single dnesifier module --- lightly/models/modules/sparse_spark.py | 64 +++++++++++++++----------- 1 file changed, 38 insertions(+), 26 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 6fdfbedec..7e96f181b 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -491,19 +491,46 @@ def forward(self, bcff: torch.Tensor, cur_active: torch.BoolTensor): class SparKDensifier(nn.Module): - def __init__(self, blocks: List[SparKDensfiyBlock]): + def __init__( + self, + encoder_in_channels: list[int], + decoder_in_channel: int, + densify_norm_str: str = "bn", + sbn: bool = False, + ): super().__init__() - self.blocks = nn.ModuleList(blocks) + 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): + # from the smallest feat map to the largest; i=0: the last feat map; i=1: the second last feat map ... + # fork arguments that depend on the position (previously used idx inside the block) + use_identity = i == 0 and e_width == d_width + kernel_size = 1 if i <= 0 else 3 + # build densify block and append + block = SparKDensfiyBlock( + 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) + # todo: the decoder's width follows a simple halfing rule; you can change it to any other rule + d_width //= 2 - def forward(self, fea_bcffs: List[torch.Tensor], cur_active: torch.BoolTensor): + def forward(self, fea_bcffs: List[torch.Tensor]): to_dec = [] + global _cur_active for i, bcff in enumerate( fea_bcffs ): # from the smallest feature map to the largest if bcff is not None: - bcff = self.blocks[i](bcff, cur_active) + bcff = self.blocks[i](bcff, _cur_active) to_dec.append(bcff) - cur_active = cur_active.repeat_interleave(2, dim=2).repeat_interleave( + _cur_active = _cur_active.repeat_interleave(2, dim=2).repeat_interleave( 2, dim=3 ) # dilate the mask map, from (B, 1, f, f) to (B, 1, H, W) return to_dec @@ -543,28 +570,13 @@ def __init__( self.sparse_encoder.enc_feat_map_chs, self.dense_decoder.width, ) - e_widths: List[int] - densify_blocks = [] - for i in range(self.hierarchy): - # from the smallest feat map to the largest; i=0: the last feat map; i=1: the second last feat map ... - e_width = e_widths.pop() - # fork arguments that depend on the position (previously used idx inside the block) - use_identity = i == 0 and e_width == d_width - kernel_size = 1 if i <= 0 else 3 - # build densify block and append - block = SparKDensfiyBlock( - e_width=e_width, - d_width=d_width, - densify_norm_str=self.densify_norm_str, - sbn=self.sbn, - use_identity_proj=use_identity, - kernel_size=kernel_size, - ) - densify_blocks.append(block) - # todo: the decoder's width follows a simple halfing rule; you can change it to any other rule - d_width //= 2 - self.densifier = SparKDensifier(densify_blocks) + self.densifier = SparKDensifier( + encoder_in_channels=e_widths, + decoder_in_channel=d_width, + densify_norm_str=self.densify_norm_str, + sbn=self.sbn, + ) print( f"[SparK.__init__] dims of mask_tokens={tuple(b.mask_token.numel() for b in self.densify_blocks)}" ) From 445bda76b78e0336d71611cd3f8e8cc22ac9d71f Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 09:04:46 -0300 Subject: [PATCH 21/85] refactor: cleaning code. --- lightly/models/modules/sparse_spark.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 7e96f181b..26eb7dc4b 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -562,18 +562,11 @@ def __init__( self.dense_decoder = dense_decoder self.sbn = sbn - self.hierarchy = len(sparse_encoder.enc_feat_map_chs) self.densify_norm_str = densify_norm.lower() - # build the `densify` layers - e_widths, d_width = ( - self.sparse_encoder.enc_feat_map_chs, - self.dense_decoder.width, - ) - self.densifier = SparKDensifier( - encoder_in_channels=e_widths, - decoder_in_channel=d_width, + encoder_in_channels=self.sparse_encoder.enc_feat_map_chs, + decoder_in_channel=self.dense_decoder.width, densify_norm_str=self.densify_norm_str, sbn=self.sbn, ) @@ -679,7 +672,6 @@ def get_config(self): "mask_ratio": self.mask_ratio, "densify_norm_str": self.densify_norm_str, "sbn": self.sbn, - "hierarchy": self.hierarchy, # enc "sparse_encoder.input_size": self.sparse_encoder.input_size, # dec From 741d5310dbb6ec2c2cd7b86d08897b1bae57e893 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 09:12:19 -0300 Subject: [PATCH 22/85] refactor: letting sparse encoder be repsonsible for sizes and etc --- lightly/models/modules/sparse_spark.py | 34 ++++++++++++-------------- 1 file changed, 15 insertions(+), 19 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 26eb7dc4b..8c6d95ea0 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -241,6 +241,9 @@ def __init__(self, cnn, input_size, sbn=False, verbose=False): get_downsample_ratio_from_timm_model(cnn), get_enc_feat_map_chs_from_timm_model(cnn), ) + # feature-map spatial size (height, width) at encoder output (patch grid) + self.fmap_h = input_size // self.downsample_ratio + self.fmap_w = input_size // self.downsample_ratio @staticmethod def dense_model_to_sparse(m: nn.Module, verbose=False, sbn=False): @@ -546,17 +549,11 @@ def __init__( sbn=False, ): super().__init__() - input_size, downsample_ratio = ( - sparse_encoder.input_size, - sparse_encoder.downsample_ratio, - ) - self.downsample_ratio = downsample_ratio - self.fmap_h, self.fmap_w = ( - input_size // downsample_ratio, - input_size // downsample_ratio, - ) + # spatial and size info moved to SparseEncoder self.mask_ratio = mask_ratio - self.len_keep = round(self.fmap_h * self.fmap_w * (1 - mask_ratio)) + self.len_keep = round( + sparse_encoder.fmap_h * sparse_encoder.fmap_w * (1 - mask_ratio) + ) self.sparse_encoder = sparse_encoder self.dense_decoder = dense_decoder @@ -581,7 +578,7 @@ def __init__( self.vis_active = self.vis_active_ex = self.vis_inp = self.vis_inp_mask = ... def mask(self, B: int, device: torch.device) -> torch.BoolTensor: - h, w = self.fmap_h, self.fmap_w + h, w = self.sparse_encoder.fmap_h, self.sparse_encoder.fmap_w index_keep, _ = random_token_mask( size=(B, h * w), mask_ratio=self.mask_ratio, device=device ) @@ -604,9 +601,10 @@ def forward( ) # (B, 1, f, f) global _cur_active _cur_active = active_b1ff # (B, 1, f, f) - active_b1hw = active_b1ff.repeat_interleave( - self.downsample_ratio, 2 - ).repeat_interleave(self.downsample_ratio, 3) # (B, 1, H, W) + ds = self.sparse_encoder.downsample_ratio + active_b1hw = active_b1ff.repeat_interleave(ds, 2).repeat_interleave( + ds, 3 + ) # (B, 1, H, W) masked_bchw = inp_bchw * active_b1hw # step2. Encode: get hierarchical encoded sparse features (a list containing 4 feature maps at 4 scales) @@ -645,14 +643,12 @@ def forward( return recon_loss def patchify(self, bchw): - p = self.downsample_ratio - h, w = self.fmap_h, self.fmap_w - B, C = bchw.shape[:2] + p = self.sparse_encoder.downsample_ratio return patchify(bchw, p) # (B, L=f*f, N=C*p*p) def unpatchify(self, bln): - p = self.downsample_ratio - h, w = self.fmap_h, self.fmap_w + p = self.sparse_encoder.downsample_ratio + h, w = self.sparse_encoder.fmap_h, self.sparse_encoder.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) From fb5c90de0547826c2868d52081527ce5db4d6859 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 09:13:06 -0300 Subject: [PATCH 23/85] feat: resnet18 --- examples/pytorch_lightning/spark.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/examples/pytorch_lightning/spark.py b/examples/pytorch_lightning/spark.py index cb612cdce..cf3864973 100644 --- a/examples/pytorch_lightning/spark.py +++ b/examples/pytorch_lightning/spark.py @@ -20,12 +20,15 @@ class SparseSpark(pl.LightningModule): def __init__(self) -> None: super().__init__() backbone = timm.create_model( - "resnet50", drop_path_rate=0.05, features_only=True + "resnet18", drop_path_rate=0.05, features_only=True ) self.sparse_encoder = SparseEncoder( backbone, input_size=416, sbn=False, verbose=True ) - self.dense_decoder = LightDecoder(self.sparse_encoder.downsample_raito) + self.dense_decoder = LightDecoder( + self.sparse_encoder.downsample_ratio, + width=self.sparse_encoder.enc_feat_map_chs[-1], + ) self.spark = SparK( sparse_encoder=self.sparse_encoder, dense_decoder=self.dense_decoder, From deeb4ef39a917999977814659a474b8afdd3161a Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 09:13:49 -0300 Subject: [PATCH 24/85] refactor: removing unused code --- lightly/models/modules/sparse_spark.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 8c6d95ea0..66e8d36f7 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -568,13 +568,10 @@ def __init__( sbn=self.sbn, ) print( - f"[SparK.__init__] dims of mask_tokens={tuple(b.mask_token.numel() for b in self.densify_blocks)}" + f"[SparK.__init__] dims of mask_tokens={tuple(b.mask_token.numel() for b in self.densifier.blocks)}" ) # these are deprecated and would never be used; can be removed. - self.register_buffer("imn_m", torch.empty(1, 3, 1, 1)) - self.register_buffer("imn_s", torch.empty(1, 3, 1, 1)) - self.register_buffer("norm_black", torch.zeros(1, 3, input_size, input_size)) self.vis_active = self.vis_active_ex = self.vis_inp = self.vis_inp_mask = ... def mask(self, B: int, device: torch.device) -> torch.BoolTensor: From 64df2e3ca83901fc099e414eb10f6885bb0aea1d Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 09:22:18 -0300 Subject: [PATCH 25/85] fix: bool tensor is inconvenient --- lightly/models/modules/sparse_spark.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 66e8d36f7..015202a07 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -574,12 +574,12 @@ def __init__( # these are deprecated and would never be used; can be removed. self.vis_active = self.vis_active_ex = self.vis_inp = self.vis_inp_mask = ... - def mask(self, B: int, device: torch.device) -> torch.BoolTensor: + def mask(self, B: int, device: torch.device) -> torch.Tensor: h, w = self.sparse_encoder.fmap_h, self.sparse_encoder.fmap_w index_keep, _ = random_token_mask( size=(B, h * w), mask_ratio=self.mask_ratio, device=device ) - return torch.BoolTensor( + 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) @@ -589,7 +589,7 @@ def mask(self, B: int, device: torch.device) -> torch.BoolTensor: def forward( self, inp_bchw: torch.Tensor, - active_b1ff: None | torch.BoolTensor = None, + active_b1ff: None | torch.Tensor = None, vis=False, ): # step1. Mask From 45b6eb8509a2223df1ebab8dd7ac27b14fde6699 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 10:09:47 -0300 Subject: [PATCH 26/85] refactor: documenting --- lightly/models/modules/sparse_spark.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 015202a07..788e146b3 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -494,6 +494,16 @@ def forward(self, bcff: torch.Tensor, cur_active: torch.BoolTensor): class SparKDensifier(nn.Module): + """ + A stack of SparKDensfiyBlocks to convert hierarchical sparse features to hierarchical dense features for decoding + + Args: + encoder_in_channels: list of channel numbers of feature maps at different scales from the encoder, in the order from shallow to deep + decoder_in_channel: channel number of the feature map at the deepest scale for decoding + densify_norm_str: the type of normalization inside SparKDensfiyBlock; can be "bn", "ln" or "none" + sbn: whether to use SyncBatchNorm (True) or BatchNorm (False) if densify_norm_str is "bn" + """ + def __init__( self, encoder_in_channels: list[int], @@ -506,7 +516,7 @@ def __init__( 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): + for i, e_width in enumerate(encoder_in_channels[::-1]): # from the smallest feat map to the largest; i=0: the last feat map; i=1: the second last feat map ... # fork arguments that depend on the position (previously used idx inside the block) use_identity = i == 0 and e_width == d_width @@ -525,7 +535,12 @@ def __init__( d_width //= 2 def forward(self, fea_bcffs: List[torch.Tensor]): + """ + Args: + fea_bcffs: a list of feature maps at different scales from the encoder, in the order from shallow to deep; each feature map is a tensor of shape (B, C, f, f) + """ to_dec = [] + fea_bcffs = fea_bcffs[::-1] # from the smallest feat map to the largest global _cur_active for i, bcff in enumerate( fea_bcffs @@ -606,11 +621,8 @@ def forward( # step2. Encode: get hierarchical encoded sparse features (a list containing 4 feature maps at 4 scales) fea_bcffs: List[torch.Tensor] = self.sparse_encoder(masked_bchw) - fea_bcffs.reverse() # after reversion: from the smallest feature map to the largest - # step3. Densify: get hierarchical dense features for decoding - to_dec = self.densifier(fea_bcffs, active_b1ff) - + to_dec = self.densifier(fea_bcffs) # step4. Decode and reconstruct rec_bchw = self.dense_decoder(to_dec) inp, rec = ( From fb7903a16c0b25b28994ee50526f492aeacd48fd Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 10:43:48 -0300 Subject: [PATCH 27/85] refactor: masking as a module --- lightly/models/modules/sparse_spark.py | 88 ++++++++++++++++++-------- 1 file changed, 60 insertions(+), 28 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 788e146b3..6db7b1f27 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -1,7 +1,7 @@ import math import sys from pprint import pformat -from typing import List +from typing import List, NamedTuple import timm import torch @@ -542,18 +542,68 @@ def forward(self, fea_bcffs: List[torch.Tensor]): to_dec = [] fea_bcffs = fea_bcffs[::-1] # from the smallest feat map to the largest global _cur_active + active_fmap_current = ( + _cur_active # (B, 1, f, f), the mask map at the current scale + ) for i, bcff in enumerate( fea_bcffs ): # from the smallest feature map to the largest if bcff is not None: - bcff = self.blocks[i](bcff, _cur_active) + bcff = self.blocks[i](bcff, active_fmap_current) to_dec.append(bcff) - _cur_active = _cur_active.repeat_interleave(2, dim=2).repeat_interleave( + active_fmap_current = active_fmap_current.repeat_interleave( + 2, dim=2 + ).repeat_interleave( 2, dim=3 ) # dilate the mask map, from (B, 1, f, f) to (B, 1, H, W) return to_dec +class SparKMaskingOuptut(NamedTuple): + masked_bchw: torch.Tensor + per_level_mask: List[torch.Tensor] + + +class SparKMasker(nn.Module): + def __init__(self, sparse_encoder: SparseEncoder, mask_ratio: float = 0.6): + super().__init__() + self.sparse_encoder = sparse_encoder + self.mask_ratio = mask_ratio + + def mask(self, B: int, device: torch.device) -> torch.Tensor: + h, w = self.sparse_encoder.fmap_h, self.sparse_encoder.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: torch.Tensor) -> SparKMaskingOuptut: + global _cur_active + _cur_active = self.mask(inp_bchw.shape[0], inp_bchw.device) # (B, 1, f, f) + active_b1ff = _cur_active.clone() + downsample_ratio = self.sparse_encoder.downsample_ratio + per_level_mask = [active_b1ff] + for i 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 + ) # (B, 1, f*ds, f*ds) + per_level_mask.append(active_b1cHcW) + + active_b1hw = per_level_mask[ + -1 + ] # the mask map at the deepest scale (the smallest feature map), which would be used for masking the input to the encoder; shape: (B, 1, f, f) + masked_bchw = inp_bchw * active_b1hw + return SparKMaskingOuptut( + masked_bchw=masked_bchw, per_level_mask=per_level_mask + ) + + class SparK(nn.Module): def __init__( self, @@ -576,6 +626,8 @@ def __init__( self.sbn = sbn self.densify_norm_str = densify_norm.lower() + self.masker = SparKMasker(sparse_encoder=sparse_encoder, mask_ratio=mask_ratio) + self.densifier = SparKDensifier( encoder_in_channels=self.sparse_encoder.enc_feat_map_chs, decoder_in_channel=self.dense_decoder.width, @@ -589,36 +641,16 @@ def __init__( # these are deprecated and would never be used; can be removed. self.vis_active = self.vis_active_ex = self.vis_inp = self.vis_inp_mask = ... - def mask(self, B: int, device: torch.device) -> torch.Tensor: - h, w = self.sparse_encoder.fmap_h, self.sparse_encoder.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: torch.Tensor, - active_b1ff: None | torch.Tensor = None, vis=False, ): # step1. Mask - active_b1ff = active_b1ff or self.mask( - inp_bchw.shape[0], inp_bchw.device - ) # (B, 1, f, f) - global _cur_active - _cur_active = active_b1ff # (B, 1, f, f) - ds = self.sparse_encoder.downsample_ratio - active_b1hw = active_b1ff.repeat_interleave(ds, 2).repeat_interleave( - ds, 3 - ) # (B, 1, H, W) - masked_bchw = inp_bchw * active_b1hw - + mask_out: SparKMaskingOuptut = 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) fea_bcffs: List[torch.Tensor] = self.sparse_encoder(masked_bchw) # step3. Densify: get hierarchical dense features for decoding @@ -637,7 +669,7 @@ def forward( ) # (B, L, C) ==mean==> (B, L) non_active = ( - active_b1ff.logical_not().int().view(active_b1ff.shape[0], -1) + active_b1fHfW.logical_not().int().view(active_b1fHfW.shape[0], -1) ) # (B, 1, f, f) => (B, L) recon_loss = l2_loss.mul_(non_active).sum() / ( non_active.sum() + 1e-8 From 138380f0d1aefa1ee1ab9555482049c7af046c17 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 10:44:59 -0300 Subject: [PATCH 28/85] refactor: removing unused variables --- lightly/models/modules/sparse_spark.py | 21 ++++----------------- 1 file changed, 4 insertions(+), 17 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 6db7b1f27..9de7b41ab 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -609,38 +609,25 @@ def __init__( self, sparse_encoder: SparseEncoder, dense_decoder: LightDecoder, - mask_ratio=0.6, - densify_norm="bn", + mask_ratio: float = 0.6, + densify_norm: str = "bn", sbn=False, ): super().__init__() # spatial and size info moved to SparseEncoder - self.mask_ratio = mask_ratio - self.len_keep = round( - sparse_encoder.fmap_h * sparse_encoder.fmap_w * (1 - mask_ratio) - ) - self.sparse_encoder = sparse_encoder self.dense_decoder = dense_decoder - - self.sbn = sbn - self.densify_norm_str = densify_norm.lower() - self.masker = SparKMasker(sparse_encoder=sparse_encoder, mask_ratio=mask_ratio) - self.densifier = SparKDensifier( encoder_in_channels=self.sparse_encoder.enc_feat_map_chs, decoder_in_channel=self.dense_decoder.width, - densify_norm_str=self.densify_norm_str, - sbn=self.sbn, + densify_norm_str=densify_norm.lower(), + sbn=sbn, ) print( f"[SparK.__init__] dims of mask_tokens={tuple(b.mask_token.numel() for b in self.densifier.blocks)}" ) - # these are deprecated and would never be used; can be removed. - self.vis_active = self.vis_active_ex = self.vis_inp = self.vis_inp_mask = ... - def forward( self, inp_bchw: torch.Tensor, From 6b3dbdf7833e010799ed59b0ef0701dcd133e4f7 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 10:47:23 -0300 Subject: [PATCH 29/85] refactor: removing unecessary module dependency --- lightly/models/modules/sparse_spark.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 9de7b41ab..ce8d8b6e4 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -565,13 +565,19 @@ class SparKMaskingOuptut(NamedTuple): class SparKMasker(nn.Module): - def __init__(self, sparse_encoder: SparseEncoder, mask_ratio: float = 0.6): + def __init__( + self, + feature_map_size: tuple[int, int], + downsample_ratio: int, + mask_ratio: float = 0.6, + ) -> None: super().__init__() - self.sparse_encoder = sparse_encoder + 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) -> torch.Tensor: - h, w = self.sparse_encoder.fmap_h, self.sparse_encoder.fmap_w + h, w = self.fmap_h, self.fmap_w index_keep, _ = random_token_mask( size=(B, h * w), mask_ratio=self.mask_ratio, device=device ) @@ -586,7 +592,7 @@ def forward(self, inp_bchw: torch.Tensor) -> SparKMaskingOuptut: global _cur_active _cur_active = self.mask(inp_bchw.shape[0], inp_bchw.device) # (B, 1, f, f) active_b1ff = _cur_active.clone() - downsample_ratio = self.sparse_encoder.downsample_ratio + downsample_ratio = self.downsample_ratio per_level_mask = [active_b1ff] for i in range(int(math.log2(downsample_ratio))): previous_mask = per_level_mask[-1] @@ -617,7 +623,11 @@ def __init__( # spatial and size info moved to SparseEncoder self.sparse_encoder = sparse_encoder self.dense_decoder = dense_decoder - self.masker = SparKMasker(sparse_encoder=sparse_encoder, mask_ratio=mask_ratio) + self.masker = SparKMasker( + feature_map_size=(self.sparse_encoder.fmap_h, self.sparse_encoder.fmap_w), + downsample_ratio=self.sparse_encoder.downsample_ratio, + mask_ratio=mask_ratio, + ) self.densifier = SparKDensifier( encoder_in_channels=self.sparse_encoder.enc_feat_map_chs, decoder_in_channel=self.dense_decoder.width, From ed8691ab7117f3afd67350289ec05127ad0cfd17 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 10:55:15 -0300 Subject: [PATCH 30/85] refactor: loss as module --- lightly/models/modules/sparse_spark.py | 63 ++++++++++++++++++++------ 1 file changed, 50 insertions(+), 13 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index ce8d8b6e4..73635d481 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -610,6 +610,52 @@ def forward(self, inp_bchw: torch.Tensor) -> SparKMaskingOuptut: ) +class SparKPatchReconLoss(nn.Module): + """Loss module that computes per-patch normalized reconstruction loss. + + Accepts the non-flattened mask `(B, 1, f, f)` and flattens internally. + """ + + def __init__(self, eps: float = 1e-6): + super().__init__() + self.eps = eps + + def forward( + self, + inp_patches: torch.Tensor, + rec_patches: torch.Tensor, + active_mask: torch.Tensor, + ): + """Compute reconstruction loss and return per-patch stats. + + Args: + inp_patches: (B, L, N) original patches + rec_patches: (B, L, N) reconstructed patches + active_mask: (B, 1, f, f) boolean mask (required) + + Returns: + recon_loss: scalar tensor + mean: (B, L, 1) per-patch mean + var: (B, L, 1) per-patch std + """ + 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_norm = (inp_patches - mean) / var + + l2_loss = ((rec_patches - inp_norm) ** 2).mean(dim=2) # (B, L) + + non_active = active_mask.logical_not().int().view(active_mask.shape[0], -1) + + recon_loss = l2_loss.mul_(non_active).sum() / (non_active.sum() + 1e-8) + return recon_loss, mean, var + + class SparK(nn.Module): def __init__( self, @@ -637,6 +683,8 @@ def __init__( print( f"[SparK.__init__] dims of mask_tokens={tuple(b.mask_token.numel() for b in self.densifier.blocks)}" ) + # loss module for patch reconstruction + self.recon_loss_fn = SparKPatchReconLoss() def forward( self, @@ -658,19 +706,8 @@ def forward( self.patchify(inp_bchw), self.patchify(rec_bchw), ) # inp and rec: (B, L = f*f, N = C*downsample_raito**2) - mean = inp.mean(dim=-1, keepdim=True) - var = (inp.var(dim=-1, keepdim=True) + 1e-6) ** 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_b1fHfW.logical_not().int().view(active_b1fHfW.shape[0], -1) - ) # (B, 1, f, f) => (B, L) - recon_loss = l2_loss.mul_(non_active).sum() / ( - non_active.sum() + 1e-8 - ) # loss only on masked (non-active) patches + + recon_loss, mean, var = self.recon_loss_fn(inp, rec, active_b1fHfW) if vis: masked_bchw = inp_bchw * active_b1hw From 3dcd5c32d03b6f1f942ebe6fd7c7c0decc0e721d Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 10:58:54 -0300 Subject: [PATCH 31/85] refactor: spark visualization decoding logic as module --- lightly/models/modules/sparse_spark.py | 56 ++++++++++++++++++++++++-- 1 file changed, 52 insertions(+), 4 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 73635d481..cf77a6d00 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -656,6 +656,51 @@ def forward( return recon_loss, mean, var +class SparKOutputDecoder(nn.Module): + """Handles de-normalizing reconstructed patches and producing visualization tensors. + + Usage: call with `rec_patches, mean, var, inp_bchw, active_mask_full` where + `rec_patches` is (B, L, N), `mean`/`var` are (B, L, 1), `inp_bchw` is original image + and `active_mask_full` is (B, 1, H, W) boolean mask indicating visible patches. + The decoder is configured with only the minimal spatial properties: `fmap_h`, + `fmap_w` and `downsample_ratio` (no encoder object required). + """ + + def __init__(self, fmap_h: int, fmap_w: int, downsample_ratio: int): + super().__init__() + self.fmap_h = fmap_h + self.fmap_w = fmap_w + self.downsample_ratio = downsample_ratio + + def unpatchify(self, bln: torch.Tensor) -> torch.Tensor: + 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: torch.Tensor, + mean: torch.Tensor, + var: torch.Tensor, + inp_bchw: torch.Tensor, + active_mask_full: torch.Tensor, + ): + # de-normalize and unpatchify + rec_bchw = self.unpatchify(rec_patches * var + mean) + + # masked input at full resolution + masked_bchw = inp_bchw * active_mask_full + + # combine: use original where visible, reconstructed where masked + rec_or_inp = torch.where(active_mask_full, inp_bchw, rec_bchw) + + return inp_bchw, masked_bchw, rec_or_inp + + class SparK(nn.Module): def __init__( self, @@ -685,6 +730,12 @@ def __init__( ) # loss module for patch reconstruction self.recon_loss_fn = SparKPatchReconLoss() + # output decoder for visualization (pass minimal spatial props) + self.output_decoder = SparKOutputDecoder( + self.sparse_encoder.fmap_h, + self.sparse_encoder.fmap_w, + self.sparse_encoder.downsample_ratio, + ) def forward( self, @@ -710,10 +761,7 @@ def forward( recon_loss, mean, var = self.recon_loss_fn(inp, rec, active_b1fHfW) if vis: - masked_bchw = inp_bchw * active_b1hw - rec_bchw = self.unpatchify(rec * var + mean) - rec_or_inp = torch.where(active_b1hw, inp_bchw, rec_bchw) - return inp_bchw, masked_bchw, rec_or_inp + return self.output_decoder(rec, mean, var, inp_bchw, active_b1hw) else: return recon_loss From 1c316da2bad93d7f2bab8e3a75b82497b615ac8a Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 10:59:19 -0300 Subject: [PATCH 32/85] refactor: remove unused --- lightly/models/modules/sparse_spark.py | 9 --------- 1 file changed, 9 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index cf77a6d00..ddee9fccf 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -769,15 +769,6 @@ def patchify(self, bchw): p = self.sparse_encoder.downsample_ratio return patchify(bchw, p) # (B, L=f*f, N=C*p*p) - def unpatchify(self, bln): - p = self.sparse_encoder.downsample_ratio - h, w = self.sparse_encoder.fmap_h, self.sparse_encoder.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 __repr__(self): return ( f"\n" From 3ba183902e75cbda021e430cc0238c91a69ed501 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 11:07:15 -0300 Subject: [PATCH 33/85] refactor --- lightly/models/modules/sparse_spark.py | 14 ++++---------- 1 file changed, 4 insertions(+), 10 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index ddee9fccf..7319bd898 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -631,7 +631,7 @@ def forward( Args: inp_patches: (B, L, N) original patches rec_patches: (B, L, N) reconstructed patches - active_mask: (B, 1, f, f) boolean mask (required) + active_mask: (B, L, f, f) boolean mask (required) Returns: recon_loss: scalar tensor @@ -725,9 +725,7 @@ def __init__( densify_norm_str=densify_norm.lower(), sbn=sbn, ) - print( - f"[SparK.__init__] dims of mask_tokens={tuple(b.mask_token.numel() for b in self.densifier.blocks)}" - ) + self.downsample_ratio = self.sparse_encoder.downsample_ratio # loss module for patch reconstruction self.recon_loss_fn = SparKPatchReconLoss() # output decoder for visualization (pass minimal spatial props) @@ -754,8 +752,8 @@ def forward( # step4. Decode and reconstruct rec_bchw = self.dense_decoder(to_dec) inp, rec = ( - self.patchify(inp_bchw), - self.patchify(rec_bchw), + patchify(inp_bchw, self.downsample_ratio), + patchify(rec_bchw, self.downsample_ratio), ) # inp and rec: (B, L = f*f, N = C*downsample_raito**2) recon_loss, mean, var = self.recon_loss_fn(inp, rec, active_b1fHfW) @@ -765,10 +763,6 @@ def forward( else: return recon_loss - def patchify(self, bchw): - p = self.sparse_encoder.downsample_ratio - return patchify(bchw, p) # (B, L=f*f, N=C*p*p) - def __repr__(self): return ( f"\n" From 70d32dd55abbcd6d0eb6d6b8d39993fb69b27ca3 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 11:14:10 -0300 Subject: [PATCH 34/85] refactor: removed big module and refactored timm funcs --- lightly/models/modules/sparse_spark.py | 130 +++---------------------- 1 file changed, 11 insertions(+), 119 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 7319bd898..ace5b5227 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -211,14 +211,6 @@ def __repr__(self): return super().__repr__()[:-1] + f", sp={self.sparse})" -def get_downsample_ratio_from_timm_model(model: nn.Module) -> int: - return model.feature_info[-1]["reduction"] - - -def get_enc_feat_map_chs_from_timm_model(model: nn.Module) -> List[int]: - return [fi["num_chs"] for fi in model.feature_info] - - class SparseEncoder(nn.Module): """ Converts a dense CNN model to a sparse CNN model by replacing standard layers @@ -231,15 +223,23 @@ class SparseEncoder(nn.Module): enc_feat_map_chs: List[int] - def __init__(self, cnn, input_size, sbn=False, verbose=False): + def __init__( + self, + cnn, + input_size, + downsample_ratio: int, + feature_map_channels: list[int], + sbn=False, + verbose=False, + ): super().__init__() self.sp_cnn = SparseEncoder.dense_model_to_sparse( m=cnn, verbose=verbose, sbn=sbn ) self.input_size, self.downsample_ratio, self.enc_feat_map_chs = ( input_size, - get_downsample_ratio_from_timm_model(cnn), - get_enc_feat_map_chs_from_timm_model(cnn), + downsample_ratio, + feature_map_channels, ) # feature-map spatial size (height, width) at encoder output (patch grid) self.fmap_h = input_size // self.downsample_ratio @@ -699,111 +699,3 @@ def forward( rec_or_inp = torch.where(active_mask_full, inp_bchw, rec_bchw) return inp_bchw, masked_bchw, rec_or_inp - - -class SparK(nn.Module): - def __init__( - self, - sparse_encoder: SparseEncoder, - dense_decoder: LightDecoder, - mask_ratio: float = 0.6, - densify_norm: str = "bn", - sbn=False, - ): - super().__init__() - # spatial and size info moved to SparseEncoder - self.sparse_encoder = sparse_encoder - self.dense_decoder = dense_decoder - self.masker = SparKMasker( - feature_map_size=(self.sparse_encoder.fmap_h, self.sparse_encoder.fmap_w), - downsample_ratio=self.sparse_encoder.downsample_ratio, - mask_ratio=mask_ratio, - ) - self.densifier = SparKDensifier( - encoder_in_channels=self.sparse_encoder.enc_feat_map_chs, - decoder_in_channel=self.dense_decoder.width, - densify_norm_str=densify_norm.lower(), - sbn=sbn, - ) - self.downsample_ratio = self.sparse_encoder.downsample_ratio - # loss module for patch reconstruction - self.recon_loss_fn = SparKPatchReconLoss() - # output decoder for visualization (pass minimal spatial props) - self.output_decoder = SparKOutputDecoder( - self.sparse_encoder.fmap_h, - self.sparse_encoder.fmap_w, - self.sparse_encoder.downsample_ratio, - ) - - def forward( - self, - inp_bchw: torch.Tensor, - vis=False, - ): - # step1. Mask - mask_out: SparKMaskingOuptut = 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) - fea_bcffs: List[torch.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 = ( - patchify(inp_bchw, self.downsample_ratio), - patchify(rec_bchw, self.downsample_ratio), - ) # inp and rec: (B, L = f*f, N = C*downsample_raito**2) - - recon_loss, mean, var = self.recon_loss_fn(inp, rec, active_b1fHfW) - - if vis: - return self.output_decoder(rec, mean, var, inp_bchw, active_b1hw) - else: - return recon_loss - - def __repr__(self): - return ( - f"\n" - f"[SparK.config]: {pformat(self.get_config(), indent=2, width=250)}\n" - f"[SparK.structure]: {super().__repr__().replace(SparK.__name__, '')}" - ) - - def get_config(self): - return { - # self - "mask_ratio": self.mask_ratio, - "densify_norm_str": self.densify_norm_str, - "sbn": self.sbn, - # enc - "sparse_encoder.input_size": self.sparse_encoder.input_size, - # dec - "dense_decoder.width": self.dense_decoder.width, - } - - def state_dict( - self, destination=None, prefix="", keep_vars=False, with_config=False - ) -> dict[str, torch.Tensor]: - state = super().state_dict( - destination=destination, prefix=prefix, keep_vars=keep_vars - ) - if with_config: - state["config"] = self.get_config() - return state - - def load_state_dict( - self, state_dict: dict[str, torch.Tensor], strict=True - ) -> dict[str, torch.Tensor]: - config: dict = state_dict.pop("config", None) - incompatible_keys = super().load_state_dict(state_dict, strict=strict) - if config is not None: - for k, v in self.get_config().items(): - ckpt_v = config.get(k, None) - if ckpt_v != v: - err = f"[SparseMIM.load_state_dict] config mismatch: this.{k}={v} (ckpt.{k}={ckpt_v})" - if strict: - raise AttributeError(err) - else: - print(err, file=sys.stderr) - return incompatible_keys From 2f73d15dd53ac7d04ddae07d84d68e8b5cfb57e7 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 11:17:49 -0300 Subject: [PATCH 35/85] feat: example script --- examples/pytorch_lightning/spark.py | 113 ++++++++++++++++++++++++---- 1 file changed, 97 insertions(+), 16 deletions(-) diff --git a/examples/pytorch_lightning/spark.py b/examples/pytorch_lightning/spark.py index cf3864973..11ef7b0a3 100644 --- a/examples/pytorch_lightning/spark.py +++ b/examples/pytorch_lightning/spark.py @@ -9,33 +9,103 @@ import timm import torch import torchvision -from pytorch_lightning.callbacks import RichProgressBar +from pytorch_lightning.callbacks import ModelCheckpoint, RichProgressBar +from torch import nn from torchvision.transforms import v2 ## The global projection head is the same as the Barlow Twins one -from lightly.models.modules.sparse_spark import LightDecoder, SparK, SparseEncoder +from lightly.models.modules.sparse_spark import ( + LightDecoder, + SparKDensifier, + SparKMasker, + SparKMaskingOuptut, + SparKOutputDecoder, + SparKPatchReconLoss, + SparseEncoder, +) +from lightly.models.utils import patchify + + +def get_downsample_ratio_from_timm_model(model: nn.Module) -> int: + return model.feature_info[-1]["reduction"] -class SparseSpark(pl.LightningModule): - def __init__(self) -> None: +def get_enc_feat_map_chs_from_timm_model(model: nn.Module) -> list[int]: + return [fi["num_chs"] for fi in model.feature_info] + + +class SparseSparK(pl.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( "resnet18", drop_path_rate=0.05, features_only=True ) self.sparse_encoder = SparseEncoder( - backbone, input_size=416, sbn=False, verbose=True + backbone, + downsample_ratio=get_downsample_ratio_from_timm_model(backbone), + feature_map_channels=get_enc_feat_map_chs_from_timm_model(backbone), + input_size=input_size, + sbn=sbn, + verbose=True, ) self.dense_decoder = LightDecoder( self.sparse_encoder.downsample_ratio, width=self.sparse_encoder.enc_feat_map_chs[-1], ) - self.spark = SparK( - sparse_encoder=self.sparse_encoder, - dense_decoder=self.dense_decoder, + self.masker = SparKMasker( + feature_map_size=(self.sparse_encoder.fmap_h, self.sparse_encoder.fmap_w), + downsample_ratio=self.sparse_encoder.downsample_ratio, + mask_ratio=mask_ratio, + ) + self.densifier = SparKDensifier( + encoder_in_channels=self.sparse_encoder.enc_feat_map_chs, + decoder_in_channel=self.dense_decoder.width, + densify_norm_str=densify_norm.lower(), + sbn=sbn, + ) + self.downsample_ratio = self.sparse_encoder.downsample_ratio + # loss module for patch reconstruction + self.recon_loss_fn = SparKPatchReconLoss() + # output decoder for visualization (pass minimal spatial props) + self.output_decoder = SparKOutputDecoder( + self.sparse_encoder.fmap_h, + self.sparse_encoder.fmap_w, + self.sparse_encoder.downsample_ratio, ) - def forward(self, x): - return self.spark(x) + def forward( + self, + inp_bchw: torch.Tensor, + vis=False, + ): + # step1. Mask + mask_out: SparKMaskingOuptut = 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) + fea_bcffs: list[torch.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 = ( + patchify(inp_bchw, self.downsample_ratio), + patchify(rec_bchw, self.downsample_ratio), + ) # inp and rec: (B, L = f*f, N = C*downsample_raito**2) + + recon_loss, mean, var = self.recon_loss_fn(inp, rec, active_b1fHfW) + + if vis: + return self.output_decoder(rec, mean, var, inp_bchw, active_b1hw) + else: + return recon_loss def training_step(self, batch, batch_index) -> torch.Tensor: img, target = batch @@ -52,11 +122,12 @@ def training_step(self, batch, batch_index) -> torch.Tensor: return recon_loss def configure_optimizers(self): - optim = torch.optim.SGD(self.parameters(), lr=0.01, momentum=0.9) - return optim + return torch.optim.SGD( + self.parameters(), lr=0.03, momentum=0.9, weight_decay=1e-4 + ) -model = SparseSpark() +model = SparseSparK(input_size=416) # we ignore object detection annotations by setting target_transform to return 0 @@ -72,6 +143,7 @@ def target_transform(t): v2.Resize((416, 416)), v2.RGB(), v2.ToTensor(), + v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ] ), target_transform=target_transform, @@ -84,12 +156,21 @@ def target_transform(t): batch_size=4, shuffle=True, drop_last=True, - num_workers=0, + num_workers=8, ) + accelerator = "gpu" if torch.cuda.is_available() else "cpu" trainer = pl.Trainer( - max_epochs=10, devices=1, accelerator=accelerator, callbacks=[RichProgressBar()] + max_epochs=30, + devices=1, + accelerator=accelerator, + callbacks=[ + RichProgressBar(), + ], +) +trainer.fit( + model=model, + train_dataloaders=dataloader, ) -trainer.fit(model=model, train_dataloaders=dataloader) From b48d4eb4db1d78b263bbfe24284e4eb1b40af243 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 9 Feb 2026 11:25:21 -0300 Subject: [PATCH 36/85] refactor: removed unused code and added opyrights --- lightly/models/modules/sparse_spark.py | 45 ++------------------------ 1 file changed, 3 insertions(+), 42 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index ace5b5227..0b0a57122 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -1,15 +1,14 @@ +# Copyright (c) 2023 Keyu Tian +# Copyright (c) ByteDance, Inc. and its affiliates. import math -import sys -from pprint import pformat from typing import List, NamedTuple -import timm import torch import torch.nn as nn from timm.models.layers import DropPath, trunc_normal_ from torch.nn.common_types import _size_2_t -from lightly.models.utils import patchify, random_token_mask +from lightly.models.utils import random_token_mask def is_pow2n(x): @@ -321,13 +320,6 @@ def forward(self, x): return self.sp_cnn(x) -# Copyright (c) ByteDance, Inc. and its affiliates. -# All rights reserved. -# -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - - class UNetBlock(nn.Module): def __init__(self, cin, cout, bn2d): """ @@ -403,37 +395,6 @@ def initialize(self) -> None: nn.init.constant_(m.bias, 0) nn.init.constant_(m.weight, 1.0) - # Copyright (c) ByteDance, Inc. and its affiliates. - - -# All rights reserved. -# -# This source code is licensed under the license found in the -# LICENSE file in the root directory of this source tree. - -pretrain_default_model_kwargs = { - "your_convnet": dict(), - "resnet50": dict(drop_path_rate=0.05), - "resnet101": dict(drop_path_rate=0.08), - "resnet152": dict(drop_path_rate=0.10), - "resnet200": dict(drop_path_rate=0.15), - "convnext_small": dict(sparse=True, drop_path_rate=0.2), - "convnext_base": dict(sparse=True, drop_path_rate=0.3), - "convnext_large": dict(sparse=True, drop_path_rate=0.4), -} - - -def build_sparse_encoder( - name: str, input_size: int, sbn=False, drop_path_rate=0.0, verbose=False -): - kwargs = pretrain_default_model_kwargs[name] - if drop_path_rate != 0: - kwargs["drop_path_rate"] = drop_path_rate - print(f"[build_sparse_encoder] model kwargs={kwargs}") - cnn = timm.create_model(name, **kwargs) - - return SparseEncoder(cnn, input_size=input_size, sbn=sbn, verbose=verbose) - class SparKDensfiyBlock(nn.Module): def __init__( From 25687d3c0165a3c8e921ddffc67ce2056e788699 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 12 Feb 2026 08:38:20 -0300 Subject: [PATCH 37/85] doc: improved documentation and type hinting --- lightly/models/modules/sparse_spark.py | 532 +++++++++++++++++++------ 1 file changed, 406 insertions(+), 126 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 0b0a57122..4803ad0a1 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -1,7 +1,7 @@ # Copyright (c) 2023 Keyu Tian # Copyright (c) ByteDance, Inc. and its affiliates. import math -from typing import List, NamedTuple +from typing import List, NamedTuple, Optional, Tuple, Union import torch import torch.nn as nn @@ -11,11 +11,30 @@ from lightly.models.utils import random_token_mask -def is_pow2n(x): +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, ...]) -> _size_2_t: + """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 elif len(t) == 1: @@ -24,11 +43,35 @@ def coalesce_to_size_2_t(t: tuple[int, ...]) -> _size_2_t: raise ValueError(f"Invalid tuple length: {len(t)}; expected 1 or 2.") -_cur_active: torch.Tensor = None # B1ff +_cur_active: Optional[torch.Tensor] = ( + None # B1ff - Global active/mask tensor tracked during forward passes +) + +def _get_active_ex_or_ii( + H: int, W: int, returning_active_ex: bool = True +) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: + """Get active indices or expanded active mask from global _cur_active. -# todo: try to use `gather` for speed? -def _get_active_ex_or_ii(H: int, W: int, returning_active_ex=True): + 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). + """ + assert _cur_active is not None, ( + "_cur_active must be set before calling this function" + ) h_repeat, w_repeat = H // _cur_active.shape[-2], W // _cur_active.shape[-1] active_ex = _cur_active.repeat_interleave(h_repeat, dim=2).repeat_interleave( w_repeat, dim=3 @@ -37,27 +80,46 @@ def _get_active_ex_or_ii(H: int, W: int, returning_active_ex=True): active_ex if returning_active_ex else active_ex.squeeze(1).nonzero(as_tuple=True) - ) # ii: bi, hi, wi + ) + +def sp_conv_forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass for sparse convolution/pooling layers. -def sp_conv_forward(self, x: torch.Tensor): + 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. + x: Input tensor of shape (B, C, H, W). + + Returns: + Masked output tensor of same shape as input, with inactive spatial positions zeroed. + """ x = super(type(self), self).forward(x) - x *= _get_active_ex_or_ii( - H=x.shape[2], W=x.shape[3], returning_active_ex=True - ) # (BCHW) *= (B1HW), mask the output of conv + x *= _get_active_ex_or_ii(H=x.shape[2], W=x.shape[3], returning_active_ex=True) return x -def sp_bn_forward(self, x: torch.Tensor): +def sp_bn_forward(self, x: torch.Tensor) -> torch.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. + x: 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=x.shape[2], W=x.shape[3], returning_active_ex=False) bhwc = x.permute(0, 2, 3, 1) - nc = bhwc[ - ii - ] # select the features on non-masked positions to form a flatten feature `nc` - nc = super(type(self), self).forward( - nc - ) # use BN1d to normalize this flatten feature `nc` + nc = bhwc[ii] + nc = super(type(self), self).forward(nc) bchw = torch.zeros_like(bhwc) bchw[ii] = nc @@ -66,42 +128,83 @@ def sp_bn_forward(self, x: torch.Tensor): class SparseConv2d(nn.Conv2d): - forward = sp_conv_forward # hack: override the forward function; see `sp_conv_forward` above for more details + """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): - forward = sp_conv_forward # hack: override the forward function; see `sp_conv_forward` above for more details + """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): - forward = sp_conv_forward # hack: override the forward function; see `sp_conv_forward` above for more details + """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): - forward = sp_bn_forward # hack: override the forward function; see `sp_bn_forward` above for more details + """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): - forward = sp_bn_forward # hack: override the forward function; see `sp_bn_forward` above for more details + """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). + + When sparse=True, only applies normalization to active (unmasked) spatial positions, + using indices from the global _cur_active mask. """ def __init__( - self, normalized_shape, eps=1e-6, data_format="channels_last", sparse=True - ): + self, + normalized_shape: int, + eps: float = 1e-6, + data_format: str = "channels_last", + sparse: bool = True, + ) -> 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 self.sparse = sparse - def forward(self, input: torch.Tensor): + def forward(self, input: torch.Tensor) -> torch.Tensor: if input.ndim == 4: # BHWC or BCHW if self.data_format == "channels_last": # BHWC if self.sparse: @@ -150,28 +253,34 @@ def __repr__(self): class SparseConvNeXtBlock(nn.Module): - r"""ConvNeXt Block. There are two equivalent implementations: + r"""ConvNeXt Block with optional 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 + + We use (2) as we find it slightly faster in PyTorch. When sparse=True, applies masking to the output. Args: - dim (int): Number of input channels. - drop_path (float): Stochastic depth rate. Default: 0.0 - layer_scale_init_value (float): Init value for Layer Scale. Default: 1e-6. + 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. + sparse: Whether to apply sparse masking. Default: True. + ks: Kernel size for depthwise convolution. Default: 7. """ def __init__( - self, dim, drop_path=0.0, layer_scale_init_value=1e-6, sparse=True, ks=7 - ): + self, + dim: int, + drop_path: float = 0.0, + layer_scale_init_value: float = 1e-6, + sparse: bool = True, + ks: int = 7, + ) -> None: super().__init__() - self.dwconv = nn.Conv2d( - dim, dim, kernel_size=ks, padding=ks // 2, groups=dim - ) # depthwise conv + self.dwconv = nn.Conv2d(dim, dim, kernel_size=ks, padding=ks // 2, groups=dim) self.norm = SparseConvNeXtLayerNorm(dim, eps=1e-6, sparse=sparse) - self.pwconv1 = nn.Linear( - dim, 4 * dim - ) # pointwise/1x1 convs, implemented with linear layers + self.pwconv1 = nn.Linear(dim, 4 * dim) self.act = nn.GELU() self.pwconv2 = nn.Linear(4 * dim, dim) self.gamma = ( @@ -184,19 +293,17 @@ def __init__( ) self.sparse = sparse - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: input = x x = self.dwconv(x) - x = x.permute(0, 2, 3, 1) # (N, C, H, W) -> (N, H, W, C) + x = x.permute(0, 2, 3, 1) x = self.norm(x) x = self.pwconv1(x) - x = self.act( - x - ) # GELU(0) == (0), so there is no need to mask x (no need to `x *= _get_active_ex_or_ii`) + 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) # (N, H, W, C) -> (N, C, H, W) + x = x.permute(0, 3, 1, 2) if self.sparse: x *= _get_active_ex_or_ii( @@ -211,26 +318,31 @@ def __repr__(self): class SparseEncoder(nn.Module): - """ - Converts a dense CNN model to a sparse CNN model by replacing standard layers - - Attributes: - enc_feat_map_chs: List[int]: list of channel numbers of feature maps at different scales, in the order from shallow to deep + """Converts a dense CNN model to a sparse CNN model by replacing standard layers. + Recursively traverses a model and replaces standard layers (Conv2d, Pooling, BatchNorm, LayerNorm) + with their sparse counterparts. Sparse layers respect a global active mask (_cur_active) that + indicates which spatial regions should be processed. + Attributes: + enc_feat_map_chs: List of channel numbers of feature maps at different scales, shallow to deep. + input_size: Original input image size (assumed square). + downsample_ratio: Total spatial downsampling ratio from input to deepest feature map. + fmap_h: Height of feature map at encoder output (patch grid). + fmap_w: Width of feature map at encoder output (patch grid). """ enc_feat_map_chs: List[int] def __init__( self, - cnn, - input_size, + cnn: nn.Module, + input_size: int, downsample_ratio: int, feature_map_channels: list[int], - sbn=False, - verbose=False, - ): + sbn: bool = False, + verbose: bool = False, + ) -> None: super().__init__() self.sp_cnn = SparseEncoder.dense_model_to_sparse( m=cnn, verbose=verbose, sbn=sbn @@ -240,12 +352,29 @@ def __init__( downsample_ratio, feature_map_channels, ) - # feature-map spatial size (height, width) at encoder output (patch grid) self.fmap_h = input_size // self.downsample_ratio self.fmap_w = input_size // self.downsample_ratio @staticmethod - def dense_model_to_sparse(m: nn.Module, verbose=False, sbn=False): + 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. + + 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 = m if isinstance(m, nn.Conv2d): bias = m.bias is not None @@ -316,15 +445,31 @@ def dense_model_to_sparse(m: nn.Module, verbose=False, sbn=False): del m return oup - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.Tensor: + """Forward pass through the sparse CNN encoder. + + Args: + x: Input image tensor. + + Returns: + Output feature tensor from the sparse CNN. + """ return self.sp_cnn(x) class UNetBlock(nn.Module): - def __init__(self, cin, cout, bn2d): - """ - a UNet block with 2x up sampling - """ + """U-Net upsampling block with 2x spatial upsampling and conv refinement. + + 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 @@ -337,22 +482,40 @@ def __init__(self, cin, cout, bn2d): bn2d(cout), ) - def forward(self, x): + def forward(self, x: torch.Tensor) -> torch.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) class LightDecoder(nn.Module): + """Lightweight hierarchical decoder for feature map reconstruction. + + 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, width=768, sbn=True - ): # todo: the decoder's width follows a simple halfing rule; you can change it to any other rule + 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) - ] # todo: the decoder's width follows a simple halfing rule; you can change it to any other rule + channels = [self.width // 2**i for i in range(n + 1)] bn2d = nn.SyncBatchNorm if sbn else nn.BatchNorm2d self.dec = nn.ModuleList( [ @@ -364,9 +527,17 @@ def __init__( self.initialize() - def forward(self, to_dec: List[torch.Tensor]): + def forward(self, to_dec: List[torch.Tensor]) -> torch.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 = 0 - for i, d in enumerate(self.dec): + 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) @@ -376,6 +547,13 @@ 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) + """ for m in self.modules(): if isinstance(m, nn.Linear): trunc_normal_(m.weight, std=0.02) @@ -397,6 +575,20 @@ def initialize(self) -> None: class SparKDensfiyBlock(nn.Module): + """Block for densifying sparse features by filling masked regions with learned tokens. + + 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, @@ -405,7 +597,7 @@ def __init__( sbn: bool = False, use_identity_proj: bool = False, kernel_size: int = 3, - ): + ) -> None: super().__init__() # mask token p = nn.Parameter(torch.zeros(1, e_width, 1, 1)) @@ -444,7 +636,18 @@ def __init__( ) self.densify_proj = densify_proj - def forward(self, bcff: torch.Tensor, cur_active: torch.BoolTensor): + def forward( + self, bcff: torch.Tensor, cur_active: torch.Tensor + ) -> Optional[torch.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, or None if input is None. + """ if bcff is None: return None bcff = self.densify_norm(bcff) @@ -455,14 +658,16 @@ def forward(self, bcff: torch.Tensor, cur_active: torch.BoolTensor): class SparKDensifier(nn.Module): - """ - A stack of SparKDensfiyBlocks to convert hierarchical sparse features to hierarchical dense features for decoding + """Stack of densify blocks to convert sparse hierarchical features to dense features. + + 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 of feature maps at different scales from the encoder, in the order from shallow to deep - decoder_in_channel: channel number of the feature map at the deepest scale for decoding - densify_norm_str: the type of normalization inside SparKDensfiyBlock; can be "bn", "ln" or "none" - sbn: whether to use SyncBatchNorm (True) or BatchNorm (False) if densify_norm_str is "bn" + 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__( @@ -471,18 +676,15 @@ def __init__( 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]): - # from the smallest feat map to the largest; i=0: the last feat map; i=1: the second last feat map ... - # fork arguments that depend on the position (previously used idx inside the block) use_identity = i == 0 and e_width == d_width kernel_size = 1 if i <= 0 else 3 - # build densify block and append block = SparKDensfiyBlock( e_width=e_width, d_width=d_width, @@ -492,40 +694,56 @@ def __init__( kernel_size=kernel_size, ) self.blocks.append(block) - # todo: the decoder's width follows a simple halfing rule; you can change it to any other rule d_width //= 2 - def forward(self, fea_bcffs: List[torch.Tensor]): - """ + def forward(self, fea_bcffs: List[torch.Tensor]) -> List[Optional[torch.Tensor]]: + """Convert sparse features to dense by filling masked regions. + Args: - fea_bcffs: a list of feature maps at different scales from the encoder, in the order from shallow to deep; each feature map is a tensor of shape (B, C, f, f) + 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] # from the smallest feat map to the largest + fea_bcffs = fea_bcffs[::-1] global _cur_active - active_fmap_current = ( - _cur_active # (B, 1, f, f), the mask map at the current scale - ) - for i, bcff in enumerate( - fea_bcffs - ): # from the smallest feature map to the largest + active_fmap_current = _cur_active + for i, bcff in enumerate(fea_bcffs): if bcff is not None: 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 - ) # dilate the mask map, from (B, 1, f, f) to (B, 1, H, W) + ).repeat_interleave(2, dim=3) return to_dec class SparKMaskingOuptut(NamedTuple): + """Output container for SparKMasker forward pass. + + Attributes: + masked_bchw: Input image with masks applied at full resolution. + per_level_mask: List of binary masks at each hierarchical level. + """ + masked_bchw: torch.Tensor per_level_mask: List[torch.Tensor] class SparKMasker(nn.Module): + """Generates hierarchical random token masks for sparse feature processing. + + 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], @@ -538,6 +756,15 @@ def __init__( self.mask_ratio = mask_ratio def mask(self, B: int, device: torch.device) -> torch.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 @@ -550,21 +777,31 @@ def mask(self, B: int, device: torch.device) -> torch.Tensor: ) def forward(self, inp_bchw: torch.Tensor) -> SparKMaskingOuptut: + """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: + SparKMaskingOuptut containing: + - masked_bchw: Input image masked at full resolution. + - per_level_mask: List of masks at each hierarchical level. + """ global _cur_active - _cur_active = self.mask(inp_bchw.shape[0], inp_bchw.device) # (B, 1, f, f) + _cur_active = self.mask(inp_bchw.shape[0], inp_bchw.device) active_b1ff = _cur_active.clone() downsample_ratio = self.downsample_ratio per_level_mask = [active_b1ff] - for i in range(int(math.log2(downsample_ratio))): + 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 - ) # (B, 1, f*ds, f*ds) + ) per_level_mask.append(active_b1cHcW) - active_b1hw = per_level_mask[ - -1 - ] # the mask map at the deepest scale (the smallest feature map), which would be used for masking the input to the encoder; shape: (B, 1, f, f) + active_b1hw = per_level_mask[-1] masked_bchw = inp_bchw * active_b1hw return SparKMaskingOuptut( masked_bchw=masked_bchw, per_level_mask=per_level_mask @@ -572,12 +809,16 @@ def forward(self, inp_bchw: torch.Tensor) -> SparKMaskingOuptut: class SparKPatchReconLoss(nn.Module): - """Loss module that computes per-patch normalized reconstruction loss. + """Computes per-patch normalized reconstruction loss for masked regions. - Accepts the non-flattened mask `(B, 1, f, f)` and flattens internally. + 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): + def __init__(self, eps: float = 1e-6) -> None: super().__init__() self.eps = eps @@ -586,18 +827,27 @@ def forward( inp_patches: torch.Tensor, rec_patches: torch.Tensor, active_mask: torch.Tensor, - ): - """Compute reconstruction loss and return per-patch stats. + ) -> Tuple[torch.Tensor, torch.Tensor, torch.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: (B, L, N) original patches - rec_patches: (B, L, N) reconstructed patches - active_mask: (B, L, f, f) boolean mask (required) + inp_patches: Original patches of shape (B, L, N) where B=batch, L=levels, 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: - recon_loss: scalar tensor - mean: (B, L, 1) per-patch mean - var: (B, L, 1) per-patch std + 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( @@ -609,7 +859,7 @@ def forward( inp_norm = (inp_patches - mean) / var - l2_loss = ((rec_patches - inp_norm) ** 2).mean(dim=2) # (B, L) + l2_loss = ((rec_patches - inp_norm) ** 2).mean(dim=2) non_active = active_mask.logical_not().int().view(active_mask.shape[0], -1) @@ -618,22 +868,39 @@ def forward( class SparKOutputDecoder(nn.Module): - """Handles de-normalizing reconstructed patches and producing visualization tensors. + """Decodes reconstructed patches back to image space and combines with original. + + Handles denormalization and unpatchifying of reconstructed patches, then performs + per-pixel blending: uses original pixels where visible (active), reconstructed pixels + where masked (inactive). - Usage: call with `rec_patches, mean, var, inp_bchw, active_mask_full` where - `rec_patches` is (B, L, N), `mean`/`var` are (B, L, 1), `inp_bchw` is original image - and `active_mask_full` is (B, 1, H, W) boolean mask indicating visible patches. - The decoder is configured with only the minimal spatial properties: `fmap_h`, - `fmap_w` and `downsample_ratio` (no encoder object required). + Minimal configuration: only requires spatial properties (fmap_h, fmap_w, downsample_ratio). + No encoder object needed. + + 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): + 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: torch.Tensor) -> torch.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 @@ -649,14 +916,27 @@ def forward( var: torch.Tensor, inp_bchw: torch.Tensor, active_mask_full: torch.Tensor, - ): - # de-normalize and unpatchify - rec_bchw = self.unpatchify(rec_patches * var + mean) + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Reconstruct image by blending original and reconstructed regions. - # masked input at full resolution - masked_bchw = inp_bchw * active_mask_full + 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. - # combine: use original where visible, reconstructed where masked + 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 From 96739fe723bdc31fec30264a54fbefcaa5cb43be Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 12 Feb 2026 09:34:53 -0300 Subject: [PATCH 38/85] fix: type hinting --- lightly/models/modules/sparse_spark.py | 23 +++++++++++++---------- 1 file changed, 13 insertions(+), 10 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 4803ad0a1..55cc662a3 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -1,7 +1,10 @@ # Copyright (c) 2023 Keyu Tian # Copyright (c) ByteDance, Inc. and its affiliates. + +from __future__ import annotations + import math -from typing import List, NamedTuple, Optional, Tuple, Union +from typing import NamedTuple import torch import torch.nn as nn @@ -43,14 +46,14 @@ def coalesce_to_size_2_t(t: tuple[int, ...]) -> _size_2_t: raise ValueError(f"Invalid tuple length: {len(t)}; expected 1 or 2.") -_cur_active: Optional[torch.Tensor] = ( +_cur_active: torch.Tensor | None = ( None # B1ff - Global active/mask tensor tracked during forward passes ) def _get_active_ex_or_ii( H: int, W: int, returning_active_ex: bool = True -) -> Union[torch.Tensor, Tuple[torch.Tensor, ...]]: +) -> torch.Tensor | tuple[torch.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). @@ -332,7 +335,7 @@ class SparseEncoder(nn.Module): fmap_w: Width of feature map at encoder output (patch grid). """ - enc_feat_map_chs: List[int] + enc_feat_map_chs: list[int] def __init__( self, @@ -527,7 +530,7 @@ def __init__( self.initialize() - def forward(self, to_dec: List[torch.Tensor]) -> torch.Tensor: + def forward(self, to_dec: list[torch.Tensor]) -> torch.Tensor: """Progressively upsample and combine feature maps. Args: @@ -638,7 +641,7 @@ def __init__( def forward( self, bcff: torch.Tensor, cur_active: torch.Tensor - ) -> Optional[torch.Tensor]: + ) -> torch.Tensor | None: """Densify sparse features by filling masked regions with learned tokens. Args: @@ -696,7 +699,7 @@ def __init__( self.blocks.append(block) d_width //= 2 - def forward(self, fea_bcffs: List[torch.Tensor]) -> List[Optional[torch.Tensor]]: + def forward(self, fea_bcffs: list[torch.Tensor]) -> list[torch.Tensor | None]: """Convert sparse features to dense by filling masked regions. Args: @@ -729,7 +732,7 @@ class SparKMaskingOuptut(NamedTuple): """ masked_bchw: torch.Tensor - per_level_mask: List[torch.Tensor] + per_level_mask: list[torch.Tensor] class SparKMasker(nn.Module): @@ -827,7 +830,7 @@ def forward( inp_patches: torch.Tensor, rec_patches: torch.Tensor, active_mask: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Compute reconstruction loss and per-patch statistics. Normalizes original patches based on per-patch mean and variance, then computes @@ -916,7 +919,7 @@ def forward( var: torch.Tensor, inp_bchw: torch.Tensor, active_mask_full: torch.Tensor, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Reconstruct image by blending original and reconstructed regions. Denormalizes reconstructed patches, unpatchifies them, and performs pixel-wise From ed70f9ab86d78f28efc1c2ed8f4b0339f6c398e8 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 13 Feb 2026 08:00:08 -0300 Subject: [PATCH 39/85] tests: testing active ex --- lightly/models/modules/sparse_spark.py | 6 ++ tests/models/modules/test_sparse_spark.py | 87 +++++++++++++++++++++++ 2 files changed, 93 insertions(+) create mode 100644 tests/models/modules/test_sparse_spark.py diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 55cc662a3..ae86c43a8 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -76,6 +76,12 @@ def _get_active_ex_or_ii( "_cur_active must be set before calling this function" ) h_repeat, w_repeat = H // _cur_active.shape[-2], W // _cur_active.shape[-1] + assert h_repeat > 0, ( + f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" + ) + assert w_repeat > 0, ( + f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" + ) active_ex = _cur_active.repeat_interleave(h_repeat, dim=2).repeat_interleave( w_repeat, dim=3 ) diff --git a/tests/models/modules/test_sparse_spark.py b/tests/models/modules/test_sparse_spark.py new file mode 100644 index 000000000..3b3162613 --- /dev/null +++ b/tests/models/modules/test_sparse_spark.py @@ -0,0 +1,87 @@ +from contextlib import contextmanager +from typing import Generator + +import pytest +import torch + +import lightly.models.modules.sparse_spark as sparse_spark + + +@contextmanager +def _cleanup_curr_active() -> Generator[None, None, None]: + try: + yield + finally: + sparse_spark._cur_active = None + + +def test__get_active_ex_or_ii_expands_mask() -> None: + with _cleanup_curr_active(): + H, W = 32, 32 + sparse_spark._cur_active = torch.tensor( + [ + [ + [ + [1, 0], + [0, 1], + ] + ] + ], + dtype=torch.bool, + ) + + 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: + with _cleanup_curr_active(): + H, W = 4, 4 + sparse_spark._cur_active = 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: + with _cleanup_curr_active(): + H, W = 32, 32 + sparse_spark._cur_active = None + 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_returning_ex_false_correct_values() -> None: + with _cleanup_curr_active(): + H, W = 32, 32 + sparse_spark._cur_active = torch.tensor( + [ + [ + [ + [1, 0], + [0, 1], + ] + ] + ], + dtype=torch.bool, + ) + + 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) From 4fb16d71aa44ebc74c2c31cf1ee58f65d589be04 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 13 Feb 2026 08:01:39 -0300 Subject: [PATCH 40/85] tests: sp conv forward test --- tests/models/modules/test_sparse_spark.py | 35 +++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/tests/models/modules/test_sparse_spark.py b/tests/models/modules/test_sparse_spark.py index 3b3162613..4254b9d15 100644 --- a/tests/models/modules/test_sparse_spark.py +++ b/tests/models/modules/test_sparse_spark.py @@ -85,3 +85,38 @@ def test__get_active_ex_or_ii_returning_ex_false_correct_values() -> None: active_ex_scattered[active_b, :, active_h, active_w] = 1 assert torch.equal(active_ex, active_ex_scattered) + + +def test_sp_conv_forward() -> None: + with _cleanup_curr_active(): + H, W = 32, 32 + sparse_spark._cur_active = torch.tensor( + [ + [ + [ + [1, 0], + [0, 1], + ] + ] + ], + dtype=torch.bool, + ) + + 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() From 6188aa45c12080b47aa41a4d9f5c870b6321f04e Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 13 Feb 2026 08:09:31 -0300 Subject: [PATCH 41/85] refactor: using fixture instead of context manager --- tests/models/modules/test_sparse_spark.py | 161 ++++++++++------------ 1 file changed, 76 insertions(+), 85 deletions(-) diff --git a/tests/models/modules/test_sparse_spark.py b/tests/models/modules/test_sparse_spark.py index 4254b9d15..12691bad9 100644 --- a/tests/models/modules/test_sparse_spark.py +++ b/tests/models/modules/test_sparse_spark.py @@ -1,4 +1,3 @@ -from contextlib import contextmanager from typing import Generator import pytest @@ -7,116 +6,108 @@ import lightly.models.modules.sparse_spark as sparse_spark -@contextmanager -def _cleanup_curr_active() -> Generator[None, None, None]: - try: - yield - finally: - sparse_spark._cur_active = None +@pytest.fixture(autouse=True) +def _cleanup_cur_active() -> Generator[None, None, None]: + sparse_spark._cur_active = None + yield + sparse_spark._cur_active = None def test__get_active_ex_or_ii_expands_mask() -> None: - with _cleanup_curr_active(): - H, W = 32, 32 - sparse_spark._cur_active = torch.tensor( + H, W = 32, 32 + sparse_spark._cur_active = torch.tensor( + [ [ [ - [ - [1, 0], - [0, 1], - ] + [1, 0], + [0, 1], ] - ], - dtype=torch.bool, - ) + ] + ], + dtype=torch.bool, + ) - active = sparse_spark._get_active_ex_or_ii(H=H, W=W, returning_active_ex=True) - assert not isinstance(active, tuple) + 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() + 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: - with _cleanup_curr_active(): - H, W = 4, 4 - sparse_spark._cur_active = 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) + H, W = 4, 4 + sparse_spark._cur_active = 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: - with _cleanup_curr_active(): - H, W = 32, 32 - sparse_spark._cur_active = None - with pytest.raises( - AssertionError, - ): - sparse_spark._get_active_ex_or_ii(H=H, W=W, returning_active_ex=False) + H, W = 32, 32 + sparse_spark._cur_active = None + 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_returning_ex_false_correct_values() -> None: - with _cleanup_curr_active(): - H, W = 32, 32 - sparse_spark._cur_active = torch.tensor( + H, W = 32, 32 + sparse_spark._cur_active = torch.tensor( + [ [ [ - [ - [1, 0], - [0, 1], - ] + [1, 0], + [0, 1], ] - ], - dtype=torch.bool, - ) + ] + ], + dtype=torch.bool, + ) - 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_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 + 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) + assert torch.equal(active_ex, active_ex_scattered) def test_sp_conv_forward() -> None: - with _cleanup_curr_active(): - H, W = 32, 32 - sparse_spark._cur_active = torch.tensor( + H, W = 32, 32 + sparse_spark._cur_active = torch.tensor( + [ [ [ - [ - [1, 0], - [0, 1], - ] + [1, 0], + [0, 1], ] - ], - dtype=torch.bool, - ) - - 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() + ] + ], + dtype=torch.bool, + ) + + 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() From 7cfc5089b22bd0b08aba5f22aa8df4fcaacb86a6 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 13 Feb 2026 15:28:37 -0300 Subject: [PATCH 42/85] format: formatting --- lightly/models/modules/sparse_spark.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index ae86c43a8..45754a505 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -72,16 +72,16 @@ def _get_active_ex_or_ii( Note: Optimization opportunity: Consider using gather() for better performance (see TODO). """ - assert _cur_active is not None, ( - "_cur_active must be set before calling this function" - ) + assert ( + _cur_active is not None + ), "_cur_active must be set before calling this function" h_repeat, w_repeat = H // _cur_active.shape[-2], W // _cur_active.shape[-1] - assert h_repeat > 0, ( - f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" - ) - assert w_repeat > 0, ( - f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" - ) + assert ( + h_repeat > 0 + ), f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" + assert ( + w_repeat > 0 + ), f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" active_ex = _cur_active.repeat_interleave(h_repeat, dim=2).repeat_interleave( w_repeat, dim=3 ) From 0098425ae94f6dc0eef4c9d21834ec378d255c6f Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 14 Feb 2026 09:58:42 -0300 Subject: [PATCH 43/85] feat: moved patch recon loss to loss module --- examples/pytorch_lightning/spark.py | 3 +- lightly/loss/sparse_spark.py | 61 ++++++++++++++++++++ lightly/models/modules/sparse_spark.py | 77 +++----------------------- 3 files changed, 72 insertions(+), 69 deletions(-) create mode 100644 lightly/loss/sparse_spark.py diff --git a/examples/pytorch_lightning/spark.py b/examples/pytorch_lightning/spark.py index 11ef7b0a3..04149cdfb 100644 --- a/examples/pytorch_lightning/spark.py +++ b/examples/pytorch_lightning/spark.py @@ -13,6 +13,8 @@ from torch import nn from torchvision.transforms import v2 +from lightly.loss.sparse_spark import SparKPatchReconLoss + ## The global projection head is the same as the Barlow Twins one from lightly.models.modules.sparse_spark import ( LightDecoder, @@ -20,7 +22,6 @@ SparKMasker, SparKMaskingOuptut, SparKOutputDecoder, - SparKPatchReconLoss, SparseEncoder, ) from lightly.models.utils import patchify diff --git a/lightly/loss/sparse_spark.py b/lightly/loss/sparse_spark.py new file mode 100644 index 000000000..ec3b2cfab --- /dev/null +++ b/lightly/loss/sparse_spark.py @@ -0,0 +1,61 @@ +import torch +from torch import nn + + +class SparKPatchReconLoss(nn.Module): + """Computes per-patch normalized reconstruction loss for masked regions. + + 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) -> None: + super().__init__() + self.eps = eps + + def forward( + self, + inp_patches: torch.Tensor, + rec_patches: torch.Tensor, + active_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.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=levels, 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_norm = (inp_patches - mean) / var + + l2_loss = ((rec_patches - inp_norm) ** 2).mean(dim=2) + + non_active = active_mask.logical_not().int().view(active_mask.shape[0], -1) + + recon_loss = l2_loss.mul_(non_active).sum() / (non_active.sum() + 1e-8) + return recon_loss, mean, var diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 45754a505..6ea21ca3b 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -72,16 +72,16 @@ def _get_active_ex_or_ii( Note: Optimization opportunity: Consider using gather() for better performance (see TODO). """ - assert ( - _cur_active is not None - ), "_cur_active must be set before calling this function" + assert _cur_active is not None, ( + "_cur_active must be set before calling this function" + ) h_repeat, w_repeat = H // _cur_active.shape[-2], W // _cur_active.shape[-1] - assert ( - h_repeat > 0 - ), f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" - assert ( - w_repeat > 0 - ), f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" + assert h_repeat > 0, ( + f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" + ) + assert w_repeat > 0, ( + f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" + ) active_ex = _cur_active.repeat_interleave(h_repeat, dim=2).repeat_interleave( w_repeat, dim=3 ) @@ -817,65 +817,6 @@ def forward(self, inp_bchw: torch.Tensor) -> SparKMaskingOuptut: ) -class SparKPatchReconLoss(nn.Module): - """Computes per-patch normalized reconstruction loss for masked regions. - - 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) -> None: - super().__init__() - self.eps = eps - - def forward( - self, - inp_patches: torch.Tensor, - rec_patches: torch.Tensor, - active_mask: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.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=levels, 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_norm = (inp_patches - mean) / var - - l2_loss = ((rec_patches - inp_norm) ** 2).mean(dim=2) - - non_active = active_mask.logical_not().int().view(active_mask.shape[0], -1) - - recon_loss = l2_loss.mul_(non_active).sum() / (non_active.sum() + 1e-8) - return recon_loss, mean, var - - class SparKOutputDecoder(nn.Module): """Decodes reconstructed patches back to image space and combines with original. From dea0173a2e5530b2053828ad84b1469d03f811db Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 14 Feb 2026 10:02:23 -0300 Subject: [PATCH 44/85] fix: no need of this sparse argument since its always sparse. --- lightly/models/modules/sparse_spark.py | 84 +++++++++----------------- 1 file changed, 27 insertions(+), 57 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 6ea21ca3b..5a171fb92 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -196,8 +196,8 @@ class SparseConvNeXtLayerNorm(nn.LayerNorm): shape (batch_size, height, width, channels) while channels_first corresponds to inputs with shape (batch_size, channels, height, width). - When sparse=True, only applies normalization to active (unmasked) spatial positions, - using indices from the global _cur_active mask. + This sparse implementation only applies normalization to active (unmasked) spatial + positions using indices from the global _cur_active mask. """ def __init__( @@ -205,76 +205,53 @@ def __init__( normalized_shape: int, eps: float = 1e-6, data_format: str = "channels_last", - sparse: bool = True, ) -> 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 - self.sparse = sparse def forward(self, input: torch.Tensor) -> torch.Tensor: if input.ndim == 4: # BHWC or BCHW if self.data_format == "channels_last": # BHWC - if self.sparse: - 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: - return super(SparseConvNeXtLayerNorm, self).forward(input) + 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 - if self.sparse: - 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: - u = input.mean(1, keepdim=True) - s = (input - u).pow(2).mean(1, keepdim=True) - input = (input - u) / torch.sqrt(s + self.eps) - input = ( - self.weight[:, None, None] * input + self.bias[:, None, None] - ) - return input + 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 - if self.sparse: - raise NotImplementedError - else: - return super().forward(input) - - def __repr__(self): - return ( - super().__repr__()[:-1] - + f", ch={self.data_format.split('_')[-1]}, sp={self.sparse})" - ) + raise NotImplementedError class SparseConvNeXtBlock(nn.Module): - r"""ConvNeXt Block with optional sparse computation support. + 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. When sparse=True, applies masking to the output. + 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. - sparse: Whether to apply sparse masking. Default: True. ks: Kernel size for depthwise convolution. Default: 7. """ @@ -283,12 +260,11 @@ def __init__( dim: int, drop_path: float = 0.0, layer_scale_init_value: float = 1e-6, - sparse: bool = True, ks: int = 7, ) -> None: super().__init__() self.dwconv = nn.Conv2d(dim, dim, kernel_size=ks, padding=ks // 2, groups=dim) - self.norm = SparseConvNeXtLayerNorm(dim, eps=1e-6, sparse=sparse) + 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) @@ -300,7 +276,6 @@ def __init__( self.drop_path: nn.Module = ( DropPath(drop_path) if drop_path > 0.0 else nn.Identity() ) - self.sparse = sparse def forward(self, x: torch.Tensor) -> torch.Tensor: input = x @@ -314,16 +289,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: x = self.gamma * x x = x.permute(0, 3, 1, 2) - if self.sparse: - x *= _get_active_ex_or_ii( - H=x.shape[2], W=x.shape[3], returning_active_ex=True - ) + x *= _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 __repr__(self): - return super().__repr__()[:-1] + f", sp={self.sparse})" class SparseEncoder(nn.Module): From 83bdd60251d560f89bfa9bd3758d0e86d9fc11c7 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 14 Feb 2026 10:02:56 -0300 Subject: [PATCH 45/85] format --- lightly/models/modules/sparse_spark.py | 19 +++++++++---------- 1 file changed, 9 insertions(+), 10 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 5a171fb92..1fe18540a 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -72,16 +72,16 @@ def _get_active_ex_or_ii( Note: Optimization opportunity: Consider using gather() for better performance (see TODO). """ - assert _cur_active is not None, ( - "_cur_active must be set before calling this function" - ) + assert ( + _cur_active is not None + ), "_cur_active must be set before calling this function" h_repeat, w_repeat = H // _cur_active.shape[-2], W // _cur_active.shape[-1] - assert h_repeat > 0, ( - f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" - ) - assert w_repeat > 0, ( - f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" - ) + assert ( + h_repeat > 0 + ), f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" + assert ( + w_repeat > 0 + ), f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" active_ex = _cur_active.repeat_interleave(h_repeat, dim=2).repeat_interleave( w_repeat, dim=3 ) @@ -295,7 +295,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x - class SparseEncoder(nn.Module): """Converts a dense CNN model to a sparse CNN model by replacing standard layers. From 4b7f1192edbaa1a0fb84bc1ddd592f3309deb48a Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 14 Feb 2026 10:03:35 -0300 Subject: [PATCH 46/85] feat: init module access --- lightly/loss/__init__.py | 3 ++- lightly/models/modules/__init__.py | 7 ++++++- 2 files changed, 8 insertions(+), 2 deletions(-) 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/models/modules/__init__.py b/lightly/models/modules/__init__.py index f5c7c39f4..df9c60506 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, + SparseEncoder, +) from lightly.utils import dependency as _dependency if _dependency.torchvision_vit_available(): From 86386b744d212eba7ec084cfd9dda53af99249ee Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 14 Feb 2026 10:10:12 -0300 Subject: [PATCH 47/85] feat: removed sparse encoder since it adds no necessary logic. --- lightly/models/modules/__init__.py | 2 +- lightly/models/modules/sparse_spark.py | 227 ++++++++++--------------- 2 files changed, 88 insertions(+), 141 deletions(-) diff --git a/lightly/models/modules/__init__.py b/lightly/models/modules/__init__.py index df9c60506..d4be20e0d 100644 --- a/lightly/models/modules/__init__.py +++ b/lightly/models/modules/__init__.py @@ -33,7 +33,7 @@ SparKDensifier, SparKMasker, SparKOutputDecoder, - SparseEncoder, + dense_model_to_sparse, ) from lightly.utils import dependency as _dependency diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 1fe18540a..5a96c5fbc 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -72,16 +72,16 @@ def _get_active_ex_or_ii( Note: Optimization opportunity: Consider using gather() for better performance (see TODO). """ - assert ( - _cur_active is not None - ), "_cur_active must be set before calling this function" + assert _cur_active is not None, ( + "_cur_active must be set before calling this function" + ) h_repeat, w_repeat = H // _cur_active.shape[-2], W // _cur_active.shape[-1] - assert ( - h_repeat > 0 - ), f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" - assert ( - w_repeat > 0 - ), f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" + assert h_repeat > 0, ( + f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" + ) + assert w_repeat > 0, ( + f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" + ) active_ex = _cur_active.repeat_interleave(h_repeat, dim=2).repeat_interleave( w_repeat, dim=3 ) @@ -295,144 +295,91 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x -class SparseEncoder(nn.Module): - """Converts a dense CNN model to a sparse CNN model by replacing standard layers. +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. - Recursively traverses a model and replaces standard layers (Conv2d, Pooling, BatchNorm, LayerNorm) - with their sparse counterparts. Sparse layers respect a global active mask (_cur_active) that - indicates which spatial regions should be processed. + Handles Conv2d, MaxPool2d, AvgPool2d, BatchNorm2d, SyncBatchNorm, and LayerNorm layers. + Copies weight and state tensors to maintain the original model's parameters. - Attributes: - enc_feat_map_chs: List of channel numbers of feature maps at different scales, shallow to deep. - input_size: Original input image size (assumed square). - downsample_ratio: Total spatial downsampling ratio from input to deepest feature map. - fmap_h: Height of feature map at encoder output (patch grid). - fmap_w: Width of feature map at encoder output (patch grid). - """ + 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. - enc_feat_map_chs: list[int] + Returns: + Sparse version of the input model with converted layers. - def __init__( - self, - cnn: nn.Module, - input_size: int, - downsample_ratio: int, - feature_map_channels: list[int], - sbn: bool = False, - verbose: bool = False, - ) -> None: - super().__init__() - self.sp_cnn = SparseEncoder.dense_model_to_sparse( - m=cnn, verbose=verbose, sbn=sbn - ) - self.input_size, self.downsample_ratio, self.enc_feat_map_chs = ( - input_size, - downsample_ratio, - feature_map_channels, + Raises: + NotImplementedError: If Conv1d layers are encountered. + """ + oup = m + if isinstance(m, nn.Conv2d): + bias = m.bias is not None + + oup = 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, ) - self.fmap_h = input_size // self.downsample_ratio - self.fmap_w = input_size // self.downsample_ratio - - @staticmethod - 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. - - 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. + oup.weight.data.copy_(m.weight.data) - Returns: - Sparse version of the input model with converted layers. - - Raises: - NotImplementedError: If Conv1d layers are encountered. - """ - oup = m - if isinstance(m, nn.Conv2d): - bias = m.bias is not None - - oup = 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, - ) - oup.weight.data.copy_(m.weight.data) - - if m.bias is not None and oup.bias is not None: - oup.bias.data.copy_(m.bias.data) - - 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)): - oup = (SparseSyncBatchNorm2d if sbn else SparseBatchNorm2d)( - m.weight.shape[0], - eps=m.eps, - momentum=m.momentum, - affine=m.affine, - track_running_stats=m.track_running_stats, - ) - oup.weight.data.copy_(m.weight.data) - oup.bias.data.copy_(m.bias.data) - oup.running_mean.data.copy_(m.running_mean.data) - oup.running_var.data.copy_(m.running_var.data) - oup.num_batches_tracked.data.copy_(m.num_batches_tracked.data) - if hasattr(m, "qconfig"): - oup.qconfig = m.qconfig - elif isinstance(m, nn.LayerNorm) and not isinstance(m, SparseConvNeXtLayerNorm): - oup = SparseConvNeXtLayerNorm(m.weight.shape[0], eps=m.eps) - oup.weight.data.copy_(m.weight.data) + if m.bias is not None and oup.bias is not None: oup.bias.data.copy_(m.bias.data) - elif isinstance(m, (nn.Conv1d,)): - raise NotImplementedError - - for name, child in m.named_children(): - oup.add_module( - name, - SparseEncoder.dense_model_to_sparse(child, verbose=verbose, sbn=sbn), - ) - del m - return oup - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """Forward pass through the sparse CNN encoder. - - Args: - x: Input image tensor. - Returns: - Output feature tensor from the sparse CNN. - """ - return self.sp_cnn(x) + 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)): + oup = (SparseSyncBatchNorm2d if sbn else SparseBatchNorm2d)( + m.weight.shape[0], + eps=m.eps, + momentum=m.momentum, + affine=m.affine, + track_running_stats=m.track_running_stats, + ) + oup.weight.data.copy_(m.weight.data) + oup.bias.data.copy_(m.bias.data) + oup.running_mean.data.copy_(m.running_mean.data) + oup.running_var.data.copy_(m.running_var.data) + oup.num_batches_tracked.data.copy_(m.num_batches_tracked.data) + if hasattr(m, "qconfig"): + oup.qconfig = m.qconfig + elif isinstance(m, nn.LayerNorm) and not isinstance(m, SparseConvNeXtLayerNorm): + oup = SparseConvNeXtLayerNorm(m.weight.shape[0], eps=m.eps) + oup.weight.data.copy_(m.weight.data) + oup.bias.data.copy_(m.bias.data) + 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)) + del m + return oup class UNetBlock(nn.Module): @@ -589,7 +536,7 @@ def __init__( ) elif densify_norm_str == "ln": self.densify_norm = SparseConvNeXtLayerNorm( - e_width, data_format="channels_first", sparse=True + e_width, data_format="channels_first" ) else: self.densify_norm = nn.Identity() From 92ebd99ff298fefd39214d56191b13b9571b61da Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 14 Feb 2026 10:12:20 -0300 Subject: [PATCH 48/85] refactor: removing dense model to sparse --- examples/pytorch_lightning/spark.py | 33 +++++++++++++---------------- 1 file changed, 15 insertions(+), 18 deletions(-) diff --git a/examples/pytorch_lightning/spark.py b/examples/pytorch_lightning/spark.py index 04149cdfb..d421d1d01 100644 --- a/examples/pytorch_lightning/spark.py +++ b/examples/pytorch_lightning/spark.py @@ -22,7 +22,7 @@ SparKMasker, SparKMaskingOuptut, SparKOutputDecoder, - SparseEncoder, + dense_model_to_sparse, ) from lightly.models.utils import patchify @@ -47,37 +47,34 @@ def __init__( backbone = timm.create_model( "resnet18", drop_path_rate=0.05, features_only=True ) - self.sparse_encoder = SparseEncoder( - backbone, - downsample_ratio=get_downsample_ratio_from_timm_model(backbone), - feature_map_channels=get_enc_feat_map_chs_from_timm_model(backbone), - input_size=input_size, - sbn=sbn, - verbose=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 = dense_model_to_sparse(backbone, sbn=sbn, verbose=True) + self.fmap_h = input_size // downsample_ratio + self.fmap_w = input_size // downsample_ratio self.dense_decoder = LightDecoder( - self.sparse_encoder.downsample_ratio, - width=self.sparse_encoder.enc_feat_map_chs[-1], + downsample_ratio, + width=enc_feat_map_chs[-1], ) self.masker = SparKMasker( - feature_map_size=(self.sparse_encoder.fmap_h, self.sparse_encoder.fmap_w), - downsample_ratio=self.sparse_encoder.downsample_ratio, + feature_map_size=(self.fmap_h, self.fmap_w), + downsample_ratio=downsample_ratio, mask_ratio=mask_ratio, ) self.densifier = SparKDensifier( - encoder_in_channels=self.sparse_encoder.enc_feat_map_chs, + 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 = self.sparse_encoder.downsample_ratio + 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( - self.sparse_encoder.fmap_h, - self.sparse_encoder.fmap_w, - self.sparse_encoder.downsample_ratio, + self.fmap_h, + self.fmap_w, + downsample_ratio, ) def forward( From 93c7ee50d28a166955361dc821a50c73d521cae3 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 14 Feb 2026 10:12:47 -0300 Subject: [PATCH 49/85] fix: renamin to unet decoder --- examples/pytorch_lightning/spark.py | 4 ++-- lightly/models/modules/sparse_spark.py | 20 ++++++++++---------- 2 files changed, 12 insertions(+), 12 deletions(-) diff --git a/examples/pytorch_lightning/spark.py b/examples/pytorch_lightning/spark.py index d421d1d01..8f5d07bb5 100644 --- a/examples/pytorch_lightning/spark.py +++ b/examples/pytorch_lightning/spark.py @@ -17,11 +17,11 @@ ## The global projection head is the same as the Barlow Twins one from lightly.models.modules.sparse_spark import ( - LightDecoder, SparKDensifier, SparKMasker, SparKMaskingOuptut, SparKOutputDecoder, + UNetDecoder, dense_model_to_sparse, ) from lightly.models.utils import patchify @@ -52,7 +52,7 @@ def __init__( self.sparse_encoder = dense_model_to_sparse(backbone, sbn=sbn, verbose=True) self.fmap_h = input_size // downsample_ratio self.fmap_w = input_size // downsample_ratio - self.dense_decoder = LightDecoder( + self.dense_decoder = UNetDecoder( downsample_ratio, width=enc_feat_map_chs[-1], ) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 5a96c5fbc..1ae2f93c9 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -72,16 +72,16 @@ def _get_active_ex_or_ii( Note: Optimization opportunity: Consider using gather() for better performance (see TODO). """ - assert _cur_active is not None, ( - "_cur_active must be set before calling this function" - ) + assert ( + _cur_active is not None + ), "_cur_active must be set before calling this function" h_repeat, w_repeat = H // _cur_active.shape[-2], W // _cur_active.shape[-1] - assert h_repeat > 0, ( - f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" - ) - assert w_repeat > 0, ( - f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" - ) + assert ( + h_repeat > 0 + ), f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" + assert ( + w_repeat > 0 + ), f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" active_ex = _cur_active.repeat_interleave(h_repeat, dim=2).repeat_interleave( w_repeat, dim=3 ) @@ -420,7 +420,7 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return self.conv(x) -class LightDecoder(nn.Module): +class UNetDecoder(nn.Module): """Lightweight hierarchical decoder for feature map reconstruction. Applies a series of UNetBlocks to progressively upsample feature maps from deep to shallow, From 2710f6b0ea1a1993e69bfcb3cec64a0c1b9ddfae Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 14 Feb 2026 10:13:30 -0300 Subject: [PATCH 50/85] format --- lightly/models/modules/sparse_spark.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 1ae2f93c9..c458a2cfe 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -72,16 +72,16 @@ def _get_active_ex_or_ii( Note: Optimization opportunity: Consider using gather() for better performance (see TODO). """ - assert ( - _cur_active is not None - ), "_cur_active must be set before calling this function" + assert _cur_active is not None, ( + "_cur_active must be set before calling this function" + ) h_repeat, w_repeat = H // _cur_active.shape[-2], W // _cur_active.shape[-1] - assert ( - h_repeat > 0 - ), f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" - assert ( - w_repeat > 0 - ), f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" + assert h_repeat > 0, ( + f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" + ) + assert w_repeat > 0, ( + f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" + ) active_ex = _cur_active.repeat_interleave(h_repeat, dim=2).repeat_interleave( w_repeat, dim=3 ) From 3a59d6a571c66a49a8c5d7f2886210dab7c13f97 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Tue, 24 Feb 2026 18:07:57 -0300 Subject: [PATCH 51/85] fix: removed inplace operation --- lightly/loss/sparse_spark.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lightly/loss/sparse_spark.py b/lightly/loss/sparse_spark.py index ec3b2cfab..af6318c68 100644 --- a/lightly/loss/sparse_spark.py +++ b/lightly/loss/sparse_spark.py @@ -57,5 +57,5 @@ def forward( non_active = active_mask.logical_not().int().view(active_mask.shape[0], -1) - recon_loss = l2_loss.mul_(non_active).sum() / (non_active.sum() + 1e-8) + recon_loss = (l2_loss * non_active).sum() / (non_active.sum() + 1e-8) return recon_loss, mean, var From 34ed5f67383b32fbce779db98b197fdfea7ea631 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Tue, 24 Feb 2026 18:08:44 -0300 Subject: [PATCH 52/85] fix: using proper eps --- lightly/loss/sparse_spark.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lightly/loss/sparse_spark.py b/lightly/loss/sparse_spark.py index af6318c68..ab1f74023 100644 --- a/lightly/loss/sparse_spark.py +++ b/lightly/loss/sparse_spark.py @@ -57,5 +57,5 @@ def forward( non_active = active_mask.logical_not().int().view(active_mask.shape[0], -1) - recon_loss = (l2_loss * non_active).sum() / (non_active.sum() + 1e-8) + recon_loss = (l2_loss * non_active).sum() / (non_active.sum() + self.eps) return recon_loss, mean, var From 2ff68e98a38056df2eda94d5ae53a64dcd7f90ca Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Tue, 24 Feb 2026 18:10:01 -0300 Subject: [PATCH 53/85] docs: rst for loss --- docs/source/lightly.loss.rst | 3 +++ 1 file changed, 3 insertions(+) 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: From de92784fa84a1e10fd24f858e21cb30ab01e5611 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Tue, 24 Feb 2026 18:11:43 -0300 Subject: [PATCH 54/85] feat: annotations --- lightly/loss/sparse_spark.py | 1 + 1 file changed, 1 insertion(+) diff --git a/lightly/loss/sparse_spark.py b/lightly/loss/sparse_spark.py index ab1f74023..bcf37a21c 100644 --- a/lightly/loss/sparse_spark.py +++ b/lightly/loss/sparse_spark.py @@ -1,3 +1,4 @@ +from __future__ import annotations import torch from torch import nn From ad4c9ac7a9f429a7c76c204d5ab065d0d30cf8c4 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Tue, 24 Feb 2026 18:13:41 -0300 Subject: [PATCH 55/85] typos --- lightly/models/modules/sparse_spark.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index c458a2cfe..cc020599f 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -499,7 +499,7 @@ def initialize(self) -> None: nn.init.constant_(m.weight, 1.0) -class SparKDensfiyBlock(nn.Module): +class SparKDensifiyBlock(nn.Module): """Block for densifying sparse features by filling masked regions with learned tokens. Applies normalization to sparse features, then uses learned mask tokens to fill inactive @@ -610,7 +610,7 @@ def __init__( 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 = SparKDensfiyBlock( + block = SparKDensifiyBlock( e_width=e_width, d_width=d_width, densify_norm_str=densify_norm_str, @@ -645,7 +645,7 @@ def forward(self, fea_bcffs: list[torch.Tensor]) -> list[torch.Tensor | None]: return to_dec -class SparKMaskingOuptut(NamedTuple): +class SparKMaskingOutput(NamedTuple): """Output container for SparKMasker forward pass. Attributes: @@ -701,7 +701,7 @@ def mask(self, B: int, device: torch.device) -> torch.Tensor: .bool() ) - def forward(self, inp_bchw: torch.Tensor) -> SparKMaskingOuptut: + def forward(self, inp_bchw: torch.Tensor) -> SparKMaskingOutput: """Generate hierarchical masks for the input. Creates masks at multiple scales from the patch level up to full input resolution. @@ -728,7 +728,7 @@ def forward(self, inp_bchw: torch.Tensor) -> SparKMaskingOuptut: active_b1hw = per_level_mask[-1] masked_bchw = inp_bchw * active_b1hw - return SparKMaskingOuptut( + return SparKMaskingOutput( masked_bchw=masked_bchw, per_level_mask=per_level_mask ) From 8b2eca5b6bef46b0b8ca32e0b2cb9d4ac56e1762 Mon Sep 17 00:00:00 2001 From: Lionel Date: Mon, 2 Mar 2026 18:10:03 +0100 Subject: [PATCH 56/85] typing issues --- examples/pytorch_lightning/spark.py | 83 ++++++++++++++++++----------- 1 file changed, 51 insertions(+), 32 deletions(-) diff --git a/examples/pytorch_lightning/spark.py b/examples/pytorch_lightning/spark.py index 8f5d07bb5..8c772442d 100644 --- a/examples/pytorch_lightning/spark.py +++ b/examples/pytorch_lightning/spark.py @@ -1,5 +1,5 @@ # This example requires the following dependencies to be installed: -# pip install lightly +# 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 @@ -9,33 +9,43 @@ import timm import torch import torchvision -from pytorch_lightning.callbacks import ModelCheckpoint, RichProgressBar -from torch import nn +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, - SparKMaskingOuptut, + SparKMaskingOutput, SparKOutputDecoder, UNetDecoder, - dense_model_to_sparse, + sparse_layer_context, ) -from lightly.models.utils import patchify -def get_downsample_ratio_from_timm_model(model: nn.Module) -> int: +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: nn.Module) -> list[int]: +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(pl.LightningModule): +class SparseSparK(LightningModule): def __init__( self, input_size: int = 416, @@ -45,15 +55,17 @@ def __init__( ): super().__init__() backbone = timm.create_model( - "resnet18", drop_path_rate=0.05, features_only=True + 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 ) - 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 = dense_model_to_sparse(backbone, sbn=sbn, verbose=True) self.fmap_h = input_size // downsample_ratio self.fmap_w = input_size // downsample_ratio self.dense_decoder = UNetDecoder( - downsample_ratio, + up_sample_ratio=downsample_ratio, width=enc_feat_map_chs[-1], ) self.masker = SparKMasker( @@ -72,41 +84,51 @@ def __init__( self.recon_loss_fn = SparKPatchReconLoss() # output decoder for visualization (pass minimal spatial props) self.output_decoder = SparKOutputDecoder( - self.fmap_h, - self.fmap_w, - downsample_ratio, + fmap_h=self.fmap_h, + fmap_w=self.fmap_w, + downsample_ratio=downsample_ratio, ) def forward( self, - inp_bchw: torch.Tensor, + inp_bchw: Tensor, vis=False, ): # step1. Mask - mask_out: SparKMaskingOuptut = self.masker(inp_bchw) + 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) - fea_bcffs: list[torch.Tensor] = self.sparse_encoder(masked_bchw) - # step3. Densify: get hierarchical dense features for decoding - to_dec = self.densifier(fea_bcffs) + # 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 = ( - patchify(inp_bchw, self.downsample_ratio), - patchify(rec_bchw, self.downsample_ratio), - ) # inp and rec: (B, L = f*f, N = C*downsample_raito**2) + 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, rec, active_b1fHfW) + 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, mean, var, inp_bchw, active_b1hw) + 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, batch_index) -> torch.Tensor: - img, target = batch + 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( @@ -164,9 +186,6 @@ def target_transform(t): max_epochs=30, devices=1, accelerator=accelerator, - callbacks=[ - RichProgressBar(), - ], ) trainer.fit( model=model, From b16c6b34ccb1081d851921de84125fc37378c5c7 Mon Sep 17 00:00:00 2001 From: Lionel Date: Mon, 2 Mar 2026 18:10:14 +0100 Subject: [PATCH 57/85] formatting --- lightly/loss/sparse_spark.py | 1 + 1 file changed, 1 insertion(+) diff --git a/lightly/loss/sparse_spark.py b/lightly/loss/sparse_spark.py index bcf37a21c..2acb3f6d0 100644 --- a/lightly/loss/sparse_spark.py +++ b/lightly/loss/sparse_spark.py @@ -1,4 +1,5 @@ from __future__ import annotations + import torch from torch import nn From 1b9f53a08c2882a56444d44ae193a4dab8d82a53 Mon Sep 17 00:00:00 2001 From: Lionel Date: Mon, 2 Mar 2026 18:10:44 +0100 Subject: [PATCH 58/85] context manager for global tensor + typing --- lightly/models/modules/sparse_spark.py | 102 ++++++++++++++++--------- 1 file changed, 66 insertions(+), 36 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index cc020599f..e1357ef74 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -4,12 +4,13 @@ from __future__ import annotations import math +from contextlib import contextmanager +from contextvars import ContextVar from typing import NamedTuple import torch import torch.nn as nn from timm.models.layers import DropPath, trunc_normal_ -from torch.nn.common_types import _size_2_t from lightly.models.utils import random_token_mask @@ -26,7 +27,7 @@ def is_pow2n(x: int) -> bool: return x > 0 and (x & (x - 1) == 0) -def coalesce_to_size_2_t(t: tuple[int, ...]) -> _size_2_t: +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: @@ -46,9 +47,21 @@ def coalesce_to_size_2_t(t: tuple[int, ...]) -> _size_2_t: raise ValueError(f"Invalid tuple length: {len(t)}; expected 1 or 2.") -_cur_active: torch.Tensor | None = ( - None # B1ff - Global active/mask tensor tracked during forward passes -) +_cur_active: ContextVar[torch.Tensor | None] = ContextVar("_cur_active", default=None) + + +@contextmanager +def sparse_layer_context(active_mask: torch.Tensor): + """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( @@ -72,17 +85,20 @@ def _get_active_ex_or_ii( Note: Optimization opportunity: Consider using gather() for better performance (see TODO). """ - assert _cur_active is not None, ( - "_cur_active must be set before calling this function" - ) - h_repeat, w_repeat = H // _cur_active.shape[-2], W // _cur_active.shape[-1] - assert h_repeat > 0, ( - f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" - ) - assert w_repeat > 0, ( - f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" - ) - active_ex = _cur_active.repeat_interleave(h_repeat, dim=2).repeat_interleave( + 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 ( @@ -92,7 +108,9 @@ def _get_active_ex_or_ii( ) -def sp_conv_forward(self, x: torch.Tensor) -> torch.Tensor: +def sp_conv_forward( + module: nn.AvgPool2d | nn.MaxPool2d | nn.Conv2d, input: torch.Tensor +) -> torch.Tensor: """Forward pass for sparse convolution/pooling layers. Applies the parent class forward operation and masks the output using the global @@ -100,17 +118,21 @@ def sp_conv_forward(self, x: torch.Tensor) -> torch.Tensor: Args: self: ConvTranspose2d, MaxPool2d, or AvgPool2d instance. - x: Input tensor of shape (B, C, H, W). + input: Input tensor of shape (B, C, H, W). Returns: Masked output tensor of same shape as input, with inactive spatial positions zeroed. """ - x = super(type(self), self).forward(x) - x *= _get_active_ex_or_ii(H=x.shape[2], W=x.shape[3], returning_active_ex=True) + x: torch.Tensor = super(type(module), module).forward(input) # type: ignore[arg-type] + x *= _get_active_ex_or_ii( + H=input.shape[2], W=input.shape[3], returning_active_ex=True + ) return x -def sp_bn_forward(self, x: torch.Tensor) -> torch.Tensor: +def sp_bn_forward( + module: nn.BatchNorm1d | nn.SyncBatchNorm, input: torch.Tensor +) -> torch.Tensor: """Forward pass for sparse batch normalization layers. Applies batch norm only to active (unmasked) spatial positions, efficiently handling @@ -119,16 +141,18 @@ def sp_bn_forward(self, x: torch.Tensor) -> torch.Tensor: Args: self: BatchNorm1d or SyncBatchNorm instance. - x: Input tensor of shape (B, C, H, W) in channels_first format. + 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=x.shape[2], W=x.shape[3], returning_active_ex=False) + ii = _get_active_ex_or_ii( + H=input.shape[2], W=input.shape[3], returning_active_ex=False + ) - bhwc = x.permute(0, 2, 3, 1) + bhwc = input.permute(0, 2, 3, 1) nc = bhwc[ii] - nc = super(type(self), self).forward(nc) + nc = super(type(module), module).forward(nc) # type: ignore[arg-type] bchw = torch.zeros_like(bhwc) bchw[ii] = nc @@ -461,7 +485,7 @@ def forward(self, to_dec: list[torch.Tensor]) -> torch.Tensor: Returns: Upsampled feature tensor with 3 output channels. """ - x = 0 + 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] @@ -542,6 +566,7 @@ def __init__( self.densify_norm = nn.Identity() # densify proj + self.densify_proj: nn.Identity | nn.Conv2d if use_identity_proj: self.densify_proj = nn.Identity() print( @@ -562,7 +587,7 @@ def __init__( self.densify_proj = densify_proj def forward( - self, bcff: torch.Tensor, cur_active: torch.Tensor + self, bcff: torch.Tensor | None, cur_active: torch.Tensor ) -> torch.Tensor | None: """Densify sparse features by filling masked regions with learned tokens. @@ -576,6 +601,7 @@ def forward( if bcff is None: return None bcff = self.densify_norm(bcff) + assert bcff is not None mask_tokens = self.mask_token.expand_as(bcff) bcff = torch.where(cur_active.expand_as(bcff), bcff, mask_tokens) bcff = self.densify_proj(bcff) @@ -621,20 +647,26 @@ def __init__( self.blocks.append(block) d_width //= 2 - def forward(self, fea_bcffs: list[torch.Tensor]) -> list[torch.Tensor | None]: + def forward(self, fea_bcffs: list[torch.Tensor]) -> list[torch.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. + 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). + List of densified feature tensors for decoder processing (order reversed and + dilated). """ to_dec = [] fea_bcffs = fea_bcffs[::-1] - global _cur_active - active_fmap_current = _cur_active + 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): if bcff is not None: bcff = self.blocks[i](bcff, active_fmap_current) @@ -714,9 +746,7 @@ def forward(self, inp_bchw: torch.Tensor) -> SparKMaskingOutput: - masked_bchw: Input image masked at full resolution. - per_level_mask: List of masks at each hierarchical level. """ - global _cur_active - _cur_active = self.mask(inp_bchw.shape[0], inp_bchw.device) - active_b1ff = _cur_active.clone() + 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))): From 4e9cde15c5df6b4d3037854c0364fac72597c926 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 2 Mar 2026 14:30:51 -0300 Subject: [PATCH 59/85] format --- lightly/loss/sparse_spark.py | 1 + lightly/models/modules/sparse_spark.py | 18 +++++++++--------- 2 files changed, 10 insertions(+), 9 deletions(-) diff --git a/lightly/loss/sparse_spark.py b/lightly/loss/sparse_spark.py index bcf37a21c..2acb3f6d0 100644 --- a/lightly/loss/sparse_spark.py +++ b/lightly/loss/sparse_spark.py @@ -1,4 +1,5 @@ from __future__ import annotations + import torch from torch import nn diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index cc020599f..921f00a4b 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -72,16 +72,16 @@ def _get_active_ex_or_ii( Note: Optimization opportunity: Consider using gather() for better performance (see TODO). """ - assert _cur_active is not None, ( - "_cur_active must be set before calling this function" - ) + assert ( + _cur_active is not None + ), "_cur_active must be set before calling this function" h_repeat, w_repeat = H // _cur_active.shape[-2], W // _cur_active.shape[-1] - assert h_repeat > 0, ( - f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" - ) - assert w_repeat > 0, ( - f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" - ) + assert ( + h_repeat > 0 + ), f"Target height {H} must be >= mask height {_cur_active.shape[-2]}" + assert ( + w_repeat > 0 + ), f"Target width {W} must be >= mask width {_cur_active.shape[-1]}" active_ex = _cur_active.repeat_interleave(h_repeat, dim=2).repeat_interleave( w_repeat, dim=3 ) From 6abe8fc035bb560c3ff9e0b5940552bef68e8b27 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 2 Mar 2026 14:37:07 -0300 Subject: [PATCH 60/85] refactor: better naming --- lightly/loss/sparse_spark.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lightly/loss/sparse_spark.py b/lightly/loss/sparse_spark.py index 2acb3f6d0..63a48e0d4 100644 --- a/lightly/loss/sparse_spark.py +++ b/lightly/loss/sparse_spark.py @@ -53,9 +53,9 @@ def forward( mean = inp_patches.mean(dim=-1, keepdim=True) var = (inp_patches.var(dim=-1, keepdim=True) + self.eps) ** 0.5 - inp_norm = (inp_patches - mean) / var + inp_normalized = (inp_patches - mean) / var - l2_loss = ((rec_patches - inp_norm) ** 2).mean(dim=2) + l2_loss = ((rec_patches - inp_normalized) ** 2).mean(dim=2) non_active = active_mask.logical_not().int().view(active_mask.shape[0], -1) From c9c99d21a90b07b8e3f705cbd0517079a660eb7b Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 2 Mar 2026 15:35:48 -0300 Subject: [PATCH 61/85] refactor: lenght, not levels --- lightly/loss/sparse_spark.py | 29 ++++++++++++++++++++++++++--- 1 file changed, 26 insertions(+), 3 deletions(-) diff --git a/lightly/loss/sparse_spark.py b/lightly/loss/sparse_spark.py index 63a48e0d4..138d0c544 100644 --- a/lightly/loss/sparse_spark.py +++ b/lightly/loss/sparse_spark.py @@ -2,6 +2,8 @@ import torch from torch import nn +from lightly.utils.dist import gather +import torch.distributed as dist class SparKPatchReconLoss(nn.Module): @@ -14,9 +16,16 @@ class SparKPatchReconLoss(nn.Module): eps: Small value for numerical stability. Default: 1e-6. """ - def __init__(self, eps: float = 1e-6) -> None: + 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, @@ -31,7 +40,7 @@ def forward( only over masked (active_mask=False) patches. Args: - inp_patches: Original patches of shape (B, L, N) where B=batch, L=levels, N=patch_dim. + 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). @@ -59,5 +68,19 @@ def forward( non_active = active_mask.logical_not().int().view(active_mask.shape[0], -1) - recon_loss = (l2_loss * non_active).sum() / (non_active.sum() + self.eps) + 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 From 73da36e3306bf76e028d6ee47079280617a94995 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 2 Mar 2026 15:42:11 -0300 Subject: [PATCH 62/85] fix: remove unecessary target transform --- examples/pytorch_lightning/spark.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/examples/pytorch_lightning/spark.py b/examples/pytorch_lightning/spark.py index 8c772442d..dfca5b220 100644 --- a/examples/pytorch_lightning/spark.py +++ b/examples/pytorch_lightning/spark.py @@ -151,10 +151,6 @@ def configure_optimizers(self): # we ignore object detection annotations by setting target_transform to return 0 -def target_transform(t): - return 0 - - dataset = torchvision.datasets.Caltech101( "datasets/caltech101", download=True, @@ -166,7 +162,6 @@ def target_transform(t): v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), ] ), - target_transform=target_transform, ) # or create a dataset from a folder containing images or videos: # dataset = LightlyDataset("path/to/folder") From 708a18e01a7356fb8a96e71d76a2c3301c9493df Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 2 Mar 2026 16:07:47 -0300 Subject: [PATCH 63/85] test: testing loss, comparing with reference and distributed testing --- tests/loss/test_sparse_spark.py | 73 +++++++++++++++++++++++++++++++++ 1 file changed, 73 insertions(+) create mode 100644 tests/loss/test_sparse_spark.py diff --git a/tests/loss/test_sparse_spark.py b/tests/loss/test_sparse_spark.py new file mode 100644 index 000000000..c3ad7ff1e --- /dev/null +++ b/tests/loss/test_sparse_spark.py @@ -0,0 +1,73 @@ +import torch.distributed as dist +import torch +import pytest +from lightly.loss import SparKPatchReconLoss +from pytest_mock import MockerFixture + + +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: + 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 From 0c10ea022287c008568d033ec5ad1feba2a1e38c Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 2 Mar 2026 16:11:45 -0300 Subject: [PATCH 64/85] doc: where i took the loss from --- tests/loss/test_sparse_spark.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/loss/test_sparse_spark.py b/tests/loss/test_sparse_spark.py index c3ad7ff1e..43c7d7bb7 100644 --- a/tests/loss/test_sparse_spark.py +++ b/tests/loss/test_sparse_spark.py @@ -59,6 +59,13 @@ def _reference_loss_implementation( 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 From bbe0b07b7e7c284af2faf42f5e6b3e8a0b815bfe Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Mon, 2 Mar 2026 16:50:32 -0300 Subject: [PATCH 65/85] refactor: make example simpler --- examples/pytorch_lightning/spark.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/examples/pytorch_lightning/spark.py b/examples/pytorch_lightning/spark.py index dfca5b220..62101dd5f 100644 --- a/examples/pytorch_lightning/spark.py +++ b/examples/pytorch_lightning/spark.py @@ -149,8 +149,6 @@ def configure_optimizers(self): model = SparseSparK(input_size=416) - -# we ignore object detection annotations by setting target_transform to return 0 dataset = torchvision.datasets.Caltech101( "datasets/caltech101", download=True, @@ -174,14 +172,10 @@ def configure_optimizers(self): num_workers=8, ) - -accelerator = "gpu" if torch.cuda.is_available() else "cpu" - trainer = pl.Trainer( max_epochs=30, - devices=1, - accelerator=accelerator, ) + trainer.fit( model=model, train_dataloaders=dataloader, From bde52922717ae7c6d040d541dabf0263a26abe6d Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Wed, 4 Mar 2026 14:35:17 -0300 Subject: [PATCH 66/85] typo: spark masking output --- lightly/models/modules/sparse_spark.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index a3821dadc..012720491 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -742,7 +742,7 @@ def forward(self, inp_bchw: torch.Tensor) -> SparKMaskingOutput: inp_bchw: Input image tensor of shape (B, C, H, W). Returns: - SparKMaskingOuptut containing: + SparKMaskingOutput containing: - masked_bchw: Input image masked at full resolution. - per_level_mask: List of masks at each hierarchical level. """ From 73b0bb160edbaa4c0fa7e20cfbe8897224b2735f Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Wed, 4 Mar 2026 14:38:11 -0300 Subject: [PATCH 67/85] fix: imports according to #1895 --- lightly/loss/sparse_spark.py | 9 +++-- lightly/models/modules/sparse_spark.py | 53 ++++++++++++-------------- 2 files changed, 30 insertions(+), 32 deletions(-) diff --git a/lightly/loss/sparse_spark.py b/lightly/loss/sparse_spark.py index 138d0c544..918f944d8 100644 --- a/lightly/loss/sparse_spark.py +++ b/lightly/loss/sparse_spark.py @@ -4,6 +4,7 @@ from torch import nn from lightly.utils.dist import gather import torch.distributed as dist +from torch import Tensor class SparKPatchReconLoss(nn.Module): @@ -29,10 +30,10 @@ def __init__(self, eps: float = 1e-6, gather_distributed: bool = False) -> None: def forward( self, - inp_patches: torch.Tensor, - rec_patches: torch.Tensor, - active_mask: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + 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 diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 012720491..f230ac5c2 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -10,6 +10,7 @@ import torch import torch.nn as nn +from torch import Tensor from timm.models.layers import DropPath, trunc_normal_ from lightly.models.utils import random_token_mask @@ -47,11 +48,11 @@ def coalesce_to_size_2_t(t: tuple[int, ...]) -> tuple[int, int]: raise ValueError(f"Invalid tuple length: {len(t)}; expected 1 or 2.") -_cur_active: ContextVar[torch.Tensor | None] = ContextVar("_cur_active", default=None) +_cur_active: ContextVar[Tensor | None] = ContextVar("_cur_active", default=None) @contextmanager -def sparse_layer_context(active_mask: torch.Tensor): +def sparse_layer_context(active_mask: Tensor): """Context manager to set the active mask for sparse layers. Args: @@ -66,7 +67,7 @@ def sparse_layer_context(active_mask: torch.Tensor): def _get_active_ex_or_ii( H: int, W: int, returning_active_ex: bool = True -) -> torch.Tensor | tuple[torch.Tensor, ...]: +) -> 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). @@ -109,8 +110,8 @@ def _get_active_ex_or_ii( def sp_conv_forward( - module: nn.AvgPool2d | nn.MaxPool2d | nn.Conv2d, input: torch.Tensor -) -> torch.Tensor: + 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 @@ -123,16 +124,14 @@ def sp_conv_forward( Returns: Masked output tensor of same shape as input, with inactive spatial positions zeroed. """ - x: torch.Tensor = super(type(module), module).forward(input) # type: ignore[arg-type] + x: Tensor = super(type(module), module).forward(input) # type: ignore[arg-type] x *= _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: torch.Tensor -) -> torch.Tensor: +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 @@ -235,7 +234,7 @@ def __init__( super().__init__(normalized_shape, eps, elementwise_affine=True) self.data_format = data_format - def forward(self, input: torch.Tensor) -> torch.Tensor: + 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( @@ -301,7 +300,7 @@ def __init__( DropPath(drop_path) if drop_path > 0.0 else nn.Identity() ) - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward(self, x: Tensor) -> Tensor: input = x x = self.dwconv(x) x = x.permute(0, 2, 3, 1) @@ -431,7 +430,7 @@ def __init__(self, cin: int, cout: int, bn2d: type) -> None: bn2d(cout), ) - def forward(self, x: torch.Tensor) -> torch.Tensor: + def forward(self, x: Tensor) -> Tensor: """Apply upsampling and convolution refinement. Args: @@ -476,7 +475,7 @@ def __init__( self.initialize() - def forward(self, to_dec: list[torch.Tensor]) -> torch.Tensor: + def forward(self, to_dec: list[Tensor]) -> Tensor: """Progressively upsample and combine feature maps. Args: @@ -586,9 +585,7 @@ def __init__( ) self.densify_proj = densify_proj - def forward( - self, bcff: torch.Tensor | None, cur_active: torch.Tensor - ) -> torch.Tensor | None: + def forward(self, bcff: Tensor | None, cur_active: Tensor) -> Tensor | None: """Densify sparse features by filling masked regions with learned tokens. Args: @@ -647,7 +644,7 @@ def __init__( self.blocks.append(block) d_width //= 2 - def forward(self, fea_bcffs: list[torch.Tensor]) -> list[torch.Tensor]: + def forward(self, fea_bcffs: list[Tensor]) -> list[Tensor]: """Convert sparse features to dense by filling masked regions. Args: @@ -685,8 +682,8 @@ class SparKMaskingOutput(NamedTuple): per_level_mask: List of binary masks at each hierarchical level. """ - masked_bchw: torch.Tensor - per_level_mask: list[torch.Tensor] + masked_bchw: Tensor + per_level_mask: list[Tensor] class SparKMasker(nn.Module): @@ -712,7 +709,7 @@ def __init__( self.downsample_ratio = downsample_ratio self.mask_ratio = mask_ratio - def mask(self, B: int, device: torch.device) -> torch.Tensor: + def mask(self, B: int, device: torch.device) -> Tensor: """Generate a random binary mask for features. Args: @@ -733,7 +730,7 @@ def mask(self, B: int, device: torch.device) -> torch.Tensor: .bool() ) - def forward(self, inp_bchw: torch.Tensor) -> SparKMaskingOutput: + 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. @@ -785,7 +782,7 @@ def __init__(self, fmap_h: int, fmap_w: int, downsample_ratio: int) -> None: self.fmap_w = fmap_w self.downsample_ratio = downsample_ratio - def unpatchify(self, bln: torch.Tensor) -> torch.Tensor: + 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) @@ -807,12 +804,12 @@ def unpatchify(self, bln: torch.Tensor) -> torch.Tensor: def forward( self, - rec_patches: torch.Tensor, - mean: torch.Tensor, - var: torch.Tensor, - inp_bchw: torch.Tensor, - active_mask_full: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + 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 From 86ea8a93f8164add15cd13f2ba747b59f04a3a35 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Wed, 4 Mar 2026 14:39:03 -0300 Subject: [PATCH 68/85] fix: format --- lightly/loss/sparse_spark.py | 6 +++--- lightly/models/modules/sparse_spark.py | 14 +++++++------- tests/loss/test_sparse_spark.py | 7 ++++--- 3 files changed, 14 insertions(+), 13 deletions(-) diff --git a/lightly/loss/sparse_spark.py b/lightly/loss/sparse_spark.py index 918f944d8..31c4371a7 100644 --- a/lightly/loss/sparse_spark.py +++ b/lightly/loss/sparse_spark.py @@ -1,10 +1,10 @@ from __future__ import annotations import torch -from torch import nn -from lightly.utils.dist import gather import torch.distributed as dist -from torch import Tensor +from torch import Tensor, nn + +from lightly.utils.dist import gather class SparKPatchReconLoss(nn.Module): diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index f230ac5c2..2b0acfb38 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -10,8 +10,8 @@ import torch import torch.nn as nn -from torch import Tensor from timm.models.layers import DropPath, trunc_normal_ +from torch import Tensor from lightly.models.utils import random_token_mask @@ -93,12 +93,12 @@ def _get_active_ex_or_ii( "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]}" - ) + 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 ) diff --git a/tests/loss/test_sparse_spark.py b/tests/loss/test_sparse_spark.py index 43c7d7bb7..38e41ab15 100644 --- a/tests/loss/test_sparse_spark.py +++ b/tests/loss/test_sparse_spark.py @@ -1,9 +1,10 @@ -import torch.distributed as dist -import torch import pytest -from lightly.loss import SparKPatchReconLoss +import torch +import torch.distributed as dist from pytest_mock import MockerFixture +from lightly.loss import SparKPatchReconLoss + class TestSparKPatchReconLoss: def test__gather_distributed(self, mocker: MockerFixture) -> None: From edbfdcf69195a6b839f473b39ec8c0cc1043ea6d Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Wed, 4 Mar 2026 17:44:38 -0300 Subject: [PATCH 69/85] fix: tests now uses context manager for sparse mask --- tests/models/modules/test_sparse_spark.py | 123 +++++++++------------- 1 file changed, 47 insertions(+), 76 deletions(-) diff --git a/tests/models/modules/test_sparse_spark.py b/tests/models/modules/test_sparse_spark.py index 12691bad9..f58013813 100644 --- a/tests/models/modules/test_sparse_spark.py +++ b/tests/models/modules/test_sparse_spark.py @@ -1,21 +1,11 @@ -from typing import Generator - import pytest import torch import lightly.models.modules.sparse_spark as sparse_spark -@pytest.fixture(autouse=True) -def _cleanup_cur_active() -> Generator[None, None, None]: - sparse_spark._cur_active = None - yield - sparse_spark._cur_active = None - - -def test__get_active_ex_or_ii_expands_mask() -> None: - H, W = 32, 32 - sparse_spark._cur_active = torch.tensor( +def _create_mask() -> torch.Tensor: + return torch.tensor( [ [ [ @@ -27,87 +17,68 @@ def test__get_active_ex_or_ii_expands_mask() -> None: dtype=torch.bool, ) - 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_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 - sparse_spark._cur_active = 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) + 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 - sparse_spark._cur_active = None - with pytest.raises( - AssertionError, - ): + 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 - sparse_spark._cur_active = torch.tensor( - [ - [ - [ - [1, 0], - [0, 1], - ] - ] - ], - dtype=torch.bool, - ) - - 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) + 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 + 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) + assert torch.equal(active_ex, active_ex_scattered) def test_sp_conv_forward() -> None: H, W = 32, 32 - sparse_spark._cur_active = torch.tensor( - [ - [ - [ - [1, 0], - [0, 1], - ] - ] - ], - dtype=torch.bool, - ) - - 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() + 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() From e28391a310b7a5b2dc0f6a018249701c5d237a01 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 7 Mar 2026 12:15:13 -0300 Subject: [PATCH 70/85] fix: using Tensor instead of torch.Tensor --- tests/models/modules/test_sparse_spark.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/models/modules/test_sparse_spark.py b/tests/models/modules/test_sparse_spark.py index f58013813..8568c261b 100644 --- a/tests/models/modules/test_sparse_spark.py +++ b/tests/models/modules/test_sparse_spark.py @@ -2,9 +2,10 @@ import torch import lightly.models.modules.sparse_spark as sparse_spark +from torch import Tensor -def _create_mask() -> torch.Tensor: +def _create_mask() -> Tensor: return torch.tensor( [ [ @@ -61,7 +62,7 @@ def test__get_active_ex_or_ii_returning_ex_false_correct_values() -> None: assert torch.equal(active_ex, active_ex_scattered) -def test_sp_conv_forward() -> None: +def test__sp_conv_forward() -> None: H, W = 32, 32 with sparse_spark.sparse_layer_context(_create_mask()): conv = sparse_spark.SparseConv2d( From a1f82a0d18eb93833855cfa6274579abe47241fe Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 7 Mar 2026 12:48:32 -0300 Subject: [PATCH 71/85] doc: referencing original author --- lightly/loss/sparse_spark.py | 2 ++ lightly/models/modules/sparse_spark.py | 16 ++++++++++++++++ 2 files changed, 18 insertions(+) diff --git a/lightly/loss/sparse_spark.py b/lightly/loss/sparse_spark.py index 31c4371a7..664147a57 100644 --- a/lightly/loss/sparse_spark.py +++ b/lightly/loss/sparse_spark.py @@ -10,6 +10,8 @@ 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. diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 2b0acfb38..8d7a609df 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -323,6 +323,8 @@ def dense_model_to_sparse( ) -> 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. @@ -408,6 +410,8 @@ def dense_model_to_sparse( 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. @@ -446,6 +450,8 @@ def forward(self, x: Tensor) -> Tensor: 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. @@ -525,6 +531,8 @@ def initialize(self) -> None: 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. @@ -608,6 +616,8 @@ def forward(self, bcff: Tensor | None, cur_active: Tensor) -> Tensor | None: 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. @@ -677,6 +687,8 @@ def forward(self, fea_bcffs: list[Tensor]) -> list[Tensor]: 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. @@ -689,6 +701,8 @@ class SparKMaskingOutput(NamedTuple): 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. @@ -763,6 +777,8 @@ def forward(self, inp_bchw: Tensor) -> SparKMaskingOutput: 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). From 9ce0f803553b9c2e035dcaf23b1660f56e1c5b9a Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 7 Mar 2026 12:55:22 -0300 Subject: [PATCH 72/85] test: spark masking test --- tests/models/modules/test_sparse_spark.py | 64 +++++++++++++++++++++++ 1 file changed, 64 insertions(+) diff --git a/tests/models/modules/test_sparse_spark.py b/tests/models/modules/test_sparse_spark.py index 8568c261b..3172379cd 100644 --- a/tests/models/modules/test_sparse_spark.py +++ b/tests/models/modules/test_sparse_spark.py @@ -2,6 +2,7 @@ import torch import lightly.models.modules.sparse_spark as sparse_spark +from lightly.models.modules.sparse_spark import SparKMasker, SparKMaskingOutput from torch import Tensor @@ -83,3 +84,66 @@ def test__sp_conv_forward() -> None: 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, 65, dtype=torch.float32).view(1, 1, 8, 8) + 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]) From e66d84a09b0dc97915d5f994d89b706b31e41dc7 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 7 Mar 2026 13:04:44 -0300 Subject: [PATCH 73/85] doc: removed unecessary comment --- lightly/models/modules/sparse_spark.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 8d7a609df..193645f50 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -783,9 +783,6 @@ class SparKOutputDecoder(nn.Module): per-pixel blending: uses original pixels where visible (active), reconstructed pixels where masked (inactive). - Minimal configuration: only requires spatial properties (fmap_h, fmap_w, downsample_ratio). - No encoder object needed. - Args: fmap_h: Height of feature map at patch level. fmap_w: Width of feature map at patch level. From 86934d29ce1a49ecb4a94c9cdfcfedbdccde8b5a Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 7 Mar 2026 14:48:35 -0300 Subject: [PATCH 74/85] feat: testing densify block --- lightly/models/modules/sparse_spark.py | 18 +++- tests/models/modules/test_sparse_spark.py | 104 ++++++++++++++++++++-- 2 files changed, 112 insertions(+), 10 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 193645f50..7342f1ead 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -593,6 +593,20 @@ def __init__( ) 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 | None, cur_active: Tensor) -> Tensor | None: """Densify sparse features by filling masked regions with learned tokens. @@ -606,9 +620,7 @@ def forward(self, bcff: Tensor | None, cur_active: Tensor) -> Tensor | None: if bcff is None: return None bcff = self.densify_norm(bcff) - assert bcff is not None - mask_tokens = self.mask_token.expand_as(bcff) - bcff = torch.where(cur_active.expand_as(bcff), bcff, mask_tokens) + bcff = self._fill_with_mask_tokens(bcff, cur_active) bcff = self.densify_proj(bcff) return bcff diff --git a/tests/models/modules/test_sparse_spark.py b/tests/models/modules/test_sparse_spark.py index 3172379cd..c02f2373e 100644 --- a/tests/models/modules/test_sparse_spark.py +++ b/tests/models/modules/test_sparse_spark.py @@ -1,9 +1,13 @@ import pytest import torch +from torch import Tensor import lightly.models.modules.sparse_spark as sparse_spark -from lightly.models.modules.sparse_spark import SparKMasker, SparKMaskingOutput -from torch import Tensor +from lightly.models.modules.sparse_spark import ( + SparKDensifiyBlock, + SparKMasker, + SparKMaskingOutput, +) def _create_mask() -> Tensor: @@ -85,6 +89,7 @@ def test__sp_conv_forward() -> None: assert out[:, :, 16:, :16].logical_not().all() assert out[:, :, 16:, 16:].all() + class TestSparKMasker: def test_forward(self) -> None: masker = SparKMasker( @@ -99,15 +104,18 @@ def test_forward(self) -> None: 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) + 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): + 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] + 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() + 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( @@ -147,3 +155,85 @@ def test_batch_independence(self) -> None: 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]) From a984e0d3f57665c3a23b7e6964eec1dd64d6557b Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Sat, 7 Mar 2026 15:03:23 -0300 Subject: [PATCH 75/85] fix: removed unecessary none handling --- lightly/models/modules/sparse_spark.py | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 7342f1ead..e4b8b99a2 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -607,7 +607,7 @@ def _fill_with_mask_tokens(self, features: Tensor, active_mask: Tensor) -> Tenso active_expanded: Tensor = active_mask.expand(features.size()) return torch.where(active_expanded, features, mask_tokens) - def forward(self, bcff: Tensor | None, cur_active: Tensor) -> Tensor | None: + def forward(self, bcff: Tensor, cur_active: Tensor) -> Tensor: """Densify sparse features by filling masked regions with learned tokens. Args: @@ -615,10 +615,8 @@ def forward(self, bcff: Tensor | None, cur_active: Tensor) -> Tensor | None: cur_active: Boolean mask tensor of shape (B, 1, H, W) indicating active regions. Returns: - Densified and projected feature tensor, or None if input is None. + Densified and projected feature tensor of shape (B, C', H, W). """ - if bcff is None: - return None bcff = self.densify_norm(bcff) bcff = self._fill_with_mask_tokens(bcff, cur_active) bcff = self.densify_proj(bcff) @@ -687,8 +685,7 @@ def forward(self, fea_bcffs: list[Tensor]) -> list[Tensor]: ) active_fmap_current = active_mask for i, bcff in enumerate(fea_bcffs): - if bcff is not None: - bcff = self.blocks[i](bcff, active_fmap_current) + bcff = self.blocks[i](bcff, active_fmap_current) to_dec.append(bcff) active_fmap_current = active_fmap_current.repeat_interleave( 2, dim=2 From 09103d66318c1c321adde4b890d328c542e5ed61 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Wed, 1 Apr 2026 09:02:35 -0300 Subject: [PATCH 76/85] fix(test): sparse spark test with incorrect shapes --- tests/models/modules/test_sparse_spark.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/models/modules/test_sparse_spark.py b/tests/models/modules/test_sparse_spark.py index c02f2373e..847e640f4 100644 --- a/tests/models/modules/test_sparse_spark.py +++ b/tests/models/modules/test_sparse_spark.py @@ -122,7 +122,7 @@ def test_masked_bchw_applies_mask(self) -> None: feature_map_size=(4, 4), downsample_ratio=8, ) - x = torch.arange(1, 65, dtype=torch.float32).view(1, 1, 8, 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 From e43d7a3d0b41f5a5a7e748938aee3a0bc745c265 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Wed, 1 Apr 2026 09:08:44 -0300 Subject: [PATCH 77/85] test: spark densifier --- tests/models/modules/test_sparse_spark.py | 170 ++++++++++++++++++++++ 1 file changed, 170 insertions(+) diff --git a/tests/models/modules/test_sparse_spark.py b/tests/models/modules/test_sparse_spark.py index 847e640f4..7730777e8 100644 --- a/tests/models/modules/test_sparse_spark.py +++ b/tests/models/modules/test_sparse_spark.py @@ -237,3 +237,173 @@ def test_fill_with_mask_tokens_batch_independence(self) -> None: # 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 From ef88d62051c703089f3afbd91fd0b821682ccbd5 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 2 Apr 2026 09:29:26 -0300 Subject: [PATCH 78/85] format --- tests/models/modules/test_sparse_spark.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/tests/models/modules/test_sparse_spark.py b/tests/models/modules/test_sparse_spark.py index 7730777e8..03b86f7d9 100644 --- a/tests/models/modules/test_sparse_spark.py +++ b/tests/models/modules/test_sparse_spark.py @@ -290,7 +290,7 @@ def test_forward_upsamples_active_mask(self) -> None: # 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 + 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.""" @@ -324,9 +324,9 @@ def test_forward_to_dec_contains_all_feature_maps(self) -> None: # 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]}" - ) + 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).""" @@ -339,7 +339,7 @@ def test_forward_reverses_fea_bcffs_order(self) -> None: ) # Use distinct shapes to verify order fea_bcffs: list[Tensor] = [ - torch.randn(1, 64, 8, 8), # shallowest + torch.randn(1, 64, 8, 8), # shallowest torch.randn(1, 128, 4, 4), torch.randn(1, 256, 2, 2), # deepest ] @@ -372,7 +372,9 @@ def test_forward_complete_integration(self) -> None: 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) + 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) From 5e8f28efb803b5af690e5d500fd310300ef7c0c5 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 2 Apr 2026 09:30:22 -0300 Subject: [PATCH 79/85] feat: notebook example --- .../notebooks/pytorch_lightning/spark.ipynb | 302 ++++++++++++++++++ 1 file changed, 302 insertions(+) create mode 100644 examples/notebooks/pytorch_lightning/spark.ipynb 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 +} From 441e6086f66e09d22a473a8e540b7e0349c6c40c Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Thu, 2 Apr 2026 10:02:56 -0300 Subject: [PATCH 80/85] format and mypy --- lightly/models/modules/sparse_spark.py | 94 +++++++++++++++++--------- 1 file changed, 62 insertions(+), 32 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index e4b8b99a2..e7212141b 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -6,7 +6,7 @@ import math from contextlib import contextmanager from contextvars import ContextVar -from typing import NamedTuple +from typing import Generator, NamedTuple, cast import torch import torch.nn as nn @@ -41,9 +41,9 @@ def coalesce_to_size_2_t(t: tuple[int, ...]) -> tuple[int, int]: ValueError: If tuple length is not 1 or 2. """ if len(t) == 2: - return t + return (t[0], t[1]) elif len(t) == 1: - return t[0], t[0] + return (t[0], t[0]) else: raise ValueError(f"Invalid tuple length: {len(t)}; expected 1 or 2.") @@ -52,7 +52,9 @@ def coalesce_to_size_2_t(t: tuple[int, ...]) -> tuple[int, int]: @contextmanager -def sparse_layer_context(active_mask: Tensor): +def sparse_layer_context( + active_mask: Tensor, +) -> Generator[None, None, None]: """Context manager to set the active mask for sparse layers. Args: @@ -125,8 +127,11 @@ def sp_conv_forward( 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 *= _get_active_ex_or_ii( - H=input.shape[2], W=input.shape[3], returning_active_ex=True + x *= cast( + Tensor, + _get_active_ex_or_ii( + H=input.shape[2], W=input.shape[3], returning_active_ex=True + ), ) return x @@ -312,7 +317,10 @@ def forward(self, x: Tensor) -> Tensor: x = self.gamma * x x = x.permute(0, 3, 1, 2) - x *= _get_active_ex_or_ii(H=x.shape[2], W=x.shape[3], returning_active_ex=True) + 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 @@ -339,11 +347,11 @@ def dense_model_to_sparse( Raises: NotImplementedError: If Conv1d layers are encountered. """ - oup = m + oup: nn.Module = m if isinstance(m, nn.Conv2d): bias = m.bias is not None - oup = SparseConv2d( + sparse_conv: SparseConv2d = SparseConv2d( m.in_channels, m.out_channels, kernel_size=coalesce_to_size_2_t(m.kernel_size), @@ -356,10 +364,11 @@ def dense_model_to_sparse( bias=bias, padding_mode=m.padding_mode, ) - oup.weight.data.copy_(m.weight.data) + sparse_conv.weight.data.copy_(m.weight.data) - if m.bias is not None and oup.bias is not None: - oup.bias.data.copy_(m.bias.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( @@ -380,29 +389,50 @@ def dense_model_to_sparse( divisor_override=m.divisor_override, ) elif isinstance(m, (nn.BatchNorm2d, nn.SyncBatchNorm)): - oup = (SparseSyncBatchNorm2d if sbn else SparseBatchNorm2d)( - m.weight.shape[0], - eps=m.eps, - momentum=m.momentum, - affine=m.affine, - track_running_stats=m.track_running_stats, - ) - oup.weight.data.copy_(m.weight.data) - oup.bias.data.copy_(m.bias.data) - oup.running_mean.data.copy_(m.running_mean.data) - oup.running_var.data.copy_(m.running_var.data) - oup.num_batches_tracked.data.copy_(m.num_batches_tracked.data) - if hasattr(m, "qconfig"): - oup.qconfig = m.qconfig + 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): - oup = SparseConvNeXtLayerNorm(m.weight.shape[0], eps=m.eps) - oup.weight.data.copy_(m.weight.data) - oup.bias.data.copy_(m.bias.data) + 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)) + oup.add_module(name, dense_model_to_sparse(child, verbose=verbose, sbn=sbn)) # type: ignore[union-attr] del m return oup @@ -444,7 +474,7 @@ def forward(self, x: Tensor) -> Tensor: Upsampled and refined feature tensor. """ x = self.up_sample(x) - return self.conv(x) + return self.conv(x) # type: ignore[no-any-return] class UNetDecoder(nn.Module): @@ -495,7 +525,7 @@ def forward(self, to_dec: list[Tensor]) -> Tensor: 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) + return self.proj(x) # type: ignore[no-any-return] def extra_repr(self) -> str: return f"width={self.width}" From f21faaa73201ab3c30096b5192f658b38b122884 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 3 Apr 2026 09:30:32 -0300 Subject: [PATCH 81/85] fix: skipping tests on minimal and scoped imports --- lightly/models/modules/sparse_spark.py | 5 ++++- tests/loss/test_sparse_spark.py | 2 ++ tests/models/modules/test_sparse_spark.py | 2 ++ 3 files changed, 8 insertions(+), 1 deletion(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index e7212141b..e30ce3397 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -10,7 +10,6 @@ import torch import torch.nn as nn -from timm.models.layers import DropPath, trunc_normal_ from torch import Tensor from lightly.models.utils import random_token_mask @@ -290,6 +289,8 @@ def __init__( 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) @@ -538,6 +539,8 @@ def initialize(self) -> None: - 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) diff --git a/tests/loss/test_sparse_spark.py b/tests/loss/test_sparse_spark.py index 38e41ab15..b5166361c 100644 --- a/tests/loss/test_sparse_spark.py +++ b/tests/loss/test_sparse_spark.py @@ -3,6 +3,8 @@ import torch.distributed as dist from pytest_mock import MockerFixture +pytest.importorskip("timm.models.layers") + from lightly.loss import SparKPatchReconLoss diff --git a/tests/models/modules/test_sparse_spark.py b/tests/models/modules/test_sparse_spark.py index 03b86f7d9..ab5dde196 100644 --- a/tests/models/modules/test_sparse_spark.py +++ b/tests/models/modules/test_sparse_spark.py @@ -2,6 +2,8 @@ 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, From 3e962a81945563c8de7cf77a5384183985d384f0 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 3 Apr 2026 16:12:33 -0300 Subject: [PATCH 82/85] fix: missing trunc_normal_ import --- lightly/models/modules/sparse_spark.py | 1 + 1 file changed, 1 insertion(+) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index e30ce3397..c8c6e2796 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -587,6 +587,7 @@ def __init__( 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)) From 4e5e06761672f4fdf28473485c14b26bf88c8727 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Fri, 3 Apr 2026 16:24:55 -0300 Subject: [PATCH 83/85] Fix Python 3.7 type hints compatibility in sparse_spark tests Use List[T] instead of list[T] for type hints to support Python 3.7. Add type: ignore for mypy no-any-return in reference loss implementation. --- tests/loss/test_sparse_spark.py | 2 +- tests/models/modules/test_sparse_spark.py | 33 ++++++++++++----------- 2 files changed, 18 insertions(+), 17 deletions(-) diff --git a/tests/loss/test_sparse_spark.py b/tests/loss/test_sparse_spark.py index b5166361c..766878455 100644 --- a/tests/loss/test_sparse_spark.py +++ b/tests/loss/test_sparse_spark.py @@ -80,4 +80,4 @@ def _reference_loss_implementation( 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 + 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 index ab5dde196..bbaafa30b 100644 --- a/tests/models/modules/test_sparse_spark.py +++ b/tests/models/modules/test_sparse_spark.py @@ -1,6 +1,7 @@ import pytest import torch from torch import Tensor +from typing import List pytest.importorskip("timm.models.layers") @@ -251,7 +252,7 @@ def test_forward_raises_without_active_mask(self) -> None: sbn=False, ) # Create dummy feature maps (shallow to deep) - fea_bcffs: list[Tensor] = [ + fea_bcffs: List[Tensor] = [ torch.randn(1, 64, 8, 8), torch.randn(1, 128, 4, 4), torch.randn(1, 256, 2, 2), @@ -270,7 +271,7 @@ def test_forward_upsamples_active_mask(self) -> None: sbn=False, ) # Create dummy feature maps at different scales (shallow to deep) - fea_bcffs: list[Tensor] = [ + fea_bcffs: List[Tensor] = [ torch.randn(1, 64, 8, 8), torch.randn(1, 128, 4, 4), torch.randn(1, 256, 2, 2), @@ -281,7 +282,7 @@ def test_forward_upsamples_active_mask(self) -> None: ) with sparse_spark.sparse_layer_context(initial_mask): - to_dec: list[Tensor] = densifier(fea_bcffs) + to_dec: List[Tensor] = densifier(fea_bcffs) # Check that we have 3 output feature maps (one per block) assert len(to_dec) == 3 @@ -296,7 +297,7 @@ def test_forward_upsamples_active_mask(self) -> None: 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] + 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, @@ -304,7 +305,7 @@ def test_forward_to_dec_contains_all_feature_maps(self) -> None: densify_norm_str="bn", sbn=False, ) - fea_bcffs: list[Tensor] = [ + fea_bcffs: List[Tensor] = [ torch.randn(1, 64, 8, 8), torch.randn(1, 128, 4, 4), torch.randn(1, 256, 2, 2), @@ -312,7 +313,7 @@ def test_forward_to_dec_contains_all_feature_maps(self) -> None: 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) + to_dec: List[Tensor] = densifier(fea_bcffs) # Should have one output per block assert len(to_dec) == len(densifier.blocks) @@ -324,7 +325,7 @@ def test_forward_to_dec_contains_all_feature_maps(self) -> None: # 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] + expected_channels: List[int] = [256, 128, 64] for i, td in enumerate(to_dec): assert ( td.shape[1] == expected_channels[i] @@ -332,7 +333,7 @@ def test_forward_to_dec_contains_all_feature_maps(self) -> None: 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] + 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, @@ -340,7 +341,7 @@ def test_forward_reverses_fea_bcffs_order(self) -> None: sbn=False, ) # Use distinct shapes to verify order - fea_bcffs: list[Tensor] = [ + fea_bcffs: List[Tensor] = [ torch.randn(1, 64, 8, 8), # shallowest torch.randn(1, 128, 4, 4), torch.randn(1, 256, 2, 2), # deepest @@ -348,7 +349,7 @@ def test_forward_reverses_fea_bcffs_order(self) -> None: 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) + 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) @@ -359,7 +360,7 @@ def test_forward_reverses_fea_bcffs_order(self) -> None: def test_forward_complete_integration(self) -> None: """Integration test for complete forward pass with proper context.""" - encoder_in_channels: list[int] = [64, 128, 256] + 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, @@ -368,7 +369,7 @@ def test_forward_complete_integration(self) -> None: sbn=False, ) batch_size: int = 2 - fea_bcffs: list[Tensor] = [ + 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), @@ -379,7 +380,7 @@ def test_forward_complete_integration(self) -> None: ) with sparse_spark.sparse_layer_context(initial_mask): - to_dec: list[Tensor] = densifier(fea_bcffs) + to_dec: List[Tensor] = densifier(fea_bcffs) # Verify output structure assert len(to_dec) == 3 @@ -392,14 +393,14 @@ def test_forward_complete_integration(self) -> None: 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 + 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] = [ + fea_bcffs: List[Tensor] = [ torch.randn(1, 256, 8, 8), torch.randn(1, 128, 4, 4), torch.randn(1, 64, 2, 2), @@ -407,7 +408,7 @@ def test_forward_with_identity_proj_first_block(self) -> None: 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) + to_dec: List[Tensor] = densifier(fea_bcffs) # First block should preserve channels (identity projection) assert to_dec[0].shape[1] == 256 From e5f0b04a411e750e5e4bd8abe636337936d70266 Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Tue, 7 Apr 2026 17:02:31 -0300 Subject: [PATCH 84/85] format --- lightly/models/modules/sparse_spark.py | 1 + tests/models/modules/test_sparse_spark.py | 3 ++- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index c8c6e2796..74c2e0dd8 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -588,6 +588,7 @@ def __init__( 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)) diff --git a/tests/models/modules/test_sparse_spark.py b/tests/models/modules/test_sparse_spark.py index bbaafa30b..8119e4fc7 100644 --- a/tests/models/modules/test_sparse_spark.py +++ b/tests/models/modules/test_sparse_spark.py @@ -1,7 +1,8 @@ +from typing import List + import pytest import torch from torch import Tensor -from typing import List pytest.importorskip("timm.models.layers") From 96870358288ae7448770ce18d6fc0d6e830020be Mon Sep 17 00:00:00 2001 From: gabrielfruet Date: Wed, 8 Apr 2026 09:26:23 -0300 Subject: [PATCH 85/85] doc: removed copyright docstring, we already reference authoer on the classes --- lightly/models/modules/sparse_spark.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/lightly/models/modules/sparse_spark.py b/lightly/models/modules/sparse_spark.py index 74c2e0dd8..89f44afa7 100644 --- a/lightly/models/modules/sparse_spark.py +++ b/lightly/models/modules/sparse_spark.py @@ -1,6 +1,3 @@ -# Copyright (c) 2023 Keyu Tian -# Copyright (c) ByteDance, Inc. and its affiliates. - from __future__ import annotations import math