torch-batteries

Image generated by AI

A lightweight, event-driven training layer for PyTorch that provides reusable utilities for training, evaluation, metrics, checkpointing, and experiment management.

Keep the forward pass and loss in your model while Battery handles device placement, optimization, metrics, callbacks, checkpoints, progress, evaluation, and prediction.

Features

• Explicit train, validation, test, and prediction steps with @charge
• Event-driven DataPack construction for reusable datasets and DataLoaders
• Single-forward-pass automatic, stateful, and full-phase metrics
• Early stopping and Top-K model checkpoints
• Resumable model, optimizer, callback, metric, DataPack, and history state
• Gradient accumulation, clipping, mixed precision, and scheduler callbacks
• Structured batch transfer and structured/streaming prediction
• Optional Weights & Biases experiment tracking

Installation

pip install torch-batteries

Python 3.12+ and PyTorch 2.9+ are required.

For Weights & Biases integration:

pip install "torch-batteries[wandb]"

For the notebook dependencies:

pip install "torch-batteries[example]"

Quick start

Define workflow steps directly on a PyTorch model with @charge. Return StepOutput to expose the same predictions used for the loss to automatic metrics:

import torch
from torch import nn
from torch.nn import functional as F
from torch.utils.data import DataLoader, TensorDataset

from torch_batteries import Battery, Event, EventContext, StepOutput, charge


class Model(nn.Module):
    def __init__(self) -> None:
        super().__init__()
        self.linear = nn.Linear(10, 1)

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.linear(inputs)

    @charge(Event.TRAIN_STEP)
    def training_step(self, context: EventContext) -> StepOutput:
        inputs, targets = context["batch"]
        predictions = self(inputs)
        return StepOutput(
            loss=F.mse_loss(predictions, targets),
            predictions=predictions,
            targets=targets,
        )


inputs = torch.randn(64, 10)
targets = torch.randn(64, 1)
train_loader = DataLoader(TensorDataset(inputs, targets), batch_size=16)

model = Model()
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
battery = Battery(
    model,
    optimizer=optimizer,
    metrics={"mae": lambda pred, target: F.l1_loss(pred, target)},
)
results = battery.train(train_loader, epochs=3, verbose=0)
print(results["train_loss"])

Automatic metrics are the functions passed to Battery(metrics=...). They use the predictions and targets returned in StepOutput and are averaged across each phase without an additional model forward pass. Add charged validation, test, or prediction methods only for the workflows you use.

Documentation

• Installation
• Quick start
• Core concepts
• Guides
• DataPack workflows
• API reference

Examples

Explore practical examples demonstrating torch-batteries features:

Example	Description	Notebook	Colab
Function Fitting with MLP	Train a neural network to approximate a polynomial function using the event-driven training approach	function_fitting.ipynb
Iris Classification with MLP	Classify the Hugging Face Iris dataset with a tiny MLP and implicit DataPack workflows	iris_classification.ipynb
Image Classification with CNN	Build MNIST datasets and DataLoaders through an event-driven DataPack	image_classification.ipynb
FashionMNIST Diffusion	Train a class-conditioned Diffusers U-Net with optimization callbacks and streaming prediction	fashion_mnist_diffusion.ipynb
CartPole Reinforcement Learning	Train a compact DQN from replay transitions with optimization events and stateful metrics	cartpole_reinforcement_learning.ipynb
CIFAR-10 ResNet18 Transfer Learning	Fine-tune a pretrained ResNet18 with resumable training state, full-phase metrics, and structured prediction	cifar10_transfer_learning.ipynb
Learning Rate Sweep with Early Stopping	Conduct a learning rate sweep on MNIST classification with aggressive early stopping and log results to Weights & Biases	lr_sweep_early_stopping.ipynb
USPS and SEMEION Named Datasets	Train and evaluate across two digit sources with named metrics and early stopping	usps_semeion_multi_dataset.ipynb
Ames Housing With Two Targets	Train a shared model with regression and classification heads using DataPack and MetricSpec	ames_housing_multi_target.ipynb
FrozenLake A2C with Local Tracking	Train and evaluate an on-policy agent through DataPack and save metrics with LocalTracker	frozenlake_a2c_local_tracking.ipynb
Overhead over Vanilla PyTorch	Measure the epoch and per-batch time Battery.train adds over an equivalent vanilla PyTorch loop on ResNets, optionally alongside PyTorch Lightning, PyTorch Ignite, fastai, and Keras with the PyTorch backend	overhead_over_vanilla_pytorch.ipynb

Development

See CONTRIBUTING.md for environment setup, quality checks, and the release workflow. Issues and feature requests are welcome in the GitHub repository.