Đóng góp của bài báo

Các mô hình ViT nếu như không pretrained trước đó sẽ có hiệu suất kém hơn so với các mô hình ConvNets. Lý do là các mô hình Transformer có model capacity cao với inductive bias thấp, điều này dẫn tới tình trạng overfitting. Bài toán đặt ra là làm như thế nào để kết hợp hiệu quả tương tác global và local trong mô hình Transformer với mục tiêu cân bằng model capacity và tính tổng quát hóa trong một lượng tài nguyên tính toán hợp lý.

Bài báo đề xuất một loại Transformer module mới có tên là multi-axis self-attention (Max-SA) có khả năng thực hiện đồng thời cả tương tác local và global (local and global spatial interactions) trong một block. Nếu so sánh với self-attention đầy đủ thì Max-SA linh hoạt và hiệu quả hơn (ví dụ như có thể áp dụng cho các input có độ dài khác nhau với độ phức tạp tuyến tính). Trái ngược với window attention hay local attention, Max-SA cho phép model capacity mạnh hơn bằng cách đề xuất một global receptive field. Ngoài ra, với độ phức tạp tuyến tính, Max-SA có thể được sử dụng như một module attention tổng quát tại bất kì layer nào trong mô hình, kể cả những layer ở stage đầu hay input với high resolution.

Bên cạnh đó, nhóm tác giả cũng thiết kế một backbone đơn giản mà hiệu quả có tên là Multi-axis Vision Transformer (MaxViT) bằng cách stack các block mà thành phần là Max-SA và các lớp convolution. Mô hình MaxViT cải thiện kết quả SOTA ở nhiều task liên quan tới hình ảnh. Tương quan giữa độ chính xác với FLOPs và lượng tham số giữa các mô hình được thể hiện trong 2 biểu đồ dưới:

Phương pháp

Attention

Thành phần attention được sử dụng trong mô hình là Relative self-attention mà không phải self-attention tiêu chuẩn (được giới thiệu trong bài báo của mô hình ViT). Lý do là self-attention tiêu chuẩn không có thông tin về vị trí (location), chính vì vậy mà chúng không có tính chất translation equivariant, đây là một inductive bias quan trọng xuất hiện trong các mô hình ConvNet. Relative attention mà mô hình sử dụng có công thức như sau:

trong đó $Q, K, V \in \mathbb{R}^{(H \times W) \times C}$ là ma trận query, key và value. $d$ là hidden dimension. Các trọng số attention được quyết định bởi một ma trận learnable location-aware $B$ có chứa thông tin về vị trí và input-adaptive attention $Q K^T / \sqrt{d}$ . Relative attention có các ưu điểm là input-adaptivity, translation equivariance và global interaction.

Multi-axis Attention

Global interaction là một ưu điểm vượt trội của self-attention so với local convolution. Tuy nhiên, nếu áp dụng trực tiếp attention sẽ không khả thi do yêu cầu độ phức tạp bậc 2 với spatial size. Để giải quyết vấn đề này, nhóm tác giả đề xuất một cách tiếp cận multi-axis, mục tiêu là tách full-size attention thành 2 thành phần local và global bằng cách tách thành các chiều spatial khác nhau. Gọi $X \in \mathbb{R}^{H \times W \times C}$ là input feature map, thay vì sử dụng attention cho flattened spatial dimension $HW$ , nhóm tác giả tạo một tensor có shape là $\left(\frac{H}{P} \times \frac{W}{P}, P \times P, C\right)$ . Tensor này đại diện cho các window không chồng nhau có kích thước là $P \times P$ . Việc sử dụng self-attention cho các cửa số kích thước $P \times P$ chính là tương đương với việc ta thực hiện local interaction, nhóm tác giả gọi đây là block attention.

Tuy khắc phục được vấn đề về tài nguyên tính toán của full self-attention nhưng local-attention lại gặp vấn đề là underfit trên các bộ dataset lớn. Đề giải quyết vấn đề này, nhóm tác giả đề xuất một ý tưởng có tên là grid attention. Thay vì chia feature map thành các window size cố định, nhóm tác giả thực hiện chia lưới (grid) tensor thành tensor có shape $(G \times G, \frac{H}{G} \times \frac{W}{G}, C)$ sử dụng $G \times G$ grid, kết quả là ta có các window có kích thước là $\frac{H}{G} \times \frac{W}{G}$ . Với việc sử dụng window size và grid size giống nhau ( $P = G = 7$ ), ta có thể cân bằng tài nguyên cho các phép tính local và global (cả 2 đều có độ phức tạp tuyến tính với spatial size hoặc sequence length).

