refactor shared sampling code
@@ -26,6 +26,7 @@ class KgeSampler(Configurable):
dataset.num_relations() if slot == P else dataset.num_entities()
)
self.dataset = dataset
+ self.shared = self.get_option("shared")
# auto config
for slot, copy_from in [(S, O), (P, None), (O, S)]:
@@ -64,12 +65,14 @@ class KgeSampler(Configurable):
"""
if num_samples is None:
num_samples = self.num_samples[slot]
- if self.get_option("shared"):
- negative_samples = self._sample_shared(positive_triples, slot, num_samples)
+ if self.shared:
+ negative_samples = self._sample_shared(
+ positive_triples, slot, num_samples).expand(
+ positive_triples.size(0), num_samples)
else:
negative_samples = self._sample(positive_triples, slot, num_samples)
if self.filter_positives[slot]:
- if self.get_option("shared"):
+ if self.shared:
raise ValueError(
"Filtering is not supported with shared negative sampling.")
negative_samples = self._filter(negative_samples, slot, positive_triples)
@@ -136,6 +139,4 @@ class KgeUniformSampler(KgeSampler):
def _sample_shared(
self, positive_triples: torch.Tensor, slot: int, num_samples: int
):
- return torch.randint(
- self.vocabulary_size[slot], (1, num_samples)
- ).expand(positive_triples.size(0), num_samples)
+ return self._sample(torch.empty(1), slot, num_samples)
