Repository: SforAiDl/vformer (language: python)

Issue:
Restructuring AbsolutePositionEmbedding class
AbsolutePositionEmbedding class was structured specifically for the PVT, but we can use it in other models too if we re-structure it properly, it should also support sinusoidal position embedding or a separate class for Sinusoidal embedding also works. 




Discussion / hints:
Should I assign this to you since you are the best equipped to take it up?
sure

Fix patch:
diff --git a/docs/api/models/classification/cct.rst b/docs/api/models/classification/cct.rst
new file mode 100644
index 00000000..80f6d68b
--- /dev/null
+++ b/docs/api/models/classification/cct.rst
@@ -0,0 +1,5 @@
+Compact Convolutional Transformer
+================================
+
+.. automodule:: vformer.models.classification.cct
+    :members:
diff --git a/docs/api/models/classification/cvt.rst b/docs/api/models/classification/cvt.rst
new file mode 100644
index 00000000..bae128bc
--- /dev/null
+++ b/docs/api/models/classification/cvt.rst
@@ -0,0 +1,5 @@
+Compact Vision Transformer
+================================
+
+.. automodule:: vformer.models.classification.cvt
+    :members:
diff --git a/docs/api/models/classification/index.rst b/docs/api/models/classification/index.rst
index 43ef872e..e4f8254d 100644
--- a/docs/api/models/classification/index.rst
+++ b/docs/api/models/classification/index.rst
@@ -5,7 +5,9 @@ Classification
     :caption: Contents
     :maxdepth: 2
 
+    cct
     cross
+    cvt
     pyramid
     swin
