Upload folder using huggingface_hub

.ipynb_checkpoints/__init__-checkpoint.py

@@ -0,0 +1 @@
+from .vision_transformer import VisionTransformer, vit_tiny, vit_small, vit_base, vit_large

.ipynb_checkpoints/head-checkpoint.py

@@ -0,0 +1,241 @@
+import torch
+import torch.nn as nn
+import utils
+
+from utils import trunc_normal_
+
+class CSyncBatchNorm(nn.SyncBatchNorm):
+    def __init__(self,
+                 *args,
+                 with_var=False,
+                 **kwargs):
+        super(CSyncBatchNorm, self).__init__(*args, **kwargs)
+        self.with_var = with_var
+
+    def forward(self, x):
+        # center norm
+        self.training = False
+        if not self.with_var:
+            self.running_var = torch.ones_like(self.running_var)
+        normed_x = super(CSyncBatchNorm, self).forward(x)
+        # udpate center
+        self.training = True
+        _ = super(CSyncBatchNorm, self).forward(x)
+        return normed_x
+
+class PSyncBatchNorm(nn.SyncBatchNorm):
+    def __init__(self,
+                 *args,
+                 bunch_size,
+                 **kwargs):
+        procs_per_bunch = min(bunch_size, utils.get_world_size())
+        assert utils.get_world_size() % procs_per_bunch == 0
+        n_bunch = utils.get_world_size() // procs_per_bunch
+        #
+        ranks = list(range(utils.get_world_size()))
+        print('---ALL RANKS----\n{}'.format(ranks))
+        rank_groups = [ranks[i*procs_per_bunch: (i+1)*procs_per_bunch] for i in range(n_bunch)]
+        print('---RANK GROUPS----\n{}'.format(rank_groups))
+        process_groups = [torch.distributed.new_group(pids) for pids in rank_groups]
+        bunch_id = utils.get_rank() // procs_per_bunch
+        process_group = process_groups[bunch_id]
+        print('---CURRENT GROUP----\n{}'.format(process_group))
+        super(PSyncBatchNorm, self).__init__(*args, process_group=process_group, **kwargs)
+
+class CustomSequential(nn.Sequential):
+    bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm)
+
+    def forward(self, input):
+        for module in self:
+            dim = len(input.shape)
+            if isinstance(module, self.bn_types) and dim > 2:
+                perm = list(range(dim - 1)); perm.insert(1, dim - 1)
+                inv_perm = list(range(dim)) + [1]; inv_perm.pop(1)
+                input = module(input.permute(*perm)).permute(*inv_perm)
+            else:
+                input = module(input)
+        return input
+
+class DINOHead(nn.Module):
+    def __init__(self, in_dim, out_dim, norm=None, act='gelu', last_norm=None, 
+                 nlayers=3, hidden_dim=2048, bottleneck_dim=256, norm_last_layer=True, **kwargs):
+        super().__init__()
+        norm = self._build_norm(norm, hidden_dim)
+        last_norm = self._build_norm(last_norm, out_dim, affine=False, **kwargs)
+        act = self._build_act(act)
+
+        nlayers = max(nlayers, 1)
+        if nlayers == 1:
+            if bottleneck_dim > 0:
+                self.mlp = nn.Linear(in_dim, bottleneck_dim)
+            else:
+                self.mlp = nn.Linear(in_dim, out_dim)
+        else:
+            layers = [nn.Linear(in_dim, hidden_dim)]
+            if norm is not None:
+                layers.append(norm)
+            layers.append(act)
+            for _ in range(nlayers - 2):
+                layers.append(nn.Linear(hidden_dim, hidden_dim))
+                if norm is not None:
+                    layers.append(norm)
+                layers.append(act)
+            if bottleneck_dim > 0:
+                layers.append(nn.Linear(hidden_dim, bottleneck_dim))
+            else:
+                layers.append(nn.Linear(hidden_dim, out_dim))
+            self.mlp = CustomSequential(*layers)
+        self.apply(self._init_weights)
+        
+        if bottleneck_dim > 0:
+            self.last_layer = nn.utils.weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))
+            self.last_layer.weight_g.data.fill_(1)
+            if norm_last_layer:
+                self.last_layer.weight_g.requires_grad = False
+        else:
+            self.last_layer = None
+
+        self.last_norm = last_norm
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x = self.mlp(x)
+        if self.last_layer is not None:
+            x = nn.functional.normalize(x, dim=-1, p=2)
+            x = self.last_layer(x)
+        if self.last_norm is not None:
+            x = self.last_norm(x)
+        return x
+
+    def _build_norm(self, norm, hidden_dim, **kwargs):
+        if norm == 'bn':
+            norm = nn.BatchNorm1d(hidden_dim, **kwargs)
+        elif norm == 'syncbn':
+            norm = nn.SyncBatchNorm(hidden_dim, **kwargs)
+        elif norm == 'csyncbn':
+            norm = CSyncBatchNorm(hidden_dim, **kwargs)
+        elif norm == 'psyncbn':
+            norm =  PSyncBatchNorm(hidden_dim, **kwargs)
+        elif norm == 'ln':
+            norm = nn.LayerNorm(hidden_dim, **kwargs)
+        else:
+            assert norm is None, "unknown norm type {}".format(norm)
+        return norm
+
+    def _build_act(self, act):
+        if act == 'relu':
+            act = nn.ReLU()
+        elif act == 'gelu':
+            act = nn.GELU()
+        else:
+            assert False, "unknown act type {}".format(act)
+        return act
+
+class iBOTHead(DINOHead):
+
+    def __init__(self, *args, patch_out_dim=8192, norm=None, act='gelu', last_norm=None, 
+                 nlayers=3, hidden_dim=2048, bottleneck_dim=256, norm_last_layer=True, 
+                 shared_head=False, **kwargs):
+        
+        super(iBOTHead, self).__init__(*args,
+                                        norm=norm,
+                                        act=act,
+                                        last_norm=last_norm,
+                                        nlayers=nlayers,
+                                        hidden_dim=hidden_dim,
+                                        bottleneck_dim=bottleneck_dim,
+                                        norm_last_layer=norm_last_layer, 
+                                        **kwargs)
+
+        if not shared_head:
+            if bottleneck_dim > 0:
+                self.last_layer2 = nn.utils.weight_norm(nn.Linear(bottleneck_dim, patch_out_dim, bias=False))
+                self.last_layer2.weight_g.data.fill_(1)
+                if norm_last_layer:
+                    self.last_layer2.weight_g.requires_grad = False
+            else:
+                self.mlp2 = nn.Linear(hidden_dim, patch_out_dim)
+                self.last_layer2 = None
+
+            self.last_norm2 = self._build_norm(last_norm, patch_out_dim, affine=False, **kwargs)
+        else:
+            if bottleneck_dim > 0:
+                self.last_layer2 = self.last_layer
+            else:
+                self.mlp2 = self.mlp[-1]
+                self.last_layer2 = None
+
+            self.last_norm2 = self.last_norm
+
+    def forward(self, x):
+        if len(x.shape) == 2:
+            return super(iBOTHead, self).forward(x)
+
+        if self.last_layer is not None:
+            x = self.mlp(x)
+            x = nn.functional.normalize(x, dim=-1, p=2)
+            x1 = self.last_layer(x[:, 0])
+            x2 = self.last_layer2(x[:, 1:])
+        else:
+            x = self.mlp[:-1](x)
+            x1 = self.mlp[-1](x[:, 0])
+            x2 = self.mlp2(x[:, 1:])
+        
+        if self.last_norm is not None:
+            x1 = self.last_norm(x1)
+            x2 = self.last_norm2(x2)
+        
+        return x1, x2
+
+
+
+class TemporalSideContext(nn.Module):
+    def __init__(self, D, max_len=64, n_layers=6, n_head=8, dropout=0.1):
+        super().__init__()
+        #self.pos_t = nn.Embedding(max_len, D)  # learnable embedding for positions
+        layer = nn.TransformerEncoderLayer(D, n_head, 4*D,
+                                           dropout=dropout, batch_first=True)
+        self.enc = nn.TransformerEncoder(layer, n_layers)
+
+    def forward(self, x):          # x [B,T,D]
+        B,T,D = x.shape
+        device = x.device
+        # Generate relative frame positions [0, 1, ..., T-1]
+        #pos_ids = torch.arange(T, device=device).unsqueeze(0)  # [1, T]
+        #pos_embed = self.pos_t(pos_ids)                        # [1, T, D]
+        #x = x + pos_embed
+        return self.enc(x)          # [B,T,D]
+
+
+
+class TemporalHead(nn.Module):
+    """
+    Converts backbone features [B,T,D] → logits [B,T,1] for Plackett–Luce.
+    """
+    def __init__(self, backbone_dim: int, hidden_mul: float = 0.5, max_len: int = 64):
+        super().__init__()
+        hidden_dim = int(backbone_dim * hidden_mul)
+
+        self.reduce = nn.Sequential(
+            nn.Linear(backbone_dim, hidden_dim),
+            nn.GELU()
+        )
+        self.temporal = TemporalSideContext(hidden_dim, max_len=max_len)
+        self.scorer  = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.GELU(),
+            nn.Linear(hidden_dim // 2, 1)
+        )
+
+    def forward(self, x: torch.Tensor):     # x : [B,T,D]
+        x = self.reduce(x)                  # [B,T,hidden]
+        x = self.temporal(x)                # [B,T,hidden]
+        return self.scorer(x)               # [B,T,1]
+
+
+

