tidy up

Author: baxtree

pykg2vec/models/pointwise.py

@@ -208,10 +208,10 @@ def get_reg(self, h, r, t, reg_type="F2"):
 
         if reg_type.lower() == 'f2':
             regul_term = torch.mean(torch.sum(h_e_real ** 2, -1) + torch.sum(h_e_img ** 2, -1) + torch.sum(r_e_real ** 2, -1) +
-                torch.sum(r_e_img ** 2, -1) + torch.sum(t_e_real ** 2, -1) + torch.sum(t_e_img ** 2, -1))
+                                    torch.sum(r_e_img ** 2, -1) + torch.sum(t_e_real ** 2, -1) + torch.sum(t_e_img ** 2, -1))
         elif reg_type.lower() == 'n3':
             regul_term = torch.mean(torch.sum(h_e_real ** 3, -1) + torch.sum(h_e_img ** 3, -1) + torch.sum(r_e_real ** 3, -1) +
-                torch.sum(r_e_img ** 3, -1) + torch.sum(t_e_real ** 3, -1) + torch.sum(t_e_img ** 3, -1))
+                                    torch.sum(r_e_img ** 3, -1) + torch.sum(t_e_real ** 3, -1) + torch.sum(t_e_img ** 3, -1))
         else:
             raise NotImplementedError('Unknown regularizer type: %s' % reg_type)
 

pykg2vec/models/projection.py

@@ -214,9 +214,9 @@ def forward(self, e, r, er_e2, direction="tail"):
         emb_hr_r = self.rel_embeddings(r)  # [m, k]
 
         if direction == "tail":
-            ere2_sigmoid = self.g(torch.dropout(self.f1(emb_hr_e, emb_hr_r), p=self.hidden_dropout, train=True), self.ent_embeddings.weight)
+            ere2_sigmoid = ProjE_pointwise.g(torch.dropout(self.f1(emb_hr_e, emb_hr_r), p=self.hidden_dropout, train=True), self.ent_embeddings.weight)
         else:
-            ere2_sigmoid = self.g(torch.dropout(self.f2(emb_hr_e, emb_hr_r), p=self.hidden_dropout, train=True), self.ent_embeddings.weight)
+            ere2_sigmoid = ProjE_pointwise.g(torch.dropout(self.f2(emb_hr_e, emb_hr_r), p=self.hidden_dropout, train=True), self.ent_embeddings.weight)
 
         ere2_loss_left = -torch.sum((torch.log(torch.clamp(ere2_sigmoid, 1e-10, 1.0)) * torch.max(torch.FloatTensor([0]).to(self.device), er_e2)))
         ere2_loss_right = -torch.sum((torch.log(torch.clamp(1 - ere2_sigmoid, 1e-10, 1.0)) * torch.max(torch.FloatTensor([0]).to(self.device), torch.neg(er_e2))))
@@ -243,21 +243,11 @@ def f2(self, t, r):
         """
         return torch.tanh(t * self.De2.weight + r * self.Dr2.weight + self.bc2.weight)
 
-    def g(self, f, w):
-        """Defines activation layer.
-
-            Args:
-               f (Tensor): output of the forward layers.
-               w (Tensor): Matrix for multiplication.
-        """
-        # [b, k] [k, tot_ent]
-        return torch.sigmoid(torch.matmul(f, w.T))
-
     def predict_tail_rank(self, h, r, topk=-1):
         emb_h = self.ent_embeddings(h)  # [1, k]
         emb_r = self.rel_embeddings(r)  # [1, k]
 
-        hrt_sigmoid = -self.g(self.f1(emb_h, emb_r), self.ent_embeddings.weight)
+        hrt_sigmoid = -ProjE_pointwise.g(self.f1(emb_h, emb_r), self.ent_embeddings.weight)
         _, rank = torch.topk(hrt_sigmoid, k=topk)
 
         return rank
@@ -266,11 +256,21 @@ def predict_head_rank(self, t, r, topk=-1):
         emb_t = self.ent_embeddings(t)  # [m, k]
         emb_r = self.rel_embeddings(r)  # [m, k]
 
-        hrt_sigmoid = -self.g(self.f2(emb_t, emb_r), self.ent_embeddings.weight)
+        hrt_sigmoid = -ProjE_pointwise.g(self.f2(emb_t, emb_r), self.ent_embeddings.weight)
         _, rank = torch.topk(hrt_sigmoid, k=topk)
 
         return rank
 
+    @staticmethod
+    def g(f, w):
+        """Defines activation layer.
+
+            Args:
+               f (Tensor): output of the forward layers.
+               w (Tensor): Matrix for multiplication.
+        """
+        # [b, k] [k, tot_ent]
+        return torch.sigmoid(torch.matmul(f, w.T))
 
 class TuckER(ProjectionModel):
     """

pykg2vec/utils/trainer.py

@@ -173,13 +173,7 @@ def train_step_projection(self, h, r, t, hr_t, tr_h):
 
         return loss
 
-    def train_step_pointwise(self, h, r, t, target):
-        preds = self.model(h, r, t)
-        loss = self.model.loss(preds, target)
-        loss += self.model.get_reg(h, r, t)
-        return loss
-
-    def train_step_hyperbolic(self, h, r, t, target):
+    def train_step_pointwise_hyperbolic(self, h, r, t, target):
         preds = self.model(h, r, t)
         loss = self.model.loss(preds, target)
         loss += self.model.get_reg(h, r, t)
@@ -289,7 +283,7 @@ def train_model_epoch(self, epoch_idx, tuning=False):
                 r = torch.LongTensor(data[1]).to(self.config.device)
                 t = torch.LongTensor(data[2]).to(self.config.device)
                 y = torch.LongTensor(data[3]).to(self.config.device)
-                loss = self.train_step_pointwise(h, r, t, y)
+                loss = self.train_step_pointwise_hyperbolic(h, r, t, y)
             elif self.model.training_strategy == TrainingStrategy.PAIRWISE_BASED:
                 pos_h = torch.LongTensor(data[0]).to(self.config.device)
                 pos_r = torch.LongTensor(data[1]).to(self.config.device)
@@ -303,7 +297,7 @@ def train_model_epoch(self, epoch_idx, tuning=False):
                 r = torch.cat((torch.LongTensor(data[1]).to(self.config.device), torch.LongTensor(data[4]).to(self.config.device)), dim=-1)
                 t = torch.cat((torch.LongTensor(data[2]).to(self.config.device), torch.LongTensor(data[5]).to(self.config.device)), dim=-1)
                 y = torch.cat((torch.ones(np.array(data[0]).shape).to(self.config.device), torch.zeros(np.array(data[3]).shape).to(self.config.device)), dim=-1)
-                loss = self.train_step_hyperbolic(h, r, t, y)
+                loss = self.train_step_pointwise_hyperbolic(h, r, t, y)
             else:
                 raise NotImplementedError("Unknown training strategy: %s" % self.model.training_strategy)