score only against needed entities in negsam sp_po
@@ -818,6 +818,70 @@ class TrainingJobNegativeSampling(TrainingJob):
).view(-1)
prepare_time += time.time()
+ # now pick the scores we need
+ forward_time -= time.time()
+ scores = all_scores[row_indexes, column_indexes].view(batch_size, -1)
+ forward_time += time.time()
+ elif self._implementation == "sp_po_unique":
+ # compute all scores for slot
+ forward_time -= time.time()
+ unique_triples, column_indexes = torch.unique(
+ torch.cat((triples[:, [slot]], negative_samples[slot]), 1).view(-1),
+ return_inverse=True)
+ if slot == S:
+ all_scores = self.model.score_po(triples[:, P], triples[:, O],
+ unique_triples)
+ elif slot == O:
+ all_scores = self.model.score_sp(triples[:, S], triples[:, P],
+ unique_triples)
+ else:
+ raise NotImplementedError
+ forward_time += time.time()
+
+ # determine indexes of relevant scores in scoring matrix
+ prepare_time -= time.time()
+ row_indexes = (
+ torch.arange(batch_size, device=self.device)
+ .unsqueeze(1)
+ .repeat(1, 1 + num_samples)
+ .view(-1)
+ ) # 000 111 222; each 1+num_negative times (here: 3)
+ prepare_time += time.time()
+
+ # now pick the scores we need
+ forward_time -= time.time()
+ scores = all_scores[row_indexes, column_indexes].view(batch_size, -1)
+ forward_time += time.time()
+ elif self._implementation == "sp_po_unique_all":
+ # compute all scores for slot
+ forward_time -= time.time()
+ unique_triples, column_indexes = torch.unique(
+ torch.cat((triples[:, [slot]], negative_samples[slot]), 1).view(-1),
+ return_inverse=True)
+ if slot == S:
+ unique_po, row_index = torch.unique(triples[:, [P, O]], dim=0,
+ return_inverse=True)
+ all_scores = self.model.score_po(unique_po[:, 0], unique_po[:, 1],
+ unique_triples)
+ elif slot == O:
+ unique_sp, row_index = torch.unique(triples[:, [S, P]], dim=0,
+ return_inverse=True)
+ all_scores = self.model.score_sp(unique_sp[:, 0], unique_sp[:, 1],
+ unique_triples)
+ else:
+ raise NotImplementedError
+ forward_time += time.time()
+
+ # determine indexes of relevant scores in scoring matrix
+ prepare_time -= time.time()
+ row_indexes = (
+ row_index
+ .unsqueeze(1)
+ .repeat(1, 1 + num_samples)
+ .view(-1)
+ ) # 000 111 222; each 1+num_negative times (here: 3)
+ prepare_time += time.time()
+
# now pick the scores we need
forward_time -= time.time()
scores = all_scores[row_indexes, column_indexes].view(batch_size, -1)
