Repository navigation
Adds a PyTorch backend. Backward compatible: the Keras API is unchanged and no existing checkpoint needs converting.
TorchTrainer
PyTorch has no fit, so you hand over the training step and the trainer calls it once per round:
from infinite_training import TorchTrainer, Target
def step() -> dict[str, float]:
optimizer.zero_grad()
loss = loss_fn(model(x), y)
loss.backward()
optimizer.step()
return {"loss": loss.item()}
trainer = TorchTrainer(model=model, target=Target("loss", True, 1e-4), timeout=60)
trainer.train(step)
predictions, best_loss = trainer.predict_best(x)Everything around the step — the Target, the round loop, the timeout, Ctrl+C handling, checkpointing and resume — is shared with the Keras trainer. One round is one call to your step function, so you choose the granularity at which the target and timeout are checked.
There is no compile(): the shadow copy holding the best weights is built in the constructor, so inference works straight away.
Install
pip install "infinite-training[torch]"
Checkpoints are NumPy state dicts rather than torch.save archives, so a checkpoint written on a GPU resumes on a CPU with no device map, and stays readable without PyTorch installed.
Backends are now imported lazily
import infinite_training no longer imports TensorFlow, and referencing TorchTrainer never imports it either — one framework installed is enough. Asking for a trainer whose framework is missing raises an ImportError that names the extra:
TorchTrainer requires torch, which is not installed.
Install it with: pip install 'infinite-training[torch]'
TensorFlow remains a hard dependency for the whole 2.x line, so no existing pip install infinite-training breaks. It moves to the [tensorflow] extra in 3.0.0. A torch-only user can opt out today with pip install --no-deps infinite-training plus numpy torch.
Heads up if you install both frameworks
The default PyTorch wheel bundles CUDA libraries that clash with the ones TensorFlow loads, and importing both into one process can segfault the interpreter. This is not specific to this package — import tensorflow; import torch is enough to trigger it on Linux. If you need both, install the CPU build:
pip install torch --index-url https://download.pytorch.org/whl/cpu
The lazy imports mean you will not hit this while using a single backend.
Also in this release
- The loop, target and timeout bookkeeping, value history and checkpoint paths moved into an internal
_basemodule shared by both backends. No public name changed. examples/mnist_torch.py, the PyTorch counterpart of the Keras MNIST example.- 93 tests, 38 of them new, plus a CI job that installs each framework on its own and runs the suite to keep the lazy-import claim honest.
Full changelog: https://github.com/vyncint/infinite-training/blob/main/CHANGELOG.md