How mit-han-lab/tiny-training works
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: train_cls.py Training entry · mobilenet_v2.py Quantized OFA models · cls_trainer.py Classification trainer · dataset_entry.py Vision data pipeline
Components 17
PyTorch Training
Configuration-driven quantized vision training.
- Training entry · Selects and runs configured experiments.
algorithm/train_cls.py - Experiment configs · Default, transfer-learning, and VWW experiment definitions.
algorithm/configs/default.yaml - QAS and sparse scripts · Launch comparison and sparse-update variants.
algorithm/scripts/compare_qas.sh - Vision data pipeline · Builds image-folder datasets and transforms.
algorithm/core/dataset/dataset_entry.py - Model dispatch · Selects OFA-derived network families and assets.
algorithm/core/model/model_entry.py - Quantized OFA models · MobileNetV2 and Proxyless model families.
algorithm/core/ofa_nn/networks/mobilenet_v2.py - MCU model asset
algorithm/assets/mcu_models/mbv2-w0.35.pkl - Quantized operators and QAS · Implements quantized behavior and scale-aware gradients.
algorithm/quantize/quantized_ops.py - Classification trainer · Owns forward, backward, optimizer, and scheduler execution.
algorithm/core/trainer/cls_trainer.py - Sparse partial backward · 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 · Starts PyTorch/MCUNetV3-to-training-IR compilation.
compilation/mcu_ir_gen.py - PyTorch-to-IR conversion · Wraps and loads MCUNetV3 models into internal IR.
compilation/convert/pth2ir.py - Training IR · Deployment-oriented representation of model training.
- Compile-time autodiff · Maps forward IR ops to MCU and int8 gradient operations.
compilation/autodiff/auto_diff.py - Sparse graph transform · Prunes and reorders backward, storage, and update work.
compilation/mod.py - Training graph JSON · Serialized pruned training graph.
compilation/ir2json.py
MCU Deployment
External TinyEngine execution environment.
- TinyEngine backend · Generates, compiles, and deploys MCU training code.
Connections 18
- 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