# mit-han-lab/tiny-training architecture

> Architecture diagram of the GitHub repository mit-han-lab/tiny-training, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-08-14). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/mit-han-lab/tiny-training
- Repository: https://github.com/mit-han-lab/tiny-training

## Overview

#### Purpose and top-level flow
This repository implements the algorithm/compiler portion of MCUNetV3 on-device training under 256 KB. It trains quantized vision models with quantization-aware scaling (QAS) and sparse updates, then compiles selected models into a pruned training graph JSON for TinyEngine MCU deployment. The MCU backend itself is external.

#### Algorithm entry point and experiment control
`algorithm/train_cls.py` is the primary training entry point. YAML experiment definitions select datasets, models, optimization, transfer learning, and visual wake words workflows:
- `algorithm/configs/default.yaml`
- `algorithm/configs/transfer.yaml`
- `algorithm/configs/vww.yaml`

Shell entry points drive the two key experimental variants:
- `algorithm/scripts/compare_qas.sh`
- `algorithm/scripts/sparse_update.sh`

#### Training assembly layer
The algorithm package constructs the PyTorch training pipeline from configuration: dataset dispatch, model dispatch, optimizer/scheduler setup, distributed utilities, and classification training lifecycle. The principal control boundary is the trainer, which owns forward/backward/update execution and applies sparse/partial-backward behavior:
- `algorithm/core/trainer/base_trainer.py`
- `algorithm/core/trainer/cls_trainer.py`
- `algorithm/core/utils/partial_backward.py`

Data enters through image-folder datasets and transforms, then batches are passed to the selected model/trainer:
- `algorithm/core/dataset/dataset_entry.py`
- `algorithm/core/dataset/vision/image_folder.py`
- `algorithm/core/dataset/vision/transform/transform.py`

#### Models and quantized training
Model selection exposes OFA-derived MobileNetV2 and Proxyless network families, with pretrained MCU model assets used by experiments:
- `algorithm/core/model/model_entry.py`
- `algorithm/core/ofa_nn/networks/mobilenet_v2.py`
- `algorithm/core/ofa_nn/networks/proxyless_nets.py`
- `algorithm/assets/mcu_models/mbv2-w0.35.pkl`

Custom quantized operators/formats implement real quantized graph behavior rather than fake-quantization-only training. QAS adjusts gradients to account for quantization scale; sparse update limits gradient computation and parameter updates to chosen layers/tensors:
- `algorithm/quantize/quantize_helper.py`
- `algorithm/quantize/quantized_ops.py`
- `algorithm/quantize/custom_quantized_format.py`

#### Compiler entry and PyTorch-to-IR boundary
The compilation pipeline consumes a PyTorch/MCUNetV3 model and produces an internal graph representation. `compilation/mcu_ir_gen.py` is the high-level compiler entry; conversion wrappers/loaders bridge model definitions and checkpoint semantics into IR:
- `compilation/mcu_ir_gen.py`
- `compilation/convert/pth2ir.py`
- `compilation/convert/mcunetv3_wrapper.py`
- `compilation/convert/load_mcunetv3.py`

This is the boundary between Python/PyTorch model execution and deployment-oriented graph processing.

#### Compile-time autodiff and sparse graph transformation
Rather than retaining runtime autograd on the MCU, the compiler generates backward operations at compile time. It maps forward IR operations to gradient implementations, including MCU-specific operations and int8-gradient support:
- `compilation/autodiff/auto_diff.py`
- `compilation/autodiff/op2grad.py`
- `compilation/autodiff/mcuop.py`
- `compilation/autodiff/int8_grad.py`

Graph modification/reordering applies the sparse-update policy, reducing required backward computation, activation storage, and update work:
- `compilation/mod.py`
- `compilation/ir_utils/general_func.py`

#### Deployment artifact and external runtime
The final training graph is serialized as JSON for the external TinyEngine/Tiny Training Engine backend, which handles MCU code generation, kernels, compilation, and device deployment:
- `compilation/ir2json.py`
- `compilation/serialize.py`
- `compilation/utils.py`

Material runtimes are Python/PyTorch for training and compilation, plus the external TinyEngine MCU runtime/codegen environment. No hosted service or web application is part of this repository.

