torchlight 0.1.0__tar.gz

torchlight-0.1.0/.gitignore

@@ -0,0 +1,6 @@
+.venv/
+__pycache__/
+*.egg-info/
+.pytest_cache/
+build/
+dist/

torchlight-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 nehz
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torchlight-0.1.0/PKG-INFO

@@ -0,0 +1,202 @@
+Metadata-Version: 2.5
+Name: torchlight
+Version: 0.1.0
+Summary: Torch Light: a lightweight PyTorch companion with a minimal Trainer, callbacks, metrics and checkpointing.
+Author: nehz
+License: MIT
+License-File: LICENSE
+Keywords: deep-learning,pytorch,torch,trainer,training
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Typing :: Typed
+Requires-Python: >=3.9
+Requires-Dist: torch>=2.0
+Provides-Extra: test
+Requires-Dist: pytest>=7; extra == 'test'
+Description-Content-Type: text/markdown
+
+# torchlight
+
+**Torch Light**: a lightweight companion for PyTorch.
+
+`torchlight` gives you a small, readable training loop plus the extras you end
+up rewriting in every project: callbacks, metrics, early stopping,
+checkpointing and seeding/device helpers. It is a few hundred lines of plain
+PyTorch with no other dependencies, so you can read all of it in one sitting.
+When you outgrow it, copy the parts you need.
+
+## Features
+
+- **`Trainer`**: `fit` / `evaluate` / `predict` over any iterable of
+  `(inputs, targets)` batches, with automatic device placement (including
+  nested dicts, lists and tuples), optional gradient-norm clipping, and a
+  per-epoch `history`.
+- **Callbacks**: hook into fit start, epoch start, batch end, epoch end and
+  fit end. Built in: `EarlyStopping` (with `restore_best_weights`),
+  `ModelCheckpoint` (best-only or every epoch, templated filenames), and
+  `PrintLogger`.
+- **Metrics**: accumulate-then-compute `Accuracy` (multi-class logits or
+  binary thresholds) and `MeanAbsoluteError`. Subclass `Metric` to add your own.
+- **Checkpoints**: `save_checkpoint` / `load_checkpoint` for model and
+  optimizer state plus any extra values.
+- **Utilities**: `seed_everything`, `get_device`, `to_device`,
+  `count_parameters`.
+
+## Install
+
+```bash
+pip install torchlight
+```
+
+`torch>=2.0` is the only dependency. To use the CPU-only PyTorch build, install
+it first, for example
+`pip install torch --index-url https://download.pytorch.org/whl/cpu`.
+
+From a source checkout:
+
+```bash
+python -m venv .venv
+.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu
+.venv/bin/pip install -e ".[test]"
+.venv/bin/python -m pytest
+```
+
+## Quickstart
+
+```python
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+import torchlight as tl
+
+tl.seed_everything(42)
+
+# Synthetic two-class data
+x = torch.randn(512, 2)
+y = (x[:, 0] + x[:, 1] > 0).long()
+train = DataLoader(TensorDataset(x[:400], y[:400]), batch_size=32, shuffle=True)
+val = DataLoader(TensorDataset(x[400:], y[400:]), batch_size=64)
+
+model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 2))
+trainer = tl.Trainer(
+    model,
+    optimizer=torch.optim.Adam(model.parameters(), lr=1e-2),
+    loss_fn=nn.CrossEntropyLoss(),
+    metrics={"accuracy": tl.Accuracy()},
+    callbacks=[
+        tl.EarlyStopping(monitor="val_loss", patience=3, restore_best_weights=True),
+        tl.ModelCheckpoint("checkpoints/best.pt", monitor="val_accuracy", mode="max"),
+        tl.PrintLogger(),
+    ],
+)
+
+history = trainer.fit(train, val, epochs=20)
+print(history["val_accuracy"][-1])
+
+print(trainer.evaluate(val))        # {'loss': ..., 'accuracy': ...}
+probs = trainer.predict(val).softmax(dim=1)
+
+# Later: restore the best checkpoint
+ckpt = tl.load_checkpoint("checkpoints/best.pt", model)
+print(ckpt["epoch"], ckpt["logs"])
+```
+
+## API overview
+
+Everything below is importable from the top-level `torchlight` package.
+
+### `Trainer`
+
+```python
+Trainer(model, optimizer, loss_fn, metrics=None, callbacks=None, device=None, grad_clip=None)
+```
+
+| Member | Description |
+| --- | --- |
+| `fit(train_loader, val_loader=None, epochs=1) -> dict[str, list[float]]` | Train for up to `epochs` epochs and return `history`. Logs `loss` and each metric name; with `val_loader`, also `val_loss` and `val_<metric>`. |
+| `evaluate(loader) -> dict[str, float]` | Loss and metrics under `torch.no_grad()`, in eval mode, with no `val_` prefix. |
+| `predict(loader) -> torch.Tensor` | Model outputs concatenated on the CPU. Batches may be `(inputs, targets)` or bare `inputs`. |
+| `history` | `dict[str, list[float]]` of per-epoch values. |
+| `should_stop` | Set to `True` (usually from a callback) to stop after the current epoch. |
+| `model`, `optimizer`, `loss_fn`, `metrics`, `callbacks`, `device`, `grad_clip`, `epochs` | The configured objects. `epochs` is the value passed to the latest `fit`. |
+
+Batches must be `(inputs, targets)` pairs. `inputs` goes straight to
+`model(inputs)`, and both halves are moved to `trainer.device`. The loss is
+averaged over samples, with each batch weighted by its size. `device`
+defaults to `get_device()`.
+
+### Callbacks
+
+```python
+class Callback:
+    def on_fit_start(self, trainer): ...
+    def on_epoch_start(self, trainer, epoch): ...           # epoch is 0-based
+    def on_batch_end(self, trainer, batch_idx, loss): ...   # loss is a float
+    def on_epoch_end(self, trainer, epoch, logs): ...       # logs: dict[str, float]
+    def on_fit_end(self, trainer): ...                      # also runs after an early stop
+```
+
+- `EarlyStopping(monitor="val_loss", patience=3, mode="min", min_delta=0.0, restore_best_weights=False)`
+  sets `trainer.should_stop` after `patience` epochs without improvement.
+  Attributes: `best`, `best_epoch`, `wait`, `stopped_epoch`.
+- `ModelCheckpoint(path, monitor="val_loss", mode="min", save_best_only=True)`
+  writes `save_checkpoint(path, model, optimizer, epoch=..., logs=...)`.
+  `path` may use `{epoch}` and any log key as format fields, e.g.
+  `"ckpt/ep{epoch:02d}-{val_loss:.3f}.pt"`. Attributes: `best`, `best_path`.
+- `PrintLogger(precision=4)` prints `epoch 3/10  loss=0.1234  val_loss=...`.
+
+If a monitored key is missing from the epoch logs (for example `val_loss`
+without a `val_loader`), the callback raises `KeyError`.
+
+### Metrics
+
+```python
+class Metric:
+    def update(self, preds, targets) -> None: ...
+    def compute(self) -> float: ...
+    def reset(self) -> None: ...
+```
+
+- `Accuracy(threshold=0.5)`: uses `argmax(dim=1)` when `preds` has one more
+  dimension than `targets`. Otherwise `preds >= threshold`; use
+  `threshold=0.0` for raw binary logits.
+- `MeanAbsoluteError()`: `preds` and `targets` must have the same shape.
+- `Mean()`: weighted running mean, `update(value, weight=1)`. The trainer uses
+  it for the loss.
+
+`compute()` returns `nan` if nothing has been accumulated.
+
+### Checkpoints
+
+- `save_checkpoint(path, model, optimizer=None, **extra) -> pathlib.Path`
+  saves `{"model": ..., "optimizer": ... or None, **extra}` and creates
+  parent directories.
+- `load_checkpoint(path, model, optimizer=None, map_location="cpu") -> dict`
+  loads weights (and optimizer state, if both are present) in place and
+  returns the full dict. It uses `torch.load(..., weights_only=True)`.
+
+### Utilities
+
+- `seed_everything(seed, deterministic=False) -> int` seeds `random`, NumPy
+  (if installed) and torch (CPU and CUDA).
+- `get_device(prefer=None) -> torch.device` returns `prefer` if given,
+  otherwise the first available of CUDA, MPS and CPU.
+- `to_device(data, device)` recursively moves tensors in dicts, lists, tuples
+  and named tuples.
+- `count_parameters(model, trainable_only=True) -> int`.
+
+## License
+
+MIT

