AxiomML is a modular, high-performance ML experiment framework built for modern deep learning. First-class configuration, W&B logging, distributed training — no boilerplate.
from axioml import Experiment # Configure in YAML — zero boilerplate exp = Experiment("configs/mnist.yaml") # Train with auto-logging, checkpoints & metrics exp.train() # Evaluate and export metrics = exp.evaluate() exp.export("model.onnx")
From config-driven experiment management to distributed training — AxiomML gives you a production-ready foundation.
Declarative YAML/JSON configs with hierarchical merging, presets, and CLI overrides. Keep every experiment reproducible.
Automatic metric tracking, hyperparameter logging, artifact versioning, and rich media (figures, distributions) to Weights & Biases.
DDP (Distributed DataParallel) with automatic rank allocation, gradient checkpointing, and mixed-precision via native PyTorch amp.
Plug-and-play components: models, datasets, optimizers, schedulers, transforms. Swap parts without touching training logic.
Auto-save, resume, and version checkpoints. Keep top-K by metric. Includes optimizer state, epoch, and config snapshot.
Built-in transforms, stratified splitting, batching, shuffling, and caching. Supports custom datasets via simple registry.
# configs/mnist.yaml
model: mlp
dataset: mnist
trainer:
epochs: 10
batch_size: 64
learning_rate: 0.001
optimizer: adam
scheduler: cosine
logging:
project: axioml-demo
tags: [mlp, mnist]
checkpoint:
dir: ./checkpoints
top_k: 3
metric: val_acc
from axioml import Experiment exp = Experiment("configs/mnist.yaml") exp.train() # full training loop # Logs: loss, acc, lr, grad_norm → W&B + console # Checkpoints saved at each epoch
# Launch with torchrun # torchrun --nproc_per_node=4 train.py from axioml import Experiment exp = Experiment("configs/mnist.yaml", distributed=True) exp.train() # auto DDP wrapping # Gradient checkpointing, mixed precision, sync BN
exp.evaluate() # on test split # Returns: {loss, acc, f1, ...} exp.export("model.onnx") # export to ONNX exp.export("model.pt") # or raw TorchScript # Load checkpoint for inference exp.load_checkpoint("checkpoints/epoch=9-val_acc=0.99.pt")
pip install axioml — that's it. Python 3.10+ required, PyTorch installed automatically.
Create a config.yaml with your model, dataset, and training settings. Use presets for common architectures.
python train.py — one command to train, log, checkpoint, and evaluate. Distribute with torchrun.