Metadata-Version: 2.4
Name: torch-batteries
Version: 0.11.0
Summary: A lightweight Python package that supplies batteries-included abstractions for PyTorch workflows
Author: Michal Szczygiel, Arkadiusz Paterak, Antoni Zięciak
License-Expression: Apache-2.0
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
Classifier: Development Status :: 1 - Planning
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.12
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
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: matplotlib>=3.10.7; extra == "example"
Requires-Dist: numpy>=2.3.4; 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: datasets<6.0.0,>=5.0.1; extra == "all"
Requires-Dist: diffusers>=0.39.0; extra == "all"
Requires-Dist: gymnasium[classic-control]<2.0.0,>=1.2.0; extra == "all"
Requires-Dist: ipykernel>=7.1.0; extra == "all"
Requires-Dist: ipywidgets>=8.1.7; extra == "all"
Requires-Dist: matplotlib>=3.10.7; extra == "all"
Requires-Dist: numpy>=2.3.4; extra == "all"
Requires-Dist: scikit-learn>=1.7.2; extra == "all"
Requires-Dist: seaborn>=0.13.2; extra == "all"
Requires-Dist: torchvision>=0.24.1; extra == "all"
Requires-Dist: wandb>=0.16.0; 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="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>
</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. 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) |

## 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).
