def test_instantiate_distmult(self):
"""Test that DistMult can be instantiated."""
distmult = DistMult(**DISTMULT_CONFIG)
distmult.num_entities = DISTMULT_CONFIG[NUM_ENTITIES]
distmult.num_relations = DISTMULT_CONFIG[NUM_RELATIONS]
self.assertIsNotNone(distmult)
self.assertEqual(distmult.num_entities, 5)
self.assertEqual(distmult.num_relations, 5)
self.assertEqual(distmult.embedding_dim, 5)
self.assertEqual(distmult.margin_loss, 4)
def test_compute_scores_distmult(self):
"""Test that DistMult's socore function computes the scores correct."""
distmult = DistMult(config=DISTMULT_CONFIG)
h_embs = torch.tensor([[1., 1.], [1., 1.]], dtype=torch.float)
r_embs = torch.tensor([[1., 1.], [2., 2.]], dtype=torch.float)
t_embs = torch.tensor([[2., 2.], [4., 4.]], dtype=torch.float)
scores = distmult._compute_scores(h_embs, r_embs,
t_embs).cpu().numpy().tolist()
self.assertEqual(scores, [-4., -16.])
class DistMultUnit(EnsembleUnit):
def __init__(self,
relation_embeddings,
num_entities,
num_relations,
entity_embeddings,
preferred_device='cpu'):
super(DistMultUnit, self).__init__()
if preferred_device == 'cuda':
preferred_device = 'gpu'
self.model = DistMult(preferred_device=preferred_device,
random_seed=1234,
embedding_dim=entity_embeddings.shape[1],
margin_loss=1,
optimizer=torch.optim.Adam)
self.model.num_entities = num_entities
self.model.num_relations = num_relations
self.model.relation_embeddings = nn.Embedding.from_pretrained(
relation_embeddings) # noqa
self.model.entity_embeddings = nn.Embedding.from_pretrained(
entity_embeddings) # noqa
def forward(self, x: torch.tensor) -> torch.tensor:
"""
Compute the scores for given s, p and every o.
:param x: torch.tensor, dtype: torch.long, shape: (batch_size, 2)
The s, p indices.
:return:
scores: torch.tensor, dtype: torch.float, shape: (batch_size, num_entities) # noqa
The scores for every possible entity as o.
"""
batch_size = x.shape[0]
head_embeddings = self.model.entity_embeddings(x[:, 0:1]).view(
batch_size, 1, -1) # noqa
relation_embeddings = self.model.relation_embeddings(x[:, 1:2]).view(
batch_size, 1, -1) # noqa
num_entities, entity_embedding_size = self.model.entity_embeddings.weight.shape # noqa
tail_embeddings = self.model.entity_embeddings.weight.view(
1, num_entities, entity_embedding_size) # noqa
scores = -torch.sum(
head_embeddings * relation_embeddings * tail_embeddings,
dim=-1) # noqa
return scores
def test_instantiate_distmult(self):
"""Test that DistMult can be instantiated."""
distmult = DistMult(config=DISTMULT_CONFIG)
self.assertIsNotNone(distmult)
self.assertEqual(distmult.num_entities, 5)
self.assertEqual(distmult.num_relations, 5)
self.assertEqual(distmult.embedding_dim, 5)
self.assertEqual(distmult.margin_loss, 4)
def __init__(self,
relation_embeddings,
num_entities,
num_relations,
entity_embeddings,
preferred_device='cpu'):
super(DistMultUnit, self).__init__()
if preferred_device == 'cuda':
preferred_device = 'gpu'
self.model = DistMult(preferred_device=preferred_device,
random_seed=1234,
embedding_dim=entity_embeddings.shape[1],
margin_loss=1,
optimizer=torch.optim.Adam)
self.model.num_entities = num_entities
self.model.num_relations = num_relations
self.model.relation_embeddings = nn.Embedding.from_pretrained(
relation_embeddings) # noqa
self.model.entity_embeddings = nn.Embedding.from_pretrained(
entity_embeddings) # noqa
def get_model(num_entities, num_relations, embedding_size, margin_loss):
'''Initializes a DistMult model.
Parameters
----------
num_entities - int: The total number of distinct entities in the
dataset.
num_relations - int: The total number of distinct realtions in the
dataset.
'''
model = DistMult(preferred_device='gpu',
random_seed=1234,
embedding_dim=embedding_size,
margin_loss=margin_loss,
optimizer=torch.optim.Adam)
model.num_entities = num_entities
model.num_relations = num_relations
#pylint: disable=protected-access
model._init_embeddings()
return model