TransE is now more memory efficient.
@@ -10,11 +10,27 @@ class TransEScorer(RelationalScorer):
super().__init__(config, dataset)
self._norm = config.get("transe.l_norm")
- def score_emb_spo(self, s_emb, p_emb, o_emb):
+ def score_emb(self, s_emb, p_emb, o_emb, combine: str):
n = p_emb.size(0)
+ if combine == "spo":
out = torch.norm(s_emb + p_emb - o_emb,
p=self._norm,
dim=1) * -1
+ elif combine == "sp*":
+ out = torch.zeros(n, o_emb.size(0)).to(self.config.get("job.device"))
+ for i in range(n):
+ out[i,:] = torch.norm((s_emb[i,:] + p_emb[i,:]) - o_emb,
+ p=self._norm,
+ dim=1) * -1
+ elif combine == "*po":
+ out = torch.zeros(n, s_emb.size(0)).to(self.config.get("job.device"))
+ for i in range(n):
+ out[i,:] = torch.norm(s_emb + (p_emb[i,:] - o_emb[i,:]),
+ p=self._norm,
+ dim=1) * -1
+ else:
+ raise ValueError('cannot handle combine="{}".format(combine)')
+
return out.view(n, -1)
