Merge pull request #202 from baxtree/development-criterion

group loss functions into the single utitilty class

Author: louisccc

pykg2vec/models/KGMeta.py

@@ -26,12 +26,17 @@ def forward(self, h, r, t):
         raise NotImplementedError
 
     def load_params(self, param_list, kwargs):
+        """Function to load the hyperparameters"""
         for param_name in param_list:
             if param_name not in kwargs:
                 raise Exception("hyperparameter %s not found!" % param_name)
             self.database[param_name] = kwargs[param_name]
         return self.database
 
+    def get_reg(self, h, r, t, **kwargs):
+        """Function to override if regularization is needed"""
+        return 0.0
+
 
 class PairwiseModel(nn.Module, Model):
     """ Meta Class for KGE models with translational distance"""
@@ -74,6 +79,7 @@ def __init__(self, model_name):
         self.training_strategy = TrainingStrategy.PROJECTION_BASED
         self.database = {}  # dict to store model-specific hyperparameter
 
+
 class HyperbolicSpaceModel(nn.Module, Model):
     """ Meta Class for KGE models of hyperbolic space"""
 

pykg2vec/models/hyperbolic.py

@@ -4,6 +4,7 @@
 
 from pykg2vec.models.KGMeta import HyperbolicSpaceModel
 from pykg2vec.models.Domain import NamedEmbedding
+from pykg2vec.utils.criterion import Criterion
 
 
 class MuRP(HyperbolicSpaceModel):
@@ -50,6 +51,8 @@ def __init__(self, **kwargs):
             self.rel_embeddings,
         ]
 
+        self.loss = Criterion.bce
+
     def embed(self, h, r, t):
         """Function to get the embedding value.
 

pykg2vec/models/pairwise.py

@@ -6,6 +6,7 @@
 import numpy as np
 from pykg2vec.models.KGMeta import PairwiseModel
 from pykg2vec.models.Domain import NamedEmbedding
+from pykg2vec.utils.criterion import Criterion
 
 
 class TransE(PairwiseModel):
@@ -58,6 +59,8 @@ def __init__(self, **kwargs):
             self.rel_embeddings,
         ]
 
+        self.loss = Criterion.pairwise_hinge
+
     def forward(self, h, r, t):
         """Function to get the embedding value.
 
@@ -151,6 +154,8 @@ def __init__(self, **kwargs):
             self.w,
         ]
 
+        self.loss = Criterion.pairwise_hinge
+
     def forward(self, h, r, t):
         h_e, r_e, t_e = self.embed(h, r, t)
 
@@ -243,6 +248,8 @@ def __init__(self, **kwargs):
             self.rel_mappings,
         ]
 
+        self.loss = Criterion.pairwise_hinge
+
     def embed(self, h, r, t):
         """Function to get the embedding value.
 
@@ -345,6 +352,8 @@ def __init__(self, **kwargs):
             self.rel_embeddings,
         ]
 
+        self.loss = Criterion.pairwise_hinge
+
     def forward(self, h, r, t):
         """Function to get the embedding value.
 
@@ -431,6 +440,8 @@ def __init__(self, **kwargs):
             self.rel_matrix,
         ]
 
+        self.loss = Criterion.pairwise_hinge
+
     def transform(self, e, matrix):
         matrix = matrix.view(-1, self.ent_hidden_size, self.rel_hidden_size)
         if e.shape[0] != matrix.shape[0]:
@@ -541,6 +552,8 @@ def __init__(self, **kwargs):
             self.mr2,
         ]
 
+        self.loss = Criterion.pairwise_hinge
+
     def embed(self, h, r, t):
         """Function to get the embedding value.
 
@@ -639,6 +652,8 @@ def __init__(self, **kwargs):
             self.bv,
         ]
 
+        self.loss = Criterion.pairwise_hinge
+
     def embed(self, h, r, t):
         """Function to get the embedding value.
 
@@ -665,9 +680,9 @@ def _gu_linear(self, h, r):
             Returns:
                 Tensors: Returns the bilinear loss.
         """
-        mu1h = torch.matmul(self.mu1.weight, self.transpose(h)) # [k, b]
-        mu2r = torch.matmul(self.mu2.weight, self.transpose(r)) # [k, b]
-        return self.transpose(mu1h + mu2r + self.bu.weight)  # [b, k]
+        mu1h = torch.matmul(self.mu1.weight, h.T)  # [k, b]
+        mu2r = torch.matmul(self.mu2.weight, r.T)  # [k, b]
+        return (mu1h + mu2r + self.bu.weight).T  # [b, k]
 
     def _gv_linear(self, r, t):
         """Function to calculate linear loss.
@@ -679,9 +694,9 @@ def _gv_linear(self, r, t):
             Returns:
                 Tensors: Returns the bilinear loss.
         """
-        mv1t = torch.matmul(self.mv1.weight, self.transpose(t)) # [k, b]
-        mv2r = torch.matmul(self.mv2.weight, self.transpose(r)) # [k, b]
-        return self.transpose(mv1t + mv2r + self.bv.weight)  # [b, k]
+        mv1t = torch.matmul(self.mv1.weight, t.T)  # [k, b]
+        mv2r = torch.matmul(self.mv2.weight, r.T)  # [k, b]
+        return (mv1t + mv2r + self.bv.weight).T  # [b, k]
 
     def forward(self, h, r, t):
         """Function to that performs semanting matching.
