torch-batteries

Image generated by AI

A lightweight, event-driven training layer for PyTorch. 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

Development

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