torchloop 0.2.3__tar.gz → 0.3.0__tar.gz

{torchloop-0.2.3 → torchloop-0.3.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: torchloop
-Version: 0.2.3
+Version: 0.3.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
@@ -47,19 +47,17 @@ Requires-Dist: torchvision>=0.15.0
 Requires-Dist: tqdm>=4.65.0
 Provides-Extra: all
 Requires-Dist: black>=23.0; extra == 'all'
-Requires-Dist: coremltools>=7.0; (platform_system == 'Darwin') and extra == 'all'
 Requires-Dist: hatch>=1.7.0; extra == 'all'
+Requires-Dist: mlflow>=2.0.0; extra == 'all'
 Requires-Dist: mypy>=1.0; extra == 'all'
 Requires-Dist: onnx>=1.14.0; extra == 'all'
-Requires-Dist: onnx>=1.15; extra == 'all'
 Requires-Dist: onnxruntime>=1.15.0; extra == 'all'
-Requires-Dist: onnxruntime>=1.16; extra == 'all'
 Requires-Dist: pytest-cov>=4.0; extra == 'all'
 Requires-Dist: pytest>=7.0; extra == 'all'
 Requires-Dist: ruff>=0.1.0; extra == 'all'
-Requires-Dist: tensorflow-lite-runtime>=2.15; (platform_system != 'Darwin') and extra == 'all'
 Requires-Dist: tensorflow>=2.13.0; extra == 'all'
 Requires-Dist: torchinfo>=1.8; extra == 'all'
+Requires-Dist: wandb>=0.15.0; extra == 'all'
 Provides-Extra: dev
 Requires-Dist: black>=23.0; extra == 'dev'
 Requires-Dist: hatch>=1.7.0; extra == 'dev'
@@ -68,6 +66,10 @@ Requires-Dist: pytest-cov>=4.0; extra == 'dev'
 Requires-Dist: pytest>=7.0; extra == 'dev'
 Requires-Dist: ruff>=0.1.0; extra == 'dev'
 Requires-Dist: torchinfo>=1.8; extra == 'dev'
+Provides-Extra: docs
+Requires-Dist: mkdocs-material>=9.0.0; extra == 'docs'
+Requires-Dist: mkdocs>=1.5.0; extra == 'docs'
+Requires-Dist: mkdocstrings[python]>=0.24.0; extra == 'docs'
 Provides-Extra: edge
 Requires-Dist: coremltools>=7.0; (platform_system == 'Darwin') and extra == 'edge'
 Requires-Dist: onnx>=1.15; extra == 'edge'
@@ -77,6 +79,9 @@ 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'
+Provides-Extra: logging
+Requires-Dist: mlflow>=2.0.0; extra == 'logging'
+Requires-Dist: wandb>=0.15.0; extra == 'logging'
 Description-Content-Type: text/markdown
 
 # torchloop

{torchloop-0.2.3 → torchloop-0.3.0}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "torchloop"
-version = "0.2.3"
+version = "0.3.0"
 description = "Lightweight PyTorch utility library for training, evaluation, and TFLite export — without the framework lock-in."
 readme = "README.md"
 license = { file = "LICENSE" }
@@ -41,6 +41,15 @@ edge = [
   "tensorflow-lite-runtime>=2.15; platform_system != 'Darwin'",
   "coremltools>=7.0; platform_system == 'Darwin'",
 ]
+logging = [
+  "wandb>=0.15.0",
+  "mlflow>=2.0.0",
+]
+docs = [
+  "mkdocs>=1.5.0",
+  "mkdocs-material>=9.0.0",
+  "mkdocstrings[python]>=0.24.0",
+]
 export = [
   "onnx>=1.14.0",
   "onnxruntime>=1.15.0",
@@ -55,7 +64,7 @@ dev = [
   "ruff>=0.1.0",          # linter + formatter
   "hatch>=1.7.0",
 ]
-all = ["torchloop[edge,export,dev]"]
+all = ["torchloop[export,dev,logging]"]
 
 [project.urls]
 Homepage = "https://github.com/Tharun007-TK/torchloop"

{torchloop-0.2.3 → torchloop-0.3.0}/src/torchloop/__init__.py

@@ -7,16 +7,10 @@ Modules:
     exporter  : PyTorch → ONNX → TFLite with optional quantization
 """
 
-__version__ = "0.2.3"
+__version__ = "0.3.0"
 __author__ = "Tharun Kumar"
 
-from torchloop.callbacks import (
-    CSVLogger,
-    Callback,
-    EarlyStopping,
-    ModelCheckpoint,
-    StopTraining,
-)
+from torchloop.callbacks import Callback, MLflowLogger, WandBLogger
 from torchloop.edge import deploy_to_edge, estimate_model
 from torchloop.evaluator import Evaluator
 from torchloop.exporter import Exporter
@@ -27,10 +21,8 @@ __all__ = [
     "Evaluator",
     "Exporter",
     "Callback",
-    "EarlyStopping",
-    "ModelCheckpoint",
-    "CSVLogger",
-    "StopTraining",
+    "WandBLogger",
+    "MLflowLogger",
     "deploy_to_edge",
     "estimate_model",
     "__version__",

torchloop-0.3.0/src/torchloop/callbacks/__init__.py

@@ -0,0 +1,7 @@
+"""Callback integrations for torchloop."""
+
+from torchloop.callbacks.base import Callback
+from torchloop.callbacks.mlflow_logger import MLflowLogger
+from torchloop.callbacks.wandb_logger import WandBLogger
+
+__all__ = ["Callback", "WandBLogger", "MLflowLogger"]

torchloop-0.3.0/src/torchloop/callbacks/base.py

@@ -0,0 +1,19 @@
+"""Base callback abstractions for torchloop training events."""
+
+from __future__ import annotations
+
+
+class Callback:
+    """Base callback with optional training lifecycle hooks.
+
+    Subclasses can override any hook they need.
+    """
+
+    def on_train_begin(self, logs: dict) -> None:
+        """Run before training starts."""
+
+    def on_epoch_end(self, epoch: int, logs: dict) -> None:
+        """Run at the end of each epoch."""
+
+    def on_train_end(self, logs: dict) -> None:
+        """Run after training finishes."""

torchloop-0.3.0/src/torchloop/callbacks/mlflow_logger.py

@@ -0,0 +1,72 @@
+"""MLflow callback integration for torchloop."""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import Optional
+
+from torchloop.callbacks.base import Callback
+
+
+class MLflowLogger(Callback):
+    """Log training metrics to MLflow.
+
+    Args:
+        experiment_name: MLflow experiment name.
+        tracking_uri: Optional tracking server URI.
+        run_name: Optional MLflow run name.
+    """
+
+    def __init__(
+        self,
+        experiment_name: str,
+        tracking_uri: Optional[str] = None,
+        run_name: Optional[str] = None,
+    ) -> None:
+        self.experiment_name = experiment_name
+        self.tracking_uri = tracking_uri
+        self.run_name = run_name
+
+    def on_train_begin(self, logs: dict) -> None:
+        """Initialize MLflow experiment and run."""
+        try:
+            mlflow = import_module("mlflow")
+        except ImportError as exc:
+            raise ImportError(
+                "mlflow is required for MLflowLogger. "
+                "Install with: pip install torchloop[logging]"
+            ) from exc
+
+        if self.tracking_uri:
+            mlflow.set_tracking_uri(self.tracking_uri)
+        mlflow.set_experiment(self.experiment_name)
+        mlflow.start_run(run_name=self.run_name)
+
+    def on_epoch_end(self, epoch: int, logs: dict) -> None:
+        """Log epoch metrics to MLflow."""
+        try:
+            mlflow = import_module("mlflow")
+        except ImportError as exc:
+            raise ImportError(
+                "mlflow is required for MLflowLogger. "
+                "Install with: pip install torchloop[logging]"
+            ) from exc
+
+        numeric_logs = {
+            key: value
+            for key, value in logs.items()
+            if isinstance(value, (int, float))
+        }
+        mlflow.log_metrics(numeric_logs, step=epoch)
+
+    def on_train_end(self, logs: dict) -> None:
+        """End the active MLflow run."""
+        try:
+            mlflow = import_module("mlflow")
+        except ImportError as exc:
+            raise ImportError(
+                "mlflow is required for MLflowLogger. "
+                "Install with: pip install torchloop[logging]"
+            ) from exc
+
+        mlflow.end_run()

torchloop-0.3.0/src/torchloop/callbacks/wandb_logger.py

@@ -0,0 +1,64 @@
+"""Weights & Biases callback integration for torchloop."""
+
+from __future__ import annotations
+
+from importlib import import_module
+from typing import Any, Optional
+
+from torchloop.callbacks.base import Callback
+
+
+class WandBLogger(Callback):
+    """Log training metrics to Weights & Biases.
+
+    Args:
+        project: Weights & Biases project name.
+        name: Optional run name.
+        config: Optional run configuration dictionary.
+    """
+
+    def __init__(
+        self,
+        project: str,
+        name: Optional[str] = None,
+        config: Optional[dict[str, Any]] = None,
+    ) -> None:
+        self.project = project
+        self.name = name
+        self.config = config or {}
+
+    def on_train_begin(self, logs: dict) -> None:
+        """Initialize a W&B run."""
+        try:
+            wandb = import_module("wandb")
+        except ImportError as exc:
+            raise ImportError(
+                "wandb is required for WandBLogger. "
+                "Install with: pip install torchloop[logging]"
+            ) from exc
+
+        wandb.init(project=self.project, name=self.name, config=self.config)
+
+    def on_epoch_end(self, epoch: int, logs: dict) -> None:
+        """Log epoch metrics to W&B."""
+        try:
+            wandb = import_module("wandb")
+        except ImportError as exc:
+            raise ImportError(
+                "wandb is required for WandBLogger. "
+                "Install with: pip install torchloop[logging]"
+            ) from exc
+
+        wandb.log(logs, step=epoch)
+
+    def on_train_end(self, logs: dict) -> None:
+        """Finish the active W&B run."""
+        try:
+            wandb = import_module("wandb")
+        except ImportError as exc:
+            raise ImportError(
+                "wandb is required for WandBLogger. "
+                "Install with: pip install torchloop[logging]"
+            ) from exc
+
+        wandb.finish()

{torchloop-0.2.3 → torchloop-0.3.0}/src/torchloop/edge/deploy.py

@@ -61,7 +61,10 @@ def deploy_to_edge(
     artifact_format = "onnx"
 
     if target in {"esp32", "android"}:
-        artifact_path = output if output.suffix == ".tflite" else output.with_suffix(".tflite")
+        artifact_path = (
+            output if output.suffix == ".tflite"
+            else output.with_suffix(".tflite")
+        )
         if target == "esp32" and quantize_type != "int8":
             warnings.warn(
                 "ESP32 is typically best with int8 quantization.",
@@ -79,7 +82,10 @@ def deploy_to_edge(
             _convert_to_tflite(onnx_path, artifact_path, quantize, quantize_type)
             artifact_format = "tflite"
     elif target == "ios":
-        artifact_path = output if output.suffix == ".mlpackage" else output.with_suffix(".mlpackage")
+        artifact_path = (
+            output if output.suffix == ".mlpackage"
+            else output.with_suffix(".mlpackage")
+        )
         _convert_to_coreml(model, input_shape, artifact_path)
         artifact_format = "coreml"
     elif target == "jetson":
@@ -109,7 +115,9 @@ def _export_to_onnx(
         onnx_name = "".join(["on", "nx"])
         onnx = import_module(onnx_name)
     except ImportError as exc:
-        raise ImportError("onnx is required. Install with: pip install torchloop[edge]") from exc
+        raise ImportError(
+            "onnx is required. Install with: pip install torchloop[edge]"
+        ) from exc
 
     model.eval()
     device = _get_model_device(model)
@@ -184,7 +192,9 @@ def _convert_to_coreml(
         coremltools_name = "".join(["coreml", "tools"])
         ct = import_module(coremltools_name)
     except ImportError as exc:
-        raise ImportError("coremltools is required. Install with: pip install torchloop[edge]") from exc
+        raise ImportError(
+            "coremltools is required. Install with: pip install torchloop[edge]"
+        ) from exc
 
     model_cpu = model.to("cpu").eval()
     traced = torch.jit.trace(model_cpu, torch.randn(*input_shape))

{torchloop-0.2.3 → torchloop-0.3.0}/src/torchloop/edge/estimate.py

@@ -63,19 +63,34 @@ def _estimate_activation_mb(input_shape: tuple[int, ...]) -> float:
     return (numel * 4 * 2) / (1024 * 1024)
 
 
-def _estimate_flops(model: torch.nn.Module, input_shape: tuple[int, ...]) -> int:
+def _estimate_flops(
+    model: torch.nn.Module,
+    input_shape: tuple[int, ...],
+) -> int:
     """Estimate FLOPs using hooks for common modules (Conv2d, Linear)."""
     flops = 0
     hooks = []
 
-    def conv_hook(module: torch.nn.Conv2d, _inp: tuple[torch.Tensor], out: torch.Tensor) -> None:
+    def conv_hook(
+        module: torch.nn.Conv2d,
+        _inp: tuple[torch.Tensor],
+        out: torch.Tensor,
+    ) -> None:
         nonlocal flops
         batch = out.shape[0]
         out_h, out_w = out.shape[2], out.shape[3]
-        kernel_ops = module.kernel_size[0] * module.kernel_size[1] * (module.in_channels / module.groups)
+        kernel_ops = (
+            module.kernel_size[0]
+            * module.kernel_size[1]
+            * (module.in_channels / module.groups)
+        )
         flops += int(batch * out_h * out_w * module.out_channels * kernel_ops * 2)
 
-    def linear_hook(module: torch.nn.Linear, inp: tuple[torch.Tensor], _out: torch.Tensor) -> None:
+    def linear_hook(
+        module: torch.nn.Linear,
+        inp: tuple[torch.Tensor],
+        _out: torch.Tensor,
+    ) -> None:
         nonlocal flops
         batch = inp[0].shape[0]
         flops += int(batch * module.in_features * module.out_features * 2)

{torchloop-0.2.3 → torchloop-0.3.0}/src/torchloop/trainer.py

@@ -32,7 +32,7 @@ import torch.nn as nn
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
-from torchloop.callbacks import Callback, StopTraining
+from torchloop.callbacks import Callback
 
 
 class Trainer:
@@ -129,10 +129,9 @@ class Trainer:
         Returns:
             history dict with train_loss, val_loss, val_metric, lr per epoch.
         """
-        self._trigger_hook("on_train_begin", logs=self.history)
+        self._run_callbacks("on_train_begin", dict(self.history))
         for epoch in range(1, epochs + 1):
             t0 = time.time()
-            self._trigger_hook("on_epoch_begin", epoch=epoch, logs=self.history)
             train_loss = self._train_epoch(train_loader)
             self.history["train_loss"].append(train_loss)
 
@@ -153,13 +152,13 @@ class Trainer:
             )
 
             epoch_logs = {
+                "epoch": epoch,
                 "train_loss": train_loss,
                 "val_loss": val_loss,
                 "val_metric": val_metric,
                 "lr": current_lr,
-                "model": self.model,
             }
-            self._trigger_hook("on_epoch_end", epoch=epoch, logs=epoch_logs)
+            self._run_callbacks("on_epoch_end", epoch_logs)
 
             if self._should_stop():
                 print(f"  Early stopping triggered at epoch {epoch}.")
@@ -169,7 +168,7 @@ class Trainer:
             self.model.load_state_dict(self._best_state)
             print("  Restored best model weights.")
 
-        self._trigger_hook("on_train_end", logs=self.history)
+        self._run_callbacks("on_train_end", dict(self.history))
         return self.history
 
     def add_callback(self, callback: Callback) -> None:
@@ -222,11 +221,6 @@ class Trainer:
                 pending_steps = 0
 
             total_loss += raw_loss.item() * inputs.size(0)
-            self._trigger_hook(
-                "on_batch_end",
-                batch=batch_idx,
-                logs={"loss": float(raw_loss.item())},
-            )
 
         if pending_steps > 0:
             self._optimizer_step()
@@ -296,15 +290,15 @@ class Trainer:
             )
         )
 
