ultralytics 8.3.142__py3-none-any.whl → 8.3.144__py3-none-any.whl

ultralytics/models/sam/modules/tiny_encoder.py

@@ -10,7 +10,7 @@
 # --------------------------------------------------------
 
 import itertools
-from typing import Tuple
+from typing import List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -24,32 +24,46 @@ class Conv2d_BN(torch.nn.Sequential):
     """
     A sequential container that performs 2D convolution followed by batch normalization.
 
+    This module combines a 2D convolution layer with batch normalization, providing a common building block
+    for convolutional neural networks. The batch normalization weights and biases are initialized to specific
+    values for optimal training performance.
+
     Attributes:
         c (torch.nn.Conv2d): 2D convolution layer.
         bn (torch.nn.BatchNorm2d): Batch normalization layer.
 
-    Methods:
-        __init__: Initializes the Conv2d_BN with specified parameters.
-
-    Args:
-        a (int): Number of input channels.
-        b (int): Number of output channels.
-        ks (int): Kernel size for the convolution. Defaults to 1.
-        stride (int): Stride for the convolution. Defaults to 1.
-        pad (int): Padding for the convolution. Defaults to 0.
-        dilation (int): Dilation factor for the convolution. Defaults to 1.
-        groups (int): Number of groups for the convolution. Defaults to 1.
-        bn_weight_init (float): Initial value for batch normalization weight. Defaults to 1.
-
     Examples:
         >>> conv_bn = Conv2d_BN(3, 64, ks=3, stride=1, pad=1)
         >>> input_tensor = torch.randn(1, 3, 224, 224)
         >>> output = conv_bn(input_tensor)
         >>> print(output.shape)
+        torch.Size([1, 64, 224, 224])
     """
 
-    def __init__(self, a, b, ks=1, stride=1, pad=0, dilation=1, groups=1, bn_weight_init=1):
-        """Initializes a sequential container with 2D convolution followed by batch normalization."""
+    def __init__(
+        self,
+        a: int,
+        b: int,
+        ks: int = 1,
+        stride: int = 1,
+        pad: int = 0,
+        dilation: int = 1,
+        groups: int = 1,
+        bn_weight_init: float = 1,
+    ):
+        """
+        Initialize a sequential container with 2D convolution followed by batch normalization.
+
+        Args:
+            a (int): Number of input channels.
+            b (int): Number of output channels.
+            ks (int, optional): Kernel size for the convolution.
+            stride (int, optional): Stride for the convolution.
+            pad (int, optional): Padding for the convolution.
+            dilation (int, optional): Dilation factor for the convolution.
+            groups (int, optional): Number of groups for the convolution.
+            bn_weight_init (float, optional): Initial value for batch normalization weight.
+        """
         super().__init__()
         self.add_module("c", torch.nn.Conv2d(a, b, ks, stride, pad, dilation, groups, bias=False))
         bn = torch.nn.BatchNorm2d(b)
@@ -60,7 +74,10 @@ class Conv2d_BN(torch.nn.Sequential):
 
 class PatchEmbed(nn.Module):
     """
-    Embeds images into patches and projects them into a specified embedding dimension.
+    Embed images into patches and project them into a specified embedding dimension.
+
+    This module converts input images into patch embeddings using a sequence of convolutional layers,
+    effectively downsampling the spatial dimensions while increasing the channel dimension.
 
     Attributes:
         patches_resolution (Tuple[int, int]): Resolution of the patches after embedding.
@@ -69,19 +86,25 @@ class PatchEmbed(nn.Module):
         embed_dim (int): Dimension of the embedding.
         seq (nn.Sequential): Sequence of convolutional and activation layers for patch embedding.
 
-    Methods:
-        forward: Processes the input tensor through the patch embedding sequence.
-
     Examples:
         >>> import torch
         >>> patch_embed = PatchEmbed(in_chans=3, embed_dim=96, resolution=224, activation=nn.GELU)
         >>> x = torch.randn(1, 3, 224, 224)
         >>> output = patch_embed(x)
         >>> print(output.shape)
+        torch.Size([1, 96, 56, 56])
     """
 
-    def __init__(self, in_chans, embed_dim, resolution, activation):
-        """Initializes patch embedding with convolutional layers for image-to-patch conversion and projection."""
+    def __init__(self, in_chans: int, embed_dim: int, resolution: int, activation):
+        """
+        Initialize patch embedding with convolutional layers for image-to-patch conversion and projection.
+
+        Args:
+            in_chans (int): Number of input channels.
+            embed_dim (int): Dimension of the embedding.
+            resolution (int): Input image resolution.
+            activation (nn.Module): Activation function to use between convolutions.
+        """
         super().__init__()
         img_size: Tuple[int, int] = to_2tuple(resolution)
         self.patches_resolution = (img_size[0] // 4, img_size[1] // 4)
@@ -95,8 +118,8 @@ class PatchEmbed(nn.Module):
             Conv2d_BN(n // 2, n, 3, 2, 1),
         )
 
-    def forward(self, x):
-        """Processes input tensor through patch embedding sequence, converting images to patch embeddings."""
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Process input tensor through patch embedding sequence, converting images to patch embeddings."""
         return self.seq(x)
 
 
@@ -104,21 +127,21 @@ class MBConv(nn.Module):
     """
     Mobile Inverted Bottleneck Conv (MBConv) layer, part of the EfficientNet architecture.
 
+    This module implements the mobile inverted bottleneck convolution with expansion, depthwise convolution,
+    and projection phases, along with residual connections for improved gradient flow.
+
     Attributes:
         in_chans (int): Number of input channels.
-        hidden_chans (int): Number of hidden channels.
+        hidden_chans (int): Number of hidden channels after expansion.
         out_chans (int): Number of output channels.
-        conv1 (Conv2d_BN): First convolutional layer.
+        conv1 (Conv2d_BN): First convolutional layer for channel expansion.
         act1 (nn.Module): First activation function.
         conv2 (Conv2d_BN): Depthwise convolutional layer.
         act2 (nn.Module): Second activation function.
-        conv3 (Conv2d_BN): Final convolutional layer.
+        conv3 (Conv2d_BN): Final convolutional layer for projection.
         act3 (nn.Module): Third activation function.
         drop_path (nn.Module): Drop path layer (Identity for inference).
 
-    Methods:
-        forward: Performs the forward pass through the MBConv layer.
-
     Examples:
         >>> in_chans, out_chans = 32, 64
         >>> mbconv = MBConv(in_chans, out_chans, expand_ratio=4, activation=nn.ReLU, drop_path=0.1)
@@ -128,8 +151,17 @@ class MBConv(nn.Module):
         torch.Size([1, 64, 56, 56])
     """
 
-    def __init__(self, in_chans, out_chans, expand_ratio, activation, drop_path):
-        """Initializes the MBConv layer with specified input/output channels, expansion ratio, and activation."""
+    def __init__(self, in_chans: int, out_chans: int, expand_ratio: float, activation, drop_path: float):
+        """
+        Initialize the MBConv layer with specified input/output channels, expansion ratio, and activation.
+
+        Args:
+            in_chans (int): Number of input channels.
+            out_chans (int): Number of output channels.
+            expand_ratio (float): Channel expansion ratio for the hidden layer.
+            activation (nn.Module): Activation function to use.
+            drop_path (float): Drop path rate for stochastic depth.
+        """
         super().__init__()
         self.in_chans = in_chans
         self.hidden_chans = int(in_chans * expand_ratio)
@@ -148,8 +180,8 @@ class MBConv(nn.Module):
         # self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity()
         self.drop_path = nn.Identity()
 
-    def forward(self, x):
-        """Implements the forward pass of MBConv, applying convolutions and skip connection."""
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Implement the forward pass of MBConv, applying convolutions and skip connection."""
         shortcut = x
         x = self.conv1(x)
         x = self.act1(x)
@@ -163,10 +195,11 @@ class MBConv(nn.Module):
 
 class PatchMerging(nn.Module):
     """
-    Merges neighboring patches in the feature map and projects to a new dimension.
+    Merge neighboring patches in the feature map and project to a new dimension.
 
     This class implements a patch merging operation that combines spatial information and adjusts the feature
-    dimension. It uses a series of convolutional layers with batch normalization to achieve this.
+    dimension using a series of convolutional layers with batch normalization. It effectively reduces spatial
+    resolution while potentially increasing channel dimensions.
 
     Attributes:
         input_resolution (Tuple[int, int]): The input resolution (height, width) of the feature map.
@@ -177,19 +210,25 @@ class PatchMerging(nn.Module):
         conv2 (Conv2d_BN): The second convolutional layer for spatial merging.
         conv3 (Conv2d_BN): The third convolutional layer for final projection.
 
-    Methods:
-        forward: Applies the patch merging operation to the input tensor.
-
     Examples:
         >>> input_resolution = (56, 56)
         >>> patch_merging = PatchMerging(input_resolution, dim=64, out_dim=128, activation=nn.ReLU)
         >>> x = torch.randn(4, 64, 56, 56)
         >>> output = patch_merging(x)
         >>> print(output.shape)
+        torch.Size([4, 3136, 128])
     """
 
-    def __init__(self, input_resolution, dim, out_dim, activation):
-        """Initializes the PatchMerging module for merging and projecting neighboring patches in feature maps."""
+    def __init__(self, input_resolution: Tuple[int, int], dim: int, out_dim: int, activation):
+        """
+        Initialize the PatchMerging module for merging and projecting neighboring patches in feature maps.
+
+        Args:
+            input_resolution (Tuple[int, int]): The input resolution (height, width) of the feature map.
+            dim (int): The input dimension of the feature map.
+            out_dim (int): The output dimension after merging and projection.
+            activation (nn.Module): The activation function used between convolutions.
+        """
         super().__init__()
 
         self.input_resolution = input_resolution
@@ -201,8 +240,8 @@ class PatchMerging(nn.Module):
         self.conv2 = Conv2d_BN(out_dim, out_dim, 3, stride_c, 1, groups=out_dim)
         self.conv3 = Conv2d_BN(out_dim, out_dim, 1, 1, 0)
 
-    def forward(self, x):
-        """Applies patch merging and dimension projection to the input feature map."""
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply patch merging and dimension projection to the input feature map."""
         if x.ndim == 3:
             H, W = self.input_resolution
             B = len(x)
@@ -222,7 +261,8 @@ class ConvLayer(nn.Module):
     """
     Convolutional Layer featuring multiple MobileNetV3-style inverted bottleneck convolutions (MBConv).
 
-    This layer optionally applies downsample operations to the output and supports gradient checkpointing.
+    This layer optionally applies downsample operations to the output and supports gradient checkpointing
+    for memory efficiency during training.
 
     Attributes:
         dim (int): Dimensionality of the input and output.
@@ -230,32 +270,30 @@ class ConvLayer(nn.Module):
         depth (int): Number of MBConv layers in the block.
         use_checkpoint (bool): Whether to use gradient checkpointing to save memory.
         blocks (nn.ModuleList): List of MBConv layers.
-        downsample (Optional[Callable]): Function for downsampling the output.
-
-    Methods:
-        forward: Processes the input through the convolutional layers.
+        downsample (Optional[nn.Module]): Function for downsampling the output.
 
     Examples:
         >>> input_tensor = torch.randn(1, 64, 56, 56)
         >>> conv_layer = ConvLayer(64, (56, 56), depth=3, activation=nn.ReLU)
         >>> output = conv_layer(input_tensor)
         >>> print(output.shape)
+        torch.Size([1, 3136, 128])
     """
 
     def __init__(
         self,
-        dim,
-        input_resolution,
-        depth,
+        dim: int,
+        input_resolution: Tuple[int, int],
+        depth: int,
         activation,
-        drop_path=0.0,
-        downsample=None,
-        use_checkpoint=False,
-        out_dim=None,
-        conv_expand_ratio=4.0,
+        drop_path: Union[float, List[float]] = 0.0,
+        downsample: Optional[nn.Module] = None,
+        use_checkpoint: bool = False,
+        out_dim: Optional[int] = None,
+        conv_expand_ratio: float = 4.0,
     ):
         """
-        Initializes the ConvLayer with the given dimensions and settings.
+        Initialize the ConvLayer with the given dimensions and settings.
 
         This layer consists of multiple MobileNetV3-style inverted bottleneck convolutions (MBConv) and
         optionally applies downsampling to the output.
@@ -265,17 +303,11 @@ class ConvLayer(nn.Module):
             input_resolution (Tuple[int, int]): The resolution of the input image.
             depth (int): The number of MBConv layers in the block.
             activation (nn.Module): Activation function applied after each convolution.
-            drop_path (float | List[float]): Drop path rate. Single float or a list of floats for each MBConv.
-            downsample (Optional[nn.Module]): Function for downsampling the output. None to skip downsampling.
-            use_checkpoint (bool): Whether to use gradient checkpointing to save memory.
-            out_dim (Optional[int]): The dimensionality of the output. None means it will be the same as `dim`.
-            conv_expand_ratio (float): Expansion ratio for the MBConv layers.
-
-        Examples:
-            >>> input_tensor = torch.randn(1, 64, 56, 56)
-            >>> conv_layer = ConvLayer(64, (56, 56), depth=3, activation=nn.ReLU)
-            >>> output = conv_layer(input_tensor)
-            >>> print(output.shape)
+            drop_path (float | List[float], optional): Drop path rate. Single float or a list of floats for each MBConv.
+            downsample (Optional[nn.Module], optional): Function for downsampling the output. None to skip downsampling.
+            use_checkpoint (bool, optional): Whether to use gradient checkpointing to save memory.
+            out_dim (Optional[int], optional): The dimensionality of the output. None means it will be the same as `dim`.
+            conv_expand_ratio (float, optional): Expansion ratio for the MBConv layers.
         """
         super().__init__()
         self.dim = dim
@@ -304,19 +336,19 @@ class ConvLayer(nn.Module):
             else downsample(input_resolution, dim=dim, out_dim=out_dim, activation=activation)
         )
 
-    def forward(self, x):
-        """Processes input through convolutional layers, applying MBConv blocks and optional downsampling."""
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Process input through convolutional layers, applying MBConv blocks and optional downsampling."""
         for blk in self.blocks:
             x = torch.utils.checkpoint(blk, x) if self.use_checkpoint else blk(x)  # warn: checkpoint is slow import
         return x if self.downsample is None else self.downsample(x)
 
 
-class Mlp(nn.Module):
+class MLP(nn.Module):
     """
     Multi-layer Perceptron (MLP) module for transformer architectures.
 
     This module applies layer normalization, two fully-connected layers with an activation function in between,
-    and dropout. It is commonly used in transformer-based architectures.
+    and dropout. It is commonly used in transformer-based architectures for processing token embeddings.
 
     Attributes:
         norm (nn.LayerNorm): Layer normalization applied to the input.
@@ -325,32 +357,45 @@ class Mlp(nn.Module):
         act (nn.Module): Activation function applied after the first fully-connected layer.
         drop (nn.Dropout): Dropout layer applied after the activation function.
 
-    Methods:
-        forward: Applies the MLP operations on the input tensor.
-
     Examples:
         >>> import torch
         >>> from torch import nn
-        >>> mlp = Mlp(in_features=256, hidden_features=512, out_features=256, act_layer=nn.GELU, drop=0.1)
+        >>> mlp = MLP(in_features=256, hidden_features=512, out_features=256, activation=nn.GELU, drop=0.1)
         >>> x = torch.randn(32, 100, 256)
         >>> output = mlp(x)
         >>> print(output.shape)
         torch.Size([32, 100, 256])
     """
 
-    def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.0):
-        """Initializes a multi-layer perceptron with configurable input, hidden, and output dimensions."""
+    def __init__(
+        self,
+        in_features: int,
+        hidden_features: Optional[int] = None,
+        out_features: Optional[int] = None,
+        activation=nn.GELU,
+        drop: float = 0.0,
+    ):
+        """
+        Initialize a multi-layer perceptron with configurable input, hidden, and output dimensions.
+
+        Args:
+            in_features (int): Number of input features.
+            hidden_features (Optional[int], optional): Number of hidden features.
+            out_features (Optional[int], optional): Number of output features.
+            activation (nn.Module): Activation function applied after the first fully-connected layer.
+            drop (float, optional): Dropout probability.
+        """
         super().__init__()
         out_features = out_features or in_features
         hidden_features = hidden_features or in_features
         self.norm = nn.LayerNorm(in_features)
         self.fc1 = nn.Linear(in_features, hidden_features)
         self.fc2 = nn.Linear(hidden_features, out_features)
-        self.act = act_layer()
+        self.act = activation()
         self.drop = nn.Dropout(drop)
 
-    def forward(self, x):
-        """Applies MLP operations: layer norm, FC layers, activation, and dropout to the input tensor."""
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply MLP operations: layer norm, FC layers, activation, and dropout to the input tensor."""
         x = self.norm(x)
         x = self.fc1(x)
         x = self.act(x)
@@ -379,12 +424,8 @@ class Attention(torch.nn.Module):
         qkv (nn.Linear): Linear layer for computing query, key, and value projections.
         proj (nn.Linear): Linear layer for final projection.
         attention_biases (nn.Parameter): Learnable attention biases.
-        attention_bias_idxs (Tensor): Indices for attention biases.
-        ab (Tensor): Cached attention biases for inference, deleted during training.
-
-    Methods:
-        train: Sets the module in training mode and handles the 'ab' attribute.
-        forward: Performs the forward pass of the attention mechanism.
+        attention_bias_idxs (torch.Tensor): Indices for attention biases.
+        ab (torch.Tensor): Cached attention biases for inference, deleted during training.
 
     Examples:
         >>> attn = Attention(dim=256, key_dim=64, num_heads=8, resolution=(14, 14))
@@ -396,14 +437,14 @@ class Attention(torch.nn.Module):
 
     def __init__(
         self,
-        dim,
-        key_dim,
-        num_heads=8,
-        attn_ratio=4,
-        resolution=(14, 14),
+        dim: int,
+        key_dim: int,
+        num_heads: int = 8,
+        attn_ratio: float = 4,
+        resolution: Tuple[int, int] = (14, 14),
     ):
         """
-        Initializes the Attention module for multi-head attention with spatial awareness.
+        Initialize the Attention module for multi-head attention with spatial awareness.
 
         This module implements a multi-head attention mechanism with support for spatial awareness, applying
         attention biases based on spatial resolution. It includes trainable attention biases for each unique
@@ -412,16 +453,9 @@ class Attention(torch.nn.Module):
         Args:
             dim (int): The dimensionality of the input and output.
             key_dim (int): The dimensionality of the keys and queries.
-            num_heads (int): Number of attention heads.
-            attn_ratio (float): Attention ratio, affecting the dimensions of the value vectors.
-            resolution (Tuple[int, int]): Spatial resolution of the input feature map.
-
-        Examples:
-            >>> attn = Attention(dim=256, key_dim=64, num_heads=8, resolution=(14, 14))
-            >>> x = torch.randn(1, 196, 256)
-            >>> output = attn(x)
-            >>> print(output.shape)
-            torch.Size([1, 196, 256])
+            num_heads (int, optional): Number of attention heads.
+            attn_ratio (float, optional): Attention ratio, affecting the dimensions of the value vectors.
+            resolution (Tuple[int, int], optional): Spatial resolution of the input feature map.
         """
         super().__init__()
 
@@ -453,16 +487,16 @@ class Attention(torch.nn.Module):
         self.register_buffer("attention_bias_idxs", torch.LongTensor(idxs).view(N, N), persistent=False)
 
     @torch.no_grad()
-    def train(self, mode=True):
-        """Performs multi-head attention with spatial awareness and trainable attention biases."""
+    def train(self, mode: bool = True):
+        """Set the module in training mode and handle the 'ab' attribute for cached attention biases."""
         super().train(mode)
         if mode and hasattr(self, "ab"):
             del self.ab
         else:
             self.ab = self.attention_biases[:, self.attention_bias_idxs]
 
-    def forward(self, x):  # x
-        """Applies multi-head attention with spatial awareness and trainable attention biases."""
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply multi-head attention with spatial awareness and trainable attention biases."""
         B, N, _ = x.shape  # B, N, C
 
         # Normalization
@@ -490,7 +524,8 @@ class TinyViTBlock(nn.Module):
     TinyViT Block that applies self-attention and a local convolution to the input.
 
     This block is a key component of the TinyViT architecture, combining self-attention mechanisms with
-    local convolutions to process input features efficiently.
+    local convolutions to process input features efficiently. It supports windowed attention for
+    computational efficiency and includes residual connections.
 
     Attributes:
         dim (int): The dimensionality of the input and output.
@@ -500,13 +535,9 @@ class TinyViTBlock(nn.Module):
         mlp_ratio (float): Ratio of MLP hidden dimension to embedding dimension.
         drop_path (nn.Module): Stochastic depth layer, identity function during inference.
         attn (Attention): Self-attention module.
-        mlp (Mlp): Multi-layer perceptron module.
+        mlp (MLP): Multi-layer perceptron module.
         local_conv (Conv2d_BN): Depth-wise local convolution layer.
 
-    Methods:
-        forward: Processes the input through the TinyViT block.
-        extra_repr: Returns a string with extra information about the block's parameters.
-
     Examples:
         >>> input_tensor = torch.randn(1, 196, 192)
         >>> block = TinyViTBlock(dim=192, input_resolution=(14, 14), num_heads=3)
@@ -517,18 +548,18 @@ class TinyViTBlock(nn.Module):
 
     def __init__(
         self,
-        dim,
-        input_resolution,
-        num_heads,
-        window_size=7,
-        mlp_ratio=4.0,
-        drop=0.0,
-        drop_path=0.0,
-        local_conv_size=3,
+        dim: int,
+        input_resolution: Tuple[int, int],
+        num_heads: int,
+        window_size: int = 7,
+        mlp_ratio: float = 4.0,
+        drop: float = 0.0,
+        drop_path: float = 0.0,
+        local_conv_size: int = 3,
         activation=nn.GELU,
     ):
         """
-        Initializes a TinyViT block with self-attention and local convolution.
+        Initialize a TinyViT block with self-attention and local convolution.
 
         This block is a key component of the TinyViT architecture, combining self-attention mechanisms with
         local convolutions to process input features efficiently.
@@ -537,23 +568,12 @@ class TinyViTBlock(nn.Module):
             dim (int): Dimensionality of the input and output features.
             input_resolution (Tuple[int, int]): Spatial resolution of the input feature map (height, width).
             num_heads (int): Number of attention heads.
-            window_size (int): Size of the attention window. Must be greater than 0.
-            mlp_ratio (float): Ratio of MLP hidden dimension to embedding dimension.
-            drop (float): Dropout rate.
-            drop_path (float): Stochastic depth rate.
-            local_conv_size (int): Kernel size of the local convolution.
-            activation (torch.nn.Module): Activation function for MLP.
-
-        Raises:
-            AssertionError: If window_size is not greater than 0.
-            AssertionError: If dim is not divisible by num_heads.
-
-        Examples:
-            >>> block = TinyViTBlock(dim=192, input_resolution=(14, 14), num_heads=3)
-            >>> input_tensor = torch.randn(1, 196, 192)
-            >>> output = block(input_tensor)
-            >>> print(output.shape)
-            torch.Size([1, 196, 192])
+            window_size (int, optional): Size of the attention window. Must be greater than 0.
+            mlp_ratio (float, optional): Ratio of MLP hidden dimension to embedding dimension.
+            drop (float, optional): Dropout rate.
+            drop_path (float, optional): Stochastic depth rate.
+            local_conv_size (int, optional): Kernel size of the local convolution.
+            activation (nn.Module): Activation function for MLP.
         """
         super().__init__()
         self.dim = dim
@@ -575,13 +595,13 @@ class TinyViTBlock(nn.Module):
 
         mlp_hidden_dim = int(dim * mlp_ratio)
         mlp_activation = activation
-        self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=mlp_activation, drop=drop)
+        self.mlp = MLP(in_features=dim, hidden_features=mlp_hidden_dim, activation=mlp_activation, drop=drop)
 
         pad = local_conv_size // 2
         self.local_conv = Conv2d_BN(dim, dim, ks=local_conv_size, stride=1, pad=pad, groups=dim)
 
-    def forward(self, x):
-        """Applies self-attention, local convolution, and MLP operations to the input tensor."""
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Apply self-attention, local convolution, and MLP operations to the input tensor."""
         h, w = self.input_resolution
         b, hw, c = x.shape  # batch, height*width, channels
         assert hw == h * w, "input feature has wrong size"
@@ -624,7 +644,7 @@ class TinyViTBlock(nn.Module):
 
     def extra_repr(self) -> str:
         """
-        Returns a string representation of the TinyViTBlock's parameters.
+        Return a string representation of the TinyViTBlock's parameters.
 
         This method provides a formatted string containing key information about the TinyViTBlock, including its
         dimension, input resolution, number of attention heads, window size, and MLP ratio.
@@ -648,7 +668,8 @@ class BasicLayer(nn.Module):
     A basic TinyViT layer for one stage in a TinyViT architecture.
 
     This class represents a single layer in the TinyViT model, consisting of multiple TinyViT blocks
-    and an optional downsampling operation.
+    and an optional downsampling operation. It processes features at a specific resolution and
+    dimensionality within the overall architecture.
 
     Attributes:
         dim (int): The dimensionality of the input and output features.
@@ -658,10 +679,6 @@ class BasicLayer(nn.Module):
         blocks (nn.ModuleList): List of TinyViT blocks that make up this layer.
         downsample (nn.Module | None): Downsample layer at the end of the layer, if specified.
 
-    Methods:
-        forward: Processes the input through the layer's blocks and optional downsampling.
-        extra_repr: Returns a string with the layer's parameters for printing.
-
     Examples:
         >>> input_tensor = torch.randn(1, 3136, 192)
         >>> layer = BasicLayer(dim=192, input_resolution=(56, 56), depth=2, num_heads=3, window_size=7)
@@ -672,22 +689,22 @@ class BasicLayer(nn.Module):
 
     def __init__(
         self,
-        dim,
-        input_resolution,
-        depth,
-        num_heads,
-        window_size,
-        mlp_ratio=4.0,
-        drop=0.0,
-        drop_path=0.0,
-        downsample=None,
-        use_checkpoint=False,
-        local_conv_size=3,
+        dim: int,
+        input_resolution: Tuple[int, int],
+        depth: int,
+        num_heads: int,
+        window_size: int,
+        mlp_ratio: float = 4.0,
+        drop: float = 0.0,
+        drop_path: Union[float, List[float]] = 0.0,
+        downsample: Optional[nn.Module] = None,
+        use_checkpoint: bool = False,
+        local_conv_size: int = 3,
         activation=nn.GELU,
-        out_dim=None,
+        out_dim: Optional[int] = None,
     ):
         """
-        Initializes a BasicLayer in the TinyViT architecture.
+        Initialize a BasicLayer in the TinyViT architecture.
 
         This layer consists of multiple TinyViT blocks and an optional downsampling operation. It is designed to
         process feature maps at a specific resolution and dimensionality within the TinyViT model.
@@ -698,23 +715,14 @@ class BasicLayer(nn.Module):
             depth (int): Number of TinyViT blocks in this layer.
             num_heads (int): Number of attention heads in each TinyViT block.
             window_size (int): Size of the local window for attention computation.
-            mlp_ratio (float): Ratio of MLP hidden dimension to embedding dimension.
-            drop (float): Dropout rate.
-            drop_path (float | List[float]): Stochastic depth rate. Can be a float or a list of floats for each block.
-            downsample (nn.Module | None): Downsampling layer at the end of the layer. None to skip downsampling.
-            use_checkpoint (bool): Whether to use gradient checkpointing to save memory.
-            local_conv_size (int): Kernel size for the local convolution in each TinyViT block.
+            mlp_ratio (float, optional): Ratio of MLP hidden dimension to embedding dimension.
+            drop (float, optional): Dropout rate.
+            drop_path (float | List[float], optional): Stochastic depth rate. Can be a float or a list of floats for each block.
+            downsample (nn.Module | None, optional): Downsampling layer at the end of the layer. None to skip downsampling.
+            use_checkpoint (bool, optional): Whether to use gradient checkpointing to save memory.
+            local_conv_size (int, optional): Kernel size for the local convolution in each TinyViT block.
             activation (nn.Module): Activation function used in the MLP.
-            out_dim (int | None): Output dimension after downsampling. None means it will be the same as `dim`.
-
-        Raises:
-            ValueError: If `drop_path` is a list and its length doesn't match `depth`.
-
-        Examples:
-            >>> layer = BasicLayer(dim=96, input_resolution=(56, 56), depth=2, num_heads=3, window_size=7)
-            >>> x = torch.randn(1, 56 * 56, 96)
-            >>> output = layer(x)
-            >>> print(output.shape)
+            out_dim (int | None, optional): Output dimension after downsampling. None means it will be the same as `dim`.
         """
         super().__init__()
         self.dim = dim
