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.11.0.tar.gz (15.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.11.0-py3-none-any.whl (16.8 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: trainloop-0.11.0.tar.gz
  • Upload date:
  • Size: 15.0 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.11.0.tar.gz
Algorithm Hash digest
SHA256 a35945a6c7271dd98a5a8fd82a85718f0c90c0e524a1c11ac929a3ce35ed089b
MD5 0bfed29a53e3dc905011d850a84bc34f
BLAKE2b-256 49fa12234f6ebdbaa78756c294fa95d72530769c5f08db24dffd59ed5fafbe0c

See more details on using hashes here.

File details

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

File metadata

  • Download URL: trainloop-0.11.0-py3-none-any.whl
  • Upload date:
  • Size: 16.8 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.11.0-py3-none-any.whl
Algorithm Hash digest
SHA256 b011370ea5c7eba8175706b1e0e27bfc61cfc32f03f5ac1a99fbad04820570b0
MD5 3d557f14985c62919376f759f032546b
BLAKE2b-256 dca94dec4c5b96cadfd3c486735ccd1d481760c5ff9b67f713b433db801734cc

See more details on using hashes here.

Release history Release notifications | RSS feed

This release

0.11.0 This release

2 files

0.10.0

2 files

0.9.0

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