congrads 1.1.2__py3-none-any.whl → 1.2.0__py3-none-any.whl

congrads/core/congradscore.py

@@ -0,0 +1,271 @@
+"""This module provides the core CongradsCore class for the main training functionality.
+
+It is designed to integrate constraint-guided optimization into neural network training.
+It extends traditional training processes by enforcing specific constraints
+on the model's outputs, ensuring that the network satisfies domain-specific
+requirements during both training and evaluation.
+
+The `CongradsCore` class serves as the central engine for managing the
+training, validation, and testing phases of a neural network model,
+incorporating constraints that influence the loss function and model updates.
+The model is trained with standard loss functions while also incorporating
+constraint-based adjustments, which are tracked and logged
+throughout the process.
+
+Key features:
+- Support for various constraints that can influence the training process.
+- Integration with PyTorch's `DataLoader` for efficient batch processing.
+- Metric management for tracking loss and constraint satisfaction.
+- Checkpoint management for saving and evaluating model states.
+
+"""
+
+from collections.abc import Callable
+
+import torch
+from torch import Tensor, sum
+from torch.nn import Module
+from torch.nn.modules.loss import _Loss
+from torch.optim import Optimizer
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from congrads.utils.utility import LossWrapper
+
+from ..callbacks.base import CallbackManager
+from ..checkpoints import CheckpointManager
+from ..constraints.base import Constraint
+from ..core.batch_runner import BatchRunner
+from ..core.constraint_engine import ConstraintEngine
+from ..core.epoch_runner import EpochRunner
+from ..descriptor import Descriptor
+from ..metrics import MetricManager
+
+
+class CongradsCore:
+    """The CongradsCore class is the central training engine for constraint-guided optimization.
+
+    It integrates standard neural network training
+    with additional constraint-driven adjustments to the loss function, ensuring
+    that the network satisfies domain-specific constraints during training.
+    """
+
+    def __init__(
+        self,
+        descriptor: Descriptor,
+        constraints: list[Constraint],
+        network: Module,
+        criterion: _Loss,
+        optimizer: Optimizer,
+        device: torch.device,
+        dataloader_train: DataLoader,
+        dataloader_valid: DataLoader | None = None,
+        dataloader_test: DataLoader | None = None,
+        metric_manager: MetricManager | None = None,
+        callback_manager: CallbackManager | None = None,
+        checkpoint_manager: CheckpointManager | None = None,
+        network_uses_grad: bool = False,
+        epsilon: float = 1e-6,
+        constraint_aggregator: Callable[..., Tensor] = sum,
+        enforce_all: bool = True,
+        disable_progress_bar_epoch: bool = False,
+        disable_progress_bar_batch: bool = False,
+        epoch_runner_cls: type["EpochRunner"] | None = None,
+        batch_runner_cls: type["BatchRunner"] | None = None,
+        constraint_engine_cls: type["ConstraintEngine"] | None = None,
+    ):
+        """Initialize the CongradsCore object.
+
+        Args:
+            descriptor (Descriptor): Describes variable layers in the network.
+            constraints (list[Constraint]): List of constraints to guide training.
+            network (Module): The neural network model to train.
+            criterion (callable): The loss function used for
+                training and validation.
+            optimizer (Optimizer): The optimizer used for updating model parameters.
+            device (torch.device): The device (e.g., CPU or GPU) for computations.
+            dataloader_train (DataLoader): DataLoader for training data.
+            dataloader_valid (DataLoader, optional): DataLoader for validation data.
+                If not provided, validation is skipped.
+            dataloader_test (DataLoader, optional): DataLoader for test data.
+                If not provided, testing is skipped.
+            metric_manager (MetricManager, optional): Manages metric tracking and recording.
+            callback_manager (CallbackManager, optional): Manages training callbacks.
+            checkpoint_manager (CheckpointManager, optional): Manages
+                    checkpointing. If not set, no checkpointing is done.
+            network_uses_grad (bool, optional): A flag indicating if the network
+                contains gradient calculation computations. Default is False.
+            epsilon (float, optional): A small value to avoid division by zero
+                in gradient calculations. Default is 1e-10.
+            constraint_aggregator (Callable[..., Tensor], optional): A function
+                to aggregate the constraint rescale loss. Default is `sum`.
+            enforce_all (bool, optional): If set to False, constraints will only be monitored and
+                not influence the training process. Overrides constraint-specific `enforce` parameters.
+                Defaults to True.
+            disable_progress_bar_epoch (bool, optional): If set to True, the epoch
+                progress bar will not show. Defaults to False.
+            disable_progress_bar_batch (bool, optional): If set to True, the batch
+                progress bar will not show. Defaults to False.
+            epoch_runner_cls (type[EpochRunner], optional): Custom EpochRunner class.
+                If not provided, the default EpochRunner is used.
+            batch_runner_cls (type[BatchRunner], optional): Custom BatchRunner class.
+                If not provided, the default BatchRunner is used.
+            constraint_engine_cls (type[ConstraintEngine], optional): Custom ConstraintEngine class.
+                If not provided, the default ConstraintEngine is used.
+
+        Note:
+            A warning is logged if the descriptor has no variable layers,
+            as at least one variable layer is required for the constraint logic
+            to influence the training process.
+        """
+        # Init object variables
+        self.device = device
+        self.network = network.to(device)
+        self.criterion = LossWrapper(criterion)
+        self.optimizer = optimizer
+        self.descriptor = descriptor
+
+        self.constraints = constraints or []
+        self.epsilon = epsilon
+        self.aggregator = constraint_aggregator
+        self.enforce_all = enforce_all
+
+        self.metric_manager = metric_manager
+        self.callback_manager = callback_manager
+        self.checkpoint_manager = checkpoint_manager
+
+        self.disable_progress_bar_epoch = disable_progress_bar_epoch
+        self.disable_progress_bar_batch = disable_progress_bar_batch
+
+        # Initialize constraint engine
+        self.constraint_engine = (constraint_engine_cls or ConstraintEngine)(
+            constraints=self.constraints,
+            descriptor=self.descriptor,
+            device=self.device,
+            epsilon=self.epsilon,
+            aggregator=self.aggregator,
+            enforce_all=self.enforce_all,
+            metric_manager=self.metric_manager,
+        )
+
+        # Initialize runners
+        self.batch_runner = (batch_runner_cls or BatchRunner)(
+            network=self.network,
+            criterion=self.criterion,
+            optimizer=self.optimizer,
+            constraint_engine=self.constraint_engine,
+            metric_manager=self.metric_manager,
+            callback_manager=self.callback_manager,
+            device=self.device,
+        )
+
+        self.epoch_runner = (epoch_runner_cls or EpochRunner)(
+            batch_runner=self.batch_runner,
+            network=self.network,
+            train_loader=dataloader_train,
+            valid_loader=dataloader_valid,
+            test_loader=dataloader_test,
+            network_uses_grad=network_uses_grad,
+            disable_progress_bar=self.disable_progress_bar_batch,
+        )
+
+        # Initialize constraint metrics
+        if self.metric_manager is not None:
+            self._initialize_metrics()
+
+    def _initialize_metrics(self) -> None:
+        """Register metrics for loss, constraint satisfaction ratio (CSR), and constraints.
+
+        This method registers the following metrics:
+
+        - Loss/train: Training loss.
+        - Loss/valid: Validation loss.
+        - Loss/test: Test loss after training.
+        - CSR/train: Constraint satisfaction ratio during training.
+        - CSR/valid: Constraint satisfaction ratio during validation.
+        - CSR/test: Constraint satisfaction ratio after training.
+        - One metric per constraint, for both training and validation.
+
+        """
+        self.metric_manager.register("Loss/train", "during_training")
+        self.metric_manager.register("Loss/valid", "during_training")
+        self.metric_manager.register("Loss/test", "after_training")
+
+        if len(self.constraints) > 0:
+            self.metric_manager.register("CSR/train", "during_training")
+            self.metric_manager.register("CSR/valid", "during_training")
+            self.metric_manager.register("CSR/test", "after_training")
+
+        for constraint in self.constraints:
+            self.metric_manager.register(f"{constraint.name}/train", "during_training")
+            self.metric_manager.register(f"{constraint.name}/valid", "during_training")
+            self.metric_manager.register(f"{constraint.name}/test", "after_training")
+
+    def fit(
+        self,
+        *,
+        start_epoch: int = 0,
+        max_epochs: int = 100,
+        test_model: bool = True,
+        final_checkpoint_name: str = "checkpoint_final.pth",
+    ) -> None:
+        """Run the full training loop, including optional validation, testing, and checkpointing.
+
+        This method performs training over multiple epochs with the following steps:
+        1. Trigger "on_train_start" callbacks if a callback manager is present.
+        2. For each epoch:
+        - Trigger "on_epoch_start" callbacks.
+        - Run a training epoch via the EpochRunner.
+        - Run a validation epoch via the EpochRunner.
+        - Evaluate checkpoint criteria if a checkpoint manager is present.
+        - Trigger "on_epoch_end" callbacks.
+        3. Trigger "on_train_end" callbacks after all epochs.
+        4. Optionally run a test epoch via the EpochRunner if `test_model` is True,
+        with corresponding "on_test_start" and "on_test_end" callbacks.
+        5. Save a final checkpoint using the checkpoint manager.
+
+        Args:
+            start_epoch: Index of the starting epoch (default 0). Useful for resuming training.
+            max_epochs: Maximum number of epochs to run (default 100).
+            test_model: Whether to run a test epoch after training (default True).
+            final_checkpoint_name: Filename for the final checkpoint saved at the end of training
+                                (default "checkpoint_final.pth").
+
+        Returns:
+            None
+        """
+        if self.callback_manager:
+            self.callback_manager.run("on_train_start", {"epoch": start_epoch})
+
+        for epoch in tqdm(
+            range(start_epoch, max_epochs),
+            initial=start_epoch,
+            desc="Epoch",
+            disable=self.disable_progress_bar_epoch,
+        ):
+            if self.callback_manager:
+                self.callback_manager.run("on_epoch_start", {"epoch": epoch})
+
+            self.epoch_runner.train()
+            self.epoch_runner.validate()
+
+            if self.checkpoint_manager:
+                self.checkpoint_manager.evaluate_criteria(epoch)
+
+            if self.callback_manager:
+                self.callback_manager.run("on_epoch_end", {"epoch": epoch})
+
+        if self.callback_manager:
+            self.callback_manager.run("on_train_end", {"epoch": epoch})
+
+        if test_model:
+            if self.callback_manager:
+                self.callback_manager.run("on_test_start", {"epoch": epoch})
+
+            self.epoch_runner.test()
+
+            if self.callback_manager:
+                self.callback_manager.run("on_test_end", {"epoch": epoch})
+
+        if self.checkpoint_manager:
+            self.checkpoint_manager.save(epoch, final_checkpoint_name)

