module

PVT

23 members

lucid.models.vision.pvt

PVT — Pyramid Vision Transformer v2 (Wang et al., 2021/2022).

Image Classification

Paper

Wang, Wenhai, et al. "Pyramid Vision Transformer: A Versatile Backbone for Dense Prediction without Convolutions." Proceedings of the IEEE/CVF International Conference on Computer Vision, 2021, pp. 568-578.

Overview

The Pyramid Vision Transformer (PVT) introduces a hierarchical feature pyramid into a pure-transformer backbone so that it can serve as a drop-in replacement for ResNet-style backbones in dense prediction tasks such as detection and segmentation. Plain ViT produces a single-scale, low-resolution feature map and scales quadratically with the input resolution; PVT addresses both problems with two design changes.

First, the backbone has four stages, each beginning with a patch embedding convolution (overlapping $7 \times 7$ in PVT v2, non-overlapping in v1) of stride $s_i \in \{4, 2, 2, 2\}$ , producing the canonical $(\tfrac{H}{4}, \tfrac{W}{4}) \to \dots \to (\tfrac{H}{32}, \tfrac{W}{32})$ pyramid. Second, each transformer block uses Spatial-Reduction Attention (SRA): before computing keys and values, the token sequence is reshaped back to a 2-D feature map and downsampled by a stride- $R_i$ convolution, then flattened again. Concretely,

K, V \leftarrow \mathrm{LN}\!\bigl( W_R \cdot \mathrm{Reshape}_{R_i \times R_i}(x)\bigr), \qquad \mathrm{SRA}(x) = \mathrm{softmax}\!\left( \frac{Q K^\top}{\sqrt{d}}\right) V,

where $R_i$ decreases with depth (e.g. $8, 4, 2, 1$ ). This cuts the per-stage attention cost from $\mathcal{O}((HW)^2)$ to $\mathcal{O}((HW)^2 / R_i^2)$ , enabling high-resolution inputs at reasonable cost. Variants PVT-Tiny / Small / Medium / Large scale stage depths along this shared pyramidal-SRA design.

Classes

dataPVTConfig1 methods

Configuration dataclass for every PVT (v2) variant.

classPVT4 methods

PVT v2 backbone (Wang et al., 2022).

classPVTForImageClassification2 methods

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

Functions

fnpvt_tiny→ PVT

PVT-Tiny backbone — alias for pvt_v2_b1.

fnpvt_tiny_cls→ PVTForImageClassification

PVT-Tiny image classifier — alias for pvt_v2_b1_cls.

fnpvt_v2_b0→ PVT

PVT v2-B0 backbone (Wang et al., 2022).

fnpvt_v2_b0_cls→ PVTForImageClassification

PVT v2-B0 image classifier (Wang et al., 2022).

fnpvt_v2_b1→ PVT

PVT v2-B1 backbone (Wang et al., 2022).

fnpvt_v2_b1_cls→ PVTForImageClassification

PVT v2-B1 image classifier (Wang et al., 2022).

fnpvt_v2_b2→ PVT

PVT v2-B2 backbone (Wang et al., 2022).

fnpvt_v2_b2_cls→ PVTForImageClassification

PVT v2-B2 image classifier (Wang et al., 2022).

fnpvt_v2_b3→ PVT

PVT v2-B3 backbone (Wang et al., 2022).

fnpvt_v2_b3_cls→ PVTForImageClassification

PVT v2-B3 image classifier (Wang et al., 2022).

fnpvt_v2_b4→ PVT

PVT v2-B4 backbone (Wang et al., 2022).

fnpvt_v2_b4_cls→ PVTForImageClassification

PVT v2-B4 image classifier (Wang et al., 2022).

fnpvt_v2_b5→ PVT

PVT v2-B5 backbone (Wang et al., 2022).

fnpvt_v2_b5_cls→ PVTForImageClassification

PVT v2-B5 image classifier (Wang et al., 2022).

Weights

classPVTv2B0Weights

Pretrained weights for lucid.models.pvt_v2_b0_cls.

classPVTv2B1Weights

Pretrained weights for lucid.models.pvt_v2_b1_cls.

classPVTv2B2Weights

Pretrained weights for lucid.models.pvt_v2_b2_cls.

classPVTv2B3Weights

Pretrained weights for lucid.models.pvt_v2_b3_cls.

classPVTv2B4Weights

Pretrained weights for lucid.models.pvt_v2_b4_cls.

classPVTv2B5Weights

Pretrained weights for lucid.models.pvt_v2_b5_cls.