Models

create_pretrained(variant; in_chans=3, num_classes=nothing,
                  revision="main", cache_dir=hf_hub_cache_dir(),
                  prefix=()) -> (model, load)

Family-agnostic pretrained-weight entry point, mirroring timm.create_model(..., pretrained=True). Returns the model and a closure that loads the released model.safetensors into a (ps, st) pair the caller produced with Lux.setup. The closure captures variant, in_chans, num_classes, and the HF / prefix kwargs at construction time, so calling it is the only place (ps, st) need to be threaded.

model, load = create_pretrained(:resnet50_a1_in1k)
ps, st = Lux.setup(Xoshiro(0), model)
ps, st = load(ps, st)

num_classes = nothing (the default) builds the head the released checkpoint ships with — default_num_classes(variant). Pass an explicit 0 for a features-only model, or any other Int to swap in a custom-width head (the released classifier is then skipped and the warning case fires).

For composition, build model separately and pass it into an outer @compact, capturing prefix = (:backbone,) so the closure writes into the right subtree:

backbone, load_backbone = create_pretrained(:resnet50_a1_in1k;
    num_classes = 0, prefix = (:backbone,))
outer = @compact(backbone = backbone,
    head = Dense(2048 => num_outputs)) do x
    head(backbone(x))
end
ps, st = Lux.setup(rng, outer)
ps, st = load_backbone(ps, st)

create_model(variant; kwargs...) -> model

Family-agnostic random-init model constructor, mirroring timm.create_model(..., pretrained=False). Dispatches on variant to the matching family constructor and returns the bare @compact model — no parameters, no state, no pretrained weights.

Use this when you want to train from scratch, or as a building block inside an outer @compact when composing a larger model. To load the released weights for a variant, use create_pretrained instead.

model = create_model(:resnet50_a1_in1k; num_classes = 1000)
ps, st = Lux.setup(rng, model)        # random init, ready for training

kwargs are forwarded to the family constructor (in_chans, num_classes).

default_num_classes(variant) -> Int

Head dimension the released checkpoint for variant was trained at. Returns 0 for encoder-only variants (DINOv3 ConvNeXt, ConvNeXtV2 fcmae pretrains).

ResNetVariant

Architectural config for a classic timm ResNet variant.

Fields:

• name: lookup key (e.g. :resnet50_a1_in1k).
• block: residual block type, either :basic (used by r18/r34) or :bottleneck (used by r50/r101/r152).
• layers: per-stage block count (d1, d2, d3, d4).
• planes: base channel widths per stage (64, 128, 256, 512). Multiplied by 4 inside :bottleneck stages to give the actual output channel count.
• num_features: backbone output channels (planes[end] for :basic, planes[end] * 4 for :bottleneck).
• hf_repo: HuggingFace repo containing model.safetensors.
• default_num_classes: head dimension the released weights ship with.
• default_input_size: native training resolution (224 for every registered variant). Informational only: the model is fully convolutional and accepts any size.

RESNET_VARIANTS :: Dict{Symbol, ResNetVariant}

Lookup table for classic ResNet variants currently ported from timm. Keys are the timm model names with dots rewritten as underscores.

SEResNetVariant

Architectural config for a single SE-ResNet variant.

Fields:

• name: lookup key (e.g. :seresnet50_a1_in1k).
• layers: per-stage block count (d1, d2, d3, d4).
• planes: base channel widths per stage (64, 128, 256, 512). Multiplied by 4 (the bottleneck expansion) to give the actual stage output channels.
• num_features: backbone output channels (planes[end] * 4 = 2048).
• hf_repo: HuggingFace repo containing model.safetensors.
• default_num_classes: head dimension the released weights ship with.
• default_input_size: native training resolution (224). Informational only.
• se_reduction: SE bottleneck reduction divisor (16 for every variant); the SE inner width is se_make_divisible(out_ch / se_reduction, 8).

SERESNET_VARIANTS :: Dict{Symbol, SEResNetVariant}

Lookup table for the SE-ResNet variants ported from timm. Keys are the timm model name with dots rewritten as underscores.

BiTVariant

Architectural config for a single BiT ResNetV2 variant.

Fields:

• name: lookup key (e.g. :resnetv2_50x1_bit_goog_in21k).
• layers: per-stage depth tuple (3,4,6,3) for r50, (3,4,23,3) for r101, (3,8,36,3) for r152.
• width_factor: integer width multiplier from the timm name suffix (x1, x2, x3, x4).
• stem_chs: stem output channels (64 * width_factor).
• stage_chs: per-stage output channel tuple (base widths (256,512,1024,2048) scaled by width_factor).
• num_features: backbone output channels (stage_chs[end]).
• hf_repo: HuggingFace repo containing model.safetensors.
• default_num_classes: head dimension the released weights were trained with (21843 for goog_in21k, 1000 for the in1k tags).
• default_input_size: native training resolution (224 for most tags, 384 for the _384 teacher variant). The model itself is fully convolutional and accepts any input size; this is just what the released weights were tuned at.

BIT_VARIANTS :: Dict{Symbol, BiTVariant}

Lookup table for the BiT variants this package currently ports. Keys mirror the timm model name with the dot rewritten as an underscore (the dot is reserved in Julia identifiers); the full timm name with the dot lives at BIT_VARIANTS[key].hf_repo.

ConvNeXtVariant

Architectural config for a single ConvNeXt v1 variant.

Fields:

