Temporary fix for cdist bug of pytorch
@@ -3,6 +3,7 @@ from kge import Config, Dataset
from kge.model.kge_model import RelationalScorer, KgeModel
from torch.nn import functional as F
+
class TransEScorer(RelationalScorer):
r"""Implementation of the TransE KGE scorer."""
@@ -11,15 +12,33 @@ class TransEScorer(RelationalScorer):
self._norm = self.get_option("l_norm")
def score_emb(self, s_emb, p_emb, o_emb, combine: str):
+ # TODO cdist gradient computation is currently broken in batched mode; see
+ # https://github.com/pytorch/pytorch/issues/22353. A fix is in
+ # https://github.com/pytorch/pytorch/commit/c33adf539c82fe26ed678a2aea4427fbbcbd7c97,
+ # but has not been released. Once released (pytorch 1.2.1?), only keep the cdist
+ # branches below.
n = p_emb.size(0)
if combine == "spo":
- out = -F.pairwise_distance(s_emb + p_emb, o_emb, p=self._norm)
+ out = -F.pairwise_distance(s_emb + p_emb, o_emb, self._norm)
elif combine == "sp*":
- out = -torch.cdist(s_emb + p_emb, o_emb, p=self._norm)
+ sp_emb = s_emb + p_emb
+ if n > 1 and o_emb.size(0) > 1:
+ out = -torch.cdist(sp_emb, o_emb, p=self._norm)
+ else:
+ out = torch.zeros(n, o_emb.size(0)).to(self.config.get("job.device"))
+ for i in range(n):
+ out[i, :] = -F.pairwise_distance(sp_emb[i, :], o_emb, p=self._norm)
elif combine == "*po":
- out = -torch.cdist(o_emb - p_emb, s_emb, p=self._norm)
+ po_emb = -p_emb + o_emb
+ if n > 1 and s_emb.size(0) > 1:
+ out = -torch.cdist(po_emb, s_emb, self._norm)
+ else:
+ out = torch.zeros(n, s_emb.size(0)).to(self.config.get("job.device"))
+ for i in range(n):
+ out[i, :] = -F.pairwise_distance(po_emb[i, :], s_emb, p=self._norm)
else:
raise ValueError('cannot handle combine="{}".format(combine)')
+
return out.view(n, -1)
