Multiscale Vision Transformers

tests/test_attention.py

@@ -51,6 +51,22 @@ def test_CrossAttention():
     assert out.shape == test_tensor1.shape
     del attention
 
+def test_MultiScaleAttention():
+
+    test_tensor1 = torch.randn(96,8,56,56)
+    test_tensor2 = torch.randn(768,8,14,14)
+    thw  = [2,2,2]
+
+    attention = ATTENTION_REGISTRY.get("MultiScaleAttention")(dim=192)
+    out = attention(test_tensor1, thw)
+    assert out.shape == (192,8,28,28)
+    del attention
+
+    attention = ATTENTION_REGISTRY.get("VanillaSelfAttention")(dim=768)
+    out = attention(test_tensor2)
+    assert out.shape == (768,8,14,14)
+    del attention
+
 
 def test_SpatialAttention():
 

tests/test_encoder.py

@@ -56,7 +56,22 @@ def test_SwinEncoder():
     out = encoder_block(test_tensor)
     assert out.shape == test_tensor.shape
 
-
+def test_MultiScaleBlock():
+
+    test_tensor1 = torch.randn(96,8,56,56)
+    encoder1 = ENCODER_REGISTRY.get("MultiScaleBlock")(dim=192)
+    out1 = encoder1(test_tensor)
+    assert out1.shape == (192,8,28,28)  
+
+    test_tensor2 = torch.randn(768,8,14,14)
+    encoder2 = ENCODER_REGISTRY.get("MultiScaleBlock")(dim=768)
+    out2 = encoder2(test_tensor)
+    assert out2.shape == (768,8,14,14)  # shape remains same
+
+
+    del encoder1, encoder2, test_tensor1, test_tensor2
+
+
 def test_PVTEncoder():
 
     test_tensor = torch.randn(4, 3136, 64)

tests/test_models.py

@@ -159,7 +159,15 @@ def test_CrossVit():
     assert out.shape == (2, 10)
     del model
 
+def test_MultiScale():
 