torchlight-0.1.0/README.md

@@ -0,0 +1,174 @@
+# torchlight
+
+**Torch Light**: a lightweight companion for PyTorch.
+
+`torchlight` gives you a small, readable training loop plus the extras you end
+up rewriting in every project: callbacks, metrics, early stopping,
+checkpointing and seeding/device helpers. It is a few hundred lines of plain
+PyTorch with no other dependencies, so you can read all of it in one sitting.
+When you outgrow it, copy the parts you need.
+
+## Features
+
+- **`Trainer`**: `fit` / `evaluate` / `predict` over any iterable of
+  `(inputs, targets)` batches, with automatic device placement (including
+  nested dicts, lists and tuples), optional gradient-norm clipping, and a
+  per-epoch `history`.
+- **Callbacks**: hook into fit start, epoch start, batch end, epoch end and
+  fit end. Built in: `EarlyStopping` (with `restore_best_weights`),
+  `ModelCheckpoint` (best-only or every epoch, templated filenames), and
+  `PrintLogger`.
+- **Metrics**: accumulate-then-compute `Accuracy` (multi-class logits or
+  binary thresholds) and `MeanAbsoluteError`. Subclass `Metric` to add your own.
+- **Checkpoints**: `save_checkpoint` / `load_checkpoint` for model and
+  optimizer state plus any extra values.
+- **Utilities**: `seed_everything`, `get_device`, `to_device`,
+  `count_parameters`.
+
+## Install
+
+```bash
+pip install torchlight
+```
+
+`torch>=2.0` is the only dependency. To use the CPU-only PyTorch build, install
+it first, for example
+`pip install torch --index-url https://download.pytorch.org/whl/cpu`.
+
+From a source checkout:
+
+```bash
+python -m venv .venv
+.venv/bin/pip install torch --index-url https://download.pytorch.org/whl/cpu
+.venv/bin/pip install -e ".[test]"
+.venv/bin/python -m pytest
+```
+
+## Quickstart
+
+```python
+import torch
+from torch import nn
+from torch.utils.data import DataLoader, TensorDataset
+
+import torchlight as tl
+
+tl.seed_everything(42)
+
+# Synthetic two-class data
+x = torch.randn(512, 2)
+y = (x[:, 0] + x[:, 1] > 0).long()
+train = DataLoader(TensorDataset(x[:400], y[:400]), batch_size=32, shuffle=True)
+val = DataLoader(TensorDataset(x[400:], y[400:]), batch_size=64)
+
+model = nn.Sequential(nn.Linear(2, 16), nn.ReLU(), nn.Linear(16, 2))
+trainer = tl.Trainer(
+    model,
+    optimizer=torch.optim.Adam(model.parameters(), lr=1e-2),
+    loss_fn=nn.CrossEntropyLoss(),
+    metrics={"accuracy": tl.Accuracy()},
+    callbacks=[
+        tl.EarlyStopping(monitor="val_loss", patience=3, restore_best_weights=True),
+        tl.ModelCheckpoint("checkpoints/best.pt", monitor="val_accuracy", mode="max"),
+        tl.PrintLogger(),
+    ],
+)
+
+history = trainer.fit(train, val, epochs=20)
+print(history["val_accuracy"][-1])
+
+print(trainer.evaluate(val))        # {'loss': ..., 'accuracy': ...}
+probs = trainer.predict(val).softmax(dim=1)
+
+# Later: restore the best checkpoint
+ckpt = tl.load_checkpoint("checkpoints/best.pt", model)
+print(ckpt["epoch"], ckpt["logs"])
+```
+
+## API overview
+
+Everything below is importable from the top-level `torchlight` package.
+
+### `Trainer`
+
+```python
+Trainer(model, optimizer, loss_fn, metrics=None, callbacks=None, device=None, grad_clip=None)
+```
+
+| Member | Description |
+| --- | --- |
+| `fit(train_loader, val_loader=None, epochs=1) -> dict[str, list[float]]` | Train for up to `epochs` epochs and return `history`. Logs `loss` and each metric name; with `val_loader`, also `val_loss` and `val_<metric>`. |
+| `evaluate(loader) -> dict[str, float]` | Loss and metrics under `torch.no_grad()`, in eval mode, with no `val_` prefix. |
+| `predict(loader) -> torch.Tensor` | Model outputs concatenated on the CPU. Batches may be `(inputs, targets)` or bare `inputs`. |
+| `history` | `dict[str, list[float]]` of per-epoch values. |
+| `should_stop` | Set to `True` (usually from a callback) to stop after the current epoch. |
+| `model`, `optimizer`, `loss_fn`, `metrics`, `callbacks`, `device`, `grad_clip`, `epochs` | The configured objects. `epochs` is the value passed to the latest `fit`. |
+
+Batches must be `(inputs, targets)` pairs. `inputs` goes straight to
+`model(inputs)`, and both halves are moved to `trainer.device`. The loss is
+averaged over samples, with each batch weighted by its size. `device`
+defaults to `get_device()`.
+
+### Callbacks
+
+```python
+class Callback:
+    def on_fit_start(self, trainer): ...
+    def on_epoch_start(self, trainer, epoch): ...           # epoch is 0-based
+    def on_batch_end(self, trainer, batch_idx, loss): ...   # loss is a float
+    def on_epoch_end(self, trainer, epoch, logs): ...       # logs: dict[str, float]
+    def on_fit_end(self, trainer): ...                      # also runs after an early stop
+```
+
+- `EarlyStopping(monitor="val_loss", patience=3, mode="min", min_delta=0.0, restore_best_weights=False)`
+  sets `trainer.should_stop` after `patience` epochs without improvement.
+  Attributes: `best`, `best_epoch`, `wait`, `stopped_epoch`.
+- `ModelCheckpoint(path, monitor="val_loss", mode="min", save_best_only=True)`
+  writes `save_checkpoint(path, model, optimizer, epoch=..., logs=...)`.
+  `path` may use `{epoch}` and any log key as format fields, e.g.
+  `"ckpt/ep{epoch:02d}-{val_loss:.3f}.pt"`. Attributes: `best`, `best_path`.
+- `PrintLogger(precision=4)` prints `epoch 3/10  loss=0.1234  val_loss=...`.
+
+If a monitored key is missing from the epoch logs (for example `val_loss`
+without a `val_loader`), the callback raises `KeyError`.
+
+### Metrics
+
+```python
+class Metric:
+    def update(self, preds, targets) -> None: ...
+    def compute(self) -> float: ...
+    def reset(self) -> None: ...
+```
+
+- `Accuracy(threshold=0.5)`: uses `argmax(dim=1)` when `preds` has one more
+  dimension than `targets`. Otherwise `preds >= threshold`; use
+  `threshold=0.0` for raw binary logits.
+- `MeanAbsoluteError()`: `preds` and `targets` must have the same shape.
+- `Mean()`: weighted running mean, `update(value, weight=1)`. The trainer uses
+  it for the loss.
+
+`compute()` returns `nan` if nothing has been accumulated.
+
+### Checkpoints
+
+- `save_checkpoint(path, model, optimizer=None, **extra) -> pathlib.Path`
+  saves `{"model": ..., "optimizer": ... or None, **extra}` and creates
+  parent directories.
+- `load_checkpoint(path, model, optimizer=None, map_location="cpu") -> dict`
+  loads weights (and optimizer state, if both are present) in place and
+  returns the full dict. It uses `torch.load(..., weights_only=True)`.
+
+### Utilities
+
+- `seed_everything(seed, deterministic=False) -> int` seeds `random`, NumPy
+  (if installed) and torch (CPU and CUDA).
+- `get_device(prefer=None) -> torch.device` returns `prefer` if given,
+  otherwise the first available of CUDA, MPS and CPU.
+- `to_device(data, device)` recursively moves tensors in dicts, lists, tuples
+  and named tuples.
+- `count_parameters(model, trainable_only=True) -> int`.
+
+## License
+
+MIT

