Added only custom_losses.py and base_model.py

Author: Akash-bu

dicee/losses/custom_losses.py

@@ -375,61 +375,68 @@ def forward(self, inputs, targets, current_epoch):
 
         return loss
 
+
+
 class UNITEI(nn.Module):
     """
-    Implementation of UNITE-I loss
+    UNITE-I loss (pattern-aware & noise-resilient) for mini-batch training.
     """
-    def __init__(self,
-                 gamma: float = 10,
-                 sigma: float = 1000.0,
-                #  penalty_weight: float = 1.0,
-                #  tau_init: float = 1e-3
-                 ):
-        super(UNITEI, self).__init__() 
-
-        # self.num_triples = num_triples
-        self.gamma = gamma 
-        self.sigma = sigma 
-        # self.penalty_weight = penalty_weight 
-        # self.tau = tau_init 
-        # self.eps = 1e-12 
-        self.relu = nn.ReLU() 
-
-    def forward(self, pred, target, current_epoch = None):
-        """
-        This function calculates the UNITE loss.
-
-        Agrs: 
-        pred: The output logits from the KGE model.
-        target: The target labesl
-
-        """ 
-        target = target.float() 
-
-        #Separate positive and negative logits 
-        l_pos = pred[target == 1.0]
-        l_neg = pred[target == 0.0] 
-
-        loss_pos_total = torch.tensor(0.0, device=pred.device)
-        loss_neg_total = torch.tensor(0.0, device=pred.device) 
-
-        if l_pos.numel() > 0:
-            tau_pos_star = self.relu(self.gamma - l_pos) 
-            loss_pos = self.sigma * (tau_pos_star ** 2)
-            loss_pos_total = torch.sum(loss_pos)
-
-        if l_neg.numel() > 0:
-            tau_neg_star = self.relu(l_neg - self.gamma)
-            loss_neg = self.sigma * (tau_neg_star ** 2)
-            loss_neg_total = torch.sum(loss_neg) 
-
-        total_loss = loss_pos_total + loss_neg_total 
-
-        if pred.numel() == 0:
-            return torch.tensor(0.0, device=pred.device, requires_grad=True)
-
-        return total_loss / pred.numel() 
-
-
-
 
+    def __init__(
+        self,
+        gamma: float = 5,     
+        sigma: float = 1000,  
+        lambda_q: float = 0.3,  # Regulaization term
+        clamp_scores: bool = True,
+        use_tanh_scale: bool = True,
+    ):
+        super().__init__()
+        self.gamma = float(gamma)
+        self.sigma = float(sigma)
+        self.lambda_q = float(lambda_q)
+        self.clamp_scores = clamp_scores
+        self.use_tanh_scale = use_tanh_scale
+        self.eps = 1e-12
+
+    def _normalize_scores(self, pred: torch.Tensor) -> torch.Tensor:
+        # To keep gamma meaningful and gradients don’t die or explode.
+        x = pred
+        if self.use_tanh_scale:
+            # squashes to (-1,1), then scales to (-γ, γ)
+            x = torch.tanh(x) * self.gamma
+        if self.clamp_scores:
+            x = torch.clamp(x, min=-self.gamma, max=self.gamma)
+        return x
+
+    def forward(self, pred: torch.Tensor, target: torch.Tensor, current_epoch = None):
+
+        target = target.float()
+        pred = self._normalize_scores(pred)
+
+        tau_pos = F.relu(self.gamma - pred)  # only where target==1 used
+        tau_neg = F.relu(pred - self.gamma)  # only where target==0 used
+
+        loss_pos = self.sigma * (tau_pos ** 2)
+        loss_neg = self.sigma * (tau_neg ** 2)
+
+        # Constraint penalties with Q(x) = sigmoid(x)
+        # Pos: Q(gamma - f) >= Q(tau) → violation = ReLU(Q(tau) - Q(gamma - f)) where f=pred
+        # Neg: Q(f - gamma) >= Q(tau) → violation = ReLU(Q(tau) - Q(f - tau))
+        Q_tau_pos = torch.sigmoid(tau_pos)
+        Q_tau_neg = torch.sigmoid(tau_neg)
+        Q_pos = torch.sigmoid(self.gamma - pred)
+        Q_neg = torch.sigmoid(pred - self.gamma)
+
+        pos_violation = F.relu(Q_tau_pos - Q_pos)
+        neg_violation = F.relu(Q_tau_neg - Q_neg)
+
+        # Mask by labels and combine
+        pos_mask = target
+        neg_mask = 1.0 - target
+
+        loss = (
+            pos_mask * (loss_pos + self.lambda_q * pos_violation)
+            + neg_mask * (loss_neg + self.lambda_q * neg_violation)
+        )
+
+        return loss.mean()

dicee/models/base_model.py

@@ -43,12 +43,39 @@ def training_step(self, batch, batch_idx=None):
             # Default
             x_batch, y_batch = batch
             yhat_batch = self.forward(x_batch)
+
+        if batch_idx == 0:
+            with torch.no_grad():
+                preds_sample = yhat_batch.view(-1)[:30]
+                mean_pred = preds_sample.mean().item()
+                min_pred = preds_sample.min().item()
+                max_pred = preds_sample.max().item()
+                self.log_dict({
+                    "debug/pred_mean": mean_pred,
+                    "debug/pred_min": min_pred,
+                    "debug/pred_max": max_pred
+                }, prog_bar=True, on_step=True, logger=True)
+
+
         elif len(batch)==3:
             # KvsSample or 1vsSample
             x_batch, y_select, y_batch = batch
             yhat_batch = self.forward((x_batch,y_select))
-        else:
-            raise RuntimeError("Invalid batch received.")
+            if batch_idx == 0:
+                with torch.no_grad():
+                    preds_sample = yhat_batch.view(-1)[:30]
+                    mean_pred = preds_sample.mean().item()
+                    min_pred = preds_sample.min().item()
+                    max_pred = preds_sample.max().item()
+                    self.log_dict({
+                        "debug/pred_mean": mean_pred,
+                        "debug/pred_min": min_pred,
+                        "debug/pred_max": max_pred
+                    }, prog_bar=True, on_step=True, logger=True)
+
+            else:
+                raise RuntimeError("Invalid batch received.")
+
 
         #total_norm = 0
         #for param in self.parameters():
@@ -202,8 +229,8 @@ def __init__(self, args: dict):
         if self.args["loss_fn"] == "ACLS":
             self.loss = ACLS()
         if self.args["loss_fn"] == "UNITEI":
-            self.loss = UNITEI(gamma = self.args.get("unite_gamma", 5.0), 
-                               sigma = self.args.get("unite_sigma", 1000.0))
+            self.loss = UNITEI()
+
        #if self.args["loss_fn"] == "GradientBasedLSLR":
        #    self.loss = GradientBasedLSLR()
        #if self.args["loss_fn"] == "GradientBasedAdaptiveLSLR":