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

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:

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

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

Trainer(model, optimizer, loss_fn, metrics=None, callbacks=None, device=None, grad_clip=None)

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

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

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