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.8.0.tar.gz (14.0 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.8.0-py3-none-any.whl (15.9 kB view details)

Uploaded Python 3

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

Hashes for trainloop-0.8.0.tar.gz
Algorithm Hash digest
SHA256 bdf88476c8511b62ff69fc5b2b76f37c358c82d156ceb32e9b6423cb81345d4e
MD5 b330f02975aef165b1fd5d62dbcfdc6e
BLAKE2b-256 2c64ab372c22a11ba458b64150ea6755ba33b1ca41b75167a252bfb8fadf48e7

See more details on using hashes here.

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

Hashes for trainloop-0.8.0-py3-none-any.whl
Algorithm Hash digest
SHA256 0d643f82aec5dce04fd919ac571f1d8a33e23f134b329b441e229ade00a59231
MD5 3e68c5f8abc7aaecd92df5ca6d10e4c9
BLAKE2b-256 12e335ef646a2d1e520b77890df3c52313c072de5a26218bd2ed67b27c806ea6

See more details on using hashes here.

Release history Release notifications | RSS feed

0.11.0

2 files

0.10.0

2 files

0.9.0

2 files

This release

0.8.0 This release

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