congrads 0.1.0__py3-none-any.whl → 0.3.0__py3-none-any.whl

congrads/__init__.py

@@ -1,21 +1,11 @@
-# __init__.py
+try:  # noqa: D104
+    from importlib.metadata import version as get_version  # Python 3.8+
+except ImportError:
+    from pkg_resources import (
+        get_distribution as get_version,
+    )  # Fallback for older versions
 
-# Only expose the submodules, not individual classes
-from . import core
-from . import constraints
-from . import datasets
-from . import descriptor
-from . import learners
-from . import metrics
-from . import networks
-
-# Define __all__ to specify that the submodules are accessible, but not classes directly.
-__all__ = [
-    "core",
-    "constraints",
-    "datasets",
-    "descriptor",
-    "learners",
-    "metrics",
-    "networks"
-]
+try:
+    version = get_version("congrads")  # Replace with your package name
+except Exception:
+    version = "0.0.0"  # Fallback if the package isn't installed

congrads/callbacks/base.py

@@ -0,0 +1,357 @@
+"""Callback and Operation Framework for Modular Training Pipelines.
+
+This module provides a structured system for defining and executing
+callbacks and operations at different stages of a training lifecycle.
+It is designed to support:
+
+- Stateless, reusable operations that produce outputs merged into
+  the event-local data.
+- Callbacks that group operations and/or custom logic for specific
+  stages of training, epochs, batches, and steps.
+- A central CallbackManager to orchestrate multiple callbacks,
+  maintain shared context, and execute stage-specific pipelines
+  in deterministic order.
+
+Stages supported:
+    - on_train_start
+    - on_train_end
+    - on_epoch_start
+    - on_epoch_end
+    - on_batch_start
+    - on_batch_end
+    - on_test_start
+    - on_test_end
+    - on_train_batch_start
+    - on_train_batch_end
+    - on_valid_batch_start
+    - on_valid_batch_end
+    - on_test_batch_start
+    - on_test_batch_end
+    - after_train_forward
+    - after_valid_forward
+    - after_test_forward
+
+Usage:
+    1. Define Operations by subclassing `Operation` and implementing
+       the `compute` method.
+    2. Create a Callback subclass or instance and register Operations
+       to stages via `add(stage, operation)`.
+    3. Register callbacks with `CallbackManager`.
+    4. Invoke `CallbackManager.run(stage, data)` at appropriate points
+       in the training loop, passing in event-local data.
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import Iterable
+from typing import Any, Literal, Self
+
+Stage = Literal[
+    "on_train_start",
+    "on_train_end",
+    "on_epoch_start",
+    "on_epoch_end",
+    "on_test_start",
+    "on_test_end",
+    "on_batch_start",
+    "on_batch_end",
+    "on_train_batch_start",
+    "on_train_batch_end",
+    "on_valid_batch_start",
+    "on_valid_batch_end",
+    "on_test_batch_start",
+    "on_test_batch_end",
+    "after_train_forward",
+    "after_valid_forward",
+    "after_test_forward",
+]
+
+STAGES: tuple[Stage, ...] = (
+    "on_train_start",
+    "on_train_end",
+    "on_epoch_start",
+    "on_epoch_end",
+    "on_test_start",
+    "on_test_end",
+    "on_batch_start",
+    "on_batch_end",
+    "on_train_batch_start",
+    "on_train_batch_end",
+    "on_valid_batch_start",
+    "on_valid_batch_end",
+    "on_test_batch_start",
+    "on_test_batch_end",
+    "after_train_forward",
+    "after_valid_forward",
+    "after_test_forward",
+)
+
+
+class Operation(ABC):
+    """Abstract base class representing a stateless unit of work executed inside a callback stage.
+
+    Subclasses should implement the `compute` method which returns
+    a dictionary of outputs to merge into the running event data.
+    """
+
+    def __repr__(self) -> str:
+        """Return a concise string representation of the operation."""
+        return f"<{self.__class__.__name__}>"
+
+    def __call__(self, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
+        """Execute the operation with the given event-local data and shared context.
+
+        Args:
+            data (dict[str, Any]): Event-local dictionary containing data for this stage.
+            ctx (dict[str, Any]): Shared context dictionary accessible by all operations and callbacks.
+
+        Returns:
+            dict[str, Any]: Outputs produced by the operation to merge into the running data.
+                            Returns an empty dict if `compute` returns None.
+        """
+        out = self.compute(data, ctx)
+        if out is None:
+            return {}
+        if not isinstance(out, dict):
+            raise TypeError(f"{self.__class__.__name__}.compute must return dict or None")
+        return out
+
+    @abstractmethod
+    def compute(self, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any] | None:
+        """Perform the operation's computation.
+
+        Args:
+            data (dict[str, Any]): Event-local dictionary containing the current data.
+            ctx (dict[str, Any]): Shared context dictionary.
+
+        Returns:
+            dict[str, Any] or None: Outputs to merge into the running data.
+                                     Returning None is equivalent to {}.
+        """
+        raise NotImplementedError
+
+
+class Callback(ABC):  # noqa: B024
+    """Abstract base class representing a callback that can have multiple operations registered to different stages of the training lifecycle.
+
+    Each stage method executes all operations registered for that stage
+    in insertion order. Operations can modify the event-local data dictionary.
+    """
+
+    def __init__(self):
+        """Initialize the callback with empty operation lists for all stages."""
+        self._ops_by_stage: dict[Stage, list[Operation]] = {s: [] for s in STAGES}
+
+    def __repr__(self) -> str:
+        """Return a concise string showing number of operations per stage."""
+        ops_summary = {stage: len(ops) for stage, ops in self._ops_by_stage.items() if ops}
+        return f"<{self.__class__.__name__} ops={ops_summary}>"
+
+    def add(self, stage: Stage, op: Operation) -> Self:
+        """Register an operation to execute at the given stage.
+
+        Args:
+            stage (Stage): Lifecycle stage at which to run the operation.
+            op (Operation): Operation instance to add.
+
+        Returns:
+            Self: Returns self for method chaining.
+        """
+        self._ops_by_stage[stage].append(op)
+        return self
+
+    def _run_ops(self, stage: Stage, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
+        """Execute all operations registered for a specific stage.
+
+        Args:
+            stage (Stage): Lifecycle stage to execute.
+            data (dict[str, Any]): Event-local data to pass to operations.
+            ctx (dict[str, Any]): Shared context across callbacks and operations.
+
+        Returns:
+            dict[str, Any]: Merged data dictionary after executing all operations.
+
+        Notes:
+            - Operations are executed in insertion order.
+            - If an operation overwrites existing keys, a warning is issued.
+        """
+        out = dict(data)
+
+        for operation in self._ops_by_stage[stage]:
+            try:
+                produced = operation(out, ctx) or {}
+            except Exception as e:
+                raise RuntimeError(f"Error in operation {operation} at stage {stage}") from e
+
+            collisions = set(produced.keys()) & set(out.keys())
+            if collisions:
+                import warnings
+
+                warnings.warn(
+                    f"Operation {operation} at stage '{stage}' is overwriting keys: {collisions}",
+                    stacklevel=2,
+                )
+
+            out.update(produced)
+
+        return out
+
+    # --- training ---
+    def on_train_start(self, data: dict[str, Any], ctx: dict[str, Any]):
+        """Execute operations registered for the 'on_train_start' stage."""
+        self._run_ops("on_train_start", data, ctx)
+
+    def on_train_end(self, data: dict[str, Any], ctx: dict[str, Any]):
+        """Execute operations registered for the 'on_train_end' stage."""
+        self._run_ops("on_train_end", data, ctx)
+
+    # --- epoch ---
+    def on_epoch_start(self, data: dict[str, Any], ctx: dict[str, Any]):
+        """Execute operations registered for the 'on_epoch_start' stage."""
+        self._run_ops("on_epoch_start", data, ctx)
+
+    def on_epoch_end(self, data: dict[str, Any], ctx: dict[str, Any]):
+        """Execute operations registered for the 'on_epoch_end' stage."""
+        self._run_ops("on_epoch_end", data, ctx)
+
+    # --- test ---
+    def on_test_start(self, data: dict[str, Any], ctx: dict[str, Any]):
+        """Execute operations registered for the 'on_test_start' stage."""
+        self._run_ops("on_test_start", data, ctx)
+
+    def on_test_end(self, data: dict[str, Any], ctx: dict[str, Any]):
+        """Execute operations registered for the 'on_test_end' stage."""
+        self._run_ops("on_test_end", data, ctx)
+
+    # --- batch ---
+    def on_batch_start(self, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
+        """Execute operations registered for the 'on_batch_start' stage."""
+        return self._run_ops("on_batch_start", data, ctx)
+
+    def on_batch_end(self, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
+        """Execute operations registered for the 'on_batch_end' stage."""
+        return self._run_ops("on_batch_end", data, ctx)
+
+    def on_train_batch_start(self, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
+        """Execute operations registered for the 'on_train_batch_start' stage."""
+        return self._run_ops("on_train_batch_start", data, ctx)
+
+    def on_train_batch_end(self, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
+        """Execute operations registered for the 'on_train_batch_end' stage."""
+        return self._run_ops("on_train_batch_end", data, ctx)
+
+    def on_valid_batch_start(self, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
+        """Execute operations registered for the 'on_valid_batch_start' stage."""
+        return self._run_ops("on_valid_batch_start", data, ctx)
+
+    def on_valid_batch_end(self, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
+        """Execute operations registered for the 'on_valid_batch_end' stage."""
+        return self._run_ops("on_valid_batch_end", data, ctx)
+
+    def on_test_batch_start(self, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
+        """Execute operations registered for the 'on_test_batch_start' stage."""
+        return self._run_ops("on_test_batch_start", data, ctx)
+
+    def on_test_batch_end(self, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
+        """Execute operations registered for the 'on_test_batch_end' stage."""
+        return self._run_ops("on_test_batch_end", data, ctx)
+
+    def after_train_forward(self, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
+        """Execute operations registered for the 'after_train_forward' stage."""
+        return self._run_ops("after_train_forward", data, ctx)
+
+    def after_valid_forward(self, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
+        """Execute operations registered for the 'after_valid_forward' stage."""
+        return self._run_ops("after_valid_forward", data, ctx)
+
+    def after_test_forward(self, data: dict[str, Any], ctx: dict[str, Any]) -> dict[str, Any]:
+        """Execute operations registered for the 'after_test_forward' stage."""
+        return self._run_ops("after_test_forward", data, ctx)
+
+
+class CallbackManager:
+    """Orchestrates multiple callbacks and executes them at specific lifecycle stages.
+
+    - Callbacks are executed in registration order.
+    - Event-local data flows through all callbacks.
+    - Shared context is available for cross-callback communication.
+    """
+
+    def __init__(self, callbacks: Iterable[Callback] | None = None):
+        """Initialize a CallbackManager instance.
+
+        Args:
+        callbacks (Iterable[Callback] | None): Optional initial callbacks to register.
+            If None, starts with an empty callback list.
+
+        Attributes:
+        _callbacks (list[Callback]): Internal list of registered callbacks.
+        ctx (dict[str, Any]): Shared context dictionary accessible to all callbacks
+            and operations for cross-event communication.
+        """
+        self._callbacks: list[Callback] = list(callbacks) if callbacks else []
+        self.ctx: dict[str, Any] = {}
+
+    def __repr__(self) -> str:
+        """Return a concise representation showing registered callbacks and ctx keys."""
+        names = [cb.__class__.__name__ for cb in self._callbacks]
+        return f"<CallbackManager callbacks={names} ctx_keys={list(self.ctx.keys())}>"
+
+    def add(self, callback: Callback) -> Self:
+        """Register a single callback.
+
+        Args:
+            callback (Callback): Callback instance to add.
+
+        Returns:
+            Self: Returns self for fluent chaining.
+        """
+        self._callbacks.append(callback)
+        return self
+
+    def extend(self, callbacks: Iterable[Callback]) -> None:
+        """Register multiple callbacks at once.
+
+        Args:
+            callbacks (Iterable[Callback]): Iterable of callbacks to add.
+        """
+        self._callbacks.extend(callbacks)
+
+    def run(self, stage: Stage, data: dict[str, Any]) -> dict[str, Any]:
+        """Execute all registered callbacks for a specific stage.
+
+        Args:
+            stage (Stage): Lifecycle stage to run (e.g., "on_batch_start").
+            data (dict[str, Any]): Event-local data dictionary to pass through callbacks.
+
+        Returns:
+            dict[str, Any]: The final merged data dictionary after executing all callbacks.
+
+        Raises:
+            ValueError: If a callback does not implement the requested stage.
+            RuntimeError: If any callback raises an exception during execution.
+        """
+        for cb in self._callbacks:
+            if not hasattr(cb, stage):
+                raise ValueError(
+                    f"Callback {cb.__class__.__name__} has no handler for stage {stage}"
+                )
+            handler = getattr(cb, stage)
+
+            try:
+                new_data = handler(data, self.ctx)
+                if new_data is not None:
+                    data = new_data
+
+            except Exception as e:
+                raise RuntimeError(f"Error in callback {cb.__class__.__name__}.{stage}") from e
+
+        return data
+
+    @property
+    def callbacks(self) -> tuple[Callback, ...]:
+        """Return a read-only tuple of registered callbacks.
+
+        Returns:
+            tuple[Callback, ...]: Registered callbacks.
+        """
+        return tuple(self._callbacks)

congrads/callbacks/registry.py

@@ -0,0 +1,106 @@
+"""Holds all callback implementations for use in the training workflow.
+
+This module acts as a central registry for defining and storing different
+callback classes, such as logging, checkpointing, or custom behaviors
+triggered during training, validation, or testing. It is intended to
+collect all callback implementations in one place for easy reference
+and import, and can be extended as new callbacks are added.
+"""
+
+from torch.utils.tensorboard import SummaryWriter
+
+from ..callbacks.base import Callback
+from ..metrics import MetricManager
+from ..utils.utility import CSVLogger
+
+
+class LoggerCallback(Callback):
+    """Callback to log metrics to TensorBoard and CSV after each epoch or test.
+
+    This callback queries a MetricManager for aggregated metrics, writes them
+    to TensorBoard using SummaryWriter, and logs them to a CSV file via CSVLogger.
+    It also flushes loggers and resets metrics after logging.
+
+    Methods implemented:
+    - on_epoch_end: Logs metrics at the end of a training epoch.
+    - on_test_end: Logs metrics at the end of testing.
+    """
+
+    def __init__(
+        self,
+        metric_manager: MetricManager,
+        tensorboard_logger: SummaryWriter,
+        csv_logger: CSVLogger,
+    ):
+        """Initialize the LoggerCallback.
+
+        Args:
+            metric_manager: Instance of MetricManager used to collect metrics.
+            tensorboard_logger: TensorBoard SummaryWriter instance for logging scalars.
+            csv_logger: CSVLogger instance for logging metrics to CSV files.
+        """
+        super().__init__()
+        self.metric_manager = metric_manager
+        self.tensorboard_logger = tensorboard_logger
+        self.csv_logger = csv_logger
+
+    def on_epoch_end(self, data: dict[str, any], ctx: dict[str, any]):
+        """Log training metrics at the end of an epoch.
+
+        Aggregates metrics from the MetricManager under the 'during_training' category,
+        writes them to TensorBoard and CSV, flushes the loggers, and resets the metrics
+        for the next epoch.
+
+        Args:
+            data: Dictionary containing batch/epoch context (must include 'epoch').
+            ctx: Additional context dictionary (unused in this implementation).
+
+        Returns:
+            data: The same input dictionary, unmodified.
+        """
+        epoch = data["epoch"]
+
+        # Log training metrics
+        metrics = self.metric_manager.aggregate("during_training")
+        for name, value in metrics.items():
+            self.tensorboard_logger.add_scalar(name, value.item(), epoch)
+            self.csv_logger.add_value(name, value.item(), epoch)
+
+        # Flush/save
+        self.tensorboard_logger.flush()
+        self.csv_logger.save()
+
+        # Reset metric manager for training
+        self.metric_manager.reset("during_training")
+
+        return data
+
+    def on_test_end(self, data: dict[str, any], ctx: dict[str, any]):
+        """Log test metrics at the end of testing.
+
+        Aggregates metrics from the MetricManager under the 'after_training' category,
+        writes them to TensorBoard and CSV, flushes the loggers, and resets the metrics.
+
+        Args:
+            data: Dictionary containing test context (must include 'epoch').
+            ctx: Additional context dictionary (unused in this implementation).
+
+        Returns:
+            data: The same input dictionary, unmodified.
+        """
+        epoch = data["epoch"]
+
+        # Log test metrics
+        metrics = self.metric_manager.aggregate("after_training")
+        for name, value in metrics.items():
+            self.tensorboard_logger.add_scalar(name, value.item(), epoch)
+            self.csv_logger.add_value(name, value.item(), epoch)
+
+        # Flush/save
+        self.tensorboard_logger.flush()
+        self.csv_logger.save()
+
+        # Reset metric manager for test
+        self.metric_manager.reset("after_training")
+
+        return data

congrads/checkpoints.py

@@ -0,0 +1,178 @@
+"""Module for managing PyTorch model checkpoints.
+
+Provides the `CheckpointManager` class to save and load model and optimizer
+states during training, track the best metric values, and optionally report
+checkpoint events.
+"""
+
+import os
+from collections.abc import Callable
+from pathlib import Path
+
+from torch import Tensor, load, save
+from torch.nn import Module
+from torch.optim import Optimizer
+
+from .metrics import MetricManager
+from .utils.validation import validate_callable, validate_type
+
+
+class CheckpointManager:
+    """Manage saving and loading checkpoints for PyTorch models and optimizers.
+
+    Handles checkpointing based on a criteria function, restores metric
+    states, and optionally reports when a checkpoint is saved.
+    """
+
+    def __init__(
+        self,
+        criteria_function: Callable[[dict[str, Tensor], dict[str, Tensor]], bool],
+        network: Module,
+        optimizer: Optimizer,
+        metric_manager: MetricManager,
+        save_dir: str = "checkpoints",
+        create_dir: bool = False,
+        report_save: bool = False,
+    ):
+        """Initialize the CheckpointManager.
+
+        Args:
+            criteria_function (Callable[[dict[str, Tensor], dict[str, Tensor]], bool]):
+                Function that determines if the current checkpoint should be
+                saved based on the current and best metric values.
+            network (torch.nn.Module): The model to save/load.
+            optimizer (torch.optim.Optimizer): The optimizer to save/load.
+            metric_manager (MetricManager): Manages metric states for checkpointing.
+            save_dir (str, optional): Directory to save checkpoints. Defaults to 'checkpoints'.
+            create_dir (bool, optional): Whether to create `save_dir` if it does not exist.
+                Defaults to False.
+            report_save (bool, optional): Whether to report when a checkpoint is saved.
+                Defaults to False.
+
+        Raises:
+            TypeError: If any provided attribute has an incompatible type.
+            FileNotFoundError: If `save_dir` does not exist and `create_dir` is False.
+        """
+        # Type checking
+        validate_callable("criteria_function", criteria_function)
+        validate_type("network", network, Module)
+        validate_type("optimizer", optimizer, Optimizer)
+        validate_type("metric_manager", metric_manager, MetricManager)
+        validate_type("create_dir", create_dir, bool)
+        validate_type("report_save", report_save, bool)
+
+        # Create path or raise error if create_dir is not found
+        if not os.path.exists(save_dir):
+            if not create_dir:
+                raise FileNotFoundError(
+                    f"Save directory '{save_dir}' configured in checkpoint manager is not found."
+                )
+            Path(save_dir).mkdir(parents=True, exist_ok=True)
+
+        # Initialize objects variables
+        self.criteria_function = criteria_function
+        self.network = network
+        self.optimizer = optimizer
+        self.metric_manager = metric_manager
+        self.save_dir = save_dir
+        self.report_save = report_save
+
+        self.best_metric_values: dict[str, Tensor] = {}
+
+    def evaluate_criteria(self, epoch: int, metric_group: str = "during_training"):
+        """Evaluate the criteria function to determine if a better model is found.
+
+        Aggregates the current metric values during training and applies the
+        criteria function. If the criteria function indicates improvement, the
+        best metric values are updated, a checkpoint is saved, and a message is
+        optionally printed.
+
+        Args:
+            epoch (int): The current epoch number.
+            metric_group (str, optional): The metric group to evaluate. Defaults to 'during_training'.
+        """
+        current_metric_values = self.metric_manager.aggregate(metric_group)
+        if self.criteria_function is not None and self.criteria_function(
+            current_metric_values, self.best_metric_values
+        ):
+            # Print message if a new checkpoint is saved
+            if self.report_save:
+                print(f"New checkpoint saved at epoch {epoch}.")
+
+            # Update current best metric values
+            for metric_name, metric_value in current_metric_values.items():
+                self.best_metric_values[metric_name] = metric_value
+
+            # Save the current state
+            self.save(epoch)
+
+    def resume(self, filename: str = "checkpoint.pth", ignore_missing: bool = False) -> int:
+        """Resumes training from a saved checkpoint file.
+
+        Args:
+            filename (str): The name of the checkpoint file to load.
+                Defaults to "checkpoint.pth".
+            ignore_missing (bool): If True, does not raise an error if the
+                checkpoint file is missing and continues without loading,
+                starting from epoch 0. Defaults to False.
+
+        Returns:
+            int: The epoch number from the loaded checkpoint, or 0 if
+                ignore_missing is True and no checkpoint was found.
+
+        Raises:
+            TypeError: If a provided attribute has an incompatible type.
+            FileNotFoundError: If the specified checkpoint file does not exist.
+        """
+        # Type checking
+        validate_type("filename", filename, str)
+        validate_type("ignore_missing", ignore_missing, bool)
+
+        # Return starting epoch, either from checkpoint file or default
+        filepath = os.path.join(self.save_dir, filename)
+        if os.path.exists(filepath):
+            checkpoint = self.load(filename)
+            return checkpoint["epoch"]
+        elif ignore_missing:
+            return 0
+        else:
+            raise FileNotFoundError(f"A checkpoint was not found at {filepath} to resume training.")
+
+    def save(self, epoch: int, filename: str = "checkpoint.pth"):
+        """Save a checkpoint.
+
+        Args:
+            epoch (int): Current epoch number.
+            filename (str): Name of the checkpoint file. Defaults to
+                'checkpoint.pth'.
+        """
+        state = {
+            "epoch": epoch,
+            "network_state": self.network.state_dict(),
+            "optimizer_state": self.optimizer.state_dict(),
+            "best_metrics": self.best_metric_values,
+        }
+        filepath = os.path.join(self.save_dir, filename)
+        save(state, filepath)
+
+    def load(self, filename: str):
+        """Load a checkpoint and restore the training state.
+
+        Loads the checkpoint from the specified file and restores the network
+        weights, optimizer state, and best metric values.
+
+        Args:
+            filename (str): Name of the checkpoint file.
+
+        Returns:
+            dict: A dictionary containing the loaded checkpoint information,
+                including epoch, loss, and other relevant training state.
+        """
+        filepath = os.path.join(self.save_dir, filename)
+
+        checkpoint = load(filepath, weights_only=True)
+        self.network.load_state_dict(checkpoint["network_state"])
+        self.optimizer.load_state_dict(checkpoint["optimizer_state"])
+        self.best_metric_values = checkpoint["best_metrics"]
+
+        return checkpoint