Documentation

Complete reference for the AxiomML experiment framework. All configuration parameters, API methods, and advanced usage covered.

Overview

AxiomML is a modular, config-driven ML experiment framework built on PyTorch. It separates experiment logic from configuration, enabling reproducible runs with zero boilerplate. Features include:

• YAML/JSON config — Declare every experiment parameter in a single file.
• W&B integration — Automatic metric logging, artifact tracking, and rich media.
• Distributed training — DDP with mixed precision, gradient checkpointing, and sync BN.
• Checkpoint management — Top-K saving by metric, auto-resume, versioned artifacts.
• Modular components — Swap models, datasets, optimizers, schedulers without touching training code.

Installation

# From PyPI
pip install axioml

# From source (latest)
pip install git+https://github.com/anomalyco/AxiomML.git

Requires Python 3.10+ and PyTorch 2.0+. CUDA is optional but recommended.

Quick Start

Create a config file and run an experiment in three steps:

# 1. Create config.yaml
model: mlp
dataset: mnist
trainer:
  epochs: 5
  batch_size: 64
  learning_rate: 0.001

# 2. Write train.py
from axioml import Experiment

exp = Experiment("config.yaml")
exp.train()
metrics = exp.evaluate()

# 3. Run
python train.py

Configuration System

Configs are YAML files (or JSON) that define every aspect of an experiment. The system supports:

• Hierarchical merge — Base configs with overrides.
• CLI overrides — python train.py --config trainer.epochs=20
• Presets — Predefined configs for common architectures.
• Type coercion — Automatic type conversion from YAML to Python types.

Below is the complete parameter reference organized by section.

Experiment Config

Top-level experiment settings.

Data / Dataset Config

Configure dataset loading, transforms, batching, and splitting.

Model Config

Select model architecture and load pretrained weights.

Optimizer Config

Scheduler Config

Trainer Config

Logging Config

Checkpoint Config

Distributed Config

Full Example Config

# configs/full_example.yaml
# Top-level
name: my_experiment
seed: 42
device: cuda
mixed_precision: fp16
output_dir: ./outputs

# Data
data:
  dataset: cifar10
  root: ./data
  transforms:
    - RandomHorizontalFlip
    - RandomCrop: {size: 32, padding: 4}
    - ToTensor
    - Normalize: {mean: [0.5, 0.5, 0.5], std: [0.5, 0.5, 0.5]}
  batch_size: 128
  num_workers: 8
  shuffle: true
  splits:
    train: 0.8
    val: 0.1
    test: 0.1

# Model
model:
  name: resnet18
  params:
    num_classes: 10
    pretrained: false
  compile: true

# Optimizer
optimizer:
  name: adamw
  lr: 0.0003
  weight_decay: 0.01
  betas: [0.9, 0.95]

# Scheduler
scheduler:
  name: warmup_cosine
  params:
    T_max: 100
    eta_min: 1e-6
  warmup_epochs: 5

# Trainer
trainer:
  epochs: 100
  grad_accumulation: 2
  validate_every: 1
  early_stop_patience: 20

# Logging
logging:
  project: axioml-demo
  tags: [resnet18, cifar10, adamw]
  log_images: true
  log_histograms: true

# Checkpoint
checkpoint:
  dir: ./checkpoints
  top_k: 3
  metric: val_acc
  mode: max

# Distributed
distributed:
  enabled: false
  sync_bn: false

Experiment API

The Experiment class is the main entry point. It accepts a config (path, dict, or parsed config object).

Constructor

Experiment(
    config: str | dict | Config,
    distributed: bool = False,
    resume: str | None = None,
)

Methods

Trainer API

The Trainer handles the training loop, validation, checkpointing, and logging. Accessible via exp.trainer.

Properties

Data / Dataset API

Built-in Datasets

Custom Datasets

Register any torch.utils.data.Dataset subclass:

from axioml.data import register_dataset

@register_dataset("my_dataset")
class MyDataset(torch.utils.data.Dataset):
    def __init__(self, root, split="train", transforms=None):
        # ...
        pass
    def __len__(self): pass
    def __getitem__(self, idx): pass

Transforms

Transforms are specified by name in config. Built-in transforms include:

Models (Model Zoo)

AxiomML includes a registry of built-in model architectures. Custom models can be registered similarly to datasets.

Built-in Models

Registering Custom Models

from axioml.models import register_model

@register_model("my_net")
class MyNet(nn.Module):
    def __init__(self, num_classes=10, **kwargs):
        super().__init__()
        # ...

    def forward(self, x):
        return x

Logging API

AxiomML uses a Logger abstraction that writes to multiple backends simultaneously: W&B, console (Rich), and file.

Callbacks

AxiomML supports a callback system for hooking into the training lifecycle. Callbacks are classes that implement one or more of the following methods:

from axioml.callbacks import Callback

class MyCallback(Callback):
    def on_epoch_end(self, trainer):
        print(f"Epoch {trainer.current_epoch} done")

Built-in callbacks (auto-enabled):

• CheckpointCallback — Saves top-K checkpoints by metric.
• LoggingCallback — Metrics logging to W&B and console.
• EarlyStoppingCallback — Stop on plateau.
• ProgressBarCallback — Rich progress bar during training.
• LRMonitorCallback — Log learning rate each step.

Metrics

AxiomML tracks standard metrics automatically. Custom metrics can be added:

Export & Inference

After training, export models for production deployment:

Inference example:

import torch
from axioml import Experiment

exp = Experiment("config.yaml")
exp.load_checkpoint("checkpoints/best.pt")

# Single prediction
x = torch.randn(1, 3, 32, 32)  # example input
pred = exp.predict(x)

# Batch prediction
batch = torch.randn(16, 3, 32, 32)
predictions = exp.predict(batch)
# predictions: torch.Tensor of shape (16, num_classes)