congrads/core/constraint_engine.py

@@ -0,0 +1,170 @@
+"""Manages the evaluation and optional enforcement of constraints on neural network outputs.
+
+Responsibilities:
+- Compute and log Constraint Satisfaction Rate (CSR) for training, validation, and test batches.
+- Optionally adjust loss during training based on constraint directions and rescale factors.
+- Handle gradient computation and CGGD application.
+"""
+
+import torch
+from torch import Tensor, no_grad
+from torch.linalg import vector_norm
+
+from ..constraints.base import Constraint
+from ..descriptor import Descriptor
+from ..metrics import MetricManager
+
+
+class ConstraintEngine:
+    """Manages constraint evaluation and enforcement for a neural network.
+
+    The ConstraintEngine coordinates constraints defined in Constraint objects,
+    computes gradients for layers that affect the loss, logs metrics, and optionally
+    modifies the loss during training according to the constraints. It supports
+    separate phases for training, validation, and testing.
+    """
+
+    def __init__(
+        self,
+        *,
+        constraints: list[Constraint],
+        descriptor: Descriptor,
+        metric_manager: MetricManager,
+        device: torch.device,
+        epsilon: float,
+        aggregator: callable,
+        enforce_all: bool,
+    ) -> None:
+        """Initialize the ConstraintEngine.
+
+        Args:
+            constraints: List of Constraint objects to evaluate and optionally enforce.
+            descriptor: Descriptor containing metadata about network layers and which
+                        variables affect the loss.
+            metric_manager: MetricManager instance for logging CSR metrics.
+            device: Torch device where tensors will be allocated (CPU or GPU).
+            epsilon: Small positive value to avoid division by zero in gradient norms.
+            aggregator: Callable used to reduce per-layer constraint contributions
+                        to a scalar loss adjustment.
+            enforce_all: Whether to enforce all constraints during training.
+        """
+        self.constraints = constraints
+        self.descriptor = descriptor
+        self.metric_manager = metric_manager
+        self.device = device
+        self.epsilon = epsilon
+        self.enforce_all = enforce_all
+        self.aggregator = aggregator
+
+    def train(self, data: dict[str, Tensor], loss: Tensor) -> Tensor:
+        """Apply all active constraints during training.
+
+        Computes the original loss gradients for layers that affect the loss,
+        evaluates each constraint, logs the Constraint Satisfaction Rate (CSR),
+        and adjusts the loss according to constraint satisfaction.
+
+        Args:
+            data: Dictionary containing input and prediction tensors for the batch.
+            loss: The original loss tensor computed from the network output.
+
+        Returns:
+            Tensor: The loss tensor after applying constraint-based adjustments.
+        """
+        return self._apply_constraints(data, loss, phase="train", enforce=True)
+
+    def validate(self, data: dict[str, Tensor], loss: Tensor) -> Tensor:
+        """Evaluate constraints during validation without modifying the loss.
+
+        Computes and logs the Constraint Satisfaction Rate (CSR) for each constraint,
+        but does not apply rescale adjustments to the loss.
+
+        Args:
+            data: Dictionary containing input and prediction tensors for the batch.
+            loss: The original loss tensor computed from the network output.
+
+        Returns:
+            Tensor: The original loss tensor, unchanged.
+        """
+        return self._apply_constraints(data, loss, phase="valid", enforce=False)
+
+    def test(self, data: dict[str, Tensor], loss: Tensor) -> Tensor:
+        """Evaluate constraints during testing without modifying the loss.
+
+        Computes and logs the Constraint Satisfaction Rate (CSR) for each constraint,
+        but does not apply rescale adjustments to the loss.
+
+        Args:
+            data: Dictionary containing input and prediction tensors for the batch.
+            loss: The original loss tensor computed from the network output.
+
+        Returns:
+            Tensor: The original loss tensor, unchanged.
+        """
+        return self._apply_constraints(data, loss, phase="test", enforce=False)
+
+    def _apply_constraints(
+        self, data: dict[str, Tensor], loss: Tensor, phase: str, enforce: bool
+    ) -> Tensor:
+        """Evaluate constraints, log CSR, and optionally adjust the loss.
+
+        During training, computes loss gradients for variable layers that affect the loss.
+        Iterates over all constraints, logging the Constraint Satisfaction Rate (CSR)
+        and, if enforcement is enabled, adjusts the loss using constraint-specific
+        directions and rescale factors.
+
+        Args:
+            data: Dictionary containing input and prediction tensors for the batch.
+            loss: Original loss tensor computed from the network output.
+            phase: Current phase, one of "train", "valid", or "test".
+            enforce: If True, constraint-based adjustments are applied to the loss.
+
+        Returns:
+            Tensor: The combined loss after applying constraints (or the original loss
+            if enforce is False or not in training phase).
+        """
+        total_rescale_loss = torch.tensor(0.0, device=self.device, dtype=loss.dtype)
+
+        # Precompute gradients for variable layers affecting the loss
+        if phase == "train":
+            norm_loss_grad = {}
+            variable_keys = self.descriptor.variable_keys & self.descriptor.affects_loss_keys
+
+            for key in variable_keys:
+                grad = torch.autograd.grad(
+                    outputs=loss, inputs=data[key], retain_graph=True, allow_unused=True
+                )[0]
+
+                if grad is None:
+                    raise RuntimeError(
+                        f"Unable to compute loss gradients for layer '{key}'. "
+                        "Set has_loss=False in Descriptor if this layer does not affect loss."
+                    )
+
+                grad_flat = grad.view(grad.shape[0], -1)
+                norm_loss_grad[key] = (
+                    vector_norm(grad_flat, dim=1, ord=2, keepdim=True)
+                    .clamp(min=self.epsilon)
+                    .detach()
+                )
+
+        # Iterate constraints
+        for constraint in self.constraints:
+            checks, mask = constraint.check_constraint(data)
+            directions = constraint.calculate_direction(data)
+
+            # Log CSR
+            csr = (torch.sum(checks * mask) / torch.sum(mask)).unsqueeze(0)
+            self.metric_manager.accumulate(f"{constraint.name}/{phase}", csr)
+            self.metric_manager.accumulate(f"CSR/{phase}", csr)
+
+            # Skip adjustment if not enforcing
+            if not enforce or not constraint.enforce or not self.enforce_all or phase != "train":
+                continue
+
+            # Compute constraint-based rescale loss
+            for key in constraint.layers & self.descriptor.variable_keys:
+                with no_grad():
+                    rescale = (1 - checks) * directions[key] * constraint.rescale_factor
+                total_rescale_loss += self.aggregator(data[key] * rescale * norm_loss_grad[key])
+
+        return loss + total_rescale_loss

