torchloop 0.1.0__tar.gz

torchloop-0.1.0/.github/workflows/ci.yml

@@ -0,0 +1,32 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+    branches: [main]
+
+jobs:
+  test:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.9", "3.10", "3.11"]
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v5
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install dependencies
+        run: |
+          pip install -e ".[dev]"
+
+      - name: Lint with ruff
+        run: ruff check src/
+
+      - name: Run tests
+        run: pytest

torchloop-0.1.0/.github/workflows/publish.yml

@@ -0,0 +1,30 @@
+name: Publish to PyPI
+
+on:
+  push:
+    tags:
+      - "v*.*.*"
+
+jobs:
+  publish:
+    runs-on: ubuntu-latest
+    environment: pypi
+    permissions:
+      id-token: write
+
+    steps:
+      - uses: actions/checkout@v4
+
+      - name: Set up Python
+        uses: actions/setup-python@v5
+        with:
+          python-version: "3.11"
+
+      - name: Install hatch
+        run: pip install hatch
+
+      - name: Build package
+        run: hatch build
+
+      - name: Publish to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1

torchloop-0.1.0/.gitignore

@@ -0,0 +1,30 @@
+# Python bytecode and caches
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Virtual environments
+venv/
+.venv/
+env/
+ENV/
+
+# Packaging/build artifacts
+build/
+dist/
+*.egg-info/
+.eggs/
+
+# Test, coverage, and tooling artifacts
+.pytest_cache/
+.coverage
+.coverage.*
+htmlcov/
+.mypy_cache/
+.ruff_cache/
+
+# Editor/OS files
+.vscode/
+.idea/
+.DS_Store
+Thumbs.db

torchloop-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Tharun K
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torchloop-0.1.0/PKG-INFO

@@ -0,0 +1,162 @@
+Metadata-Version: 2.4
+Name: torchloop
+Version: 0.1.0
+Summary: Lightweight PyTorch utility library for training, evaluation, and TFLite export — without the framework lock-in.
+Project-URL: Homepage, https://github.com/Tharun007-TK/torchloop
+Project-URL: Repository, https://github.com/Tharun007-TK/torchloop
+Project-URL: Issues, https://github.com/Tharun007-TK/torchloop/issues
+Author-email: Tharun Kumar <tharunkumarvmt@gmail.com>
+License: MIT License
+        
+        Copyright (c) 2025 Tharun K
+        
+        Permission is hereby granted, free of charge, to any person obtaining a copy
+        of this software and associated documentation files (the "Software"), to deal
+        in the Software without restriction, including without limitation the rights
+        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+        copies of the Software, and to permit persons to whom the Software is
+        furnished to do so, subject to the following conditions:
+        
+        The above copyright notice and this permission notice shall be included in all
+        copies or substantial portions of the Software.
+        
+        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+        SOFTWARE.
+License-File: LICENSE
+Keywords: deep learning,export,ml utilities,pytorch,tflite,training
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Developers
+Classifier: Intended Audience :: Science/Research
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Requires-Dist: matplotlib>=3.7.0
+Requires-Dist: numpy>=1.24.0
+Requires-Dist: scikit-learn>=1.3.0
+Requires-Dist: torch>=2.0.0
+Requires-Dist: torchvision>=0.15.0
+Requires-Dist: tqdm>=4.65.0
+Provides-Extra: all
+Requires-Dist: hatch>=1.7.0; extra == 'all'
+Requires-Dist: onnx>=1.14.0; extra == 'all'
+Requires-Dist: onnxruntime>=1.15.0; extra == 'all'
+Requires-Dist: pytest-cov>=4.1.0; extra == 'all'
+Requires-Dist: pytest>=7.4.0; extra == 'all'
+Requires-Dist: ruff>=0.1.0; extra == 'all'
+Requires-Dist: tensorflow>=2.13.0; extra == 'all'
+Provides-Extra: dev
+Requires-Dist: hatch>=1.7.0; extra == 'dev'
+Requires-Dist: pytest-cov>=4.1.0; extra == 'dev'
+Requires-Dist: pytest>=7.4.0; extra == 'dev'
+Requires-Dist: ruff>=0.1.0; extra == 'dev'
+Provides-Extra: export
+Requires-Dist: onnx>=1.14.0; extra == 'export'
+Requires-Dist: onnxruntime>=1.15.0; extra == 'export'
+Requires-Dist: tensorflow>=2.13.0; extra == 'export'
+Description-Content-Type: text/markdown
+
+# torchloop
+
+> Lightweight PyTorch utility library for training, evaluation, and TFLite export — without the framework lock-in.
+
+[![CI](https://github.com/Tharun007-TK/torchloop/actions/workflows/ci.yml/badge.svg)](https://github.com/Tharun007-TK/torchloop/actions)
+[![PyPI](https://img.shields.io/pypi/v/torchloop)](https://pypi.org/project/torchloop/)
+[![Python](https://img.shields.io/pypi/pyversions/torchloop)](https://pypi.org/project/torchloop/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+
+---
+
+## The Problem
+
+You write the same PyTorch training loop in every project. Same checkpoint logic. Same metric assembly. Same TFLite export steps. It's tedious and inconsistent.
+
+`torchloop` abstracts exactly that — nothing more.
+
+---
+
+## Install
+
+```bash
+pip install torchloop
+
+# With TFLite export support
+pip install torchloop[export]
+```
+
+---
+
+## Usage
+
+### Training
+
+```python
+from torchloop import Trainer
+
+trainer = Trainer(
+    model,
+    optimizer=torch.optim.Adam(model.parameters()),
+    criterion=torch.nn.CrossEntropyLoss(),
+    device="cuda",
+    patience=5,           # early stopping
+)
+
+history = trainer.fit(train_loader, val_loader, epochs=30)
+trainer.save("best.pt")
+```
+
+### Evaluation
+
+```python
+from torchloop import Evaluator
+
+ev = Evaluator(model, device="cuda")
+results = ev.report(val_loader, class_names=["No Damage", "Minor", "Major", "Destroyed"])
+# prints sklearn classification report
+
+fig = ev.confusion_matrix(val_loader)
+fig.savefig("cm.png")
+
+per_class = ev.f1_per_class(val_loader)
+# {'No Damage': 0.91, 'Minor': 0.78, ...}
+```
+
+### Export
+
+```python
+from torchloop.exporter import Exporter
+
+exp = Exporter(model, input_shape=(1, 3, 224, 224))
+exp.to_onnx("model.onnx")
+exp.to_tflite("model.tflite", quantize=True)
+```
+
+---
+
+## Design Principles
+
+- **No lock-in**: Works with any nn.Module. No subclassing required.
+- **Minimal surface area**: Three modules. That's it.
+- **You own the model**: torchloop wraps your loop, doesn't replace your architecture.
+
+---
+
+## Roadmap
+
+- [ ] `v0.1.0` — Trainer, Evaluator, Exporter
+- [ ] `v0.2.0` — LR scheduler support, mixed precision (AMP)
+- [ ] `v0.3.0` — W&B / MLflow logging hooks
+
+---
+
+## License
+
+MIT

torchloop-0.1.0/README.md

@@ -0,0 +1,96 @@
+# torchloop
+
+> Lightweight PyTorch utility library for training, evaluation, and TFLite export — without the framework lock-in.
+
+[![CI](https://github.com/Tharun007-TK/torchloop/actions/workflows/ci.yml/badge.svg)](https://github.com/Tharun007-TK/torchloop/actions)
+[![PyPI](https://img.shields.io/pypi/v/torchloop)](https://pypi.org/project/torchloop/)
+[![Python](https://img.shields.io/pypi/pyversions/torchloop)](https://pypi.org/project/torchloop/)
+[![License: MIT](https://img.shields.io/badge/License-MIT-yellow.svg)](LICENSE)
+
+---
+
+## The Problem
+
+You write the same PyTorch training loop in every project. Same checkpoint logic. Same metric assembly. Same TFLite export steps. It's tedious and inconsistent.
+
+`torchloop` abstracts exactly that — nothing more.
+
+---
+
+## Install
+
+```bash
+pip install torchloop
+
+# With TFLite export support
+pip install torchloop[export]
+```
+
+---
+
+## Usage
+
+### Training
+
+```python
+from torchloop import Trainer
+
+trainer = Trainer(
+    model,
+    optimizer=torch.optim.Adam(model.parameters()),
+    criterion=torch.nn.CrossEntropyLoss(),
+    device="cuda",
+    patience=5,           # early stopping
+)
+
+history = trainer.fit(train_loader, val_loader, epochs=30)
+trainer.save("best.pt")
+```
+
+### Evaluation
+
+```python
+from torchloop import Evaluator
+
+ev = Evaluator(model, device="cuda")
+results = ev.report(val_loader, class_names=["No Damage", "Minor", "Major", "Destroyed"])
+# prints sklearn classification report
+
+fig = ev.confusion_matrix(val_loader)
+fig.savefig("cm.png")
+
+per_class = ev.f1_per_class(val_loader)
+# {'No Damage': 0.91, 'Minor': 0.78, ...}
+```
+
+### Export
+
+```python
+from torchloop.exporter import Exporter
+
+exp = Exporter(model, input_shape=(1, 3, 224, 224))
+exp.to_onnx("model.onnx")
+exp.to_tflite("model.tflite", quantize=True)
+```
+
+---
+
+## Design Principles
+
+- **No lock-in**: Works with any nn.Module. No subclassing required.
+- **Minimal surface area**: Three modules. That's it.
+- **You own the model**: torchloop wraps your loop, doesn't replace your architecture.
+
+---
+
+## Roadmap
+
+- [ ] `v0.1.0` — Trainer, Evaluator, Exporter
+- [ ] `v0.2.0` — LR scheduler support, mixed precision (AMP)
+- [ ] `v0.3.0` — W&B / MLflow logging hooks
+
+---
+
+## License
+
+MIT

torchloop-0.1.0/pyproject.toml

@@ -0,0 +1,68 @@
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project]
+name = "torchloop"
+version = "0.1.0"
+description = "Lightweight PyTorch utility library for training, evaluation, and TFLite export — without the framework lock-in."
+readme = "README.md"
+license = { file = "LICENSE" }
+requires-python = ">=3.9"
+authors = [
+  { name = "Tharun Kumar", email = "tharunkumarvmt@gmail.com" }
+]
+keywords = ["pytorch", "deep learning", "training", "tflite", "export", "ml utilities"]
+classifiers = [
+  "Development Status :: 3 - Alpha",
+  "Intended Audience :: Developers",
+  "Intended Audience :: Science/Research",
+  "License :: OSI Approved :: MIT License",
+  "Programming Language :: Python :: 3",
+  "Programming Language :: Python :: 3.9",
+  "Programming Language :: Python :: 3.10",
+  "Programming Language :: Python :: 3.11",
+  "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+
+dependencies = [
+  "torch>=2.0.0",
+  "torchvision>=0.15.0",
+  "scikit-learn>=1.3.0",
+  "numpy>=1.24.0",
+  "matplotlib>=3.7.0",
+  "tqdm>=4.65.0",
+]
+
+[project.optional-dependencies]
+export = [
+  "onnx>=1.14.0",
+  "onnxruntime>=1.15.0",
+  "tensorflow>=2.13.0",   # needed for TFLite conversion
+]
+dev = [
+  "pytest>=7.4.0",
+  "pytest-cov>=4.1.0",
+  "ruff>=0.1.0",          # linter + formatter
+  "hatch>=1.7.0",
+]
+all = ["torchloop[export,dev]"]
+
+[project.urls]
+Homepage = "https://github.com/Tharun007-TK/torchloop"
+Repository = "https://github.com/Tharun007-TK/torchloop"
+Issues = "https://github.com/Tharun007-TK/torchloop/issues"
+
+[tool.hatch.build.targets.wheel]
+packages = ["src/torchloop"]
+
+[tool.ruff]
+line-length = 88
+target-version = "py39"
+
+[tool.ruff.lint]
+select = ["E", "F", "I"]   # pycodestyle + pyflakes + isort
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+addopts = "--cov=src/torchloop --cov-report=term-missing"

torchloop-0.1.0/src/torchloop/__init__.py

@@ -0,0 +1,16 @@
+"""
+torchloop — Lightweight PyTorch utility library.
+
+Modules:
+    trainer   : Training loop, metric logging, checkpoint management
+    evaluator : Classification report, confusion matrix, per-class F1
+    exporter  : PyTorch → ONNX → TFLite with optional quantization
+"""
+
+__version__ = "0.1.0"
+__author__ = "Tharun Kumar"
+
+from torchloop.evaluator import Evaluator
+from torchloop.trainer import Trainer
+
+__all__ = ["Trainer", "Evaluator", "__version__"]

torchloop-0.1.0/src/torchloop/evaluator.py

@@ -0,0 +1,128 @@
+"""
+torchloop.evaluator
+-------------------
+One-call classification diagnostics. No more assembling sklearn +
+matplotlib calls manually across every project.
+
+Usage:
+    from torchloop import Evaluator
+
+    ev = Evaluator(model, device="cuda")
+    ev.report(val_loader, class_names=["cat", "dog"])
+    ev.confusion_matrix(val_loader)
+"""
+
+from __future__ import annotations
+
+from typing import Optional
+
+import matplotlib.pyplot as plt
+import numpy as np
+import torch
+import torch.nn as nn
+from sklearn.metrics import (
+    ConfusionMatrixDisplay,
+    classification_report,
+    confusion_matrix,
+    f1_score,
+)
+from torch.utils.data import DataLoader
+
+
+class Evaluator:
+    """
+    Classification model evaluator.
+
+    Args:
+        model       : Trained nn.Module.
+        device      : 'cuda', 'cpu', or 'mps'. Auto-detects if None.
+    """
+
+    def __init__(self, model: nn.Module, device: Optional[str] = None):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = model.to(self.device)
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def report(
+        self,
+        loader: DataLoader,
+        class_names: Optional[list[str]] = None,
+    ) -> dict:
+        """
+        Print full sklearn classification report.
+
+        Returns:
+            dict with keys: accuracy, macro_f1, weighted_f1, per_class_f1
+        """
+        preds, targets = self._infer(loader)
+        report = classification_report(
+            targets, preds, target_names=class_names, zero_division=0
+        )
+        print(report)
+        per_class = f1_score(targets, preds, average=None, zero_division=0).tolist()
+        return {
+            "accuracy": float((np.array(preds) == np.array(targets)).mean()),
+            "macro_f1": float(
+                f1_score(targets, preds, average="macro", zero_division=0)
+            ),
+            "weighted_f1": float(
+                f1_score(targets, preds, average="weighted", zero_division=0)
+            ),
+            "per_class_f1": {
+                (class_names[i] if class_names else str(i)): round(v, 4)
+                for i, v in enumerate(per_class)
+            },
+        }
+
+    def confusion_matrix(
+        self,
+        loader: DataLoader,
+        class_names: Optional[list[str]] = None,
+        normalize: Optional[str] = "true",   # 'true' | 'pred' | 'all' | None
+        figsize: tuple = (8, 6),
+    ) -> plt.Figure:
+        """
+        Plot and return confusion matrix figure.
+        """
+        preds, targets = self._infer(loader)
+        cm = confusion_matrix(targets, preds, normalize=normalize)
+        fig, ax = plt.subplots(figsize=figsize)
+        disp = ConfusionMatrixDisplay(cm, display_labels=class_names)
+        disp.plot(ax=ax, colorbar=True, cmap="Blues")
+        ax.set_title("Confusion Matrix")
+        plt.tight_layout()
+        return fig
+
+    def f1_per_class(
+        self,
+        loader: DataLoader,
+        class_names: Optional[list[str]] = None,
+    ) -> dict[str, float]:
+        """
+        Returns per-class F1 as a dict. Clean for logging to W&B or MLflow.
+        """
+        preds, targets = self._infer(loader)
+        scores = f1_score(targets, preds, average=None, zero_division=0)
+        return {
+            (class_names[i] if class_names else str(i)): round(float(s), 4)
+            for i, s in enumerate(scores)
+        }
+
+    # ------------------------------------------------------------------
+    # Internal
+    # ------------------------------------------------------------------
+
+    def _infer(self, loader: DataLoader) -> tuple[list, list]:
+        self.model.eval()
+        all_preds, all_targets = [], []
+        with torch.no_grad():
+            for inputs, targets in loader:
+                inputs = inputs.to(self.device)
+                outputs = self.model(inputs)
+                preds = outputs.argmax(dim=1).cpu().tolist()
+                all_preds.extend(preds)
+                all_targets.extend(targets.tolist())
+        return all_preds, all_targets

torchloop-0.1.0/src/torchloop/exporter.py

@@ -0,0 +1,137 @@
+"""
+torchloop.exporter
+------------------
+PyTorch → ONNX → TFLite in one place.
+Requires: pip install torchloop[export]
+
+Usage:
+    from torchloop.exporter import Exporter
+
+    exp = Exporter(model, input_shape=(1, 3, 224, 224))
+    exp.to_onnx("model.onnx")
+    exp.to_tflite("model.tflite", quantize=True)
+"""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+
+class Exporter:
+    """
+    Handles model export from PyTorch to ONNX and TFLite.
+
+    Args:
+        model       : Trained nn.Module (will be set to eval mode).
+        input_shape : Tuple describing one sample input e.g. (1, 3, 224, 224).
+        device      : Device to run dummy forward pass on.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        input_shape: tuple,
+        device: Optional[str] = None,
+    ):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = model.to(self.device).eval()
+        self.input_shape = input_shape
+        self._dummy = torch.randn(*input_shape).to(self.device)
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def to_onnx(self, path: str | Path, opset: int = 17) -> Path:
+        """
+        Export model to ONNX format.
+
+        Args:
+            path  : Output .onnx file path.
+            opset : ONNX opset version. Default 17 covers most torch ops.
+
+        Returns:
+            Resolved path to exported file.
+        """
+        try:
+            import onnx
+        except ImportError:
+            raise ImportError(
+                "onnx is not installed. Run: pip install torchloop[export]"
+            )
+
+        path = Path(path)
+        torch.onnx.export(
+            self.model,
+            self._dummy,
+            str(path),
+            opset_version=opset,
+            input_names=["input"],
+            output_names=["output"],
+            dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}},
+        )
+        model_onnx = onnx.load(str(path))
+        onnx.checker.check_model(model_onnx)
+        print(f"  ONNX export verified → {path}")
+        return path
+
+    def to_tflite(
+        self,
+        path: str | Path,
+        quantize: bool = False,
+        onnx_path: Optional[str | Path] = None,
+    ) -> Path:
+        """
+        Export model to TFLite via ONNX → TF → TFLite pipeline.
+
+        Args:
+            path      : Output .tflite file path.
+            quantize  : If True, applies dynamic range quantization.
+            onnx_path : Intermediate .onnx file path. Auto-generated if None.
+
+        Returns:
+            Resolved path to exported .tflite file.
+
+        Note:
+            Requires tensorflow and onnx2tf installed.
+            pip install torchloop[export] onnx2tf
+        """
+        try:
+            import onnx2tf
+            import tensorflow as tf
+        except ImportError:
+            raise ImportError(
+                "tensorflow or onnx2tf not installed.\n"
+                "Run: pip install torchloop[export] onnx2tf"
+            )
+
+        path = Path(path)
+
+        # Step 1: Export to ONNX first
+        _onnx_path = Path(onnx_path) if onnx_path else path.with_suffix(".onnx")
+        self.to_onnx(_onnx_path)
+
+        # Step 2: ONNX → SavedModel via onnx2tf
+        saved_model_dir = path.parent / "_tflite_savedmodel_tmp"
+        onnx2tf.convert(
+            input_onnx_file_path=str(_onnx_path),
+            output_folder_path=str(saved_model_dir),
+            not_use_onnxsim=False,
+            verbosity="error",
+        )
+
+        # Step 3: SavedModel → TFLite
+        converter = tf.lite.TFLiteConverter.from_saved_model(str(saved_model_dir))
+        if quantize:
+            converter.optimizations = [tf.lite.Optimize.DEFAULT]
+            print("  Quantization: dynamic range enabled.")
+
+        tflite_model = converter.convert()
+        path.write_bytes(tflite_model)
+        size_kb = path.stat().st_size / 1024
+        print(f"  TFLite export → {path}  ({size_kb:.1f} KB)")
+        return path

torchloop-0.1.0/src/torchloop/trainer.py

@@ -0,0 +1,176 @@
+"""
+torchloop.trainer
+-----------------
+Wraps the PyTorch training loop so you stop rewriting it.
+
+Usage:
+    from torchloop import Trainer
+
+    trainer = Trainer(model, optimizer, criterion, device="cuda")
+    trainer.fit(train_loader, val_loader, epochs=20)
+    trainer.save("best.pt")
+"""
+
+from __future__ import annotations
+
+import time
+from pathlib import Path
+from typing import Callable, Optional
+
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+
+class Trainer:
+    """
+    Minimal, opinionated PyTorch training loop.
+
+    Args:
+        model       : nn.Module to train.
+        optimizer   : Any torch.optim optimizer.
+        criterion   : Loss function (nn.Module or callable).
+        device      : 'cuda', 'cpu', or 'mps'. Auto-detects if None.
+        metric_fn   : Optional callable(preds, targets) → float for val metric.
+        patience    : Early stopping patience (epochs). None = disabled.
+    """
+
+    def __init__(
+        self,
+        model: nn.Module,
+        optimizer: torch.optim.Optimizer,
+        criterion: Callable,
+        device: Optional[str] = None,
+        metric_fn: Optional[Callable] = None,
+        patience: Optional[int] = None,
+    ):
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = model.to(self.device)
+        self.optimizer = optimizer
+        self.criterion = criterion
+        self.metric_fn = metric_fn
+        self.patience = patience
+
+        self.history: dict[str, list] = {
+            "train_loss": [],
+            "val_loss": [],
+            "val_metric": [],
+        }
+        self._best_val_loss = float("inf")
+        self._best_state: Optional[dict] = None
+        self._no_improve_count = 0
+
+    # ------------------------------------------------------------------
+    # Public API
+    # ------------------------------------------------------------------
+
+    def fit(
+        self,
+        train_loader: DataLoader,
+        val_loader: Optional[DataLoader] = None,
+        epochs: int = 10,
+    ) -> dict:
+        """
+        Train the model.
+
+        Returns:
+            history dict with train_loss, val_loss, val_metric per epoch.
+        """
+        for epoch in range(1, epochs + 1):
+            t0 = time.time()
+            train_loss = self._train_epoch(train_loader)
+            self.history["train_loss"].append(train_loss)
+
+            val_loss, val_metric = None, None
+            if val_loader is not None:
+                val_loss, val_metric = self._val_epoch(val_loader)
+                self.history["val_loss"].append(val_loss)
+                self.history["val_metric"].append(val_metric)
+                self._checkpoint(val_loss)
+
+            self._log(epoch, epochs, train_loss, val_loss, val_metric, time.time() - t0)
+
+            if self._should_stop():
+                print(f"  Early stopping triggered at epoch {epoch}.")
+                break
+
+        if self._best_state is not None:
+            self.model.load_state_dict(self._best_state)
+            print("  Restored best model weights.")
+
+        return self.history
+
+    def save(self, path: str | Path) -> None:
+        """Save model state dict to path."""
+        torch.save(self.model.state_dict(), path)
+        print(f"  Saved → {path}")
+
+    def load(self, path: str | Path) -> None:
+        """Load model state dict from path."""
+        self.model.load_state_dict(torch.load(path, map_location=self.device))
+        print(f"  Loaded ← {path}")
+
+    # ------------------------------------------------------------------
+    # Internal
+    # ------------------------------------------------------------------
+
+    def _train_epoch(self, loader: DataLoader) -> float:
+        self.model.train()
+        total_loss = 0.0
+        for inputs, targets in tqdm(loader, desc="  train", leave=False):
+            inputs, targets = inputs.to(self.device), targets.to(self.device)
+            self.optimizer.zero_grad()
+            outputs = self.model(inputs)
+            loss = self.criterion(outputs, targets)
+            loss.backward()
+            self.optimizer.step()
+            total_loss += loss.item() * inputs.size(0)
+        return total_loss / len(loader.dataset)
+
+    def _val_epoch(self, loader: DataLoader) -> tuple[float, Optional[float]]:
+        self.model.eval()
+        total_loss = 0.0
+        all_preds, all_targets = [], []
+        with torch.no_grad():
+            for inputs, targets in tqdm(loader, desc="  val  ", leave=False):
+                inputs, targets = inputs.to(self.device), targets.to(self.device)
+                outputs = self.model(inputs)
+                loss = self.criterion(outputs, targets)
+                total_loss += loss.item() * inputs.size(0)
+                if self.metric_fn is not None:
+                    all_preds.append(outputs.cpu())
+                    all_targets.append(targets.cpu())
+        avg_loss = total_loss / len(loader.dataset)
+        metric = None
+        if self.metric_fn is not None and all_preds:
+            metric = self.metric_fn(
+                torch.cat(all_preds), torch.cat(all_targets)
+            )
+        return avg_loss, metric
+
+    def _checkpoint(self, val_loss: float) -> None:
+        if val_loss < self._best_val_loss:
+            self._best_val_loss = val_loss
+            self._best_state = {
+                k: v.clone() for k, v in self.model.state_dict().items()
+            }
+            self._no_improve_count = 0
+        else:
+            self._no_improve_count += 1
+
+    def _should_stop(self) -> bool:
+        return (
+            self.patience is not None
+            and self._no_improve_count >= self.patience
+        )
+
+    @staticmethod
+    def _log(epoch, epochs, train_loss, val_loss, val_metric, elapsed) -> None:
+        parts = [f"Epoch [{epoch:>3}/{epochs}]", f"train_loss={train_loss:.4f}"]
+        if val_loss is not None:
+            parts.append(f"val_loss={val_loss:.4f}")
+        if val_metric is not None:
+            parts.append(f"val_metric={val_metric:.4f}")
+        parts.append(f"({elapsed:.1f}s)")
+        print("  " + "  ".join(parts))

torchloop-0.1.0/tests/test_evaluator.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, TensorDataset
+
+from torchloop import Evaluator
+
+
+def _make_loader():
+    X = torch.randn(64, 16)
+    y = torch.randint(0, 3, (64,))
+    return DataLoader(TensorDataset(X, y), batch_size=16)
+
+
+def _make_model():
+    return nn.Sequential(nn.Linear(16, 32), nn.ReLU(), nn.Linear(32, 3))
+
+
+def test_report_returns_keys():
+    model = _make_model()
+    ev = Evaluator(model, device="cpu")
+    result = ev.report(_make_loader(), class_names=["a", "b", "c"])
+    assert "macro_f1" in result
+    assert "per_class_f1" in result
+    assert set(result["per_class_f1"].keys()) == {"a", "b", "c"}
+
+
+def test_f1_per_class_length():
+    model = _make_model()
+    ev = Evaluator(model, device="cpu")
+    scores = ev.f1_per_class(_make_loader())
+    assert len(scores) == 3
+
+
+def test_confusion_matrix_returns_figure():
+    import matplotlib.pyplot as plt
+    model = _make_model()
+    ev = Evaluator(model, device="cpu")
+    fig = ev.confusion_matrix(_make_loader())
+    assert isinstance(fig, plt.Figure)
+    plt.close(fig)

torchloop-0.1.0/tests/test_trainer.py

@@ -0,0 +1,62 @@
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, TensorDataset
+
+from torchloop import Trainer
+
+
+def _make_loader(n=64, features=16, classes=3, batch=16):
+    X = torch.randn(n, features)
+    y = torch.randint(0, classes, (n,))
+    return DataLoader(TensorDataset(X, y), batch_size=batch)
+
+
+def _make_model(features=16, classes=3):
+    return nn.Sequential(nn.Linear(features, 32), nn.ReLU(), nn.Linear(32, classes))
+
+
+def test_trainer_fit_returns_history():
+    model = _make_model()
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+    criterion = nn.CrossEntropyLoss()
+    trainer = Trainer(model, optimizer, criterion, device="cpu")
+    history = trainer.fit(_make_loader(), _make_loader(), epochs=3)
+    assert "train_loss" in history
+    assert len(history["train_loss"]) == 3
+
+
+def test_trainer_early_stopping():
+    model = _make_model()
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+    criterion = nn.CrossEntropyLoss()
+    trainer = Trainer(model, optimizer, criterion, device="cpu", patience=2)
+    history = trainer.fit(_make_loader(), _make_loader(), epochs=20)
+    assert len(history["train_loss"]) <= 20
+
+
+def test_trainer_save_load(tmp_path):
+    model = _make_model()
+    optimizer = torch.optim.Adam(model.parameters())
+    criterion = nn.CrossEntropyLoss()
+    trainer = Trainer(model, optimizer, criterion, device="cpu")
+    trainer.fit(_make_loader(), epochs=1)
+    save_path = tmp_path / "model.pt"
+    trainer.save(save_path)
+    assert save_path.exists()
+    trainer.load(save_path)
+
+def test_trainer_with_metric_fn():
+    from sklearn.metrics import f1_score as skf1
+
+    def metric_fn(preds, targets):
+        p = preds.argmax(dim=1).numpy()
+        t = targets.numpy()
+        return skf1(t, p, average="macro", zero_division=0)
+
+    model = _make_model()
+    optimizer = torch.optim.Adam(model.parameters())
+    criterion = nn.CrossEntropyLoss()
+    trainer = Trainer(model, optimizer, criterion, device="cpu", metric_fn=metric_fn)
+    history = trainer.fit(_make_loader(), _make_loader(), epochs=2)
+    assert "val_metric" in history
+    assert len(history["val_metric"]) == 2