|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| import itertools
|
| from typing import Tuple
|
|
|
| import torch
|
| import torch.nn as nn
|
| import torch.nn.functional as F
|
| import torch.utils.checkpoint as checkpoint
|
|
|
| from ultralytics.nn.modules import LayerNorm2d
|
| from ultralytics.utils.instance import to_2tuple
|
|
|
|
|
| class Conv2d_BN(torch.nn.Sequential):
|
| """
|
| A sequential container that performs 2D convolution followed by batch normalization.
|
|
|
| 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)
|
| """
|
|
|
| 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."""
|
| super().__init__()
|
| self.add_module("c", torch.nn.Conv2d(a, b, ks, stride, pad, dilation, groups, bias=False))
|
| bn = torch.nn.BatchNorm2d(b)
|
| torch.nn.init.constant_(bn.weight, bn_weight_init)
|
| torch.nn.init.constant_(bn.bias, 0)
|
| self.add_module("bn", bn)
|
|
|
|
|
| class PatchEmbed(nn.Module):
|
| """
|
| Embeds images into patches and projects them into a specified embedding dimension.
|
|
|
| Attributes:
|
| patches_resolution (Tuple[int, int]): Resolution of the patches after embedding.
|
| num_patches (int): Total number of patches.
|
| in_chans (int): Number of input channels.
|
| 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)
|
| """
|
|
|
| def __init__(self, in_chans, embed_dim, resolution, activation):
|
| """Initializes patch embedding with convolutional layers for image-to-patch conversion and projection."""
|
| super().__init__()
|
| img_size: Tuple[int, int] = to_2tuple(resolution)
|
| self.patches_resolution = (img_size[0] // 4, img_size[1] // 4)
|
| self.num_patches = self.patches_resolution[0] * self.patches_resolution[1]
|
| self.in_chans = in_chans
|
| self.embed_dim = embed_dim
|
| n = embed_dim
|
| self.seq = nn.Sequential(
|
| Conv2d_BN(in_chans, n // 2, 3, 2, 1),
|
| activation(),
|
| Conv2d_BN(n // 2, n, 3, 2, 1),
|
| )
|
|
|
| def forward(self, x):
|
| """Processes input tensor through patch embedding sequence, converting images to patch embeddings."""
|
| return self.seq(x)
|
|
|
|
|
| class MBConv(nn.Module):
|
| """
|
| Mobile Inverted Bottleneck Conv (MBConv) layer, part of the EfficientNet architecture.
|
|
|
| Attributes:
|
| in_chans (int): Number of input channels.
|
| hidden_chans (int): Number of hidden channels.
|
| out_chans (int): Number of output channels.
|
| conv1 (Conv2d_BN): First convolutional layer.
|
| act1 (nn.Module): First activation function.
|
| conv2 (Conv2d_BN): Depthwise convolutional layer.
|
| act2 (nn.Module): Second activation function.
|
| conv3 (Conv2d_BN): Final convolutional layer.
|
| 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)
|
| >>> x = torch.randn(1, in_chans, 56, 56)
|
| >>> output = mbconv(x)
|
| >>> print(output.shape)
|
| 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."""
|
| super().__init__()
|
| self.in_chans = in_chans
|
| self.hidden_chans = int(in_chans * expand_ratio)
|
| self.out_chans = out_chans
|
|
|
| self.conv1 = Conv2d_BN(in_chans, self.hidden_chans, ks=1)
|
| self.act1 = activation()
|
|
|
| self.conv2 = Conv2d_BN(self.hidden_chans, self.hidden_chans, ks=3, stride=1, pad=1, groups=self.hidden_chans)
|
| self.act2 = activation()
|
|
|
| self.conv3 = Conv2d_BN(self.hidden_chans, out_chans, ks=1, bn_weight_init=0.0)
|
| self.act3 = activation()
|
|
|
|
|
|
|
| self.drop_path = nn.Identity()
|
|
|
| def forward(self, x):
|
| """Implements the forward pass of MBConv, applying convolutions and skip connection."""
|
| shortcut = x
|
| x = self.conv1(x)
|
| x = self.act1(x)
|
| x = self.conv2(x)
|
| x = self.act2(x)
|
| x = self.conv3(x)
|
| x = self.drop_path(x)
|
| x += shortcut
|
| return self.act3(x)
|
|
|
|
|
| class PatchMerging(nn.Module):
|
| """
|
| Merges neighboring patches in the feature map and projects 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.
|
|
|
| Attributes:
|
| 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.
|
| act (nn.Module): The activation function used between convolutions.
|
| conv1 (Conv2d_BN): The first convolutional layer for dimension projection.
|
| 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)
|
| """
|
|
|
| def __init__(self, input_resolution, dim, out_dim, activation):
|
| """Initializes the PatchMerging module for merging and projecting neighboring patches in feature maps."""
|
| super().__init__()
|
|
|
| self.input_resolution = input_resolution
|
| self.dim = dim
|
| self.out_dim = out_dim
|
| self.act = activation()
|
| self.conv1 = Conv2d_BN(dim, out_dim, 1, 1, 0)
|
| stride_c = 1 if out_dim in {320, 448, 576} else 2
|
| 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."""
|
| if x.ndim == 3:
|
| H, W = self.input_resolution
|
| B = len(x)
|
|
|
| x = x.view(B, H, W, -1).permute(0, 3, 1, 2)
|
|
|
| x = self.conv1(x)
|
| x = self.act(x)
|
|
|
| x = self.conv2(x)
|
| x = self.act(x)
|
| x = self.conv3(x)
|
| return x.flatten(2).transpose(1, 2)
|
|
|
|
|
| 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.
|
|
|
| Attributes:
|
| dim (int): Dimensionality of the input and output.
|
| input_resolution (Tuple[int, int]): Resolution of the input image.
|
| 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.
|
|
|
| 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)
|
| """
|
|
|
| def __init__(
|
| self,
|
| dim,
|
| input_resolution,
|
| depth,
|
| activation,
|
| drop_path=0.0,
|
| downsample=None,
|
| use_checkpoint=False,
|
| out_dim=None,
|
| conv_expand_ratio=4.0,
|
| ):
|
| """
|
| Initializes 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.
|
|
|
| Args:
|
| dim (int): The dimensionality of the input and output.
|
| 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)
|
| """
|
| super().__init__()
|
| self.dim = dim
|
| self.input_resolution = input_resolution
|
| self.depth = depth
|
| self.use_checkpoint = use_checkpoint
|
|
|
|
|
| self.blocks = nn.ModuleList(
|
| [
|
| MBConv(
|
| dim,
|
| dim,
|
| conv_expand_ratio,
|
| activation,
|
| drop_path[i] if isinstance(drop_path, list) else drop_path,
|
| )
|
| for i in range(depth)
|
| ]
|
| )
|
|
|
|
|
| self.downsample = (
|
| None
|
| if downsample is None
|
| 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."""
|
| for blk in self.blocks:
|
| x = checkpoint.checkpoint(blk, x) if self.use_checkpoint else blk(x)
|
| return x if self.downsample is None else self.downsample(x)
|
|
|
|
|
| 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.
|
|
|
| Attributes:
|
| norm (nn.LayerNorm): Layer normalization applied to the input.
|
| fc1 (nn.Linear): First fully-connected layer.
|
| fc2 (nn.Linear): Second fully-connected layer.
|
| 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)
|
| >>> 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."""
|
| 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.drop = nn.Dropout(drop)
|
|
|
| def forward(self, x):
|
| """Applies 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)
|
| x = self.drop(x)
|
| x = self.fc2(x)
|
| return self.drop(x)
|
|
|
|
|
| class Attention(torch.nn.Module):
|
| """
|
| Multi-head attention module with spatial awareness and trainable attention biases.
|
|
|
| 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
|
| offset between spatial positions in the resolution grid.
|
|
|
| Attributes:
|
| num_heads (int): Number of attention heads.
|
| scale (float): Scaling factor for attention scores.
|
| key_dim (int): Dimensionality of the keys and queries.
|
| nh_kd (int): Product of num_heads and key_dim.
|
| d (int): Dimensionality of the value vectors.
|
| dh (int): Product of d and num_heads.
|
| attn_ratio (float): Attention ratio affecting the dimensions of the value vectors.
|
| norm (nn.LayerNorm): Layer normalization applied to input.
|
| 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.
|
|
|
| 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])
|
| """
|
|
|
| def __init__(
|
| self,
|
| dim,
|
| key_dim,
|
| num_heads=8,
|
| attn_ratio=4,
|
| resolution=(14, 14),
|
| ):
|
| """
|
| Initializes 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
|
| offset between spatial positions in the resolution grid.
|
|
|
| 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])
|
| """
|
| super().__init__()
|
|
|
| assert isinstance(resolution, tuple) and len(resolution) == 2, "'resolution' argument not tuple of length 2"
|
| self.num_heads = num_heads
|
| self.scale = key_dim**-0.5
|
| self.key_dim = key_dim
|
| self.nh_kd = nh_kd = key_dim * num_heads
|
| self.d = int(attn_ratio * key_dim)
|
| self.dh = int(attn_ratio * key_dim) * num_heads
|
| self.attn_ratio = attn_ratio
|
| h = self.dh + nh_kd * 2
|
|
|
| self.norm = nn.LayerNorm(dim)
|
| self.qkv = nn.Linear(dim, h)
|
| self.proj = nn.Linear(self.dh, dim)
|
|
|
| points = list(itertools.product(range(resolution[0]), range(resolution[1])))
|
| N = len(points)
|
| attention_offsets = {}
|
| idxs = []
|
| for p1 in points:
|
| for p2 in points:
|
| offset = (abs(p1[0] - p2[0]), abs(p1[1] - p2[1]))
|
| if offset not in attention_offsets:
|
| attention_offsets[offset] = len(attention_offsets)
|
| idxs.append(attention_offsets[offset])
|
| self.attention_biases = torch.nn.Parameter(torch.zeros(num_heads, len(attention_offsets)))
|
| 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."""
|
| 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):
|
| """Applies multi-head attention with spatial awareness and trainable attention biases."""
|
| B, N, _ = x.shape
|
|
|
|
|
| x = self.norm(x)
|
|
|
| qkv = self.qkv(x)
|
|
|
| q, k, v = qkv.view(B, N, self.num_heads, -1).split([self.key_dim, self.key_dim, self.d], dim=3)
|
|
|
| q = q.permute(0, 2, 1, 3)
|
| k = k.permute(0, 2, 1, 3)
|
| v = v.permute(0, 2, 1, 3)
|
| self.ab = self.ab.to(self.attention_biases.device)
|
|
|
| attn = (q @ k.transpose(-2, -1)) * self.scale + (
|
| self.attention_biases[:, self.attention_bias_idxs] if self.training else self.ab
|
| )
|
| attn = attn.softmax(dim=-1)
|
| x = (attn @ v).transpose(1, 2).reshape(B, N, self.dh)
|
| return self.proj(x)
|
|
|
|
|
| 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.
|
|
|
| Attributes:
|
| dim (int): The dimensionality of the input and output.
|
| input_resolution (Tuple[int, int]): Spatial resolution of the input feature map.
|
| num_heads (int): Number of attention heads.
|
| window_size (int): Size of the attention window.
|
| 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.
|
| 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)
|
| >>> output = block(input_tensor)
|
| >>> print(output.shape)
|
| torch.Size([1, 196, 192])
|
| """
|
|
|
| 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,
|
| activation=nn.GELU,
|
| ):
|
| """
|
| Initializes 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.
|
|
|
| Args:
|
| 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])
|
| """
|
| super().__init__()
|
| self.dim = dim
|
| self.input_resolution = input_resolution
|
| self.num_heads = num_heads
|
| assert window_size > 0, "window_size must be greater than 0"
|
| self.window_size = window_size
|
| self.mlp_ratio = mlp_ratio
|
|
|
|
|
|
|
| self.drop_path = nn.Identity()
|
|
|
| assert dim % num_heads == 0, "dim must be divisible by num_heads"
|
| head_dim = dim // num_heads
|
|
|
| window_resolution = (window_size, window_size)
|
| self.attn = Attention(dim, head_dim, num_heads, attn_ratio=1, resolution=window_resolution)
|
|
|
| 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)
|
|
|
| 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."""
|
| h, w = self.input_resolution
|
| b, hw, c = x.shape
|
| assert hw == h * w, "input feature has wrong size"
|
| res_x = x
|
| if h == self.window_size and w == self.window_size:
|
| x = self.attn(x)
|
| else:
|
| x = x.view(b, h, w, c)
|
| pad_b = (self.window_size - h % self.window_size) % self.window_size
|
| pad_r = (self.window_size - w % self.window_size) % self.window_size
|
| padding = pad_b > 0 or pad_r > 0
|
| if padding:
|
| x = F.pad(x, (0, 0, 0, pad_r, 0, pad_b))
|
|
|
| pH, pW = h + pad_b, w + pad_r
|
| nH = pH // self.window_size
|
| nW = pW // self.window_size
|
|
|
|
|
| x = (
|
| x.view(b, nH, self.window_size, nW, self.window_size, c)
|
| .transpose(2, 3)
|
| .reshape(b * nH * nW, self.window_size * self.window_size, c)
|
| )
|
| x = self.attn(x)
|
|
|
|
|
| x = x.view(b, nH, nW, self.window_size, self.window_size, c).transpose(2, 3).reshape(b, pH, pW, c)
|
| if padding:
|
| x = x[:, :h, :w].contiguous()
|
|
|
| x = x.view(b, hw, c)
|
|
|
| x = res_x + self.drop_path(x)
|
| x = x.transpose(1, 2).reshape(b, c, h, w)
|
| x = self.local_conv(x)
|
| x = x.view(b, c, hw).transpose(1, 2)
|
|
|
| return x + self.drop_path(self.mlp(x))
|
|
|
| def extra_repr(self) -> str:
|
| """
|
| Returns 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.
|
|
|
| Returns:
|
| (str): A formatted string containing the block's parameters.
|
|
|
| Examples:
|
| >>> block = TinyViTBlock(dim=192, input_resolution=(14, 14), num_heads=3, window_size=7, mlp_ratio=4.0)
|
| >>> print(block.extra_repr())
|
| dim=192, input_resolution=(14, 14), num_heads=3, window_size=7, mlp_ratio=4.0
|
| """
|
| return (
|
| f"dim={self.dim}, input_resolution={self.input_resolution}, num_heads={self.num_heads}, "
|
| f"window_size={self.window_size}, mlp_ratio={self.mlp_ratio}"
|
| )
|
|
|
|
|
| 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.
|
|
|
| Attributes:
|
| dim (int): The dimensionality of the input and output features.
|
| input_resolution (Tuple[int, int]): Spatial resolution of the input feature map.
|
| depth (int): Number of TinyViT blocks in this layer.
|
| use_checkpoint (bool): Whether to use gradient checkpointing to save memory.
|
| 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)
|
| >>> output = layer(input_tensor)
|
| >>> print(output.shape)
|
| torch.Size([1, 784, 384])
|
| """
|
|
|
| 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,
|
| activation=nn.GELU,
|
| out_dim=None,
|
| ):
|
| """
|
| Initializes 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.
|
|
|
| Args:
|
| dim (int): Dimensionality of the input and output features.
|
| input_resolution (Tuple[int, int]): Spatial resolution of the input feature map (height, width).
|
| 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.
|
| 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)
|
| """
|
| super().__init__()
|
| self.dim = dim
|
| self.input_resolution = input_resolution
|
| self.depth = depth
|
| self.use_checkpoint = use_checkpoint
|
|
|
|
|
| self.blocks = nn.ModuleList(
|
| [
|
| TinyViTBlock(
|
| dim=dim,
|
| input_resolution=input_resolution,
|
| num_heads=num_heads,
|
| window_size=window_size,
|
| mlp_ratio=mlp_ratio,
|
| drop=drop,
|
| drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,
|
| local_conv_size=local_conv_size,
|
| activation=activation,
|
| )
|
| for i in range(depth)
|
| ]
|
| )
|
|
|
|
|
| self.downsample = (
|
| None
|
| if downsample is None
|
| else downsample(input_resolution, dim=dim, out_dim=out_dim, activation=activation)
|
| )
|
|
|
| def forward(self, x):
|
| """Processes input through TinyViT blocks and optional downsampling."""
|
| for blk in self.blocks:
|
| x = checkpoint.checkpoint(blk, x) if self.use_checkpoint else blk(x)
|
| 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 f"dim={self.dim}, input_resolution={self.input_resolution}, depth={self.depth}"
|
|
|
|
|
| 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.
|
|
|
| Attributes:
|
| img_size (int): Input image size.
|
| num_classes (int): Number of classification classes.
|
| depths (List[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.
|
| patches_resolution (Tuple[int, int]): Resolution of embedded patches.
|
| layers (nn.ModuleList): List of network layers.
|
| norm_head (nn.LayerNorm): Layer normalization for the classifier head.
|
| 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])
|
| """
|
|
|
| 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,
|
| ):
|
| """
|
| Initializes 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])
|
| """
|
| super().__init__()
|
| self.img_size = img_size
|
| self.num_classes = num_classes
|
| self.depths = depths
|
| self.num_layers = len(depths)
|
| self.mlp_ratio = mlp_ratio
|
|
|
| activation = nn.GELU
|
|
|
| self.patch_embed = PatchEmbed(
|
| in_chans=in_chans, embed_dim=embed_dims[0], resolution=img_size, activation=activation
|
| )
|
|
|
| patches_resolution = self.patch_embed.patches_resolution
|
| self.patches_resolution = patches_resolution
|
|
|
|
|
| dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]
|
|
|
|
|
| self.layers = nn.ModuleList()
|
| for i_layer in range(self.num_layers):
|
| kwargs = dict(
|
| dim=embed_dims[i_layer],
|
| input_resolution=(
|
| patches_resolution[0] // (2 ** (i_layer - 1 if i_layer == 3 else i_layer)),
|
| patches_resolution[1] // (2 ** (i_layer - 1 if i_layer == 3 else i_layer)),
|
| ),
|
|
|
|
|
| depth=depths[i_layer],
|
| drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])],
|
| downsample=PatchMerging if (i_layer < self.num_layers - 1) else None,
|
| use_checkpoint=use_checkpoint,
|
| out_dim=embed_dims[min(i_layer + 1, len(embed_dims) - 1)],
|
| activation=activation,
|
| )
|
| if i_layer == 0:
|
| layer = ConvLayer(conv_expand_ratio=mbconv_expand_ratio, **kwargs)
|
| else:
|
| layer = BasicLayer(
|
| num_heads=num_heads[i_layer],
|
| window_size=window_sizes[i_layer],
|
| mlp_ratio=self.mlp_ratio,
|
| drop=drop_rate,
|
| local_conv_size=local_conv_size,
|
| **kwargs,
|
| )
|
| self.layers.append(layer)
|
|
|
|
|
| self.norm_head = nn.LayerNorm(embed_dims[-1])
|
| self.head = nn.Linear(embed_dims[-1], num_classes) if num_classes > 0 else torch.nn.Identity()
|
|
|
|
|
| self.apply(self._init_weights)
|
| self.set_layer_lr_decay(layer_lr_decay)
|
| self.neck = nn.Sequential(
|
| nn.Conv2d(
|
| embed_dims[-1],
|
| 256,
|
| kernel_size=1,
|
| bias=False,
|
| ),
|
| LayerNorm2d(256),
|
| nn.Conv2d(
|
| 256,
|
| 256,
|
| kernel_size=3,
|
| padding=1,
|
| bias=False,
|
| ),
|
| LayerNorm2d(256),
|
| )
|
|
|
| def set_layer_lr_decay(self, layer_lr_decay):
|
| """Sets layer-wise learning rate decay for the TinyViT model based on depth."""
|
| decay_rate = layer_lr_decay
|
|
|
|
|
| depth = sum(self.depths)
|
| 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."""
|
| for p in m.parameters():
|
| p.lr_scale = scale
|
|
|
| self.patch_embed.apply(lambda x: _set_lr_scale(x, lr_scales[0]))
|
| i = 0
|
| for layer in self.layers:
|
| for block in layer.blocks:
|
| block.apply(lambda x: _set_lr_scale(x, lr_scales[i]))
|
| i += 1
|
| if layer.downsample is not None:
|
| layer.downsample.apply(lambda x: _set_lr_scale(x, lr_scales[i - 1]))
|
| assert i == depth
|
| for m in [self.norm_head, self.head]:
|
| m.apply(lambda x: _set_lr_scale(x, lr_scales[-1]))
|
|
|
| for k, p in self.named_parameters():
|
| p.param_name = k
|
|
|
| def _check_lr_scale(m):
|
| """Checks 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
|
|
|
| self.apply(_check_lr_scale)
|
|
|
| @staticmethod
|
| def _init_weights(m):
|
| """Initializes weights for linear and normalization layers in the TinyViT model."""
|
| if isinstance(m, nn.Linear):
|
|
|
|
|
| if m.bias is not None:
|
| nn.init.constant_(m.bias, 0)
|
| elif isinstance(m, nn.LayerNorm):
|
| nn.init.constant_(m.bias, 0)
|
| nn.init.constant_(m.weight, 1.0)
|
|
|
| @torch.jit.ignore
|
| def no_weight_decay_keywords(self):
|
| """Returns 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."""
|
| x = self.patch_embed(x)
|
|
|
| x = self.layers[0](x)
|
| start_i = 1
|
|
|
| for i in range(start_i, len(self.layers)):
|
| layer = self.layers[i]
|
| x = layer(x)
|
| batch, _, channel = x.shape
|
| x = x.view(batch, self.patches_resolution[0] // 4, self.patches_resolution[1] // 4, channel)
|
| 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."""
|
| return self.forward_features(x)
|
|
|
| def set_imgsz(self, imgsz=[1024, 1024]):
|
| """Set image size to make model compatible with different image sizes."""
|
| imgsz = [s // 4 for s in imgsz]
|
| self.patches_resolution = imgsz
|
| for i, layer in enumerate(self.layers):
|
| input_resolution = (
|
| imgsz[0] // (2 ** (i - 1 if i == 3 else i)),
|
| imgsz[1] // (2 ** (i - 1 if i == 3 else i)),
|
| )
|
| layer.input_resolution = input_resolution
|
| if layer.downsample is not None:
|
| layer.downsample.input_resolution = input_resolution
|
| if isinstance(layer, BasicLayer):
|
| for b in layer.blocks:
|
| b.input_resolution = input_resolution
|
|
|