-    def _trigger_hook(self, hook_name: str, **kwargs: Any) -> None:
-        for callback in self.callbacks:
-            hook = getattr(callback, hook_name, None)
+    def _run_callbacks(self, event: str, logs: dict[str, Any]) -> None:
+        for callback in (self.callbacks or []):
+            hook = getattr(callback, event, None)
             if hook is None:
                 continue
-            try:
-                hook(**kwargs)
-            except StopTraining:
-                self._stop_early = True
+            if event == "on_epoch_end":
+                hook(int(logs.get("epoch", 0)), logs)
+            else:
+                hook(logs)
 
     @staticmethod
     def _log(

torchloop-0.3.0/tests/test_callbacks.py

@@ -0,0 +1,98 @@
+"""Tests for torchloop callback integrations."""
+
+from __future__ import annotations
+
+import sys
+
+import pytest
+import torch
+import torch.nn as nn
+from torch.utils.data import DataLoader, TensorDataset
+
+from torchloop import Trainer
+from torchloop.callbacks import Callback, MLflowLogger, WandBLogger
+
+
+def _make_loader(n: int = 32, features: int = 8, classes: int = 3) -> DataLoader:
+    x_data = torch.randn(n, features)
+    y_data = torch.randint(0, classes, (n,))
+    dataset = TensorDataset(x_data, y_data)
+    return DataLoader(dataset, batch_size=8)
+
+
+def _make_model(features: int = 8, classes: int = 3) -> nn.Module:
+    return nn.Sequential(nn.Linear(features, 16), nn.ReLU(), nn.Linear(16, classes))
+
+
+def test_callback_base_methods_are_noop() -> None:
+    """Ensure base callback methods are safe no-op defaults."""
+    callback = Callback()
+    callback.on_train_begin({"start": True})
+    callback.on_epoch_end(1, {"loss": 1.0})
+    callback.on_train_end({"done": True})
+
+
+def test_wandb_logger_raises_if_not_installed(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Raise ImportError when wandb is unavailable at runtime."""
+    monkeypatch.setitem(sys.modules, "wandb", None)
+    logger = WandBLogger(project="proj")
+
+    with pytest.raises(ImportError, match="wandb is required"):
+        logger.on_train_begin({"train": True})
+
+
+
+def test_mlflow_logger_raises_if_not_installed(monkeypatch: pytest.MonkeyPatch) -> None:
+    """Raise ImportError when mlflow is unavailable at runtime."""
+    monkeypatch.setitem(sys.modules, "mlflow", None)
+    logger = MLflowLogger(experiment_name="exp")
+
+    with pytest.raises(ImportError, match="mlflow is required"):
+        logger.on_train_begin({"train": True})
+
+
+
+def test_trainer_accepts_empty_callbacks_list() -> None:
+    """Allow explicit empty callback lists during fit execution."""
+    model = _make_model()
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+    criterion = nn.CrossEntropyLoss()
+
+    trainer = Trainer(
+        model=model,
+        optimizer=optimizer,
+        criterion=criterion,
+        device="cpu",
+        callbacks=[],
+    )
+
+    history = trainer.fit(_make_loader(), _make_loader(), epochs=2)
+    assert len(history["train_loss"]) == 2
+
+
+
+def test_trainer_calls_on_epoch_end() -> None:
+    """Invoke callback epoch-end hook once per epoch with epoch indices."""
+
+    class TestCallback(Callback):
+        def __init__(self) -> None:
+            self.epochs: list[int] = []
+
+        def on_epoch_end(self, epoch: int, logs: dict) -> None:
+            self.epochs.append(epoch)
+
+    callback = TestCallback()
+    model = _make_model()
+    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
+    criterion = nn.CrossEntropyLoss()
+
+    trainer = Trainer(
+        model=model,
+        optimizer=optimizer,
+        criterion=criterion,
+        device="cpu",
+        callbacks=[callback],
+    )
+    trainer.fit(_make_loader(), _make_loader(), epochs=3)
+
+    assert callback.epochs == [1, 2, 3]

{torchloop-0.2.3 → torchloop-0.3.0}/tests/test_trainer.py

@@ -4,7 +4,7 @@ import pytest
 from torch.utils.data import DataLoader, TensorDataset
 
 from torchloop import Trainer
-from torchloop.callbacks import Callback, EarlyStopping
+from torchloop.callbacks import Callback
 
 
 def _make_loader(n=64, features=16, classes=3, batch=16):
@@ -140,20 +140,12 @@ def test_trainer_callbacks_triggered():
     class Recorder(Callback):
         def __init__(self):
             self.train_begin = 0
-            self.epoch_begin = 0
-            self.batch_end = 0
             self.epoch_end = 0
             self.train_end = 0
 
         def on_train_begin(self, logs=None):
             self.train_begin += 1
 
-        def on_epoch_begin(self, epoch, logs=None):
-            self.epoch_begin += 1
-
-        def on_batch_end(self, batch, logs=None):
-            self.batch_end += 1
-
         def on_epoch_end(self, epoch, logs=None):
             self.epoch_end += 1
 
@@ -175,17 +167,22 @@ def test_trainer_callbacks_triggered():
     trainer.fit(_make_loader(n=32, batch=8), epochs=2)
 
     assert recorder.train_begin == 1
-    assert recorder.epoch_begin == 2
-    assert recorder.batch_end == 8
     assert recorder.epoch_end == 2
     assert recorder.train_end == 1
 
 
-def test_early_stopping_callback_stops_training():
+def test_callback_receives_epoch_end_for_each_epoch():
+    class EpochRecorder(Callback):
+        def __init__(self):
+            self.epochs = []
+
+        def on_epoch_end(self, epoch, logs=None):
+            self.epochs.append(epoch)
+
     model = _make_model()
     optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
     criterion = nn.CrossEntropyLoss()
-    callback = EarlyStopping(patience=1, monitor="val_loss", min_delta=0.1)
+    callback = EpochRecorder()
 
     trainer = Trainer(
         model,
@@ -194,8 +191,8 @@ def test_early_stopping_callback_stops_training():
         device="cpu",
         callbacks=[callback],
     )
-    history = trainer.fit(_make_loader(), _make_loader(), epochs=10)
-    assert len(history["train_loss"]) < 10
+    trainer.fit(_make_loader(), _make_loader(), epochs=3)
+    assert callback.epochs == [1, 2, 3]
 
 
 @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")

torchloop-0.2.3/src/torchloop/callbacks.py

@@ -1,216 +0,0 @@
-"""Callback utilities for training lifecycle hooks."""
-
-from __future__ import annotations
-
-import csv
-from pathlib import Path
-from typing import Any, Optional
-
-import torch
-
-
-class StopTraining(Exception):
-    """Raised by callbacks to request early stop of training."""
-
-
-class Callback:
-    """Base callback class.
-
-    Subclass this and override hook methods to customize training behavior.
-    """
-
-    def on_train_begin(self, logs: Optional[dict[str, Any]] = None) -> None:
-        """Called before the first epoch starts.
-
-        Args:
-            logs: Optional training state dictionary.
-        """
-
-    def on_train_end(self, logs: Optional[dict[str, Any]] = None) -> None:
-        """Called after training ends.
-
-        Args:
-            logs: Optional training state dictionary.
-        """
-
-    def on_epoch_begin(
-        self,
-        epoch: int,
-        logs: Optional[dict[str, Any]] = None,
-    ) -> None:
-        """Called at the start of each epoch.
-
-        Args:
-            epoch: Current epoch number.
-            logs: Optional training state dictionary.
-        """
-
-    def on_epoch_end(
-        self,
-        epoch: int,
-        logs: Optional[dict[str, Any]] = None,
-    ) -> None:
-        """Called at the end of each epoch.
-
-        Args:
-            epoch: Current epoch number.
-            logs: Optional training state dictionary.
-        """
-
-    def on_batch_end(
-        self,
-        batch: int,
-        logs: Optional[dict[str, Any]] = None,
-    ) -> None:
-        """Called after each training batch.
-
-        Args:
-            batch: Zero-based batch index.
-            logs: Optional batch-level metrics.
-        """
-
-
-class EarlyStopping(Callback):
-    """Stop training when a monitored metric stops improving.
-
-    Args:
-        patience: Number of epochs to wait for an improvement.
-        monitor: Metric name expected in logs.
-        min_delta: Minimum improvement required to reset patience.
-    """
-
-    def __init__(
-        self,
-        patience: int = 5,
-        monitor: str = "val_loss",
-        min_delta: float = 0.0,
-    ) -> None:
-        if patience < 1:
-            raise ValueError("patience must be >= 1")
-        if min_delta < 0:
-            raise ValueError("min_delta must be >= 0")
-
-        self.patience = patience
-        self.monitor = monitor
-        self.min_delta = min_delta
-        self.best = float("inf")
-        self.counter = 0
-
-    def on_epoch_end(
-        self,
-        epoch: int,
-        logs: Optional[dict[str, Any]] = None,
-    ) -> None:
-        """Check metric value and raise StopTraining when patience is exceeded.
-
-        Args:
-            epoch: Current epoch number.
-            logs: Metrics dictionary from trainer.
-        """
-        logs = logs or {}
-        current = logs.get(self.monitor)
-        if current is None:
-            return
-
-        if current <= self.best - self.min_delta:
-            self.best = float(current)
-            self.counter = 0
-            return
-
-        self.counter += 1
-        if self.counter >= self.patience:
-            raise StopTraining(f"Early stopping triggered at epoch {epoch}")
-
-
-class ModelCheckpoint(Callback):
-    """Save model checkpoints based on monitored metric.
-
-    Args:
-        monitor: Metric name expected in logs.
-        save_best_only: Save only when monitored metric improves.
-        filepath: Path to write checkpoint file.
-    """
-
-    def __init__(
-        self,
-        monitor: str = "val_loss",
-        save_best_only: bool = True,
-        filepath: str = "checkpoint.pt",
-    ) -> None:
-        self.monitor = monitor
-        self.save_best_only = save_best_only
-        self.filepath = Path(filepath)
-        self.best = float("inf")
-
-    def on_epoch_end(
-        self,
-        epoch: int,
-        logs: Optional[dict[str, Any]] = None,
-    ) -> None:
-        """Persist model checkpoint if criteria are met.
-
-        Args:
-            epoch: Current epoch number.
-            logs: Metrics dictionary from trainer.
-        """
-        logs = logs or {}
-        model = logs.get("model")
-        if model is None:
-            return
-
-        current = logs.get(self.monitor)
-        if current is None:
-            return
-
-        if not self.save_best_only:
-            torch.save(model.state_dict(), self.filepath)
-            return
-
-        if float(current) < self.best:
-            self.best = float(current)
-            torch.save(model.state_dict(), self.filepath)
-
-
-class CSVLogger(Callback):
-    """Log epoch metrics to a CSV file.
-
-    Args:
-        log_dir: Directory where CSV logs are written.
-        filename: Name of CSV file.
-    """
-
-    def __init__(self, log_dir: str = "./logs", filename: str = "metrics.csv") -> None:
-        self.log_dir = Path(log_dir)
-        self.filename = filename
-        self._initialized = False
-        self._fieldnames: list[str] = []
-
-    def on_epoch_end(
-        self,
-        epoch: int,
-        logs: Optional[dict[str, Any]] = None,
-    ) -> None:
-        """Append epoch metrics to CSV.
-
-        Args:
-            epoch: Current epoch number.
-            logs: Metrics dictionary from trainer.
-        """
-        logs = dict(logs or {})
-        logs["epoch"] = epoch
-
-        self.log_dir.mkdir(parents=True, exist_ok=True)
-        file_path = self.log_dir / self.filename
-
-        if not self._initialized:
-            self._fieldnames = sorted(logs.keys())
-            with file_path.open("w", newline="", encoding="utf-8") as f:
-                writer = csv.DictWriter(f, fieldnames=self._fieldnames)
-                writer.writeheader()
-                writer.writerow({k: logs.get(k) for k in self._fieldnames})
-            self._initialized = True
-            return
-
-        with file_path.open("a", newline="", encoding="utf-8") as f:
-            writer = csv.DictWriter(f, fieldnames=self._fieldnames)
-            writer.writerow({k: logs.get(k) for k in self._fieldnames})

torchloop-0.2.3/tests/test_callbacks.py

@@ -1,191 +0,0 @@
-"""Tests for callback utilities."""
-
-from __future__ import annotations
-
-import csv
-from pathlib import Path
-
-import pytest
-import torch
-import torch.nn as nn
-from torch.utils.data import DataLoader, TensorDataset
-
-from torchloop import Trainer
-from torchloop.callbacks import CSVLogger, Callback, EarlyStopping, ModelCheckpoint, StopTraining
-
-
-class TinyNet(nn.Module):
-    """Simple network for callback tests."""
-
-    def __init__(self) -> None:
-        super().__init__()
-        self.fc = nn.Linear(10, 2)
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        return self.fc(x)
-
-
-@pytest.fixture
-def loaders() -> tuple[DataLoader, DataLoader]:
-    """Create tiny train and validation loaders."""
-    x_train = torch.randn(12, 10)
-    y_train = torch.randint(0, 2, (12,))
-    x_val = torch.randn(8, 10)
-    y_val = torch.randint(0, 2, (8,))
-
-    train_loader = DataLoader(TensorDataset(x_train, y_train), batch_size=4)
-    val_loader = DataLoader(TensorDataset(x_val, y_val), batch_size=4)
-    return train_loader, val_loader
-
-
-@pytest.fixture
-def trainer_parts() -> tuple[nn.Module, torch.optim.Optimizer, nn.Module]:
-    """Create model, optimizer, and criterion for Trainer setup."""
-    model = TinyNet()
-    optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)
-    criterion = nn.CrossEntropyLoss()
-    return model, optimizer, criterion
-
-
-def test_early_stopping_init_values() -> None:
-    """Validate early stopping initialization fields."""
-    callback = EarlyStopping(patience=3, monitor="val_loss", min_delta=0.1)
-    assert callback.patience == 3
-    assert callback.monitor == "val_loss"
-    assert callback.min_delta == 0.1
-    assert callback.best == float("inf")
-    assert callback.counter == 0
-
-
-def test_early_stopping_validation_errors() -> None:
-    """Validate constructor error handling for invalid config."""
-    with pytest.raises(ValueError, match="patience must be >= 1"):
-        EarlyStopping(patience=0)
-
-    with pytest.raises(ValueError, match="min_delta must be >= 0"):
-        EarlyStopping(min_delta=-0.1)
-
-
-def test_early_stopping_triggers_after_patience() -> None:
-    """Raise StopTraining when monitored metric does not improve."""
-    callback = EarlyStopping(patience=2, monitor="val_loss")
-    callback.on_epoch_end(epoch=0, logs={"val_loss": 1.0})
-    callback.on_epoch_end(epoch=1, logs={"val_loss": 1.1})
-
-    with pytest.raises(StopTraining):
-        callback.on_epoch_end(epoch=2, logs={"val_loss": 1.2})
-
-
-def test_model_checkpoint_save_best_only(tmp_path: Path) -> None:
-    """Save checkpoint only when monitored metric improves."""
-    model = TinyNet()
-    path = tmp_path / "best.pt"
-    callback = ModelCheckpoint(filepath=str(path), save_best_only=True)
-
-    callback.on_epoch_end(epoch=0, logs={"model": model, "val_loss": 1.0})
-    first_mtime = path.stat().st_mtime
-
-    callback.on_epoch_end(epoch=1, logs={"model": model, "val_loss": 1.2})
-    assert path.stat().st_mtime == first_mtime
-
-    callback.on_epoch_end(epoch=2, logs={"model": model, "val_loss": 0.8})
-    assert path.stat().st_mtime >= first_mtime
-
-
-def test_model_checkpoint_save_all(tmp_path: Path) -> None:
-    """Always save checkpoint when save_best_only is disabled."""
-    model = TinyNet()
-    path = tmp_path / "all.pt"
-    callback = ModelCheckpoint(filepath=str(path), save_best_only=False)
-
-    callback.on_epoch_end(epoch=0, logs={"model": model, "val_loss": 2.0})
-    first_mtime = path.stat().st_mtime
-    callback.on_epoch_end(epoch=1, logs={"model": model, "val_loss": 3.0})
-    assert path.stat().st_mtime >= first_mtime
-
-
-def test_model_checkpoint_noop_when_logs_missing(tmp_path: Path) -> None:
-    """Do not crash when required keys are missing in logs."""
-    path = tmp_path / "noop.pt"
-    callback = ModelCheckpoint(filepath=str(path))
-
-    callback.on_epoch_end(epoch=0, logs={"val_loss": 1.0})
-    callback.on_epoch_end(epoch=1, logs={"model": TinyNet()})
-    assert not path.exists()
-
-
-def test_csv_logger_writes_header_and_rows(tmp_path: Path) -> None:
-    """Write CSV header once and append per-epoch values."""
-    callback = CSVLogger(log_dir=str(tmp_path), filename="metrics.csv")
-    callback.on_epoch_end(epoch=0, logs={"loss": 1.0, "val_loss": 1.1})
-    callback.on_epoch_end(epoch=1, logs={"loss": 0.9, "val_loss": 1.0})
-
-    csv_path = tmp_path / "metrics.csv"
-    with csv_path.open("r", encoding="utf-8") as f:
-        rows = list(csv.reader(f))
-
-    assert len(rows) == 3
-    assert "epoch" in rows[0]
-
-
-def test_trainer_calls_custom_callback_hooks(
-    trainer_parts: tuple[nn.Module, torch.optim.Optimizer, nn.Module],
-    loaders: tuple[DataLoader, DataLoader],
-) -> None:
-    """Ensure custom callback receives train, epoch, and batch hooks."""
-
-    class Recorder(Callback):
-        def __init__(self) -> None:
-            self.events: list[str] = []
-
-        def on_train_begin(self, logs=None) -> None:
-            self.events.append("train_begin")
-
-        def on_epoch_begin(self, epoch: int, logs=None) -> None:
-            self.events.append(f"epoch_begin_{epoch}")
-
-        def on_batch_end(self, batch: int, logs=None) -> None:
-            self.events.append("batch_end")
-
-        def on_epoch_end(self, epoch: int, logs=None) -> None:
-            self.events.append(f"epoch_end_{epoch}")
-
-        def on_train_end(self, logs=None) -> None:
-            self.events.append("train_end")
-
-    model, optimizer, criterion = trainer_parts
-    train_loader, val_loader = loaders
-    callback = Recorder()
-    trainer = Trainer(model, optimizer, criterion, callbacks=[callback], device="cpu")
-    trainer.fit(train_loader, val_loader, epochs=2)
-
-    assert "train_begin" in callback.events
-    assert "train_end" in callback.events
-    assert "epoch_begin_1" in callback.events
-    assert "epoch_end_2" in callback.events
-    assert "batch_end" in callback.events
-
-
-def test_trainer_handles_stop_training_callback(
-    trainer_parts: tuple[nn.Module, torch.optim.Optimizer, nn.Module],
-    loaders: tuple[DataLoader, DataLoader],
-) -> None:
-    """StopTraining raised by callback should halt training gracefully."""
-
-    class StopAfterFirstEpoch(Callback):
-        def on_epoch_end(self, epoch: int, logs=None) -> None:
-            if epoch >= 1:
-                raise StopTraining("stop")
-
-    model, optimizer, criterion = trainer_parts
-    train_loader, val_loader = loaders
-    trainer = Trainer(
-        model,
-        optimizer,
-        criterion,
-        callbacks=[StopAfterFirstEpoch()],
-        device="cpu",
-    )
-
-    history = trainer.fit(train_loader, val_loader, epochs=5)
-    assert len(history["train_loss"]) == 1