MaxViTForImageClassification
ImageClassificationModelClassificationHeadMixinMaxViTForImageClassification(config: MaxViTConfig)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:
Pass labels to forward to compute the cross-entropy
loss in the same pass.
Parameters
configMaxViTConfignum_classes to the
desired number of output categories. See MaxViTConfig.Attributes
stemnn.Modulestagesnn.ModuleListhead_HeadNormNotes
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
1Instance methods
2forward(x: Tensor, labels: Tensor | None = None)Return bias tensor of shape (1, num_heads, ws², ws²).
Replace the final projection of the norm-MLP head.
The head is a _HeadNorm (matching timm's NormMlpClassifierHead
layout), so only its inner fc is class-dependent; the norm and
pre-logits stages carry over. The mixin's default looks for
self.classifier and raised AttributeError here.