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

congrads/core.py

@@ -1,44 +1,28 @@
 import logging
-from typing import Dict
-from lightning import LightningModule
 from torch import Tensor, float32, no_grad, norm, tensor
-from torchmetrics import Metric
-from torch.nn import ModuleDict
+from torch.optim import Optimizer
+from torch.nn import Module
+from torch.utils.data import DataLoader
+from time import time
 
+from .metrics import MetricManager
 from .constraints import Constraint
-from .metrics import ConstraintSatisfactionRatio
 from .descriptor import Descriptor
 
 
-class CGGDModule(LightningModule):
-    """
-    A PyTorch Lightning module that integrates constraint-guided optimization into the training and validation steps.
+class CongradsCore:
 
-    This module extends the `LightningModule` and incorporates constraints on the neural network's predictions 
-    by adjusting the loss using a rescale factor. The constraints are checked, and the loss is modified to guide 
-    the optimization process based on these constraints.
-
-    Attributes:
-        descriptor (Descriptor): The object that describes the layers and neurons of the network, including 
-                                 the categorization of variable layers.
-        constraints (list[Constraint]): A list of constraints that define the conditions to guide the optimization.
-        train_csr (Dict[str, Metric]): A dictionary of `ConstraintSatisfactionRatio` metrics to track constraint satisfaction 
-                                       during training, indexed by constraint name.
-        valid_csr (Dict[str, Metric]): A dictionary of `ConstraintSatisfactionRatio` metrics to track constraint satisfaction 
-                                       during validation, indexed by constraint name.
-    """
-
-    def __init__(self, descriptor: Descriptor, constraints: list[Constraint]):
-        """
-        Initializes the CGGDModule with a descriptor and a list of constraints.
-
-        Args:
-            descriptor (Descriptor): The object that describes the network's layers and neurons, including their categorization.
-            constraints (list[Constraint]): A list of constraints that will guide the optimization process.
-
-        Raises:
-            Warning if there are no variable layers in the descriptor, as constraints will not be applied.
-        """
+    def __init__(
+        self,
+        descriptor: Descriptor,
+        constraints: list[Constraint],
+        loaders: tuple[DataLoader, DataLoader, DataLoader],
+        network: Module,
+        criterion: callable,
+        optimizer: Optimizer,
+        metric_manager: MetricManager,
+        device,
+    ):
 
         # Init parent class
         super().__init__()
@@ -46,6 +30,14 @@ class CGGDModule(LightningModule):
         # Init object variables
         self.descriptor = descriptor
         self.constraints = constraints
+        self.train_loader = loaders[0]
+        self.valid_loader = loaders[1]
+        self.test_loader = loaders[2]
+        self.network = network
+        self.criterion = criterion
+        self.optimizer = optimizer
+        self.metric_manager = metric_manager
+        self.device = device
 
         # Perform checks
         if len(self.descriptor.variable_layers) == 0:
@@ -53,128 +45,162 @@ class CGGDModule(LightningModule):
                 "The descriptor object has no variable layers. The constraint guided loss adjustment is therefore not used. Is this the intended behaviour?"
             )
 
-        # Assign descriptor to constraints
+        # Initialize constraint metrics
+        metric_manager.register("Loss/train")
+        metric_manager.register("Loss/valid")
+        metric_manager.register("CSR/train")
+        metric_manager.register("CSR/valid")
+
         for constraint in self.constraints:
