Pytest unit tests, new losses support, and normalization enhancement

deel/torchlip/functional.py

@@ -278,6 +278,11 @@ def lipschitz_prelu(
 
 
 # Losses
+def apply_reduction(val: torch.Tensor, reduction: str) -> torch.Tensor:
+    red = getattr(torch, reduction, None)
+    if red is None:
+        return val
+    return red(val)
 
 
 def kr_loss(
@@ -305,12 +310,19 @@ def kr_loss(
     """
 
     target = target.view(input.shape)
+    pos_target = (target > 0).to(input.dtype)
+    mean_pos = torch.mean(pos_target, dim=0)
+    # pos factor = batch_size/number of positive samples
+    pos_factor = torch.nan_to_num(1.0 / mean_pos)
+    # neg factor = batch_size/number of negative samples
+    neg_factor = -torch.nan_to_num(1.0 / (1.0 - mean_pos))
 
-    c1 = torch.mean(input[target > 0])
-    c1 = torch.nan_to_num(c1)
-    c2 = torch.mean(input[target <= 0])
-    c2 = torch.nan_to_num(c2)
-    return c1 - c2
+    weighted_input = torch.where(target > 0, pos_factor, neg_factor) * input
+    # Since element-wise KR terms are averaged by loss reduction later on, it is needed
+    # to multiply by batch_size here.
+    # In binary case (`y_true` of shape (batch_size, 1)), `tf.reduce_mean(axis=-1)`
+    # behaves like `tf.squeeze()` to return element-wise loss of shape (batch_size, ).
+    return torch.mean(weighted_input, dim=-1)
 
 
 def neg_kr_loss(
@@ -359,8 +371,8 @@ def hinge_margin_loss(
     """
     target = target.view(input.shape)
     sign_target = torch.where(target > 0, 1.0, -1.0).to(input.dtype)
-
-    return torch.mean(F.relu(min_margin / 2.0 - sign_target * input))
+    hinge = F.relu(min_margin / 2.0 - sign_target * input)
+    return torch.mean(hinge, dim=-1)
 
 
 def hkr_loss(
@@ -395,8 +407,9 @@ def hkr_loss(
         return -kr_loss(input, target)
     # true value: positive value should be the first to be coherent with the
     # hinge loss (positive y_pred)
-    return alpha * hinge_margin_loss(input, target, min_margin)
-    -(1 - alpha) * kr_loss(input, target)
+    return alpha * hinge_margin_loss(input, target, min_margin) - (1 - alpha) * kr_loss(
+        input, target
+    )
 
 
 def kr_multiclass_loss(

deel/torchlip/modules/loss.py

@@ -35,18 +35,20 @@ class KRLoss(torch.nn.Module):
     duality.
     """
 
-    def __init__(self, true_values=None):
+    def __init__(self, reduction: str = "mean", true_values=None):
         """
         Args:
             true_values: tuple containing the two label for each predicted class.
         """
         super().__init__()
+        self.reduction = reduction
         assert (
             true_values is None
         ), "depreciated true_values should be None (use target>0)"
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        return F.kr_loss(input, target)
+        loss_batch = F.kr_loss(input, target)
+        return F.apply_reduction(loss_batch, self.reduction)
 
 
 class NegKRLoss(torch.nn.Module):
@@ -55,35 +57,39 @@ class NegKRLoss(torch.nn.Module):
     the Kantorovich-Rubinstein duality.
     """
 
-    def __init__(self, true_values=None):
+    def __init__(self, reduction: str = "mean", true_values=None):
         """
         Args:
             true_values: tuple containing the two label for each predicted class.
         """
         super().__init__()
+        self.reduction = reduction
         assert (
             true_values is None
         ), "depreciated true_values should be None (use target>0)"
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        return F.neg_kr_loss(input, target)
+        loss_batch = F.neg_kr_loss(input, target)
+        return F.apply_reduction(loss_batch, self.reduction)
 
 
 class HingeMarginLoss(torch.nn.Module):
     """
     Hinge margin loss.
     """
 
-    def __init__(self, min_margin: float = 1.0):
+    def __init__(self, min_margin: float = 1.0, reduction: str = "mean"):
         """
         Args:
             min_margin: The minimal margin to enforce.
         """
         super().__init__()
+        self.reduction = reduction
         self.min_margin = min_margin
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        return F.hinge_margin_loss(input, target, self.min_margin)
+        loss_batch = F.hinge_margin_loss(input, target, self.min_margin)
+        return F.apply_reduction(loss_batch, self.reduction)
 
 
 class HKRLoss(torch.nn.Module):
@@ -96,6 +102,7 @@ def __init__(
         self,
         alpha: float,
         min_margin: float = 1.0,
+        reduction: str = "mean",
         true_values=None,
     ):
         """
@@ -105,6 +112,7 @@ def __init__(
             true_values: tuple containing the two label for each predicted class.
         """
         super().__init__()
+        self.reduction = reduction
         if (alpha >= 0) and (alpha <= 1):
             self.alpha = alpha
         else:
@@ -119,7 +127,8 @@ def __init__(
         ), "depreciated true_values should be None (use target>0)"
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
-        return F.hkr_loss(input, target, self.alpha, self.min_margin)
+        loss_batch = F.hkr_loss(input, target, self.alpha, self.min_margin)
+        return F.apply_reduction(loss_batch, self.reduction)
 
 
 class KRMulticlassLoss(torch.nn.Module):

tests/test_losses.py

@@ -437,7 +437,8 @@ def test_softhkrmulticlass_loss():
             },
             binary_data(y_true1b),
             binary_data(y_pred1b),
-            [0.25, -2.2, 2.95, -0.65, -2.6, 0.7, 3.4, -1.55],
+            np.float64([0.25, -2.2, 2.95, -0.65, -2.6, 0.7, 3.4, -1.55])
+            * uft.scaleAlpha(2.5),
             1e-7,
         ),
         (
@@ -545,7 +546,7 @@ def test_no_reduction_loss_generic(
         expected_loss,
         rtol=rtol,
         atol=5e-6,
-        err_msg=f"Loss {loss.name} failed",
+        err_msg=f"Loss {loss} failed",
     )
 
 
@@ -749,7 +750,7 @@ def test_minibatches_binary_loss_generic(
             expected_loss,
             rtol=rtol,
             atol=5e-6,
-            err_msg=f"Loss {loss.name} failed",
+            err_msg=f"Loss {loss} failed",
         )
     loss_val_minibatches = 0
     for i in range(len(segments) - 1):
@@ -764,7 +765,7 @@ def test_minibatches_binary_loss_generic(
         loss_val_minibatches,
         rtol=rtol,  # 5e-6,
         atol=5e-6,
-        err_msg=f"Loss {loss.name} failed for hardcoded mini-batches",
+        err_msg=f"Loss {loss} failed for hardcoded mini-batches",
     )
 
 
@@ -870,5 +871,5 @@ def test_multilabel_loss_generic(loss_instance, loss_params, rtol):
         mean_loss_vals,
         rtol=rtol,  # 5e-6,
         atol=1e-4,
-        err_msg=f"Loss {loss.name} failed",
+        err_msg=f"Loss {loss} failed",
     )

tests/utils_framework.py

@@ -219,11 +219,17 @@ def get_instance_withcheck(
     ScaledAvgPool2d: partial(
         get_instance_withreplacement, dict_keys_replace={"data_format": None}
     ),
-    KRLoss: partial(get_instance_withcheck, list_keys_notimplemented=["reduction"]),
-    HingeMarginLoss: partial(
-        get_instance_withcheck, list_keys_notimplemented=["reduction"]
+    KRLoss: partial(
+        get_instance_withcheck,
+        dict_keys_replace={"name": None},
+        list_keys_notimplemented=["multi_gpu"],
+    ),
+    HingeMarginLoss: partial(get_instance_withcheck, dict_keys_replace={"name": None}),
+    HKRLoss: partial(
+        get_instance_withcheck,
+        dict_keys_replace={"name": None},
+        list_keys_notimplemented=["multi_gpu"],
     ),
-    HKRLoss: partial(get_instance_withcheck, list_keys_notimplemented=["reduction"]),
     HingeMulticlassLoss: partial(
         get_instance_withcheck, list_keys_notimplemented=["reduction"]
     ),