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.

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