.ipynb_checkpoints/puzzle_decoder-checkpoint.py

@@ -0,0 +1,62 @@
+import torch.nn as nn
+
+class AttentionBlock(nn.Module):
+    def __init__(self, attn_mul=1, embed_dim=384, num_heads=16):
+        super().__init__()
+        self.attn_mul = attn_mul
+        self.auxiliary_cross_attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
+        self.auxiliary_layer_norm1 = nn.LayerNorm(embed_dim)
+        self.auxiliary_self_attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
+        self.auxiliary_layer_norm2 = nn.LayerNorm(embed_dim)
+        self.auxiliary_linear = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.auxiliary_layer_norm3 = nn.LayerNorm(embed_dim)
+
+    def forward(self, current_patch, neighbor_patch):
+        # shape of aux_crop and stu_crop: (batch_size, seq_len, embed_dim)
+
+        # MultiheadAttention takes in the query, key, value. Here we use stu_crop to attend to aux_crop.
+        cross_attn_output, cross_attn_output_weights = self.auxiliary_cross_attn(current_patch, neighbor_patch, neighbor_patch)
+        cross_attn_output = self.auxiliary_layer_norm1(self.attn_mul * cross_attn_output + current_patch) # layer norm with skip connection
+        
+        # Then we use cross_attn_output to attend to cross_attn_output itself
+        self_attn_output, self_attn_output_weights = self.auxiliary_self_attn(cross_attn_output, cross_attn_output, cross_attn_output)
+        self_attn_output = self.auxiliary_layer_norm2(self_attn_output + cross_attn_output) # layer norm with skip connection
+
+        # Finally, apply feed forward.
+        output = self.auxiliary_linear(self_attn_output)
+        output = self.auxiliary_layer_norm3(output + self_attn_output) # layer norm with skip connection
+        
+        return output, cross_attn_output_weights, self_attn_output_weights
+
+
+
+class PuzzleDecoder(nn.Module):
+    def __init__(self, attn_mul=4, num_blocks=1, embed_dim=384, num_heads=16):
+        super().__init__()
+        self.decoder = nn.ModuleList([AttentionBlock(attn_mul, embed_dim, num_heads) for _ in range(num_blocks)])
+        self.scorer = nn.Sequential(
+            nn.LayerNorm(embed_dim),
+            nn.Linear(embed_dim, embed_dim),
+            nn.GELU(),
+            nn.Linear(embed_dim, 1),
+        )
+        nn.init.trunc_normal_(self.scorer[-1].weight, std=0.02)
+        nn.init.zeros_(self.scorer[-1].bias)
+
+    def forward(self, current_patch, neighbor_patch, return_feats=False, return_attn=False):
+        x = current_patch
+        cross_ws, self_ws = [], []
+        for block in self.decoder:
+            x, cw, sw = block(x, neighbor_patch)
+            if return_attn:
+                cross_ws.append(cw); self_ws.append(sw)
+
+        scores = self.scorer(x).squeeze(-1)        # [B, N]
+        if return_feats or return_attn:
+            out = {"scores": scores}
+            if return_feats: out["feats"] = x
+            if return_attn:  out["cross_w"] = cross_ws; out["self_w"] = self_ws
+            return out
+
+        return scores
+

