class

MaxViTForImageClassification

extendsPretrainedModelClassificationHeadMixin

MaxViTForImageClassification(config: MaxViTConfig)

source edit

MaxViT with a linear classification head (Tu et al., 2022).

Wraps the same multi-axis attention trunk as MaxViT (stem + four MaxViT stages) and adds the reference-recipe head: global average pool → LayerNorm → Linear(dim) + Tanh → final linear classifier:

\text{logits} = W_{\text{cls}}\, \mathrm{Tanh}\!\bigl(W_{\text{pre}}\, \mathrm{LN}(\mathrm{GAP}(z^{L}))\bigr) + b_{\text{cls}}.

Pass labels to forward to compute the cross-entropy loss in the same pass.

Parameters

configMaxViTConfig

Architecture specification. Must set num_classes to the desired number of output categories. See MaxViTConfig.

Attributes

stemnn.Module

Two-layer stride-2 convolutional stem.

stagesnn.ModuleList

Four MaxViT stages.

head_HeadNorm

LayerNorm → pre_logits.fc (Linear) → Tanh → fc (Linear).

Notes

Reference: Zhengzhong Tu et al., "MaxViT: Multi-Axis Vision Transformer", ECCV 2022. MaxViT-Tiny / Small / Base / Large / XLarge reach 83.6 / 84.5 / 84.9 / 85.2 / 85.5 % top-1 on ImageNet-1k (Tu et al., 2022, Table 6).

Examples

End-to-end inference with the default MaxViT-Tiny classifier:
>>> import lucid
>>> from lucid.models.vision.maxvit import (
...     MaxViTConfig, MaxViTForImageClassification,
... )
>>> model = MaxViTForImageClassification(MaxViTConfig())
>>> x = lucid.randn(1, 3, 224, 224)
>>> out = model(x)
>>> out.logits.shape
(1, 1000)

Used by 2

Constructors

dunder

init

→None

__init__(config: MaxViTConfig)

source edit

Instance methods

forward

→ImageClassificationOutput

forward(x: Tensor, labels: Tensor | None = None)

source edit

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

End-to-end inference with the default MaxViT-Tiny classifier: >>> import lucid >>> from lucid.models.vision.maxvit import ( ... MaxViTConfig, MaxViTForImageClassification, ... ) >>> model = MaxViTForImageClassification(MaxViTConfig()) >>> x = lucid.randn(1, 3, 224, 224) >>> out = model(x) >>> out.logits.shape (1, 1000)