calculate entity penalty only once if s_embedder == o_embedder
@@ -591,19 +591,31 @@ class KgeModel(KgeBase):
# weighted).
if "batch" in kwargs and "triples" in kwargs["batch"]:
triples = kwargs["batch"]["triples"].to(self.config.get("job.device"))
- return (
- super().penalty(**kwargs)
- + self.get_s_embedder().penalty(indexes=triples[:, S], **kwargs)
- + self.get_p_embedder().penalty(indexes=triples[:, P], **kwargs)
- + self.get_o_embedder().penalty(indexes=triples[:, O], **kwargs)
- )
+ penalty_result = super().penalty(**kwargs) + self.get_p_embedder().penalty(indexes=triples[:, P], **kwargs)
+ if self.get_s_embedder() == self.get_o_embedder():
+ entity_penalty_result = self.get_s_embedder().penalty(indexes=torch.cat((triples[:, S].view(-1, 1), triples[:, O].view(-1, 1)), dim=1), **kwargs)
+ if not self.get_s_embedder().get_option("regularize_args.weighted"):
+ # backwards compatibility
+ for penalty in entity_penalty_result:
+ for p in penalty:
+ p *= 2
+ penalty_result += entity_penalty_result
else:
- return (
- super().penalty(**kwargs)
- + self.get_s_embedder().penalty(**kwargs)
- + self.get_p_embedder().penalty(**kwargs)
- + self.get_o_embedder().penalty(**kwargs)
- )
+ penalty_result += self.get_s_embedder().penalty(indexes=triples[:, S], **kwargs)
+ penalty_result += self.get_o_embedder().penalty(indexes=triples[:, O], **kwargs)
+ return penalty_result
+ else:
+ penalty_result = super().penalty(**kwargs) + self.get_p_embedder().penalty(**kwargs)
+ if self.get_s_embedder() == self.get_o_embedder():
+ entity_penalty_result = self.get_s_embedder().penalty(**kwargs)
+ for penalty in entity_penalty_result:
+ for p in penalty:
+ p *= 2
+ penalty_result += entity_penalty_result
+ else:
+ penalty_result += self.get_s_embedder().penalty(**kwargs)
+ penalty_result += self.get_o_embedder().penalty(**kwargs)
+ return penalty_result
def get_s_embedder(self) -> KgeEmbedder:
return self._entity_embedder
