CoAtNet backbone (Dai et al., 2021).
CoAtNet is a hybrid backbone that interleaves depthwise convolutions and relative self-attention in a single four-stage pyramid (preceded by a two-layer convolutional stem). The first two stages () use MBConv blocks (squeeze-and- excitation, expansion ratio 4) and the last two stages () use relative-attention transformer blocks operating on flattened token sequences:
where is a learned bias indexed by the relative spatial offset between tokens. This recovers the translation equivariance that convolutions provide while still permitting global, data-dependent mixing. Each stage downsamples , so the final feature map is .
forward_features returns the raw spatial feature map from
the last attention stage. Use this backbone when you need
multi-scale or spatial features for detection / segmentation; for
end-to-end classification use
CoAtNetForImageClassification.
Parameters
configCoAtNetConfigblocks_per_stage, dims,
stem_width, attn_heads, mbconv_expand,
image_size, and in_channels. See CoAtNetConfig.Attributes
stemnn.Sequentials3_TransformerStages4_TransformerStagefeature_infolist[FeatureInfo]Notes
Reference: Zihang Dai et al., "CoAtNet: Marrying Convolution and Attention for All Data Sizes", NeurIPS 2021, arXiv:2106.04803.
Examples
Build a CoAtNet-0 backbone and run a forward pass:
>>> import lucid
>>> from lucid.models.vision.coatnet import CoAtNet, CoAtNetConfig
>>> model = CoAtNet(CoAtNetConfig())
>>> x = lucid.randn(1, 3, 224, 224)
>>> feat = model.forward_features(x)
>>> feat.shape # (B, dims[-1], H/32, W/32)
(1, 768, 7, 7)