congrads 0.2.0__py3-none-any.whl → 0.3.1__py3-none-any.whl

congrads/checkpoints.py

@@ -0,0 +1,180 @@
+"""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
+
+__all__ = ["CheckpointManager"]
+
+
+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

congrads/constraints/base.py

@@ -0,0 +1,244 @@
+"""Defines the abstract base class `Constraint` for specifying constraints on neural network outputs.
+
+A `Constraint` monitors whether the network predictions satisfy certain
+conditions during training, validation, and testing. It can optionally
+adjust the loss to enforce constraints, and logs the relevant metrics.
+
+Responsibilities:
+- Track which network layers/tags the constraint applies to
+- Check constraint satisfaction for a batch of predictions
+- Compute adjustment directions to enforce the constraint
+- Provide a rescale factor and enforcement flag to influence loss adjustment
+
+Subclasses must implement the abstract methods:
+- `check_constraint(data)`: Evaluate constraint satisfaction for a batch
+- `calculate_direction(data)`: Compute directions to adjust predictions
+"""
+
+import random
+import string
+import warnings
+from abc import ABC, abstractmethod
+from numbers import Number
+from typing import Literal
+
+from torch import Tensor
+
+from congrads.descriptor import Descriptor
+from congrads.utils.validation import validate_iterable, validate_type
+
+__all__ = ["Constraint", "MonotonicityConstraint"]
+
+
+class Constraint(ABC):
+    """Abstract base class for defining constraints applied to neural networks.
+
+    A `Constraint` specifies conditions that the neural network outputs
+    should satisfy. It supports monitoring constraint satisfaction
+    during training and can adjust loss to enforce constraints. Subclasses
+    must implement the `check_constraint` and `calculate_direction` methods.
+
+    Args:
+        tags (set[str]): Tags referencing parts of the network where this constraint applies to.
+        name (str, optional): A unique name for the constraint. If not provided,
+            a name is generated based on the class name and a random suffix.
+        enforce (bool, optional): If False, only monitor the constraint
+            without adjusting the loss. Defaults to True.
+        rescale_factor (Number, optional): Factor to scale the
+            constraint-adjusted loss. Defaults to 1.5. Should be greater
+            than 1 to give weight to the constraint.
+
+    Raises:
+        TypeError: If a provided attribute has an incompatible type.
+        ValueError: If any tag in `tags` is not
+            defined in the `descriptor`.
+
+    Note:
+        - If `rescale_factor <= 1`, a warning is issued.
+        - If `name` is not provided, a name is auto-generated,
+          and a warning is logged.
+
+    """
+
+    descriptor: Descriptor = None
+    device = None
+
+    def __init__(
+        self, tags: set[str], name: str = None, enforce: bool = True, rescale_factor: Number = 1.5
+    ) -> None:
+        """Initializes a new Constraint instance.
+
+        Args:
+            tags (set[str]): Tags referencing parts of the network where this constraint applies to.
+            name (str, optional): A unique name for the constraint. If not
+                provided, a name is generated based on the class name and a
+                random suffix.
+            enforce (bool, optional): If False, only monitor the constraint
+                without adjusting the loss. Defaults to True.
+            rescale_factor (Number, optional): Factor to scale the
+                constraint-adjusted loss. Defaults to 1.5. Should be greater
+                than 1 to give weight to the constraint.
+
+        Raises:
+            TypeError: If a provided attribute has an incompatible type.
+            ValueError: If any tag in `tags` is not defined in the `descriptor`.
+
+        Note:
+            - If `rescale_factor <= 1`, a warning is issued.
+            - If `name` is not provided, a name is auto-generated, and a
+            warning is logged.
+        """
+        # Init parent class
+        super().__init__()
+
+        # Type checking
+        validate_iterable("tags", tags, str)
+        validate_type("name", name, str, allow_none=True)
+        validate_type("enforce", enforce, bool)
+        validate_type("rescale_factor", rescale_factor, Number)
+
+        # Init object variables
+        self.tags = tags
+        self.rescale_factor = rescale_factor
+        self.initial_rescale_factor = rescale_factor
+        self.enforce = enforce
+
+        # Perform checks
+        if rescale_factor <= 1:
+            warnings.warn(
+                f"Rescale factor for constraint {name} is <= 1. The network "
+                "will favor general loss over the constraint-adjusted loss. "
+                "Is this intended behavior? Normally, the rescale factor "
+                "should always be larger than 1.",
+                stacklevel=2,
+            )
+
+        # If no constraint_name is set, generate one based
+        # on the class name and a random suffix
+        if name:
+            self.name = name
+        else:
+            random_suffix = "".join(random.choices(string.ascii_uppercase + string.digits, k=6))
+            self.name = f"{self.__class__.__name__}_{random_suffix}"
+            warnings.warn(f"Name for constraint is not set. Using {self.name}.", stacklevel=2)
+
+        # Infer layers from descriptor and tags
+        self.layers = set()
+        for tag in self.tags:
+            if not self.descriptor.exists(tag):
+                raise ValueError(
+                    f"The tag {tag} used with constraint "
+                    f"{self.name} is not defined in the descriptor. Please "
+                    "add it to the correct layer using "
+                    "descriptor.add('layer', ...)."
+                )
+
+            layer, _ = self.descriptor.location(tag)
+            self.layers.add(layer)
+
+    @abstractmethod
+    def check_constraint(self, data: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
+        """Evaluates whether the given model predictions satisfy the constraint.
+
+        1 IS SATISFIED, 0 IS NOT SATISFIED
+
+        Args:
+            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
+
+        Returns:
+            tuple[Tensor, Tensor]: A tuple where the first element is a tensor of floats
+            indicating whether the constraint is satisfied (with value 1.0
+            for satisfaction, and 0.0 for non-satisfaction, and the second element is a tensor
+            mask that indicates the relevance of each sample (`True` for relevant
+            samples and `False` for irrelevant ones).
+        """
+        pass
+
+    @abstractmethod
+    def calculate_direction(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Compute adjustment directions to better satisfy the constraint.
+
+        Given the model predictions, input batch, and context, this method calculates the direction
+        in which the predictions referenced by a tag should be adjusted to satisfy the constraint.
+
+        Args:
+            data (dict[str, Tensor]): Dictionary that holds batch data, model predictions and context.
+
+        Returns:
+            dict[str, Tensor]: Dictionary mapping network layers to tensors that
+                specify the adjustment direction for each tag.
+        """
+        pass
+
+
+class MonotonicityConstraint(Constraint, ABC):
+    """Abstract base class for monotonicity constraints.
+
+    Subclasses must define how monotonicity is evaluated and how corrective
+    directions are computed.
+    """
+
+    def __init__(
+        self,
+        tag_prediction: str,
+        tag_reference: str,
+        rescale_factor_lower: float = 1.5,
+        rescale_factor_upper: float = 1.75,
+        stable: bool = True,
+        direction: Literal["ascending", "descending"] = "ascending",
+        name: str = None,
+        enforce: bool = True,
+    ):
+        """Constraint that enforces monotonicity on a predicted output.
+
+        This constraint ensures that the activations of a prediction tag (`tag_prediction`)
+        are monotonically ascending or descending with respect to a target tag (`tag_reference`).
+
+        Args:
+            tag_prediction (str): Name of the tag whose activations should follow the monotonic relationship.
+            tag_reference (str): Name of the tag that acts as the monotonic reference.
+            rescale_factor_lower (float, optional): Lower bound for rescaling rank differences. Defaults to 1.5.
+            rescale_factor_upper (float, optional): Upper bound for rescaling rank differences. Defaults to 1.75.
+            stable (bool, optional): Whether to use stable sorting when ranking. Defaults to True.
+            direction (str, optional): Direction of monotonicity to enforce, either 'ascending' or 'descending'. Defaults to 'ascending'.
+            name (str, optional): Custom name for the constraint. If None, a descriptive name is auto-generated.
+            enforce (bool, optional): If False, the constraint is only monitored (not enforced). Defaults to True.
+        """
+        # Type checking
+        validate_type("rescale_factor_lower", rescale_factor_lower, float)
+        validate_type("rescale_factor_upper", rescale_factor_upper, float)
+        validate_type("stable", stable, bool)
+        validate_type("direction", direction, str)
+
+        # Compose constraint name
+        if name is None:
+            name = f"{tag_prediction} monotonically {direction} by {tag_reference}"
+
+        # Init parent class
+        super().__init__({tag_prediction}, name, enforce, 1.0)
+
+        # Init variables
+        self.tag_prediction = tag_prediction
+        self.tag_reference = tag_reference
+        self.rescale_factor_lower = rescale_factor_lower
+        self.rescale_factor_upper = rescale_factor_upper
+        self.stable = stable
+        self.direction = direction
+        self.descending = direction == "descending"
+
+        # Init member variables
+        self.compared_rankings: Tensor = None
+
+    @abstractmethod
+    def check_constraint(self, data: dict[str, Tensor]) -> tuple[Tensor, Tensor]:
+        """Evaluate whether the monotonicity constraint is satisfied.
+
+        Implementations must set `self.compared_rankings` with per-sample
+        correction directions.
+        """
+        pass
+
+    @abstractmethod
+    def calculate_direction(self, data: dict[str, Tensor]) -> dict[str, Tensor]:
+        """Return directions for monotonicity enforcement."""
+        pass