congrads/core/epoch_runner.py

@@ -0,0 +1,119 @@
+"""Defines the EpochRunner class for running full training, validation, and test epochs.
+
+This module handles:
+- Switching the network between training and evaluation modes
+- Iterating over DataLoaders with optional progress bars
+- Delegating per-batch processing to a BatchRunner instance
+- Optional gradient tracking control for evaluation phases
+- Warnings when validation or test loaders are not provided
+"""
+
+import warnings
+
+from torch import set_grad_enabled
+from torch.nn import Module
+from torch.utils.data import DataLoader
+from tqdm import tqdm
+
+from ..core.batch_runner import BatchRunner
+
+
+class EpochRunner:
+    """Runs full epochs over DataLoaders.
+
+    Responsibilities:
+    - Model mode switching
+    - Iteration over DataLoader
+    - Delegation to BatchRunner
+    - Progress bars
+    """
+
+    def __init__(
+        self,
+        network: Module,
+        batch_runner: BatchRunner,
+        train_loader: DataLoader,
+        valid_loader: DataLoader | None = None,
+        test_loader: DataLoader | None = None,
+        *,
+        network_uses_grad: bool = False,
+        disable_progress_bar: bool = False,
+    ):
+        """Initialize the EpochRunner.
+
+        Args:
+            network: The neural network module to train/validate/test.
+            batch_runner: The BatchRunner instance for processing batches.
+            train_loader: DataLoader for training data.
+            valid_loader: DataLoader for validation data. Defaults to None.
+            test_loader: DataLoader for test data. Defaults to None.
+            network_uses_grad: Whether the network uses gradient computation. Defaults to False.
+            disable_progress_bar: Whether to disable progress bar display. Defaults to False.
+        """
+        self.network = network
+        self.batch_runner = batch_runner
+        self.train_loader = train_loader
+        self.valid_loader = valid_loader
+        self.test_loader = test_loader
+        self.network_uses_grad = network_uses_grad
+        self.disable_progress_bar = disable_progress_bar
+
+    def train(self) -> None:
+        """Run a training epoch over the training DataLoader.
+
+        Sets the network to training mode and iterates over batches,
+        delegating each batch to the BatchRunner for processing.
+        """
+        self.network.train()
+
+        for batch in tqdm(
+            self.train_loader,
+            desc="Training batches",
+            leave=False,
+            disable=self.disable_progress_bar,
+        ):
+            self.batch_runner.train_batch(batch)
+
+    def validate(self) -> None:
+        """Run a validation epoch over the validation DataLoader.
+
+        Sets the network to evaluation mode and iterates over batches,
+        delegating each batch to the BatchRunner for processing.
+        Skips validation if no valid_loader is provided.
+        """
+        if self.valid_loader is None:
+            warnings.warn("Validation skipped: no valid_loader provided.", stacklevel=2)
+            return
+
+        with set_grad_enabled(self.network_uses_grad):
+            self.network.eval()
+
+            for batch in tqdm(
+                self.valid_loader,
+                desc="Validation batches",
+                leave=False,
+                disable=self.disable_progress_bar,
+            ):
+                self.batch_runner.valid_batch(batch)
+
+    def test(self) -> None:
+        """Run a test epoch over the test DataLoader.
+
+        Sets the network to evaluation mode and iterates over batches,
+        delegating each batch to the BatchRunner for processing.
+        Skips testing if no test_loader is provided.
+        """
+        if self.test_loader is None:
+            warnings.warn("Testing skipped: no test_loader provided.", stacklevel=2)
+            return
+
+        with set_grad_enabled(self.network_uses_grad):
+            self.network.eval()
+
+            for batch in tqdm(
+                self.test_loader,
+                desc="Test batches",
+                leave=False,
+                disable=self.disable_progress_bar,
+            ):
+                self.batch_runner.test_batch(batch)

