# wkcn/tinyvit architecture

> Architecture diagram of the GitHub repository wkcn/tinyvit, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-27). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/wkcn/tinyvit
- Repository: https://github.com/wkcn/tinyvit

## Overview

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`](https://github.com/wkcn/tinyvit/blob/main/main.py) (Train entry), [`data/build.py`](https://github.com/wkcn/tinyvit/blob/main/data/build.py) (Loader builder), [`inference.py`](https://github.com/wkcn/tinyvit/blob/main/inference.py) (Inference script), [`save_logits.py`](https://github.com/wkcn/tinyvit/blob/main/save_logits.py) (Save teacher logits)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_training["Training and evaluation"]
  node_main["Train entry<br/>[main.py]"]
  node_engine["Train and validate<br/>[main.py]"]
  node_optimizer["Optimizer<br/>[optimizer.py]"]
  node_scheduler["LR scheduler<br/>[lr_scheduler.py]"]
  node_logger["Run logging<br/>[logger.py]"]
end

subgraph group_data["Image data pipeline"]
  node_loader["Loader builder<br/>[build.py]"]
  node_dataset["Dataset sources"]
  node_transforms["Image transforms"]
end

subgraph group_models["Models and inference"]
  node_clip["CLIP teacher<br/>[clip.py]"]
  node_model_builder["Model builder<br/>[build.py]"]
  node_tinyvit["TinyViT<br/>[tiny_vit.py]"]
  node_inference["Inference script<br/>[inference.py]"]
  node_classnames["ImageNet labels"]
end

subgraph group_distillation["Distillation logits"]
  node_wrapper["Logits wrapper<br/>[build.py]"]
  node_logit_manager["Logit manager<br/>[manager.py]"]
  node_logit_store[("Sparse logits")]
  node_save_logits["Save teacher logits<br/>[save_logits.py]"]
end

node_user(("Researcher"))
node_wandb(("Weights and Biases"))
node_image_user(("Image input"))
node_checkpoint["Model checkpoints"]

node_user -->|"starts training"| node_main
node_main -->|"builds loaders"| node_loader
node_loader -->|"loads samples"| node_dataset
node_loader -->|"builds transforms"| node_transforms
node_loader -.->|"wraps training data"| node_wrapper
node_wrapper -->|"reads or writes"| node_logit_manager
node_logit_manager -->|"accesses logits"| node_logit_store
node_main -->|"builds model"| node_model_builder
node_model_builder -->|"constructs"| node_tinyvit
node_main -->|"runs epochs"| node_engine
node_engine -->|"updates parameters"| node_optimizer
node_engine -->|"steps schedule"| node_scheduler
node_main -->|"records run"| node_logger
node_logger -.->|"logs metrics"| node_wandb
node_main -->|"saves model"| node_checkpoint
node_user -.->|"starts logit export"| node_save_logits
node_save_logits -.->|"runs teacher"| node_clip
node_save_logits -.->|"writes logits"| node_logit_store
node_image_user -->|"supplies image"| node_inference
node_inference -->|"preprocesses image"| node_transforms
node_inference -->|"classifies image"| node_tinyvit
node_inference -->|"maps class indices"| node_classnames

click node_main "https://github.com/wkcn/tinyvit/blob/main/main.py"
click node_engine "https://github.com/wkcn/tinyvit/blob/main/main.py"
click node_optimizer "https://github.com/wkcn/tinyvit/blob/main/optimizer.py"
click node_scheduler "https://github.com/wkcn/tinyvit/blob/main/lr_scheduler.py"
click node_logger "https://github.com/wkcn/tinyvit/blob/main/logger.py"
click node_loader "https://github.com/wkcn/tinyvit/blob/main/data/build.py"
click node_dataset "https://github.com/wkcn/tinyvit/tree/main/data/augmentation/parsers"
click node_transforms "https://github.com/wkcn/tinyvit/blob/main/data/augmentation/transforms_factory.py"
click node_wrapper "https://github.com/wkcn/tinyvit/blob/main/data/build.py"
click node_logit_manager "https://github.com/wkcn/tinyvit/blob/main/data/augmentation/manager.py"
click node_save_logits "https://github.com/wkcn/tinyvit/blob/main/save_logits.py"
click node_clip "https://github.com/wkcn/tinyvit/blob/main/models/clip.py"
click node_model_builder "https://github.com/wkcn/tinyvit/blob/main/models/build.py"
click node_tinyvit "https://github.com/wkcn/tinyvit/blob/main/models/tiny_vit.py"
click node_inference "https://github.com/wkcn/tinyvit/blob/main/inference.py"
click node_classnames "https://github.com/wkcn/tinyvit/blob/main/data/imagenet_classnames.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_main,node_engine,node_optimizer,node_scheduler,node_logger toneBlue
class node_loader,node_dataset,node_transforms toneAmber
class node_clip,node_model_builder,node_tinyvit,node_inference,node_classnames toneMint
class node_wrapper,node_logit_manager,node_logit_store,node_save_logits toneRose
class node_user,node_wandb,node_image_user,node_checkpoint toneIndigo
```

## Components

### Training and evaluation

- [Train entry](https://github.com/wkcn/tinyvit/blob/main/main.py): `main.py`
- [Train and validate](https://github.com/wkcn/tinyvit/blob/main/main.py): `main.py`
- [Optimizer](https://github.com/wkcn/tinyvit/blob/main/optimizer.py): `optimizer.py`
- [LR scheduler](https://github.com/wkcn/tinyvit/blob/main/lr_scheduler.py): `lr_scheduler.py`
- [Run logging](https://github.com/wkcn/tinyvit/blob/main/logger.py): `logger.py`

### Image data pipeline

- [Loader builder](https://github.com/wkcn/tinyvit/blob/main/data/build.py): `data/build.py`
- [Dataset sources](https://github.com/wkcn/tinyvit/tree/main/data/augmentation/parsers): `data/augmentation/parsers`
- [Image transforms](https://github.com/wkcn/tinyvit/blob/main/data/augmentation/transforms_factory.py): `data/augmentation/transforms_factory.py`

### Models and inference

- [CLIP teacher](https://github.com/wkcn/tinyvit/blob/main/models/clip.py): `models/clip.py`
- [Model builder](https://github.com/wkcn/tinyvit/blob/main/models/build.py): `models/build.py`
- [TinyViT](https://github.com/wkcn/tinyvit/blob/main/models/tiny_vit.py): `models/tiny_vit.py`
- [Inference script](https://github.com/wkcn/tinyvit/blob/main/inference.py): `inference.py`
- [ImageNet labels](https://github.com/wkcn/tinyvit/blob/main/data/imagenet_classnames.py): `data/imagenet_classnames.py`

### Distillation logits

- [Logits wrapper](https://github.com/wkcn/tinyvit/blob/main/data/build.py): `data/build.py`
- [Logit manager](https://github.com/wkcn/tinyvit/blob/main/data/augmentation/manager.py): `data/augmentation/manager.py`
- **Sparse logits**
- [Save teacher logits](https://github.com/wkcn/tinyvit/blob/main/save_logits.py): `save_logits.py`

### Other components

- **Researcher**
- **Weights and Biases**
- **Image input**
- **Model checkpoints**

## Connections

- 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