Models

create_pretrained(variant; in_chans=3, num_classes=nothing,
                  features_only=false, out_indices=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)

features_only = true builds the feature-pyramid model described in create_model and loads the released backbone weights into it. The parameter tree is identical to the num_classes = 0 feature extractor, so this is the same load path; only the forward differs:

backbone, load = create_pretrained(:resnet18_a1_in1k;
    in_chans = 1, features_only = true, prefix = (:backbone,))
info = feature_info(:resnet18_a1_in1k)   # decoder widths

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, features_only, out_indices).

Feature-pyramid mode

features_only = true is the Luximm analog of timm's features_only=True: the forward returns a tuple of intermediate feature maps ordered by increasing reduction (highest resolution first) instead of a single map. This is what a UNet or FPN decoder consumes.

model = create_model(:resnet18_a1_in1k; features_only = true)
ps, st = Lux.setup(rng, model)
feats, st = model(x, ps, st)          # NTuple{5, Array{Float32, 4}}
size.(feats, 3)                       # (64, 64, 128, 256, 512)

out_indices selects a subset of the family's taps. It is 1-based, so timm's out_indices=(1, 2, 3, 4) is Luximm's (2, 3, 4, 5). Indices must be strictly increasing; nothing (the default) returns every tap. Query the tap table with feature_info:

feature_info(:resnet18_a1_in1k).reductions          # (2, 4, 8, 16, 32)
model = create_model(:resnet18_a1_in1k;
                     features_only = true, out_indices = (2, 3, 4, 5))

A features-only model always has num_classes = 0 (passing anything else is an error) and builds the same parameter tree as the plain num_classes = 0 feature extractor, so create_pretrained loads released weights into it unchanged. Selecting fewer taps changes only the forward, never the tree: every stage is still built and still runs.

Supported families: ResNet, SE-ResNet, BiT ResNetV2, ConvNeXt, ConvNeXt V2, ViT. VGG and CoAtNet raise an error explaining why.

For the CNN families the taps form a true pyramid (increasing reduction). A ViT is single-scale, so its taps all sit at the patch-size reduction; they mirror timm's features_only=True for vit_* exactly — each tap is a selected block's raw output with the class token dropped and the patch tokens reshaped to a (W, H, embed_dim, N) grid (timm applies no final LayerNorm to the intermediates, and neither do we). timm's vit_* default out_indices = 3 (the last three blocks) is Luximm's out_indices = (depth-2, depth-1, depth).

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).

feature_info(variant; out_indices=nothing) -> FeatureInfo

Tap table for variant in feature-pyramid mode: the reduction (spatial stride) and channel count of each feature map a create_model(variant; features_only = true) model returns, ordered by increasing reduction. The Luximm analog of timm's model.feature_info.

Call this to size a decoder before building it: a UNet needs the skip channel counts, an FPN needs them to size its lateral 1×1 convolutions.

info = feature_info(:resnet18_a1_in1k)
info.reductions              # (2, 4, 8, 16, 32)
info.channels                # (64, 64, 128, 256, 512)

info = feature_info(:resnet18_a1_in1k; out_indices = (2, 3, 4, 5))
info.channels                # (64, 128, 256, 512)

out_indices is validated exactly as create_model validates it, so the returned info always describes the tuple that model's forward produces. Families without a pyramid (VGG, CoAtNet) raise an error.

FeatureInfo{K}

Metadata for the K feature maps a features-only model returns, ordered by increasing reduction (highest resolution first). The Luximm analog of timm's model.feature_info.

Fields:

• names: the tap's module name in timm, e.g. (:act1, :layer1, …). Purely informational; useful when cross-checking against a timm model.
• reductions: spatial stride of each map relative to the input, e.g. (2, 4, 8, 16, 32). A map at reduction r has spatial size (W ÷ r, H ÷ r).
• channels: channel count of each map, i.e. size(feats[i], 3). This is what a UNet/FPN decoder needs to size its skip projections.
• indices: which entries of the family's full tap list these are. For an unfiltered feature_info(variant) this is (1, 2, …, K); after selecting with out_indices it echoes that selection back.

info = feature_info(:resnet18_a1_in1k)
info.reductions   # (2, 4, 8, 16, 32)
info.channels     # (64, 64, 128, 256, 512)

info = feature_info(:resnet18_a1_in1k; out_indices = (2, 3, 4, 5))
info.channels     # (64, 128, 256, 512)
info.indices      # (2, 3, 4, 5)

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).
• pre_norm: whether a norm_pre LayerNorm runs after pos_embed and before the encoder blocks (CLIP image towers; timm pre_norm=True).
• norm_eps: LayerNorm epsilon for every norm in the model — 1f-6 for timm's ImageNet ViTs, 1f-5 for the CLIP towers.
• stem_bias: whether the patch-embed conv has a bias (true for timm's ImageNet ViTs; false for CLIP, whose patch_embed.proj is bias-free).

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.