torchlight-0.1.0/pyproject.toml

@@ -0,0 +1,39 @@
+[build-system]
+requires = ["hatchling>=1.18"]
+build-backend = "hatchling.build"
+
+[project]
+name = "torchlight"
+version = "0.1.0"
+description = "Torch Light: a lightweight PyTorch companion with a minimal Trainer, callbacks, metrics and checkpointing."
+readme = "README.md"
+requires-python = ">=3.9"
+license = { text = "MIT" }
+authors = [{ name = "nehz" }]
+keywords = ["pytorch", "torch", "training", "trainer", "deep-learning"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Developers",
+    "Intended Audience :: Science/Research",
+    "License :: OSI Approved :: MIT License",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.9",
+    "Programming Language :: Python :: 3.10",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Typing :: Typed",
+]
+dependencies = ["torch>=2.0"]
+
+[project.optional-dependencies]
+test = ["pytest>=7"]
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/torchlight"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

torchlight-0.1.0/src/torchlight/__init__.py

@@ -0,0 +1,28 @@
+"""torchlight: a lightweight companion for PyTorch training loops."""
+
+from .callbacks import Callback, EarlyStopping, ModelCheckpoint, PrintLogger
+from .checkpoint import load_checkpoint, save_checkpoint
+from .metrics import Accuracy, Mean, MeanAbsoluteError, Metric
+from .trainer import Trainer
+from .utils import count_parameters, get_device, seed_everything, to_device
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Trainer",
+    "Callback",
+    "EarlyStopping",
+    "ModelCheckpoint",
+    "PrintLogger",
+    "Metric",
+    "Mean",
+    "Accuracy",
+    "MeanAbsoluteError",
+    "save_checkpoint",
+    "load_checkpoint",
+    "seed_everything",
+    "get_device",
+    "to_device",
+    "count_parameters",
+    "__version__",
+]

