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.
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 |
- Installation
- Quick start
- How it works
- Resuming a session
- API reference
- Choosing a target
- Security note on checkpoints
- Migrating from 2.0.x
- Contributing
- License
pip install infinite-trainingFor the PyTorch backend:
pip install "infinite-training[torch]"To run the bundled Keras example as well:
pip install "infinite-training[example]"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-trainingitself, so that an existingpip install infinite-trainingkeeps working unchanged. If you only want PyTorch you can skip it today withpip install --no-deps infinite-trainingfollowed bypip install numpy torch. In 3.0.0 TensorFlow moves to thetensorflowextra 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 torchis 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.
import numpy as np
import tensorflow as tf
from infinite_training import InfiniteTrainer, Target
x = np.random.rand(256, 2)
y = x.sum(axis=1, keepdims=True)
model = tf.keras.Sequential(
[
tf.keras.layers.Input(shape=(2,)),
tf.keras.layers.Dense(16, activation="relu"),
tf.keras.layers.Dense(1),
]
)
trainer = InfiniteTrainer(
model=model,
target=Target(name="loss", smaller_is_better=True, target_value=1e-4),
timeout=60, # seconds
)
trainer.compile(optimizer="adam", loss="mse")
trainer.train(x, y, epochs=5, verbose=0) # loops until the target, the timeout, or Ctrl+C
predictions, best_loss = trainer.predict_best(x, verbose=0)
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 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.
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.
Each iteration of the loop is one call to Model.fit:
- Fit — call
Model.fit(*args, **kwargs)once. With the defaultepochs=1, one round is one epoch; passepochs=5to check the target less often and reduce overhead. - Read — take the final value of
target.namefrom the returnedHistoryand append it tovalue_history. - Remember — if it beats the best value so far, copy the current weights into the best-weights model.
- Decide — stop if the target is reached, or if
timeoutseconds have elapsed. - Checkpoint — when the loop ends, for any reason, write best weights, last weights, best value and history to disk.
Ctrl+C is caught and treated as a normal stop, so the checkpoint still gets written.
Timeout granularity. The elapsed-time check happens between rounds, never inside
fit. A session can therefore overruntimeoutby up to the duration of one round. Use a smallerepochsfor a tighter bound.
Checkpoints are written to four files in the working directory by default. Point them anywhere you like:
trainer = InfiniteTrainer(
model=model,
best_weights_path="runs/exp1/best_weights.npy",
last_weights_path="runs/exp1/last_weights.npy",
best_value_path="runs/exp1/best_value.npy",
value_history_path="runs/exp1/value_history.npy",
)Parent directories are created automatically. Construct a trainer with the same paths and it reloads the previous state:
trainer = InfiniteTrainer(model=build_model(), best_weights_path="runs/exp1/best_weights.npy", ...)
trainer.rounds_completed # e.g. 12 — carried over from the previous session
trainer.last_value # available immediately, no re-training needed
trainer.compile(optimizer="adam", loss="mse") # restores both weight sets
trainer.train(x, y) # continues, appending to the same historycompile() loads the last weights into model and the best weights into best_model, so training continues from where it stopped while the best result stays intact.
The stopping criterion.
| Argument | Type | Default | Description |
|---|---|---|---|
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.
| Argument | Type | Default | Description |
|---|---|---|---|
model |
tf.keras.Model |
required | Must be clonable by tf.keras.models.clone_model. |
target |
Target |
Target() |
Stopping criterion. |
timeout |
float |
math.inf |
Wall-clock budget in seconds, checked between rounds. |
best_weights_path |
str |
"optimize_weight.npy" |
Best weights checkpoint. |
last_weights_path |
str |
"last_weight.npy" |
Most recent weights checkpoint. |
best_value_path |
str |
"optimize_value.npy" |
Best observed value. |
value_history_path |
str |
"list_value.npy" |
Per-round value history. |
| Method | Description |
|---|---|
compile(*args, **kwargs) |
Forwards to Model.compile, then restores both weight sets. Call before train. |
train(*args, **kwargs) |
Forwards to Model.fit, looping until the target, the timeout, or Ctrl+C. Always checkpoints. |
save() |
Write all four checkpoints immediately. |
predict_best(*args, **kwargs) |
(predictions, best_value) using the best weights. |
predict_last(*args, **kwargs) |
(predictions, last_value) using the most recent weights. |
show_result(*args, **kwargs) |
Print both sets of predictions side by side. |
| Property | Description |
|---|---|
best_value |
Best value observed, across all sessions. |
last_value |
Value from the most recent round, or None before any round has run. |
value_history |
NumPy array with one entry per round. |
rounds_completed |
Number of recorded rounds. |
best_weights / last_weights |
Weight lists. |
best_model |
Shadow model holding the best weights (available after compile). |
| 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.
Target("loss", smaller_is_better=True, target_value=0.01) # stop below 0.01 loss
Target("val_accuracy", smaller_is_better=False, target_value=0.98) # stop above 98%
Target("loss", smaller_is_better=True, target_value=0.0) # stop at a negative loss
Target() # no target: timeout or Ctrl+C onlyA val_* target requires passing validation_data to train, otherwise the key is absent from the history and the trainer raises a RuntimeError listing the keys that are available.
Checkpoints are .npy files written with allow_pickle=True, which is required because Keras weight lists are ragged. Loading a checkpoint executes pickle, so only load checkpoint files you produced yourself. Never point a trainer at checkpoint paths supplied by an untrusted party.
Version 2.1.0 is backward compatible: the 2.0.x API still works and emits DeprecationWarning. The old names will be removed in 3.0.0.
| 2.0.x | 2.1.0 |
|---|---|
InfinityTraining |
InfiniteTrainer |
optimize_weight |
best_weights |
last_weight |
last_weights |
optimize_value |
best_value |
list_value |
value_history |
optimize_model |
best_model |
predict_optimize() |
predict_best() |
optimize_weight_path= |
best_weights_path= |
last_weight_path= |
last_weights_path= |
optimize_value_path= |
best_value_path= |
list_value_path= |
value_history_path= |
Two behaviour fixes may affect you, both listed in the CHANGELOG:
Target(target_value=0)is now honoured. Previously0was treated as "unset", turning the target into an unreachable bound.last_valuereturnsNonebefore the first round instead of raisingAttributeError.
See CONTRIBUTING.md for development setup, tests and coding standards. Bug reports and pull requests are welcome on the issue tracker.
Released under the MIT License.