data
CoAtNetConfig
extends
ModelConfigCoAtNetConfig(num_classes: int = 1000, in_channels: int = 3, image_size: int = 224, variant: str = 'coatnet_0', blocks_per_stage: tuple[int, ...] = (2, 3, 5, 2), dims: tuple[int, ...] = (96, 192, 384, 768), stem_width: int = 64, attn_heads: tuple[int, ...] = (12, 24), mbconv_expand: int = 4, head_hidden_size: int | None = 768, dropout: float = 0.0, mixed_s3: tuple[int, int, int] | None = None)Configuration dataclass for every CoAtNet variant.
CoAtNetConfig is an immutable container that fully specifies
the architecture of a CoAtNet (Dai et al., 2021). CoAtNet uses
four stages after a two-conv stem, with the first two stages
being depthwise-convolutional (MBConv) and the last two stages
being relative-attention transformer blocks.
Parameters
num_classesint= 1000Number of output classes for the classification head. Defaults
to
1000 (ImageNet-1k).in_channelsint= 3Number of input image channels. Defaults to
3 (RGB).image_sizeint= 224Spatial side length of the square input image (used to build
the fixed-size relative-position bias tables in T-stages).
Defaults to
224.variantstr= 'coatnet_0'Informational label identifying the canonical CoAtNet variant
(
"coatnet_0" through "coatnet_7"). Defaults to
"coatnet_0".blocks_per_stagetuple of int= (2, 3, 5, 2)Number of blocks in each of the four stages
. Defaults to
(2, 3, 5, 2)
(CoAtNet-0).dimstuple of int= (96, 192, 384, 768)Channel widths for each of the four stages. Defaults to
(96, 192, 384, 768) (CoAtNet-0).stem_widthint= 64Output channels of the two-layer convolutional stem (before
). Defaults to
64.attn_headstuple of int= (12, 24)Number of attention heads for each of the two T-stages
. Defaults to
(12, 24), derived from
dim / 32 with head dim 32.mbconv_expandint= 4Expansion ratio of the inverted bottleneck inside MBConv blocks
(
mid = out_ch * expand). Defaults to 4.head_hidden_sizeint or None= 768If not
None, inserts an extra pre-classifier linear layer
of the given width with a tanh activation between the pooled
feature and the final classifier. Defaults to 768
(matches the reference recipe).dropoutfloat= 0.0Dropout probability inside the classifier head. Defaults to
0.0.Attributes
model_typeClassVar[str]Constant string
"coatnet" used by the model registry.Notes
The canonical CoAtNet-0 (registered as coatnet_0) has
approximately 25.6M parameters and reaches 81.6% top-1 on
ImageNet-1k at 224x224 (Dai et al., 2021, Table 5).
Reference: Zihang Dai et al., "CoAtNet: Marrying Convolution and Attention for All Data Sizes", NeurIPS 2021, arXiv:2106.04803.
Examples
Build a CoAtNet-0 configuration with a 10-class head for CIFAR-10:
>>> from lucid.models.vision.coatnet import CoAtNetConfig
>>> cfg = CoAtNetConfig(num_classes=10)
>>> cfg.dims, cfg.blocks_per_stage, cfg.num_classes
((96, 192, 384, 768), (2, 3, 5, 2), 10)Used by 3
Constructors
1dunder
__init__
→None__init__(num_classes: int = 1000, in_channels: int = 3, image_size: int = 224, variant: str = 'coatnet_0', blocks_per_stage: tuple[int, ...] = (2, 3, 5, 2), dims: tuple[int, ...] = (96, 192, 384, 768), stem_width: int = 64, attn_heads: tuple[int, ...] = (12, 24), mbconv_expand: int = 4, head_hidden_size: int | None = 768, dropout: float = 0.0, mixed_s3: tuple[int, int, int] | None = None)