Source code for torchvision.models.maxvit
import math
from functools import partial
from typing import Any, Callable, List, Optional, OrderedDict, 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 parameteres
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[str, Any] = OrderedDict() # type: ignore
# 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))
return model
[docs]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,
}
},
"_docs": """These weights reproduce closely the results of the paper using a similar training recipe.""",
},
)
DEFAULT = IMAGENET1K_V1
[docs]@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,
)
