This is an attempt to create a minimal extendable framework for development and research on classification models.
All code in mini_trainer should follow the following core principles:
- There should be NO additional dependencies beyond core
Python,PyTorch(torch,torchvision, etc.),matplotlibandtqdm. - The required portion of any API should be as minimal as possible (i.e. to train a model we only require
mt_train -i <TRAINING_DATA>) - All hyperparameters and system configuration should have smart defaults that are as general as possible
- All functionality should be extendable to custom model architectures, loss functions, training regimes, data formats etc.
We recommend using uv for package and environment management.
See Install uv for instructions.
# Recommended installation (includes logging, visualization, and optional utilities)
uv pip install "mini_trainer[recommended]" --torch-backend=auto
# or standard pip
pip install "mini_trainer[recommended]"
# Installation with all features (timm, transformers, BioCLIP, etc.)
uv pip install "mini_trainer[all]" --torch-backend=auto
# or standard pip
pip install "mini_trainer[all]"
# Minimal installation (core training & inference loop only)
uv pip install mini_trainer --torch-backend=auto
# or standard pip
pip install mini_trainergit clone ssh://git@github.com:asgersvenning/mini_trainer.git
cd mini_trainer
# Sync with recommended extras:
uv sync --extra recommended --extra [cpu/cu126/cu130/cu132]
# Or sync with all features (timm, transformers, BioCLIP):
uv sync --extra all --extra [cpu/cu126/cu130/cu132]
source .venv/bin/activateTip
We highly recommend installing torch and torchvision with native CUDA support via either uv sync ... --extra [cpu/cu126/cu130/cu132] or uv pip install ... --torch-backend=auto, and crucially running scripts or tools associated with your uv virtual environment by activating the venv:
source .venv/bin/activateUsing uv run ... is likely to automatically install CUDA-incompatible wheels. If you really want to use uv run, we suggest using the --no-sync flag every time.
Note that if you are "lucky" you might have the default CUDA version on your system, meaning that uv run might in fact use the correct wheels. This is, however, not guaranteed.
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.
Feel free to contribute, but here are a few tips:
- Follow the installation guide to setup a proper dev environment.
- Setup linting via
ruff; verify with:ruff check mini_trainer. - Please avoid adding new dependencies 🙂