class

PVTForImageClassification

extendsPretrainedModelClassificationHeadMixin

PVTForImageClassification(config: PVTConfig)

source edit

PVT v2 with a linear classification head (Wang et al., 2022).

Wraps the same hierarchical four-stage trunk as PVT (overlapping patch embeddings + Spatial-Reduction Attention) and adds a global mean pool + single nn.Linear classification head:

\text{logits} = W_{\text{head}}\, \mathrm{Mean}(z^{L}) + b_{\text{head}}, \qquad W_{\text{head}} \in \mathbb{R}^{C_{\text{out}} \times d_{L}}.

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

Parameters

configPVTConfig

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

Attributes

patch_embed_OverlapPatchEmbed

Top-level

7 \times 7

stride-4 patch embedding.

stagesnn.ModuleList

Four PVT v2 stages.

classifiernn.Linear

Final linear projection of width (num_classes, embed_dims[-1]).

Notes

Reference: Wang et al., "PVT v2: Improved Baselines with Pyramid Vision Transformer", CVMJ 2022. PVT v2-B0 / B1 / B2 / B3 / B4 / B5 reach 70.5 / 78.7 / 82.0 / 83.1 / 83.6 / 83.8 % top-1 on ImageNet-1k at 224x224 (Wang et al., 2022, Table 1).

Examples

End-to-end inference with the default PVT v2-B1 classifier:
>>> import lucid
>>> from lucid.models.vision.pvt import (
...     PVTConfig, PVTForImageClassification,
... )
>>> model = PVTForImageClassification(PVTConfig())
>>> x = lucid.randn(1, 3, 224, 224)
>>> out = model(x)
>>> out.logits.shape
(1, 1000)

Used by 2

Constructors

dunder

init

→None

__init__(config: PVTConfig)

source edit

Instance methods

forward

→ImageClassificationOutput

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

source edit

End-to-end inference with the default PVT v2-B1 classifier: >>> import lucid >>> from lucid.models.vision.pvt import ( ... PVTConfig, PVTForImageClassification, ... ) >>> model = PVTForImageClassification(PVTConfig()) >>> x = lucid.randn(1, 3, 224, 224) >>> out = model(x) >>> out.logits.shape (1, 1000)