torch-batteries

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
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
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

Release files for torch-batteries 0.8.0

For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.

Source distribution (sdist)

Source distribution for torch-batteries 0.8.0
File	Size	Uploaded
torch_batteries-0.8.0.tar.gz	51.4 kB	Aug 2, 2026	Details

Built distribution (wheel)

Table of built distributions (wheels) for torch-batteries 0.8.0
File	Interpreter	ABI	Platform	Reset
torch_batteries-0.8.0-py3-none-any.whl 62.9 kB Aug 2, 2026	Python 3	none	any	Details

Total release size: 114.3 kB

Release files / torch_batteries-0.8.0.tar.gz

Download URL	torch_batteries-0.8.0.tar.gz
Size	51.4 kB
Tags	Source
SHA-256 checksum How to use checksums	`4679474ef63fd42e93ab77e2a711a3459d2415ee059f559cb717e57dcd9cd61e`
BLAKE2b-256 checksum How to use checksums	`3b81cd5029335d229e02c1e7d6ac7879e6e91d87cf9413d81eaa852a03b6d37c`
Upload date	Aug 2, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/7.0.0 CPython/3.12.13`

Release files / torch_batteries-0.8.0-py3-none-any.whl

Download URL	torch_batteries-0.8.0-py3-none-any.whl
Size	62.9 kB
Tags	Python 3
SHA-256 checksum How to use checksums	`8514a16f6793082edc89d377795f7ca6f0248856477123939122833111749283`
BLAKE2b-256 checksum How to use checksums	`ff2aaa79b52f793c37067f2b5c45713dea5bf098c523ebf5e18f12da0d5e8b17`
Upload date	Aug 2, 2026
Uploaded using Trusted Publishing? What is trusted publishing?	No
Uploaded via	`twine/7.0.0 CPython/3.12.13`