weighting refactory

Co-authored-by: Dario Coscia <[email protected]>

Author: GiovanniCanali

docs/source/_rst/callback/linear_weight_update_callback.rst

@@ -4,11 +4,9 @@
     "SwitchOptimizer",
     "MetricTracker",
     "PINAProgressBar",
-    "LinearWeightUpdate",
     "R3Refinement",
 ]
 
 from .optimizer_callback import SwitchOptimizer
 from .processing_callback import MetricTracker, PINAProgressBar
-from .linear_weight_update_callback import LinearWeightUpdate
 from .refinement import R3Refinement

pina/callback/__init__.py

@@ -2,7 +2,7 @@
 
 import torch
 from .weighting_interface import WeightingInterface
-from ..utils import check_consistency
+from ..utils import check_consistency, in_range
 
 
 class NeuralTangentKernelWeighting(WeightingInterface):
@@ -20,32 +20,34 @@ class NeuralTangentKernelWeighting(WeightingInterface):
 
     """
 
-    def __init__(self, alpha=0.5):
+    def __init__(self, update_every_n_epochs=1, alpha=0.5):
         """
         Initialization of the :class:`NeuralTangentKernelWeighting` class.
 
+        :param int update_every_n_epochs: The number of training epochs between
+            weight updates. If set to 1, the weights are updated at every epoch.
+            Default is 1.
         :param float alpha: The alpha parameter.
         :raises ValueError: If ``alpha`` is not between 0 and 1 (inclusive).
         """
-        super().__init__()
+        super().__init__(update_every_n_epochs=update_every_n_epochs)
 
         # Check consistency
         check_consistency(alpha, float)
-        if alpha < 0 or alpha > 1:
-            raise ValueError("alpha should be a value between 0 and 1")
+        if not in_range(alpha, [0, 1], strict=False):
+            raise ValueError("alpha must be in range (0, 1).")
 
         # Initialize parameters
         self.alpha = alpha
         self.weights = {}
-        self.default_value_weights = 1.0
 
-    def aggregate(self, losses):
+    def weights_update(self, losses):
         """
-        Weight the losses according to the Neural Tangent Kernel algorithm.
+        Update the weighting scheme based on the given losses.
 
-        :param dict(torch.Tensor) input: The dictionary of losses.
-        :return: The aggregation of the losses. It should be a scalar Tensor.
-        :rtype: torch.Tensor
+        :param dict losses: The dictionary of losses.
+        :return: The updated weights.
+        :rtype: dict
         """
         # Define a dictionary to store the norms of the gradients
         losses_norm = {}
@@ -60,14 +62,10 @@ def aggregate(self, losses):
 
         # Update the weights
         self.weights = {
-            condition: self.alpha
-            * self.weights.get(condition, self.default_value_weights)
+            condition: self.alpha * self.weights.get(condition, 1)
             + (1 - self.alpha)
             * losses_norm[condition]
             / sum(losses_norm.values())
             for condition in losses
         }
-
-        return sum(
-            self.weights[condition] * loss for condition, loss in losses.items()
-        )
+        return self.weights

pina/callback/linear_weight_update_callback.py

@@ -4,22 +4,6 @@
 from ..utils import check_consistency
 
 
-class _NoWeighting(WeightingInterface):
-    """
-    Weighting scheme that does not apply any weighting to the losses.
-    """
-
-    def aggregate(self, losses):
-        """
-        Aggregate the losses.
-
-        :param dict losses: The dictionary of losses.
-        :return: The aggregated losses.
-        :rtype: torch.Tensor
-        """
-        return sum(losses.values())
-
-
 class ScalarWeighting(WeightingInterface):
     """
     Weighting scheme that assigns a scalar weight to each loss term.
@@ -36,28 +20,42 @@ def __init__(self, weights):
             dictionary, the default value is used.
         :type weights: float | int | dict
         """
-        super().__init__()
+        super().__init__(update_every_n_epochs=1, aggregator="sum")
 
         # Check consistency
         check_consistency([weights], (float, dict, int))
 
-        # Weights initialization
-        if isinstance(weights, (float, int)):
+        # Initialization
+        if isinstance(weights, dict):
+            self.values = weights
+            self.default_value_weights = 1
+        elif isinstance(weights, (float, int)):
+            self.values = {}
             self.default_value_weights = weights
-            self.weights = {}
         else:
-            self.default_value_weights = 1.0
-            self.weights = weights
+            raise ValueError
 
-    def aggregate(self, losses):
+    def weights_update(self, losses):
         """