congrads/descriptor.py

@@ -12,7 +12,7 @@ indices, and optional attributes, such as whether the data is constant or variab
 
 from torch import Tensor
 
-from .utils import validate_type
+from .utils.validation import validate_type
 
 
 class Descriptor:

congrads/metrics.py

@@ -9,7 +9,7 @@ from collections.abc import Callable
 
 from torch import Tensor, cat, nanmean, tensor
 
-from .utils import validate_callable, validate_type
+from .utils.validation import validate_callable, validate_type
 
 
 class Metric:

congrads/transformations/base.py

@@ -0,0 +1,37 @@
+"""Module defining transformations and components."""
+
+from abc import ABC, abstractmethod
+
+from torch import Tensor
+
+from ..utils.validation import validate_type
+
+
+class Transformation(ABC):
+    """Abstract base class for tag data transformations."""
+
+    def __init__(self, tag: str):
+        """Initialize a Transformation.
+
+        Args:
+            tag (str): Tag this transformation applies to.
+        """
+        validate_type("tag", tag, str)
+
+        super().__init__()
+        self.tag = tag
+
+    @abstractmethod
+    def __call__(self, data: Tensor) -> Tensor:
+        """Apply the transformation to the input tensor.
+
+        Args:
+            data (Tensor): Input tensor representing network data.
+
+        Returns:
+            Tensor: Transformed tensor.
+
+        Raises:
+            NotImplementedError: Must be implemented by subclasses.
+        """
+        raise NotImplementedError

congrads/{transformations.py → transformations/registry.py}

@@ -1,41 +1,11 @@
-"""Module defining transformations and components."""
+"""Module holding specific transformation implementations."""
 
-from abc import ABC, abstractmethod
 from numbers import Number
 
 from torch import Tensor
 
-from .utils import validate_callable, validate_type
-
-
-class Transformation(ABC):
-    """Abstract base class for tag data transformations."""
-
-    def __init__(self, tag: str):
-        """Initialize a Transformation.
-
-        Args:
-            tag (str): Tag this transformation applies to.
-        """
-        validate_type("tag", tag, str)
-
-        super().__init__()
-        self.tag = tag
-
-    @abstractmethod
-    def __call__(self, data: Tensor) -> Tensor:
-        """Apply the transformation to the input tensor.
-
-        Args:
-            data (Tensor): Input tensor representing network data.
-
-        Returns:
-            Tensor: Transformed tensor.
-
-        Raises:
-            NotImplementedError: Must be implemented by subclasses.
-        """
-        raise NotImplementedError
+from ..utils.validation import validate_callable, validate_type
+from .base import Transformation
 
 
 class IdentityTransformation(Transformation):