Skip to content

Repository files navigation

Documentation Documentation CI

Julia ports of timm (PyTorch Image Models, by Ross Wightman) backbones for Lux.jl, with pretrained weights loaded directly from HuggingFace Hub in .safetensors format. The name is an homage to the project we port from.

Status

Most of Luximm was written by AI agents driving the porting workflow encoded in skills/timm-to-lux/, with human review at each phase and the parity tests as the correctness backstop. The code is already being used in real projects, so the registered backbones work for forward inference with the released weights. That said: expect bugs and rough edges, especially around anything the parity tests do not exercise (custom training loops, mixed-precision paths, exotic input shapes). File issues and PRs.

For LLMs & Agents

This library ships agent skills. In Claude Code:

/plugin marketplace add csvance/Luximm.jl
/plugin install luximm

That installs timm-to-lux, the porting workflow behind this library (converting PyTorch timm models into numerically-equivalent Lux.jl implementations with parity-tested weight loading), and tachikoma-tui, a bootstrap guide for Tachikoma.jl TUI apps. Both live in skills/ as plain Markdown.

Available backbones

Family Variant prefix Weights Weight License Commercial Use
ResNet :resnet* 5 Apache 2.0
SE-ResNet :seresnet* 1 Apache 2.0
BiT ResNetV2 :resnetv2_*_bit_* 15 Apache 2.0
ConvNeXt :convnext_* 19 Apache 2.0
ConvNeXt (DINOv3) :convnext_* 4 DINOv3 License ⚠️
ConvNeXt V2 :convnextv2_* 26 CC BY-NC 4.0
VGG :vgg* 8 CC BY 4.0
ViT :vit_* 4 Apache 2.0
CoAtNet :coatnet_* 5 Apache 2.0

Basic usage

using Luximm, Lux, Random

# ResNet50 with the trained 1000-class ImageNet head.
# `create_pretrained` is family-agnostic; the symbol selects the family.
# It returns the model and a closure that loads the released weights
# into `(ps, st)` once you've run `Lux.setup`.
model, load = create_pretrained(:resnet50_a1_in1k)
ps, st = Lux.setup(Xoshiro(0), model)
ps, st = load(ps, st)
st = Lux.testmode(st)                     # BatchNorm/Dropout in eval mode

x = randn(Float32, 224, 224, 3, 1)
logits, _ = model(x, ps, st)              # (1000, 1)
top1 = argmax(vec(logits))                # ImageNet class index

create_model(variant; ...) (without weight loading) is also exported for from-scratch training. For the full walkthrough, including feature-extractor mode (num_classes = 0), single-channel inputs (in_chans = 1), and the HuggingFace cache layout, see the Getting Started docs page.

Composing with a pretrained backbone

Drop a feature-extractor backbone into your own @compact block and let the loader fill in just the backbone's subtree. The prefix = (:backbone,) tuple matches the slot name in the outer model, so load_backbone writes only into ps.backbone.* and st.backbone.*, leaving the head at its random initialization for downstream training:

using Luximm, Lux, NNlib, Random

backbone, load_backbone = create_pretrained(:resnet50_a1_in1k;
    num_classes = 0, prefix = (:backbone,))

model = @compact(
    backbone = backbone,
    head     = Dense(2048 => 10),   # custom 10-class head
) do x
    feats  = backbone(x)                                   # (7, 7, 2048, N)
    pooled = NNlib.meanpool(feats, size(feats)[1:2])       # (1, 1, 2048, N)
    head(reshape(pooled, size(pooled, 3), size(pooled, 4)))
end

ps, st = Lux.setup(Xoshiro(0), model)
ps, st = load_backbone(ps, st)
st = Lux.testmode(st)

x = randn(Float32, 224, 224, 3, 1)
logits, _ = model(x, ps, st)                               # (10, 1)

For multi-backbone composition and deeper nesting patterns, see the Getting Started docs page.

Feature pyramids for dense prediction

features_only = true mirrors timm's features_only=True: the forward returns a tuple of intermediate feature maps ordered by increasing reduction, which is what a UNet or FPN decoder consumes. feature_info is the tap table, so the decoder can be sized before it is built.

using Luximm, Lux, Random

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

model, load = create_pretrained(:resnet18_a1_in1k; features_only = true)
ps, st = Lux.setup(Xoshiro(0), model)
ps, st = load(ps, st)
st = Lux.testmode(st)

x = randn(Float32, 224, 224, 3, 1)
feats, _ = model(x, ps, st)                # 5-tuple, finest first
size(feats[1]), size(feats[end])           # (112,112,64,1), (7,7,512,1)

out_indices selects a subset of the taps (1-based and strictly increasing, so timm's out_indices=(1,2,3,4) is Luximm's (2,3,4,5)). A features-only model builds the same parameter tree as the plain num_classes = 0 extractor, so the released weights load into it unchanged.

Family Pyramid Reductions
ResNet 2, 4, 8, 16, 32
SE-ResNet 2, 4, 8, 16, 32
BiT ResNetV2 2, 4, 8, 16, 32
ConvNeXt 4, 8, 16, 32
ConvNeXt V2 4, 8, 16, 32
ViT patch (per block)
VGG n/a
CoAtNet n/a

The unsupported families raise an error explaining why. Three caveats are worth knowing, all matching timm: the ConvNeXt families have no reduction-2 tap (their patch stem strides by 4 in one convolution); BiT's reduction-32 tap is the raw pre-activation stage4 output rather than the final_norm-applied map the same model returns at num_classes = 0; and a ViT is single-scale — every tap sits at the patch-size reduction, one per encoder block, each the raw post-block output with the class token dropped and the patch tokens reshaped to a grid (timm applies no final LayerNorm to its intermediates either). timm's vit_* default out_indices = 3 (the last three blocks) is Luximm's out_indices = (depth-2, depth-1, depth); pass nothing to get every block. See the Getting Started page for details.

License and attribution

Luximm.jl is licensed under the Apache License, Version 2.0 (see LICENSE and NOTICE). The license matches upstream timm. The Julia code in this repository is original, but layer naming, hyperparameters, padding ordering, and state_dict key layout are deliberately taken from timm so pretrained weights load directly.

Acknowledgements

Thanks to Ross Wightman for timm, to HuggingFace for hosting the .safetensors weights that Luximm.jl loads at runtime, to the Julia ML ecosystem maintainers whose work makes a port like this plausible, and to Medical Metrics Inc. for allowing me to work on and open-source the project.

About

Lux.jl Image Models

Topics

Resources

Stars

12 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages