add tinyvit

Author: seefun

timm/layers/__init__.py

@@ -36,7 +36,7 @@
 from .patch_dropout import PatchDropout
 from .patch_embed import PatchEmbed, PatchEmbedWithSize, resample_patch_embed
 from .pool2d_same import AvgPool2dSame, create_pool2d
-from .pos_embed import resample_abs_pos_embed, resample_abs_pos_embed_nhwc
+from .pos_embed import resample_abs_pos_embed, resample_abs_pos_embed_nhwc, resample_relative_position_bias_table
 from .pos_embed_rel import RelPosMlp, RelPosBias, RelPosBiasTf, gen_relative_position_index, gen_relative_log_coords, \
     resize_rel_pos_bias_table, resize_rel_pos_bias_table_simple
 from .pos_embed_sincos import pixel_freq_bands, freq_bands, build_sincos2d_pos_embed, build_fourier_pos_embed, \

timm/layers/pos_embed.py

@@ -78,3 +78,38 @@ def resample_abs_pos_embed_nhwc(
         _logger.info(f'Resized position embedding: {posemb.shape[-3:-1]} to {new_size}.')
 
     return posemb
+
+
+def resample_relative_position_bias_table(
+        position_bias_table,
+        new_size,
+        interpolation: str = 'bicubic',
+        antialias: bool = True,
+        verbose: bool = False
+):
+    """
+    Resample relative position bias table suggested in LeVit
+    Adapted from: https://github.com/microsoft/Cream/blob/main/TinyViT/utils.py
+    """
+    L1, nH1 = position_bias_table.size()
+    L2, nH2 = new_size
+    assert nH1 == nH2
+    if L1 != L2:
+        orig_dtype = position_bias_table.dtype
+        position_bias_table = position_bias_table.float()
+        # bicubic interpolate relative_position_bias_table if not match
+        S1 = int(L1 ** 0.5)
+        S2 = int(L2 ** 0.5)
+        relative_position_bias_table_resized = F.interpolate(
+            position_bias_table.permute(1, 0).view(1, nH1, S1, S1),
+            size=(S2, S2),
+            mode=interpolation,
+            antialias=antialias)
+        relative_position_bias_table_resized = \
+            relative_position_bias_table_resized.view(nH2, L2).permute(1, 0)
+        relative_position_bias_table_resized.to(orig_dtype)
+        if not torch.jit.is_scripting() and verbose:
+            _logger.info(f'Resized position bias: {L1, nH1} to {L2, nH2}.')
+        return relative_position_bias_table_resized
+    else:
+        return position_bias_table

timm/models/__init__.py

@@ -58,6 +58,7 @@
 from .swin_transformer import *
 from .swin_transformer_v2 import *
 from .swin_transformer_v2_cr import *
+from .tiny_vit import *
 from .tnt import *
 from .tresnet import *
 from .twins import *

timm/models/efficientvit_msra.py

@@ -15,7 +15,7 @@
 import torch.nn as nn
 
 from timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
-from timm.models.layers import SqueezeExcite, SelectAdaptivePool2d, trunc_normal_, _assert
+from timm.layers import SqueezeExcite, SelectAdaptivePool2d, trunc_normal_, _assert
 from ._builder import build_model_with_cfg
 from ._manipulate import checkpoint_seq
 from ._registry import register_model, generate_default_cfgs