Skip to main content

trainloop

PyPI version

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.9.0.tar.gz (14.7 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

trainloop-0.9.0-py3-none-any.whl (16.5 kB view details)

Uploaded Python 3

File details

Details for the file trainloop-0.9.0.tar.gz.

File metadata

  • Download URL: trainloop-0.9.0.tar.gz
  • Upload date:
  • Size: 14.7 kB
  • Tags: Source
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: uv/0.12.7 {"installer":{"name":"uv","version":"0.12.7","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

Hashes for trainloop-0.9.0.tar.gz
Algorithm Hash digest
SHA256 26834d0d422499da8806a27a05d20506179b3a03440e6b93b3b940610cb84536
MD5 61a939bc80e6a1f59e4bcaf4d0f75ad5
BLAKE2b-256 245eafd683b74755fd120d3082f0625df3eb6199eab42dee53d4c01085f6dc18

See more details on using hashes here.

File details

Details for the file trainloop-0.9.0-py3-none-any.whl.

File metadata

  • Download URL: trainloop-0.9.0-py3-none-any.whl
  • Upload date:
  • Size: 16.5 kB
  • Tags: Python 3
  • Uploaded using Trusted Publishing? Yes
  • Uploaded via: uv/0.12.7 {"installer":{"name":"uv","version":"0.12.7","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

Hashes for trainloop-0.9.0-py3-none-any.whl
Algorithm Hash digest
SHA256 2c3a9098fc6bc579917ecab1cae05a37bf7f3f77e351d90ba1cae3ff51152435
MD5 f88446bf18d2e4e2a81e281b6431d598
BLAKE2b-256 786a54655b7a66ce8251158e56e91be888880453f9b4e1244bc6798fdbc1186b

See more details on using hashes here.

Release history Release notifications | RSS feed

0.11.0

2 files

0.10.0

2 files

This release

0.9.0 This release

2 files

0.8.0

2 files

0.7.0

2 files

0.6.0

2 files

0.5.2

2 files

0.5.1

2 files

0.5.0

2 files

0.4.0

2 files

0.3.0

2 files

0.2.0

2 files

0.1.0

2 files

Anthropic, PBC Visionary sponsor Bloomberg Visionary sponsor Hudson River Trading Visionary sponsor Meta Visionary sponsor NVIDIA Visionary sponsor Microsoft Sustainability sponsor Depot Continuous Integration AWS Cloud computing and Security Sponsor Datadog Monitoring Fastly CDN Google Download Analytics Sentry Error logging StatusPage Status page