Main entry points: [`algorithm/train_cls.py`](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/train_cls.py) (Training entry), [`algorithm/core/ofa_nn/networks/mobilenet_v2.py`](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/core/ofa_nn/networks/mobilenet_v2.py) (Quantized OFA models), [`algorithm/core/trainer/cls_trainer.py`](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/core/trainer/cls_trainer.py) (Classification trainer), [`algorithm/core/dataset/dataset_entry.py`](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/core/dataset/dataset_entry.py) (Vision data pipeline)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_training["PyTorch Training"]
  node_train_entry["Training entry<br/>Python CLI<br/>[train_cls.py]"]
  node_experiment_configs["Experiment configs<br/>YAML configs<br/>[default.yaml]"]
  node_experiment_scripts["QAS and sparse scripts<br/>Shell entrypoints<br/>[compare_qas.sh]"]
  node_dataset_pipeline["Vision data pipeline<br/>Dataset dispatch<br/>[dataset_entry.py]"]
  node_model_dispatch["Model dispatch<br/>Model factory<br/>[model_entry.py]"]
  node_quantized_models["Quantized OFA models<br/>PyTorch models<br/>[mobilenet_v2.py]"]
  node_mcu_checkpoint["MCU model asset<br/>Pretrained checkpoint<br/>[mbv2-w0.35.pkl]"]
  node_quantization["Quantized operators and QAS<br/>Quantization layer<br/>[quantized_ops.py]"]
  node_classification_trainer["Classification trainer<br/>Training lifecycle<br/>[cls_trainer.py]"]
  node_partial_backward["Sparse partial backward<br/>Gradient policy"]
end

subgraph group_compiler["Training Graph Compiler"]
  node_compiler_entry["IR generation entry<br/>Python compiler CLI<br/>[mcu_ir_gen.py]"]
  node_model_to_ir["PyTorch-to-IR conversion<br/>Conversion boundary<br/>[pth2ir.py]"]
  node_training_ir{{"Training IR<br/>Internal graph"}}
  node_compile_time_autodiff["Compile-time autodiff<br/>Backward graph generation<br/>[auto_diff.py]"]
  node_sparse_graph_transform["Sparse graph transform<br/>IR optimization<br/>[mod.py]"]
  node_training_graph_json["Training graph JSON<br/>Deployment artifact<br/>[ir2json.py]"]
end

subgraph group_deployment["MCU Deployment"]
  node_tinyengine["TinyEngine backend<br/>External MCU runtime"]
end

node_experiment_configs -->|"configure"| node_train_entry
node_experiment_scripts -->|"launch variants"| node_train_entry
node_train_entry -->|"build data"| node_dataset_pipeline
node_train_entry -->|"select model"| node_model_dispatch
node_train_entry -->|"run lifecycle"| node_classification_trainer
node_dataset_pipeline -->|"batches"| node_classification_trainer
node_model_dispatch -->|"instantiate"| node_quantized_models
node_mcu_checkpoint -->|"initialize"| node_quantized_models
node_quantized_models -->|"uses"| node_quantization
node_quantization -->|"quantized forward and QAS"| node_classification_trainer
node_partial_backward -->|"constrains backward"| node_classification_trainer
node_quantized_models -.->|"selected model"| node_compiler_entry
node_compiler_entry -->|"convert"| node_model_to_ir
node_model_to_ir -->|"emit forward IR"| node_training_ir
node_training_ir -->|"forward graph"| node_compile_time_autodiff
node_compile_time_autodiff -->|"backward IR"| node_sparse_graph_transform
node_sparse_graph_transform -->|"pruned graph"| node_training_graph_json
node_training_graph_json -->|"consume for codegen"| node_tinyengine