.ipynb_checkpoints/vision_transformer-checkpoint.py

@@ -0,0 +1,283 @@
+import math
+import torch
+import torch.nn as nn
+
+from functools import partial
+from utils import trunc_normal_
+from timm.models.registry import register_model
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    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).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x, attn
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., 
+                 attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, init_values=0):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if init_values > 0:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+    def forward(self, x, return_attention=False):
+        y, attn = self.attn(self.norm1(x))
+        if return_attention:
+            return attn
+        if self.gamma_1 is None:
+            x = x + self.drop_path(y)
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            x = x + self.drop_path(self.gamma_1 * y)
+            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        num_patches = (img_size // patch_size) * (img_size // patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+            
+    def forward(self, x):
+        B, C, H, W = x.shape
+        return self.proj(x)
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer """
+    def __init__(self, img_size=[224], patch_size=16, in_chans=3, num_classes=0, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=partial(nn.LayerNorm, eps=1e-6), return_all_tokens=False, 
+                 init_values=0, use_mean_pooling=False, masked_im_modeling=False):
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim
+        self.return_all_tokens = return_all_tokens
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size[0], patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, 
+                init_values=init_values)
+            for i in range(depth)])
+
+        self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
+        self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None
+        # Classifier head
+        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+
+
+        # masked image modeling
+        self.masked_im_modeling = masked_im_modeling
+        if masked_im_modeling:
+            self.masked_embed = nn.Parameter(torch.zeros(1, embed_dim))
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        class_pos_embed = self.pos_embed[:, 0]
+        patch_pos_embed = self.pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_embed.patch_size
+        h0 = h // self.patch_embed.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            mode='bicubic',
+        )
+        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def prepare_tokens(self, x, no_pe = False, mask=None):
+        B, nc, w, h = x.shape
+        # patch linear embedding
+        x = self.patch_embed(x)
+
+        # mask image modeling
+        if mask is not None:
+            x = self.mask_model(x, mask)
+        x = x.flatten(2).transpose(1, 2)
+
+        # add the [CLS] token to the embed patch tokens
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # add positional encoding to each token
+        if not no_pe:
+            x = x + self.interpolate_pos_encoding(x, w, h)
+
+        return self.pos_drop(x)
+
+    def forward(self, x, return_all_tokens=None, mask=None, no_pe=False):
+        # mim
+        if self.masked_im_modeling:
+            assert mask is not None
+            x = self.prepare_tokens(x, no_pe=no_pe, mask=mask)
+        else:
+            x = self.prepare_tokens(x, no_pe=no_pe)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x = self.norm(x)
+        if self.fc_norm is not None:
+            x[:, 0] = self.fc_norm(x[:, 1:, :].mean(1))
+        
+        return_all_tokens = self.return_all_tokens if \
+            return_all_tokens is None else return_all_tokens
+        if return_all_tokens:
+            return x
+        return x[:, 0]
+
+    def get_last_selfattention(self, x):
+        x = self.prepare_tokens(x)
+        for i, blk in enumerate(self.blocks):
+            if i < len(self.blocks) - 1:
+                x = blk(x)
+            else:
+                # return attention of the last block
+                return blk(x, return_attention=True)
+
+    def get_intermediate_layers(self, x, n=1):
+        x = self.prepare_tokens(x)
+        # we return the output tokens from the `n` last blocks
+        output = []
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if len(self.blocks) - i <= n:
+                output.append(self.norm(x))
+        return output
+        
+    def get_num_layers(self):
+        return len(self.blocks)
+
+    def mask_model(self, x, mask):
+        x.permute(0, 2, 3, 1)[mask, :] = self.masked_embed.to(x.dtype)
+        return x
+
+def vit_tiny(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4,
+        qkv_bias=True, **kwargs)
+    return model
+
+def vit_small(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4,
+        qkv_bias=True, **kwargs)
+    return model
+
+def vit_base(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
+        qkv_bias=True, **kwargs)
+    return model
+
+def vit_large(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4,
+        qkv_bias=True, **kwargs)
+    return model

models/__init__.py

@@ -0,0 +1 @@
+from .vision_transformer import VisionTransformer, vit_tiny, vit_small, vit_base, vit_large

models/head.py

@@ -0,0 +1,241 @@
+import torch
+import torch.nn as nn
+import utils
+
+from utils import trunc_normal_
+
+class CSyncBatchNorm(nn.SyncBatchNorm):
+    def __init__(self,
+                 *args,
+                 with_var=False,
+                 **kwargs):
+        super(CSyncBatchNorm, self).__init__(*args, **kwargs)
+        self.with_var = with_var
+
+    def forward(self, x):
+        # center norm
+        self.training = False
+        if not self.with_var:
+            self.running_var = torch.ones_like(self.running_var)
+        normed_x = super(CSyncBatchNorm, self).forward(x)
+        # udpate center
+        self.training = True
+        _ = super(CSyncBatchNorm, self).forward(x)
+        return normed_x
+
+class PSyncBatchNorm(nn.SyncBatchNorm):
+    def __init__(self,
+                 *args,
+                 bunch_size,
+                 **kwargs):
+        procs_per_bunch = min(bunch_size, utils.get_world_size())
+        assert utils.get_world_size() % procs_per_bunch == 0
+        n_bunch = utils.get_world_size() // procs_per_bunch
+        #
+        ranks = list(range(utils.get_world_size()))
+        print('---ALL RANKS----\n{}'.format(ranks))
+        rank_groups = [ranks[i*procs_per_bunch: (i+1)*procs_per_bunch] for i in range(n_bunch)]
+        print('---RANK GROUPS----\n{}'.format(rank_groups))
+        process_groups = [torch.distributed.new_group(pids) for pids in rank_groups]
+        bunch_id = utils.get_rank() // procs_per_bunch
+        process_group = process_groups[bunch_id]
+        print('---CURRENT GROUP----\n{}'.format(process_group))
+        super(PSyncBatchNorm, self).__init__(*args, process_group=process_group, **kwargs)
+
+class CustomSequential(nn.Sequential):
+    bn_types = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, nn.SyncBatchNorm)
+
+    def forward(self, input):
+        for module in self:
+            dim = len(input.shape)
+            if isinstance(module, self.bn_types) and dim > 2:
+                perm = list(range(dim - 1)); perm.insert(1, dim - 1)
+                inv_perm = list(range(dim)) + [1]; inv_perm.pop(1)
+                input = module(input.permute(*perm)).permute(*inv_perm)
+            else:
+                input = module(input)
+        return input
+
+class DINOHead(nn.Module):
+    def __init__(self, in_dim, out_dim, norm=None, act='gelu', last_norm=None, 
+                 nlayers=3, hidden_dim=2048, bottleneck_dim=256, norm_last_layer=True, **kwargs):
+        super().__init__()
+        norm = self._build_norm(norm, hidden_dim)
+        last_norm = self._build_norm(last_norm, out_dim, affine=False, **kwargs)
+        act = self._build_act(act)
+
+        nlayers = max(nlayers, 1)
+        if nlayers == 1:
+            if bottleneck_dim > 0:
+                self.mlp = nn.Linear(in_dim, bottleneck_dim)
+            else:
+                self.mlp = nn.Linear(in_dim, out_dim)
+        else:
+            layers = [nn.Linear(in_dim, hidden_dim)]
+            if norm is not None:
+                layers.append(norm)
+            layers.append(act)
+            for _ in range(nlayers - 2):
+                layers.append(nn.Linear(hidden_dim, hidden_dim))
+                if norm is not None:
+                    layers.append(norm)
+                layers.append(act)
+            if bottleneck_dim > 0:
+                layers.append(nn.Linear(hidden_dim, bottleneck_dim))
+            else:
+                layers.append(nn.Linear(hidden_dim, out_dim))
+            self.mlp = CustomSequential(*layers)
+        self.apply(self._init_weights)
+        
+        if bottleneck_dim > 0:
+            self.last_layer = nn.utils.weight_norm(nn.Linear(bottleneck_dim, out_dim, bias=False))
+            self.last_layer.weight_g.data.fill_(1)
+            if norm_last_layer:
+                self.last_layer.weight_g.requires_grad = False
+        else:
+            self.last_layer = None
+
+        self.last_norm = last_norm
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+
+    def forward(self, x):
+        x = self.mlp(x)
+        if self.last_layer is not None:
+            x = nn.functional.normalize(x, dim=-1, p=2)
+            x = self.last_layer(x)
+        if self.last_norm is not None:
+            x = self.last_norm(x)
+        return x
+
+    def _build_norm(self, norm, hidden_dim, **kwargs):
+        if norm == 'bn':
+            norm = nn.BatchNorm1d(hidden_dim, **kwargs)
+        elif norm == 'syncbn':
+            norm = nn.SyncBatchNorm(hidden_dim, **kwargs)
+        elif norm == 'csyncbn':
+            norm = CSyncBatchNorm(hidden_dim, **kwargs)
+        elif norm == 'psyncbn':
+            norm =  PSyncBatchNorm(hidden_dim, **kwargs)
+        elif norm == 'ln':
+            norm = nn.LayerNorm(hidden_dim, **kwargs)
+        else:
+            assert norm is None, "unknown norm type {}".format(norm)
+        return norm
+
+    def _build_act(self, act):
+        if act == 'relu':
+            act = nn.ReLU()
+        elif act == 'gelu':
+            act = nn.GELU()
+        else:
+            assert False, "unknown act type {}".format(act)
+        return act
+
+class iBOTHead(DINOHead):
+
+    def __init__(self, *args, patch_out_dim=8192, norm=None, act='gelu', last_norm=None, 
+                 nlayers=3, hidden_dim=2048, bottleneck_dim=256, norm_last_layer=True, 
+                 shared_head=False, **kwargs):
+        
+        super(iBOTHead, self).__init__(*args,
+                                        norm=norm,
+                                        act=act,
+                                        last_norm=last_norm,
+                                        nlayers=nlayers,
+                                        hidden_dim=hidden_dim,
+                                        bottleneck_dim=bottleneck_dim,
+                                        norm_last_layer=norm_last_layer, 
+                                        **kwargs)
+
+        if not shared_head:
+            if bottleneck_dim > 0:
+                self.last_layer2 = nn.utils.weight_norm(nn.Linear(bottleneck_dim, patch_out_dim, bias=False))
+                self.last_layer2.weight_g.data.fill_(1)
+                if norm_last_layer:
+                    self.last_layer2.weight_g.requires_grad = False
+            else:
+                self.mlp2 = nn.Linear(hidden_dim, patch_out_dim)
+                self.last_layer2 = None
+
+            self.last_norm2 = self._build_norm(last_norm, patch_out_dim, affine=False, **kwargs)
+        else:
+            if bottleneck_dim > 0:
+                self.last_layer2 = self.last_layer
+            else:
+                self.mlp2 = self.mlp[-1]
+                self.last_layer2 = None
+
+            self.last_norm2 = self.last_norm
+
+    def forward(self, x):
+        if len(x.shape) == 2:
+            return super(iBOTHead, self).forward(x)
+
+        if self.last_layer is not None:
+            x = self.mlp(x)
+            x = nn.functional.normalize(x, dim=-1, p=2)
+            x1 = self.last_layer(x[:, 0])
+            x2 = self.last_layer2(x[:, 1:])
+        else:
+            x = self.mlp[:-1](x)
+            x1 = self.mlp[-1](x[:, 0])
+            x2 = self.mlp2(x[:, 1:])
+        
+        if self.last_norm is not None:
+            x1 = self.last_norm(x1)
+            x2 = self.last_norm2(x2)
+        
+        return x1, x2
+
+
+
+class TemporalSideContext(nn.Module):
+    def __init__(self, D, max_len=64, n_layers=6, n_head=8, dropout=0.1):
+        super().__init__()
+        #self.pos_t = nn.Embedding(max_len, D)  # learnable embedding for positions
+        layer = nn.TransformerEncoderLayer(D, n_head, 4*D,
+                                           dropout=dropout, batch_first=True)
+        self.enc = nn.TransformerEncoder(layer, n_layers)
+
+    def forward(self, x):          # x [B,T,D]
+        B,T,D = x.shape
+        device = x.device
+        # Generate relative frame positions [0, 1, ..., T-1]
+        #pos_ids = torch.arange(T, device=device).unsqueeze(0)  # [1, T]
+        #pos_embed = self.pos_t(pos_ids)                        # [1, T, D]
+        #x = x + pos_embed
+        return self.enc(x)          # [B,T,D]
+
+
+
+class TemporalHead(nn.Module):
+    """
+    Converts backbone features [B,T,D] → logits [B,T,1] for Plackett–Luce.
+    """
+    def __init__(self, backbone_dim: int, hidden_mul: float = 0.5, max_len: int = 64):
+        super().__init__()
+        hidden_dim = int(backbone_dim * hidden_mul)
+
+        self.reduce = nn.Sequential(
+            nn.Linear(backbone_dim, hidden_dim),
+            nn.GELU()
+        )
+        self.temporal = TemporalSideContext(hidden_dim, max_len=max_len)
+        self.scorer  = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim // 2),
+            nn.GELU(),
+            nn.Linear(hidden_dim // 2, 1)
+        )
+
+    def forward(self, x: torch.Tensor):     # x : [B,T,D]
+        x = self.reduce(x)                  # [B,T,hidden]
+        x = self.temporal(x)                # [B,T,hidden]
+        return self.scorer(x)               # [B,T,1]
+
+
+

models/puzzle_decoder.py

@@ -0,0 +1,62 @@
+import torch.nn as nn
+
+class AttentionBlock(nn.Module):
+    def __init__(self, attn_mul=1, embed_dim=384, num_heads=16):
+        super().__init__()
+        self.attn_mul = attn_mul
+        self.auxiliary_cross_attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
+        self.auxiliary_layer_norm1 = nn.LayerNorm(embed_dim)
+        self.auxiliary_self_attn = nn.MultiheadAttention(embed_dim, num_heads, batch_first=True)
+        self.auxiliary_layer_norm2 = nn.LayerNorm(embed_dim)
+        self.auxiliary_linear = nn.Linear(embed_dim, embed_dim, bias=True)
+        self.auxiliary_layer_norm3 = nn.LayerNorm(embed_dim)
+
+    def forward(self, current_patch, neighbor_patch):
+        # shape of aux_crop and stu_crop: (batch_size, seq_len, embed_dim)
+
+        # MultiheadAttention takes in the query, key, value. Here we use stu_crop to attend to aux_crop.
+        cross_attn_output, cross_attn_output_weights = self.auxiliary_cross_attn(current_patch, neighbor_patch, neighbor_patch)
+        cross_attn_output = self.auxiliary_layer_norm1(self.attn_mul * cross_attn_output + current_patch) # layer norm with skip connection
+        
+        # Then we use cross_attn_output to attend to cross_attn_output itself
+        self_attn_output, self_attn_output_weights = self.auxiliary_self_attn(cross_attn_output, cross_attn_output, cross_attn_output)
+        self_attn_output = self.auxiliary_layer_norm2(self_attn_output + cross_attn_output) # layer norm with skip connection
+
+        # Finally, apply feed forward.
+        output = self.auxiliary_linear(self_attn_output)
+        output = self.auxiliary_layer_norm3(output + self_attn_output) # layer norm with skip connection
+        
+        return output, cross_attn_output_weights, self_attn_output_weights
+
+
+
+class PuzzleDecoder(nn.Module):
+    def __init__(self, attn_mul=4, num_blocks=1, embed_dim=384, num_heads=16):
+        super().__init__()
+        self.decoder = nn.ModuleList([AttentionBlock(attn_mul, embed_dim, num_heads) for _ in range(num_blocks)])
+        self.scorer = nn.Sequential(
+            nn.LayerNorm(embed_dim),
+            nn.Linear(embed_dim, embed_dim),
+            nn.GELU(),
+            nn.Linear(embed_dim, 1),
+        )
+        nn.init.trunc_normal_(self.scorer[-1].weight, std=0.02)
+        nn.init.zeros_(self.scorer[-1].bias)
+
+    def forward(self, current_patch, neighbor_patch, return_feats=False, return_attn=False):
+        x = current_patch
+        cross_ws, self_ws = [], []
+        for block in self.decoder:
+            x, cw, sw = block(x, neighbor_patch)
+            if return_attn:
+                cross_ws.append(cw); self_ws.append(sw)
+
+        scores = self.scorer(x).squeeze(-1)        # [B, N]
+        if return_feats or return_attn:
+            out = {"scores": scores}
+            if return_feats: out["feats"] = x
+            if return_attn:  out["cross_w"] = cross_ws; out["self_w"] = self_ws
+            return out
+
+        return scores
+

models/vision_transformer.py

@@ -0,0 +1,283 @@
+import math
+import torch
+import torch.nn as nn
+
+from functools import partial
+from utils import trunc_normal_
+from timm.models.registry import register_model
+
+def drop_path(x, drop_prob: float = 0., training: bool = False):
+    if drop_prob == 0. or not training:
+        return x
+    keep_prob = 1 - drop_prob
+    shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
+    random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device)
+    random_tensor.floor_()  # binarize
+    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).
+    """
+    def __init__(self, drop_prob=None):
+        super(DropPath, self).__init__()
+        self.drop_prob = drop_prob
+
+    def forward(self, x):
+        return drop_path(x, self.drop_prob, self.training)
+
+
+class Mlp(nn.Module):
+    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.):
+        super().__init__()
+        out_features = out_features or in_features
+        hidden_features = hidden_features or in_features
+        self.fc1 = nn.Linear(in_features, hidden_features)
+        self.act = act_layer()
+        self.fc2 = nn.Linear(hidden_features, out_features)
+        self.drop = nn.Dropout(drop)
+
+    def forward(self, x):
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.drop(x)
+        x = self.fc2(x)
+        x = self.drop(x)
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, dim, num_heads=8, qkv_bias=False, qk_scale=None, attn_drop=0., proj_drop=0.):
+        super().__init__()
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = qk_scale or head_dim ** -0.5
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+        self.attn_drop = nn.Dropout(attn_drop)
+        self.proj = nn.Linear(dim, dim)
+        self.proj_drop = nn.Dropout(proj_drop)
+
+    def forward(self, x):
+        B, N, C = x.shape
+        qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+        q, k, v = qkv[0], qkv[1], qkv[2]
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+        attn = self.attn_drop(attn)
+
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        x = self.proj_drop(x)
+        return x, attn
+
+class Block(nn.Module):
+    def __init__(self, dim, num_heads, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop=0., 
+                 attn_drop=0., drop_path=0., act_layer=nn.GELU, norm_layer=nn.LayerNorm, init_values=0):
+        super().__init__()
+        self.norm1 = norm_layer(dim)
+        self.attn = Attention(
+            dim, num_heads=num_heads, qkv_bias=qkv_bias, qk_scale=qk_scale, attn_drop=attn_drop, proj_drop=drop)
+        self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
+        self.norm2 = norm_layer(dim)
+        mlp_hidden_dim = int(dim * mlp_ratio)
+        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop)
+
+        if init_values > 0:
+            self.gamma_1 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+            self.gamma_2 = nn.Parameter(init_values * torch.ones((dim)), requires_grad=True)
+        else:
+            self.gamma_1, self.gamma_2 = None, None
+
+    def forward(self, x, return_attention=False):
+        y, attn = self.attn(self.norm1(x))
+        if return_attention:
+            return attn
+        if self.gamma_1 is None:
+            x = x + self.drop_path(y)
+            x = x + self.drop_path(self.mlp(self.norm2(x)))
+        else:
+            x = x + self.drop_path(self.gamma_1 * y)
+            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
+        return x
+
+class PatchEmbed(nn.Module):
+    """ Image to Patch Embedding
+    """
+    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
+        super().__init__()
+        num_patches = (img_size // patch_size) * (img_size // patch_size)
+        self.img_size = img_size
+        self.patch_size = patch_size
+        self.num_patches = num_patches
+
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size)
+            
+    def forward(self, x):
+        B, C, H, W = x.shape
+        return self.proj(x)
+
+class VisionTransformer(nn.Module):
+    """ Vision Transformer """
+    def __init__(self, img_size=[224], patch_size=16, in_chans=3, num_classes=0, embed_dim=768, depth=12,
+                 num_heads=12, mlp_ratio=4., qkv_bias=False, qk_scale=None, drop_rate=0., attn_drop_rate=0.,
+                 drop_path_rate=0., norm_layer=partial(nn.LayerNorm, eps=1e-6), return_all_tokens=False, 
+                 init_values=0, use_mean_pooling=False, masked_im_modeling=False):
+        super().__init__()
+        self.num_features = self.embed_dim = embed_dim
+        self.return_all_tokens = return_all_tokens
+
+        self.patch_embed = PatchEmbed(
+            img_size=img_size[0], patch_size=patch_size, in_chans=in_chans, embed_dim=embed_dim)
+        num_patches = self.patch_embed.num_patches
+
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
+        self.pos_drop = nn.Dropout(p=drop_rate)
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)]  # stochastic depth decay rule
+        self.blocks = nn.ModuleList([
+            Block(
+                dim=embed_dim, num_heads=num_heads, mlp_ratio=mlp_ratio, qkv_bias=qkv_bias, qk_scale=qk_scale,
+                drop=drop_rate, attn_drop=attn_drop_rate, drop_path=dpr[i], norm_layer=norm_layer, 
+                init_values=init_values)
+            for i in range(depth)])
+
+        self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
+        self.fc_norm = norm_layer(embed_dim) if use_mean_pooling else None
+        # Classifier head
+        self.head = nn.Linear(embed_dim, num_classes) if num_classes > 0 else nn.Identity()
+
+        trunc_normal_(self.pos_embed, std=.02)
+        trunc_normal_(self.cls_token, std=.02)
+        self.apply(self._init_weights)
+
+
+        # masked image modeling
+        self.masked_im_modeling = masked_im_modeling
+        if masked_im_modeling:
+            self.masked_embed = nn.Parameter(torch.zeros(1, embed_dim))
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, nn.LayerNorm):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def interpolate_pos_encoding(self, x, w, h):
+        npatch = x.shape[1] - 1
+        N = self.pos_embed.shape[1] - 1
+        if npatch == N and w == h:
+            return self.pos_embed
+        class_pos_embed = self.pos_embed[:, 0]
+        patch_pos_embed = self.pos_embed[:, 1:]
+        dim = x.shape[-1]
+        w0 = w // self.patch_embed.patch_size
+        h0 = h // self.patch_embed.patch_size
+        # we add a small number to avoid floating point error in the interpolation
+        # see discussion at https://github.com/facebookresearch/dino/issues/8
+        w0, h0 = w0 + 0.1, h0 + 0.1
+        patch_pos_embed = nn.functional.interpolate(
+            patch_pos_embed.reshape(1, int(math.sqrt(N)), int(math.sqrt(N)), dim).permute(0, 3, 1, 2),
+            scale_factor=(w0 / math.sqrt(N), h0 / math.sqrt(N)),
+            mode='bicubic',
+        )
+        assert int(w0) == patch_pos_embed.shape[-2] and int(h0) == patch_pos_embed.shape[-1]
+        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
+        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
+
+    def prepare_tokens(self, x, no_pe = False, mask=None):
+        B, nc, w, h = x.shape
+        # patch linear embedding
+        x = self.patch_embed(x)
+
+        # mask image modeling
+        if mask is not None:
+            x = self.mask_model(x, mask)
+        x = x.flatten(2).transpose(1, 2)
+
+        # add the [CLS] token to the embed patch tokens
+        cls_tokens = self.cls_token.expand(B, -1, -1)
+        x = torch.cat((cls_tokens, x), dim=1)
+
+        # add positional encoding to each token
+        if not no_pe:
+            x = x + self.interpolate_pos_encoding(x, w, h)
+
+        return self.pos_drop(x)
+
+    def forward(self, x, return_all_tokens=None, mask=None, no_pe=False):
+        # mim
+        if self.masked_im_modeling:
+            assert mask is not None
+            x = self.prepare_tokens(x, no_pe=no_pe, mask=mask)
+        else:
+            x = self.prepare_tokens(x, no_pe=no_pe)
+
+        for blk in self.blocks:
+            x = blk(x)
+
+        x = self.norm(x)
+        if self.fc_norm is not None:
+            x[:, 0] = self.fc_norm(x[:, 1:, :].mean(1))
+        
+        return_all_tokens = self.return_all_tokens if \
+            return_all_tokens is None else return_all_tokens
+        if return_all_tokens:
+            return x
+        return x[:, 0]
+
+    def get_last_selfattention(self, x):
+        x = self.prepare_tokens(x)
+        for i, blk in enumerate(self.blocks):
+            if i < len(self.blocks) - 1:
+                x = blk(x)
+            else:
+                # return attention of the last block
+                return blk(x, return_attention=True)
+
+    def get_intermediate_layers(self, x, n=1):
+        x = self.prepare_tokens(x)
+        # we return the output tokens from the `n` last blocks
+        output = []
+        for i, blk in enumerate(self.blocks):
+            x = blk(x)
+            if len(self.blocks) - i <= n:
+                output.append(self.norm(x))
+        return output
+        
+    def get_num_layers(self):
+        return len(self.blocks)
+
+    def mask_model(self, x, mask):
+        x.permute(0, 2, 3, 1)[mask, :] = self.masked_embed.to(x.dtype)
+        return x
+
+def vit_tiny(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=192, depth=12, num_heads=3, mlp_ratio=4,
+        qkv_bias=True, **kwargs)
+    return model
+
+def vit_small(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=384, depth=12, num_heads=6, mlp_ratio=4,
+        qkv_bias=True, **kwargs)
+    return model
+
+def vit_base(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=768, depth=12, num_heads=12, mlp_ratio=4,
+        qkv_bias=True, **kwargs)
+    return model
+
+def vit_large(patch_size=16, **kwargs):
+    model = VisionTransformer(
+        patch_size=patch_size, embed_dim=1024, depth=24, num_heads=16, mlp_ratio=4,
+        qkv_bias=True, **kwargs)
+    return model