Merge pull request #1812 from seefun/master

add ViT for Segment-Anything Model

Author: rwightman

tests/test_models.py

@@ -41,7 +41,7 @@
     'vit_*', 'tnt_*', 'pit_*', 'coat_*', 'cait_*', '*mixer_*', 'gmlp_*', 'resmlp_*', 'twins_*',
     'convit_*', 'levit*', 'visformer*', 'deit*', 'jx_nest_*', 'nest_*', 'xcit_*', 'crossvit_*', 'beit*',
     'poolformer_*', 'volo_*', 'sequencer2d_*', 'pvt_v2*', 'mvitv2*', 'gcvit*', 'efficientformer*',
-    'eva_*', 'flexivit*', 'eva02*'
+    'eva_*', 'flexivit*', 'eva02*', 'samvit_*'
 ]
 NUM_NON_STD = len(NON_STD_FILTERS)
 

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
+from .pos_embed import resample_abs_pos_embed, resample_abs_pos_embed_nhwc
 from .pos_embed_rel import RelPosMlp, RelPosBias, RelPosBiasTf, gen_relative_position_index, gen_relative_log_coords
 from .pos_embed_sincos import pixel_freq_bands, freq_bands, build_sincos2d_pos_embed, build_fourier_pos_embed, \
     build_rotary_pos_embed, apply_rot_embed, apply_rot_embed_cat, apply_rot_embed_list, apply_keep_indices_nlc, \

timm/layers/patch_embed.py

@@ -37,6 +37,7 @@ def __init__(
             flatten: bool = True,
             output_fmt: Optional[str] = None,
             bias: bool = True,
+            strict_img_size: bool = True,
     ):
         super().__init__()
         self.patch_size = to_2tuple(patch_size)
@@ -56,15 +57,26 @@ def __init__(
             # flatten spatial dim and transpose to channels last, kept for bwd compat
             self.flatten = flatten
             self.output_fmt = Format.NCHW
+        self.strict_img_size = strict_img_size
 
         self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=bias)
         self.norm = norm_layer(embed_dim) if norm_layer else nn.Identity()
 
     def forward(self, x):
         B, C, H, W = x.shape
         if self.img_size is not None:
-            _assert(H == self.img_size[0], f"Input image height ({H}) doesn't match model ({self.img_size[0]}).")
-            _assert(W == self.img_size[1], f"Input image width ({W}) doesn't match model ({self.img_size[1]}).")
+            if self.strict_img_size:
+                _assert(H == self.img_size[0], f"Input height ({H}) doesn't match model ({self.img_size[0]}).")
+                _assert(W == self.img_size[1], f"Input width ({W}) doesn't match model ({self.img_size[1]}).")
+            else:
+                _assert(
+                    H % self.patch_size[0] == 0,
+                    f"Input height ({H}) should be divisible by patch size ({self.patch_size[0]})."
+                )
+                _assert(
+                    W % self.patch_size[1] == 0,
+                    f"Input width ({W}) should be divisible by patch size ({self.patch_size[1]})."
+                )
 
         x = self.proj(x)
         if self.flatten:

timm/layers/pos_embed.py

@@ -52,3 +52,24 @@ def resample_abs_pos_embed(
         _logger.info(f'Resized position embedding: {old_size} to {new_size}.')
 
     return posemb
+
+
+def resample_abs_pos_embed_nhwc(
+        posemb,
+        new_size: List[int],
+        interpolation: str = 'bicubic',
+        antialias: bool = True,
+        verbose: bool = False,
+):
+    if new_size[0] == posemb.shape[-3] and new_size[1] == posemb.shape[-2]:
+        return posemb
+
+    # do the interpolation
+    posemb = posemb.reshape(1, posemb.shape[-3], posemb.shape[-2], posemb.shape[-1]).permute(0, 3, 1, 2)
+    posemb = F.interpolate(posemb, size=new_size, mode=interpolation, antialias=antialias)
+    posemb = posemb.permute(0, 2, 3, 1)
+
+    if not torch.jit.is_scripting() and verbose:
+        _logger.info(f'Resized position embedding: {posemb.shape[-3:-1]} to {new_size}.')
+
+    return posemb

timm/models/__init__.py

@@ -60,6 +60,7 @@
 from .vision_transformer import *
 from .vision_transformer_hybrid import *
 from .vision_transformer_relpos import *
+from .vision_transformer_sam import *
 from .volo import *
 from .vovnet import *
 from .xception import *