-        Aggregate the losses.
+        Update the weighting scheme based on the given losses.
 
         :param dict losses: The dictionary of losses.
-        :return: The aggregated losses.
-        :rtype: torch.Tensor
+        :return: The updated weights.
+        :rtype: dict
+        """
+        return {
+            condition: self.values.get(condition, self.default_value_weights)
+            for condition in losses.keys()
+        }
+
+
+class _NoWeighting(ScalarWeighting):
+    """
+    Weighting scheme that does not apply any weighting to the losses.
+    """
+
+    def __init__(self):
+        """
+        Initialization of the :class:`_NoWeighting` class.
         """
-        return sum(
-            self.weights.get(condition, self.default_value_weights) * loss
-            for condition, loss in losses.items()
-        )
+        super().__init__(weights=1)

pina/loss/ntk_weighting.py

@@ -2,7 +2,6 @@
 
 import torch
 from .weighting_interface import WeightingInterface
-from ..utils import check_positive_integer
 
 
 class SelfAdaptiveWeighting(WeightingInterface):
@@ -22,59 +21,37 @@ class SelfAdaptiveWeighting(WeightingInterface):
 
     """
 
-    def __init__(self, k=100):
+    def __init__(self, update_every_n_epochs=1):
         """
         Initialization of the :class:`SelfAdaptiveWeighting` class.
 
-        :param int k: The number of epochs after which the weights are updated.
-            Default is 100.
-
-        :raises ValueError: If ``k`` is not a positive integer.
+        :param int update_every_n_epochs: The number of training epochs between
+            weight updates. If set to 1, the weights are updated at every epoch.
+            Default is 1.
         """
-        super().__init__()
-
-        # Check consistency
-        check_positive_integer(value=k, strict=True)
+        super().__init__(update_every_n_epochs=update_every_n_epochs)
 
-        # Initialize parameters
-        self.k = k
-        self.weights = {}
-        self.default_value_weights = 1.0
-
-    def aggregate(self, losses):
+    def weights_update(self, losses):
         """
-        Weight the losses according to the self-adaptive algorithm.
+        Update the weighting scheme based on the given losses.
 
-        :param dict(torch.Tensor) losses: The dictionary of losses.
-        :return: The aggregation of the losses. It should be a scalar Tensor.
-        :rtype: torch.Tensor
+        :param dict losses: The dictionary of losses.
+        :return: The updated weights.
+        :rtype: dict
         """
-        # If weights have not been initialized, set them to 1
-        if not self.weights:
-            self.weights = {
-                condition: self.default_value_weights for condition in losses
-            }
-
-        # Update every k epochs
-        if self.solver.trainer.current_epoch % self.k == 0:
-
-            # Define a dictionary to store the norms of the gradients
-            losses_norm = {}
-
-            # Compute the gradient norms for each loss component
-            for condition, loss in losses.items():
-                loss.backward(retain_graph=True)
-                grads = torch.cat(
-                    [p.grad.flatten() for p in self.solver.model.parameters()]
-                )
-                losses_norm[condition] = grads.norm()
-
-            # Update the weights
-            self.weights = {
-                condition: sum(losses_norm.values()) / losses_norm[condition]
-                for condition in losses
-            }
-
-        return sum(
-            self.weights[condition] * loss for condition, loss in losses.items()
-        )
+        # Define a dictionary to store the norms of the gradients
+        losses_norm = {}
+
+        # Compute the gradient norms for each loss component
+        for condition, loss in losses.items():
+            loss.backward(retain_graph=True)
+            grads = torch.cat(
+                [p.grad.flatten() for p in self.solver.model.parameters()]
+            )
+            losses_norm[condition] = grads.norm()
+
+        # Update the weights
+        return {
+            condition: sum(losses_norm.values()) / losses_norm[condition]
+            for condition in losses
+        }