@@ -701,11 +716,6 @@ def forward(self, h, r, t):
 
         return -torch.sum(self._gu_linear(norm_h, norm_r) * self._gv_linear(norm_r, norm_t), 1)
 
-    @staticmethod
-    def transpose(tensor):
-        dims = tuple(range(len(tensor.shape)-1, -1, -1))    # (rank-1...0)
-        return tensor.permute(dims)
-
 
 class SME_BL(SME):
     """ `A Semantic Matching Energy Function for Learning with Multi-relational Data`_
@@ -729,6 +739,7 @@ class SME_BL(SME):
     def __init__(self, **kwargs):
         super(SME_BL, self).__init__(**kwargs)
         self.model_name = self.__class__.__name__.lower()
+        self.loss = Criterion.pairwise_hinge
 
     def _gu_bilinear(self, h, r):
         """Function to calculate bilinear loss.
@@ -740,9 +751,9 @@ def _gu_bilinear(self, h, r):
             Returns:
                 Tensors: Returns the bilinear loss.
         """
-        mu1h = torch.matmul(self.mu1.weight, self.transpose(h)) # [k, b]
-        mu2r = torch.matmul(self.mu2.weight, self.transpose(r)) # [k, b]
-        return self.transpose(mu1h * mu2r + self.bu.weight)  # [b, k]
+        mu1h = torch.matmul(self.mu1.weight, h.T)  # [k, b]
+        mu2r = torch.matmul(self.mu2.weight, r.T)  # [k, b]
+        return (mu1h * mu2r + self.bu.weight).T  # [b, k]
 
     def _gv_bilinear(self, r, t):
         """Function to calculate bilinear loss.
@@ -754,9 +765,9 @@ def _gv_bilinear(self, r, t):
             Returns:
                 Tensors: Returns the bilinear loss.
         """
-        mv1t = torch.matmul(self.mv1.weight, self.transpose(t)) # [k, b]
-        mv2r = torch.matmul(self.mv2.weight, self.transpose(r)) # [k, b]
-        return self.transpose(mv1t * mv2r + self.bv.weight)  # [b, k]
+        mv1t = torch.matmul(self.mv1.weight, t.T)  # [k, b]
+        mv2r = torch.matmul(self.mv2.weight, r.T)  # [k, b]
+        return (mv1t * mv2r + self.bv.weight).T  # [b, k]
 
     def forward(self, h, r, t):
         """Function to that performs semanting matching.
@@ -821,6 +832,8 @@ def __init__(self, **kwargs):
             self.rel_embeddings,
         ]
 
+        self.loss = Criterion.pariwise_logistic
+
     def embed(self, h, r, t):
         """Function to get the embedding value.
 
@@ -891,6 +904,8 @@ def __init__(self, **kwargs):
             self.rel_matrices,
         ]
 
+        self.loss = Criterion.pairwise_hinge
+
     def embed(self, h, r, t):
         """ Function to get the embedding value.
 
@@ -987,6 +1002,8 @@ def __init__(self, **kwargs):
             self.mr,
         ]
 
+        self.loss = Criterion.pairwise_hinge
+
     def train_layer(self, h, t):
         """ Defines the forward pass training layers of the algorithm.
 
@@ -1030,7 +1047,7 @@ def forward(self, h, r, t):
         norm_t = F.normalize(t_e, p=2, dim=-1)
         return -torch.sum(norm_r*self.train_layer(norm_h, norm_t), -1)
 
-    def get_reg(self):
+    def get_reg(self, h, r, t):
         return self.lmbda*torch.sqrt(sum([torch.sum(torch.pow(var.weight, 2)) for var in self.parameter_list]))
 
 
@@ -1095,6 +1112,8 @@ def __init__(self, **kwargs):
         min_rel = torch.min(torch.FloatTensor().new_full(self.rel_embeddings_sigma.weight.shape, self.cmax), torch.add(self.rel_embeddings_sigma.weight, 1.0))
         self.rel_embeddings_sigma.weight = nn.Parameter(torch.max(torch.FloatTensor().new_full(self.rel_embeddings_sigma.weight.shape, self.cmin), min_rel))
 
+        self.loss = Criterion.pairwise_hinge
+
     def forward(self, h, r, t):
         h_mu, h_sigma, r_mu, r_sigma, t_mu, t_sigma = self.embed(h, r, t)
         return self._cal_score_kl_divergence(h_mu, h_sigma, r_mu, r_sigma, t_mu, t_sigma)
@@ -1199,6 +1218,8 @@ def __init__(self, **kwargs):
             self.rel_embeddings,
         ]
 
+        self.loss = Criterion.pairwise_hinge
+
     def forward(self, h, r, t):
         h_e, r_e, t_e = self.embed(h, r, t)
         r_e = F.normalize(r_e, p=2, dim=-1)