@@ -747,14 +755,14 @@ class BasicLayer(nn.Module):
             else downsample(input_resolution, dim=dim, out_dim=out_dim, activation=activation)
         )
 
-    def forward(self, x):
-        """Processes input through TinyViT blocks and optional downsampling."""
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Process input through TinyViT blocks and optional downsampling."""
         for blk in self.blocks:
             x = torch.utils.checkpoint(blk, x) if self.use_checkpoint else blk(x)  # warn: checkpoint is slow import
         return x if self.downsample is None else self.downsample(x)
 
     def extra_repr(self) -> str:
-        """Returns a string with the layer's parameters for printing."""
+        """Return a string with the layer's parameters for printing."""
         return f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
 
 
@@ -763,12 +771,13 @@ class TinyViT(nn.Module):
     TinyViT: A compact vision transformer architecture for efficient image classification and feature extraction.
 
     This class implements the TinyViT model, which combines elements of vision transformers and convolutional
-    neural networks for improved efficiency and performance on vision tasks.
+    neural networks for improved efficiency and performance on vision tasks. It features hierarchical processing
+    with patch embedding, multiple stages of attention and convolution blocks, and a feature refinement neck.
 
     Attributes:
         img_size (int): Input image size.
         num_classes (int): Number of classification classes.
-        depths (List[int]): Number of blocks in each stage.
+        depths (Tuple[int, int, int, int]): Number of blocks in each stage.
         num_layers (int): Total number of layers in the network.
         mlp_ratio (float): Ratio of MLP hidden dimension to embedding dimension.
         patch_embed (PatchEmbed): Module for patch embedding.
@@ -778,66 +787,52 @@ class TinyViT(nn.Module):
         head (nn.Linear): Linear layer for final classification.
         neck (nn.Sequential): Neck module for feature refinement.
 
-    Methods:
-        set_layer_lr_decay: Sets layer-wise learning rate decay.
-        _init_weights: Initializes weights for linear and normalization layers.
-        no_weight_decay_keywords: Returns keywords for parameters that should not use weight decay.
-        forward_features: Processes input through the feature extraction layers.
-        forward: Performs a forward pass through the entire network.
-
     Examples:
         >>> model = TinyViT(img_size=224, num_classes=1000)
         >>> x = torch.randn(1, 3, 224, 224)
         >>> features = model.forward_features(x)
         >>> print(features.shape)
-        torch.Size([1, 256, 64, 64])
+        torch.Size([1, 256, 56, 56])
     """
 
     def __init__(
         self,
-        img_size=224,
-        in_chans=3,
-        num_classes=1000,
-        embed_dims=(96, 192, 384, 768),
-        depths=(2, 2, 6, 2),
-        num_heads=(3, 6, 12, 24),
-        window_sizes=(7, 7, 14, 7),
-        mlp_ratio=4.0,
-        drop_rate=0.0,
-        drop_path_rate=0.1,
-        use_checkpoint=False,
-        mbconv_expand_ratio=4.0,
-        local_conv_size=3,
-        layer_lr_decay=1.0,
+        img_size: int = 224,
+        in_chans: int = 3,
+        num_classes: int = 1000,
+        embed_dims: Tuple[int, int, int, int] = (96, 192, 384, 768),
+        depths: Tuple[int, int, int, int] = (2, 2, 6, 2),
+        num_heads: Tuple[int, int, int, int] = (3, 6, 12, 24),
+        window_sizes: Tuple[int, int, int, int] = (7, 7, 14, 7),
+        mlp_ratio: float = 4.0,
+        drop_rate: float = 0.0,
+        drop_path_rate: float = 0.1,
+        use_checkpoint: bool = False,
+        mbconv_expand_ratio: float = 4.0,
+        local_conv_size: int = 3,
+        layer_lr_decay: float = 1.0,
     ):
         """
-        Initializes the TinyViT model.
+        Initialize the TinyViT model.
 
         This constructor sets up the TinyViT architecture, including patch embedding, multiple layers of
         attention and convolution blocks, and a classification head.
 
         Args:
-            img_size (int): Size of the input image.
-            in_chans (int): Number of input channels.
-            num_classes (int): Number of classes for classification.
-            embed_dims (Tuple[int, int, int, int]): Embedding dimensions for each stage.
-            depths (Tuple[int, int, int, int]): Number of blocks in each stage.
-            num_heads (Tuple[int, int, int, int]): Number of attention heads in each stage.
-            window_sizes (Tuple[int, int, int, int]): Window sizes for each stage.
-            mlp_ratio (float): Ratio of MLP hidden dim to embedding dim.
-            drop_rate (float): Dropout rate.
-            drop_path_rate (float): Stochastic depth rate.
-            use_checkpoint (bool): Whether to use checkpointing to save memory.
-            mbconv_expand_ratio (float): Expansion ratio for MBConv layer.
-            local_conv_size (int): Kernel size for local convolutions.
-            layer_lr_decay (float): Layer-wise learning rate decay factor.
-
-        Examples:
-            >>> model = TinyViT(img_size=224, num_classes=1000)
-            >>> x = torch.randn(1, 3, 224, 224)
-            >>> output = model(x)
-            >>> print(output.shape)
-            torch.Size([1, 1000])
+            img_size (int, optional): Size of the input image.
+            in_chans (int, optional): Number of input channels.
+            num_classes (int, optional): Number of classes for classification.
+            embed_dims (Tuple[int, int, int, int], optional): Embedding dimensions for each stage.
+            depths (Tuple[int, int, int, int], optional): Number of blocks in each stage.
+            num_heads (Tuple[int, int, int, int], optional): Number of attention heads in each stage.
+            window_sizes (Tuple[int, int, int, int], optional): Window sizes for each stage.
+            mlp_ratio (float, optional): Ratio of MLP hidden dim to embedding dim.
+            drop_rate (float, optional): Dropout rate.
+            drop_path_rate (float, optional): Stochastic depth rate.
+            use_checkpoint (bool, optional): Whether to use checkpointing to save memory.
+            mbconv_expand_ratio (float, optional): Expansion ratio for MBConv layer.
+            local_conv_size (int, optional): Kernel size for local convolutions.
+            layer_lr_decay (float, optional): Layer-wise learning rate decay factor.
         """
         super().__init__()
         self.img_size = img_size
@@ -914,8 +909,8 @@ class TinyViT(nn.Module):
             LayerNorm2d(256),
         )
 
-    def set_layer_lr_decay(self, layer_lr_decay):
-        """Sets layer-wise learning rate decay for the TinyViT model based on depth."""
+    def set_layer_lr_decay(self, layer_lr_decay: float):
+        """Set layer-wise learning rate decay for the TinyViT model based on depth."""
         decay_rate = layer_lr_decay
 
         # Layers -> blocks (depth)
@@ -923,7 +918,7 @@ class TinyViT(nn.Module):
         lr_scales = [decay_rate ** (depth - i - 1) for i in range(depth)]
 
         def _set_lr_scale(m, scale):
-            """Sets the learning rate scale for each layer in the model based on the layer's depth."""
+            """Set the learning rate scale for each layer in the model based on the layer's depth."""
             for p in m.parameters():
                 p.lr_scale = scale
 
@@ -943,7 +938,7 @@ class TinyViT(nn.Module):
             p.param_name = k
 
         def _check_lr_scale(m):
-            """Checks if the learning rate scale attribute is present in module's parameters."""
+            """Check if the learning rate scale attribute is present in module's parameters."""
             for p in m.parameters():
                 assert hasattr(p, "lr_scale"), p.param_name
 
@@ -951,7 +946,7 @@ class TinyViT(nn.Module):
 
     @staticmethod
     def _init_weights(m):
-        """Initializes weights for linear and normalization layers in the TinyViT model."""
+        """Initialize weights for linear and normalization layers in the TinyViT model."""
         if isinstance(m, nn.Linear):
             # NOTE: This initialization is needed only for training.
             # trunc_normal_(m.weight, std=.02)
@@ -963,11 +958,11 @@ class TinyViT(nn.Module):
 
     @torch.jit.ignore
     def no_weight_decay_keywords(self):
-        """Returns a set of keywords for parameters that should not use weight decay."""
+        """Return a set of keywords for parameters that should not use weight decay."""
         return {"attention_biases"}
 
-    def forward_features(self, x):
-        """Processes input through feature extraction layers, returning spatial features."""
+    def forward_features(self, x: torch.Tensor) -> torch.Tensor:
+        """Process input through feature extraction layers, returning spatial features."""
         x = self.patch_embed(x)  # x input is (N, C, H, W)
 
         x = self.layers[0](x)
@@ -981,11 +976,11 @@ class TinyViT(nn.Module):
         x = x.permute(0, 3, 1, 2)
         return self.neck(x)
 
-    def forward(self, x):
-        """Performs the forward pass through the TinyViT model, extracting features from the input image."""
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Perform the forward pass through the TinyViT model, extracting features from the input image."""
         return self.forward_features(x)
 
-    def set_imgsz(self, imgsz=[1024, 1024]):
+    def set_imgsz(self, imgsz: List[int] = [1024, 1024]):
         """Set image size to make model compatible with different image sizes."""
         imgsz = [s // 4 for s in imgsz]
         self.patches_resolution = imgsz