Training Loop made easy

projects
1

:fire: Introducing Torchflow — A Lightweight Library to Streamline PyTorch Training, Tuning, and Callbacks

Hello PyTorch Community :waving_hand:,

I’m excited to share my first open-source project, Torchflow, now available on PyPI!
After months of learning, iterating, and debugging (and many nights without stable power :high_voltage:), I’ve built this library to make training and experimentation in PyTorch a lot easier — especially for learners, researchers, and data scientists working on rapid prototyping.

:light_bulb: What is Torchflow?

Torchflow is a modular library designed to help you:

  • Simplify training loops in PyTorch

  • Implement and customize callbacks for logging, early stopping, and checkpoints

  • Integrate hyperparameter tuning with minimal setup

  • Maintain a clean, readable, and reproducible experiment structure

It’s meant to reduce boilerplate code while keeping the full flexibility of native PyTorch.

:rocket: Quick Install

pip install torchflow-core

:books: Learn More

This snippet shows how a user can define a model, use Torchflow’s training utilities, and integrate callbacks — all in just a few lines.

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# Import Torchflow Trainer
from torchflow import Trainer, EarlyStopping, ModelCheckpoint

# 1️⃣ Simple dataset
X = torch.randn(500, 10)
y = torch.randint(0, 2, (500,))
dataset = TensorDataset(X, y)
loader = DataLoader(dataset, batch_size=32, shuffle=True)

# 2️⃣ Simple model
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(10, 16),
            nn.ReLU(),
            nn.Linear(16, 2)
        )

    def forward(self, x):
        return self.fc(x)

model = SimpleNet()

# 3️⃣ Define training setup
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)

# 4️⃣ Initialize Trainer with callbacks
trainer = Trainer(
    model=model,
    criterion=criterion,
    optimizer=optimizer,
    callbacks=[
        EarlyStopping(patience=3),
        ModelCheckpoint(save_best=True)
    ]
)

# 5️⃣ Train and evaluate
trainer.fit(loader, epochs=10)
trainer.evaluate(loader)

:handshake: Get Involved

I’d love your feedback, ideas, or contributions!
Check out the CONTRIBUTION.md file in the repo to see how you can help improve Torchflow.

This project grew out of my experience as a Tech4Dev Developer Foundry Fellow (AI/ML Track) — a program that helped me turn knowledge into impactful projects. I hope it helps others in the same way.

Thanks to the PyTorch community for being such an inspiring and supportive space. :folded_hands: