Metadata-Version: 2.4
Name: torch-batteries
Version: 1.0.1
Summary: A lightweight, event-driven training layer for PyTorch that provides reusable utilities for training, evaluation, metrics, checkpointing, and experiment management.
Author: Michal Szczygiel, Arkadiusz Paterak, Antoni Zięciak
License-Expression: Apache-2.0
Project-URL: Documentation, https://michalszc.github.io/torch-batteries/
Project-URL: Homepage, https://github.com/michalszc/torch-batteries
Project-URL: Repository, https://github.com/michalszc/torch-batteries
Project-URL: Issues, https://github.com/michalszc/torch-batteries/issues
Keywords: pytorch,machine learning,deep learning,training,evaluation,datasets,metrics,callbacks,checkpointing,experiment tracking
Classifier: Development Status :: 5 - Production/Stable
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
Classifier: Programming Language :: Python :: 3.13
Classifier: Programming Language :: Python :: 3.14
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.12
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: torch>=2.9.0
Requires-Dist: tqdm>=4.67.1
Requires-Dist: PyYAML>=6.0
Requires-Dist: typing-extensions>=4.10.0
Provides-Extra: example
Requires-Dist: datasets<6.0.0,>=5.0.1; extra == "example"
Requires-Dist: diffusers>=0.39.0; extra == "example"
Requires-Dist: gymnasium[classic-control]<2.0.0,>=1.2.0; extra == "example"
Requires-Dist: ipykernel>=7.1.0; extra == "example"
Requires-Dist: ipywidgets>=8.1.7; extra == "example"
Requires-Dist: kagglehub>=0.3.0; extra == "example"
Requires-Dist: matplotlib>=3.10.7; extra == "example"
Requires-Dist: numpy>=2.3.4; extra == "example"
Requires-Dist: pandas>=2.2.0; extra == "example"
Requires-Dist: scikit-learn>=1.7.2; extra == "example"
Requires-Dist: seaborn>=0.13.2; extra == "example"
Requires-Dist: torchvision>=0.24.1; extra == "example"
Provides-Extra: wandb
Requires-Dist: wandb>=0.16.0; extra == "wandb"
Provides-Extra: all
Requires-Dist: torch-batteries[example,wandb]; extra == "all"
Dynamic: license-file

# torch-batteries

<p align="center">
   <a href="https://pypi.org/project/torch-batteries/"><img src="https://img.shields.io/pypi/v/torch-batteries" alt="PyPI version"></a>
   <a href="https://pypi.org/project/torch-batteries/"><img src="https://img.shields.io/pypi/dm/torch-batteries" alt="PyPI monthly downloads"></a>
   <a href="https://github.com/michalszc/torch-batteries/blob/master/pyproject.toml"><img src="https://img.shields.io/badge/python-3.12%2B-important" alt="Python 3.12 or newer"></a>
   <a href="https://github.com/michalszc/torch-batteries/actions/workflows/ci.yml"><img src="https://github.com/michalszc/torch-batteries/actions/workflows/ci.yml/badge.svg?branch=master" alt="CI on master"></a>
   <a href="LICENSE"><img src="https://img.shields.io/badge/license-Apache%202.0-blue.svg" alt="License: Apache 2.0"></a>
   <a href="https://michalszc.github.io/torch-batteries/"><img src="https://img.shields.io/badge/docs-online-blue.svg" alt="Docs"></a>
   <a href="https://zenodo.org/badge/latestdoi/1103141433"><img src="https://zenodo.org/badge/1103141433.svg" alt="Latest Zenodo DOI"></a>
</p>

<p align="center">
   <img src="https://raw.githubusercontent.com/michalszc/torch-batteries/refs/heads/master/assets/logo.png" alt="torch-batteries" style="max-height:512px; height:auto;" />
   <br/>
   <em>Image generated by AI</em>
</p>

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

```bash
pip install torch-batteries
```

Python 3.12+ and PyTorch 2.9+ are required.

For Weights & Biases integration:

```bash
pip install "torch-batteries[wandb]"
```

For the notebook dependencies:

```bash
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:

```python
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](https://michalszc.github.io/torch-batteries/getting-started/installation/)
- [Quick start](https://michalszc.github.io/torch-batteries/getting-started/quickstart/)
- [Core concepts](https://michalszc.github.io/torch-batteries/getting-started/core-concepts/)
- [Guides](https://michalszc.github.io/torch-batteries/guides/)
- [DataPack workflows](https://michalszc.github.io/torch-batteries/guides/data-pack/)
- [API reference](https://michalszc.github.io/torch-batteries/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](notebooks/function_fitting.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/michalszc/torch-batteries/blob/master/notebooks/function_fitting.ipynb) |
| **Iris Classification with MLP** | Classify the Hugging Face Iris dataset with a tiny MLP and implicit DataPack workflows | [iris_classification.ipynb](notebooks/iris_classification.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/michalszc/torch-batteries/blob/master/notebooks/iris_classification.ipynb) |
| **Image Classification with CNN** | Build MNIST datasets and DataLoaders through an event-driven DataPack | [image_classification.ipynb](notebooks/image_classification.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/michalszc/torch-batteries/blob/master/notebooks/image_classification.ipynb) |
| **FashionMNIST Diffusion** | Train a class-conditioned Diffusers U-Net with optimization callbacks and streaming prediction | [fashion_mnist_diffusion.ipynb](notebooks/fashion_mnist_diffusion.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/michalszc/torch-batteries/blob/master/notebooks/fashion_mnist_diffusion.ipynb) |
| **CartPole Reinforcement Learning** | Train a compact DQN from replay transitions with optimization events and stateful metrics | [cartpole_reinforcement_learning.ipynb](notebooks/cartpole_reinforcement_learning.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/michalszc/torch-batteries/blob/master/notebooks/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](notebooks/cifar10_transfer_learning.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/michalszc/torch-batteries/blob/master/notebooks/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](notebooks/lr_sweep_early_stopping.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/michalszc/torch-batteries/blob/master/notebooks/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](notebooks/usps_semeion_multi_dataset.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/michalszc/torch-batteries/blob/master/notebooks/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](notebooks/ames_housing_multi_target.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/michalszc/torch-batteries/blob/master/notebooks/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](notebooks/frozenlake_a2c_local_tracking.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/michalszc/torch-batteries/blob/master/notebooks/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](notebooks/overhead_over_vanilla_pytorch.ipynb) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/michalszc/torch-batteries/blob/master/notebooks/overhead_over_vanilla_pytorch.ipynb) |

## Development

See [CONTRIBUTING.md](CONTRIBUTING.md) for environment setup, quality checks, and
the release workflow. Issues and feature requests are welcome in the
[GitHub repository](https://github.com/michalszc/torch-batteries/issues).
