Initial NaFlex ViT model and training support#2466

Draft

rwightman wants to merge 16 commits intomainfrom

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