add batching support for self-adaptive pinns

Author: GiovanniCanali

pina/solver/physics_informed_solver/self_adaptive_pinn.py

@@ -15,15 +15,20 @@ class Weights(torch.nn.Module):
     :class:`SelfAdaptivePINN` solver.
     """
 
-    def __init__(self, func):
+    def __init__(self, func, num_points):
         """
         Initialization of the :class:`Weights` class.
 
         :param torch.nn.Module func: the mask model.
+        :param int num_points: the number of input points.
         """
         super().__init__()
+
+        # Check consistency
         check_consistency(func, torch.nn.Module)
-        self.sa_weights = torch.nn.Parameter(torch.Tensor())
+
+        # Initialize the weights as a learnable parameter
+        self.sa_weights = torch.nn.Parameter(torch.zeros(num_points, 1))
         self.func = func
 
     def forward(self):
@@ -140,73 +145,163 @@ def __init__(
             If ``None``, the :class:`torch.nn.MSELoss` loss is used.
             Default is `None`.
         """
-        # check consistency weitghs_function
+        # Check consistency
         check_consistency(weight_function, torch.nn.Module)
 
-        # create models for weights
-        weights_dict = {}
-        for condition_name in problem.conditions:
-            weights_dict[condition_name] = Weights(weight_function)
-        weights_dict = torch.nn.ModuleDict(weights_dict)
+        # Define a ModuleDict for the weights
+        weights = {}
+        for cond, data in problem.input_pts.items():
+            weights[cond] = Weights(func=weight_function, num_points=len(data))
+        weights = torch.nn.ModuleDict(weights)
 
         super().__init__(
-            models=[model, weights_dict],
+            models=[model, weights],
             problem=problem,
             optimizers=[optimizer_model, optimizer_weights],
             schedulers=[scheduler_model, scheduler_weights],
             weighting=weighting,
             loss=loss,
         )
 
-        self._vectorial_loss = deepcopy(self.loss)
-        self._vectorial_loss.reduction = "none"
+        # Extract the reduction method from the loss function
+        self._reduction = self._loss_fn.reduction
 
-    def forward(self, x):
-        """
-        Forward pass.
+        # Set the loss function to return non-aggregated losses
+        self._loss_fn = type(self._loss_fn)(reduction="none")
 
-        :param LabelTensor x: Input tensor.
-        :return: The output of the neural network.
-        :rtype: LabelTensor
+    def training_step(self, batch, batch_idx, **kwargs):
         """
-        return self.model(x)
-
-    def training_step(self, batch):
-        """
-        Solver training step, overridden to perform manual optimization.
+        Solver training step. It computes the optimization cycle and aggregates
+        the losses using the ``weighting`` attribute.
 
         :param list[tuple[str, dict]] batch: A batch of data. Each element is a
             tuple containing a condition name and a dictionary of points.
-        :return: The aggregated loss.
-        :rtype: LabelTensor
+        :param int batch_idx: The index of the current batch.
+        :param dict kwargs: Additional keyword arguments passed to
+            ``optimization_cycle``.
+        :return: The loss of the training step.
+        :rtype: torch.Tensor
         """
         # Weights optimization
         self.optimizer_weights.instance.zero_grad()
-        loss = super().training_step(batch)
+        loss = self._optimization_cycle(
+            batch=batch, batch_idx=batch_idx, **kwargs
+        )
         self.manual_backward(-loss)
         self.optimizer_weights.instance.step()
         self.scheduler_weights.instance.step()
 
         # Model optimization
         self.optimizer_model.instance.zero_grad()
-        loss = super().training_step(batch)
+        loss = self._optimization_cycle(
+            batch=batch, batch_idx=batch_idx, **kwargs
+        )
         self.manual_backward(loss)
         self.optimizer_model.instance.step()
         self.scheduler_model.instance.step()
 
+        # Log the loss
+        self.store_log("train_loss", loss, self.get_batch_size(batch))
+
+        return loss
+
+    @torch.set_grad_enabled(True)
+    def validation_step(self, batch, **kwargs):
+        """
+        The validation step for the Self-Adaptive PINN solver. It returns the
+        average residual computed with the ``loss`` function not aggregated.
+
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
+        :param dict kwargs: Additional keyword arguments passed to
+            ``optimization_cycle``.
+        :return: The loss of the validation step.
+        :rtype: torch.Tensor
+        """
+        losses = self.optimization_cycle(batch=batch, **kwargs)
+
+        # Aggregate losses for each condition
+        for cond, loss in losses.items():
+            losses[cond] = self._apply_reduction(loss=losses[cond])
+
+        loss = (sum(losses.values()) / len(losses)).as_subclass(torch.Tensor)
+        self.store_log("val_loss", loss, self.get_batch_size(batch))
         return loss
 
+    @torch.set_grad_enabled(True)
+    def test_step(self, batch, **kwargs):
+        """
+        The test step for the Self-Adaptive PINN solver. It returns the average
+        residual computed with the ``loss`` function not aggregated.
+
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
+        :param dict kwargs: Additional keyword arguments passed to
+            ``optimization_cycle``.
+        :return: The loss of the test step.
+        :rtype: torch.Tensor
+        """
+        losses = self.optimization_cycle(batch=batch, **kwargs)
+
+        # Aggregate losses for each condition
+        for cond, loss in losses.items():
+            losses[cond] = self._apply_reduction(loss=losses[cond])
+
+        loss = (sum(losses.values()) / len(losses)).as_subclass(torch.Tensor)
+        self.store_log("test_loss", loss, self.get_batch_size(batch))
+        return loss
+
+    def loss_phys(self, samples, equation):
+        """
+        Computes the physics loss for the physics-informed solver based on the
+        provided samples and equation.
+
+        :param LabelTensor samples: The samples to evaluate the physics loss.
+        :param EquationInterface equation: The governing equation.
+        :return: The computed physics loss.
+        :rtype: LabelTensor
+        """
+        residuals = self.compute_residual(samples, equation)
+        return self._loss_fn(residuals, torch.zeros_like(residuals))
+
+    def loss_data(self, input, target):
+        """
+        Compute the data loss for the Self-Adaptive PINN solver by evaluating
+        the loss between the network's output and the true solution. This method
+        should not be overridden, if not intentionally.
+
+        :param input: The input to the neural network.
+        :type input: LabelTensor | torch.Tensor
+        :param target: The target to compare with the network's output.
+        :type target: LabelTensor | torch.Tensor
+        :return: The supervised loss, averaged over the number of observations.
+        :rtype: LabelTensor | torch.Tensor
+        """
+        return self._loss_fn(self.forward(input), target)
+
+    def forward(self, x):
+        """
+        Forward pass.
+
+        :param x: Input tensor.
+        :type x: torch.Tensor | LabelTensor
+        :return: The output of the neural network.
+        :rtype: torch.Tensor | LabelTensor
+        """
+        return self.model(x)
+
     def configure_optimizers(self):
         """
         Optimizer configuration.
 
         :return: The optimizers and the schedulers
         :rtype: tuple[list[Optimizer], list[Scheduler]]
         """
-        # If the problem is an InverseProblem, add the unknown parameters
-        # to the parameters to be optimized
+        # Hook the optimizers to the models
         self.optimizer_model.hook(self.model.parameters())
-        self.optimizer_weights.hook(self.weights_dict.parameters())
+        self.optimizer_weights.hook(self.weights.parameters())
+
+        # Add unknown parameters to optimization list in case of InverseProblem
         if isinstance(self.problem, InverseProblem):
             self.optimizer_model.instance.add_param_group(
                 {
@@ -216,110 +311,88 @@ def configure_optimizers(self):
                     ]
                 }
             )
+
+        # Hook the schedulers to the optimizers
         self.scheduler_model.hook(self.optimizer_model)
         self.scheduler_weights.hook(self.optimizer_weights)
+
         return (
             [self.optimizer_model.instance, self.optimizer_weights.instance],
             [self.scheduler_model.instance, self.scheduler_weights.instance],
         )
 
-    def on_train_start(self):
+    def _optimization_cycle(self, batch, batch_idx, **kwargs):
         """
-        This method is called at the start of the training process to set the
-        self-adaptive weights as parameters of the mask model.
-
-        :raises NotImplementedError: If the batch size is not ``None``.
-        """
-        if self.trainer.batch_size is not None:
-            raise NotImplementedError(
-                "SelfAdaptivePINN only works with full "
-                "batch size, set batch_size=None inside "
-                "the Trainer to use the solver."
-            )
-        device = torch.device(
-            self.trainer._accelerator_connector._accelerator_flag
-        )
+        Aggregate the loss for each condition in the batch.
 
-        # Initialize the self adaptive weights only for training points
-        for (
-            condition_name,
-            tensor,
-        ) in self.trainer.data_module.train_dataset.input.items():
-            self.weights_dict[condition_name].sa_weights.data = torch.rand(
-                (tensor.shape[0], 1), device=device
+        :param list[tuple[str, dict]] batch: A batch of data. Each element is a
+            tuple containing a condition name and a dictionary of points.
+        :param int batch_idx: The index of the current batch.
+        :param dict kwargs: Additional keyword arguments passed to
+            ``optimization_cycle``.
+        :return: The losses computed for all conditions in the batch, casted
+            to a subclass of :class:`torch.Tensor`. It should return a dict
+            containing the condition name and the associated scalar loss.
+        :rtype: dict
+        """
+        # Compute non-aggregated residuals
+        residuals = self.optimization_cycle(batch)
+
+        # Compute losses
+        losses = {}
+        for cond, res in residuals.items():
+
+            weight_tensor = self.weights[cond]()
+
+            # Get the correct indices for the weights. Modulus is used according
+            # to the number of points in the condition, as in the PinaDataset.
+            len_res = len(res)
+            idx = torch.arange(
+                batch_idx * len_res,
+                (batch_idx + 1) * len_res,
+                device=res.device,
+            ) % len(self.problem.input_pts[cond])
+
+            # Apply the weights to the residuals
+            losses[cond] = self._apply_reduction(
+                loss=(res * weight_tensor[idx])
             )
-        return super().on_train_start()
 
-    def on_load_checkpoint(self, checkpoint):
-        """
-        Override of the Pytorch Lightning ``on_load_checkpoint`` method to
-        handle checkpoints for Self-Adaptive Weights. This method should not be
-        overridden, if not intentionally.
-
-        :param dict checkpoint: Pytorch Lightning checkpoint dict.
-        """
-        # First initialize self-adaptive weights with correct shape,
-        # then load the values from the checkpoint.
-        for condition_name, _ in self.problem.input_pts.items():
-            shape = checkpoint["state_dict"][
-                f"_pina_models.1.{condition_name}.sa_weights"
-            ].shape
-            self.weights_dict[condition_name].sa_weights.data = torch.rand(
-                shape
+            # Store log
+            self.store_log(
+                f"{cond}_loss", losses[cond].item(), self.get_batch_size(batch)
             )
-        return super().on_load_checkpoint(checkpoint)
 
-    def loss_phys(self, samples, equation):
-        """
-        Computes the physics loss for the physics-informed solver based on the
-        provided samples and equation.
+        # Clamp unknown parameters in InverseProblem (if needed)
+        self._clamp_params()
 
-        :param LabelTensor samples: The samples to evaluate the physics loss.
-        :param EquationInterface equation: The governing equation.
-        :return: The computed physics loss.
-        :rtype: LabelTensor
-        """
-        residual = self.compute_residual(samples, equation)
-        weights = self.weights_dict[self.current_condition_name].forward()
-        loss_value = self._vectorial_loss(
-            torch.zeros_like(residual, requires_grad=True), residual
-        )
-        return self._vect_to_scalar(weights * loss_value)
-
-    def loss_data(self, input, target):
-        """
-        Compute the data loss for the PINN solver by evaluating the loss
-        between the network's output and the true solution. This method should
-        not be overridden, if not intentionally.
+        # Aggregate
+        loss = self.weighting.aggregate(losses).as_subclass(torch.Tensor)
 
-        :param input: The input to the neural network.
-        :type input: LabelTensor
-        :param target: The target to compare with the network's output.
-        :type target: LabelTensor
-        :return: The supervised loss, averaged over the number of observations.
-        :rtype: LabelTensor
-        """
-        return self._loss_fn(self.forward(input), target)
+        return loss
 
-    def _vect_to_scalar(self, loss_value):
+    def _apply_reduction(self, loss):
         """
-        Computation of the scalar loss.
+        Apply the specified reduction to the loss. The reduction is deferred
+        until the end of the optimization cycle to allow self-adaptive weights
+        to be applied to each point beforehand.
 
-        :param LabelTensor loss_value: the tensor of pointwise losses.
-        :raises RuntimeError: If the loss reduction is not ``mean`` or ``sum``.
-        :return: The computed scalar loss.
-        :rtype: LabelTensor
-        """
-        if self.loss.reduction == "mean":
-            ret = torch.mean(loss_value)
-        elif self.loss.reduction == "sum":
-            ret = torch.sum(loss_value)
-        else:
-            raise RuntimeError(
-                f"Invalid reduction, got {self.loss.reduction} "
-                "but expected mean or sum."
-            )
-        return ret
+        :param torch.Tensor loss: The loss tensor to be reduced.
+        :return: The reduced loss tensor.
+        :rtype: torch.Tensor
+        :raises ValueError: If the reduction method is neither "mean" nor "sum".
+        """
+        # Apply the specified reduction method
+        if self._reduction == "mean":
+            return loss.mean()
+        if self._reduction == "sum":
+            return loss.sum()
+
+        # Raise an error if the reduction method is not recognized
+        raise ValueError(
+            f"Unknown reduction: {self._reduction}."
+            " Supported reductions are 'mean' and 'sum'."
+        )
 
     @property
     def model(self):
@@ -332,7 +405,7 @@ def model(self):
         return self.models[0]
 
     @property
-    def weights_dict(self):
+    def weights(self):
         """
         The self-adaptive weights.