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.10.0.tar.gz (14.9 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.10.0-py3-none-any.whl (16.7 kB view details)

Uploaded Python 3

File details

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

File metadata

  • Download URL: trainloop-0.10.0.tar.gz
  • Upload date:
  • Size: 14.9 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.10.0.tar.gz
Algorithm Hash digest
SHA256 c664f357306adda4e859168a81da0cb6a67315ef526c65c26c33396d8baac569
MD5 8a967d224c80ae884e38a557b5a6115f
BLAKE2b-256 57f8bbc6382b58856814a449256d8a117610d7e0f4a1715ea8d1713d3de83ae5

See more details on using hashes here.

File details

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

File metadata

  • Download URL: trainloop-0.10.0-py3-none-any.whl
  • Upload date:
  • Size: 16.7 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.10.0-py3-none-any.whl
Algorithm Hash digest
SHA256 a3227d7b020936b1b0afe63cb563047c9d446b87cceb82cc084655ed7b622982
MD5 9536e35f8734dc3f5e08ab382feb3b3f
BLAKE2b-256 0194126405572627ff00032c00ca4efd971f6c75f30ecd9b55a58db05b9d4886

See more details on using hashes here.

Release history Release notifications | RSS feed

0.11.0

2 files

This release

0.10.0 This release

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