Mini trainer is an extensible PyTorch framework for training, evaluating and
exporting image classifiers. It supports flat labels and hierarchical class
structures on your own datasets; biological taxonomy is one application.
The training framework is published on PyPI as mt-trainer.
Python imports remain mini_trainer.
- Flat and hierarchical classification: train classifiers and predict at multiple levels of a label hierarchy. Start with the MNIST and hierarchical Blair examples.
- Choice of image backbone: use Torchvision models, optional timm, Transformers or BioCLIP backends, and customize training components through builders.
- Training and data: image folders or supported metadata Parquet inputs, mixed precision, distributed training, checkpoint loading and configurable augmentation, optimizers and schedules. See the training API, distributed setup and runtime contracts and limitations.
- Evaluation and inspection: classification metrics, confusion matrices, hierarchy visualizations and optional Weights & Biases logging. The benchmark guide explains maintained evaluation coverage.
- Deployment: native PyTorch inference and optional ONNX export with class mappings and preprocessing metadata. See the separate quantization guide for opt-in x86 INT8 support and its limits.
These guides and examples serve as the documentation entry points while a standalone documentation site is pending.
Model packages contain a particular trained model and its integration interface; they are versioned separately from the general training framework. Additional models can cover other datasets, label sets and application domains.
| Model | Scope | Use it |
|---|---|---|
| MAMBO V3 | Moths and butterflies (Lepidoptera); species, genus and family predictions | PyPI package · Hugging Face weights/model card · Interactive demo · Release |
The MAMBO V3 deployment guide
covers its Python API and CLI, PyTorch/ONNX runtimes, geographic presets, custom
class lists, embeddings, optional test-time augmentation and measured quality/speed.
ONNX/CPU integration needs neither the training package nor a GPU. MAMBO weights
are CC BY-NC-SA 4.0 (non-commercial, share-alike); framework and adapter code are MIT.
The model guide is pinned to the released source; master continues framework
development independently of the release/mambo-v3 maintenance branch.
| Task | Start here |
|---|---|
| Integrate MAMBO V3 | Released deployment API and CLI |
| Prepare data or adapt an example | Examples |
| Change the trainer or run checks | Development guide, test map |
| Compare models, backends or training settings | Benchmarks |
| Run on UCloud or reproduce research | UCloud, research experiments |
| Choose the next development task | Roadmap |
Use uv for environment and package management. Choose a published package or a source checkout.
The published distribution is mt-trainer; Python imports remain mini_trainer.
The similarly named mini-trainer / mini_trainer PyPI project is unrelated.
uv venv --python 3.12
source .venv/bin/activate
uv pip install "mt-trainer[recommended]" --torch-backend=auto| Package choice | Includes |
|---|---|
mt-trainer |
Core training and inference |
mt-trainer[recommended] |
Core plus logging, visualization and optional utilities |
mt-trainer[all] |
Recommended extras plus notebooks, model backends and ONNX export |
Substitute the desired package in the install command. Standard pip install also
works; select its PyTorch CPU/CUDA installation separately for your environment.
Choose one backend: cpu, cu126, cu130 or cu132. The example selects CUDA 13.0;
change TORCH_BACKEND to match your intended environment before synchronizing.
git clone https://github.com/asgersvenning/mini_trainer.git
cd mini_trainer
TORCH_BACKEND=cu130
uv sync --extra recommended --extra "$TORCH_BACKEND"
source .venv/bin/activateReplace recommended with all for the additional backends/export tools above.
Activate the environment, use its executables directly, or use uv run --no-sync.
An implicit sync can replace the deliberately selected PyTorch backend. Select the
backend explicitly whenever installing or synchronizing dependencies.
Defaults use process CPU availability, affinity, visible cgroup quotas and Slurm
allocation limits. Shared resources may still need an explicit per-process budget.
Set --num_workers N for loading (0 runs in the main process), and
--cache-workers N for training-cache preparation. CUDA-cached datasets use zero
DataLoader workers. See automatic budgets
for caps and fallback behavior; cache readers and DataLoader workers are separate.
mini_trainer supports logging your training runs, including metrics, confusion matrices, and the probabilistic dendrogram, directly to Weights & Biases.
To use this feature you must install mini_trainer with the recommended extras. See Installation for more information.
- Login to your wandb account:
wandb login
- Train with the
--wandbflag: Simply append the--wandbflag to your training command.mt_train -i path/to/dataset --wandb
This repository draws inspiration from https://github.com/pytorch/vision/tree/main/references/classification.
Keep required dependencies and configuration small, provide useful defaults, and keep architectures, losses, data formats and training components extensible. Third-party integrations should remain optional.
Feel free to contribute, but here are a few tips:
- Follow the installation guide to setup a proper dev environment.
- Run
bash dev/check.sh staticfor lint, formatting, and architecture checks. - Run
bash dev/check.sh testfor the test suite; see the development guide for focused checks and environment details. - Run
bash dev/check-wheel.shto validate a minimal wheel installation in a disposable CPU environment. - Please avoid adding new dependencies 🙂
Repository agents should start with AGENTS.md. Planned improvements and their acceptance criteria are tracked in the roadmap. Remaining quantization work has a focused target-qualification roadmap.
Export trained models with mt_export --weights weights.pt --output exported-model
after installing the export extra and the relevant model backend. The generic
exporter preserves evaluation outputs and verifies ONNX Runtime parity. See the
export guide for the Python API, preprocessing contract and coverage.
Follow the benchmark results and coverage and continuous run history. The suite progresses from an exact synthetic oracle to MNIST and hierarchical Blair, with separate CPU and GPU profiles, visible summaries, and retained reproduction artifacts.
For configured GPU runners, the opt-in TensorRT deployment workflow rebuilds engines on the target and reports paired quality, latency and memory.
EMA (--ema / ema=True) is currently nonfunctional: classifier caches populated
by evaluation can break later EMA updates. Leave it disabled. Enabling it emits a
runtime warning; its API and checkpoint compatibility are retained, and repair is
deferred. See known limitations.
An opt-in PTQ and QAT Python API targets native x86 INT8 inference. This is an initial backend increment; CPU float32 QAT, integer inference and ordinary AMP are distinct capabilities.
Opt-in CUDA INT8 training supports Linear weights, integer forward/backward products, checkpoint restoration and CUDA inference. See the validation audit for measured memory, speed and loading benefits, supported configurations and limitations. x86 PTQ/QAT inference is a separate backend.