feat: add a PyTorch backend alongside Keras

.github/workflows/ci.yml

@@ -48,13 +48,65 @@ jobs:
           uv venv --python ${{ matrix.python-version }}
           source .venv/bin/activate
           uv pip install -e ".[dev]"
+          # CPU-only torch. The default wheel bundles CUDA libraries that clash
+          # with the ones TensorFlow loads, which segfaults the interpreter as
+          # soon as both are imported in one process. The runners have no GPU,
+          # so the CPU build is the right one here regardless.
+          uv pip install torch --index-url https://download.pytorch.org/whl/cpu
       - name: Run tests
         env:
           TF_CPP_MIN_LOG_LEVEL: "3"
         run: |
           source .venv/bin/activate
           pytest tests -v
 
+  # Each backend has to work on its own, so prove the package is usable with
+  # only one framework installed. The suite skips the modules whose framework
+  # is absent.
+  backend-isolation:
+    runs-on: ubuntu-latest
+    strategy:
+      fail-fast: false
+      matrix:
+        include:
+          - backend: torch
+            # CPU-only, matching the test job and the GPU-less runners.
+            install: "torch --index-url https://download.pytorch.org/whl/cpu"
+          - backend: tensorflow
+            install: "tensorflow"
+    steps:
+      - uses: actions/checkout@v4
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: true
+      - name: Install with only ${{ matrix.backend }}
+        run: |
+          uv venv --python 3.12
+          source .venv/bin/activate
+          # --no-deps keeps the other framework out; the package still declares
+          # tensorflow as a runtime dependency for the 2.x line.
+          uv pip install --no-deps -e .
+          uv pip install numpy pytest
+          # Installed separately so a custom index applies only to the backend.
+          uv pip install ${{ matrix.install }}
+      - name: Confirm the other framework is absent
+        run: |
+          source .venv/bin/activate
+          python - <<'PY'
+          import importlib.util
+          other = {"torch": "tensorflow", "tensorflow": "torch"}["${{ matrix.backend }}"]
+          assert importlib.util.find_spec(other) is None, f"{other} unexpectedly installed"
+          assert importlib.util.find_spec("${{ matrix.backend }}") is not None
+          print(f"only ${{ matrix.backend }} present, as intended")
+          PY
+      - name: Run the tests that apply
+        env:
+          TF_CPP_MIN_LOG_LEVEL: "3"
+        run: |
+          source .venv/bin/activate
+          pytest tests -v
+
   build:
     runs-on: ubuntu-latest
     steps:

CHANGELOG.md

@@ -5,6 +5,35 @@ All notable changes to this project are documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [2.2.0] - 2026-07-28
+
+Adds a PyTorch backend. Fully backward compatible: the Keras API is unchanged
+and no existing checkpoint needs converting.
+
+### Added
+
+- `TorchTrainer`, a PyTorch backend sharing `Target`, the round loop, the
+  timeout, the `Ctrl+C` handling and the resume behaviour with the Keras
+  trainer. PyTorch has no `fit`, so you pass a step function that runs one round
+  and returns its metrics.
+- A `torch` extra: `pip install "infinite-training[torch]"`.
+- A `tensorflow` extra, so new code can already name the dependency it wants.
+  TensorFlow remains a hard requirement for the whole 2.x line; it moves to the
+  extra in 3.0.0.
+- `examples/mnist_torch.py`, the PyTorch counterpart of the Keras MNIST example.
+- Tests for the PyTorch backend and for backend isolation (39 new, 94 in total),
+  plus a CI job that installs each framework on its own and runs the suite.
+
+### Changed
+
+- Backends are imported lazily. `import infinite_training` no longer imports
+  TensorFlow, and referencing `TorchTrainer` never imports it either — so having
+  only one of the two frameworks installed is enough. Asking for a trainer whose
+  framework is missing raises an `ImportError` naming the extra to install.
+- The loop, the target and timeout bookkeeping, the value history and the
+  checkpoint paths moved into an internal `_base` module shared by both
+  backends. No public name changed.
+
 ## [2.1.0] - 2026-07-25
 
 Backward compatible with 2.0.0: existing code keeps working and warns where an
@@ -89,5 +118,6 @@ The following still work and will be removed in 3.0.0:
 
 - Upgraded to the latest TensorFlow and switched the release workflow to `uv`.
 
+[2.2.0]: https://github.com/vyncint/infinite-training/releases/tag/v2.2.0
 [2.1.0]: https://github.com/vyncint/infinite-training/releases/tag/v2.1.0
 [2.0.0]: https://github.com/vyncint/infinite-training/releases/tag/v2.0.0