+    model = MODEL_REGISTRY.get("MultiScaleViT")()
+    out = model(img_3channels_224)
+    assert out.shape == (8, 400)
+    del model
+
+
+
 def test_pvt():
     # classification
     model = MODEL_REGISTRY.get("PVTClassification")(

vformer/attention/__init__.py

@@ -2,3 +2,4 @@
 from .spatial import SpatialAttention
 from .vanilla import VanillaSelfAttention
 from .window import WindowAttention
+from .multiscale import MultiScaleAttention

vformer/attention/multiscale.py

@@ -0,0 +1,270 @@
+import numpy
+import torch
+import torch.nn as nn
+from ..utils import ATTENTION_REGISTRY
+
+def attention_pool(tensor, pool, thw_shape, has_cls_embed=True, norm=None):
+    """
+    Attention pooling
+    Parameters:
+    -----------
+    tensor: tensor
+            Input tensor
+    pool: nn.Module
+          Pooling function
+    thw_shape: list of int
+               Reduced space-time resolution
+    has_cls_embed: bool, optional
+                   Set to true if classification embeddding is provided
+    norm : nn.Module, optional
+           Normalization function
+    """
+
+    if pool is None:
+        return tensor, thw_shape
+    tensor_dim = tensor.ndim
+    if tensor_dim == 4:
+        pass
+    elif tensor_dim == 3:
+        tensor = tensor.unsqueeze(1)
+    else:
+        raise NotImplementedError(f"Unsupported input dimension {tensor.shape}")
+
+    if has_cls_embed:
+        cls_tok, tensor = tensor[:, :, :1, :], tensor[:, :, 1:, :]
+
+    B, N, L, C = tensor.shape
+    T, H, W = thw_shape
+    tensor = tensor.reshape(B * N, T, H, W, C).permute(0, 4, 1, 2, 3).contiguous()
+
+    tensor = pool(tensor)
+
+    thw_shape = [tensor.shape[2], tensor.shape[3], tensor.shape[4]]
+    L_pooled = tensor.shape[2] * tensor.shape[3] * tensor.shape[4]
+    tensor = tensor.reshape(B, N, C, L_pooled).transpose(2, 3)
+    if has_cls_embed:
+        tensor = torch.cat((cls_tok, tensor), dim=2)
+    if norm is not None:
+        tensor = norm(tensor)
+    # Assert tensor_dim in [3, 4]
+    if tensor_dim == 4:
+        pass
+    else:  #  tensor_dim == 3:
+        tensor = tensor.squeeze(1)
+    return tensor, thw_shape
+
+@ATTENTION_REGISTRY.register()
+class MultiScaleAttention(nn.Module):
+    """
+    Multiscale Attention
+    Parameters:
+    -----------
+    dim: int
+         Dimension of the embedding
+    num_heads: int
+               Number of attention heads
+    qkv_bias :bool, optional
+              If True, add a learnable bias to query, key, value
+    drop_rate: float, optional
+               Dropout rate
+    kernel_q: tuple of int, optional
+              Kernel size of query
+    kernel_kv: tuple of int, optional
+               Kernel size of key and value
+    stride_q: tuple of int, optional
+              Kernel size of query
+    stride_kv: tuple of int, optional
+               Kernel size of key and value
+    norm_layer: nn.Module, optional
+                Normalization function
+    has_cls_embed: bool, optional
+                   Set to true if classification embeddding is provided
+    mode: str, optional
+          Pooling function to be used. Options include `conv`, `avg`, and `max'
+    pool_first: bool, optional
+                Set to True to perform pool before projection
+    """
+
+    def __init__(
+        self,
+        dim,
+        num_heads=8,
+        qkv_bias=False,
+        drop_rate=0.0,
+        kernel_q=(1, 1, 1),
+        kernel_kv=(1, 1, 1),
+        stride_q=(1, 1, 1),
+        stride_kv=(1, 1, 1),
+        norm_layer=nn.LayerNorm,
+        has_cls_embed=True,
+        mode="conv",
+        pool_first=False,
+    ):
+        super().__init__()
+        self.pool_first = pool_first
+        self.drop_rate = drop_rate
+        self.num_heads = num_heads
+        head_dim = dim // num_heads
+        self.scale = head_dim ** -0.5
+        self.has_cls_embed = has_cls_embed
+        padding_q = [int(q // 2) for q in kernel_q]
+        padding_kv = [int(kv // 2) for kv in kernel_kv]
+
+        self.q = nn.Linear(dim, dim, bias=qkv_bias)
+        self.k = nn.Linear(dim, dim, bias=qkv_bias)
+        self.v = nn.Linear(dim, dim, bias=qkv_bias)
+        self.proj = nn.Linear(dim, dim)
+        if drop_rate > 0.0:
+            self.proj_drop = nn.Dropout(drop_rate)
+
+        # Skip pooling with kernel and stride size of (1, 1, 1).
+        if numpy.prod(kernel_q) == 1 and numpy.prod(stride_q) == 1:
+            kernel_q = ()
+        if numpy.prod(kernel_kv) == 1 and numpy.prod(stride_kv) == 1:
+            kernel_kv = ()
+
+        if mode in ("avg", "max"):
+            pool_op = nn.MaxPool3d if mode == "max" else nn.AvgPool3d
+            self.pool_q = (
+                pool_op(kernel_q, stride_q, padding_q, ceil_mode=False)
+                if len(kernel_q) > 0
+                else None
+            )
+            self.pool_k = (
+                pool_op(kernel_kv, stride_kv, padding_kv, ceil_mode=False)
+                if len(kernel_kv) > 0
+                else None
+            )
+            self.pool_v = (
+                pool_op(kernel_kv, stride_kv, padding_kv, ceil_mode=False)
+                if len(kernel_kv) > 0
+                else None
+            )
+        elif mode == "conv":
+            self.pool_q = (
+                nn.Conv3d(
+                    head_dim,
+                    head_dim,
+                    kernel_q,
+                    stride=stride_q,
+                    padding=padding_q,
+                    groups=head_dim,
+                    bias=False,
+                )
+                if len(kernel_q) > 0
+                else None
+            )
+            self.norm_q = norm_layer(head_dim) if len(kernel_q) > 0 else None
+            self.pool_k = (
+                nn.Conv3d(
+                    head_dim,
+                    head_dim,
+                    kernel_kv,
+                    stride=stride_kv,
+                    padding=padding_kv,
+                    groups=head_dim,
+                    bias=False,
+                )
+                if len(kernel_kv) > 0
+                else None
+            )
+            self.norm_k = norm_layer(head_dim) if len(kernel_kv) > 0 else None
+            self.pool_v = (
+                nn.Conv3d(
+                    head_dim,
+                    head_dim,
+                    kernel_kv,
+                    stride=stride_kv,
+                    padding=padding_kv,
+                    groups=head_dim,
+                    bias=False,
+                )
+                if len(kernel_kv) > 0
+                else None
+            )
+            self.norm_v = norm_layer(head_dim) if len(kernel_kv) > 0 else None
+        else:
+            raise NotImplementedError(f"Unsupported model {mode}")
+
+    def forward(self, x, thw_shape):
+        B, N, C = x.shape
+        if self.pool_first:
+            x = x.reshape(B, N, self.num_heads, C // self.num_heads).permute(0, 2, 1, 3)
+            q = k = v = x
+        else:
+            q = k = v = x
+            q = (
+                self.q(q)
+                .reshape(B, N, self.num_heads, C // self.num_heads)
+                .permute(0, 2, 1, 3)
+            )
+            k = (
+                self.k(k)
+                .reshape(B, N, self.num_heads, C // self.num_heads)
+                .permute(0, 2, 1, 3)
+            )
+            v = (
+                self.v(v)
+                .reshape(B, N, self.num_heads, C // self.num_heads)
+                .permute(0, 2, 1, 3)
+            )
+
+        q, q_shape = attention_pool(
+            q,
+            self.pool_q,
+            thw_shape,
+            has_cls_embed=self.has_cls_embed,
+            norm=self.norm_q if hasattr(self, "norm_q") else None,
+        )
+        k, k_shape = attention_pool(
+            k,
+            self.pool_k,
+            thw_shape,
+            has_cls_embed=self.has_cls_embed,
+            norm=self.norm_k if hasattr(self, "norm_k") else None,
+        )
+        v, v_shape = attention_pool(
+            v,
+            self.pool_v,
+            thw_shape,
+            has_cls_embed=self.has_cls_embed,
+            norm=self.norm_v if hasattr(self, "norm_v") else None,
+        )
+
+        if self.pool_first:
+            q_N = numpy.prod(q_shape) + 1 if self.has_cls_embed else numpy.prod(q_shape)
+            k_N = numpy.prod(k_shape) + 1 if self.has_cls_embed else numpy.prod(k_shape)
+            v_N = numpy.prod(v_shape) + 1 if self.has_cls_embed else numpy.prod(v_shape)
+
+            q = q.permute(0, 2, 1, 3).reshape(B, q_N, C)
+            q = (
+                self.q(q)
+                .reshape(B, q_N, self.num_heads, C // self.num_heads)
+                .permute(0, 2, 1, 3)
+            )
+
+            v = v.permute(0, 2, 1, 3).reshape(B, v_N, C)
+            v = (
+                self.v(v)
+                .reshape(B, v_N, self.num_heads, C // self.num_heads)
+                .permute(0, 2, 1, 3)
+            )
+
+            k = k.permute(0, 2, 1, 3).reshape(B, k_N, C)
+            k = (
+                self.k(k)
+                .reshape(B, k_N, self.num_heads, C // self.num_heads)
+                .permute(0, 2, 1, 3)
+            )
+
+        attn = (q @ k.transpose(-2, -1)) * self.scale
+        attn = attn.softmax(dim=-1)
+
+        N = q.shape[2]
+        x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+        x = self.proj(x)
+        if self.drop_rate > 0.0:
+            x = self.proj_drop(x)
+        return x, q_shape
+
+

vformer/encoder/__init__.py

@@ -4,3 +4,4 @@
 from .pyramid import PVTEncoder
 from .swin import SwinEncoder, SwinEncoderBlock
 from .vanilla import VanillaEncoder
+from .multiscale import MultiScaleBlock

vformer/encoder/embedding/patch_multiscale.py

@@ -0,0 +1,47 @@
+class PatchEmbed(nn.Module):
+    """
+    arameters
+    ----------
+    img_size: int
+        Image Size
+
+    dim_in: int
+        Number of input channels in the image
+    dim_out: int
+        Number of linear projection output channels
+    kernel: int
+        kernel Size
+    stride: int
+        stride Size
+    padding: int
+        padding Size
+    conv_2d : bool
+        Use nn.Conv2D if true, nn.conv3D if fals3
+    """
+
+    def __init__(
+        self,
+        dim_in=3,
+        dim_out=768,
+        kernel=(1, 16, 16),
+        stride=(1, 4, 4),
+        padding=(1, 7, 7),
+        conv_2d=False,
+    ):
+        super().__init__()
+        if conv_2d:
+            conv = nn.Conv2d
+        else:
+            conv = nn.Conv3d
+        self.proj = conv(
+            dim_in,
+            dim_out,
+            kernel_size=kernel,
+            stride=stride,
+            padding=padding,
+        )
+
+    def forward(self, x):
+        x = self.proj(x)
+        # B C (T) H W -> B (T)HW C
+        return x.flatten(2).transpose(1, 2)