class

MaxViT

extendsPretrainedModelBackboneMixin

MaxViT(config: MaxViTConfig)

source edit

MaxViT backbone (Tu et al., 2022).

MaxViT combines a convolutional MBConv stage with two complementary sparse self-attention mechanisms — block attention and grid attention — to obtain dense local plus dense global receptive fields at linear complexity in the number of tokens. The backbone is a four-stage pyramid preceded by a two-conv stem, and every MaxViT block applies (i) MobileNetV3-style MBConv, (ii) block attention, and (iii) grid attention in sequence.

Given a feature map of shape $(H, W, C)$ , block attention partitions it into non-overlapping $P \times P$ windows (default $P = 7$ ) and runs full self-attention inside each window — identical to Swin's W-MSA. Grid attention takes the same feature map and gathers tokens that share the same intra-window position across all windows. Formally, with $H = h \cdot P$ , $W = w \cdot P$ :

\mathrm{Block}: (H, W) \to (h \cdot w,\ P \cdot P),\quad \mathrm{Grid} : (H, W) \to (P \cdot P,\ h \cdot w).

forward_features returns the global-average-pooled $(B, \texttt{dims[-1]})$ feature.

Parameters

configMaxViTConfig

Frozen dataclass specifying depths, dims, window_size, num_heads, mlp_ratio, in_channels, and num_classes. See MaxViTConfig.

Attributes

stemnn.Module

Two-layer convolutional stem with stride-2 first conv.

stagesnn.ModuleList

Four MaxViT stages, each containing a stack of MaxViT blocks.

head_HeadNorm

Pre-classification head: LayerNorm → Linear + Tanh → final Linear (also present on the backbone for parameter sharing).

feature_infolist[FeatureInfo]

Four-stage feature description with reductions

(4, 8, 16, 32)

Notes

Reference: Zhengzhong Tu et al., "MaxViT: Multi-Axis Vision Transformer", ECCV 2022, arXiv:2204.01697.

Examples

Build a MaxViT-Tiny backbone and run a forward pass:
>>> import lucid
>>> from lucid.models.vision.maxvit import MaxViT, MaxViTConfig
>>> model = MaxViT(MaxViTConfig())
>>> x = lucid.randn(1, 3, 224, 224)
>>> feat = model.forward_features(x)
>>> feat.shape                       # (B, dims[-1])
(1, 512)

Used by 2

Constructors

dunder

init

→None

__init__(config: MaxViTConfig)

source edit

Properties

prop

feature_info

→list[FeatureInfo]

feature_info: list[FeatureInfo]

source edit

Instance methods

forward

→BaseModelOutput

forward(x: Tensor)

source edit

Return bias tensor of shape (1, num_heads, ws², ws²).

forward_features

→Tensor

forward_features(x: Tensor)

source edit

Build a MaxViT-Tiny backbone and run a forward pass: >>> import lucid >>> from lucid.models.vision.maxvit import MaxViT, MaxViTConfig >>> model = MaxViT(MaxViTConfig()) >>> x = lucid.randn(1, 3, 224, 224) >>> feat = model.forward_features(x) >>> feat.shape # (B, dims[-1]) (1, 512)