Table of Contents
Shortcuts

# Source code for torchvision.models.maxvit

    import math
    from collections import OrderedDict
    from functools import partial
    from typing import Any, Callable, List, Optional, Sequence, Tuple

    import numpy as np
    import torch
    import torch.nn.functional as F
    from torch import nn, Tensor
    from torchvision.models._api import register_model, Weights, WeightsEnum
    from torchvision.models._meta import _IMAGENET_CATEGORIES
    from torchvision.models._utils import _ovewrite_named_param, handle_legacy_interface
    from torchvision.ops.misc import Conv2dNormActivation, SqueezeExcitation
    from torchvision.ops.stochastic_depth import StochasticDepth
    from torchvision.transforms._presets import ImageClassification, InterpolationMode
    from torchvision.utils import _log_api_usage_once

    __all__ = [
        "MaxVit",
        "MaxVit_T_Weights",
        "maxvit_t",
    ]

    def _get_conv_output_shape(input_size: Tuple[int, int], kernel_size: int, stride: int, padding: int) -> Tuple[int, int]:
        return (
            (input_size[0] - kernel_size + 2 * padding) // stride + 1,
            (input_size[1] - kernel_size + 2 * padding) // stride + 1,
        )

    def _make_block_input_shapes(input_size: Tuple[int, int], n_blocks: int) -> List[Tuple[int, int]]:
        """Util function to check that the input size is correct for a MaxVit configuration."""
        shapes = []
        block_input_shape = _get_conv_output_shape(input_size, 3, 2, 1)
        for _ in range(n_blocks):
            block_input_shape = _get_conv_output_shape(block_input_shape, 3, 2, 1)
            shapes.append(block_input_shape)
        return shapes

    def _get_relative_position_index(height: int, width: int) -> torch.Tensor:
        coords = torch.stack(torch.meshgrid([torch.arange(height), torch.arange(width)]))
        coords_flat = torch.flatten(coords, 1)
        relative_coords = coords_flat[:, :, None] - coords_flat[:, None, :]
        relative_coords = relative_coords.permute(1, 2, 0).contiguous()
        relative_coords[:, :, 0] += height - 1
        relative_coords[:, :, 1] += width - 1
        relative_coords[:, :, 0] *= 2 * width - 1
        return relative_coords.sum(-1)

    class MBConv(nn.Module):
        """MBConv: Mobile Inverted Residual Bottleneck.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            expansion_ratio (float): Expansion ratio in the bottleneck.
            squeeze_ratio (float): Squeeze ratio in the SE Layer.
            stride (int): Stride of the depthwise convolution.
            activation_layer (Callable[..., nn.Module]): Activation function.
            norm_layer (Callable[..., nn.Module]): Normalization function.
            p_stochastic_dropout (float): Probability of stochastic depth.
        """

        def __init__(
            self,
            in_channels: int,
            out_channels: int,
            expansion_ratio: float,
            squeeze_ratio: float,
            stride: int,
            activation_layer: Callable[..., nn.Module],
            norm_layer: Callable[..., nn.Module],
            p_stochastic_dropout: float = 0.0,
        ) -> None:
            super().__init__()

            proj: Sequence[nn.Module]
            self.proj: nn.Module

            should_proj = stride != 1 or in_channels != out_channels
            if should_proj:
                proj = [nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, bias=True)]
                if stride == 2:
                    proj = [nn.AvgPool2d(kernel_size=3, stride=stride, padding=1)] + proj  # type: ignore
                self.proj = nn.Sequential(*proj)
            else:
                self.proj = nn.Identity()  # type: ignore

            mid_channels = int(out_channels * expansion_ratio)
            sqz_channels = int(out_channels * squeeze_ratio)

            if p_stochastic_dropout:
                self.stochastic_depth = StochasticDepth(p_stochastic_dropout, mode="row")  # type: ignore
            else:
                self.stochastic_depth = nn.Identity()  # type: ignore

            _layers = OrderedDict()
            _layers["pre_norm"] = norm_layer(in_channels)
            _layers["conv_a"] = Conv2dNormActivation(
                in_channels,
                mid_channels,
                kernel_size=1,
                stride=1,
                padding=0,
                activation_layer=activation_layer,
                norm_layer=norm_layer,
                inplace=None,
            )
            _layers["conv_b"] = Conv2dNormActivation(
                mid_channels,
                mid_channels,
                kernel_size=3,
                stride=stride,
                padding=1,
                activation_layer=activation_layer,
                norm_layer=norm_layer,
                groups=mid_channels,
                inplace=None,
            )
            _layers["squeeze_excitation"] = SqueezeExcitation(mid_channels, sqz_channels, activation=nn.SiLU)
            _layers["conv_c"] = nn.Conv2d(in_channels=mid_channels, out_channels=out_channels, kernel_size=1, bias=True)

            self.layers = nn.Sequential(_layers)

        def forward(self, x: Tensor) -> Tensor:
            """
            Args:
                x (Tensor): Input tensor with expected layout of [B, C, H, W].
            Returns:
                Tensor: Output tensor with expected layout of [B, C, H / stride, W / stride].
            """
            res = self.proj(x)
            x = self.stochastic_depth(self.layers(x))
            return res + x

    class RelativePositionalMultiHeadAttention(nn.Module):
        """Relative Positional Multi-Head Attention.

        Args:
            feat_dim (int): Number of input features.
            head_dim (int): Number of features per head.
            max_seq_len (int): Maximum sequence length.
        """

        def __init__(
            self,
            feat_dim: int,
            head_dim: int,
            max_seq_len: int,
        ) -> None:
            super().__init__()

            if feat_dim % head_dim != 0:
                raise ValueError(f"feat_dim: {feat_dim} must be divisible by head_dim: {head_dim}")

            self.n_heads = feat_dim // head_dim
            self.head_dim = head_dim
            self.size = int(math.sqrt(max_seq_len))
            self.max_seq_len = max_seq_len

            self.to_qkv = nn.Linear(feat_dim, self.n_heads * self.head_dim * 3)
            self.scale_factor = feat_dim**-0.5

            self.merge = nn.Linear(self.head_dim * self.n_heads, feat_dim)
            self.relative_position_bias_table = nn.parameter.Parameter(
                torch.empty(((2 * self.size - 1) * (2 * self.size - 1), self.n_heads), dtype=torch.float32),
            )

            self.register_buffer("relative_position_index", _get_relative_position_index(self.size, self.size))
            # initialize with truncated normal the bias
            torch.nn.init.trunc_normal_(self.relative_position_bias_table, std=0.02)

        def get_relative_positional_bias(self) -> torch.Tensor:
            bias_index = self.relative_position_index.view(-1)  # type: ignore
            relative_bias = self.relative_position_bias_table[bias_index].view(self.max_seq_len, self.max_seq_len, -1)  # type: ignore
            relative_bias = relative_bias.permute(2, 0, 1).contiguous()
            return relative_bias.unsqueeze(0)

        def forward(self, x: Tensor) -> Tensor:
            """
            Args:
                x (Tensor): Input tensor with expected layout of [B, G, P, D].
            Returns:
                Tensor: Output tensor with expected layout of [B, G, P, D].
            """
            B, G, P, D = x.shape
            H, DH = self.n_heads, self.head_dim

            qkv = self.to_qkv(x)
            q, k, v = torch.chunk(qkv, 3, dim=-1)

            q = q.reshape(B, G, P, H, DH).permute(0, 1, 3, 2, 4)
            k = k.reshape(B, G, P, H, DH).permute(0, 1, 3, 2, 4)
            v = v.reshape(B, G, P, H, DH).permute(0, 1, 3, 2, 4)

            k = k * self.scale_factor
            dot_prod = torch.einsum("B G H I D, B G H J D -> B G H I J", q, k)
            pos_bias = self.get_relative_positional_bias()

            dot_prod = F.softmax(dot_prod + pos_bias, dim=-1)

            out = torch.einsum("B G H I J, B G H J D -> B G H I D", dot_prod, v)
            out = out.permute(0, 1, 3, 2, 4).reshape(B, G, P, D)

            out = self.merge(out)
            return out

    class SwapAxes(nn.Module):
        """Permute the axes of a tensor."""

        def __init__(self, a: int, b: int) -> None:
            super().__init__()
            self.a = a
            self.b = b

        def forward(self, x: torch.Tensor) -> torch.Tensor:
            res = torch.swapaxes(x, self.a, self.b)
            return res

    class WindowPartition(nn.Module):
        """
        Partition the input tensor into non-overlapping windows.
        """

        def __init__(self) -> None:
            super().__init__()

        def forward(self, x: Tensor, p: int) -> Tensor:
            """
            Args:
                x (Tensor): Input tensor with expected layout of [B, C, H, W].
                p (int): Number of partitions.
            Returns:
                Tensor: Output tensor with expected layout of [B, H/P, W/P, P*P, C].
            """
            B, C, H, W = x.shape
            P = p
            # chunk up H and W dimensions
            x = x.reshape(B, C, H // P, P, W // P, P)
            x = x.permute(0, 2, 4, 3, 5, 1)
            # colapse P * P dimension
            x = x.reshape(B, (H // P) * (W // P), P * P, C)
            return x

    class WindowDepartition(nn.Module):
        """
        Departition the input tensor of non-overlapping windows into a feature volume of layout [B, C, H, W].
        """

        def __init__(self) -> None:
            super().__init__()

        def forward(self, x: Tensor, p: int, h_partitions: int, w_partitions: int) -> Tensor:
            """
            Args:
                x (Tensor): Input tensor with expected layout of [B, (H/P * W/P), P*P, C].
                p (int): Number of partitions.
                h_partitions (int): Number of vertical partitions.
                w_partitions (int): Number of horizontal partitions.
            Returns:
                Tensor: Output tensor with expected layout of [B, C, H, W].
            """
            B, G, PP, C = x.shape
            P = p
            HP, WP = h_partitions, w_partitions
            # split P * P dimension into 2 P tile dimensionsa
            x = x.reshape(B, HP, WP, P, P, C)
            # permute into B, C, HP, P, WP, P
            x = x.permute(0, 5, 1, 3, 2, 4)
            # reshape into B, C, H, W
            x = x.reshape(B, C, HP * P, WP * P)
            return x

    class PartitionAttentionLayer(nn.Module):
        """
        Layer for partitioning the input tensor into non-overlapping windows and applying attention to each window.

        Args:
            in_channels (int): Number of input channels.
            head_dim (int): Dimension of each attention head.
            partition_size (int): Size of the partitions.
            partition_type (str): Type of partitioning to use. Can be either "grid" or "window".
            grid_size (Tuple[int, int]): Size of the grid to partition the input tensor into.
            mlp_ratio (int): Ratio of the  feature size expansion in the MLP layer.
            activation_layer (Callable[..., nn.Module]): Activation function to use.
            norm_layer (Callable[..., nn.Module]): Normalization function to use.
            attention_dropout (float): Dropout probability for the attention layer.
            mlp_dropout (float): Dropout probability for the MLP layer.
            p_stochastic_dropout (float): Probability of dropping out a partition.
        """

        def __init__(
            self,
            in_channels: int,
            head_dim: int,
            # partitioning parameters
            partition_size: int,
            partition_type: str,
            # grid size needs to be known at initialization time
            # because we need to know hamy relative offsets there are in the grid
            grid_size: Tuple[int, int],
            mlp_ratio: int,
            activation_layer: Callable[..., nn.Module],
            norm_layer: Callable[..., nn.Module],
            attention_dropout: float,
            mlp_dropout: float,
            p_stochastic_dropout: float,
        ) -> None:
            super().__init__()

            self.n_heads = in_channels // head_dim
            self.head_dim = head_dim
            self.n_partitions = grid_size[0] // partition_size
            self.partition_type = partition_type
            self.grid_size = grid_size

            if partition_type not in ["grid", "window"]:
                raise ValueError("partition_type must be either 'grid' or 'window'")

            if partition_type == "window":
                self.p, self.g = partition_size, self.n_partitions
            else:
                self.p, self.g = self.n_partitions, partition_size

            self.partition_op = WindowPartition()
            self.departition_op = WindowDepartition()
            self.partition_swap = SwapAxes(-2, -3) if partition_type == "grid" else nn.Identity()
            self.departition_swap = SwapAxes(-2, -3) if partition_type == "grid" else nn.Identity()

            self.attn_layer = nn.Sequential(
                norm_layer(in_channels),
                # it's always going to be partition_size ** 2 because
                # of the axis swap in the case of grid partitioning
                RelativePositionalMultiHeadAttention(in_channels, head_dim, partition_size**2),
                nn.Dropout(attention_dropout),
            )

            # pre-normalization similar to transformer layers
            self.mlp_layer = nn.Sequential(
                nn.LayerNorm(in_channels),
                nn.Linear(in_channels, in_channels * mlp_ratio),
                activation_layer(),
                nn.Linear(in_channels * mlp_ratio, in_channels),
                nn.Dropout(mlp_dropout),
            )

            # layer scale factors
            self.stochastic_dropout = StochasticDepth(p_stochastic_dropout, mode="row")

        def forward(self, x: Tensor) -> Tensor:
            """
            Args:
                x (Tensor): Input tensor with expected layout of [B, C, H, W].
            Returns:
                Tensor: Output tensor with expected layout of [B, C, H, W].
            """

            # Undefined behavior if H or W are not divisible by p
            # https://github.com/google-research/maxvit/blob/da76cf0d8a6ec668cc31b399c4126186da7da944/maxvit/models/maxvit.py#L766
            gh, gw = self.grid_size[0] // self.p, self.grid_size[1] // self.p
            torch._assert(
                self.grid_size[0] % self.p == 0 and self.grid_size[1] % self.p == 0,
                "Grid size must be divisible by partition size. Got grid size of {} and partition size of {}".format(
                    self.grid_size, self.p
                ),
            )

            x = self.partition_op(x, self.p)
            x = self.partition_swap(x)
            x = x + self.stochastic_dropout(self.attn_layer(x))
            x = x + self.stochastic_dropout(self.mlp_layer(x))
            x = self.departition_swap(x)
            x = self.departition_op(x, self.p, gh, gw)

            return x

    class MaxVitLayer(nn.Module):
        """
        MaxVit layer consisting of a MBConv layer followed by a PartitionAttentionLayer with `window` and a PartitionAttentionLayer with `grid`.

        Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            expansion_ratio (float): Expansion ratio in the bottleneck.
            squeeze_ratio (float): Squeeze ratio in the SE Layer.
            stride (int): Stride of the depthwise convolution.
            activation_layer (Callable[..., nn.Module]): Activation function.
            norm_layer (Callable[..., nn.Module]): Normalization function.
            head_dim (int): Dimension of the attention heads.
            mlp_ratio (int): Ratio of the MLP layer.
            mlp_dropout (float): Dropout probability for the MLP layer.
            attention_dropout (float): Dropout probability for the attention layer.
            p_stochastic_dropout (float): Probability of stochastic depth.
            partition_size (int): Size of the partitions.
            grid_size (Tuple[int, int]): Size of the input feature grid.
        """

        def __init__(
            self,
            # conv parameters
            in_channels: int,
            out_channels: int,
            squeeze_ratio: float,
            expansion_ratio: float,
            stride: int,
            # conv + transformer parameters
            norm_layer: Callable[..., nn.Module],
            activation_layer: Callable[..., nn.Module],
            # transformer parameters
            head_dim: int,
            mlp_ratio: int,
            mlp_dropout: float,
            attention_dropout: float,
            p_stochastic_dropout: float,
            # partitioning parameters
            partition_size: int,
            grid_size: Tuple[int, int],
        ) -> None:
            super().__init__()

            layers: OrderedDict = OrderedDict()

            # convolutional layer
            layers["MBconv"] = MBConv(
                in_channels=in_channels,
                out_channels=out_channels,
                expansion_ratio=expansion_ratio,
                squeeze_ratio=squeeze_ratio,
                stride=stride,
                activation_layer=activation_layer,
                norm_layer=norm_layer,
                p_stochastic_dropout=p_stochastic_dropout,
            )
            # attention layers, block -> grid
            layers["window_attention"] = PartitionAttentionLayer(
                in_channels=out_channels,
                head_dim=head_dim,
                partition_size=partition_size,
                partition_type="window",
                grid_size=grid_size,
                mlp_ratio=mlp_ratio,
                activation_layer=activation_layer,
                norm_layer=nn.LayerNorm,
                attention_dropout=attention_dropout,
                mlp_dropout=mlp_dropout,
                p_stochastic_dropout=p_stochastic_dropout,
            )
            layers["grid_attention"] = PartitionAttentionLayer(
                in_channels=out_channels,
                head_dim=head_dim,
                partition_size=partition_size,
                partition_type="grid",
                grid_size=grid_size,
                mlp_ratio=mlp_ratio,
                activation_layer=activation_layer,
                norm_layer=nn.LayerNorm,
                attention_dropout=attention_dropout,
                mlp_dropout=mlp_dropout,
                p_stochastic_dropout=p_stochastic_dropout,
            )
            self.layers = nn.Sequential(layers)

        def forward(self, x: Tensor) -> Tensor:
            """
            Args:
                x (Tensor): Input tensor of shape (B, C, H, W).
            Returns:
                Tensor: Output tensor of shape (B, C, H, W).
            """
            x = self.layers(x)
            return x

    class MaxVitBlock(nn.Module):
        """
        A MaxVit block consisting of `n_layers` MaxVit layers.

         Args:
            in_channels (int): Number of input channels.
            out_channels (int): Number of output channels.
            expansion_ratio (float): Expansion ratio in the bottleneck.
            squeeze_ratio (float): Squeeze ratio in the SE Layer.
            activation_layer (Callable[..., nn.Module]): Activation function.
            norm_layer (Callable[..., nn.Module]): Normalization function.
            head_dim (int): Dimension of the attention heads.
            mlp_ratio (int): Ratio of the MLP layer.
            mlp_dropout (float): Dropout probability for the MLP layer.
            attention_dropout (float): Dropout probability for the attention layer.
            p_stochastic_dropout (float): Probability of stochastic depth.
            partition_size (int): Size of the partitions.
            input_grid_size (Tuple[int, int]): Size of the input feature grid.
            n_layers (int): Number of layers in the block.
            p_stochastic (List[float]): List of probabilities for stochastic depth for each layer.
        """

        def __init__(
            self,
            # conv parameters
            in_channels: int,
            out_channels: int,
            squeeze_ratio: float,
            expansion_ratio: float,
            # conv + transformer parameters
            norm_layer: Callable[..., nn.Module],
            activation_layer: Callable[..., nn.Module],
            # transformer parameters
            head_dim: int,
            mlp_ratio: int,
            mlp_dropout: float,
            attention_dropout: float,
            # partitioning parameters
            partition_size: int,
            input_grid_size: Tuple[int, int],
            # number of layers
            n_layers: int,
            p_stochastic: List[float],
        ) -> None:
            super().__init__()
            if not len(p_stochastic) == n_layers:
                raise ValueError(f"p_stochastic must have length n_layers={n_layers}, got p_stochastic={p_stochastic}.")

            self.layers = nn.ModuleList()
            # account for the first stride of the first layer
            self.grid_size = _get_conv_output_shape(input_grid_size, kernel_size=3, stride=2, padding=1)

            for idx, p in enumerate(p_stochastic):
                stride = 2 if idx == 0 else 1
                self.layers += [
                    MaxVitLayer(
                        in_channels=in_channels if idx == 0 else out_channels,
                        out_channels=out_channels,
                        squeeze_ratio=squeeze_ratio,
                        expansion_ratio=expansion_ratio,
                        stride=stride,
                        norm_layer=norm_layer,
                        activation_layer=activation_layer,
                        head_dim=head_dim,
                        mlp_ratio=mlp_ratio,
                        mlp_dropout=mlp_dropout,
                        attention_dropout=attention_dropout,
                        partition_size=partition_size,
                        grid_size=self.grid_size,
                        p_stochastic_dropout=p,
                    ),
                ]

        def forward(self, x: Tensor) -> Tensor:
            """
            Args:
                x (Tensor): Input tensor of shape (B, C, H, W).
            Returns:
                Tensor: Output tensor of shape (B, C, H, W).
            """
            for layer in self.layers:
                x = layer(x)
            return x

    class MaxVit(nn.Module):
        """
        Implements MaxVit Transformer from the `MaxViT: Multi-Axis Vision Transformer <https://arxiv.org/abs/2204.01697>`_ paper.
        Args:
            input_size (Tuple[int, int]): Size of the input image.
            stem_channels (int): Number of channels in the stem.
            partition_size (int): Size of the partitions.
            block_channels (List[int]): Number of channels in each block.
            block_layers (List[int]): Number of layers in each block.
            stochastic_depth_prob (float): Probability of stochastic depth. Expands to a list of probabilities for each layer that scales linearly to the specified value.
            squeeze_ratio (float): Squeeze ratio in the SE Layer. Default: 0.25.
            expansion_ratio (float): Expansion ratio in the MBConv bottleneck. Default: 4.
            norm_layer (Callable[..., nn.Module]): Normalization function. Default: None (setting to None will produce a `BatchNorm2d(eps=1e-3, momentum=0.99)`).
            activation_layer (Callable[..., nn.Module]): Activation function Default: nn.GELU.
            head_dim (int): Dimension of the attention heads.
            mlp_ratio (int): Expansion ratio of the MLP layer. Default: 4.
            mlp_dropout (float): Dropout probability for the MLP layer. Default: 0.0.
            attention_dropout (float): Dropout probability for the attention layer. Default: 0.0.
            num_classes (int): Number of classes. Default: 1000.
        """

        def __init__(
            self,
            # input size parameters
            input_size: Tuple[int, int],
            # stem and task parameters
            stem_channels: int,
            # partitioning parameters
            partition_size: int,
            # block parameters
            block_channels: List[int],
            block_layers: List[int],
            # attention head dimensions
            head_dim: int,
            stochastic_depth_prob: float,
            # conv + transformer parameters
            # norm_layer is applied only to the conv layers
            # activation_layer is applied both to conv and transformer layers
            norm_layer: Optional[Callable[..., nn.Module]] = None,
            activation_layer: Callable[..., nn.Module] = nn.GELU,
            # conv parameters
            squeeze_ratio: float = 0.25,
            expansion_ratio: float = 4,
            # transformer parameters
            mlp_ratio: int = 4,
            mlp_dropout: float = 0.0,
            attention_dropout: float = 0.0,
            # task parameters
            num_classes: int = 1000,
        ) -> None:
            super().__init__()
            _log_api_usage_once(self)

            input_channels = 3

            # https://github.com/google-research/maxvit/blob/da76cf0d8a6ec668cc31b399c4126186da7da944/maxvit/models/maxvit.py#L1029-L1030
            # for the exact parameters used in batchnorm
            if norm_layer is None:
                norm_layer = partial(nn.BatchNorm2d, eps=1e-3, momentum=0.99)

            # Make sure input size will be divisible by the partition size in all blocks
            # Undefined behavior if H or W are not divisible by p
            # https://github.com/google-research/maxvit/blob/da76cf0d8a6ec668cc31b399c4126186da7da944/maxvit/models/maxvit.py#L766
            block_input_sizes = _make_block_input_shapes(input_size, len(block_channels))
            for idx, block_input_size in enumerate(block_input_sizes):
                if block_input_size[0] % partition_size != 0 or block_input_size[1] % partition_size != 0:
                    raise ValueError(
                        f"Input size {block_input_size} of block {idx} is not divisible by partition size {partition_size}. "
                        f"Consider changing the partition size or the input size.\n"
                        f"Current configuration yields the following block input sizes: {block_input_sizes}."
                    )

            # stem
            self.stem = nn.Sequential(
                Conv2dNormActivation(
                    input_channels,
                    stem_channels,
                    3,
                    stride=2,
                    norm_layer=norm_layer,
                    activation_layer=activation_layer,
                    bias=False,
                    inplace=None,
                ),
                Conv2dNormActivation(
                    stem_channels, stem_channels, 3, stride=1, norm_layer=None, activation_layer=None, bias=True
                ),
            )

            # account for stem stride
            input_size = _get_conv_output_shape(input_size, kernel_size=3, stride=2, padding=1)
            self.partition_size = partition_size

            # blocks
            self.blocks = nn.ModuleList()
            in_channels = [stem_channels] + block_channels[:-1]
            out_channels = block_channels

            # precompute the stochastich depth probabilities from 0 to stochastic_depth_prob
            # since we have N blocks with L layers, we will have N * L probabilities uniformly distributed
            # over the range [0, stochastic_depth_prob]
            p_stochastic = np.linspace(0, stochastic_depth_prob, sum(block_layers)).tolist()

            p_idx = 0
            for in_channel, out_channel, num_layers in zip(in_channels, out_channels, block_layers):
                self.blocks.append(
                    MaxVitBlock(
                        in_channels=in_channel,
                        out_channels=out_channel,
                        squeeze_ratio=squeeze_ratio,
                        expansion_ratio=expansion_ratio,
                        norm_layer=norm_layer,
                        activation_layer=activation_layer,
                        head_dim=head_dim,
                        mlp_ratio=mlp_ratio,
                        mlp_dropout=mlp_dropout,
                        attention_dropout=attention_dropout,
                        partition_size=partition_size,
                        input_grid_size=input_size,
                        n_layers=num_layers,
                        p_stochastic=p_stochastic[p_idx : p_idx + num_layers],
                    ),
                )
                input_size = self.blocks[-1].grid_size
                p_idx += num_layers

            # see https://github.com/google-research/maxvit/blob/da76cf0d8a6ec668cc31b399c4126186da7da944/maxvit/models/maxvit.py#L1137-L1158
            # for why there is Linear -> Tanh -> Linear
            self.classifier = nn.Sequential(
                nn.AdaptiveAvgPool2d(1),
                nn.Flatten(),
                nn.LayerNorm(block_channels[-1]),
                nn.Linear(block_channels[-1], block_channels[-1]),
                nn.Tanh(),
                nn.Linear(block_channels[-1], num_classes, bias=False),
            )

            self._init_weights()

        def forward(self, x: Tensor) -> Tensor:
            x = self.stem(x)
            for block in self.blocks:
                x = block(x)
            x = self.classifier(x)
            return x

        def _init_weights(self):
            for m in self.modules():
                if isinstance(m, nn.Conv2d):
                    nn.init.normal_(m.weight, std=0.02)
                    if m.bias is not None:
                        nn.init.zeros_(m.bias)
                elif isinstance(m, nn.BatchNorm2d):
                    nn.init.constant_(m.weight, 1)
                    nn.init.constant_(m.bias, 0)
                elif isinstance(m, nn.Linear):
                    nn.init.normal_(m.weight, std=0.02)
                    if m.bias is not None:
                        nn.init.zeros_(m.bias)

    def _maxvit(
        # stem parameters
        stem_channels: int,
        # block parameters
        block_channels: List[int],
        block_layers: List[int],
        stochastic_depth_prob: float,
        # partitioning parameters
        partition_size: int,
        # transformer parameters
        head_dim: int,
        # Weights API
        weights: Optional[WeightsEnum] = None,
        progress: bool = False,
        # kwargs,
        **kwargs: Any,
    ) -> MaxVit:

        if weights is not None:
            _ovewrite_named_param(kwargs, "num_classes", len(weights.meta["categories"]))
            assert weights.meta["min_size"][0] == weights.meta["min_size"][1]
            _ovewrite_named_param(kwargs, "input_size", weights.meta["min_size"])

        input_size = kwargs.pop("input_size", (224, 224))

        model = MaxVit(
            stem_channels=stem_channels,
            block_channels=block_channels,
            block_layers=block_layers,
            stochastic_depth_prob=stochastic_depth_prob,
            head_dim=head_dim,
            partition_size=partition_size,
            input_size=input_size,
            **kwargs,
        )

        if weights is not None:
            model.load_state_dict(weights.get_state_dict(progress=progress, check_hash=True))

        return model

    class MaxVit_T_Weights(WeightsEnum):
        IMAGENET1K_V1 = Weights(
            # URL empty until official release
            url="https://download.pytorch.org/models/maxvit_t-bc5ab103.pth",
            transforms=partial(
                ImageClassification, crop_size=224, resize_size=224, interpolation=InterpolationMode.BICUBIC
            ),
            meta={
                "categories": _IMAGENET_CATEGORIES,
                "num_params": 30919624,
                "min_size": (224, 224),
                "recipe": "https://github.com/pytorch/vision/tree/main/references/classification#maxvit",
                "_metrics": {
                    "ImageNet-1K": {
                        "acc@1": 83.700,
                        "acc@5": 96.722,
                    }
                },
                "_ops": 5.558,
                "_file_size": 118.769,
                "_docs": """These weights reproduce closely the results of the paper using a similar training recipe.""",
            },
        )
        DEFAULT = IMAGENET1K_V1

    @register_model()
    @handle_legacy_interface(weights=("pretrained", MaxVit_T_Weights.IMAGENET1K_V1))
    def maxvit_t(*, weights: Optional[MaxVit_T_Weights] = None, progress: bool = True, **kwargs: Any) -> MaxVit:
        """
        Constructs a maxvit_t architecture from
        `MaxViT: Multi-Axis Vision Transformer <https://arxiv.org/abs/2204.01697>`_.

        Args:
            weights (:class:`~torchvision.models.MaxVit_T_Weights`, optional): The
                pretrained weights to use. See
                :class:`~torchvision.models.MaxVit_T_Weights` below for
                more details, and possible values. By default, no pre-trained
                weights are used.
            progress (bool, optional): If True, displays a progress bar of the
                download to stderr. Default is True.
            **kwargs: parameters passed to the ``torchvision.models.maxvit.MaxVit``
                base class. Please refer to the `source code
                <https://github.com/pytorch/vision/blob/main/torchvision/models/maxvit.py>`_
                for more details about this class.

        .. autoclass:: torchvision.models.MaxVit_T_Weights
            :members:
        """
        weights = MaxVit_T_Weights.verify(weights)

        return _maxvit(
            stem_channels=64,
            block_channels=[64, 128, 256, 512],
            block_layers=[2, 2, 5, 2],
            head_dim=32,
            stochastic_depth_prob=0.2,
            partition_size=7,
            weights=weights,
            progress=progress,
            **kwargs,
        )

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.