-    vanilla
\ No newline at end of file
+    vanilla
diff --git a/vformer/attention/spatial.py b/vformer/attention/spatial.py
index 4700f782..d809f0ae 100644
--- a/vformer/attention/spatial.py
+++ b/vformer/attention/spatial.py
@@ -7,6 +7,7 @@
 @ATTENTION_REGISTRY.register()
 class SpatialAttention(nn.Module):
     """
+
     Spatial Reduction Attention- Linear complexity attention layer
 
     Parameters
@@ -29,6 +30,7 @@ class SpatialAttention(nn.Module):
         Whether to use linear Spatial attention,default is False
     act_fn : nn.Module
         Activation function, default is False
+
     """
 
     def __init__(
@@ -74,7 +76,22 @@ def __init__(
             self.sr = nn.Conv2d(dim, dim, kernel_size=1, stride=1)
 
     def forward(self, x, H, W):
-
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Input tensor
+        H: int
+            Height  of image patches
+        W: int
+            Width of image patches
+        Returns
+        ----------
+        torch.Tensor
+            Returns output tensor by applying spatial attention on input tensor
+
+        """
         B, N, C = x.shape
         q = (
             self.q(x)
diff --git a/vformer/attention/vanilla.py b/vformer/attention/vanilla.py
index d055952d..8844f1d3 100644
--- a/vformer/attention/vanilla.py
+++ b/vformer/attention/vanilla.py
@@ -20,6 +20,7 @@ class VanillaSelfAttention(nn.Module):
         Dimension of each head
     p_dropout: float
         Dropout Probability
+
     """
 
     def __init__(self, dim, num_heads=8, head_dim=64, p_dropout=0.0):
@@ -41,7 +42,18 @@ def __init__(self, dim, num_heads=8, head_dim=64, p_dropout=0.0):
         )
 
     def forward(self, x):
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Input tensor
+        Returns
+        ----------
+        torch.Tensor
+            Returns output tensor by applying self-attention on input tensor
 
+        """
         qkv = self.to_qkv(x).chunk(3, dim=-1)
         q, k, v = map(
             lambda t: rearrange(t, "b n (h d) -> b h n d", h=self.num_heads), qkv
diff --git a/vformer/attention/window.py b/vformer/attention/window.py
index f9d39ec1..5094b783 100644
--- a/vformer/attention/window.py
+++ b/vformer/attention/window.py
@@ -59,6 +59,23 @@ def __init__(
         trunc_normal_(self.relative_position_bias_table, std=0.2)
 
     def forward(self, x, mask=None):
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            input Tensor
+        mask: torch.Tensor
+            Attention mask used for shifted window attention, if None, window attention will be used,
+            else attention mask will be taken into consideration.
+            for better understanding you may refer `this <https://github.com/microsoft/Swin-Transformer/issues/38>`
+
+        Returns
+        ----------
+        torch.Tensor
+            Returns output tensor by applying Window-Attention or Shifted-Window-Attention on input tensor
+
+        """
 
         B_, N, C = x.shape
         qkv = (
diff --git a/vformer/common/blocks.py b/vformer/common/blocks.py
index ea439989..b0bb4952 100644
--- a/vformer/common/blocks.py
+++ b/vformer/common/blocks.py
@@ -17,6 +17,7 @@ class DWConv(nn.Module):
         Padding added to all sides of the input, default is 1
     bias_dwconv:bool
         Whether to add learnable bias to the output,default is True.
+
     """
 
     def __init__(
@@ -39,8 +40,27 @@ def __init__(
         )
 
     def forward(self, x, H, W):
+        """
+
+        Parameters:
+        ----------
+        x: torch.Tensor
+            Input tensor
+        H: int
+            Height of image patch
+        W: int
+            Width of image patch
+
+        Returns:
+        ----------
+        torch.Tensor
+            Returns output tensor after performing depth-wise convolution operation
+
+        """
+
         B, N, C = x.shape
         x = x.transpose(1, 2).view(B, C, H, W)
         x = self.dwconv(x)
+
         x = x.flatten(2).transpose(1, 2)
         return x
diff --git a/vformer/decoder/mlp.py b/vformer/decoder/mlp.py
index b441ade0..e8abeddc 100644
--- a/vformer/decoder/mlp.py
+++ b/vformer/decoder/mlp.py
@@ -34,5 +34,16 @@ def __init__(self, config=(1024,), n_classes=10):
         self.decoder = nn.Sequential(*self.decoder)
 
     def forward(self, x):
-
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Input tensor
+        Returns
+        ----------
+        torch.Tensor
+            Returns output tensor of size `n_classes`, Note that `torch.nn.Softmax` is not applied to the output tensor.
+
+        """
         return self.decoder(x)
diff --git a/vformer/encoder/cross.py b/vformer/encoder/cross.py
index 53587b34..25e7a220 100644
--- a/vformer/encoder/cross.py
+++ b/vformer/encoder/cross.py
@@ -9,6 +9,7 @@
 @ENCODER_REGISTRY.register()
 class CrossEncoder(nn.Module):
     """
+
     Parameters
     ----------
     embedding_dim_s : int
diff --git a/vformer/encoder/embedding/__init__.py b/vformer/encoder/embedding/__init__.py
index 91277732..94e49c27 100644
--- a/vformer/encoder/embedding/__init__.py
+++ b/vformer/encoder/embedding/__init__.py
@@ -2,4 +2,4 @@
 from .linear import LinearEmbedding
 from .overlappatch import OverlapPatchEmbed
 from .patch import PatchEmbedding
-from .pos_embedding import AbsolutePositionEmbedding
+from .pos_embedding import *
diff --git a/vformer/encoder/embedding/cvt.py b/vformer/encoder/embedding/cvt.py
index 17c10ded..084599f4 100644
--- a/vformer/encoder/embedding/cvt.py
+++ b/vformer/encoder/embedding/cvt.py
@@ -90,6 +90,19 @@ def sequence_length(self, n_channels=3, height=224, width=224):
         return self.forward(torch.zeros((1, n_channels, height, width))).shape[1]
 
     def forward(self, x):
+        """
+
+        Parameters
+        ----------
+        x: torch.tensor
+            Input tensor
+
+        Returns
+        -----------
+        torch.Tensor
+            Returns output tensor (embedding) by applying multiple convolution and max-pooling operations on input tensor
+
+        """
 
         for conv2d, activation, maxpool in self.conv_layers:
             x = maxpool(activation(conv2d(x)))
diff --git a/vformer/encoder/embedding/linear.py b/vformer/encoder/embedding/linear.py
index c5f70700..0b22eeb6 100644
--- a/vformer/encoder/embedding/linear.py
+++ b/vformer/encoder/embedding/linear.py
@@ -15,6 +15,7 @@ class LinearEmbedding(nn.Module):
         Width of the patch
     patch_dim: int
         Dimension of the patch
+
     """
 
     def __init__(
@@ -36,4 +37,18 @@ def __init__(
         )
 
     def forward(self, x):
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Input tensor
+
+        Returns
+        ----------
+        torch.Tensor
+            Returns patch embeddings of size `embedding_dim`
+
+        """
+
         return self.patch_embedding(x)
diff --git a/vformer/encoder/embedding/overlappatch.py b/vformer/encoder/embedding/overlappatch.py
index 06a85a3b..0077901e 100644
--- a/vformer/encoder/embedding/overlappatch.py
+++ b/vformer/encoder/embedding/overlappatch.py
@@ -20,6 +20,7 @@ class OverlapPatchEmbed(nn.Module):
         Number of linear projection output channels,default is 768
     norm_layer: nn.Module, optional
         Normalization layer, default is nn.LayerNorm
+
     """
 
     def __init__(
@@ -51,6 +52,23 @@ def __init__(
         self.norm = norm_layer(embedding_dim)
 
     def forward(self, x):
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+              Input tensor
+
+        Returns
+        ----------
+        x: torch.Tensor
+            Input tensor
+        H: int
+            Height of Patch
+        W: int
+            Width of Patch
+
+        """
 
         x = self.proj(x)
         H, W = x.shape[2:]
diff --git a/vformer/encoder/embedding/patch.py b/vformer/encoder/embedding/patch.py
index 90d8d2af..14b303fa 100644
--- a/vformer/encoder/embedding/patch.py
+++ b/vformer/encoder/embedding/patch.py
@@ -18,6 +18,7 @@ class PatchEmbedding(nn.Module):
         Number of linear projection output channels
     norm_layer: nn.Module,
         Normalization layer, Default is `nn.LayerNorm`
+
     """
 
     def __init__(
@@ -46,7 +47,19 @@ def __init__(
         self.norm = norm_layer(embedding_dim)
 
     def forward(self, x):
+        """
+
+        Parameters
+        ----------
+        x:torch.Tensor
+            Input tensor
+
+        Returns
+        ----------
+        torch.Tensor
+            Returns output tensor by applying convolution operation with same `kernel_size` and `stride` on input tensor.
 
+        """
         B, C, H, W = x.shape
 
         assert (
diff --git a/vformer/encoder/embedding/pos_embedding.py b/vformer/encoder/embedding/pos_embedding.py
index 452233e6..4859cfc5 100644
--- a/vformer/encoder/embedding/pos_embedding.py
+++ b/vformer/encoder/embedding/pos_embedding.py
@@ -6,7 +6,7 @@
 from ...utils import pair
 
 
-class AbsolutePositionEmbedding(nn.Module):
+class PVTPosEmbedding(nn.Module):
     """
 
     Parameters
@@ -25,12 +25,12 @@ def __init__(self, pos_shape, pos_dim, p_dropout=0.0, std=0.02):
         super().__init__()
 
         pos_shape = pair(pos_shape)
-        self.pos_shape = pos_shape
-        self.pos_dim = pos_dim
-
         self.pos_embed = nn.Parameter(
             torch.zeros(1, pos_shape[0] * pos_shape[1], pos_dim)
         )
+        self.pos_shape = pos_shape
+        self.pos_dim = pos_dim
+
         self.drop = nn.Dropout(p=p_dropout)
         trunc_normal_(self.pos_embed, std=std)
 
@@ -65,7 +65,34 @@ def resize_pos_embed(self, pos_embed, shape, mode="bilinear", **kwargs):
         return pos_embed
 
     def forward(self, x, H, W, mode="bilinear"):
+        try:
+            x = x + self.pos_embed
+
+        except:
+            x = x + self.resize_pos_embed(self.pos_embed, (H, W), mode)
+
+        return self.drop(x)
 
-        pos_embed = self.resize_pos_embed(self.pos_embed, (H, W), mode)
 
-        return self.drop(x + pos_embed)
+class PosEmbedding(nn.Module):
+    def __init__(self, shape, dim, drop=None, sinusoidal=False, std=0.02):
+        super(PosEmbedding, self).__init__()
+        if not sinusoidal:
+            self.pos_embed = torch.zeros(1, shape, dim)
+        else:
+            pe = torch.FloatTensor(
+                [
+                    [p / (10000 ** (2 * (i // 2) / dim)) for i in range(dim)]
+                    for p in range(shape)
+                ]
+            )
+            pe[:, 0::2] = torch.sin(pe[:, 0::2])
+            pe[:, 1::2] = torch.cos(pe[:, 1::2])
+            self.pos_embed = pe
+            self.pos_embed.requires_grad = False
+        trunc_normal_(self.pos_embed, std=std)
+        self.drop = nn.Dropout(drop) if drop is not None else nn.Identity()
+
+    def forward(self, x):
+        x = x + self.pos_embed
+        return self.drop(x)
diff --git a/vformer/encoder/nn.py b/vformer/encoder/nn.py
index 3d6c51fe..4be25bc8 100644
--- a/vformer/encoder/nn.py
+++ b/vformer/encoder/nn.py
@@ -14,6 +14,7 @@ class FeedForward(nn.Module):
         Dimension of the output tensor
     p_dropout: float
         Dropout probability, default=0.0
+
     """
 
     def __init__(self, dim, hidden_dim=None, out_dim=None, p_dropout=0.0):
@@ -31,4 +32,18 @@ def __init__(self, dim, hidden_dim=None, out_dim=None, p_dropout=0.0):
         )
 
     def forward(self, x):
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Input tensor
+        Returns
+        ----------
+
+        torch.Tensor
+            Returns output tensor by performing linear operations and activation on input tensor
+
+        """
+
         return self.net(x)
diff --git a/vformer/encoder/pyramid.py b/vformer/encoder/pyramid.py
index c44136e5..0c95d8a0 100644
--- a/vformer/encoder/pyramid.py
+++ b/vformer/encoder/pyramid.py
@@ -27,9 +27,7 @@ class PVTFeedForward(nn.Module):
     use_dwconv: bool
         Whether to use Depth-wise convolutions, default is False
 
-    Kwargs
-    ------
-    kernel_size_dwconv: int,optional
+    kernel_size_dwconv: int
         `kernel_size` parameter for 2D convolution used in Depth wise convolution
     stride_dwconv: int
         `stride` parameter for 2D convolution used in Depth wise convolution
@@ -70,6 +68,23 @@ def __init__(
 
     def forward(self, x, **kwargs):
 
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Input tensor
+        H: int
+            Height of image patch
+        W: int
+            Width of image patch
+
+        Returns
+        --------
+        torch.Tensor
+            Returns output tensor
+
+        """
         x = self.relu(self.fc1(x))
 
         if self.use_dwconv:
@@ -94,6 +109,7 @@ class PVTEncoder(nn.Module):
     qkv_bias: bool
         Whether to add a bias vector to the q,k, and v matrices
     qk_scale:float, optional
+        Override default qk scale of head_dim ** -0.5 in Spatial Attention if set
     p_dropout: float
         Dropout probability
     attn_dropout: float
diff --git a/vformer/encoder/swin.py b/vformer/encoder/swin.py
index d483a54d..720e3d1f 100644
--- a/vformer/encoder/swin.py
+++ b/vformer/encoder/swin.py
@@ -17,6 +17,7 @@
 @ENCODER_REGISTRY.register()
 class SwinEncoderBlock(nn.Module):
     """
+
     Parameters
     ----------
     dim: int
@@ -34,14 +35,16 @@ class SwinEncoderBlock(nn.Module):
     qkv_bias: bool, default= True
         Whether to add a bias vector to the q,k, and v matrices
     qk_scale: float, Optional
+
     p_dropout: float
         Dropout rate
-    attn_drop: float
+    attn_dropout: float
         Dropout rate
-    drop_path: float
+    drop_path_rate: float
         Stochastic depth rate
     norm_layer:nn.Module
         Normalization layer, default is `nn.LayerNorm`
+
     """
 
     def __init__(
@@ -56,7 +59,7 @@ def __init__(
         qk_scale=None,
         p_dropout=0.0,
         attn_dropout=0.0,
-        drop_path=0.0,
+        drop_path_rate=0.0,
         norm_layer=nn.LayerNorm,
     ):
         super(SwinEncoderBlock, self).__init__()
@@ -88,7 +91,9 @@ def __init__(
             proj_dropout=p_dropout,
         )
 
-        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()
+        self.drop_path = (
+            DropPath(drop_path_rate) if drop_path_rate > 0.0 else nn.Identity()
+        )
         self.norm2 = norm_layer(dim)
         self.mlp = FeedForward(dim=dim, hidden_dim=hidden_dim, p_dropout=p_dropout)
 
@@ -105,6 +110,18 @@ def __init__(
         self.register_buffer("attn_mask", attn_mask)
 
     def forward(self, x):
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+
+        Returns
+        ----------
+        torch.Tensor
+            Returns output tensor
+
+        """
 
         H, W = self.input_resolution
         B, L, C = x.shape
@@ -159,11 +176,12 @@ class SwinEncoder(nn.Module):
     qkv_bias: bool, default is True
        Whether to add a bias vector to the q,k, and v matrices
     qk_scale: float, optional
+        Override default qk scale of head_dim ** -0.5 in Window Attention if set
     p_dropout: float,
         Dropout rate.
     attn_dropout: float, optional
         Attention dropout rate
-    drop_path: float or tuple[float]
+    drop_path_rate: float or tuple[float]
         Stochastic depth rate.
     norm_layer: nn.Module
         Normalization layer. default is nn.LayerNorm
@@ -209,7 +227,7 @@ def __init__(
                     qk_scale=qkv_scale,
                     p_dropout=p_dropout,
                     attn_dropout=attn_dropout,
-                    drop_path=drop_path[i]
+                    drop_path_rate=drop_path[i]
                     if isinstance(drop_path, list)
                     else drop_path,
                     norm_layer=norm_layer,
@@ -225,6 +243,18 @@ def __init__(
             self.downsample = None
 
     def forward(self, x):
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+
+        Returns
+        ----------
+        torch.Tensor
+            Returns output tensor
+
+        """
 
         for blk in self.blocks:
             if self.use_checkpoint:
diff --git a/vformer/encoder/vanilla.py b/vformer/encoder/vanilla.py
index 5a597c5d..b8885e22 100644
--- a/vformer/encoder/vanilla.py
+++ b/vformer/encoder/vanilla.py
@@ -13,13 +13,13 @@ class VanillaEncoder(nn.Module):
 
     Parameters
     ----------
-    latent_dim: int
+    embedding_dim: int
         Dimension of the embedding
     depth: int
         Number of self-attention layers
-    heads: int
+    num_heads: int
         Number of the attention heads
-    dim_head: int
+    head_dim: int
         Dimension of each head
     mlp_dim: int
         Dimension of the hidden layer in the feed-forward layer
@@ -33,10 +33,10 @@ class VanillaEncoder(nn.Module):
 
     def __init__(
         self,
-        latent_dim,
+        embedding_dim,
         depth,
         num_heads,
-        dim_head,
+        head_dim,
         mlp_dim,
         p_dropout=0.0,
         attn_dropout=0.0,
@@ -50,18 +50,20 @@ def __init__(
                 nn.ModuleList(
                     [
                         PreNorm(
-                            dim=latent_dim,
+                            dim=embedding_dim,
                             fn=VanillaSelfAttention(
-                                dim=latent_dim,
+                                dim=embedding_dim,
                                 num_heads=num_heads,
-                                head_dim=dim_head,
+                                head_dim=head_dim,
                                 p_dropout=attn_dropout,
                             ),
                         ),
                         PreNorm(
-                            dim=latent_dim,
+                            dim=embedding_dim,
                             fn=FeedForward(
-                                dim=latent_dim, hidden_dim=mlp_dim, p_dropout=p_dropout
+                                dim=embedding_dim,
+                                hidden_dim=mlp_dim,
+                                p_dropout=p_dropout,
                             ),
                         ),
                     ]
@@ -74,7 +76,17 @@ def __init__(
         )
 
     def forward(self, x):
+        """
 
+        Parameters
+        ----------
+        x: torch.Tensor
+
+        Returns
+        ----------
+        torch.Tensor
+            Returns output tensor
+        """
         for attn, ff in self.encoder:
             x = attn(x) + x
             x = self.drop_path(ff(x)) + x
diff --git a/vformer/models/classification/cct.py b/vformer/models/classification/cct.py
index b55754cf..6bbc36d3 100644
--- a/vformer/models/classification/cct.py
+++ b/vformer/models/classification/cct.py
@@ -4,7 +4,7 @@
 
 from ...common import BaseClassificationModel
 from ...decoder import MLPDecoder
-from ...encoder import CVTEmbedding, VanillaEncoder
+from ...encoder import CVTEmbedding, PosEmbedding, VanillaEncoder
 from ...utils import MODEL_REGISTRY, pair
 
 
@@ -60,7 +60,7 @@ def __init__(
         seq_pool=True,
         embedding_dim=768,
         num_layers=1,
-        dim_head=96,
+        head_dim=96,
         num_heads=1,
         mlp_ratio=4.0,
         num_classes=1000,
@@ -128,30 +128,23 @@ def __init__(
             self.attention_pool = nn.Linear(self.embedding_dim, 1)
 
         if positional_embedding != "none":
-
-            if positional_embedding == "learnable":
-                self.positional_emb = nn.Parameter(
-                    torch.zeros(1, self.sequence_length, embedding_dim),
-                    requires_grad=True,
-                )
-            else:
-                self.positional_emb = nn.Parameter(
-                    self.sinusoidal_embedding(self.sequence_length, embedding_dim),
-                    requires_grad=False,
-                )
-
+            self.positional_emb = PosEmbedding(
+                self.sequence_length,
+                dim=embedding_dim,
+                drop=p_dropout,
+                sinusoidal=True if positional_embedding is "sine" else False,
+            )
         else:
             self.positional_emb = None
 
-        self.dropout = nn.Dropout(p=p_dropout)
         dpr = [x.item() for x in torch.linspace(0, drop_path, num_layers)]
         self.encoder_blocks = nn.ModuleList(
             [
                 VanillaEncoder(
-                    latent_dim=embedding_dim,
+                    embedding_dim=embedding_dim,
                     num_heads=num_heads,
                     depth=1,
-                    dim_head=dim_head,
+                    head_dim=head_dim,
                     mlp_dim=hidden_dim,
                     p_dropout=p_dropout,
                     attn_dropout=attn_dropout,
@@ -176,7 +169,18 @@ def __init__(
             self.decoder = MLPDecoder(config=embedding_dim, n_classes=num_classes)
 
     def forward(self, x):
-
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Input tensor
+        Returns
+        ----------
+        torch.Tensor
+            Returns tensor of size `num_classes`
+
+        """
         x = self.embedding(x)
 
         if self.positional_emb is None and x.size(1) < self.sequence_length:
@@ -189,9 +193,7 @@ def forward(self, x):
             x = torch.cat((cls_token, x), dim=1)
 
         if self.positional_emb is not None:
-            x += self.positional_emb
-
-        x = self.dropout(x)
+            x = self.positional_emb(x)
 
         for blk in self.encoder_blocks:
             x = blk(x)
@@ -206,17 +208,3 @@ def forward(self, x):
         x = self.decoder(x)
 
         return x
-
-    @staticmethod
-    def sinusoidal_embedding(n_channels, dim):
-
-        pe = torch.FloatTensor(
-            [
-                [p / (10000 ** (2 * (i // 2) / dim)) for i in range(dim)]
-                for p in range(n_channels)
-            ]
-        )
-        pe[:, 0::2] = torch.sin(pe[:, 0::2])
-        pe[:, 1::2] = torch.cos(pe[:, 1::2])
-
-        return pe.unsqueeze(0)
diff --git a/vformer/models/classification/cross.py b/vformer/models/classification/cross.py
index e17b8ee8..cedea344 100644
--- a/vformer/models/classification/cross.py
+++ b/vformer/models/classification/cross.py
@@ -66,9 +66,9 @@ class CrossViT(BaseClassificationModel):
         Dimension of the hidden layer for the smaller patches
     latent_dim_l: int
         Dimension of the hidden layer for the larger patches
-    dim_head_s: int
+    head_dim_s: int
         Dimension of the head of the attention for the smaller patches
-    dim_head_l: int
+    head_dim_l: int
         Dimension of the head of the attention for the larger patches
     depth_s: int
         Number of attention layers in encoder for the smaller patches
@@ -195,7 +195,18 @@ def __init__(
             self.decoder_l = MLPDecoder(latent_dim_l, n_classes)
 
     def forward(self, img):
-
+        """
+
+        Parameters
+        ----------
+        img: torch.Tensor
+            Input tensor
+        Returns
+        ----------
+        torch.Tensor
+            Returns tensor of size `num_classes`
+
+        """
         emb_s = self.s(img)
         emb_l = self.l(img)
         emb_s, emb_l = self.encoder(emb_s, emb_l)
diff --git a/vformer/models/classification/cvt.py b/vformer/models/classification/cvt.py
index b6f5ca1b..da2662e0 100644
--- a/vformer/models/classification/cvt.py
+++ b/vformer/models/classification/cvt.py
@@ -4,7 +4,7 @@
 
 from ...common import BaseClassificationModel
 from ...decoder import MLPDecoder
-from ...encoder import CVTEmbedding, VanillaEncoder
+from ...encoder import CVTEmbedding, PosEmbedding, VanillaEncoder
 from ...utils import MODEL_REGISTRY, pair
 
 
@@ -53,7 +53,7 @@ def __init__(
         in_channels=3,
         seq_pool=True,
         embedding_dim=768,
-        dim_head=96,
+        head_dim=96,
         num_layers=1,
         num_heads=1,
         mlp_ratio=4.0,
@@ -115,31 +115,24 @@ def __init__(
             self.attention_pool = nn.Linear(self.embedding_dim, 1)
 
         if positional_embedding != "none":
-
-            if positional_embedding == "learnable":
-                self.positional_emb = nn.Parameter(
-                    torch.zeros(1, self.sequence_length, embedding_dim),
-                    requires_grad=True,
-                )
-            else:
-                self.positional_emb = nn.Parameter(
-                    self.sinusoidal_embedding(self.sequence_length, embedding_dim),
-                    requires_grad=False,
-                )
-
+            self.positional_emb = PosEmbedding(
+                shape=self.sequence_length,
+                dim=embedding_dim,
+                drop=p_dropout,
+                sinusoidal=True if positional_embedding is "sine" else False,
+            )
         else:
             self.positional_emb = None
 
-        self.dropout = nn.Dropout(p=p_dropout)
         dpr = [x.item() for x in torch.linspace(0, drop_path, num_layers)]
         self.encoder_blocks = nn.ModuleList(
             [
                 VanillaEncoder(
-                    latent_dim=embedding_dim,
+                    embedding_dim=embedding_dim,
                     num_heads=num_heads,
                     depth=1,
                     mlp_dim=hidden_dim,
-                    dim_head=dim_head,
+                    head_dim=head_dim,
                     p_dropout=p_dropout,
                     attn_dropout=attn_dropout,
                     drop_path_rate=dpr[i],
@@ -161,6 +154,18 @@ def __init__(
             self.decoder = MLPDecoder(config=embedding_dim, n_classes=num_classes)
 
     def forward(self, x):
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Input tensor
+        Returns
+        ----------
+        torch.Tensor
+            Returns tensor of size `num_classes`
+
+        """
 
         x = self.embedding(x)
 
@@ -174,9 +179,8 @@ def forward(self, x):
             x = torch.cat((cls_token, x), dim=1)
 
         if self.positional_emb is not None:
-            x += self.positional_emb
+            x = self.positional_emb(x)
 
-        x = self.dropout(x)
         for blk in self.encoder_blocks:
             x = blk(x)
 
@@ -190,17 +194,3 @@ def forward(self, x):
         x = self.decoder(x)
 
         return x
-
-    @staticmethod
-    def sinusoidal_embedding(n_channels, dim):
-
-        pe = torch.FloatTensor(
-            [
-                [p / (10000 ** (2 * (i // 2) / dim)) for i in range(dim)]
-                for p in range(n_channels)
-            ]
-        )
-        pe[:, 0::2] = torch.sin(pe[:, 0::2])
-        pe[:, 1::2] = torch.cos(pe[:, 1::2])
-
-        return pe.unsqueeze(0)
diff --git a/vformer/models/classification/pyramid.py b/vformer/models/classification/pyramid.py
index 12fcde32..a0b74c3e 100644
--- a/vformer/models/classification/pyramid.py
+++ b/vformer/models/classification/pyramid.py
@@ -4,7 +4,7 @@
 from timm.models.layers import trunc_normal_
 
 from ...decoder import MLPDecoder
-from ...encoder import AbsolutePositionEmbedding, OverlapPatchEmbed, PVTEncoder
+from ...encoder import OverlapPatchEmbed, PVTEncoder, PVTPosEmbedding
 from ...utils import MODEL_REGISTRY
 
 
@@ -35,15 +35,16 @@ class PVTClassification(nn.Module):
     qkv_bias: bool, default= True
         Adds bias to the qkv if true
     qk_scale: float, optional
+        Override default qk scale of head_dim ** -0.5 Spatial Attention if set
     p_dropout: float,
         Dropout rate,default is 0.0
-    attn_drop_rate:  float,
+    attn_dropout:  float,
         Attention dropout rate, default is 0.0
     drop_path_rate: float
         Stochastic depth rate, default is 0.1
     norm_layer:
         Normalization layer, default is nn.LayerNorm
-    sr_ratio: float
+    sr_ratios: float
         Spatial reduction ratio
     decoder_config:int or tuple[int], optional
         Configuration of the decoder. If None, the default configuration is used.
@@ -107,7 +108,7 @@ def __init__(
                     self.pos_embeds.append(
                         nn.ModuleList(
                             [
-                                AbsolutePositionEmbedding(
+                                PVTPosEmbedding(
                                     pos_shape=img_size // np.prod(patch_size[: i + 1]),
                                     pos_dim=embed_dims[i],
                                 )
@@ -163,6 +164,18 @@ def __init__(
             self.decoder = MLPDecoder(config=embed_dims[-1], n_classes=num_classes)
 
     def forward(self, x):
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Input tensor
+        Returns
+        ----------
+        torch.Tensor
+            Returns tensor of size `num_classes`
+
+        """
         B = x.shape[0]
         for i in range(len(self.depths)):
             patch_embed = self.patch_embeds[i]
@@ -216,15 +229,16 @@ class PVTClassificationV2(PVTClassification):
     qkv_bias: bool, default= True
         Adds bias to the qkv if true
     qk_scale: float, optional
+        Override default qk scale of head_dim ** -0.5 in Spatial Attention if set
     p_dropout: float,
         Dropout rate,default is 0.0
-    attn_drop_rate:  float,
+    attn_dropout:  float,
         Attention dropout rate, default is 0.0
     drop_path_rate: float
         Stochastic depth rate, default is 0.1
     norm_layer:nn.Module
         Normalization layer, default is nn.LayerNorm
-    sr_ratio: float
+    sr_ratios: float
         Spatial reduction ratio
     decoder_config:int or tuple[int], optional
         Configuration of the decoder. If None, the default configuration is used.
diff --git a/vformer/models/classification/swin.py b/vformer/models/classification/swin.py
index 36b5bea7..41cbe9b6 100644
--- a/vformer/models/classification/swin.py
+++ b/vformer/models/classification/swin.py
@@ -4,7 +4,7 @@
 
 from ...common import BaseClassificationModel
 from ...decoder import MLPDecoder
-from ...encoder import PatchEmbedding, SwinEncoder
+from ...encoder import PatchEmbedding, PosEmbedding, SwinEncoder
 from ...functional import PatchMerging
 from ...utils import MODEL_REGISTRY
 
@@ -38,13 +38,14 @@ class SwinTransformer(BaseClassificationModel):
     qkv_bias: bool, default= True
         Adds bias to the qkv if true
     qk_scale:  float, optional
-    drop_rate: float
+        Override default qk scale of head_dim ** -0.5 in Window Attention if set
+    p_dropout: float
         Dropout rate, default is 0.0
-    attn_drop_rate: float
+    attn_dropout: float
         Attention dropout rate,default is 0.0
     drop_path_rate: float
         Stochastic depth rate, default is 0.1
-    norm_layer:
+    norm_layer: nn.Module
         Normalization layer,default is nn.LayerNorm
     ape: bool, optional
         Whether to add relative/absolute position embedding to patch embedding, default is True
@@ -67,8 +68,8 @@ def __init__(
         mlp_ratio=4.0,
         qkv_bias=True,
         qk_scale=None,
-        drop_rate=0.0,
-        attn_drop_rate=0.0,
+        p_dropout=0.0,
+        attn_dropout=0.0,
         drop_path_rate=0.1,
         norm_layer=nn.LayerNorm,
         ape=True,
@@ -91,12 +92,11 @@ def __init__(
         self.ape = ape
         num_features = int(embedding_dim * 2 ** (len(depths) - 1))
 
-        if self.ape:
-            self.absolute_pos_embed = nn.Parameter(
-                torch.zeros(1, num_patches, embedding_dim)
-            )
-            trunc_normal_(self.absolute_pos_embed, std=0.02)
-
+        self.absolute_pos_embed = (
+            PosEmbedding(shape=num_patches, dim=embedding_dim, drop=p_dropout, std=0.02)
+            if ape
+            else nn.Identity()
+        )
         dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
         self.encoder = nn.ModuleList()
 
@@ -113,8 +113,8 @@ def __init__(
                 mlp_ratio=mlp_ratio,
                 qkv_bias=qkv_bias,
                 qkv_scale=qk_scale,
-                p_dropout=drop_rate,
-                attn_dropout=attn_drop_rate,
+                p_dropout=p_dropout,
+                attn_dropout=attn_dropout,
                 drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])],
                 norm_layer=norm_layer,
                 downsample=PatchMerging if i_layer < len(depths) - 1 else None,
@@ -137,16 +137,24 @@ def __init__(
 
         self.pool = nn.AdaptiveAvgPool1d(1)
         self.norm = norm_layer(num_features) if norm_layer is not None else nn.Identity
-        self.pos_drop = nn.Dropout(p=drop_rate)
+        self.pos_drop = nn.Dropout(p=p_dropout)
 
     def forward(self, x):
-
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Input tensor
+        Returns
+        ----------
+        torch.Tensor
+            Returns tensor of size `num_classes`
+
+        """
         x = self.patch_embed(x)
 
-        if self.ape:
-            x += self.absolute_pos_embed
-
-        x = self.pos_drop(x)
+        x = self.absolute_pos_embed(x)
 
         for layer in self.encoder:
             x = layer(x)
diff --git a/vformer/models/classification/vanilla.py b/vformer/models/classification/vanilla.py
index 5c19067f..eafff68b 100644
--- a/vformer/models/classification/vanilla.py
+++ b/vformer/models/classification/vanilla.py
@@ -4,7 +4,7 @@
 
 from ...common import BaseClassificationModel
 from ...decoder import MLPDecoder
-from ...encoder import LinearEmbedding, VanillaEncoder
+from ...encoder import LinearEmbedding, PosEmbedding, VanillaEncoder
 from ...utils import MODEL_REGISTRY
 
 
@@ -32,7 +32,7 @@ class VanillaViT(BaseClassificationModel):
         Number of the attention heads
     encoder_mlp_dim: int
         Dimension of hidden layer in the encoder
-    in_channel: int
+    in_channels: int
         Number of input channels
     decoder_config: int or tuple or list, optional
         Configuration of the decoder. If None, the default configuration is used.
@@ -66,19 +66,21 @@ def __init__(
             embedding_dim, self.patch_height, self.patch_width, self.patch_dim
         )
 
-        self.pos_embedding = nn.Parameter(
-            torch.randn(1, self.n_patches + 1, embedding_dim)
+        self.pos_embedding = PosEmbedding(
+            shape=self.n_patches + 1,
+            dim=embedding_dim,
+            drop=p_dropout_embedding,
+            sinusoidal=False,
         )
         self.cls_token = nn.Parameter(torch.randn(1, 1, embedding_dim))
-        self.embedding_dropout = nn.Dropout(p_dropout_embedding)
 
         self.encoder = VanillaEncoder(
-            embedding_dim,
-            depth,
-            attn_heads,
-            head_dim,
-            encoder_mlp_dim,
-            p_dropout_encoder,
+            embedding_dim=embedding_dim,
+            depth=depth,
+            num_heads=attn_heads,
+            head_dim=head_dim,
+            mlp_dim=encoder_mlp_dim,
+            p_dropout=p_dropout_encoder,
         )
         self.pool = lambda x: x.mean(dim=1) if pool == "mean" else x[:, 0]
 
@@ -97,14 +99,24 @@ def __init__(
             self.decoder = MLPDecoder(embedding_dim, n_classes)
 
     def forward(self, x):
-
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Input tensor
+        Returns
+        ----------
+        torch.Tensor
+            Returns tensor of size `num_classes`
+
+        """
         x = self.patch_embedding(x)
         b, n, _ = x.shape
 
         cls_tokens = repeat(self.cls_token, "() n d -> b n d", b=b)
         x = torch.cat((cls_tokens, x), dim=1)
-        x += self.pos_embedding[:, : (n + 1)]
-        x = self.embedding_dropout(x)
+        x = self.pos_embedding(x)
         x = self.encoder(x)
         x = self.pool(x)
         x = self.decoder(x)
diff --git a/vformer/models/dense/PVT/detection.py b/vformer/models/dense/PVT/detection.py
index a32391da..7b6a2c47 100644
--- a/vformer/models/dense/PVT/detection.py
+++ b/vformer/models/dense/PVT/detection.py
@@ -2,7 +2,7 @@
 import torch
 import torch.nn as nn
 
-from ....encoder import AbsolutePositionEmbedding, OverlapPatchEmbed, PVTEncoder
+from ....encoder import OverlapPatchEmbed, PVTEncoder, PVTPosEmbedding
 from ....utils import MODEL_REGISTRY
 
 
@@ -34,13 +34,14 @@ class PVTDetection(nn.Module):
     qkv_bias: bool, default= True
         Adds bias to the qkv if true
     qk_scale: float, optional
+        Override default qk scale of head_dim ** -0.5 in Spatial Attention if set
     p_dropout: float,
         Dropout rate,default is 0.0
     attn_dropout:  float,
         Attention dropout rate, default is 0.0
     drop_path_rate: float
         Stochastic depth rate, default is 0.1
-    sr_ratio: float
+    sr_ratios: float
         Spatial reduction ratio
     linear: bool
         Whether to use linear spatial attention, default is False
@@ -109,7 +110,7 @@ def __init__(
                 self.pos_embeds.append(
                     nn.ModuleList(
                         [
-                            AbsolutePositionEmbedding(
+                            PVTPosEmbedding(
                                 pos_shape=img_size // np.prod(patch_size[: i + 1]),
                                 pos_dim=embedding_dims[i],
                             )
@@ -144,6 +145,18 @@ def __init__(
 
     def forward(self, x):
 
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Input tensor
+        Returns
+        ----------
+        torch.Tensor
+            Returns list containing output features from all pyramid stages
+
+        """
         B = x.shape[0]
         out = []
 
@@ -197,13 +210,14 @@ class PVTDetectionV2(PVTDetection):
     qkv_bias: bool, default= True
         Adds bias to the qkv if true
     qk_scale: float, optional
+        Override default qk scale of head_dim ** -0.5 in Spatial Attention if set
     p_dropout: float,
         Dropout rate,default is 0.0
     attn_dropout:  float,
         Attention dropout rate, default is 0.0
     drop_path_rate: float
         Stochastic depth rate, default is 0.1
-    sr_ratio: float
+    sr_ratios: float
         Spatial reduction ratio
     linear: bool
         Whether to use linear spatial attention
@@ -211,8 +225,6 @@ class PVTDetectionV2(PVTDetection):
         Whether to use Depth-wise convolutions in Overlap-patch embedding
     ape: bool
         Whether to use absolute position embedding
-    return_pyramid: bool
-        Whether to return all pyramid feature layers, if false returns only last feature layer, default is True
     """
 
     def __init__(
@@ -231,7 +243,7 @@ def __init__(
         norm_layer=nn.LayerNorm,
         depths=[3, 4, 6, 3],
         sr_ratios=[8, 4, 2, 1],
-        use_abs_pos_embed=False,
+        ape=False,
         use_dwconv=True,
         linear=False,
     ):
@@ -251,6 +263,6 @@ def __init__(
             depths=depths,
             sr_ratios=sr_ratios,
             linear=linear,
-            ape=use_abs_pos_embed,
+            ape=ape,
             use_dwconv=use_dwconv,
         )
diff --git a/vformer/models/dense/PVT/segmentation.py b/vformer/models/dense/PVT/segmentation.py
index 620dd8d6..789aff66 100644
--- a/vformer/models/dense/PVT/segmentation.py
+++ b/vformer/models/dense/PVT/segmentation.py
@@ -3,7 +3,7 @@
 import torch.nn as nn
 
 from ....decoder import SegmentationHead
-from ....encoder import AbsolutePositionEmbedding, OverlapPatchEmbed, PVTEncoder
+from ....encoder import OverlapPatchEmbed, PVTEncoder, PVTPosEmbedding
 from ....utils import MODEL_REGISTRY
 
 
@@ -32,13 +32,14 @@ class PVTSegmentation(nn.Module):
     qkv_bias: bool, default= True
         Adds bias to the qkv if true
     qk_scale: float, optional
+        Override default qk scale of head_dim ** -0.5 in Spatial Attention if set
     p_dropout: float
         Dropout rate,default is 0.0
-    attn_drop_rate:  float
+    attn_dropout:  float
         Attention dropout rate, default is 0.0
     drop_path_rate: float
         Stochastic depth rate, default is 0.1
-    sr_ratio: float
+    sr_ratios: float
         Spatial reduction ratio
     linear: bool
         Whether to use linear spatial attention
@@ -110,7 +111,7 @@ def __init__(
                 self.pos_embeds.append(
                     nn.ModuleList(
                         [
-                            AbsolutePositionEmbedding(
+                            PVTPosEmbedding(
                                 pos_shape=img_size // np.prod(patch_size[: i + 1]),
                                 pos_dim=embedding_dims[i],
                             )
@@ -148,6 +149,17 @@ def __init__(
 
     def forward(self, x):
 
+        """
+
+        Parameters
+        ----------
+        x: torch.Tensor
+            Input tensor
+        Returns
+        ----------
+        torch.Tensor
+            Returns output tensor
+        """
         B = x.shape[0]
         out = []
 
@@ -202,13 +214,14 @@ class PVTSegmentationV2(PVTSegmentation):
     qkv_bias: bool, default= True
         Adds bias to the qkv if true
     qk_scale: float, optional
+        Override default qk scale of head_dim ** -0.5 in Spatial Attention if set
     p_dropout: float,
         Dropout rate,default is 0.0
     attn_dropout:  float,
         Attention dropout rate, default is 0.0
     drop_path_rate: float
         Stochastic depth rate, default is 0.1
-    sr_ratio: float
+    sr_ratios: float
         Spatial reduction ratio
     linear: bool
         Whether to use linear spatial attention, default is False


Test patch:
diff --git a/tests/test_encoder.py b/tests/test_encoder.py
index 6d7ffa8a..cae8386a 100644
--- a/tests/test_encoder.py
+++ b/tests/test_encoder.py
@@ -11,7 +11,7 @@ def test_VanillaEncoder():
 
     test_tensor = torch.randn(2, 65, 1024)
     encoder = ENCODER_REGISTRY.get("VanillaEncoder")(
-        latent_dim=1024, depth=6, num_heads=16, dim_head=64, mlp_dim=2048
+        embedding_dim=1024, depth=6, num_heads=16, head_dim=64, mlp_dim=2048
     )
     out = encoder(test_tensor)
     assert out.shape == test_tensor.shape  # shape remains same
diff --git a/tests/test_models.py b/tests/test_models.py
index 1c7483e7..5995aaf1 100644
--- a/tests/test_models.py
+++ b/tests/test_models.py
@@ -45,8 +45,8 @@ def test_SwinTransformer():
         mlp_ratio=4.0,
         qkv_bias=True,
         qk_scale=None,
-        drop_rate=0.0,
-        attn_drop_rate=0.0,
+        p_dropout=0.0,
+        attn_dropout=0.0,
         drop_path_rate=0.1,
         norm_layer=nn.LayerNorm,
         ape=False,
@@ -66,7 +66,7 @@ def test_SwinTransformer():
         depths=[2, 2, 6, 2],
         num_heads=[3, 6, 12, 24],
         window_size=7,
-        drop_rate=0.2,
+        p_dropout=0.2,
     )
     out = model(img_3channels_224)
     assert out.shape == (4, 10)
@@ -82,7 +82,7 @@ def test_SwinTransformer():
         depths=[2, 2, 6, 2],
         num_heads=[4, 8, 16, 32],
         window_size=8,
-        drop_rate=0.2,
+        p_dropout=0.2,
     )
     out = model(img_3channels_256)
     assert out.shape == (2, 10)
@@ -98,7 +98,7 @@ def test_SwinTransformer():
         depths=[2, 2, 6, 2],
         num_heads=[3, 6, 12, 24],
         window_size=7,
-        drop_rate=0.2,
+        p_dropout=0.2,
     )
     out = model(img_1channels_224)
     assert out.shape == (2, 10)
@@ -114,7 +114,7 @@ def test_SwinTransformer():
         depths=[2, 2, 6, 2],
         num_heads=[3, 6, 12, 24],
         window_size=7,
-        drop_rate=0.2,
+        p_dropout=0.2,
         decoder_config=(768, 256, 10, 2),
     )
     out = model(img_3channels_224)
@@ -131,7 +131,7 @@ def test_SwinTransformer():
         depths=[2, 2, 6, 2],
         num_heads=[3, 6, 12, 24],
         window_size=7,
-        drop_rate=0.2,
+        p_dropout=0.2,
         decoder_config=(768, 256, 10, 2),
         ape=False,
     )
@@ -161,7 +161,6 @@ def test_CrossVit():
 
 
 def test_pvt():
-
     # classification
     model = MODEL_REGISTRY.get("PVTClassification")(
         patch_size=[7, 3, 3, 3],