Bằng việc stack 2 loại attention để có cả local và global interaction, ta có MaxViT block (như trong hình trên). Trong MaxViT block vẫn có các thành phần giống Transformer như Feedforward networks (FFNs) và skip-connections. Bên cạnh đó, nhóm tác giả cũng sử dụng MBConv block với squeeze-and-excitation (SE) module để thực hiện multi-axis attention. Lý do là nhóm tác giả nhận thấy rằng việc sử dụng MBConv cùng với attention làm tăng khả năng tổng quát hóa và khả năng huấn luyện của mô hình. Việc sử dụng layer MBConv trước attention mang lại nhiều ưu điểm, depthwise convolution có thể được coi là một conditional position encoding (CPE), do đó ta không cần phải dùng một positional encoding layer rõ ràng nữa (như trong mô hình transformer tiêu chuẩn). Một điểm hay nữa là multi-axis attention có thể sử dụng cùng hoặc đứng riêng cho nhiều mục đích khác nhau: block attention cho local interaction và grid attention cho global mixing. Các thành phần này có thể tích hợp vào trong nhiều kiến trúc khác nhau, đặc biệt là trong các task high resolution để tận dụng ưu điểm là global interaction mà không sử dụng nhiều tài nguyên tính toán.

Các biến thể của kiến trúc

Kiến trúc tổng thể của MaxViT được thể hiện trong hình trên. Nhóm tác giả thiết kế một backbone có kiến trúc hierarchical. Input ban đầu được downsample sử dụng các layer Conv3x3 trong stem stage (S0). Phần thân của mô hình bao gồm 4 stage (S1-S4) với mỗi stage chứa một nửa resolution của stage trước đó và lượng channel (hidden dimension) được tăng gấp đôi. MaxViT block được sử dụng xuyên suốt trong mô hình. Tỉ lệ expansion và shrink cho inverted bottleneck và squeeze-excitation (SE) mặc định lần lượt là 4 và 0.25. Nhóm tác giả đặt attention head size là 32 cho tất cả attention block. Để thực hiện scale model, nhóm tác giả tăng số lượng block trên mỗi stage $B$ và channel dimension $C$ . Tổng hợp cấu hình của các biến thể MaxViT được thể hiện trong bảng dưới:

Coding

Code mô hình MaxViT tham khảo:

