torch-batteries 0.7.0

A lightweight Python package that supplies batteries-included abstractions for PyTorch workflows

torch-batteries

Image generated by AI

A lightweight Python package that supplies batteries-included abstractions for:

• Data loading pipelines
• Model training loops
• Evaluation workflows
• Metrics computation
• Experiment tracking (Weights & Biases)

Designed to reduce boilerplate and standardize experiment code.

Installation

pip install torch-batteries

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_batteries import Battery, Event, EventContext, StepOutput, charge


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

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


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, val_loader=None, epochs=10)

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.

Optional Dependencies

For experiment tracking with Weights & Biases:

pip install torch-batteries[wandb]

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
Image Classification with CNN	Train a Convolutional Neural Network (CNN) for image classification	image_classification.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