helios-ml 2.0.0

A light-weight system for training AI networks using PyTorch

What is Helios?

Named after Greek god of the Sun, Helios is a light-weight package for training ML networks built on top of PyTorch. It is designed to abstract all of the “boiler-plate” code involved with training. Specifically, it wraps the following common patterns:

• Creation of the dataloaders.

• Initialisation of CUDA, PyTorch, and random number states.

• Initialisation for distributed training.

• Training, validation, and testing loops.

• Saving and loading checkpoints.

• Exporting to ONNX.

It is important to note that Helios is not a fully fledged training environment similar to PyTorch Lightning. Instead, Helios is focused on providing a simple and straight-forward interface that abstracts most of the common code patterns while retaining the ability to be easily overridden to suit the individual needs of each training scheme.

Quick Start

Training with Helios requires three steps:

1. Subclass DataModule to define your datasets and dataloaders.

2. Subclass Model to implement your training logic.

3. Create a Trainer and call fit().

import helios
import helios.data as hld
import helios.model as hlm

class MyDataModule(hld.DataModule):
    def setup(self) -> None:
        dataset = MyDataset(...)
        params = hld.DataLoaderParams(batch_size=32)
        self._add_train_phase(dataset, params)

class MyModel(hlm.Model):
    def setup(self, for_inference: bool = False) -> None:
        self._net = MyNetwork()
        self._optimizer = ...
        self._criterion = ...

    def train_step(self, batch, state: helios.TrainingState) -> None:
        inputs, labels = batch
        self._optimizer.zero_grad()
        loss = self._criterion(self._net(inputs), labels)
        loss.backward()
        self._optimizer.step()

trainer = helios.Trainer()
trainer.fit(MyModel("my_model"), MyDataModule())

Main Features

Helios offers the following functionality out of the box:

1. Resume training: Helios has been built with the ability to resume training if it is paused. Specifically, Helios will ensure that the behaviour of the trained model is identical to the one it would’ve had if it had been trained without pauses.

2. Automatic detection of multi-GPU environments for distributed training. In addition, Helios also supports training using torchrun and will automatically handle the initialisation and clean up of the distributed state. It will also correctly set the devices and maps to ensure weights are mapped to the correct location.

3. Registries for creation of arbitrary types. These include: networks, loss functions, optimisers, schedulers, etc.

4. Correct handling of logging when doing distributed training (even over multiple nodes).

5. Native integration with Optuna for hyper-parameter optimisation. Also supports resuming studies and generating checkpoints to ensure reproducibility.

Why Helios?

Compared to larger frameworks, Helios prioritises explicitness and simplicity over automation. The table below compares it against PyTorch Lightning and Ignite across several areas for research and engineering work:

Feature	Helios	Lightning	Ignite
Mixed precision	Helper functions	Automatic	Manual
Reproducible resume	Built-in	Partial	Manual
Distributed training	torchrun	torchrun	Manual
Boilerplate style	Explicit	Hook-based	Event-based
Training unit	First-class (EPOCH/ITERATION)	Epoch-based	Epoch-based
Gradient accumulation	Iteration-aware	Epoch-based	Manual
Registry system	Built-in	None	None
Learning curve	Low	Medium	High
Debuggability	High	Medium	High

Registry system: Helios provides typed global registries (MODEL_REGISTRY, DATASET_REGISTRY, OPTIMIZER_REGISTRY, and others) that map string names to types. Any component can be created by name, enabling config-file-driven experiments and trivial component swaps without changing training code. Register a class with @REGISTRY.register and create an instance with create_model("MyModel", ...).

Core philosophy: Helios removes training boilerplate without hiding what is happening. The code has been designed to be easy to read and trace in the event something goes wrong. You can always read the code and know exactly what runs at every point in the training loop.

Installation

Install Helios using pip:

pip install -U helios-ml

The following optional features are available as install extras:

• Optuna integration for hyper-parameter tuning: pip install -U helios-ml[tune]

• Tensorboard logging: pip install -U helios-ml[tensorboard]

• Weights & Biases logging: pip install -U helios-ml[wandb]

If you require a specific version of CUDA, you can install with:

pip install -U helios-ml --extra-index-url https://download.pytorch.org/whl/cu<version>

Documentation

Documentation available here.

Contributing

Contributing guidelines are available here.

License

Helios is published under the BSD-3 license and can be viewed here.