-            constraint.descriptor = descriptor
-            constraint.run_init_descriptor()
-
-        # Init constraint metric logging
-        self.train_csr: Dict[str, Metric] = ModuleDict(
-            {
-                constraint.constraint_name: ConstraintSatisfactionRatio()
-                for constraint in self.constraints
-            }
-        )
-        self.train_csr["global"] = ConstraintSatisfactionRatio()
-        self.valid_csr: Dict[str, Metric] = ModuleDict(
-            {
-                constraint.constraint_name: ConstraintSatisfactionRatio()
-                for constraint in self.constraints
-            }
-        )
-        self.valid_csr["global"] = ConstraintSatisfactionRatio()
-
-    def training_step(
+            metric_manager.register(f"{constraint.name}/train")
+            metric_manager.register(f"{constraint.name}/valid")
+
+    def fit(self, max_epochs: int = 100):
+        # Loop over epochs
+        for epoch in range(max_epochs):
+
+            # Log start time
+            start_time = time()
+
+            # Training
+            for batch in self.train_loader:
+
+                # Set model in training mode
+                self.network.train()
+
+                # Get input-output pairs from batch
+                inputs, outputs = batch
+
+                # Transfer to GPU
+                inputs, outputs = inputs.to(self.device), outputs.to(self.device)
+
+                # Log preparation time
+                prepare_time = start_time - time()
+
+                # Model computations
+                prediction = self.network(inputs)
+
+                # Calculate loss
+                loss = self.criterion(prediction["output"], outputs)
+                self.metric_manager.accumulate("Loss/train", loss.unsqueeze(0))
+
+                # Adjust loss based on constraints
+                combined_loss = self.train_step(prediction, loss)
+
+                # Backpropx
+                self.optimizer.zero_grad()
+                combined_loss.backward(
+                    retain_graph=False, inputs=list(self.network.parameters())
+                )
+                self.optimizer.step()
+
+            # Validation
+            with no_grad():
+                for batch in self.valid_loader:
+
+                    # Set model in evaluation mode
+                    self.network.eval()
+
+                    # Get input-output pairs from batch
+                    inputs, outputs = batch
+
+                    # Transfer to GPU
+                    inputs, outputs = inputs.to(self.device), outputs.to(self.device)
+
+                    # Model computations
+                    prediction = self.network(inputs)
+
+                    # Calculate loss
+                    loss = self.criterion(prediction["output"], outputs)
+                    self.metric_manager.accumulate("Loss/valid", loss.unsqueeze(0))
+
+                    # Validate constraints
+                    self.valid_step(prediction, loss)
+
+            # TODO with valid loader, checkpoint model with best performance
+
+            # Save metrics
+            self.metric_manager.record(epoch)
+            self.metric_manager.reset()
+
+            # Log compute and preparation time
+            process_time = start_time - time() - prepare_time
+            print(
+                "Compute efficiency: {:.2f}, epoch: {}/{}:".format(
+                    process_time / (process_time + prepare_time), epoch, max_epochs
+                )
+            )
+            start_time = time()
+
+    def train_step(
         self,
         prediction: dict[str, Tensor],
         loss: Tensor,
     ):
-        """
-        The training step where the standard loss is combined with rescale loss based on the constraints.
-
-        For each constraint, the satisfaction ratio is checked, and the loss is adjusted by adding a rescale loss
-        based on the directions calculated by the constraint.
-
-        Args:
-            prediction (dict[str, Tensor]): The model's predictions for each layer.
-            loss (Tensor): The base loss from the model's forward pass.
-
-        Returns:
-            Tensor: The combined loss, including both the original loss and the rescale loss from the constraints.
-        """
 
         # Init scalar tensor for loss
         total_rescale_loss = tensor(0, dtype=float32, device=self.device)
+        loss_grads = {}
 
-        # Compute rescale loss without tracking gradients
+        # Precalculate loss gradients for each variable layer
         with no_grad():
+            for layer in self.descriptor.variable_layers:
+                self.optimizer.zero_grad()
+                loss.backward(retain_graph=True, inputs=prediction[layer])
+                loss_grads[layer] = prediction[layer].grad
 
-            # For each constraint, TODO split into real and validation only constraints
-            for constraint in self.constraints:
+        # For each constraint, TODO split into real and validation only constraints
+        for constraint in self.constraints:
 
-                # Check if constraints are satisfied and calculate directions
+            # Check if constraints are satisfied and calculate directions
+            with no_grad():
                 constraint_checks = constraint.check_constraint(prediction)
                 constraint_directions = constraint.calculate_direction(prediction)
 
-                # Only do direction calculations for variable layers affecting constraint
-                for layer in constraint.layers & self.descriptor.variable_layers:
+            # Only do direction calculations for variable layers affecting constraint
+            for layer in constraint.layers & self.descriptor.variable_layers:
 
+                with no_grad():
                     # Multiply direction modifiers with constraint result
                     constraint_result = (
-                        constraint_checks[layer].unsqueeze(1).type(float32)
+                        constraint_checks.unsqueeze(1).type(float32)
                         * constraint_directions[layer]
                     )
 
-                    # Multiply result with rescale factor o constraint
+                    # Multiply result with rescale factor of constraint
                     constraint_result *= constraint.rescale_factor
 
-                    # Calculate gradients of general loss for each sample
-                    loss.backward(retain_graph=True, inputs=prediction[layer])
-                    loss_grad = prediction[layer].grad
-
                     # Calculate loss gradient norm
-                    norm_loss_grad = norm(loss_grad, dim=0, p=2, keepdim=True)
-
-                    # Calculate rescale loss
-                    rescale_loss = (
-                        (prediction[layer] * constraint_result * norm_loss_grad)
-                        .sum()
-                        .abs()
-                    )
-
-                    # Store rescale loss for this reference space
-                    total_rescale_loss += rescale_loss
-
-                    # Log constraint satisfaction ratio
-                    # NOTE does this take into account spaces with different dimensions?
-                    self.train_csr[constraint.constraint_name](constraint_checks[layer])
-                    self.train_csr["global"](constraint_checks[layer])
-                    self.log(
-                        f"train_csr_{constraint.constraint_name}_{layer}",
-                        self.train_csr[constraint.constraint_name],
-                        on_step=False,
-                        on_epoch=True,
-                    )
-
-        # Log global constraint satisfaction ratio
-        self.log(
-            "train_csr_global",
-            self.train_csr["global"],
-            on_step=False,
-            on_epoch=True,
-        )
+                    norm_loss_grad = norm(loss_grads[layer], dim=1, p=2, keepdim=True)
+
+                # Calculate rescale loss
+                rescale_loss = (
+                    prediction[layer]
+                    * constraint_result
+                    * norm_loss_grad.detach().clone()
+                ).mean()
+
+                # Store rescale loss for this reference space
+                total_rescale_loss += rescale_loss
+
+            # Log constraint satisfaction ratio
+            self.metric_manager.accumulate(
+                f"{constraint.name}/train",
+                (~constraint_checks).type(float32),
+            )
+            self.metric_manager.accumulate(
+                "CSR/train",
+                (~constraint_checks).type(float32),
+            )
 
         # Return combined loss
         return loss + total_rescale_loss
 
-    def validation_step(
+    def valid_step(
         self,
         prediction: dict[str, Tensor],
         loss: Tensor,
     ):
-        """
-        The validation step where the satisfaction of constraints is checked without applying the rescale loss.
-
-        Similar to the training step, but without updating the loss, this method tracks the constraint satisfaction 
-        during validation.
-
-        Args:
-            prediction (dict[str, Tensor]): The model's predictions for each layer.
-            loss (Tensor): The base loss from the model's forward pass.
-
-        Returns:
-            Tensor: The base loss value for validation.
-        """
 
         # Compute rescale loss without tracking gradients
         with no_grad():
@@ -185,27 +211,15 @@ class CGGDModule(LightningModule):
                 # Check if constraints are satisfied for
                 constraint_checks = constraint.check_constraint(prediction)
 
-                # Only do direction calculations for variable layers affecting constraint
-                for layer in constraint.layers & self.descriptor.variable_layers:
-
-                    # Log constraint satisfaction ratio
-                    # NOTE does this take into account spaces with different dimensions?
-                    self.valid_csr[constraint.constraint_name](constraint_checks[layer])
-                    self.valid_csr["global"](constraint_checks[layer])
-                    self.log(
-                        f"valid_csr_{constraint.constraint_name}",
-                        self.valid_csr[constraint.constraint_name],
-                        on_step=False,
-                        on_epoch=True,
-                    )
-
-        # Log global constraint satisfaction ratio
-        self.log(
-            "valid_csr_global",
-            self.valid_csr["global"],
-            on_step=False,
-            on_epoch=True,
-        )
+                # Log constraint satisfaction ratio
+                self.metric_manager.accumulate(
+                    f"{constraint.name}/valid",
+                    (~constraint_checks).type(float32),
+                )
+                self.metric_manager.accumulate(
+                    "CSR/valid",
+                    (~constraint_checks).type(float32),
+                )
 
         # Return loss
         return loss