README.md

@@ -5,19 +5,31 @@
 [![Python versions](https://img.shields.io/pypi/pyversions/infinite-training.svg)](https://pypi.org/project/infinite-training/)
 [![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
 
-Train a Keras model until it is *good enough*, until you run out of time, or until you press `Ctrl+C` — then pick up exactly where you left off.
+Train a **Keras or PyTorch** model until it is *good enough*, until you run out of time, or until you press `Ctrl+C` — then pick up exactly where you left off.
 
-`Model.fit` makes you choose the number of epochs up front. Often what you actually want is "keep going until validation accuracy passes 0.98", or "train for ten minutes and keep the best result". `infinite_training` wraps `fit` in a resumable loop that does that, remembers the best weights it has seen, and checkpoints everything to disk so an interrupted run is never wasted.
+Training loops make you choose the number of epochs up front. Often what you actually want is "keep going until validation accuracy passes 0.98", or "train for ten minutes and keep the best result". `infinite_training` wraps the loop so it does that, remembers the best weights it has seen, and checkpoints everything to disk so an interrupted run is never wasted.
+
+Both backends share the same `Target`, checkpointing and resume behaviour:
+
+| Backend | Class | You provide |
+| --- | --- | --- |
+| Keras | `InfiniteTrainer` | A `tf.keras.Model`; the trainer calls `fit` |
+| PyTorch | `TorchTrainer` | A `nn.Module` and a step function returning metrics |
 
 ---
 
 ## Contents
 
 - [Installation](#installation)
 - [Quick start](#quick-start)
+  - [Keras](#keras)
+  - [PyTorch](#pytorch)
 - [How it works](#how-it-works)
 - [Resuming a session](#resuming-a-session)
 - [API reference](#api-reference)
+  - [`Target`](#target)
+  - [`InfiniteTrainer`](#infinitetrainer-keras)
+  - [`TorchTrainer`](#torchtrainer-pytorch)
 - [Choosing a target](#choosing-a-target)
 - [Security note on checkpoints](#security-note-on-checkpoints)
 - [Migrating from 2.0.x](#migrating-from-20x)
@@ -32,18 +44,44 @@ Train a Keras model until it is *good enough*, until you run out of time, or unt
 pip install infinite-training
 ```
 
-To run the bundled example as well:
+For the PyTorch backend:
+
+```bash
+pip install "infinite-training[torch]"
+```
+
+To run the bundled Keras example as well:
 
 ```bash
 pip install "infinite-training[example]"
 ```
 
-Requires Python 3.10+ and TensorFlow.
+Requires Python 3.10+.
+
+The two backends are imported lazily, so you only need the framework you
+actually use — importing `TorchTrainer` never loads TensorFlow, and vice versa.
+
+> **Note on the 2.x dependency set.** TensorFlow is still a hard requirement of
+> `infinite-training` itself, so that an existing `pip install infinite-training`
+> keeps working unchanged. If you only want PyTorch you can skip it today with
+> `pip install --no-deps infinite-training` followed by `pip install numpy torch`.
+> In 3.0.0 TensorFlow moves to the `tensorflow` extra and neither framework will
+> be installed by default.
+
+> **If you install both frameworks.** On Linux, the default PyTorch wheel bundles
+> CUDA libraries that can clash with the ones TensorFlow loads, segfaulting the
+> interpreter as soon as both are imported into the same process. This is not
+> specific to `infinite_training` — `import tensorflow; import torch` is enough
+> to trigger it. Because the backends here are imported lazily you will not hit
+> it by using one of them, but if you do need both installed, use the CPU build:
+> `pip install torch --index-url https://download.pytorch.org/whl/cpu`.
 
 ---
 
 ## Quick start
 
+### Keras
+
 ```python
 import numpy as np
 import tensorflow as tf
@@ -75,6 +113,54 @@ print(f"best loss {best_loss:.6f} after {trainer.rounds_completed} round(s)")
 
 `compile` and `train` forward their arguments to `Model.compile` and `Model.fit`, so anything you already pass to Keras keeps working.
 
+### PyTorch
+
+PyTorch has no `fit`, so you hand the trainer your training step instead. It is
+called once per round and returns the metrics for that round — whatever you want
+to target has to appear in the mapping it returns.
+
+```python
+import torch
+from torch import nn
+from infinite_training import TorchTrainer, Target
+
+x = torch.rand(256, 2)
+y = x.sum(dim=1, keepdim=True)
+
+model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 1))
+optimizer = torch.optim.Adam(model.parameters())
+loss_fn = nn.MSELoss()
+
+
+def step() -> dict[str, float]:
+    model.train()
+    optimizer.zero_grad()
+    loss = loss_fn(model(x), y)
+    loss.backward()
+    optimizer.step()
+    return {"loss": loss.item()}
+
+
+trainer = TorchTrainer(
+    model=model,
+    target=Target(name="loss", smaller_is_better=True, target_value=1e-4),
+    timeout=60,  # seconds
+)
+
+trainer.train(step)  # loops until the target, the timeout, or Ctrl+C
+
+predictions, best_loss = trainer.predict_best(x)
+print(f"best loss {best_loss:.6f} after {trainer.rounds_completed} round(s)")
+```
+
+There is no `compile()` step: the shadow copy that holds the best weights is
+built in the constructor, so inference works straight away.
+
+One round is one call to your step function, so *you* choose the granularity at
+which the target and the timeout are checked — one epoch, a fixed number of
+batches, or an epoch followed by an evaluation pass, as in
+[`examples/mnist_torch.py`](examples/mnist_torch.py).
+
 ---
 
 ## How it works
@@ -129,13 +215,13 @@ The stopping criterion.
 
 | Argument | Type | Default | Description |
 | --- | --- | --- | --- |
-| `name` | `str` | `"loss"` | Key to read from the Keras `History`. Any loss or metric, including `val_*` keys. |
+| `name` | `str` | `"loss"` | Key to read from the Keras `History`, or from the mapping your PyTorch step returns. Any loss or metric, including `val_*` keys. |
 | `smaller_is_better` | `bool` | `True` | `True` for losses and error rates, `False` for accuracy-like metrics. |
 | `target_value` | `float \| None` | `None` | Value at which training stops. `None` means an unreachable bound, so the session is limited only by `timeout` or `Ctrl+C`. |
 
 Methods: `is_improvement(candidate, incumbent)`, `is_reached(value)`, and the `worst_possible_value` property.
 
-### `InfiniteTrainer`
+### `InfiniteTrainer` (Keras)
 
 | Argument | Type | Default | Description |
 | --- | --- | --- | --- |
@@ -165,6 +251,42 @@ Methods: `is_improvement(candidate, incumbent)`, `is_reached(value)`, and the `w
 | `best_weights` / `last_weights` | Weight lists. |
 | `best_model` | Shadow model holding the best weights (available after `compile`). |
 
+### `TorchTrainer` (PyTorch)
+
+| Argument | Type | Default | Description |
+| --- | --- | --- | --- |
+| `model` | `nn.Module` | required | Deep-copied once to hold the best weights, so it must be copyable. |
+| `target` | `Target` | `Target()` | Stopping criterion. |
+| `timeout` | `float \| None` | `None` | Wall-clock budget in seconds, checked between rounds. `None` means unbounded. |
+| `best_weights_path` | `str` | `"best_weights.npy"` | Best weights checkpoint. |
+| `last_weights_path` | `str` | `"last_weights.npy"` | Most recent weights checkpoint. |
+| `best_value_path` | `str` | `"best_value.npy"` | Best observed value. |
+| `value_history_path` | `str` | `"value_history.npy"` | Per-round value history. |
+
+| Method | Description |
+| --- | --- |
+| `train(step_fn, *args, **kwargs)` | Calls `step_fn` in a loop until the target, the timeout, or `Ctrl+C`. Extra arguments are forwarded to it. Always checkpoints. |
+| `save()` | Write all four checkpoints immediately. |
+| `predict_best(*args, **kwargs)` | `(output, best_value)` using the best weights, in eval mode under `torch.no_grad()`. |
+| `predict_last(*args, **kwargs)` | `(output, last_value)` using the most recent weights. |
+
+Properties are the same as for `InfiniteTrainer`, except that `best_weights` and
+`last_weights` are state dicts of NumPy arrays rather than Keras weight lists,
+and `best_model` is ready immediately — there is no `compile()`.
+
+**The step-function contract.** `step_fn` must return a mapping of metric name
+to value, for example `{"loss": 0.31}`. A value may be a float, a
+`torch.Tensor`, a NumPy scalar, or a sequence — for a sequence the last element
+is used, matching how Keras reports a multi-epoch `History`. If the target names
+a key that is not in the mapping, the trainer raises `RuntimeError` listing the
+keys that were returned; if the step returns something that is not a mapping, it
+raises `TypeError`.
+
+**Checkpoint format.** State dicts are stored as NumPy arrays rather than
+`torch.save` archives, so a checkpoint written on a GPU resumes on a CPU with no
+device map, and dtypes such as the `int64` buffer in `BatchNorm` round-trip
+unchanged.
+
 ---
 
 ## Choosing a target