Initial NaFlex ViT model and training support#2466

Merged

rwightman merged 34 commits intomainfrom

Jun 5, 2025

Commits on Apr 9, 2025

Add loss scale arg, initial distributed loss scale. Maybe fix FX for the model.
rwightman
committed
Exclude embeds module and mask attn functions from tracing
rwightman
committed

Commits on Apr 10, 2025

A much faster resample_patch_embed, can be used at train/validation time
rwightman
committed

Commits on Apr 14, 2025

Improve several typing issues for flex vit, can (almost) work with jit if we bash h,w key into an int or str
rwightman
committed

Commits on Apr 21, 2025

Optimizations for pos embed resize, merge different mask helper fns
rwightman
committed

Commits on Apr 25, 2025

Add naflex loader support to validate.py, fix bug in naflex pos embed add, classic vit weight loading for naflex model
rwightman
committed

Commits on Apr 28, 2025

Commits on Apr 29, 2025

Move NaFlexCollate with dataset, remove stand alone collate_fn and remove redundancy
rwightman
committed

Commits on May 10, 2025

Add a WIP NaFlex compatible mixup/cutmix for testing
rwightman
committed

Commits on May 20, 2025

Mixup cleanup, add prob support and train script integration. Add working loader based patch compatible RandomErasing for NaFlex mode.
rwightman
committed

Commits on May 21, 2025

NaFlex random erasing performance improvements, python loops were slow. Remove subregion mode, not going to be worth it.
rwightman
committed

Commits on May 23, 2025

Merge remote-tracking branch 'origin/main' into naflex
rwightman
committed

Commits on May 24, 2025

Commits on May 25, 2025

Commits on May 30, 2025

Add siglip2 compatible naflex encoders. Add support to factorized pos embeds and 'aspect preserving mode' to Flex Embeds. Some more docstrings and typing.
rwightman
committed

Commits on Jun 3, 2025

Commits on Jun 4, 2025