KGEModel

model

class neuralkg.model.KGEModel.model.Model(args)

Bases: torch.nn.modules.module.Module

init_emb()

score_func(head_emb, relation_emb, tail_emb)

forward(triples, negs, mode)

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

tri2emb(triples, negs=None, mode='single')

Get embedding of triples.

This function get the embeddings of head, relation, and tail respectively. each embedding has three dimensions.

Parameters

• triples (tensor) – This tensor save triples id, which dimension is [triples number, 3].

• negs (tensor, optional) – This tenosr store the id of the entity to be replaced, which has one dimension. when negs is None, it is in the test/eval phase. Defaults to None.

• mode (str, optional) – This arg indicates that the negative entity will replace the head or tail entity. when it is ‘single’, it means that entity will not be replaced. Defaults to ‘single’.

Returns

Head entity embedding. relation_emb: Relation embedding. tail_emb: Tail entity embedding.

Return type

head_emb

training: bool

TransE

class neuralkg.model.KGEModel.TransE.TransE(args)

Bases: neuralkg.model.KGEModel.model.Model

`Translating Embeddings for Modeling Multi-relational Data`_ (TransE), which represents the relationships as translations in the embedding space.

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_emb

Entity embedding, shape:[num_ent, emb_dim].

rel_emb

Relation embedding, shape:[num_rel, emb_dim].

init_emb()

Initialize the entity and relation embeddings in the form of a uniform distribution.

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is \(\gamma - ||h + r - t||_F\)

Parameters

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

• tail_emb – The tail entity embedding.

• mode – Choose head-predict or tail-predict.

Returns

The score of triples.

Return type

score

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns

The score of triples.

Return type

score

get_score(batch, mode)

The functions used in the testing phase

Parameters

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns

The score of triples.

Return type

score

training: bool

TransH

class neuralkg.model.KGEModel.TransH.TransH(args)

Bases: neuralkg.model.KGEModel.model.Model

Knowledge Graph Embedding by Translating on Hyperplanes (TransH), which apply the translation from head to tail entity in a relational-specific hyperplane in order to address its inability to model one-to-many, many-to-one, and many-to-many relations.

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_emb

Entity embedding, shape:[num_ent, emb_dim].

rel_emb

Relation embedding, shape:[num_rel, emb_dim].

norm_vector

Relation-specific projection matrix, shape:[num_rel, emb_dim]

init_emb()

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is \(\gamma - \|e'_{h,r} + d_r - e'_{t,r}\|_{p}^2\)

Parameters

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

• tail_emb – The tail entity embedding.

• mode – Choose head-predict or tail-predict.

Returns

The score of triples.

Return type

score

forward(triples, negs=None, mode='single')

The functions used in the training phase, same as TransE

get_score(batch, mode)

The functions used in the testing phase, same as TransE

training: bool

TransR

class neuralkg.model.KGEModel.TransR.TransR(args)

Bases: neuralkg.model.KGEModel.model.Model

Learning Entity and Relation Embeddings for Knowledge Graph Completion`_ (TransR), which building entity and relation embeddings in separate entity space and relation spaces

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_emb

Entity embedding, shape:[num_ent, emb_dim].

rel_emb

Relation embedding, shape:[num_rel, emb_dim].

transfer_matrix

Transfer entity and relation embedding, shape:[num_rel, emb_dim*emb_dim]

init_emb()

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is \(\gamma - \| M_{r} {e}_h + r_r - M_{r}e_t \|_{p}^2\)

forward(triples, negs=None, mode='single')

The functions used in the training phase, calculate triple score.

get_score(batch, mode)

The functions used in the testing phase, predict triple score.

training: bool

BoxE

class neuralkg.model.KGEModel.BoxE.BoxE(args)

Bases: neuralkg.model.KGEModel.model.Model

A Box Embedding Model for Knowledge Base Completion (BoxE), which represents the bump embedding as translations in the super rectangle space.

args

Model configuration parameters.

init_emb(args)

Initialize the entity and relation embeddings in the form of a uniform distribution.

Parameters

• arity – The maximum ary of the knowledge graph.

• epsilon – Caculate embedding_range.

• order – The distance order of score.

• margin – Caculate embedding_range and loss.

• embedding_range – Uniform distribution range.

• ent_emb – Entity embedding, shape:[num_ent, emb_dim * 2].

• rel_emb – Relation_embedding, shape:[num_rel, emb_dim * arity * 2].

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns

The score of triples.

Return type

score

get_score(batch, mode)

The functions used in the testing phase

Parameters

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns

The score of triples.

Return type

score

score_func(head_emb, relation_emb, tail_emb)

Calculate the score of the triple embedding.

Parameters

• head_emb – The embedding of head entity.

• relation_emb – The embedding of relation.

• tail_emb – The embedding of tail entity.

Returns

The score of triples.

Return type

score

calc_score(ent_emb, box_low, box_hig, order=2)

Calculate the norm of distance.

Parameters

• ent_emb – The embedding of entity.

• box_low – The lower boundaries of the super rectangle.

• box_hig – The upper boundaries of the super rectangle.

• order – The order of this distance.

Returns

The norm of distance.

dist(ent_emb, lb, ub)

Calculate the distance.

This function calculate the distance between the entity and the super rectangle. If the entity is in its target box, distance inversely correlates with box size, to maintain low distance inside large boxes and provide a gradient to keep points inside; if the entity is not in its target box, box size linearly correlates with distance, to penalize points outside larger boxes more severely.

Parameters

• ent_emb – The embedding of entity.

• lb – The lower boundaries of the super rectangle.

• ub – The upper boundaries of the super rectangle.

Returns

The distance between entity and super rectangle.

training: bool

ComplEx

class neuralkg.model.KGEModel.ComplEx.ComplEx(args)

Bases: neuralkg.model.KGEModel.model.Model

init_emb()

Initialize the entity and relation embeddings in the form of a uniform distribution.

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is :math:`operatorname{Re}left(h^{ op} operatorname{diag}(r) overline{t}

ight)`

Args:

head_emb: The head entity embedding. relation_emb: The relation embedding. tail_emb: The tail entity embedding. mode: Choose head-predict or tail-predict.

Returns:

score: The score of triples.

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns

The score of triples.

Return type

score

get_score(batch, mode)

The functions used in the testing phase

Parameters

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns

The score of triples.

Return type

score

training: bool

ConvE

class neuralkg.model.KGEModel.ConvE.ConvE(args)

Bases: neuralkg.model.KGEModel.model.Model

Convolutional 2D Knowledge Graph Embeddings (ConvE), which use a 2D convolution network for embedding representation.

args

Model configuration parameters.

init_emb(args)

Initialize the convolution layer and embeddings .

Parameters

• conv1 – The convolution layer.

• fc – The full connection layer.

• bn0 – The batch Normalization layer.

• bn1 – The batch Normalization layer.

• bn2 – The batch Normalization layer.

• inp_drop – The dropout layer.

• hid_drop – The dropout layer.

• feg_drop – The dropout layer.

• emb_ent – Entity embedding, shape:[num_ent, emb_dim].

• emb_rel – Relation_embedding, shape:[num_rel, emb_dim].

score_func(head_emb, relation_emb, choose_emb=None)

Calculate the score of the triple embedding.

This function calculate the score of the embedding. First, the entity and relation embeddings are reshaped and concatenated; the resulting matrix is then used as input to a convolutional layer; the resulting feature map tensor is vectorised and projected into a k-dimensional space.

Parameters

• head_emb – The embedding of head entity.

• relation_emb – The embedding of relation.

Returns

Final score of the embedding.

Return type

score

forward(triples)

The functions used in the training phase

Parameters

triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

Returns

The score of triples.

Return type

score

get_score(batch, mode='tail_predict')

The functions used in the testing phase

Parameters

batch – A batch of data.

Returns

The score of triples.

Return type

score

training: bool

CrossE

class neuralkg.model.KGEModel.CrossE.CrossE(args)

Bases: neuralkg.model.KGEModel.model.Model

Interaction Embeddings for Prediction and Explanation in Knowledge Graphs (CrossE), which simulates crossover interactions(bi-directional effects between entities and relations) to select related information when predicting a new triple

args

Model configuration parameters.

ent_emb

Entity embedding, shape:[num_ent, emb_dim].

rel_emb

Relation embedding, shape:[num_rel, emb_dim].

rel_reverse_emb

Reverse Relation embedding, shape:[num_rel, emb_dim].

h_weighted_vector

Interaction matrix for head entities and relations, shape:[num_rel, emb_dim]

t_weighted_vector

Interaction matrix for tail entities and relations, shape:[num_rel, emb_dim]

hr_bias

Bias for head entity and relation

tr_bias

Bias for tail entity and relation

init_emb()

score_func(ent_emb, rel_emb, weighted_vector, mode)

Calculating the score of triples.

The formula for calculating the score is \(\sigma(tanh(c_r \circ h + c_r \circ h \circ r + b)t ^ T)\)

Parameters

• ent_emb – entity embedding

• rel_emb – relation embedding

• weighted_vector – Interaction matrix for entities and relations

• mode – Choose head-predict or tail-predict.

Returns

The score of triples.

Return type

score

forward(triples, mode='single')

The functions used in the training phase, calculate hr_score and tr_score simultaneously

get_score(batch, mode)

The functions used in the testing phase, predict triple score

regularize_loss(norm=2)

Add regularization to loss

training: bool

DistMult

class neuralkg.model.KGEModel.DistMult.DistMult(args)

Bases: neuralkg.model.KGEModel.model.Model

Embedding Entities and Relations for Learning and Inference in Knowledge Bases (DistMult)

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_emb

Entity embedding, shape:[num_ent, emb_dim].

rel_emb

Relation_embedding, shape:[num_rel, emb_dim].

init_emb()

Initialize the entity and relation embeddings in the form of a uniform distribution.

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is \(h^{ op} \operatorname{diag}(r) t\)

Parameters

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

• tail_emb – The tail entity embedding.

• mode – Choose head-predict or tail-predict.

Returns

The score of triples.

Return type

score

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns

The score of triples.

Return type

score

get_score(batch, mode)

The functions used in the testing phase

Parameters

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns

The score of triples.

Return type

score

training: bool

RotatE

class neuralkg.model.KGEModel.RotatE.RotatE(args)

Bases: neuralkg.model.KGEModel.model.Model

`RotatE: Knowledge Graph Embedding by Relational Rotation in Complex Space`_ (RotatE), which defines each relation as a rotation from the source entity to the target entity in the complex vector space.

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_emb

Entity embedding, shape:[num_ent, emb_dim * 2].

rel_emb

Relation_embedding, shape:[num_rel, emb_dim].

init_emb()

Initialize the entity and relation embeddings in the form of a uniform distribution.

score_func(head_emb, relation_emb, tail_emb, mode)

Calculating the score of triples.

The formula for calculating the score is \(\gamma - \|h \circ r - t\|\)

Parameters

• head_emb – The head entity embedding.

• relation_emb – The relation embedding.

• tail_emb – The tail entity embedding.

• mode – Choose head-predict or tail-predict.

Returns

The score of triples.

Return type

score

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns

The score of triples.

Return type

score

get_score(batch, mode)

The functions used in the testing phase

Parameters

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns

The score of triples.

Return type

score

training: bool

SimplE

class neuralkg.model.KGEModel.SimplE.SimplE(args)

Bases: neuralkg.model.KGEModel.model.Model

SimplE Embedding for Link Prediction in Knowledge Graphs (SimpleE), which presents a simple enhancement of CP (which we call SimplE) to allow the two embeddings of each entity to be learned dependently.

args

Model configuration parameters.

epsilon

Calculate embedding_range.

margin

Calculate embedding_range and loss.

embedding_range

Uniform distribution range.

ent_h_emb

Entity embedding, shape:[num_ent, emb_dim].

ent_t_emb

Entity embedding, shape:[num_ent, emb_dim].

rel_emb

Relation_embedding, shape:[num_rel, emb_dim].

rel_inv_emb

Inverse Relation_embedding, shape:[num_rel, emb_dim].

init_emb()

Initialize the entity and relation embeddings in the form of a uniform distribution.

score_func(hh_emb, rel_emb, tt_emb, ht_emb, rel_inv_emb, th_emb)

Calculating the score of triples.

Parameters

• hh_emb – The head entity embedding on embedding h.

• rel_emb – The relation embedding.

• tt_emb – The tail entity embedding on embedding t.

• ht_emb – The tail entity embedding on embedding h.

• rel_inv_emb – The tail entity embedding.

• th_emb – The head entity embedding on embedding t.

Returns

The score of triples.

Return type

score

l2_loss()

forward(triples, negs=None, mode='single')

The functions used in the training phase

Parameters

• triples – The triples ids, as (h, r, t), shape:[batch_size, 3].

• negs – Negative samples, defaults to None.

• mode – Choose head-predict or tail-predict, Defaults to ‘single’.

Returns

The score of triples.

Return type

score

get_score(batch, mode)

The functions used in the testing phase

Parameters

• batch – A batch of data.

• mode – Choose head-predict or tail-predict.

Returns

The score of triples.

Return type

score

get_emb(triples, negs=None, mode='single')

training: bool