From 757498442d4ca7f9beaf17b5bbe66fac38db721c Mon Sep 17 00:00:00 2001 From: Vyncint Ng <115854244+vyncint@users.noreply.github.com> Date: Tue, 28 Jul 2026 05:59:56 +0700 Subject: [PATCH 1/2] feat: add a PyTorch backend alongside Keras PyTorch has no fit(), so TorchTrainer takes a step function that runs one round and returns its metrics. Everything around that - the Target, the round loop, the timeout, the Ctrl+C handling, the checkpointing and the resume behaviour - is shared with the Keras trainer, which is unchanged. The shared parts moved into an internal _base module. No public name changed and the existing 55 tests pass untouched. Backends are now imported lazily, so having only one of the two frameworks installed is enough: importing TorchTrainer never loads TensorFlow, and asking for a trainer whose framework is missing raises an ImportError naming the extra to install. TensorFlow stays a hard requirement for the 2.x line so existing installs keep working; it moves to the `tensorflow` extra in 3.0.0. Adds 39 tests (94 total), a CI job that installs each framework on its own, and a PyTorch MNIST example. Co-Authored-By: Claude Opus 5 --- .github/workflows/ci.yml | 46 ++++- CHANGELOG.md | 30 +++ README.md | 126 +++++++++++- examples/mnist_torch.py | 107 ++++++++++ infinite_training/__init__.py | 82 +++++++- infinite_training/_base.py | 152 ++++++++++++++ infinite_training/_storage.py | 30 ++- infinite_training/torch_trainer.py | 222 ++++++++++++++++++++ infinite_training/trainer.py | 100 +++------ pyproject.toml | 25 ++- tests/conftest.py | 89 +++++++- tests/test_lazy_imports.py | 135 +++++++++++++ tests/test_torch_trainer.py | 314 +++++++++++++++++++++++++++++ 13 files changed, 1363 insertions(+), 95 deletions(-) create mode 100644 examples/mnist_torch.py create mode 100644 infinite_training/_base.py create mode 100644 infinite_training/torch_trainer.py create mode 100644 tests/test_lazy_imports.py create mode 100644 tests/test_torch_trainer.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6cedbc6..6023b54 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -47,7 +47,7 @@ jobs: run: | uv venv --python ${{ matrix.python-version }} source .venv/bin/activate - uv pip install -e ".[dev]" + uv pip install -e ".[dev,torch]" - name: Run tests env: TF_CPP_MIN_LOG_LEVEL: "3" @@ -55,6 +55,50 @@ jobs: 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 + install: "torch" + - 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 ${{ 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: diff --git a/CHANGELOG.md b/CHANGELOG.md index 4880d9e..e91957e 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/README.md b/README.md index 66729eb..85d57ae 100644 --- a/README.md +++ b/README.md @@ -5,9 +5,16 @@ [![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 | --- @@ -15,9 +22,14 @@ Train a Keras model until it is *good enough*, until you run out of time, or unt - [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,36 @@ 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. --- ## Quick start +### Keras + ```python import numpy as np import tensorflow as tf @@ -75,6 +105,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 +207,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 +243,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 diff --git a/examples/mnist_torch.py b/examples/mnist_torch.py new file mode 100644 index 0000000..cec370e --- /dev/null +++ b/examples/mnist_torch.py @@ -0,0 +1,107 @@ +"""Train an MNIST classifier in PyTorch until it reaches 98% test accuracy. + +The PyTorch counterpart of ``examples/mnist.py``. PyTorch has no ``fit``, so the +training step is written out and handed to the trainer, which calls it once per +round and takes care of the target, the timeout and the checkpointing. + +Run it with:: + + pip install "infinite_training[torch]" torchvision + python examples/mnist_torch.py + +Stop it at any point with ``Ctrl+C``: the best and most recent weights are +written to ``models_torch/`` and picked up automatically the next time you run +it. +""" + +from __future__ import annotations + +import torch +from torch import nn +from torch.utils.data import DataLoader +from torchvision import datasets, transforms + +from infinite_training import Target, TorchTrainer + +CHECKPOINT_DIR = "models_torch" +BATCH_SIZE = 128 +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + + +def build_loaders() -> tuple[DataLoader, DataLoader]: + """MNIST train and test loaders, normalised to the usual mean and stdev.""" + transform = transforms.Compose( + [transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))] + ) + train = datasets.MNIST("data", train=True, download=True, transform=transform) + test = datasets.MNIST("data", train=False, download=True, transform=transform) + return ( + DataLoader(train, batch_size=BATCH_SIZE, shuffle=True), + DataLoader(test, batch_size=1000), + ) + + +def build_model() -> nn.Module: + return nn.Sequential( + nn.Flatten(), + nn.Linear(28 * 28, 128), + nn.ReLU(), + nn.Linear(128, 10), + ) + + +def main() -> None: + train_loader, test_loader = build_loaders() + + model = build_model().to(DEVICE) + optimizer = torch.optim.Adam(model.parameters(), lr=1e-3) + loss_fn = nn.CrossEntropyLoss() + + def step() -> dict[str, float]: + """One epoch of training followed by an evaluation pass. + + The returned mapping is what the trainer watches, so anything you want + to target has to appear here. + """ + model.train() + running_loss = 0.0 + for images, labels in train_loader: + images, labels = images.to(DEVICE), labels.to(DEVICE) + optimizer.zero_grad() + loss = loss_fn(model(images), labels) + loss.backward() + optimizer.step() + running_loss += loss.item() * images.size(0) + + model.eval() + correct = 0 + with torch.no_grad(): + for images, labels in test_loader: + images, labels = images.to(DEVICE), labels.to(DEVICE) + correct += (model(images).argmax(dim=1) == labels).sum().item() + + metrics = { + "loss": running_loss / len(train_loader.dataset), + "test_accuracy": correct / len(test_loader.dataset), + } + print(f"loss={metrics['loss']:.4f} test_accuracy={metrics['test_accuracy']:.4f}") + return metrics + + trainer = TorchTrainer( + model=model, + target=Target("test_accuracy", smaller_is_better=False, target_value=0.98), + best_weights_path=f"{CHECKPOINT_DIR}/best_weights.npy", + last_weights_path=f"{CHECKPOINT_DIR}/last_weights.npy", + best_value_path=f"{CHECKPOINT_DIR}/best_value.npy", + value_history_path=f"{CHECKPOINT_DIR}/value_history.npy", + ) + + print(f"Training on {DEVICE}; resuming from round {trainer.rounds_completed}.") + trainer.train(step) + + print(f"Best test accuracy: {trainer.best_value:.4f}") + print(f"Rounds completed: {trainer.rounds_completed}") + + +if __name__ == "__main__": + main() diff --git a/infinite_training/__init__.py b/infinite_training/__init__.py index da80480..81fe2c0 100644 --- a/infinite_training/__init__.py +++ b/infinite_training/__init__.py @@ -1,10 +1,16 @@ -"""Train Keras models until a target is met, time runs out, or you interrupt. +"""Train models until a target is met, time runs out, or you interrupt. -``infinite_training`` wraps ``Model.fit`` in a resumable loop that tracks the -best weights it has seen and checkpoints them to disk, so training can be -stopped with ``Ctrl+C`` and continued in a later session. +``infinite_training`` wraps a training loop in a resumable session that tracks +the best weights it has seen and checkpoints them to disk, so training can be +stopped with ``Ctrl+C`` and continued later. -Typical use:: +Two backends are available and they share the same +:class:`~infinite_training.Target`, checkpointing and resume behaviour: + +- :class:`~infinite_training.InfiniteTrainer` drives ``tf.keras.Model.fit``. +- :class:`~infinite_training.TorchTrainer` drives a training step you write. + +Keras:: from infinite_training import InfiniteTrainer, Target @@ -16,19 +22,81 @@ trainer.compile(optimizer="adam", loss="mse") trainer.train(x, y) predictions, best_value = trainer.predict_best(x) + +PyTorch:: + + from infinite_training import TorchTrainer, Target + + trainer = TorchTrainer( + model=model, + target=Target("loss", smaller_is_better=True, target_value=1e-4), + timeout=600, + ) + trainer.train(step) # step() runs one round, returns {"loss": ...} + predictions, best_value = trainer.predict_best(x) + +The backends are imported lazily, so installing only one of TensorFlow and +PyTorch is enough — and importing this package does not pull in a framework +until you actually reference a trainer. """ from __future__ import annotations +from typing import TYPE_CHECKING, Any + from .target import Target -from .trainer import InfiniteTrainer, InfinityTraining -__version__ = "2.1.0" +__version__ = "2.2.0" __all__ = [ "InfiniteTrainer", + "TorchTrainer", "Target", # Deprecated, removed in 3.0.0. "InfinityTraining", "__version__", ] + +if TYPE_CHECKING: # pragma: no cover - import-time typing only + from .torch_trainer import TorchTrainer + from .trainer import InfiniteTrainer, InfinityTraining + +# Attribute name -> the submodule that defines it. +_LAZY_ATTRS = { + "InfiniteTrainer": ".trainer", + "InfinityTraining": ".trainer", + "TorchTrainer": ".torch_trainer", +} + +# Which framework each submodule needs, for a readable error when it is absent. +_BACKEND_REQUIREMENTS = { + ".trainer": ("tensorflow", "tensorflow"), + ".torch_trainer": ("torch", "torch"), +} + + +def __getattr__(name: str) -> Any: + """Import a trainer on first use, so only the backend you use is loaded.""" + module_name = _LAZY_ATTRS.get(name) + if module_name is None: + raise AttributeError(f"module {__name__!r} has no attribute {name!r}") + + from importlib import import_module + + try: + module = import_module(module_name, __name__) + except ImportError as exc: + package, extra = _BACKEND_REQUIREMENTS[module_name] + raise ImportError( + f"{name} requires {package}, which is not installed. " + f"Install it with: pip install 'infinite-training[{extra}]'" + ) from exc + + value = getattr(module, name) + # Cache on the package so later lookups skip __getattr__ entirely. + globals()[name] = value + return value + + +def __dir__() -> list[str]: + return sorted(__all__) diff --git a/infinite_training/_base.py b/infinite_training/_base.py new file mode 100644 index 0000000..4598704 --- /dev/null +++ b/infinite_training/_base.py @@ -0,0 +1,152 @@ +"""Framework-independent training session state and loop. + +Internal module: nothing here is part of the public API. + +Everything that does not depend on Keras or PyTorch lives here — the target and +timeout bookkeeping, the value history, the checkpoint paths, and the round loop +itself. A backend supplies four small hooks describing how to run one round and +how to move weights in and out of its own model objects. +""" + +from __future__ import annotations + +import logging +import math +import time +from typing import Any + +import numpy as np + +from . import _storage +from .target import Target + +logger = logging.getLogger(__name__) + +__all__ = ["BaseTrainer"] + + +class BaseTrainer: + """Shared state and round loop for every backend. + + Subclasses implement :meth:`_run_round`, :meth:`_current_weights`, + :meth:`_adopt_best_weights` and :meth:`_best_weights_snapshot`. + """ + + def __init__( + self, + target: Target | None, + timeout: float, + best_weights_path: str, + last_weights_path: str, + best_value_path: str, + value_history_path: str, + ) -> None: + # Target() as a default argument would be shared by every instance. + self.target = target if target is not None else Target() + self.timeout = timeout + + self.best_weights_path = best_weights_path + self.last_weights_path = last_weights_path + self.best_value_path = best_value_path + self.value_history_path = value_history_path + + self.best_value = _storage.load_scalar( + self.best_value_path, self.target.worst_possible_value + ) + self.value_history = _storage.load_series(self.value_history_path) + + # ------------------------------------------------------------------ + # Derived state + # ------------------------------------------------------------------ + @property + def last_value(self) -> float | None: + """Value from the most recent round, or ``None`` before any round runs. + + Derived from :attr:`value_history`, so it survives a restart. + """ + if self.value_history is None or len(self.value_history) == 0: + return None + return float(self.value_history[-1]) + + @property + def rounds_completed(self) -> int: + """How many training rounds have been recorded, across all sessions.""" + return 0 if self.value_history is None else int(len(self.value_history)) + + # ------------------------------------------------------------------ + # Backend hooks + # ------------------------------------------------------------------ + def _run_round(self, *args: Any, **kwargs: Any) -> float: + """Run one training round and return the watched value.""" + raise NotImplementedError + + def _current_weights(self) -> Any: + """Return the live model's weights in this backend's own format.""" + raise NotImplementedError + + def _adopt_best_weights(self) -> None: + """Copy the live model's weights into the shadow best-model.""" + raise NotImplementedError + + def _best_weights_snapshot(self) -> Any: + """Return the shadow best-model's weights.""" + raise NotImplementedError + + # ------------------------------------------------------------------ + # The loop + # ------------------------------------------------------------------ + def _run_session(self, *args: Any, **kwargs: Any) -> None: + """Run rounds until the target, the timeout, or ``Ctrl+C`` ends it. + + Whichever way it ends, checkpoints are written before returning. + """ + start = time.time() + stop_reason = "target reached" + try: + while True: + value = self._run_round(*args, **kwargs) + + self.value_history = np.append(self.value_history, value) + + if self.target.is_improvement(value, self.best_value): + self.best_value = value + self._adopt_best_weights() + + if self.target.is_reached(value): + break + if (time.time() - start) > self.timeout: + stop_reason = "timeout reached" + break + except KeyboardInterrupt: + stop_reason = "interrupted by user" + logger.warning("Training interrupted; saving checkpoints before exit.") + + self.last_weights = self._current_weights() + self.best_weights = self._best_weights_snapshot() + self.save() + logger.info( + "Training stopped (%s) after %d round(s); best %s=%s", + stop_reason, + self.rounds_completed, + self.target.name, + self.best_value, + ) + + # ------------------------------------------------------------------ + # Persistence + # ------------------------------------------------------------------ + def save(self) -> None: + """Write the best weights, last weights, best value and history to disk.""" + self._save_weights(self.best_weights_path, self.best_weights) + self._save_weights(self.last_weights_path, self.last_weights) + _storage.save_object(self.best_value_path, self.best_value) + _storage.save_object(self.value_history_path, self.value_history) + + def _save_weights(self, path: str, weights: Any) -> None: + """Persist weights in whatever shape this backend uses.""" + _storage.save_weights(path, weights) + + +def resolve_timeout(timeout: float | None) -> float: + """Treat ``None`` as no timeout, so callers can pass it through.""" + return math.inf if timeout is None else timeout diff --git a/infinite_training/_storage.py b/infinite_training/_storage.py index 78d1a03..826893e 100644 --- a/infinite_training/_storage.py +++ b/infinite_training/_storage.py @@ -19,7 +19,14 @@ T = TypeVar("T") -__all__ = ["load_or_default", "save_object", "save_weights", "load_weights"] +__all__ = [ + "load_or_default", + "save_object", + "save_weights", + "load_weights", + "load_state_dict", + "save_state_dict", +] def load_or_default(path: str, default_factory: Callable[[], T]) -> Any: @@ -71,3 +78,24 @@ def save_object(path: str, value: Any) -> None: def save_weights(path: str, weights: list[np.ndarray]) -> None: """Persist a Keras weight list.""" save_object(path, weights) + + +def save_state_dict(path: str, state: dict[str, np.ndarray]) -> None: + """Persist a PyTorch state dict that has been converted to NumPy arrays. + + Stored as a plain ``dict`` of arrays rather than via ``torch.save`` so that + a checkpoint stays readable without PyTorch, and so that a model trained on + a GPU can be resumed on a CPU without a device map. + """ + save_object(path, state) + + +def load_state_dict( + path: str, default_factory: Callable[[], dict[str, np.ndarray]] +) -> dict[str, np.ndarray]: + """Load a NumPy state dict, falling back to ``default_factory``.""" + if not os.path.exists(path): + return default_factory() + loaded = np.load(path, allow_pickle=True) + # Saved as a zero-dimensional object array wrapping the dict. + return dict(loaded.item()) diff --git a/infinite_training/torch_trainer.py b/infinite_training/torch_trainer.py new file mode 100644 index 0000000..a0b7f45 --- /dev/null +++ b/infinite_training/torch_trainer.py @@ -0,0 +1,222 @@ +"""Resumable, target-driven training loop for PyTorch models.""" + +from __future__ import annotations + +import copy +from collections.abc import Callable, Mapping +from typing import Any + +import numpy as np +import torch +from torch import nn + +from . import _storage +from ._base import BaseTrainer, resolve_timeout +from .target import Target + +__all__ = ["TorchTrainer"] + + +def _to_numpy(state: Mapping[str, torch.Tensor]) -> dict[str, np.ndarray]: + """Convert a state dict to plain arrays, detached from graph and device.""" + return {key: value.detach().cpu().numpy() for key, value in state.items()} + + +def _to_tensors(state: Mapping[str, np.ndarray]) -> dict[str, torch.Tensor]: + """Convert a NumPy state dict back to tensors.""" + return {key: torch.as_tensor(value) for key, value in state.items()} + + +class TorchTrainer(BaseTrainer): + """Train a PyTorch model until a target is met, time runs out, or you interrupt. + + PyTorch has no ``fit``, so you supply the training step yourself: a callable + that runs one round and returns the metrics for it. The trainer calls it in + a loop, records the watched quantity, and keeps a copy of the weights + whenever they are the best seen so far. The loop ends when the + :class:`~infinite_training.Target` is reached, when ``timeout`` seconds have + elapsed, or when you press ``Ctrl+C``. + + Whichever way it ends, the best weights, the most recent weights, and the + full value history are written to disk, so a later session picks up where + the previous one stopped. + + Unlike :class:`~infinite_training.InfiniteTrainer` there is no ``compile()`` + step — the shadow copy holding the best weights is built in the constructor. + + Args: + model: Any ``torch.nn.Module``. It is deep-copied once to hold the best + weights, so it must be copyable. + target: Stopping criterion. Defaults to a :class:`Target` that watches + ``"loss"`` and is never reached, so the session is bounded only by + ``timeout`` or ``Ctrl+C``. + timeout: Wall-clock budget in seconds. The check happens *between* + rounds, so a session can overrun by up to one round. + best_weights_path: Where the best weights are stored. + last_weights_path: Where the most recent weights are stored. + best_value_path: Where the best observed value is stored. + value_history_path: Where the per-round value history is stored. + + Example: + >>> model = nn.Sequential(nn.Linear(2, 1)) + >>> optimizer = torch.optim.Adam(model.parameters()) + >>> loss_fn = nn.MSELoss() + >>> + >>> def step(): + ... 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("loss", smaller_is_better=True, target_value=1e-4), + ... timeout=60, + ... ) + >>> trainer.train(step) + >>> predictions, value = trainer.predict_best(x) + + Note: + Checkpoints are pickled NumPy files, not ``torch.save`` archives. Only + load checkpoints you produced yourself; see the security note in the + README. + """ + + def __init__( + self, + model: nn.Module, + target: Target | None = None, + timeout: float | None = None, + best_weights_path: str = "best_weights.npy", + last_weights_path: str = "last_weights.npy", + best_value_path: str = "best_value.npy", + value_history_path: str = "value_history.npy", + ) -> None: + super().__init__( + target=target, + timeout=resolve_timeout(timeout), + best_weights_path=best_weights_path, + last_weights_path=last_weights_path, + best_value_path=best_value_path, + value_history_path=value_history_path, + ) + + self.model = model + + self.last_weights = _storage.load_state_dict( + self.last_weights_path, lambda: _to_numpy(model.state_dict()) + ) + self.best_weights = _storage.load_state_dict( + self.best_weights_path, lambda: self.last_weights + ) + + # No compile() step in PyTorch, so the shadow copy is built here. + self.best_model = copy.deepcopy(model) + self.model.load_state_dict(_to_tensors(self.last_weights)) + self.best_model.load_state_dict(_to_tensors(self.best_weights)) + + # ------------------------------------------------------------------ + # Training + # ------------------------------------------------------------------ + def train(self, step_fn: Callable[..., Mapping[str, Any]], *args: Any, **kwargs: Any) -> None: + """Call ``step_fn`` repeatedly until the session ends. + + Each call is one round, so ``step_fn`` decides the granularity at which + the target and the timeout are checked — typically one epoch, but a + fixed number of batches works just as well. + + Args: + step_fn: Runs one round and returns a mapping of metric name to + value, for example ``{"loss": 0.31}``. A value may also be a + sequence, in which case the last element is used, matching how + Keras reports a multi-epoch ``History``. + *args: Forwarded to ``step_fn`` on every call. + **kwargs: Forwarded to ``step_fn`` on every call. + + Raises: + RuntimeError: If the watched quantity is absent from the returned + metrics. + TypeError: If ``step_fn`` does not return a mapping. + """ + self._run_session(step_fn, *args, **kwargs) + + def _run_round( + self, step_fn: Callable[..., Mapping[str, Any]], *args: Any, **kwargs: Any + ) -> float: + metrics = step_fn(*args, **kwargs) + return self._read_target_value(metrics) + + def _read_target_value(self, metrics: Any) -> float: + """Extract the watched quantity from a step's metrics, or explain why not.""" + if not isinstance(metrics, Mapping): + raise TypeError( + f"The training step must return a mapping of metric name to value, " + f"for example {{{self.target.name!r}: 0.31}}, but it returned " + f"{type(metrics).__name__}." + ) + try: + value = metrics[self.target.name] + except KeyError: + available = ", ".join(sorted(map(str, metrics))) or "none" + raise RuntimeError( + f"Target {self.target.name!r} is not in the metrics returned by the " + f"training step. Available keys: {available}." + ) from None + # Accept a per-epoch sequence as well as a scalar, mirroring Keras. + if isinstance(value, (list, tuple)) or (isinstance(value, np.ndarray) and value.ndim > 0): + if len(value) == 0: + raise RuntimeError( + f"Target {self.target.name!r} came back as an empty sequence; " + "the training step must report at least one value per round." + ) + value = value[-1] + if isinstance(value, torch.Tensor): + value = value.detach().cpu().item() + return float(value) + + # ------------------------------------------------------------------ + # Backend hooks + # ------------------------------------------------------------------ + def _current_weights(self) -> dict[str, np.ndarray]: + return _to_numpy(self.model.state_dict()) + + def _adopt_best_weights(self) -> None: + self.best_model.load_state_dict(self.model.state_dict()) + + def _best_weights_snapshot(self) -> dict[str, np.ndarray]: + return _to_numpy(self.best_model.state_dict()) + + def _save_weights(self, path: str, weights: dict[str, np.ndarray]) -> None: + _storage.save_state_dict(path, weights) + + # ------------------------------------------------------------------ + # Inference + # ------------------------------------------------------------------ + def predict_best(self, *args: Any, **kwargs: Any) -> tuple[Any, float]: + """Run the best weights in eval mode, without building a graph. + + Returns: + The model output and the best observed value. + """ + return self._infer(self.best_model, *args, **kwargs), self.best_value + + def predict_last(self, *args: Any, **kwargs: Any) -> tuple[Any, float | None]: + """Run the most recent weights in eval mode, without building a graph. + + Returns: + The model output and the most recent value (``None`` if untrained). + """ + return self._infer(self.model, *args, **kwargs), self.last_value + + @staticmethod + def _infer(model: nn.Module, *args: Any, **kwargs: Any) -> Any: + """Forward pass in eval mode, restoring the previous mode afterwards.""" + was_training = model.training + model.eval() + try: + with torch.no_grad(): + return model(*args, **kwargs) + finally: + model.train(was_training) diff --git a/infinite_training/trainer.py b/infinite_training/trainer.py index 45595ef..67140c3 100644 --- a/infinite_training/trainer.py +++ b/infinite_training/trainer.py @@ -4,14 +4,13 @@ import logging import math -import time import warnings from typing import Any -import numpy as np import tensorflow as tf from . import _storage +from ._base import BaseTrainer from .target import Target __all__ = ["InfiniteTrainer", "InfinityTraining"] @@ -29,7 +28,7 @@ } -class InfiniteTrainer: +class InfiniteTrainer(BaseTrainer): """Train a Keras model until a target is met, time runs out, or you interrupt. The trainer calls :meth:`tf.keras.Model.fit` in a loop. After each round it @@ -105,43 +104,22 @@ def __init__( unexpected = ", ".join(sorted(deprecated)) raise TypeError(f"Unexpected keyword argument(s): {unexpected}") - self.model = model - # Target() as a default argument would be shared by every instance. - self.target = target if target is not None else Target() - self.timeout = timeout + super().__init__( + target=target, + timeout=timeout, + best_weights_path=best_weights_path, + last_weights_path=last_weights_path, + best_value_path=best_value_path, + value_history_path=value_history_path, + ) - self.best_weights_path = best_weights_path - self.last_weights_path = last_weights_path - self.best_value_path = best_value_path - self.value_history_path = value_history_path + self.model = model self.last_weights = _storage.load_weights(self.last_weights_path, model.get_weights) self.best_weights = _storage.load_weights(self.best_weights_path, lambda: self.last_weights) - self.best_value = _storage.load_scalar( - self.best_value_path, self.target.worst_possible_value - ) - self.value_history = _storage.load_series(self.value_history_path) self.best_model: tf.keras.Model | None = None - # ------------------------------------------------------------------ - # Derived state - # ------------------------------------------------------------------ - @property - def last_value(self) -> float | None: - """Value from the most recent round, or ``None`` before any round runs. - - Derived from :attr:`value_history`, so it survives a restart. - """ - if self.value_history is None or len(self.value_history) == 0: - return None - return float(self.value_history[-1]) - - @property - def rounds_completed(self) -> int: - """How many training rounds have been recorded, across all sessions.""" - return 0 if self.value_history is None else int(len(self.value_history)) - # ------------------------------------------------------------------ # Keras passthrough # ------------------------------------------------------------------ @@ -170,38 +148,23 @@ def train(self, *args: Any, **kwargs: Any) -> None: if self.best_model is None: raise RuntimeError("compile() must be called before train().") - start = time.time() - stop_reason = "target reached" - try: - while True: - history = self.model.fit(*args, **kwargs) - value = self._read_target_value(history) - - self.value_history = np.append(self.value_history, value) - - if self.target.is_improvement(value, self.best_value): - self.best_value = value - self.best_model.set_weights(self.model.get_weights()) - - if self.target.is_reached(value): - break - if (time.time() - start) > self.timeout: - stop_reason = "timeout reached" - break - except KeyboardInterrupt: - stop_reason = "interrupted by user" - logger.warning("Training interrupted; saving checkpoints before exit.") - - self.last_weights = self.model.get_weights() - self.best_weights = self.best_model.get_weights() - self.save() - logger.info( - "Training stopped (%s) after %d round(s); best %s=%s", - stop_reason, - self.rounds_completed, - self.target.name, - self.best_value, - ) + self._run_session(*args, **kwargs) + + # ------------------------------------------------------------------ + # Backend hooks + # ------------------------------------------------------------------ + def _run_round(self, *args: Any, **kwargs: Any) -> float: + history = self.model.fit(*args, **kwargs) + return self._read_target_value(history) + + def _current_weights(self) -> Any: + return self.model.get_weights() + + def _adopt_best_weights(self) -> None: + self.best_model.set_weights(self.model.get_weights()) + + def _best_weights_snapshot(self) -> Any: + return self.best_model.get_weights() def _read_target_value(self, history: tf.keras.callbacks.History) -> float: """Extract the watched quantity from a ``History``, or explain why not.""" @@ -216,13 +179,6 @@ def _read_target_value(self, history: tf.keras.callbacks.History) -> float: ) from None return float(series[-1]) - def save(self) -> None: - """Write the best weights, last weights, best value and history to disk.""" - _storage.save_weights(self.best_weights_path, self.best_weights) - _storage.save_weights(self.last_weights_path, self.last_weights) - _storage.save_object(self.best_value_path, self.best_value) - _storage.save_object(self.value_history_path, self.value_history) - # ------------------------------------------------------------------ # Inference # ------------------------------------------------------------------ diff --git a/pyproject.toml b/pyproject.toml index 3532415..9e3cde5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,15 +4,28 @@ build-backend = "setuptools.build_meta" [project] name = "infinite_training" -version = "2.1.0" +version = "2.2.0" authors = [ { name="Nguyễn Chí Vỹ", email="vyncint@icloud.com" }, ] -description = "Train Keras models until a target metric, a timeout, or Ctrl+C stops them - and resume where you left off." +description = "Train Keras or PyTorch models until a target metric, a timeout, or Ctrl+C stops them - and resume where you left off." readme = "README.md" # Bounded by TensorFlow, which requires >=3.10 from 2.16 onward. requires-python = ">=3.10" -keywords = ["tensorflow", "keras", "training", "checkpointing", "early-stopping", "machine-learning"] +keywords = [ + "tensorflow", + "keras", + "pytorch", + "torch", + "training", + "checkpointing", + "early-stopping", + "machine-learning", +] +# TensorFlow stays a hard requirement for the 2.x line so that an existing +# `pip install infinite-training` keeps working unchanged. It moves to the +# `tensorflow` extra in 3.0.0; the extra is already declared below so that new +# code can opt into the future spelling today. dependencies = [ "tensorflow", "numpy", @@ -36,6 +49,12 @@ classifiers = [ "Changelog" = "https://github.com/vyncint/infinite-training/blob/main/CHANGELOG.md" [project.optional-dependencies] +tensorflow = [ + "tensorflow", +] +torch = [ + "torch", +] dev = [ "pytest>=7", "ruff>=0.6", diff --git a/tests/conftest.py b/tests/conftest.py index b1e9469..f57466d 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -7,16 +7,33 @@ # Keep TensorFlow quiet; must be set before the first import. os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3") +import importlib.util # noqa: E402 + import numpy as np # noqa: E402 import pytest # noqa: E402 -import tensorflow as tf # noqa: E402 -from infinite_training import InfiniteTrainer # noqa: E402 + +def _installed(name: str) -> bool: + return importlib.util.find_spec(name) is not None + + +HAS_TENSORFLOW = _installed("tensorflow") +HAS_TORCH = _installed("torch") + +# Each backend is optional, so skip the modules whose framework is absent +# rather than failing collection for everyone. +collect_ignore = [] +if not HAS_TENSORFLOW: + collect_ignore += ["test_trainer.py", "test_persistence.py", "test_deprecations.py"] +if not HAS_TORCH: + collect_ignore += ["test_torch_trainer.py"] @pytest.fixture -def model() -> tf.keras.Model: - """A two-input, one-output functional model small enough to fit instantly.""" +def model(): + """A two-input, one-output functional Keras model small enough to fit instantly.""" + import tensorflow as tf + inputs = tf.keras.Input(shape=(2,)) outputs = tf.keras.layers.Dense(1)(inputs) return tf.keras.Model(inputs=inputs, outputs=outputs) @@ -43,7 +60,8 @@ def paths(tmp_path) -> dict[str, str]: @pytest.fixture def make_trainer(model, paths): - """Build a compiled trainer, overriding any constructor argument.""" + """Build a compiled Keras trainer, overriding any constructor argument.""" + from infinite_training import InfiniteTrainer def _factory(**kwargs) -> InfiniteTrainer: options = {**paths, **kwargs} @@ -52,3 +70,64 @@ def _factory(**kwargs) -> InfiniteTrainer: return trainer return _factory + + +# ---------------------------------------------------------------------- +# PyTorch fixtures +# ---------------------------------------------------------------------- +@pytest.fixture +def torch_model(): + """A two-input, one-output PyTorch model with deterministic initial weights.""" + import torch + from torch import nn + + torch.manual_seed(0) + return nn.Sequential(nn.Linear(2, 1)) + + +@pytest.fixture +def torch_dataset(): + """The same tiny regression problem, as tensors.""" + import torch + + x = torch.tensor([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]]) + y = torch.tensor([[0.0], [1.0], [1.0], [0.0]]) + return x, y + + +@pytest.fixture +def make_torch_trainer(torch_model, paths): + """Build a TorchTrainer, overriding any constructor argument.""" + from infinite_training import TorchTrainer + + def _factory(**kwargs) -> TorchTrainer: + options = {**paths, **kwargs} + return TorchTrainer(model=options.pop("model", torch_model), **options) + + return _factory + + +@pytest.fixture +def make_step(torch_model, torch_dataset): + """A real gradient-descent step over the tiny dataset.""" + import torch + from torch import nn + + x, y = torch_dataset + + def _factory(model=None, lr: float = 0.1): + model = model if model is not None else torch_model + optimizer = torch.optim.SGD(model.parameters(), lr=lr) + 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()} + + return step + + return _factory diff --git a/tests/test_lazy_imports.py b/tests/test_lazy_imports.py new file mode 100644 index 0000000..97030a4 --- /dev/null +++ b/tests/test_lazy_imports.py @@ -0,0 +1,135 @@ +"""The two backends must stay independent of each other. + +Installing only TensorFlow, or only PyTorch, has to be enough. These tests run +each import in a subprocess so that modules already loaded by the rest of the +suite cannot mask a regression. +""" + +from __future__ import annotations + +import importlib.util +import os +import subprocess +import sys +import textwrap + +import pytest + +HAS_TENSORFLOW = importlib.util.find_spec("tensorflow") is not None +HAS_TORCH = importlib.util.find_spec("torch") is not None + +BLOCK_MODULE = """ +import sys + +class _Blocker: + def find_spec(self, name, path=None, target=None): + if name == {blocked!r} or name.startswith({blocked!r} + "."): + raise ImportError("simulated: {blocked} is not installed") + return None + +sys.meta_path.insert(0, _Blocker()) +""" + + +def run_snippet(*parts: str) -> subprocess.CompletedProcess[str]: + """Execute the concatenated parts in a clean interpreter. + + Each part is dedented on its own, so a flush-left prefix can be combined + with an indented literal without confusing ``textwrap.dedent``. + """ + code = "".join(textwrap.dedent(part) for part in parts) + return subprocess.run( + [sys.executable, "-c", code], + capture_output=True, + text=True, + env={**os.environ, "TF_CPP_MIN_LOG_LEVEL": "3"}, + ) + + +def test_importing_the_package_loads_no_framework(): + result = run_snippet( + """ + import sys + import infinite_training + assert infinite_training.Target is not None + print("tensorflow" in sys.modules, "torch" in sys.modules) + """ + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "False False" + + +@pytest.mark.skipif(not HAS_TORCH, reason="requires torch") +def test_torch_trainer_does_not_import_tensorflow(): + result = run_snippet( + """ + import sys + from infinite_training import TorchTrainer + assert TorchTrainer is not None + print("tensorflow" in sys.modules) + """ + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "False" + + +@pytest.mark.skipif(not HAS_TENSORFLOW, reason="requires tensorflow") +def test_keras_trainer_does_not_import_torch(): + result = run_snippet( + """ + import sys + from infinite_training import InfiniteTrainer + assert InfiniteTrainer is not None + print("torch" in sys.modules) + """ + ) + assert result.returncode == 0, result.stderr + assert result.stdout.strip() == "False" + + +def test_missing_torch_gives_an_actionable_error(): + result = run_snippet( + BLOCK_MODULE.format(blocked="torch"), + """ + try: + from infinite_training import TorchTrainer + except ImportError as exc: + print(exc) + else: + raise AssertionError("expected ImportError") + """, + ) + assert result.returncode == 0, result.stderr + assert "requires torch" in result.stdout + assert "infinite-training[torch]" in result.stdout + + +def test_missing_tensorflow_gives_an_actionable_error(): + result = run_snippet( + BLOCK_MODULE.format(blocked="tensorflow"), + """ + try: + from infinite_training import InfiniteTrainer + except ImportError as exc: + print(exc) + else: + raise AssertionError("expected ImportError") + """, + ) + assert result.returncode == 0, result.stderr + assert "requires tensorflow" in result.stdout + assert "infinite-training[tensorflow]" in result.stdout + + +def test_unknown_attribute_still_raises_attribute_error(): + import infinite_training + + with pytest.raises(AttributeError, match="no attribute 'NotATrainer'"): + getattr(infinite_training, "NotATrainer") # noqa: B009 + + +def test_dir_lists_the_public_api(): + import infinite_training + + assert "TorchTrainer" in dir(infinite_training) + assert "InfiniteTrainer" in dir(infinite_training) diff --git a/tests/test_torch_trainer.py b/tests/test_torch_trainer.py new file mode 100644 index 0000000..3aa33fc --- /dev/null +++ b/tests/test_torch_trainer.py @@ -0,0 +1,314 @@ +"""Tests for the PyTorch backend: loop, metrics contract, persistence, inference.""" + +from __future__ import annotations + +import math +import os + +import numpy as np +import pytest +import torch +from torch import nn + +from infinite_training import Target, TorchTrainer + + +class TestStopConditions: + def test_stops_once_the_target_is_reached(self, make_torch_trainer, make_step): + # target_value=inf while minimising is satisfied by any finite loss, + # so the loop must exit after exactly one round. + trainer = make_torch_trainer(target=Target("loss", True, math.inf)) + trainer.train(make_step()) + assert trainer.rounds_completed == 1 + + def test_stops_on_timeout(self, make_torch_trainer, make_step): + # A target that can never be reached leaves the timeout as the only exit. + trainer = make_torch_trainer(target=Target("loss", True, -math.inf), timeout=0) + trainer.train(make_step()) + # The timeout is checked between rounds, so exactly one round runs. + assert trainer.rounds_completed == 1 + + def test_runs_many_rounds_until_the_target_is_met(self, make_torch_trainer, make_step): + # A single Linear(2, 1) cannot separate XOR, so its MSE floor is 0.25. + # 0.26 is reached after roughly twenty rounds; the timeout is only a + # backstop so a regression here fails instead of hanging. + trainer = make_torch_trainer(target=Target("loss", True, 0.26), timeout=60) + trainer.train(make_step(lr=0.5)) + assert trainer.rounds_completed > 1 + assert trainer.best_value < 0.26 + + def test_timeout_defaults_to_unbounded(self, make_torch_trainer): + trainer = make_torch_trainer() + assert trainer.timeout == math.inf + + def test_keyboard_interrupt_still_checkpoints(self, make_torch_trainer, paths): + trainer = make_torch_trainer() + + calls = {"n": 0} + + def step() -> dict[str, float]: + calls["n"] += 1 + if calls["n"] == 3: + raise KeyboardInterrupt + return {"loss": 1.0 / calls["n"]} + + trainer.train(step) + + assert trainer.rounds_completed == 2 + assert os.path.exists(paths["best_weights_path"]) + assert os.path.exists(paths["value_history_path"]) + + +class TestBestWeightTracking: + def test_tracks_the_best_value(self, make_torch_trainer): + trainer = make_torch_trainer(target=Target("loss", True, 0.05)) + values = iter([0.5, 0.9, 0.2, 0.7, 0.01]) + trainer.train(lambda: {"loss": next(values)}) + + assert trainer.best_value == pytest.approx(0.01) + assert trainer.rounds_completed == 5 + + def test_best_weights_are_not_the_last_weights_when_training_worsens( + self, make_torch_trainer, torch_model + ): + trainer = make_torch_trainer(target=Target("loss", True, -math.inf), timeout=0) + + def step() -> dict[str, float]: + # Move the weights somewhere obvious after the value is recorded. + with torch.no_grad(): + for param in torch_model.parameters(): + param.fill_(7.0) + return {"loss": 0.5} + + trainer.train(step) + + # The single round improved on the initial +inf, so best adopted it. + best = trainer.best_model.state_dict() + assert all(torch.all(value == 7.0) for value in best.values()) + + def test_maximised_target_keeps_the_largest_value(self, make_torch_trainer): + trainer = make_torch_trainer( + target=Target("acc", smaller_is_better=False, target_value=0.99) + ) + values = iter([0.1, 0.8, 0.4, 0.995]) + trainer.train(lambda: {"acc": next(values)}) + + assert trainer.best_value == pytest.approx(0.995) + + +class TestMetricsContract: + def test_missing_target_key_explains_what_is_available(self, make_torch_trainer): + trainer = make_torch_trainer(target=Target("val_loss")) + + with pytest.raises(RuntimeError, match="not in the metrics"): + trainer.train(lambda: {"loss": 0.5, "acc": 0.9}) + + def test_missing_target_key_lists_the_keys(self, make_torch_trainer): + trainer = make_torch_trainer(target=Target("val_loss")) + + with pytest.raises(RuntimeError, match="acc, loss"): + trainer.train(lambda: {"loss": 0.5, "acc": 0.9}) + + def test_non_mapping_return_is_a_clear_type_error(self, make_torch_trainer): + trainer = make_torch_trainer() + + with pytest.raises(TypeError, match="must return a mapping"): + trainer.train(lambda: 0.5) + + def test_sequence_value_uses_the_last_element(self, make_torch_trainer): + trainer = make_torch_trainer(target=Target("loss", True, 0.15)) + trainer.train(lambda: {"loss": [0.9, 0.5, 0.1]}) + + assert trainer.best_value == pytest.approx(0.1) + + def test_empty_sequence_is_rejected(self, make_torch_trainer): + trainer = make_torch_trainer() + + with pytest.raises(RuntimeError, match="empty sequence"): + trainer.train(lambda: {"loss": []}) + + def test_tensor_value_is_accepted(self, make_torch_trainer): + trainer = make_torch_trainer(target=Target("loss", True, 0.5)) + trainer.train(lambda: {"loss": torch.tensor(0.25)}) + + assert trainer.best_value == pytest.approx(0.25) + + def test_tensor_value_does_not_retain_the_graph(self, make_torch_trainer, torch_model): + trainer = make_torch_trainer(target=Target("loss", True, math.inf)) + + def step(): + # A tensor that requires grad would keep the graph alive if stored. + loss = (torch_model(torch.tensor([[1.0, 1.0]])) ** 2).sum() + assert loss.requires_grad + return {"loss": loss} + + trainer.train(step) + assert isinstance(trainer.best_value, float) + + def test_numpy_scalar_value_is_accepted(self, make_torch_trainer): + trainer = make_torch_trainer(target=Target("loss", True, 0.5)) + trainer.train(lambda: {"loss": np.float32(0.25)}) + + assert trainer.best_value == pytest.approx(0.25) + + +class TestDerivedState: + def test_last_value_is_none_before_training(self, make_torch_trainer): + assert make_torch_trainer().last_value is None + + def test_rounds_completed_is_zero_before_training(self, make_torch_trainer): + assert make_torch_trainer().rounds_completed == 0 + + def test_last_value_tracks_the_most_recent_round(self, make_torch_trainer): + trainer = make_torch_trainer(target=Target("loss", True, 0.05)) + values = iter([0.9, 0.4, 0.01]) + trainer.train(lambda: {"loss": next(values)}) + + assert trainer.last_value == pytest.approx(0.01) + + def test_step_receives_forwarded_arguments(self, make_torch_trainer): + trainer = make_torch_trainer(target=Target("loss", True, math.inf)) + seen = {} + + def step(a, *, b): + seen["args"] = (a, b) + return {"loss": 0.5} + + trainer.train(step, 1, b=2) + assert seen["args"] == (1, 2) + + +class TestPersistence: + def test_checkpoints_are_written(self, make_torch_trainer, make_step, paths): + trainer = make_torch_trainer(target=Target("loss", True, math.inf)) + trainer.train(make_step()) + + for path in paths.values(): + assert os.path.exists(path), path + + def test_parent_directories_are_created(self, make_torch_trainer, make_step, tmp_path): + nested = tmp_path / "runs" / "exp1" + trainer = make_torch_trainer( + target=Target("loss", True, math.inf), + best_weights_path=str(nested / "best.npy"), + last_weights_path=str(nested / "last.npy"), + best_value_path=str(nested / "value.npy"), + value_history_path=str(nested / "history.npy"), + ) + trainer.train(make_step()) + + assert (nested / "best.npy").exists() + + def test_session_resumes_from_disk(self, paths, torch_dataset): + target = Target("loss", True, 0.02) + + torch.manual_seed(0) + first_model = nn.Sequential(nn.Linear(2, 1)) + first = TorchTrainer(model=first_model, target=target, timeout=0, **paths) + values = iter([0.9, 0.4]) + first.train(lambda: {"loss": next(values)}) + + assert first.rounds_completed == 1 + + torch.manual_seed(0) + second_model = nn.Sequential(nn.Linear(2, 1)) + second = TorchTrainer(model=second_model, target=target, **paths) + + # History, best value and last value all carry over. + assert second.rounds_completed == 1 + assert second.best_value == pytest.approx(0.9) + assert second.last_value == pytest.approx(0.9) + + def test_resumed_weights_match_what_was_saved(self, paths, torch_dataset): + torch.manual_seed(0) + model = nn.Sequential(nn.Linear(2, 1)) + trainer = TorchTrainer(model=model, target=Target("loss", True, math.inf), **paths) + + def step() -> dict[str, float]: + with torch.no_grad(): + for param in model.parameters(): + param.fill_(3.0) + return {"loss": 0.5} + + trainer.train(step) + + # A fresh model with different initial weights must be overwritten by + # the checkpoint the constructor loads. + torch.manual_seed(99) + resumed_model = nn.Sequential(nn.Linear(2, 1)) + resumed = TorchTrainer(model=resumed_model, **paths) + + assert all(torch.all(v == 3.0) for v in resumed.model.state_dict().values()) + + def test_checkpoint_survives_integer_buffers(self, paths): + """BatchNorm carries an int64 buffer, which must round-trip unchanged.""" + torch.manual_seed(0) + model = nn.Sequential(nn.Linear(2, 2), nn.BatchNorm1d(2)) + trainer = TorchTrainer(model=model, target=Target("loss", True, math.inf), **paths) + trainer.train(lambda: {"loss": 0.5}) + + resumed = TorchTrainer(model=nn.Sequential(nn.Linear(2, 2), nn.BatchNorm1d(2)), **paths) + original = model.state_dict()["1.num_batches_tracked"] + restored = resumed.model.state_dict()["1.num_batches_tracked"] + + assert restored.dtype == original.dtype + assert torch.equal(restored, original) + + +class TestInference: + def test_predict_best_returns_output_and_value(self, make_torch_trainer, torch_dataset): + x, _ = torch_dataset + trainer = make_torch_trainer(target=Target("loss", True, 0.5)) + trainer.train(lambda: {"loss": 0.25}) + + output, value = trainer.predict_best(x) + + assert output.shape == (4, 1) + assert value == pytest.approx(0.25) + + def test_predict_last_returns_none_value_before_training( + self, make_torch_trainer, torch_dataset + ): + x, _ = torch_dataset + output, value = make_torch_trainer().predict_last(x) + + assert output.shape == (4, 1) + assert value is None + + def test_inference_does_not_build_a_graph(self, make_torch_trainer, torch_dataset): + x, _ = torch_dataset + output, _ = make_torch_trainer().predict_last(x) + + assert not output.requires_grad + + def test_inference_restores_the_previous_training_mode( + self, make_torch_trainer, torch_dataset, torch_model + ): + x, _ = torch_dataset + trainer = make_torch_trainer() + + torch_model.train() + trainer.predict_last(x) + assert torch_model.training + + torch_model.eval() + trainer.predict_last(x) + assert not torch_model.training + + +class TestNoCompileStep: + def test_best_model_is_ready_without_compile(self, make_torch_trainer): + # Unlike the Keras trainer there is no compile(); the shadow copy is + # built in the constructor, so inference works immediately. + assert make_torch_trainer().best_model is not None + + def test_best_model_is_a_distinct_object(self, make_torch_trainer, torch_model): + trainer = make_torch_trainer() + assert trainer.best_model is not torch_model + + with torch.no_grad(): + for param in torch_model.parameters(): + param.fill_(5.0) + + # Mutating the live model must not reach through to the shadow copy. + assert not any(torch.all(v == 5.0) for v in trainer.best_model.state_dict().values()) From 65dc32de4a105109215663b2f55df26375f8a5f3 Mon Sep 17 00:00:00 2001 From: Vyncint Ng <115854244+vyncint@users.noreply.github.com> Date: Tue, 28 Jul 2026 06:04:05 +0700 Subject: [PATCH 2/2] ci: install CPU-only torch so it can coexist with TensorFlow The test job segfaulted (exit 139) as soon as a torch test ran after the Keras ones. The traceback is entirely inside torch._dynamo, and the loaded extension modules include cuda.bindings.*: the default PyTorch wheel bundles CUDA libraries that clash with the ones TensorFlow loads once both are imported into the same process. The runners have no GPU, so the CPU wheel is the correct choice regardless, and it is a much smaller download. The backend-isolation jobs already passed - they only ever import one framework - but they now use the CPU build too, for consistency and speed. Also documents the hazard in the README, since a user installing both frameworks can hit it independently of this package. Co-Authored-By: Claude Opus 5 --- .github/workflows/ci.yml | 14 +++++++++++--- README.md | 8 ++++++++ 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 6023b54..3587dab 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -47,7 +47,12 @@ jobs: run: | uv venv --python ${{ matrix.python-version }} source .venv/bin/activate - uv pip install -e ".[dev,torch]" + 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" @@ -65,7 +70,8 @@ jobs: matrix: include: - backend: torch - install: "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: @@ -81,7 +87,9 @@ jobs: # --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 ${{ matrix.install }} + 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 diff --git a/README.md b/README.md index 85d57ae..b50d561 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,14 @@ actually use — importing `TorchTrainer` never loads TensorFlow, and vice versa > 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