• name: lookup key (e.g. :convnext_tiny_dinov3_lvd1689m).
• depths: per-stage block count, (d1, d2, d3, d4).
• dims: per-stage channel widths, (c1, c2, c3, c4). c1 is also the stem output channels. c4 is num_features.
• hf_repo: HuggingFace repo containing model.safetensors.
• default_num_classes: head dimension the released weights ship with. 0 for the DINO encoders (no usable head).
• default_input_size: native training resolution (224, 384, …) for the released checkpoint. Informational only: the model is fully convolutional and accepts any size, so this is not enforced.
• ls_init: LayerScale init value (gamma parameter in timm). All v1 variants released so far use 1e-6; kept as a field in case future ports need a different value.

CONVNEXT_VARIANTS :: Dict{Symbol, ConvNeXtVariant}

Lookup table for the ConvNeXt v1 variants this package ports: the DINOv3 encoders and the Facebook AI checkpoints from the original ConvNeXt paper. Additional convnext_* lineages (.in12k_*, .clip_*) can be registered without touching the constructor or mapping code.

ConvNeXtV2Variant

Architectural config for a single ConvNeXtV2 variant.

Fields:

• name: lookup key (e.g. :convnextv2_atto_fcmae).
• depths: per-stage block count, (d1, d2, d3, d4).
• dims: per-stage channel widths, (c1, c2, c3, c4). c1 is also the stem output channels. c4 is num_features.
• hf_repo: HuggingFace repo containing model.safetensors.
• default_num_classes: head dimension the released weights ship with. 0 for the bare .fcmae encoders, 1000 for the ImageNet-1K and ImageNet-22k-then-1K fine-tunes.
• default_input_size: native training resolution (224, 384, or 512) for the released checkpoint. Informational only: the model is fully convolutional and accepts any size, so this is not enforced.

CONVNEXTV2_VARIANTS :: Dict{Symbol, ConvNeXtV2Variant}

Lookup table for the ConvNeXtV2 variants this package ports. The .fcmae rows are the bare encoders; all other rows ship a 1000-class ImageNet head. convnextv2_small is not included because timm only registers it as .untrained (no pretrained weights).

VGGVariant

Architectural config for a single VGG variant.

Fields:

• name: lookup key (e.g. :vgg16_tv_in1k).
• cfg: flat layer list. Each entry is either an Int (a 3x3 pad-1 conv with that output-channel count, followed by ReLU) or the Symbol :M (a 2x2 stride-2 max-pool). Matches timm's cfgs table in vgg.py.
• batch_norm: whether a BatchNorm sits between every conv and its ReLU (the *_bn checkpoints).
• hf_repo: HuggingFace repo containing model.safetensors.
• default_num_classes: head dimension the released weights ship with (1000 for every registered variant).
• default_input_size: native training resolution (224). Informational only: the model accepts any size large enough for the 7x7 pre_logits conv, but the released head was trained at 224.

VGG_VARIANTS :: Dict{Symbol, VGGVariant}

Lookup table for the VGG variants ported from timm: the four classic depths (11/13/16/19) in plain and BatchNorm flavors, all torchvision tv_in1k checkpoints. Keys are the timm model name with the dot rewritten as an underscore.

ViTVariant

Architectural config for a single Vision Transformer variant.

Fields:

• name: lookup key (e.g. :vit_base_patch16_224_augreg2_in21k_ft_in1k).
• depth: number of transformer encoder blocks.
• embed_dim: token / channel width.
• num_heads: attention heads (head_dim = embed_dim ÷ num_heads).
• patch: patch side length (16).
• img_size: native input resolution the position embedding was trained at. Enforced by the constructor, since absolute pos-embed has no interpolation path yet.
• hf_repo: HuggingFace repo containing model.safetensors.
• default_num_classes: head dimension the released weights ship with.
• default_input_size: native training resolution (== img_size).

VIT_VARIANTS :: Dict{Symbol, ViTVariant}

Lookup table for the Vision Transformer variants ported from timm. Keys are the timm model name with dots rewritten as underscores.

CoAtNetVariant

Architectural config for a single CoAtNet variant.

Fields:

• name: lookup key (e.g. :coatnet_0_rw_224_sw_in1k).
• depths: per-stage block count (d1, d2, d3, d4).
• dims: per-stage output channel widths (c1, c2, c3, c4); c4 is num_features.
• stem_width: the two stem conv widths (s1, s2); s2 feeds stage 1.
• block_types: per-stage block kind, :C (MBConv) or :T (transformer).
• stride_mode: how MBConv blocks downsample — :pool (avg-pool the main path, stride-1 convs) or :dw (stride the depthwise conv, no pool).
• attn_early: MBConv SE placement — true puts SE between the depthwise conv and norm2 (timm se_early), false after norm2 (timm se).
• se_act: MBConv SE bottleneck activation, :relu or :silu.
• transformer_shortcut_bias: whether the transformer downsample shortcut's 1x1 expand conv carries a bias.
• layer_scale: whether transformer blocks apply LayerScale (ls1/ls2 per-channel gamma) to the attention and MLP residual branches.
• img_size: native input resolution (enforced; the transformer relative-position bias is sized to the per-stage feature map).
• hf_repo: HuggingFace repo containing model.safetensors.
• default_num_classes: head dimension the released weights ship with.
• default_input_size: native training resolution (== img_size).

COATNET_VARIANTS :: Dict{Symbol, CoAtNetVariant}

Lookup table for the CoAtNet variants ported from timm. Keys are the timm model name with dots rewritten as underscores.