trainloop
Minimal PyTorch training loop with hooks for logging, checkpointing, and customization.
Docs: https://karimknaebel.github.io/trainloop/
Install
pip install trainloop
Basic example
import logging
import torch
import torch.nn as nn
from trainloop import BaseTrainer, CheckpointingHook, ProgressHook
logging.basicConfig(level=logging.INFO)
class MyTrainer(BaseTrainer):
def build_data_loader(self):
class ToyDataset(torch.utils.data.IterableDataset):
def __iter__(self):
while True:
data = torch.randn(784)
target = torch.randint(0, 10, (1,)).item()
yield data, target
return torch.utils.data.DataLoader(ToyDataset(), batch_size=32)
def build_model(self):
return nn.Sequential(
nn.Linear(784, 128),
nn.ReLU(),
nn.Linear(128, 10),
).to(self.device)
def build_optimizer(self):
return torch.optim.AdamW(self.model.parameters(), lr=3e-4)
def build_hooks(self):
return [
ProgressHook(interval=50, with_records=True),
CheckpointingHook(interval=500, keep_previous=2),
]
def forward(self, batch):
x, y = batch
x, y = x.to(self.device), y.to(self.device)
logits = self.model(x)
loss = nn.functional.cross_entropy(logits, y)
accuracy = (logits.argmax(1) == y).float().mean().item()
return loss, {"accuracy": accuracy}
trainer = MyTrainer(max_steps=2000, device="cpu", workspace="runs/demo")
trainer.train()
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
trainloop-0.8.0.tar.gz
(14.0 kB
view details)
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
trainloop-0.8.0-py3-none-any.whl
(15.9 kB
view details)
File details
Details for the file trainloop-0.8.0.tar.gz.
File metadata
- Download URL: trainloop-0.8.0.tar.gz
- Upload date:
- Size: 14.0 kB
- Tags: Source
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
uv/0.12.6 {"installer":{"name":"uv","version":"0.12.6","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
bdf88476c8511b62ff69fc5b2b76f37c358c82d156ceb32e9b6423cb81345d4e
|
|
| MD5 |
b330f02975aef165b1fd5d62dbcfdc6e
|
|
| BLAKE2b-256 |
2c64ab372c22a11ba458b64150ea6755ba33b1ca41b75167a252bfb8fadf48e7
|
File details
Details for the file trainloop-0.8.0-py3-none-any.whl.
File metadata
- Download URL: trainloop-0.8.0-py3-none-any.whl
- Upload date:
- Size: 15.9 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? Yes
- Uploaded via:
uv/0.12.6 {"installer":{"name":"uv","version":"0.12.6","subcommand":["publish"]},"python":null,"implementation":{"name":null,"version":null},"distro":{"name":"Ubuntu","version":"24.04","id":"noble","libc":null},"system":{"name":null,"release":null},"cpu":null,"openssl_version":null,"setuptools_version":null,"rustc_version":null,"ci":true}
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
0d643f82aec5dce04fd919ac571f1d8a33e23f134b329b441e229ade00a59231
|
|
| MD5 |
3e68c5f8abc7aaecd92df5ca6d10e4c9
|
|
| BLAKE2b-256 |
12e335ef646a2d1e520b77890df3c52313c072de5a26218bd2ed67b27c806ea6
|