click node_train_entry "https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/train_cls.py"
click node_experiment_configs "https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/configs/default.yaml"
click node_experiment_scripts "https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/scripts/compare_qas.sh"
click node_dataset_pipeline "https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/core/dataset/dataset_entry.py"
click node_model_dispatch "https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/core/model/model_entry.py"
click node_quantized_models "https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/core/ofa_nn/networks/mobilenet_v2.py"
click node_mcu_checkpoint "https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/assets/mcu_models/mbv2-w0.35.pkl"
click node_quantization "https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/quantize/quantized_ops.py"
click node_classification_trainer "https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/core/trainer/cls_trainer.py"
click node_partial_backward "https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/core/utils/partial_backward.py"
click node_compiler_entry "https://github.com/mit-han-lab/tiny-training/blob/main/compilation/mcu_ir_gen.py"
click node_model_to_ir "https://github.com/mit-han-lab/tiny-training/blob/main/compilation/convert/pth2ir.py"
click node_compile_time_autodiff "https://github.com/mit-han-lab/tiny-training/blob/main/compilation/autodiff/auto_diff.py"
click node_sparse_graph_transform "https://github.com/mit-han-lab/tiny-training/blob/main/compilation/mod.py"
click node_training_graph_json "https://github.com/mit-han-lab/tiny-training/blob/main/compilation/ir2json.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_train_entry,node_experiment_configs,node_experiment_scripts,node_dataset_pipeline,node_model_dispatch,node_quantized_models,node_mcu_checkpoint,node_quantization,node_classification_trainer,node_partial_backward toneBlue
class node_compiler_entry,node_model_to_ir,node_training_ir,node_compile_time_autodiff,node_sparse_graph_transform,node_training_graph_json toneAmber
class node_tinyengine toneMint
```

## Components

### PyTorch Training

Configuration-driven quantized vision training.

- [Training entry](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/train_cls.py): Python CLI · Selects and runs configured experiments. · `algorithm/train_cls.py`
- [Experiment configs](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/configs/default.yaml): YAML configs · Default, transfer-learning, and VWW experiment definitions. · `algorithm/configs/default.yaml`
- [QAS and sparse scripts](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/scripts/compare_qas.sh): Shell entrypoints · Launch comparison and sparse-update variants. · `algorithm/scripts/compare_qas.sh`
- [Vision data pipeline](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/core/dataset/dataset_entry.py): Dataset dispatch · Builds image-folder datasets and transforms. · `algorithm/core/dataset/dataset_entry.py`
- [Model dispatch](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/core/model/model_entry.py): Model factory · Selects OFA-derived network families and assets. · `algorithm/core/model/model_entry.py`
- [Quantized OFA models](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/core/ofa_nn/networks/mobilenet_v2.py): PyTorch models · MobileNetV2 and Proxyless model families. · `algorithm/core/ofa_nn/networks/mobilenet_v2.py`
- [MCU model asset](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/assets/mcu_models/mbv2-w0.35.pkl): Pretrained checkpoint · `algorithm/assets/mcu_models/mbv2-w0.35.pkl`
- [Quantized operators and QAS](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/quantize/quantized_ops.py): Quantization layer · Implements quantized behavior and scale-aware gradients. · `algorithm/quantize/quantized_ops.py`
- [Classification trainer](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/core/trainer/cls_trainer.py): Training lifecycle · Owns forward, backward, optimizer, and scheduler execution. · `algorithm/core/trainer/cls_trainer.py`
- [Sparse partial backward](https://github.com/mit-han-lab/tiny-training/blob/main/algorithm/core/utils/partial_backward.py): Gradient policy · Restricts gradient computation and parameter updates. · `algorithm/core/utils/partial_backward.py`

### Training Graph Compiler

Converts MCUNetV3 models into MCU training IR.

- [IR generation entry](https://github.com/mit-han-lab/tiny-training/blob/main/compilation/mcu_ir_gen.py): Python compiler CLI · Starts PyTorch/MCUNetV3-to-training-IR compilation. · `compilation/mcu_ir_gen.py`
- [PyTorch-to-IR conversion](https://github.com/mit-han-lab/tiny-training/blob/main/compilation/convert/pth2ir.py): Conversion boundary · Wraps and loads MCUNetV3 models into internal IR. · `compilation/convert/pth2ir.py`
- **Training IR**: Internal graph · Deployment-oriented representation of model training.
- [Compile-time autodiff](https://github.com/mit-han-lab/tiny-training/blob/main/compilation/autodiff/auto_diff.py): Backward graph generation · Maps forward IR ops to MCU and int8 gradient operations. · `compilation/autodiff/auto_diff.py`
- [Sparse graph transform](https://github.com/mit-han-lab/tiny-training/blob/main/compilation/mod.py): IR optimization · Prunes and reorders backward, storage, and update work. · `compilation/mod.py`
- [Training graph JSON](https://github.com/mit-han-lab/tiny-training/blob/main/compilation/ir2json.py): Deployment artifact · Serialized pruned training graph. · `compilation/ir2json.py`

### MCU Deployment

External TinyEngine execution environment.

- **TinyEngine backend**: External MCU runtime · Generates, compiles, and deploys MCU training code.

## Connections

- Experiment configs → Training entry: configure
- QAS and sparse scripts → Training entry: launch variants
- Training entry → Vision data pipeline: build data
- Training entry → Model dispatch: select model
- Training entry → Classification trainer: run lifecycle
- Vision data pipeline → Classification trainer: batches
- Model dispatch → Quantized OFA models: instantiate
- MCU model asset → Quantized OFA models: initialize
- Quantized OFA models → Quantized operators and QAS: uses
- Quantized operators and QAS → Classification trainer: quantized forward and QAS
- Sparse partial backward → Classification trainer: constrains backward
- Quantized OFA models → IR generation entry: selected model
- IR generation entry → PyTorch-to-IR conversion: convert
- PyTorch-to-IR conversion → Training IR: emit forward IR
- Training IR → Compile-time autodiff: forward graph
- Compile-time autodiff → Sparse graph transform: backward IR
- Sparse graph transform → Training graph JSON: pruned graph
- Training graph JSON → TinyEngine backend: consume for codegen