""" MaxViT A PyTorch implementation of the paper: `MaxViT: Multi-Axis Vision Transformer` - MaxViT: Multi-Axis Vision Transformer
"""
from typing import Type, Callable, Tuple, Optional, Set, List, Union import torch
import torch.nn as nn from timm.models.efficientnet_blocks import SqueezeExcite, DepthwiseSeparableConv
from timm.models.layers import drop_path, trunc_normal_, Mlp, DropPath def _gelu_ignore_parameters( *args, **kwargs
) -> nn.Module: """ Bad trick to ignore the inplace=True argument in the DepthwiseSeparableConv of Timm. Args: *args: Ignored. **kwargs: Ignored. Returns: activation (nn.Module): GELU activation function. """ activation = nn.GELU() return activation class MBConv(nn.Module): """ MBConv block as described in: https://arxiv.org/pdf/2204.01697.pdf. Without downsampling: x ← x + Proj(SE(DWConv(Conv(Norm(x))))) With downsampling: x ← Proj(Pool2D(x)) + Proj(SE(DWConv ↓(Conv(Norm(x))))). Conv is a 1 X 1 convolution followed by a Batch Normalization layer and a GELU activation. SE is the Squeeze-Excitation layer. Proj is the shrink 1 X 1 convolution. Note: This implementation differs slightly from the original MobileNet implementation! Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. downscale (bool, optional): If true downscale by a factor of two is performed. Default: False act_layer (Type[nn.Module], optional): Type of activation layer to be utilized. Default: nn.GELU norm_layer (Type[nn.Module], optional): Type of normalization layer to be utilized. Default: nn.BatchNorm2d drop_path (float, optional): Dropout rate to be applied during training. Default 0. """ def __init__( self, in_channels: int, out_channels: int, downscale: bool = False, act_layer: Type[nn.Module] = nn.GELU, norm_layer: Type[nn.Module] = nn.BatchNorm2d, drop_path: float = 0., ) -> None: """ Constructor method """ # Call super constructor super(MBConv, self).__init__() # Save parameter self.drop_path_rate: float = drop_path # Check parameters for downscaling if not downscale: assert in_channels == out_channels, "If downscaling is utilized input and output channels must be equal." # Ignore inplace parameter if GELU is used if act_layer == nn.GELU: act_layer = _gelu_ignore_parameters # Make main path self.main_path = nn.Sequential( norm_layer(in_channels), nn.Conv2d(in_channels=in_channels, out_channels=in_channels, kernel_size=(1, 1)), DepthwiseSeparableConv(in_chs=in_channels, out_chs=out_channels, stride=2 if downscale else 1, act_layer=act_layer, norm_layer=norm_layer, drop_path_rate=drop_path), SqueezeExcite(in_chs=out_channels, rd_ratio=0.25), nn.Conv2d(in_channels=out_channels, out_channels=out_channels, kernel_size=(1, 1)) ) # Make skip path self.skip_path = nn.Sequential( nn.MaxPool2d(kernel_size=(2, 2), stride=(2, 2)), nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=(1, 1)) ) if downscale else nn.Identity() def forward( self, input: torch.Tensor ) -> torch.Tensor: """ Forward pass. Args: input (torch.Tensor): Input tensor of the shape [B, C_in, H, W]. Returns: output (torch.Tensor): Output tensor of the shape [B, C_out, H (// 2), W (// 2)] (downscaling is optional). """ output = self.main_path(input) if self.drop_path_rate > 0.: output = drop_path(output, self.drop_path_rate, self.training) output = output + self.skip_path(input) return output def window_partition( input: torch.Tensor, window_size: Tuple[int, int] = (7, 7)
) -> torch.Tensor: """ Window partition function. Args: input (torch.Tensor): Input tensor of the shape [B, C, H, W]. window_size (Tuple[int, int], optional): Window size to be applied. Default (7, 7) Returns: windows (torch.Tensor): Unfolded input tensor of the shape [B * windows, window_size[0], window_size[1], C]. """ # Get size of input B, C, H, W = input.shape # Unfold input windows = input.view(B, C, H // window_size[0], window_size[0], W // window_size[1], window_size[1]) # Permute and reshape to [B * windows, window_size[0], window_size[1], channels] windows = windows.permute(0, 2, 4, 3, 5, 1).contiguous().view(-1, window_size[0], window_size[1], C) return windows def window_reverse( windows: torch.Tensor, original_size: Tuple[int, int], window_size: Tuple[int, int] = (7, 7)
) -> torch.Tensor: """ Reverses the window partition. Args: windows (torch.Tensor): Window tensor of the shape [B * windows, window_size[0], window_size[1], C]. original_size (Tuple[int, int]): Original shape. window_size (Tuple[int, int], optional): Window size which have been applied. Default (7, 7) Returns: output (torch.Tensor): Folded output tensor of the shape [B, C, original_size[0], original_size[1]]. """ # Get height and width H, W = original_size # Compute original batch size B = int(windows.shape[0] / (H * W / window_size[0] / window_size[1])) # Fold grid tensor output = windows.view(B, H // window_size[0], W // window_size[1], window_size[0], window_size[1], -1) output = output.permute(0, 5, 1, 3, 2, 4).contiguous().view(B, -1, H, W) return output def grid_partition( input: torch.Tensor, grid_size: Tuple[int, int] = (7, 7)
) -> torch.Tensor: """ Grid partition function. Args: input (torch.Tensor): Input tensor of the shape [B, C, H, W]. grid_size (Tuple[int, int], optional): Grid size to be applied. Default (7, 7) Returns: grid (torch.Tensor): Unfolded input tensor of the shape [B * grids, grid_size[0], grid_size[1], C]. """ # Get size of input B, C, H, W = input.shape # Unfold input grid = input.view(B, C, grid_size[0], H // grid_size[0], grid_size[1], W // grid_size[1]) # Permute and reshape [B * (H // grid_size[0]) * (W // grid_size[1]), grid_size[0], window_size[1], C] grid = grid.permute(0, 3, 5, 2, 4, 1).contiguous().view(-1, grid_size[0], grid_size[1], C) return grid def grid_reverse( grid: torch.Tensor, original_size: Tuple[int, int], grid_size: Tuple[int, int] = (7, 7)
) -> torch.Tensor: """ Reverses the grid partition. Args: Grid (torch.Tensor): Grid tensor of the shape [B * grids, grid_size[0], grid_size[1], C]. original_size (Tuple[int, int]): Original shape. grid_size (Tuple[int, int], optional): Grid size which have been applied. Default (7, 7) Returns: output (torch.Tensor): Folded output tensor of the shape [B, C, original_size[0], original_size[1]]. """ # Get height, width, and channels (H, W), C = original_size, grid.shape[-1] # Compute original batch size B = int(grid.shape[0] / (H * W / grid_size[0] / grid_size[1])) # Fold grid tensor output = grid.view(B, H // grid_size[0], W // grid_size[1], grid_size[0], grid_size[1], C) output = output.permute(0, 5, 3, 1, 4, 2).contiguous().view(B, C, H, W) return output def get_relative_position_index( win_h: int, win_w: int
) -> torch.Tensor: """ Function to generate pair-wise relative position index for each token inside the window. Taken from Timms Swin V1 implementation. Args: win_h (int): Window/Grid height. win_w (int): Window/Grid width. Returns: relative_coords (torch.Tensor): Pair-wise relative position indexes [height * width, height * width]. """ coords = torch.stack(torch.meshgrid([torch.arange(win_h), torch.arange(win_w)])) coords_flatten = torch.flatten(coords, 1) relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :] relative_coords = relative_coords.permute(1, 2, 0).contiguous() relative_coords[:, :, 0] += win_h - 1 relative_coords[:, :, 1] += win_w - 1 relative_coords[:, :, 0] *= 2 * win_w - 1 return relative_coords.sum(-1) class RelativeSelfAttention(nn.Module): """ Relative Self-Attention similar to Swin V1. Implementation inspired by Timms Swin V1 implementation. Args: in_channels (int): Number of input channels. num_heads (int, optional): Number of attention heads. Default 32 grid_window_size (Tuple[int, int], optional): Grid/Window size to be utilized. Default (7, 7) attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 drop (float, optional): Dropout ratio of output. Default: 0.0 """ def __init__( self, in_channels: int, num_heads: int = 32, grid_window_size: Tuple[int, int] = (7, 7), attn_drop: float = 0., drop: float = 0. ) -> None: """ Constructor method """ # Call super constructor super(RelativeSelfAttention, self).__init__() # Save parameters self.in_channels: int = in_channels self.num_heads: int = num_heads self.grid_window_size: Tuple[int, int] = grid_window_size self.scale: float = num_heads ** -0.5 self.attn_area: int = grid_window_size[0] * grid_window_size[1] # Init layers self.qkv_mapping = nn.Linear(in_features=in_channels, out_features=3 * in_channels, bias=True) self.attn_drop = nn.Dropout(p=attn_drop) self.proj = nn.Linear(in_features=in_channels, out_features=in_channels, bias=True) self.proj_drop = nn.Dropout(p=drop) self.softmax = nn.Softmax(dim=-1) # Define a parameter table of relative position bias, shape: 2*Wh-1 * 2*Ww-1, nH self.relative_position_bias_table = nn.Parameter( torch.zeros((2 * grid_window_size[0] - 1) * (2 * grid_window_size[1] - 1), num_heads)) # Get pair-wise relative position index for each token inside the window self.register_buffer("relative_position_index", get_relative_position_index(grid_window_size[0], grid_window_size[1])) # Init relative positional bias trunc_normal_(self.relative_position_bias_table, std=.02) def _get_relative_positional_bias( self ) -> torch.Tensor: """ Returns the relative positional bias. Returns: relative_position_bias (torch.Tensor): Relative positional bias. """ relative_position_bias = self.relative_position_bias_table[ self.relative_position_index.view(-1)].view(self.attn_area, self.attn_area, -1) relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous() return relative_position_bias.unsqueeze(0) def forward( self, input: torch.Tensor ) -> torch.Tensor: """ Forward pass. Args: input (torch.Tensor): Input tensor of the shape [B_, N, C]. Returns: output (torch.Tensor): Output tensor of the shape [B_, N, C]. """ # Get shape of input B_, N, C = input.shape # Perform query key value mapping qkv = self.qkv_mapping(input).reshape(B_, N, 3, self.num_heads, -1).permute(2, 0, 3, 1, 4) q, k, v = qkv.unbind(0) # Scale query q = q * self.scale # Compute attention maps attn = self.softmax(q @ k.transpose(-2, -1) + self._get_relative_positional_bias()) # Map value with attention maps output = (attn @ v).transpose(1, 2).reshape(B_, N, -1) # Perform final projection and dropout output = self.proj(output) output = self.proj_drop(output) return output class MaxViTTransformerBlock(nn.Module): """ MaxViT Transformer block. With block partition: x ← x + Unblock(RelAttention(Block(LN(x)))) x ← x + MLP(LN(x)) With grid partition: x ← x + Ungrid(RelAttention(Grid(LN(x)))) x ← x + MLP(LN(x)) Layer Normalization (LN) is applied after the grid/window partition to prevent multiple reshaping operations. Grid/window reverse (Unblock/Ungrid) is performed on the final output for the same reason. Args: in_channels (int): Number of input channels. partition_function (Callable): Partition function to be utilized (grid or window partition). reverse_function (Callable): Reverse function to be utilized (grid or window reverse). num_heads (int, optional): Number of attention heads. Default 32 grid_window_size (Tuple[int, int], optional): Grid/Window size to be utilized. Default (7, 7) attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 drop (float, optional): Dropout ratio of output. Default: 0.0 drop_path (float, optional): Dropout ratio of path. Default: 0.0 mlp_ratio (float, optional): Ratio of mlp hidden dim to embedding dim. Default: 4.0 act_layer (Type[nn.Module], optional): Type of activation layer to be utilized. Default: nn.GELU norm_layer (Type[nn.Module], optional): Type of normalization layer to be utilized. Default: nn.BatchNorm2d """ def __init__( self, in_channels: int, partition_function: Callable, reverse_function: Callable, num_heads: int = 32, grid_window_size: Tuple[int, int] = (7, 7), attn_drop: float = 0., drop: float = 0., drop_path: float = 0., mlp_ratio: float = 4., act_layer: Type[nn.Module] = nn.GELU, norm_layer: Type[nn.Module] = nn.LayerNorm, ) -> None: """ Constructor method """ super(MaxViTTransformerBlock, self).__init__() # Save parameters self.partition_function: Callable = partition_function self.reverse_function: Callable = reverse_function self.grid_window_size: Tuple[int, int] = grid_window_size # Init layers self.norm_1 = norm_layer(in_channels) self.attention = RelativeSelfAttention( in_channels=in_channels, num_heads=num_heads, grid_window_size=grid_window_size, attn_drop=attn_drop, drop=drop ) self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() self.norm_2 = norm_layer(in_channels) self.mlp = Mlp( in_features=in_channels, hidden_features=int(mlp_ratio * in_channels), act_layer=act_layer, drop=drop ) def forward(self, input: torch.Tensor) -> torch.Tensor: """ Forward pass. Args: input (torch.Tensor): Input tensor of the shape [B, C_in, H, W]. Returns: output (torch.Tensor): Output tensor of the shape [B, C_out, H (// 2), W (// 2)]. """ # Save original shape B, C, H, W = input.shape # Perform partition input_partitioned = self.partition_function(input, self.grid_window_size) input_partitioned = input_partitioned.view(-1, self.grid_window_size[0] * self.grid_window_size[1], C) # Perform normalization, attention, and dropout output = input_partitioned + self.drop_path(self.attention(self.norm_1(input_partitioned))) # Perform normalization, MLP, and dropout output = output + self.drop_path(self.mlp(self.norm_2(output))) # Reverse partition output = self.reverse_function(output, (H, W), self.grid_window_size) return output class MaxViTBlock(nn.Module): """ MaxViT block composed of MBConv block, Block Attention, and Grid Attention. Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. downscale (bool, optional): If true spatial downscaling is performed. Default: False num_heads (int, optional): Number of attention heads. Default 32 grid_window_size (Tuple[int, int], optional): Grid/Window size to be utilized. Default (7, 7) attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 drop (float, optional): Dropout ratio of output. Default: 0.0 drop_path (float, optional): Dropout ratio of path. Default: 0.0 mlp_ratio (float, optional): Ratio of mlp hidden dim to embedding dim. Default: 4.0 act_layer (Type[nn.Module], optional): Type of activation layer to be utilized. Default: nn.GELU norm_layer (Type[nn.Module], optional): Type of normalization layer to be utilized. Default: nn.BatchNorm2d norm_layer_transformer (Type[nn.Module], optional): Normalization layer in Transformer. Default: nn.LayerNorm """ def __init__( self, in_channels: int, out_channels: int, downscale: bool = False, num_heads: int = 32, grid_window_size: Tuple[int, int] = (7, 7), attn_drop: float = 0., drop: float = 0., drop_path: float = 0., mlp_ratio: float = 4., act_layer: Type[nn.Module] = nn.GELU, norm_layer: Type[nn.Module] = nn.BatchNorm2d, norm_layer_transformer: Type[nn.Module] = nn.LayerNorm ) -> None: """ Constructor method """ # Call super constructor super(MaxViTBlock, self).__init__() # Init MBConv block self.mb_conv = MBConv( in_channels=in_channels, out_channels=out_channels, downscale=downscale, act_layer=act_layer, norm_layer=norm_layer, drop_path=drop_path ) # Init Block and Grid Transformer self.block_transformer = MaxViTTransformerBlock( in_channels=out_channels, partition_function=window_partition, reverse_function=window_reverse, num_heads=num_heads, grid_window_size=grid_window_size, attn_drop=attn_drop, drop=drop, drop_path=drop_path, mlp_ratio=mlp_ratio, act_layer=act_layer, norm_layer=norm_layer_transformer ) self.grid_transformer = MaxViTTransformerBlock( in_channels=out_channels, partition_function=grid_partition, reverse_function=grid_reverse, num_heads=num_heads, grid_window_size=grid_window_size, attn_drop=attn_drop, drop=drop, drop_path=drop_path, mlp_ratio=mlp_ratio, act_layer=act_layer, norm_layer=norm_layer_transformer ) def forward(self, input: torch.Tensor) -> torch.Tensor: """ Forward pass. Args: input (torch.Tensor): Input tensor of the shape [B, C_in, H, W] Returns: output (torch.Tensor): Output tensor of the shape [B, C_out, H // 2, W // 2] (downscaling is optional) """ output = self.grid_transformer(self.block_transformer(self.mb_conv(input))) return output class MaxViTStage(nn.Module): """ Stage of the MaxViT. Args: depth (int): Depth of the stage. in_channels (int): Number of input channels. out_channels (int): Number of output channels. num_heads (int, optional): Number of attention heads. Default 32 grid_window_size (Tuple[int, int], optional): Grid/Window size to be utilized. Default (7, 7) attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 drop (float, optional): Dropout ratio of output. Default: 0.0 drop_path (float, optional): Dropout ratio of path. Default: 0.0 mlp_ratio (float, optional): Ratio of mlp hidden dim to embedding dim. Default: 4.0 act_layer (Type[nn.Module], optional): Type of activation layer to be utilized. Default: nn.GELU norm_layer (Type[nn.Module], optional): Type of normalization layer to be utilized. Default: nn.BatchNorm2d norm_layer_transformer (Type[nn.Module], optional): Normalization layer in Transformer. Default: nn.LayerNorm """ def __init__( self, depth: int, in_channels: int, out_channels: int, num_heads: int = 32, grid_window_size: Tuple[int, int] = (7, 7), attn_drop: float = 0., drop: float = 0., drop_path: Union[List[float], float] = 0., mlp_ratio: float = 4., act_layer: Type[nn.Module] = nn.GELU, norm_layer: Type[nn.Module] = nn.BatchNorm2d, norm_layer_transformer: Type[nn.Module] = nn.LayerNorm ) -> None: """ Constructor method """ # Call super constructor super(MaxViTStage, self).__init__() # Init blocks self.blocks = nn.Sequential(*[ MaxViTBlock( in_channels=in_channels if index == 0 else out_channels, out_channels=out_channels, downscale=index == 0, num_heads=num_heads, grid_window_size=grid_window_size, attn_drop=attn_drop, drop=drop, drop_path=drop_path if isinstance(drop_path, float) else drop_path[index], mlp_ratio=mlp_ratio, act_layer=act_layer, norm_layer=norm_layer, norm_layer_transformer=norm_layer_transformer ) for index in range(depth) ]) def forward(self, input=torch.Tensor) -> torch.Tensor: """ Forward pass. Args: input (torch.Tensor): Input tensor of the shape [B, C_in, H, W]. Returns: output (torch.Tensor): Output tensor of the shape [B, C_out, H // 2, W // 2]. """ output = self.blocks(input) return output class MaxViT(nn.Module): """ Implementation of the MaxViT proposed in: https://arxiv.org/pdf/2204.01697.pdf Args: in_channels (int, optional): Number of input channels to the convolutional stem. Default 3 depths (Tuple[int, ...], optional): Depth of each network stage. Default (2, 2, 5, 2) channels (Tuple[int, ...], optional): Number of channels in each network stage. Default (64, 128, 256, 512) num_classes (int, optional): Number of classes to be predicted. Default 1000 embed_dim (int, optional): Embedding dimension of the convolutional stem. Default 64 num_heads (int, optional): Number of attention heads. Default 32 grid_window_size (Tuple[int, int], optional): Grid/Window size to be utilized. Default (7, 7) attn_drop (float, optional): Dropout ratio of attention weight. Default: 0.0 drop (float, optional): Dropout ratio of output. Default: 0.0 drop_path (float, optional): Dropout ratio of path. Default: 0.0 mlp_ratio (float, optional): Ratio of mlp hidden dim to embedding dim. Default: 4.0 act_layer (Type[nn.Module], optional): Type of activation layer to be utilized. Default: nn.GELU norm_layer (Type[nn.Module], optional): Type of normalization layer to be utilized. Default: nn.BatchNorm2d norm_layer_transformer (Type[nn.Module], optional): Normalization layer in Transformer. Default: nn.LayerNorm global_pool (str, optional): Global polling type to be utilized. Default "avg" """ def __init__( self, in_channels: int = 3, depths: Tuple[int, ...] = (2, 2, 5, 2), channels: Tuple[int, ...] = (64, 128, 256, 512), num_classes: int = 1000, embed_dim: int = 64, num_heads: int = 32, grid_window_size: Tuple[int, int] = (7, 7), attn_drop: float = 0., drop=0., drop_path=0., mlp_ratio=4., act_layer=nn.GELU, norm_layer=nn.BatchNorm2d, norm_layer_transformer=nn.LayerNorm, global_pool: str = "avg" ) -> None: """ Constructor method """ # Call super constructor super(MaxViT, self).__init__() # Check parameters assert len(depths) == len(channels), "For each stage a channel dimension must be given." assert global_pool in ["avg", "max"], f"Only avg and max is supported but {global_pool} is given" # Save parameters self.num_classes: int = num_classes # Init convolutional stem self.stem = nn.Sequential( nn.Conv2d(in_channels=in_channels, out_channels=embed_dim, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1)), act_layer(), nn.Conv2d(in_channels=embed_dim, out_channels=embed_dim, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1)), act_layer(), ) # Init blocks drop_path = torch.linspace(0.0, drop_path, sum(depths)).tolist() stages = [] for index, (depth, channel) in enumerate(zip(depths, channels)): stages.append( MaxViTStage( depth=depth, in_channels=embed_dim if index == 0 else channels[index - 1], out_channels=channel, num_heads=num_heads, grid_window_size=grid_window_size, attn_drop=attn_drop, drop=drop, drop_path=drop_path[sum(depths[:index]):sum(depths[:index + 1])], mlp_ratio=mlp_ratio, act_layer=act_layer, norm_layer=norm_layer, norm_layer_transformer=norm_layer_transformer ) ) self.stages = nn.ModuleList(stages) self.global_pool: str = global_pool self.head = nn.Linear(channels[-1], num_classes) @torch.jit.ignore def no_weight_decay(self) -> Set[str]: """ Gets the names of parameters to not apply weight decay to. Returns: nwd (Set[str]): Set of parameter names to not apply weight decay to. """ nwd = set() for n, _ in self.named_parameters(): if "relative_position_bias_table" in n: nwd.add(n) return nwd def reset_classifier(self, num_classes: int, global_pool: Optional[str] = None) -> None: """Method results the classification head Args: num_classes (int): Number of classes to be predicted global_pool (str, optional): If not global pooling is updated """ self.num_classes: int = num_classes if global_pool is not None: self.global_pool = global_pool self.head = nn.Linear(self.num_features, num_classes) if num_classes > 0 else nn.Identity() def forward_features(self, input: torch.Tensor) -> torch.Tensor: """ Forward pass of feature extraction. Args: input (torch.Tensor): Input images of the shape [B, C, H, W]. Returns: output (torch.Tensor): Image features of the backbone. """ output = input for stage in self.stages: output = stage(output) return output def forward_head(self, input: torch.Tensor, pre_logits: bool = False): """ Forward pass of classification head. Args: input (torch.Tensor): Input features pre_logits (bool, optional): If true pre-logits are returned Returns: output (torch.Tensor): Classification output of the shape [B, num_classes]. """ if self.global_pool == "avg": input = input.mean(dim=(2, 3)) elif self.global_pool == "max": input = torch.amax(input, dim=(2, 3)) return input if pre_logits else self.head(input) def forward(self, input: torch.Tensor) -> torch.Tensor: """ Forward pass Args: input (torch.Tensor): Input images of the shape [B, C, H, W]. Returns: output (torch.Tensor): Classification output of the shape [B, num_classes]. """ output = self.forward_features(self.stem(input)) output = self.forward_head(output) return output def max_vit_tiny_224(**kwargs) -> MaxViT: """ MaxViT tiny for a resolution of 224 X 224""" return MaxViT( depths=(2, 2, 5, 2), channels=(64, 128, 256, 512), embed_dim=64, **kwargs ) def max_vit_small_224(**kwargs) -> MaxViT: """ MaxViT small for a resolution of 224 X 224""" return MaxViT( depths=(2, 2, 5, 2), channels=(96, 128, 256, 512), embed_dim=64, **kwargs ) def max_vit_base_224(**kwargs) -> MaxViT: """ MaxViT base for a resolution of 224 X 224""" return MaxViT( depths=(2, 6, 14, 2), channels=(96, 192, 384, 768), embed_dim=64, **kwargs ) def max_vit_large_224(**kwargs) -> MaxViT: """ MaxViT large for a resolution of 224 X 224""" return MaxViT( depths=(2, 6, 14, 2), channels=(128, 256, 512, 1024), embed_dim=128, **kwargs ) if __name__ == '__main__': def test_partition_and_revers() -> None: input = torch.rand(7, 3, 14, 14) windows = window_partition(input=input) windows = window_reverse(windows=windows, window_size=(7, 7), original_size=input.shape[2:]) print(torch.all(input == windows)) grid = grid_partition(input=input) grid = grid_reverse(grid=grid, grid_size=(7, 7), original_size=input.shape[2:]) print(torch.all(input == grid)) def test_relative_self_attention() -> None: relative_self_attention = RelativeSelfAttention(in_channels=128) input = torch.rand(4, 128, 14 * 14) output = relative_self_attention(input) print(output.shape) def test_transformer_block() -> None: transformer = MaxViTTransformerBlock(in_channels=128, partition_function=grid_partition, reverse_function=grid_reverse) input = torch.rand(4, 128, 7, 7) output = transformer(input) print(output.shape) transformer = MaxViTTransformerBlock(in_channels=128, partition_function=window_partition, reverse_function=window_reverse) input = torch.rand(4, 128, 7, 7) output = transformer(input) print(output.shape) def test_block() -> None: block = MaxViTBlock(in_channels=128, out_channels=256, downscale=True) input = torch.rand(1, 128, 28, 28) output = block(input) print(output.shape) def test_networks() -> None: for get_network in [max_vit_tiny_224, max_vit_small_224, max_vit_base_224, max_vit_large_224]: network = get_network(num_classes=365) input = torch.rand(1, 3, 224, 224) output = network(input) print(output.shape) test_networks()