torchlight-0.1.0/src/torchlight/callbacks.py

@@ -0,0 +1,183 @@
+"""Callbacks that hook into :meth:`torchlight.Trainer.fit`."""
+
+from __future__ import annotations
+
+import copy
+import math
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, Any, Literal
+
+from .checkpoint import save_checkpoint
+
+if TYPE_CHECKING:
+    from .trainer import Trainer
+
+__all__ = ["Callback", "EarlyStopping", "ModelCheckpoint", "PrintLogger"]
+
+Mode = Literal["min", "max"]
+
+
+class Callback:
+    """Base class for callbacks. Override any subset of the hooks.
+
+    ``logs`` passed to :meth:`on_epoch_end` is the dict of epoch results, e.g.
+    ``{"loss": 0.31, "accuracy": 0.9, "val_loss": 0.35, "val_accuracy": 0.88}``.
+    Set ``trainer.should_stop = True`` to end training after the current epoch.
+    """
+
+    def on_fit_start(self, trainer: Trainer) -> None:
+        """Called once before the first epoch."""
+
+    def on_epoch_start(self, trainer: Trainer, epoch: int) -> None:
+        """Called before each epoch (``epoch`` is 0-based)."""
+
+    def on_batch_end(self, trainer: Trainer, batch_idx: int, loss: float) -> None:
+        """Called after each training step with that batch's loss."""
+
+    def on_epoch_end(self, trainer: Trainer, epoch: int, logs: dict[str, float]) -> None:
+        """Called after training (and validation, if any) for each epoch."""
+
+    def on_fit_end(self, trainer: Trainer) -> None:
+        """Called once after the last epoch, including after an early stop."""
+
+
+def _monitored(logs: dict[str, float], monitor: str, owner: str) -> float:
+    if monitor not in logs:
+        raise KeyError(
+            f"{owner}: monitored key {monitor!r} not in epoch logs "
+            f"(available: {sorted(logs)})"
+        )
+    return logs[monitor]
+
+
+def _improved(current: float, best: float, mode: Mode, min_delta: float) -> bool:
+    if mode == "min":
+        return current < best - min_delta
+    return current > best + min_delta
+
+
+def _check_mode(mode: str) -> None:
+    if mode not in ("min", "max"):
+        raise ValueError(f"mode must be 'min' or 'max', got {mode!r}")
+
+
+class EarlyStopping(Callback):
+    """Stop training when a monitored value stops improving.
+
+    Args:
+        monitor: Key in the epoch logs to watch, e.g. ``"val_loss"``.
+        patience: Epochs without improvement to tolerate before stopping.
+        mode: ``"min"`` if lower is better, ``"max"`` if higher is better.
+        min_delta: Minimum change that counts as an improvement.
+        restore_best_weights: On stop (or fit end), load the weights from the
+            best epoch back into the model.
+    """
+
+    def __init__(
+        self,
+        monitor: str = "val_loss",
+        patience: int = 3,
+        mode: Mode = "min",
+        min_delta: float = 0.0,
+        restore_best_weights: bool = False,
+    ) -> None:
+        _check_mode(mode)
+        self.monitor = monitor
+        self.patience = patience
+        self.mode: Mode = mode
+        self.min_delta = min_delta
+        self.restore_best_weights = restore_best_weights
+        self.best: float = math.inf if mode == "min" else -math.inf
+        self.best_epoch: int = -1
+        self.wait = 0
+        self.stopped_epoch: int | None = None
+        self._best_state: dict[str, Any] | None = None
+
+    def on_fit_start(self, trainer: Trainer) -> None:
+        self.best = math.inf if self.mode == "min" else -math.inf
+        self.best_epoch = -1
+        self.wait = 0
+        self.stopped_epoch = None
+        self._best_state = None
+
+    def on_epoch_end(self, trainer: Trainer, epoch: int, logs: dict[str, float]) -> None:
+        current = _monitored(logs, self.monitor, "EarlyStopping")
+        if _improved(current, self.best, self.mode, self.min_delta):
+            self.best = current
+            self.best_epoch = epoch
+            self.wait = 0
+            if self.restore_best_weights:
+                self._best_state = copy.deepcopy(trainer.model.state_dict())
+            return
+        self.wait += 1
+        if self.wait >= self.patience:
+            self.stopped_epoch = epoch
+            trainer.should_stop = True
+
+    def on_fit_end(self, trainer: Trainer) -> None:
+        if self.restore_best_weights and self._best_state is not None:
+            trainer.model.load_state_dict(self._best_state)
+
+
+class ModelCheckpoint(Callback):
+    """Save a checkpoint at the end of epochs.
+
+    Args:
+        path: File to write. May contain ``{epoch}`` and any log key as format
+            fields, e.g. ``"ckpt/epoch{epoch:02d}-{val_loss:.3f}.pt"``.
+        monitor: Key in the epoch logs used to decide what is "best".
+        mode: ``"min"`` or ``"max"``.
+        save_best_only: If True, only write when ``monitor`` improves;
+            otherwise write every epoch.
+
+    After fitting, :attr:`best_path` holds the last file written for the best
+    value (or the latest file when ``save_best_only`` is False).
+    """
+
+    def __init__(
+        self,
+        path: str | os.PathLike[str],
+        monitor: str = "val_loss",
+        mode: Mode = "min",
+        save_best_only: bool = True,
+    ) -> None:
+        _check_mode(mode)
+        self.path = str(path)
+        self.monitor = monitor
+        self.mode: Mode = mode
+        self.save_best_only = save_best_only
+        self.best: float = math.inf if mode == "min" else -math.inf
+        self.best_path: Path | None = None
+
+    def on_fit_start(self, trainer: Trainer) -> None:
+        self.best = math.inf if self.mode == "min" else -math.inf
+        self.best_path = None
+
+    def on_epoch_end(self, trainer: Trainer, epoch: int, logs: dict[str, float]) -> None:
+        current = _monitored(logs, self.monitor, "ModelCheckpoint")
+        improved = _improved(current, self.best, self.mode, 0.0)
+        if self.save_best_only and not improved:
+            return
+        if improved:
+            self.best = current
+        target = self.path.format(epoch=epoch, **logs)
+        self.best_path = save_checkpoint(
+            target,
+            trainer.model,
+            trainer.optimizer,
+            epoch=epoch,
+            logs=dict(logs),
+        )
+
+
+class PrintLogger(Callback):
+    """Print one line of epoch results, e.g. ``epoch 3/10  loss=0.1234 ...``."""
+
+    def __init__(self, precision: int = 4) -> None:
+        self.precision = precision
+
+    def on_epoch_end(self, trainer: Trainer, epoch: int, logs: dict[str, float]) -> None:
+        parts = "  ".join(f"{k}={v:.{self.precision}f}" for k, v in logs.items())
+        total = trainer.epochs if trainer.epochs is not None else "?"
+        print(f"epoch {epoch + 1}/{total}  {parts}")