How wkcn/tinyvit works
TinyViT is a PyTorch image-classification model and training codebase centered on training small vision transformers, including distillation from precomputed teacher logits. The main workflow builds image datasets and loaders, optionally wraps training samples with stored logits, trains and validates a TinyViT, and saves checkpoints; a separate script generates teacher logits, while inference transforms an image and returns ranked class names. The graph reflects README capabilities and the supplied excerpts; some training and logits-generation wiring is documented but not visible in the sampled source.
Main entry points: main.py Train entry · build.py Loader builder · inference.py Inference script · save_logits.py Save teacher logits
Components 21
Training and evaluation
- Train entry
main.py - Train and validate
main.py - Optimizer
optimizer.py - LR scheduler
lr_scheduler.py - Run logging
logger.py
Image data pipeline
- Loader builder
data/build.py - Dataset sources
data/augmentation/parsers - Image transforms
data/augmentation/transforms_factory.py
Models and inference
- CLIP teacher
models/clip.py - Model builder
models/build.py - TinyViT
models/tiny_vit.py - Inference script
inference.py - ImageNet labels
data/imagenet_classnames.py
Distillation logits
- Logits wrapper
data/build.py - Logit manager
data/augmentation/manager.py - Sparse logits
- Save teacher logits
save_logits.py
Other components
- Researcher
- Weights and Biases
- Image input
- Model checkpoints
Connections 22
- Researcher → Train entry: starts training
- Train entry → Loader builder: builds loaders
- Loader builder → Dataset sources: loads samples
- Loader builder → Image transforms: builds transforms
- Loader builder → Logits wrapper: wraps training data
- Logits wrapper → Logit manager: reads or writes
- Logit manager → Sparse logits: accesses logits
- Train entry → Model builder: builds model
- Model builder → TinyViT: constructs
- Train entry → Train and validate: runs epochs
- Train and validate → Optimizer: updates parameters
- Train and validate → LR scheduler: steps schedule
- Train entry → Run logging: records run
- Run logging → Weights and Biases: logs metrics
- Train entry → Model checkpoints: saves model
- Researcher → Save teacher logits: starts logit export
- Save teacher logits → CLIP teacher: runs teacher
- Save teacher logits → Sparse logits: writes logits
- Image input → Inference script: supplies image
- Inference script → Image transforms: preprocesses image
- Inference script → TinyViT: classifies image
- Inference script → ImageNet labels: maps class indices