From a4b47b1649e4849883c04fcba4a2bc2acc817c08 Mon Sep 17 00:00:00 2001 From: vyncint <115854244+vyncint@users.noreply.github.com> Date: Sat, 25 Jul 2026 16:13:17 +0700 Subject: [PATCH 1/2] refactor: fix correctness bugs, add tests, and document the package The package worked for the happy path but had several defects that only appeared at the edges, no CI on commits, and a README that stopped at a single example. This reworks all three while keeping the 2.0.x API working. Correctness fixes: - Target(target_value=0) was discarded. The default was chosen with `target_value or `, so a falsy 0 became an unreachable bound and the session could never stop on that target. - last_value raised AttributeError before the first round, and after an interrupt during it. It is now derived from the persisted history, so it is also correct straight after a restart. - Target() was a mutable default argument shared by every trainer. - train()/predict before compile() raised AttributeError on an internal attribute; they now raise RuntimeError explaining what to call. - A target naming an absent metric raised a bare KeyError; it now lists the available keys and points at validation_data for val_* names. - Checkpoint parent directories were not created, so nested paths failed. - Restored values stayed as 0-d object arrays instead of floats/lists. - requires-python said >=3.9, but TensorFlow has needed >=3.10 since 2.16. - Dropped the unused polars dependency. Structure and tooling: - Split into target.py, trainer.py and _storage.py; public imports unchanged. - Renamed InfinityTraining -> InfiniteTrainer and optimize_* -> best_*, with the old names kept as aliases that emit DeprecationWarning until 3.0.0. - Added a CI workflow running lint, tests on 3.10-3.12 and a packaging check on every push and PR. Tests previously ran only when a release was cut. - The publish workflow now verifies the tag matches the packaged version. - Moved run.py to examples/mnist.py so it no longer ships in the wheel. Tests and docs: - 55 tests, up from 3, covering targets, stop conditions, persistence and resume, error handling, and the deprecation shims. - Rewrote the README with an API reference, a resume guide, a timeout caveat and a checkpoint security note; added CHANGELOG.md and CONTRIBUTING.md. Verified: 55 passed, ruff check and format clean, twine check passed, and the README quick-start and resume flows run as written. --- .github/workflows/ci.yml | 79 +++++++ .github/workflows/python-publish.yml | 12 + CHANGELOG.md | 93 ++++++++ CONTRIBUTING.md | 80 +++++++ README.md | 257 +++++++++++++++++---- examples/mnist.py | 92 ++++++++ infinite_training/__init__.py | 251 +++------------------ infinite_training/_storage.py | 72 ++++++ infinite_training/target.py | 68 ++++++ infinite_training/trainer.py | 323 +++++++++++++++++++++++++++ pyproject.toml | 36 ++- run.py | 57 ----- tests/conftest.py | 54 +++++ tests/test_deprecations.py | 119 ++++++++++ tests/test_infinite_training.py | 48 ---- tests/test_persistence.py | 120 ++++++++++ tests/test_target.py | 94 ++++++++ tests/test_trainer.py | 120 ++++++++++ 18 files changed, 1593 insertions(+), 382 deletions(-) create mode 100644 .github/workflows/ci.yml create mode 100644 CHANGELOG.md create mode 100644 CONTRIBUTING.md create mode 100644 examples/mnist.py create mode 100644 infinite_training/_storage.py create mode 100644 infinite_training/target.py create mode 100644 infinite_training/trainer.py delete mode 100644 run.py create mode 100644 tests/conftest.py create mode 100644 tests/test_deprecations.py delete mode 100644 tests/test_infinite_training.py create mode 100644 tests/test_persistence.py create mode 100644 tests/test_target.py create mode 100644 tests/test_trainer.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000..6cedbc6 --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,79 @@ +name: CI + +on: + push: + branches: [main] + pull_request: + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ci-${{ github.ref }} + cancel-in-progress: true + +jobs: + lint: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Install uv + uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + - name: Lint and format check + run: | + uv venv --python 3.12 + source .venv/bin/activate + uv pip install ruff + ruff check . + ruff format --check . + + test: + runs-on: ubuntu-latest + strategy: + fail-fast: false + matrix: + # TensorFlow wheels are the constraint on the supported range. + python-version: ["3.10", "3.11", "3.12"] + steps: + - uses: actions/checkout@v4 + - name: Install uv + uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + - name: Install package and test dependencies + run: | + uv venv --python ${{ matrix.python-version }} + source .venv/bin/activate + uv pip install -e ".[dev]" + - name: Run tests + env: + TF_CPP_MIN_LOG_LEVEL: "3" + run: | + source .venv/bin/activate + pytest tests -v + + build: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Install uv + uses: astral-sh/setup-uv@v5 + with: + enable-cache: true + - name: Build distributions + run: | + uv venv --python 3.12 + source .venv/bin/activate + uv build + - name: Check metadata + run: | + source .venv/bin/activate + uv pip install twine + twine check dist/* + - uses: actions/upload-artifact@v4 + with: + name: distributions + path: dist/ diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml index bec3cce..a67b0c6 100644 --- a/.github/workflows/python-publish.yml +++ b/.github/workflows/python-publish.yml @@ -28,11 +28,23 @@ jobs: with: enable-cache: true - name: Install dependencies and test + env: + TF_CPP_MIN_LOG_LEVEL: "3" run: | uv venv --python 3.12 source .venv/bin/activate uv pip install -e ".[dev]" pytest tests + - name: Verify the tag matches the packaged version + if: github.event_name == 'release' + run: | + source .venv/bin/activate + package_version="$(python -c 'import infinite_training; print(infinite_training.__version__)')" + tag_version="${GITHUB_REF_NAME#v}" + if [ "$package_version" != "$tag_version" ]; then + echo "::error::Tag $tag_version does not match package version $package_version" + exit 1 + fi - name: Build package run: | source .venv/bin/activate diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000..4880d9e --- /dev/null +++ b/CHANGELOG.md @@ -0,0 +1,93 @@ +# Changelog + +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.1.0] - 2026-07-25 + +Backward compatible with 2.0.0: existing code keeps working and warns where an +API has been renamed. + +### Fixed + +- `Target(target_value=0)` is now honoured. The default was selected with + `target_value or `, so a falsy `0`/`0.0` was replaced by an unreachable + bound and the session could never stop on that target. +- `last_value` no longer raises `AttributeError` when read before training, or + after an interrupt during the first round. It is derived from the persisted + value history, so it also survives a restart. +- `Target()` is no longer used as a default argument value, which meant every + trainer built without an explicit target shared one mutable instance. +- Calling `train()` or an inference method before `compile()` now raises a clear + `RuntimeError` instead of `AttributeError: 'InfiniteTrainer' object has no + attribute 'best_model'`. +- A target naming a metric that is absent from the training history now raises a + `RuntimeError` listing the available keys, instead of a bare `KeyError`. When + the missing key starts with `val_`, the message points at `validation_data`. +- Checkpoint parent directories are created automatically; previously a path + such as `runs/exp1/best.npy` failed unless the directory already existed. +- Restored checkpoint values are converted back to Python floats and plain + weight lists rather than being left as zero-dimensional object arrays. +- `KeyboardInterrupt` is logged instead of silently swallowed. +- `requires-python` corrected from `>=3.9` to `>=3.10`. TensorFlow has required + 3.10 or newer since 2.16, so the old floor advertised an uninstallable + combination. + +### Added + +- `InfiniteTrainer`, the new name for `InfinityTraining`. +- `rounds_completed` property. +- `save()` as a public method for checkpointing on demand. +- `Target.is_improvement()`, `Target.is_reached()` and + `Target.worst_possible_value`, so the comparison rules live with the target. +- `__version__` and `__all__` exports. +- A CI workflow running lint, tests on Python 3.10-3.12, and a packaging check + on every push and pull request. Previously tests ran only when a release was + published. +- The publish workflow verifies that the release tag matches the packaged + version. +- `CONTRIBUTING.md`, `CHANGELOG.md`, and a rewritten `README.md` with an API + reference, a resume guide and a security note. +- Test suite covering targets, stop conditions, persistence and resume, + error handling, and the deprecation shims (55 tests, up from 3). + +### Changed + +- The package is split into `target.py`, `trainer.py` and `_storage.py`. Public + imports are unchanged: `from infinite_training import InfiniteTrainer, Target`. +- Type hints use `float` rather than `np.double` for scalar values. +- Progress and stop reasons are reported through the `logging` module. +- `run.py` moved to `examples/mnist.py` and no longer ships inside the wheel. +- Packaging declares its packages explicitly, so `tests` is no longer picked up + by auto-discovery. + +### Removed + +- The unused `polars` runtime dependency. + +### Deprecated + +The following still work and will be removed in 3.0.0: + +| Deprecated | Replacement | +| --- | --- | +| `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=` | + +## [2.0.0] - 2026-04-03 + +- Upgraded to the latest TensorFlow and switched the release workflow to `uv`. + +[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/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..409d384 --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,80 @@ +# Contributing + +Thanks for your interest in improving Infinite Training. Bug reports, questions +and pull requests are all welcome. + +## Getting set up + +TensorFlow wheels lag the newest Python release, so develop on Python 3.10-3.12. + +```bash +git clone https://github.com/vyncint/infinite-training.git +cd infinite-training + +python3.12 -m venv .venv +source .venv/bin/activate # Windows: .venv\Scripts\activate +pip install -e ".[dev]" +``` + +## Running the checks + +The same three checks run in CI, and all of them should pass before you open a +pull request. + +```bash +pytest tests -v # test suite +ruff check . # lint +ruff format --check . # formatting +``` + +TensorFlow prints a lot of start-up noise; `TF_CPP_MIN_LOG_LEVEL=3` silences it. +The test suite sets this for you. + +To try a change end to end against a real dataset: + +```bash +pip install -e ".[example]" +python examples/mnist.py +``` + +## Making a change + +- **Keep the public API stable.** The 2.0.x names are kept as deprecated + aliases and are covered by `tests/test_deprecations.py`. If you rename + something, add an alias that emits a `DeprecationWarning`, note it in the + deprecation table in `CHANGELOG.md`, and leave the removal for the next major + version. +- **Add a test with every behaviour change.** Bug fixes should come with a test + that fails before the fix. Group tests by behaviour in a `Test*` class and + give each test a name that reads as a sentence. +- **Update the docs in the same change.** New or renamed public API belongs in + the README's API reference; user-visible changes belong in `CHANGELOG.md` + under *Unreleased*. +- **Explain the "why" in comments.** The code should say what it does; comments + are for constraints and reasons that are not obvious from reading it. + +## Project layout + +| Path | Purpose | +| --- | --- | +| `infinite_training/target.py` | `Target` — the stopping criterion and its comparison rules | +| `infinite_training/trainer.py` | `InfiniteTrainer` — the training loop, plus deprecated aliases | +| `infinite_training/_storage.py` | Checkpoint reading and writing (internal) | +| `tests/` | Test suite, one module per concern | +| `examples/` | Runnable examples, not shipped in the wheel | + +## Releasing + +Maintainers only: + +1. Bump `version` in `pyproject.toml` and `__version__` in + `infinite_training/__init__.py`. They must match. +2. Move the *Unreleased* notes in `CHANGELOG.md` under the new version. +3. Publish a GitHub release tagged `vX.Y.Z`. The publish workflow runs the + tests, verifies that the tag matches the packaged version, builds and + uploads to PyPI. + +## Code of conduct + +Be respectful and constructive. Assume good faith, and keep discussion focused +on the work. diff --git a/README.md b/README.md index 20a016c..66729eb 100644 --- a/README.md +++ b/README.md @@ -1,57 +1,222 @@ # Infinite Training -The python package for training ML models base on loss, metrics and timeout. We can create target for training or using Ctrl + C to interrupt the training session. This package can store the last weight and optimize weight for us to continue training after interrupting. +[![CI](https://github.com/vyncint/infinite-training/actions/workflows/ci.yml/badge.svg)](https://github.com/vyncint/infinite-training/actions/workflows/ci.yml) +[![PyPI](https://img.shields.io/pypi/v/infinite-training.svg)](https://pypi.org/project/infinite-training/) +[![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. + +`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. + +--- + +## Contents + +- [Installation](#installation) +- [Quick start](#quick-start) +- [How it works](#how-it-works) +- [Resuming a session](#resuming-a-session) +- [API reference](#api-reference) +- [Choosing a target](#choosing-a-target) +- [Security note on checkpoints](#security-note-on-checkpoints) +- [Migrating from 2.0.x](#migrating-from-20x) +- [Contributing](#contributing) +- [License](#license) + +--- + +## Installation + +```bash +pip install infinite-training +``` + +To run the bundled example as well: -Using as example: ```bash -uv pip install infinite-training +pip install "infinite-training[example]" ``` + +Requires Python 3.10+ and TensorFlow. + +--- + +## Quick start + ```python -""" - Apply in example from https://www.tensorflow.org/datasets/keras_example -""" -import tensorflow_datasets as tfds +import numpy as np import tensorflow as tf -from infinite_training import InfinityTraining, Target - -(ds_train, ds_test), ds_info = tfds.load( - 'mnist', - split=['train', 'test'], - shuffle_files=True, - as_supervised=True, - with_info=True, +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 +) -def normalize_img(image, label): - """Normalizes images: `uint8` -> `float32`.""" - return tf.cast(image, tf.float32) / 255., label - -ds_train = ds_train.map( - normalize_img, num_parallel_calls=tf.data.AUTOTUNE) -ds_train = ds_train.cache() -ds_train = ds_train.shuffle(ds_info.splits['train'].num_examples) -ds_train = ds_train.batch(128) -ds_train = ds_train.prefetch(tf.data.AUTOTUNE) - -ds_test = ds_test.map( - normalize_img, num_parallel_calls=tf.data.AUTOTUNE) -ds_test = ds_test.batch(128) -ds_test = ds_test.cache() -ds_test = ds_test.prefetch(tf.data.AUTOTUNE) - -model = tf.keras.models.Sequential([ - tf.keras.layers.Flatten(input_shape=(28, 28)), - tf.keras.layers.Dense(128, activation='relu'), - tf.keras.layers.Dense(10) -]) - -it = InfinityTraining(model=model, target=Target( - name="val_sparse_categorical_accuracy", smaller_is_better=False, target_value=0.98), timeout=100) -it.compile(optimizer=tf.keras.optimizers.Adam(0.001), - loss=tf.keras.losses.SparseCategoricalCrossentropy( - from_logits=True), - metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],) -it.train(ds_train, validation_data=ds_test) -it.show_result(ds_train) +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. + +--- + +## How it works + +Each iteration of the loop is one call to `Model.fit`: + +1. **Fit** — call `Model.fit(*args, **kwargs)` once. With the default `epochs=1`, one round is one epoch; pass `epochs=5` to check the target less often and reduce overhead. +2. **Read** — take the final value of `target.name` from the returned `History` and append it to `value_history`. +3. **Remember** — if it beats the best value so far, copy the current weights into the best-weights model. +4. **Decide** — stop if the target is reached, or if `timeout` seconds have elapsed. +5. **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 overrun `timeout` by up to the duration of one round. Use a smaller `epochs` for a tighter bound. + +--- + +## Resuming a session + +Checkpoints are written to four files in the working directory by default. Point them anywhere you like: + +```python +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: + +```python +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 history +``` + +`compile()` 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. + +--- + +## API reference + +### `Target` + +The stopping criterion. + +| Argument | Type | Default | Description | +| --- | --- | --- | --- | +| `name` | `str` | `"loss"` | Key to read from the Keras `History`. 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` + +| 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`). | + +--- + +## Choosing a target + +```python +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 only +``` + +A `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. + +--- + +## Security note on checkpoints + +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. + +--- + +## Migrating from 2.0.x + +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](CHANGELOG.md): + +- `Target(target_value=0)` is now honoured. Previously `0` was treated as "unset", turning the target into an unreachable bound. +- `last_value` returns `None` before the first round instead of raising `AttributeError`. + +--- + +## Contributing + +See [CONTRIBUTING.md](CONTRIBUTING.md) for development setup, tests and coding standards. Bug reports and pull requests are welcome on the [issue tracker](https://github.com/vyncint/infinite-training/issues). + +--- + +## License + +Released under the [MIT License](LICENSE). diff --git a/examples/mnist.py b/examples/mnist.py new file mode 100644 index 0000000..1e82348 --- /dev/null +++ b/examples/mnist.py @@ -0,0 +1,92 @@ +"""Train an MNIST classifier until it reaches 98% validation accuracy. + +Adapted from https://www.tensorflow.org/datasets/keras_example. + +Run it with:: + + pip install "infinite_training[example]" + python examples/mnist.py + +Stop it at any point with ``Ctrl+C``: the best and most recent weights are +written to ``models/`` and picked up automatically the next time you run it. +""" + +from __future__ import annotations + +import tensorflow as tf +import tensorflow_datasets as tfds + +from infinite_training import InfiniteTrainer, Target + +CHECKPOINT_DIR = "models" +BATCH_SIZE = 128 + + +def normalize_img(image, label): + """Normalize images: `uint8` -> `float32` in the range [0, 1].""" + return tf.cast(image, tf.float32) / 255.0, label + + +def build_datasets(): + """Load MNIST and apply the standard input pipeline.""" + (ds_train, ds_test), ds_info = tfds.load( + "mnist", + split=["train", "test"], + shuffle_files=True, + as_supervised=True, + with_info=True, + ) + + ds_train = ds_train.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE) + ds_train = ds_train.cache() + ds_train = ds_train.shuffle(ds_info.splits["train"].num_examples) + ds_train = ds_train.batch(BATCH_SIZE) + ds_train = ds_train.prefetch(tf.data.AUTOTUNE) + + ds_test = ds_test.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE) + ds_test = ds_test.batch(BATCH_SIZE) + ds_test = ds_test.cache() + ds_test = ds_test.prefetch(tf.data.AUTOTUNE) + + return ds_train, ds_test + + +def build_model() -> tf.keras.Model: + """A small dense classifier for 28x28 grayscale digits.""" + return tf.keras.models.Sequential( + [ + tf.keras.layers.Input(shape=(28, 28)), + tf.keras.layers.Flatten(), + tf.keras.layers.Dense(128, activation="relu"), + tf.keras.layers.Dense(10), + ] + ) + + +def main() -> None: + ds_train, ds_test = build_datasets() + + trainer = InfiniteTrainer( + model=build_model(), + target=Target( + name="val_sparse_categorical_accuracy", + smaller_is_better=False, + target_value=0.98, + ), + timeout=100, + 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", + ) + trainer.compile( + optimizer=tf.keras.optimizers.Adam(0.001), + loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True), + metrics=[tf.keras.metrics.SparseCategoricalAccuracy()], + ) + trainer.train(ds_train, validation_data=ds_test) + trainer.show_result(ds_train) + + +if __name__ == "__main__": + main() diff --git a/infinite_training/__init__.py b/infinite_training/__init__.py index a578c3b..da80480 100644 --- a/infinite_training/__init__.py +++ b/infinite_training/__init__.py @@ -1,233 +1,34 @@ -import os -import math -from typing import Any, Tuple -import numpy as np -import tensorflow as tf -import time +"""Train Keras 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. -class Target(): - """ - `Target` class: is the format of target that we want to stop the training session. +Typical use:: - Attributes: - name (str): can be (`loss` or any `metric`). - smaller_is_better (bool): (`True` or `False` that means we want `result of name` (lower or greater). + from infinite_training import InfiniteTrainer, Target - target_value (double): stop the training session when: (`smaller_is_better == True` and `result of name < target_value`) or (`smaller_is_better == False` and `result of name > target_value`) - Example: + trainer = InfiniteTrainer( + model=model, + target=Target("val_accuracy", smaller_is_better=False, target_value=0.98), + timeout=600, + ) + trainer.compile(optimizer="adam", loss="mse") + trainer.train(x, y) + predictions, best_value = trainer.predict_best(x) +""" - ```python - target = Target(name = "loss", smaller_is_better = True, target_value = 0.0001) - ``` - That means: the target is training the model until the loss is lower than 0.0001 - ```python - target = Target(name = "acc", smaller_is_better = False, target_value = 0.9) - ``` - That means: the target is training the model until the acc is greater than 0.9 - """ +from __future__ import annotations - def __init__( - self, - name: str = "loss", - smaller_is_better: bool = True, - target_value: np.double = None - ) -> None: - self.name = name - self.smaller_is_better = smaller_is_better - self.target_value = target_value or ( - -math.inf if self.smaller_is_better else math.inf - ) +from .target import Target +from .trainer import InfiniteTrainer, InfinityTraining +__version__ = "2.1.0" -class InfinityTraining: - """ - `InfinityTraining` class: creating training session base on `Target`, timeout or break by `Ctrl + C`. - - Create example training session: - ```python - from infinite_training import Target, InfinityTraining - import tensorflow as tf - - inputs = tf.keras.Input(shape=(2,)) - x = tf.keras.layers.Dense(4)(inputs) - outputs = tf.keras.layers.Dense(2)(x) - model = tf.keras.Model(inputs=inputs, outputs=outputs) - - target = Target(name = "loss", smaller_is_better = True, target_value = 0.0001) - it = InfinityTraining(model = model, target = target, timeout = 10) - ``` - Attributes: - model: tf.keras.Model() is required. - optimize_weight_path: the path store the optimize model weight `default` "optimize_weight.npy". - last_weight_path: the path store the last model weight `default` "last_weight.npy". - optimize_value_path: the path store the optimize value `default` "optimize_value.npy". - list_value_path: the path store the list value `default` "list_value.npy". - target: Target object. - timeout: timeout in seconds. - Default: After the training session has been ended, four files `optimize_weight.npy`, `last_weight.npy`, `optimize_value.npy` and `list_value.npy` will be created in the current folder for continuous training. - - We can also modify these file paths by: - ```python - it = InfinityTraining( - model = model - optimize_weight_path = "result/optimize_weight.npy", - last_weight_path = "result/last_weight.npy", - optimize_value_path = "result/optimize_value.npy", - list_value_path = "result/list_value.npy" - ) - ``` - After creating training session we need to `compile`, `train` and `save`. - ```python - it.compile(optimizer = optimizer, loss = loss, metrics = metrics) - it.train(x = x, y = y, batch_size = batch_size, verbose = verbose) - it.save() - ``` - Finally, we can `predict_optimize`, `predict_last` or `show_result` as: - ```python - it.show_result(x = x) - optimize_result, optimize_value = it.predict_optimize(x = x) - last_result, last_value = it.predict_last(x = x) - ``` - """ - - def __load( - self, - path: str, - default: Any - ) -> Any: - if os.path.exists(path): - return np.load(path, allow_pickle=True) - return default - - def __save(self) -> None: - np.save(self.optimize_weight_path, np.array( - self.optimize_weight, dtype="object"), allow_pickle=True) - np.save(self.last_weight_path, np.array( - self.last_weight, dtype="object"), allow_pickle=True) - np.save(self.optimize_value_path, np.array( - self.optimize_value, dtype="object"), allow_pickle=True) - np.save(self.list_value_path, np.array( - self.list_value, dtype="object"), allow_pickle=True) - - def __init__( - self, - model: tf.keras.Model, - optimize_weight_path: str = "optimize_weight.npy", - last_weight_path: str = "last_weight.npy", - optimize_value_path: str = "optimize_value.npy", - list_value_path: str = "list_value.npy", - target: Target = Target(), - timeout: np.double = math.inf - ) -> None: - self.model = model - self.optimize_weight_path = optimize_weight_path - self.last_weight_path = last_weight_path - self.optimize_value_path = optimize_value_path - self.list_value_path = list_value_path - self.target = target - self.timeout = timeout - self.last_weight = self.__load( - path=self.last_weight_path, - default=self.model.get_weights() - ) - self.optimize_weight = self.__load( - path=self.optimize_weight_path, - default=self.last_weight - ) - self.optimize_value = self.__load( - path=self.optimize_value_path, - default=math.inf if self.target.smaller_is_better else -math.inf - ) - self.list_value = self.__load( - path=self.list_value_path, - default=np.array([]) - ) - - def compile( - self, - *args, - **kwargs - ) -> None: - """ - Similar to `tensorflow.keras.Model.compile()` https://www.tensorflow.org/api_docs/python/tf/keras/Model#compile - """ - self.model.compile(*args, **kwargs) - self.optimize_model = tf.keras.models.clone_model(self.model) - self.optimize_model.set_weights(self.optimize_weight) - self.model.set_weights(self.last_weight) - - def train( - self, - *args, - **kwargs - ) -> None: - """ - Similar to `tensorflow.keras.Model.fit()` https://www.tensorflow.org/api_docs/python/tf/keras/Model#fit - """ - start = time.time() - try: - while True: - history = self.model.fit(*args, **kwargs) - current_value = history.history[self.target.name] - self.last_value = current_value[-1] - self.list_value = np.append(self.list_value, self.last_value) - - if self.target.smaller_is_better: - ok = self.optimize_value > self.last_value - is_break = self.last_value < self.target.target_value - else: - ok = self.optimize_value < self.last_value - is_break = self.last_value > self.target.target_value - - if ok: - self.optimize_value = self.last_value - self.optimize_model.set_weights(self.model.get_weights()) - - if is_break: - break - if (time.time() - start) > self.timeout: - break - - except KeyboardInterrupt: - pass - self.last_weight = self.model.get_weights() - self.optimize_weight = self.optimize_model.get_weights() - self.__save() - - def show_result( - self, - *args, - **kwargs - ) -> None: - """ - Similar to `tensorflow.keras.Model.predict()` https://www.tensorflow.org/api_docs/python/tf/keras/Model#predict - """ - print() - print("==================================================================================================") - print("Optimize result: ", self.optimize_model.predict(*args, **kwargs)) - print(f' with {self.target.name}: ', self.optimize_value) - print("==================================================================================================") - print("Last result: ", self.model.predict(*args, **kwargs)) - print(f' with {self.target.name}: ', self.last_value) - print("==================================================================================================") - - def predict_optimize( - self, - *args, - **kwargs - ) -> Tuple[Any, Any]: - """ - Similar to `tensorflow.keras.Model.predict()` https://www.tensorflow.org/api_docs/python/tf/keras/Model#predict - """ - return (self.optimize_model.predict(*args, **kwargs), self.optimize_value) - - def predict_last( - self, - *args, - **kwargs - ) -> Tuple[Any, Any]: - """ - Similar to `tensorflow.keras.Model.predict()` https://www.tensorflow.org/api_docs/python/tf/keras/Model#predict - """ - return (self.model.predict(*args, **kwargs), self.last_value) \ No newline at end of file +__all__ = [ + "InfiniteTrainer", + "Target", + # Deprecated, removed in 3.0.0. + "InfinityTraining", + "__version__", +] diff --git a/infinite_training/_storage.py b/infinite_training/_storage.py new file mode 100644 index 0000000..24e5fad --- /dev/null +++ b/infinite_training/_storage.py @@ -0,0 +1,72 @@ +"""Checkpoint persistence helpers. + +Internal module: nothing here is part of the public API. + +Checkpoints are stored as ``.npy`` files written with ``allow_pickle=True``, +because Keras weight lists are ragged sequences of arrays that NumPy cannot +represent without object arrays. Loading such a file executes pickle, so a +checkpoint must be treated as trusted input — see the security note in the +README. +""" + +from __future__ import annotations + +import os +from typing import Any, Callable, TypeVar + +import numpy as np + +T = TypeVar("T") + +__all__ = ["load_or_default", "save_object", "save_weights", "load_weights"] + + +def load_or_default(path: str, default_factory: Callable[[], T]) -> Any: + """Load ``path`` if it exists, otherwise build a default. + + The default is produced lazily so that callers never pay for building it + when a checkpoint is present. + """ + if not os.path.exists(path): + return default_factory() + return np.load(path, allow_pickle=True) + + +def load_weights(path: str, default_factory: Callable[[], list[np.ndarray]]) -> list[np.ndarray]: + """Load a Keras weight list, falling back to ``default_factory``.""" + if not os.path.exists(path): + return default_factory() + loaded = np.load(path, allow_pickle=True) + # Saved as an object array; Keras expects a plain list of arrays. + return [np.asarray(w) for w in loaded] + + +def load_scalar(path: str, default: float) -> float: + """Load a scalar checkpoint value as a Python float.""" + if not os.path.exists(path): + return default + return float(np.load(path, allow_pickle=True)) + + +def load_series(path: str) -> np.ndarray: + """Load the recorded value history, or an empty array when absent.""" + if not os.path.exists(path): + return np.array([]) + return np.asarray(np.load(path, allow_pickle=True)) + + +def _ensure_parent_dir(path: str) -> None: + parent = os.path.dirname(path) + if parent: + os.makedirs(parent, exist_ok=True) + + +def save_object(path: str, value: Any) -> None: + """Persist an arbitrary Python object as a ``.npy`` file.""" + _ensure_parent_dir(path) + np.save(path, np.array(value, dtype="object"), allow_pickle=True) + + +def save_weights(path: str, weights: list[np.ndarray]) -> None: + """Persist a Keras weight list.""" + save_object(path, weights) diff --git a/infinite_training/target.py b/infinite_training/target.py new file mode 100644 index 0000000..6d67ead --- /dev/null +++ b/infinite_training/target.py @@ -0,0 +1,68 @@ +"""Stopping criteria for an infinite training session.""" + +from __future__ import annotations + +import math +from dataclasses import dataclass + +__all__ = ["Target"] + + +@dataclass +class Target: + """A stopping criterion evaluated after every training round. + + A target describes *which* quantity to watch, in which *direction* it + improves, and the *value* at which training may stop. + + Args: + name: Key to read from the Keras ``History``. Any loss or metric name + produced by ``Model.fit`` is valid, including validation keys such + as ``"val_loss"`` or ``"val_sparse_categorical_accuracy"``. + smaller_is_better: ``True`` for quantities minimised (losses, error + rates), ``False`` for quantities maximised (accuracy, F1). + target_value: Value at which training stops. When omitted, it defaults + to an unreachable bound (``-inf`` when minimising, ``+inf`` when + maximising), so the session runs until it times out or is + interrupted. + + Examples: + Train until the loss drops below 0.0001:: + + Target(name="loss", smaller_is_better=True, target_value=0.0001) + + Train until accuracy exceeds 0.9:: + + Target(name="acc", smaller_is_better=False, target_value=0.9) + + Train indefinitely, stopping only on timeout or ``Ctrl+C``:: + + Target() + """ + + name: str = "loss" + smaller_is_better: bool = True + target_value: float | None = None + + def __post_init__(self) -> None: + # An explicit 0.0 (or 0) is a legitimate target, so test for None + # rather than truthiness. + if self.target_value is None: + self.target_value = -math.inf if self.smaller_is_better else math.inf + + @property + def worst_possible_value(self) -> float: + """The value an as-yet-unmeasured best score starts from.""" + return math.inf if self.smaller_is_better else -math.inf + + def is_improvement(self, candidate: float, incumbent: float) -> bool: + """Return ``True`` when ``candidate`` beats ``incumbent``.""" + if self.smaller_is_better: + return candidate < incumbent + return candidate > incumbent + + def is_reached(self, value: float) -> bool: + """Return ``True`` when ``value`` satisfies the stopping criterion.""" + if self.smaller_is_better: + return value < self.target_value + return value > self.target_value diff --git a/infinite_training/trainer.py b/infinite_training/trainer.py new file mode 100644 index 0000000..45595ef --- /dev/null +++ b/infinite_training/trainer.py @@ -0,0 +1,323 @@ +"""Resumable, target-driven training loop for Keras models.""" + +from __future__ import annotations + +import logging +import math +import time +import warnings +from typing import Any + +import numpy as np +import tensorflow as tf + +from . import _storage +from .target import Target + +__all__ = ["InfiniteTrainer", "InfinityTraining"] + +logger = logging.getLogger(__name__) + +_SEPARATOR = "=" * 98 + +# Old constructor keyword -> new constructor keyword. +_DEPRECATED_KWARGS = { + "optimize_weight_path": "best_weights_path", + "last_weight_path": "last_weights_path", + "optimize_value_path": "best_value_path", + "list_value_path": "value_history_path", +} + + +class InfiniteTrainer: + """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 + reads the watched quantity from the returned ``History``, records it, 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. + + Args: + model: A compiled-or-not ``tf.keras.Model``. It must be clonable by + ``tf.keras.models.clone_model`` (functional and sequential models + are; some subclassed models are not). + 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 = tf.keras.Sequential([tf.keras.layers.Dense(1)]) + >>> trainer = InfiniteTrainer( + ... model=model, + ... target=Target("loss", smaller_is_better=True, target_value=1e-4), + ... timeout=60, + ... ) + >>> trainer.compile(optimizer="adam", loss="mse") + >>> trainer.train(x, y) + >>> predictions, value = trainer.predict_best(x) + + Note: + Checkpoints are pickled NumPy files. Only load checkpoints you produced + yourself; see the security note in the README. + """ + + def __init__( + self, + model: tf.keras.Model, + target: Target | None = None, + timeout: float = math.inf, + best_weights_path: str = "optimize_weight.npy", + last_weights_path: str = "last_weight.npy", + best_value_path: str = "optimize_value.npy", + value_history_path: str = "list_value.npy", + **deprecated: Any, + ) -> None: + for old, new in _DEPRECATED_KWARGS.items(): + if old not in deprecated: + continue + warnings.warn( + f"{old!r} is deprecated and will be removed in 3.0.0; use {new!r}.", + DeprecationWarning, + stacklevel=2, + ) + # An explicitly passed legacy path wins over the new-style default. + value = deprecated.pop(old) + if old == "optimize_weight_path": + best_weights_path = value + elif old == "last_weight_path": + last_weights_path = value + elif old == "optimize_value_path": + best_value_path = value + else: + value_history_path = value + if deprecated: + 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 + + 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.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 + # ------------------------------------------------------------------ + def compile(self, *args: Any, **kwargs: Any) -> None: + """Compile the model and prepare the shadow copy holding the best weights. + + Accepts the same arguments as :meth:`tf.keras.Model.compile`. Must be + called before :meth:`train`. + """ + self.model.compile(*args, **kwargs) + self.best_model = tf.keras.models.clone_model(self.model) + self.best_model.set_weights(self.best_weights) + self.model.set_weights(self.last_weights) + + def train(self, *args: Any, **kwargs: Any) -> None: + """Repeatedly call :meth:`tf.keras.Model.fit` until the session ends. + + Accepts the same arguments as ``fit``. Each iteration is one ``fit`` + call, so ``epochs`` (default 1) controls the granularity at which the + target and timeout are checked. + + Raises: + RuntimeError: If :meth:`compile` has not been called, or if the + watched quantity is absent from the training history. + """ + 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, + ) + + def _read_target_value(self, history: tf.keras.callbacks.History) -> float: + """Extract the watched quantity from a ``History``, or explain why not.""" + try: + series = history.history[self.target.name] + except KeyError: + available = ", ".join(sorted(history.history)) or "none" + raise RuntimeError( + f"Target {self.target.name!r} is not in the training history. " + f"Available keys: {available}. " + "Validation keys (val_*) require passing validation_data to train()." + ) 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 + # ------------------------------------------------------------------ + def predict_best(self, *args: Any, **kwargs: Any) -> tuple[Any, float]: + """Predict with the best weights seen. + + Returns: + The model output and the best observed value. + """ + self._require_compiled() + return self.best_model.predict(*args, **kwargs), self.best_value + + def predict_last(self, *args: Any, **kwargs: Any) -> tuple[Any, float | None]: + """Predict with the most recent weights. + + Returns: + The model output and the most recent value (``None`` if untrained). + """ + return self.model.predict(*args, **kwargs), self.last_value + + def show_result(self, *args: Any, **kwargs: Any) -> None: + """Print predictions from both the best and the most recent weights.""" + self._require_compiled() + best_output, best_value = self.predict_best(*args, **kwargs) + last_output, last_value = self.predict_last(*args, **kwargs) + print() + print(_SEPARATOR) + print("Best result: ", best_output) + print(f" with {self.target.name}: ", best_value) + print(_SEPARATOR) + print("Last result: ", last_output) + print(f" with {self.target.name}: ", last_value) + print(_SEPARATOR) + + def _require_compiled(self) -> None: + if self.best_model is None: + raise RuntimeError("compile() must be called before inference.") + + # ------------------------------------------------------------------ + # Deprecated aliases (removed in 3.0.0) + # ------------------------------------------------------------------ + def predict_optimize(self, *args: Any, **kwargs: Any) -> tuple[Any, float]: + """Deprecated alias for :meth:`predict_best`.""" + _warn_renamed("predict_optimize()", "predict_best()") + return self.predict_best(*args, **kwargs) + + +def _warn_renamed(old: str, new: str) -> None: + warnings.warn( + f"{old} is deprecated and will be removed in 3.0.0; use {new}.", + DeprecationWarning, + stacklevel=3, + ) + + +def _deprecated_alias(new_name: str, old_name: str) -> property: + """Build a property that proxies ``old_name`` to ``new_name`` with a warning.""" + + def getter(self: InfiniteTrainer) -> Any: + _warn_renamed(f"{old_name!r}", f"{new_name!r}") + return getattr(self, new_name) + + def setter(self: InfiniteTrainer, value: Any) -> None: + _warn_renamed(f"{old_name!r}", f"{new_name!r}") + setattr(self, new_name, value) + + getter.__doc__ = f"Deprecated alias for :attr:`{new_name}`." + return property(getter, setter) + + +for _old, _new in ( + ("optimize_weight", "best_weights"), + ("last_weight", "last_weights"), + ("optimize_value", "best_value"), + ("list_value", "value_history"), + ("optimize_model", "best_model"), + ("optimize_weight_path", "best_weights_path"), + ("last_weight_path", "last_weights_path"), + ("optimize_value_path", "best_value_path"), + ("list_value_path", "value_history_path"), +): + setattr(InfiniteTrainer, _old, _deprecated_alias(_new, _old)) +del _old, _new + + +class InfinityTraining(InfiniteTrainer): + """Deprecated name for :class:`InfiniteTrainer`. + + Retained so that existing code keeps working; it will be removed in 3.0.0. + """ + + def __init__(self, *args: Any, **kwargs: Any) -> None: + warnings.warn( + "InfinityTraining is deprecated and will be removed in 3.0.0; use InfiniteTrainer.", + DeprecationWarning, + stacklevel=2, + ) + super().__init__(*args, **kwargs) diff --git a/pyproject.toml b/pyproject.toml index 712cf30..3532415 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,35 +1,59 @@ [build-system] requires = ["setuptools>=61.0"] build-backend = "setuptools.build_meta" + [project] name = "infinite_training" -version = "2.0.0" +version = "2.1.0" authors = [ { name="Nguyễn Chí Vỹ", email="vyncint@icloud.com" }, ] -description = "The python package for training ML models base on loss, metrics and timeout." +description = "Train Keras models until a target metric, a timeout, or Ctrl+C stops them - and resume where you left off." readme = "README.md" -requires-python = ">=3.9" +# Bounded by TensorFlow, which requires >=3.10 from 2.16 onward. +requires-python = ">=3.10" +keywords = ["tensorflow", "keras", "training", "checkpointing", "early-stopping", "machine-learning"] dependencies = [ "tensorflow", "numpy", - "polars", ] classifiers = [ - "Programming Language :: Python :: 3", + "Development Status :: 4 - Beta", + "Intended Audience :: Developers", + "Intended Audience :: Science/Research", "License :: OSI Approved :: MIT License", + "Programming Language :: Python :: 3", + "Programming Language :: Python :: 3.10", + "Programming Language :: Python :: 3.11", + "Programming Language :: Python :: 3.12", "Operating System :: OS Independent", + "Topic :: Scientific/Engineering :: Artificial Intelligence", ] [project.urls] "Homepage" = "https://github.com/vyncint/infinite-training" "Bug Tracker" = "https://github.com/vyncint/infinite-training/issues" +"Changelog" = "https://github.com/vyncint/infinite-training/blob/main/CHANGELOG.md" [project.optional-dependencies] dev = [ - "pytest", + "pytest>=7", + "ruff>=0.6", ] example = [ "tensorflow-datasets", "importlib_resources" ] + +[tool.setuptools.packages.find] +include = ["infinite_training*"] + +[tool.pytest.ini_options] +testpaths = ["tests"] + +[tool.ruff] +line-length = 100 +target-version = "py310" + +[tool.ruff.lint] +select = ["E", "F", "I", "UP", "B", "SIM"] diff --git a/run.py b/run.py deleted file mode 100644 index 2373976..0000000 --- a/run.py +++ /dev/null @@ -1,57 +0,0 @@ -""" - Apply in example from https://www.tensorflow.org/datasets/keras_example -""" -import tensorflow_datasets as tfds -import tensorflow as tf -import os -from infinite_training import InfinityTraining, Target - -(ds_train, ds_test), ds_info = tfds.load( - 'mnist', - split=['train', 'test'], - shuffle_files=True, - as_supervised=True, - with_info=True, -) - - -def normalize_img(image, label): - """Normalizes images: `uint8` -> `float32`.""" - return tf.cast(image, tf.float32) / 255., label - -ds_train = ds_train.map( - normalize_img, num_parallel_calls=tf.data.AUTOTUNE) -ds_train = ds_train.cache() -ds_train = ds_train.shuffle(ds_info.splits['train'].num_examples) -ds_train = ds_train.batch(128) -ds_train = ds_train.prefetch(tf.data.AUTOTUNE) - -ds_test = ds_test.map( - normalize_img, num_parallel_calls=tf.data.AUTOTUNE) -ds_test = ds_test.batch(128) -ds_test = ds_test.cache() -ds_test = ds_test.prefetch(tf.data.AUTOTUNE) - -model = tf.keras.models.Sequential([ - tf.keras.layers.Flatten(input_shape=(28, 28)), - tf.keras.layers.Dense(128, activation='relu'), - tf.keras.layers.Dense(10) -]) - -os.makedirs("models", exist_ok=True) - -it = InfinityTraining( - model=model, - target=Target(name="val_sparse_categorical_accuracy", smaller_is_better=False, target_value=0.98), - timeout=100, - optimize_weight_path="models/optimize_weight.npy", - last_weight_path="models/last_weight.npy", - optimize_value_path="models/optimize_value.npy", - list_value_path="models/list_value.npy" -) -it.compile(optimizer=tf.keras.optimizers.Adam(0.001), - loss=tf.keras.losses.SparseCategoricalCrossentropy( - from_logits=True), - metrics=[tf.keras.metrics.SparseCategoricalAccuracy()],) -it.train(ds_train, validation_data=ds_test) -it.show_result(ds_train) \ No newline at end of file diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..b1e9469 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,54 @@ +"""Shared fixtures for the test suite.""" + +from __future__ import annotations + +import os + +# Keep TensorFlow quiet; must be set before the first import. +os.environ.setdefault("TF_CPP_MIN_LOG_LEVEL", "3") + +import numpy as np # noqa: E402 +import pytest # noqa: E402 +import tensorflow as tf # noqa: E402 + +from infinite_training import InfiniteTrainer # noqa: E402 + + +@pytest.fixture +def model() -> tf.keras.Model: + """A two-input, one-output functional model small enough to fit instantly.""" + inputs = tf.keras.Input(shape=(2,)) + outputs = tf.keras.layers.Dense(1)(inputs) + return tf.keras.Model(inputs=inputs, outputs=outputs) + + +@pytest.fixture +def dataset() -> tuple[np.ndarray, np.ndarray]: + """A tiny deterministic regression dataset.""" + x = np.array([[0.0, 0.0], [0.0, 1.0], [1.0, 0.0], [1.0, 1.0]]) + y = np.array([[0.0], [1.0], [1.0], [0.0]]) + return x, y + + +@pytest.fixture +def paths(tmp_path) -> dict[str, str]: + """Checkpoint paths isolated to a temporary directory.""" + return { + "best_weights_path": str(tmp_path / "best_weights.npy"), + "last_weights_path": str(tmp_path / "last_weights.npy"), + "best_value_path": str(tmp_path / "best_value.npy"), + "value_history_path": str(tmp_path / "value_history.npy"), + } + + +@pytest.fixture +def make_trainer(model, paths): + """Build a compiled trainer, overriding any constructor argument.""" + + def _factory(**kwargs) -> InfiniteTrainer: + options = {**paths, **kwargs} + trainer = InfiniteTrainer(model=options.pop("model", model), **options) + trainer.compile(optimizer="adam", loss="mse") + return trainer + + return _factory diff --git a/tests/test_deprecations.py b/tests/test_deprecations.py new file mode 100644 index 0000000..466e2a1 --- /dev/null +++ b/tests/test_deprecations.py @@ -0,0 +1,119 @@ +"""The 2.x names must keep working, while warning that 3.0.0 removes them. + +These tests pin the backward-compatibility contract: code written against the +2.0.0 API must run unchanged. +""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest +import tensorflow as tf + +from infinite_training import InfinityTraining, Target + +X = np.array([[0.0, 0.0], [1.0, 1.0]]) +Y = np.array([[0.0], [1.0]]) + + +@pytest.fixture +def legacy_paths(tmp_path) -> dict[str, str]: + """Checkpoint paths using the 2.0.0 keyword names.""" + return { + "optimize_weight_path": str(tmp_path / "opt.npy"), + "last_weight_path": str(tmp_path / "last.npy"), + "optimize_value_path": str(tmp_path / "opt_val.npy"), + "list_value_path": str(tmp_path / "list_val.npy"), + } + + +@pytest.fixture +def model() -> tf.keras.Model: + inputs = tf.keras.Input(shape=(2,)) + outputs = tf.keras.layers.Dense(1)(inputs) + return tf.keras.Model(inputs=inputs, outputs=outputs) + + +class TestLegacyClassName: + def test_still_constructs(self, model, legacy_paths): + with pytest.deprecated_call(): + trainer = InfinityTraining(model=model, **legacy_paths) + assert trainer is not None + + def test_is_an_infinite_trainer(self, model, legacy_paths): + from infinite_training import InfiniteTrainer + + with pytest.deprecated_call(): + trainer = InfinityTraining(model=model, **legacy_paths) + assert isinstance(trainer, InfiniteTrainer) + + +class TestLegacyConstructorKeywords: + def test_legacy_paths_are_honoured(self, model, legacy_paths): + with pytest.deprecated_call(): + trainer = InfinityTraining(model=model, **legacy_paths) + assert trainer.best_weights_path == legacy_paths["optimize_weight_path"] + assert trainer.last_weights_path == legacy_paths["last_weight_path"] + assert trainer.best_value_path == legacy_paths["optimize_value_path"] + assert trainer.value_history_path == legacy_paths["list_value_path"] + + def test_checkpoints_land_at_the_legacy_paths(self, model, legacy_paths): + import os + + x, y = X, Y + with pytest.deprecated_call(): + trainer = InfinityTraining( + model=model, target=Target("loss", True, math.inf), **legacy_paths + ) + trainer.compile(optimizer="adam", loss="mse") + trainer.train(x, y, epochs=1, verbose=0) + for path in legacy_paths.values(): + assert os.path.exists(path) + + +class TestLegacyAttributes: + @pytest.mark.parametrize( + ("legacy", "current"), + [ + ("optimize_weight", "best_weights"), + ("last_weight", "last_weights"), + ("optimize_value", "best_value"), + ("list_value", "value_history"), + ("optimize_model", "best_model"), + ], + ) + def test_legacy_attribute_proxies_current_one(self, model, legacy_paths, legacy, current): + with pytest.deprecated_call(): + trainer = InfinityTraining(model=model, **legacy_paths) + with pytest.deprecated_call(): + legacy_value = getattr(trainer, legacy) + current_value = getattr(trainer, current) + + if legacy_value is None or current_value is None: + assert legacy_value is current_value + else: + assert repr(legacy_value) == repr(current_value) + + def test_legacy_attribute_assignment_still_works(self, model, legacy_paths): + with pytest.deprecated_call(): + trainer = InfinityTraining(model=model, **legacy_paths) + with pytest.deprecated_call(): + trainer.optimize_value = 0.25 + assert trainer.best_value == 0.25 + + +class TestLegacyMethods: + def test_predict_optimize_delegates_to_predict_best(self, model, legacy_paths): + x = X + with pytest.deprecated_call(): + trainer = InfinityTraining(model=model, **legacy_paths) + trainer.compile(optimizer="adam", loss="mse") + + with pytest.deprecated_call(): + legacy_output, legacy_value = trainer.predict_optimize(x, verbose=0) + current_output, current_value = trainer.predict_best(x, verbose=0) + + assert legacy_value == current_value + assert legacy_output.shape == current_output.shape diff --git a/tests/test_infinite_training.py b/tests/test_infinite_training.py deleted file mode 100644 index 76e249e..0000000 --- a/tests/test_infinite_training.py +++ /dev/null @@ -1,48 +0,0 @@ -import math -import numpy as np -import tensorflow as tf -from infinite_training import Target, InfinityTraining -import pytest - -def test_target_default(): - target = Target() - assert target.name == "loss" - assert target.smaller_is_better is True - assert target.target_value == -math.inf - -def test_target_custom(): - target = Target(name="accuracy", smaller_is_better=False, target_value=0.95) - assert target.name == "accuracy" - assert target.smaller_is_better is False - assert target.target_value == 0.95 - -def test_infinite_training_initialization(tmp_path): - inputs = tf.keras.Input(shape=(2,)) - x = tf.keras.layers.Dense(4)(inputs) - outputs = tf.keras.layers.Dense(2)(x) - model = tf.keras.Model(inputs=inputs, outputs=outputs) - - target = Target(name="loss", smaller_is_better=True, target_value=0.0001) - - optimize_weight_path = str(tmp_path / "opt.npy") - last_weight_path = str(tmp_path / "last.npy") - optimize_value_path = str(tmp_path / "opt_val.npy") - list_value_path = str(tmp_path / "list_val.npy") - - it = InfinityTraining( - model=model, - target=target, - timeout=10, - optimize_weight_path=optimize_weight_path, - last_weight_path=last_weight_path, - optimize_value_path=optimize_value_path, - list_value_path=list_value_path - ) - - assert it.timeout == 10 - assert it.optimize_value == math.inf - assert it.list_value.size == 0 - - it.compile(optimizer="adam", loss="mse", metrics=["accuracy"]) - assert it.optimize_model is not None - diff --git a/tests/test_persistence.py b/tests/test_persistence.py new file mode 100644 index 0000000..1331486 --- /dev/null +++ b/tests/test_persistence.py @@ -0,0 +1,120 @@ +"""Tests for checkpointing and resuming across sessions.""" + +from __future__ import annotations + +import math +import os + +import numpy as np +import pytest +import tensorflow as tf + +from infinite_training import InfiniteTrainer, Target + + +def _fresh_model() -> tf.keras.Model: + inputs = tf.keras.Input(shape=(2,)) + outputs = tf.keras.layers.Dense(1)(inputs) + return tf.keras.Model(inputs=inputs, outputs=outputs) + + +def _train_once(paths, dataset, **kwargs) -> InfiniteTrainer: + """Run a single round in a brand-new trainer, as a separate 'session'.""" + x, y = dataset + trainer = InfiniteTrainer( + model=_fresh_model(), + target=Target("loss", True, math.inf), + **paths, + **kwargs, + ) + trainer.compile(optimizer="adam", loss="mse") + trainer.train(x, y, epochs=1, verbose=0) + return trainer + + +class TestCheckpointFiles: + def test_all_four_checkpoints_are_written(self, paths, dataset): + _train_once(paths, dataset) + for path in paths.values(): + assert os.path.exists(path), f"missing checkpoint: {path}" + + def test_parent_directories_are_created(self, tmp_path, dataset): + nested = { + "best_weights_path": str(tmp_path / "runs" / "a" / "best.npy"), + "last_weights_path": str(tmp_path / "runs" / "a" / "last.npy"), + "best_value_path": str(tmp_path / "runs" / "a" / "value.npy"), + "value_history_path": str(tmp_path / "runs" / "a" / "history.npy"), + } + _train_once(nested, dataset) + for path in nested.values(): + assert os.path.exists(path) + + +class TestResume: + def test_history_accumulates_across_sessions(self, paths, dataset): + first = _train_once(paths, dataset) + assert first.rounds_completed == 1 + + second = _train_once(paths, dataset) + # The second session loads the first session's history and appends. + assert second.rounds_completed == 2 + + def test_best_value_is_restored(self, paths, dataset): + first = _train_once(paths, dataset) + saved_best = first.best_value + + resumed = InfiniteTrainer(model=_fresh_model(), **paths) + assert resumed.best_value == pytest.approx(saved_best) + assert isinstance(resumed.best_value, float) + + def test_weights_are_restored(self, paths, dataset): + first = _train_once(paths, dataset) + + resumed = InfiniteTrainer(model=_fresh_model(), **paths) + for restored, original in zip(resumed.last_weights, first.last_weights): + np.testing.assert_allclose(restored, original) + for restored, original in zip(resumed.best_weights, first.best_weights): + np.testing.assert_allclose(restored, original) + + def test_last_value_survives_a_restart(self, paths, dataset): + first = _train_once(paths, dataset) + + resumed = InfiniteTrainer(model=_fresh_model(), **paths) + # Derived from the persisted history, so it is available immediately + # without re-training. + assert resumed.last_value == pytest.approx(first.last_value) + + def test_restored_weights_are_loaded_into_the_models(self, paths, dataset): + _train_once(paths, dataset) + + resumed = InfiniteTrainer(model=_fresh_model(), **paths) + resumed.compile(optimizer="adam", loss="mse") + for live, checkpointed in zip(resumed.model.get_weights(), resumed.last_weights): + np.testing.assert_allclose(live, checkpointed) + for live, checkpointed in zip(resumed.best_model.get_weights(), resumed.best_weights): + np.testing.assert_allclose(live, checkpointed) + + +class TestFreshStart: + def test_absent_checkpoints_fall_back_to_model_state(self, paths): + model = _fresh_model() + trainer = InfiniteTrainer(model=model, **paths) + for loaded, original in zip(trainer.last_weights, model.get_weights()): + np.testing.assert_allclose(loaded, original) + + def test_best_value_starts_at_the_worst_possible(self, paths): + minimising = InfiniteTrainer( + model=_fresh_model(), target=Target(smaller_is_better=True), **paths + ) + assert minimising.best_value == math.inf + + def test_best_value_starts_at_negative_infinity_when_maximising(self, paths): + maximising = InfiniteTrainer( + model=_fresh_model(), target=Target(smaller_is_better=False), **paths + ) + assert maximising.best_value == -math.inf + + def test_history_starts_empty(self, paths): + trainer = InfiniteTrainer(model=_fresh_model(), **paths) + assert trainer.rounds_completed == 0 + assert len(trainer.value_history) == 0 diff --git a/tests/test_target.py b/tests/test_target.py new file mode 100644 index 0000000..9d89ebe --- /dev/null +++ b/tests/test_target.py @@ -0,0 +1,94 @@ +"""Tests for the Target stopping criterion.""" + +from __future__ import annotations + +import math + +import pytest + +from infinite_training import Target + + +class TestDefaults: + def test_defaults_watch_loss_and_minimise(self): + target = Target() + assert target.name == "loss" + assert target.smaller_is_better is True + + def test_unset_target_is_unreachable_when_minimising(self): + # No target means "run forever": a loss can never fall below -inf. + assert Target(smaller_is_better=True).target_value == -math.inf + + def test_unset_target_is_unreachable_when_maximising(self): + assert Target(smaller_is_better=False).target_value == math.inf + + def test_custom_values_are_preserved(self): + target = Target(name="accuracy", smaller_is_better=False, target_value=0.95) + assert target.name == "accuracy" + assert target.smaller_is_better is False + assert target.target_value == 0.95 + + +class TestFalsyTargetValues: + """Zero is a legitimate target and must not be treated as 'unset'.""" + + @pytest.mark.parametrize("zero", [0, 0.0]) + def test_zero_target_is_kept_when_minimising(self, zero): + assert Target(smaller_is_better=True, target_value=zero).target_value == 0 + + @pytest.mark.parametrize("zero", [0, 0.0]) + def test_zero_target_is_kept_when_maximising(self, zero): + assert Target(smaller_is_better=False, target_value=zero).target_value == 0 + + def test_zero_target_is_actually_reachable(self): + # Regression: `target_value or default` swallowed 0.0, leaving -inf, + # so the stop condition could never fire. + target = Target(smaller_is_better=True, target_value=0.0) + assert target.is_reached(-0.5) is True + + +class TestIsImprovement: + def test_lower_beats_higher_when_minimising(self): + target = Target(smaller_is_better=True) + assert target.is_improvement(0.1, 0.2) is True + assert target.is_improvement(0.2, 0.1) is False + + def test_higher_beats_lower_when_maximising(self): + target = Target(smaller_is_better=False) + assert target.is_improvement(0.9, 0.8) is True + assert target.is_improvement(0.8, 0.9) is False + + def test_equal_values_are_not_an_improvement(self): + assert Target(smaller_is_better=True).is_improvement(0.5, 0.5) is False + assert Target(smaller_is_better=False).is_improvement(0.5, 0.5) is False + + +class TestIsReached: + def test_minimising_requires_strictly_below_target(self): + target = Target(smaller_is_better=True, target_value=0.1) + assert target.is_reached(0.09) is True + assert target.is_reached(0.1) is False + assert target.is_reached(0.2) is False + + def test_maximising_requires_strictly_above_target(self): + target = Target(smaller_is_better=False, target_value=0.9) + assert target.is_reached(0.91) is True + assert target.is_reached(0.9) is False + assert target.is_reached(0.8) is False + + def test_default_target_is_never_reached(self): + assert Target(smaller_is_better=True).is_reached(-1e308) is False + assert Target(smaller_is_better=False).is_reached(1e308) is False + + +class TestWorstPossibleValue: + def test_starts_from_positive_infinity_when_minimising(self): + assert Target(smaller_is_better=True).worst_possible_value == math.inf + + def test_starts_from_negative_infinity_when_maximising(self): + assert Target(smaller_is_better=False).worst_possible_value == -math.inf + + def test_any_first_measurement_improves_on_it(self): + for smaller in (True, False): + target = Target(smaller_is_better=smaller) + assert target.is_improvement(0.5, target.worst_possible_value) is True diff --git a/tests/test_trainer.py b/tests/test_trainer.py new file mode 100644 index 0000000..5cc213c --- /dev/null +++ b/tests/test_trainer.py @@ -0,0 +1,120 @@ +"""Tests for the training loop, stop conditions and error handling.""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest + +from infinite_training import InfiniteTrainer, Target + + +class TestStopConditions: + def test_stops_once_the_target_is_reached(self, make_trainer, dataset): + x, y = dataset + # target_value=inf while minimising is satisfied by any finite loss, + # so the loop must exit after exactly one round. + trainer = make_trainer(target=Target("loss", True, math.inf)) + trainer.train(x, y, epochs=1, verbose=0) + assert trainer.rounds_completed == 1 + + def test_stops_once_the_timeout_elapses(self, make_trainer, dataset): + x, y = dataset + trainer = make_trainer(target=Target("loss", True), timeout=0) + trainer.train(x, y, epochs=1, verbose=0) + # Timeout is checked between rounds, so exactly one round runs. + assert trainer.rounds_completed == 1 + + def test_target_takes_priority_over_timeout(self, make_trainer, dataset): + x, y = dataset + trainer = make_trainer(target=Target("loss", True, math.inf), timeout=0) + trainer.train(x, y, epochs=1, verbose=0) + assert trainer.rounds_completed == 1 + + +class TestBestTracking: + def test_records_one_history_entry_per_round(self, make_trainer, dataset): + x, y = dataset + trainer = make_trainer(target=Target("loss", True, math.inf)) + trainer.train(x, y, epochs=3, verbose=0) + # One fit() call == one recorded value, regardless of epochs per call. + assert trainer.rounds_completed == 1 + assert len(trainer.value_history) == 1 + + def test_best_value_matches_the_observed_value(self, make_trainer, dataset): + x, y = dataset + trainer = make_trainer(target=Target("loss", True, math.inf)) + trainer.train(x, y, epochs=1, verbose=0) + assert trainer.best_value == pytest.approx(trainer.last_value) + + def test_best_weights_are_captured(self, make_trainer, dataset): + x, y = dataset + trainer = make_trainer(target=Target("loss", True, math.inf)) + trainer.train(x, y, epochs=1, verbose=0) + assert trainer.best_weights + assert all(isinstance(w, np.ndarray) for w in trainer.best_weights) + + def test_best_value_is_not_replaced_by_a_worse_round(self, make_trainer, dataset): + x, y = dataset + trainer = make_trainer(target=Target("loss", True, math.inf)) + trainer.train(x, y, epochs=1, verbose=0) + best_after_first = trainer.best_value + + # Force a deliberately worse observation and re-run one round. + trainer.best_value = -1.0 + trainer.train(x, y, epochs=1, verbose=0) + assert trainer.best_value == -1.0 + assert best_after_first != -1.0 + + +class TestLastValue: + def test_is_none_before_training(self, make_trainer): + # Regression: this used to raise AttributeError because last_value was + # only ever assigned inside the training loop. + assert make_trainer().last_value is None + + def test_predict_last_works_before_training(self, make_trainer, dataset): + x, _ = dataset + _, value = make_trainer().predict_last(x, verbose=0) + assert value is None + + def test_reflects_the_most_recent_round(self, make_trainer, dataset): + x, y = dataset + trainer = make_trainer(target=Target("loss", True, math.inf)) + trainer.train(x, y, epochs=1, verbose=0) + assert trainer.last_value == pytest.approx(float(trainer.value_history[-1])) + + +class TestErrorHandling: + def test_train_without_compile_is_rejected(self, model, paths, dataset): + x, y = dataset + trainer = InfiniteTrainer(model=model, **paths) + with pytest.raises(RuntimeError, match="compile\\(\\) must be called"): + trainer.train(x, y, epochs=1, verbose=0) + + def test_inference_without_compile_is_rejected(self, model, paths, dataset): + x, _ = dataset + trainer = InfiniteTrainer(model=model, **paths) + with pytest.raises(RuntimeError, match="compile\\(\\) must be called"): + trainer.predict_best(x, verbose=0) + + def test_unknown_target_name_explains_what_is_available(self, make_trainer, dataset): + x, y = dataset + trainer = make_trainer(target=Target("not_a_metric", True, math.inf)) + with pytest.raises(RuntimeError) as excinfo: + trainer.train(x, y, epochs=1, verbose=0) + message = str(excinfo.value) + assert "not_a_metric" in message + assert "Available keys" in message + assert "loss" in message + + def test_validation_target_hints_at_validation_data(self, make_trainer, dataset): + x, y = dataset + trainer = make_trainer(target=Target("val_loss", True, math.inf)) + with pytest.raises(RuntimeError, match="validation_data"): + trainer.train(x, y, epochs=1, verbose=0) + + def test_unexpected_keyword_is_rejected(self, model, paths): + with pytest.raises(TypeError, match="Unexpected keyword argument"): + InfiniteTrainer(model=model, not_a_real_option=1, **paths) From c5f111ee5bf475d802e059b76aa3cbe41fae03c7 Mon Sep 17 00:00:00 2001 From: vyncint <115854244+vyncint@users.noreply.github.com> Date: Sat, 25 Jul 2026 16:15:32 +0700 Subject: [PATCH 2/2] fix(lint): satisfy the py310 lint rules the new CI surfaced Raising target-version to py310 enabled B905, and zip() over weight lists is exactly where an unnoticed length mismatch would matter, so pass strict=True rather than silencing it. Also import Callable from collections.abc. --- infinite_training/_storage.py | 3 ++- tests/test_persistence.py | 14 +++++++++----- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/infinite_training/_storage.py b/infinite_training/_storage.py index 24e5fad..78d1a03 100644 --- a/infinite_training/_storage.py +++ b/infinite_training/_storage.py @@ -12,7 +12,8 @@ from __future__ import annotations import os -from typing import Any, Callable, TypeVar +from collections.abc import Callable +from typing import Any, TypeVar import numpy as np diff --git a/tests/test_persistence.py b/tests/test_persistence.py index 1331486..a81e0a6 100644 --- a/tests/test_persistence.py +++ b/tests/test_persistence.py @@ -71,9 +71,9 @@ def test_weights_are_restored(self, paths, dataset): first = _train_once(paths, dataset) resumed = InfiniteTrainer(model=_fresh_model(), **paths) - for restored, original in zip(resumed.last_weights, first.last_weights): + for restored, original in zip(resumed.last_weights, first.last_weights, strict=True): np.testing.assert_allclose(restored, original) - for restored, original in zip(resumed.best_weights, first.best_weights): + for restored, original in zip(resumed.best_weights, first.best_weights, strict=True): np.testing.assert_allclose(restored, original) def test_last_value_survives_a_restart(self, paths, dataset): @@ -89,9 +89,13 @@ def test_restored_weights_are_loaded_into_the_models(self, paths, dataset): resumed = InfiniteTrainer(model=_fresh_model(), **paths) resumed.compile(optimizer="adam", loss="mse") - for live, checkpointed in zip(resumed.model.get_weights(), resumed.last_weights): + for live, checkpointed in zip( + resumed.model.get_weights(), resumed.last_weights, strict=True + ): np.testing.assert_allclose(live, checkpointed) - for live, checkpointed in zip(resumed.best_model.get_weights(), resumed.best_weights): + for live, checkpointed in zip( + resumed.best_model.get_weights(), resumed.best_weights, strict=True + ): np.testing.assert_allclose(live, checkpointed) @@ -99,7 +103,7 @@ class TestFreshStart: def test_absent_checkpoints_fall_back_to_model_state(self, paths): model = _fresh_model() trainer = InfiniteTrainer(model=model, **paths) - for loaded, original in zip(trainer.last_weights, model.get_weights()): + for loaded, original in zip(trainer.last_weights, model.get_weights(), strict=True): np.testing.assert_allclose(loaded, original) def test_best_value_starts_at_the_worst_possible(self, paths):