Models

Interfaces

Model

class torchkge.models.interfaces.Model(n_entities, n_relations)[source]

Model interface to be used by any other class implementing a knowledge graph embedding model. It is only required to implement the methods scoring_function, normalize_parameters, inference_prepare_candidates and inference_scoring_function.

Parameters

n_entities (int) – Number of entities to be embedded.
n_relations (int) – Number of relations to be embedded.

n_ent

Number of entities to be embedded.

Type: int

n_rel

Number of relations to be embedded.

Type: int

forward(heads, tails, relations, negative_heads, negative_tails, negative_relations=None)[source]

Parameters

heads (torch.Tensor, dtype: torch.long, shape: (batch_size)) – Integer keys of the current batch’s heads
tails (torch.Tensor, dtype: torch.long, shape: (batch_size)) – Integer keys of the current batch’s tails.
relations (torch.Tensor, dtype: torch.long, shape: (batch_size)) – Integer keys of the current batch’s relations.
negative_heads (torch.Tensor, dtype: torch.long, shape: (batch_size)) – Integer keys of the current batch’s negatively sampled heads.
negative_tails (torch.Tensor, dtype: torch.long, shape: (batch_size)) – Integer keys of the current batch’s negatively sampled tails.ze)
negative_relations (torch.Tensor, dtype: torch.long, shape: (batch_size)) – Integer keys of the current batch’s negatively sampled relations.

Returns

positive_triplets (torch.Tensor, dtype: torch.float, shape: (b_size)) – Scoring function evaluated on true triples.
negative_triplets (torch.Tensor, dtype: torch.float, shape: (b_size)) – Scoring function evaluated on negatively sampled triples.

get_embeddings()[source]: Return the tensors representing entities and relations in current model.

inference_prepare_candidates(h_idx, t_idx, r_idx, entities=True)[source]

Link prediction evaluation helper function. Get entities and relations embeddings, along with entity candidates ready for (projected if needed). The output will be fed to the inference_scoring_function method of the model at hand.

Parameters

h_idx (torch.Tensor, shape: (b_size), dtype: torch.long) – List of heads indices.
t_idx (torch.Tensor, shape: (b_size), dtype: torch.long) – List of tails indices.
r_idx (torch.Tensor, shape: (b_size), dtype: torch.long) – List of relations indices.
entities (bool) – Boolean indicating if candidates are entities or not.

Returns

h (torch.Tensor, shape: (b_size, rel_emb_dim), dtype: torch.float) – Head vectors fed to inference_scoring_function. For translation models it is the entities embeddings projected in relation space, for example.
t (torch.Tensor, shape: (b_size, rel_emb_dim), dtype: torch.float) – Tail vectors fed to inference_scoring_function. For translation models it is the entities embeddings projected in relation space, for example.
candidates (torch.Tensor, shape: (b_size, rel_emb_dim, n_ent),) – dtype: torch.float All entities embeddings prepared from batch evaluation. Axis 0 is simply duplication.
r (torch.Tensor, shape: (b_size, rel_emb_dim), dtype: torch.float) – Relations embeddings or matrices.

inference_scoring_function(h, t, r)[source]

Link prediction evaluation helper function. Compute the scores of (h, r, c) or (c, r, t) for any candidate c. The arguments should match the ones of inference_prepare_candidates.

Parameters

h (torch.Tensor, shape: (b_size, ent_emb_dim) or (b_size, n_ent,) – ent_emb_dim), dtype: torch.float
t (torch.Tensor, shape: (b_size, ent_emb_dim) or (b_size, n_ent,) – ent_emb_dim), dtype: torch.float
r (torch.Tensor, shape: (b_size, ent_emb_dim) or (b_size, n_rel,) – ent_emb_dim), dtype: torch.float

Returns

scores – Scores of each candidate for each triple.

Return type

torch.Tensor, shape: (b_size, n_ent), dtype: torch.float

normalize_parameters()[source]: Normalize some parameters. This methods should be end at the end of each training epoch and at the end of training as well.

scoring_function(h_idx, t_idx, r_idx)[source]

Compute the scoring function for the triplets given as argument.

Parameters

h_idx (torch.Tensor, dtype: torch.long, shape: (b_size)) – Integer keys of the current batch’s heads
t_idx (torch.Tensor, dtype: torch.long, shape: (b_size)) – Integer keys of the current batch’s tails.
r_idx (torch.Tensor, dtype: torch.long, shape: (b_size)) – Integer keys of the current batch’s relations.

Returns

score – Score of each triplet.

Return type

torch.Tensor, dtype: torch.float, shape: (b_size)

TranslationalModels

class torchkge.models.interfaces.TranslationModel(n_entities, n_relations, dissimilarity_type)[source]

Model interface to be used by any other class implementing a translation knowledge graph embedding model. This interface inherits from the interface torchkge.models.interfaces.Model. It is only required to implement the methods scoring_function, normalize_parameters and inference_prepare_candidates.

Parameters

n_entities (int) – Number of entities to be embedded.
n_relations (int) – Number of relations to be embedded.
dissimilarity_type (str) – One of ‘L1’, ‘L2’, ‘toruse_L1’, ‘toruse_L2’ and ‘toruse_eL2’.

dissimilarity

Dissimilarity function.

Type: function

get_embeddings()[source]: See torchkge.models.interfaces.Models.

inference_prepare_candidates(h_idx, t_idx, r_idx, entities=True)[source]: See torchkge.models.interfaces.Models.

inference_scoring_function(proj_h, proj_t, r)[source]: This overwrites the method declared in torchkge.models.interfaces.Models. For translation models, the computed score is the dissimilarity of between projected heads + relations and projected tails. Projections are done in relation-specific subspaces.

normalize_parameters()[source]: See torchkge.models.interfaces.Models.

scoring_function(h_idx, t_idx, r_idx)[source]: See torchkge.models.interfaces.Models.

Translational Models

Parameters used to train models available in pre-trained version :

	Dataset	Dimension	Optimizer	Learning Rate	Batch Size	Loss	Margin	L2 penalization
TransE	FB15k	100	Adam	2.1e-5	32768	Margin	.651	1e-5
TransE	FB15k237	100	Adam	2.1e-5	32768	Margin	.651	1e-5
TransE	FB15k237	150	Adam	2.7e-5	32768	Margin	.648	1e-5

TransE

class torchkge.models.translation.TransEModel(emb_dim, n_entities, n_relations, dissimilarity_type='L2')[source]

Implementation of TransE model detailed in 2013 paper by Bordes et al.. This class inherits from the torchkge.models.interfaces.TranslationModel interface. It then has its attributes as well.

References

Antoine Bordes, Nicolas Usunier, Alberto Garcia-Duran, Jason Weston, and Oksana Yakhnenko. Translating Embeddings for Modeling Multi-relational Data. In Advances in Neural Information Processing Systems 26, pages 2787–2795, 2013.

Parameters

emb_dim (int) – Dimension of the embedding.
n_entities (int) – Number of entities in the current data set.
n_relations (int) – Number of relations in the current data set.
dissimilarity_type (str) – Either ‘L1’ or ‘L2’.

emb_dim

Dimension of the embedding.

Type: int

ent_emb

Embeddings of the entities, initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_ent, emb_dim)

rel_emb

Embeddings of the relations, initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_rel, emb_dim)

get_embeddings()[source]

Return the embeddings of entities and relations.

Returns

ent_emb (torch.Tensor, shape: (n_ent, emb_dim), dtype: torch.float) – Embeddings of entities.
rel_emb (torch.Tensor, shape: (n_rel, emb_dim), dtype: torch.float) – Embeddings of relations.

inference_prepare_candidates(h_idx, t_idx, r_idx, entities=True)[source]: Link prediction evaluation helper function. Get entities embeddings and relations embeddings. The output will be fed to the inference_scoring_function method. See torchkge.models.interfaces.Models for more details on the API.

normalize_parameters()[source]: Normalize the entity embeddings, as explained in original paper. This method should be called at the end of each training epoch and at the end of training as well.

scoring_function(h_idx, t_idx, r_idx)[source]: Compute the scoring function for the triplets given as argument: \(||h + r - t||_p^p\) with p being the dissimilarity type (either 1 or 2). See referenced paper for more details on the score. See torchkge.models.interfaces.Models for more details on the API.

TransH

class torchkge.models.translation.TransHModel(emb_dim, n_entities, n_relations)[source]

Implementation of TransH model detailed in 2014 paper by Wang et al.. This class inherits from the torchkge.models.interfaces.TranslationModel interface. It then has its attributes as well.

References

Zhen Wang, Jianwen Zhang, Jianlin Feng, and Zheng Chen. Knowledge Graph Embedding by Translating on Hyperplanes. In Twenty-Eighth AAAI Conference on Artificial Intelligence, June 2014.

Parameters

emb_dim (int) – Dimension of the embedding space.
n_entities (int) – Number of entities in the current data set.
n_relations (int) – Number of relations in the current data set.

emb_dim

Dimension of the embedding.

Type: int

ent_emb

Embeddings of the entities, initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_ent, emb_dim)

rel_emb

Embeddings of the relations, initialized with Xavier uniform distribution.

Type: torch.nn.Embedding, shape: (n_rel, emb_dim)

norm_vect

Normal vectors associated to each relation and used to compute the relation-specific hyperplanes entities are projected on. See paper for more details. Initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_rel, emb_dim)

evaluate_projections()[source]: Link prediction evaluation helper function. Project all entities according to each relation. Calling this method at the beginning of link prediction makes the process faster by computing projections only once.

get_embeddings()[source]

Return the embeddings of entities and relations along with relation normal vectors.

Returns

ent_emb (torch.Tensor, shape: (n_ent, emb_dim), dtype: torch.float) – Embeddings of entities.
rel_emb (torch.Tensor, shape: (n_rel, emb_dim), dtype: torch.float) – Embeddings of relations.
norm_vect (torch.Tensor, shape: (n_rel, emb_dim), dtype: torch.float) – Normal vectors defining relation-specific hyperplanes.

inference_prepare_candidates(h_idx, t_idx, r_idx, entities=True)[source]: Link prediction evaluation helper function. Get entities embeddings and relations embeddings. The output will be fed to the inference_scoring_function method. See torchkge.models.interfaces.Models for more details on the API.

normalize_parameters()[source]: Normalize the entity embeddings and relations normal vectors, as explained in original paper. This methods should be called at the end of each training epoch and at the end of training as well.

scoring_function(h_idx, t_idx, r_idx)[source]: Compute the scoring function for the triplets given as argument: \(||p_r(h) + r - p_r(t)||_2^2\). See referenced paper for more details on the score. See torchkge.models.interfaces.Models for more details on the API.

TransR

class torchkge.models.translation.TransRModel(ent_emb_dim, rel_emb_dim, n_entities, n_relations)[source]

Implementation of TransR model detailed in 2015 paper by Lin et al.. This class inherits from the torchkge.models.interfaces.TranslationModel interface. It then has its attributes as well.

References

Yankai Lin, Zhiyuan Liu, Maosong Sun, Yang Liu, and Xuan Zhu. Learning Entity and Relation Embeddings for Knowledge Graph Completion. In Twenty-Ninth AAAI Conference on Artificial Intelligence, February 2015

Parameters

ent_emb_dim (int) – Dimension of the embedding of entities.
rel_emb_dim (int) – Dimension of the embedding of relations.
n_entities (int) – Number of entities in the current data set.
n_relations (int) – Number of relations in the current data set.

ent_emb_dim

Dimension nof the embedding of entities.

Type: int

rel_emb_dim

Dimension of the embedding of relations.

Type: int

ent_emb

Embeddings of the entities, initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_ent, ent_emb_dim)

rel_emb

Embeddings of the relations, initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_rel, rel_emb_dim)

proj_mat

Relation-specific projection matrices. See paper for more details.

Type: torch.nn.Embedding, shape: (n_rel, rel_emb_dim x ent_emb_dim)

projected_entities

Contains the projection of each entity in each relation-specific sub-space.

Type: torch.nn.Parameter, shape: (n_rel, n_ent, rel_emb_dim)

evaluated_projections

Indicates whether projected_entities has been computed. This should be set to true every time a backward pass is done in train mode.

Type: bool

evaluate_projectionss()[source]: Link prediction evaluation helper function. Project all entities according to each relation. Calling this method at the beginning of link prediction makes the process faster by computing projections only once.

get_embeddings()[source]

Return the embeddings of entities and relations along with their projection matrices.

Returns

ent_emb (torch.Tensor, shape: (n_ent, ent_emb_dim), dtype: torch.float) – Embeddings of entities.
rel_emb (torch.Tensor, shape: (n_rel, rel_emb_dim), dtype: torch.float) – Embeddings of relations.
proj_mat (torch.Tensor, shape: (n_rel, rel_emb_dim, ent_emb_dim),)
dtype (torch.float) – Relation-specific projection matrices.

inference_prepare_candidates(h_idx, t_idx, r_idx, entities=True)[source]: Link prediction evaluation helper function. Get entities embeddings and relations embeddings. The output will be fed to the inference_scoring_function method. See torchkge.models.interfaces.Models for more details on the API.

normalize_parameters()[source]: Normalize the entity and relation embeddings, as explained in original paper. This methods should be called at the end of each training epoch and at the end of training as well.

scoring_function(h_idx, t_idx, r_idx)[source]: Compute the scoring function for the triplets given as argument: \(||p_r(h) + r - p_r(t)||_2^2\). See referenced paper for more details on the score. See torchkge.models.interfaces.Models for more details on the API.

TransD

class torchkge.models.translation.TransDModel(ent_emb_dim, rel_emb_dim, n_entities, n_relations)[source]

Implementation of TransD model detailed in 2015 paper by Ji et al.. This class inherits from the torchkge.models.interfaces.TranslationModel interface. It then has its attributes as well.

References

Guoliang Ji, Shizhu He, Liheng Xu, Kang Liu, and Jun Zhao. Knowledge Graph Embedding via Dynamic Mapping Matrix. In Proceedings of the 53rd Annual Meeting of the Association for Computational Linguistics and the 7th International Joint Conference on Natural Language Processing (Volume 1: Long Papers) pages 687–696, Beijing, China, July 2015. Association for Computational Linguistics.

Parameters

ent_emb_dim (int) – Dimension of the embedding of entities.
rel_emb_dim (int) – Dimension of the embedding of relations.
n_entities (int) – Number of entities in the current data set.
n_relations (int) – Number of relations in the current data set.

ent_emb_dim

Dimension of the embedding of entities.

Type: int

rel_emb_dim

Dimension of the embedding of relations.

Type: int

ent_emb

Embeddings of the entities, initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_ent, ent_emb_dim)

rel_emb

Embeddings of the relations, initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_rel, rel_emb_dim)

ent_proj_vect

Entity-specific vector used to build projection matrices. See paper for more details. Initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_ent, ent_emb_dim)

rel_proj_vect

Relation-specific vector used to build projection matrices. See paper for more details. Initialized with Xavier uniform distribution and then normalized.

Type: torch..nn.Embedding, shape: (n_rel, rel_emb_dim)

projected_entities

Contains the projection of each entity in each relation-specific sub-space.

Type: torch.nn.Parameter, shape: (n_rel, n_ent, rel_emb_dim)

evaluated_projections

Indicates whether projected_entities has been computed. This should be set to true every time a backward pass is done in train mode.

Type: bool

evaluate_projectionss()[source]: Link prediction evaluation helper function. Project all entities according to each relation. Calling this method at the beginning of link prediction makes the process faster by computing projections only once.

get_embeddings()[source]

Return the embeddings of entities and relations along with their projection vectors.

Returns

ent_emb (torch.Tensor, shape: (n_ent, ent_emb_dim), dtype: torch.float) – Embeddings of entities.
rel_emb (torch.Tensor, shape: (n_rel, rel_emb_dim), dtype: torch.float) – Embeddings of relations.
ent_proj_vect (torch.Tensor, shape: (n_ent, ent_emb_dim),)
dtype (torch.float) – Entity projection vectors.
rel_proj_vect (torch.Tensor, shape: (n_ent, rel_emb_dim),)
dtype (torch.float) – Relation projection vectors.

inference_prepare_candidates(h_idx, t_idx, r_idx, entities=True)[source]: Link prediction evaluation helper function. Get entities embeddings and relations embeddings. The output will be fed to the inference_scoring_function method. See torchkge.models.interfaces.Models for more details on the API.

normalize_parameters()[source]: Normalize the entity embeddings and relations normal vectors, as explained in original paper. This methods should be called at the end of each training epoch and at the end of training as well.

project(ent, e_proj_vect, r_proj_vect)[source]: We note that \(p_r(e)_i = e^p^Te \times r^p_i + e_i\) which is more efficient to compute than the matrix formulation in the original paper.

scoring_function(h_idx, t_idx, r_idx)[source]: Compute the scoring function for the triplets given as argument: \(||p_r(h) + r - p_r(t)||_2^2\). See referenced paper for more details on the score. See torchkge.models.interfaces.Models for more details on the API.

TorusE

class torchkge.models.translation.TorusEModel(emb_dim, n_entities, n_relations, dissimilarity_type)[source]

Implementation of TorusE model detailed in 2018 paper by Ebisu and Ichise. This class inherits from the torchkge.models.interfaces.TranslationModel interface. It then has its attributes as well.

References

Takuma Ebisu and Ryutaro Ichise TorusE: Knowledge Graph Embedding on a Lie Group. In Proceedings of the 32nd AAAI Conference on Artificial Intelligence (New Orleans, LA, USA, Feb. 2018), AAAI Press, pp. 1819–1826.

Parameters

emb_dim (int) – Embedding dimension.
n_entities (int) – Number of entities in the current data set.
n_relations (int) – Number of relations in the current data set.
dissimilarity_type (str) – One of ‘torus_L1’, ‘torus_L2’, ‘torus_eL2’.

emb_dim

Embedding dimension.

Type: int

ent_emb

Embeddings of the entities, initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_ent, emb_dim)

rel_emb

Embeddings of the relations, initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_rel, emb_dim)

get_embeddings()[source]

Return the embeddings of entities and relations.

Returns

ent_emb (torch.Tensor, shape: (n_ent, emb_dim), dtype: torch.float) – Embeddings of entities.
rel_emb (torch.Tensor, shape: (n_rel, emb_dim), dtype: torch.float) – Embeddings of relations.

inference_prepare_candidates(h_idx, t_idx, r_idx, entities=True)[source]: Link prediction evaluation helper function. Get entities embeddings and relations embeddings. The output will be fed to the inference_scoring_function method. See torchkge.models.interfaces.Models for more details on the API.

normalize_parameters()[source]: Project embeddings on torus.

scoring_function(h_idx, t_idx, r_idx)[source]: Compute the scoring function for the triplets given as argument: See referenced paper for more details on the score. See torchkge.models.interfaces.Models for more details on the API.

Bilinear Models

RESCAL

class torchkge.models.bilinear.RESCALModel(emb_dim, n_entities, n_relations)[source]

Implementation of RESCAL model detailed in 2011 paper by Nickel et al.. In the original paper, optimization is done using Alternating Least Squares (ALS). Here we use iterative gradient descent optimization. This class inherits from the torchkge.models.interfaces.BilinearModel interface. It then has its attributes as well.

References

Maximilian Nickel, Volker Tresp, and Hans-Peter Kriegel. A Three-way Model for Collective Learning on Multi-relational Data. In Proceedings of the 28th International Conference on Machine Learning, 2011.

Parameters

emb_dim (int) – Dimension of embedding space.
n_entities (int) – Number of entities in the current data set.
n_relations (int) – Number of relations in the current data set.

ent_emb

Embeddings of the entities, initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_ent, emb_dim)

rel_mat

Matrices of the relations, initialized with Xavier uniform distribution.

Type: torch.nn.Embedding, shape: (n_rel, emb_dim x emb_dim)

get_embeddings()[source]

Return the embeddings of entities and matrices of relations.

Returns

ent_emb (torch.Tensor, shape: (n_ent, emb_dim), dtype: torch.float) – Embeddings of entities.
rel_mat (torch.Tensor, shape: (n_rel, emb_dim, emb_dim),)
dtype (torch.float) – Matrices of relations.

inference_prepare_candidates(h_idx, t_idx, r_idx, entities=True)[source]: Link prediction evaluation helper function. Get entities embeddings and relations embeddings. The output will be fed to the inference_scoring_function method. See torchkge.models.interfaces.Models for more details on the API.

inference_scoring_function(h, t, r)[source]: Link prediction evaluation helper function. See torchkge.models.interfaces.Models for more details of the API.

normalize_parameters()[source]: Normalize the entity embeddings, as explained in original paper. This methods should be called at the end of each training epoch and at the end of training as well.

scoring_function(h_idx, t_idx, r_idx)[source]: Compute the scoring function for the triplets given as argument: \(h^T \cdot M_r \cdot t\). See referenced paper for more details on the score. See torchkge.models.interfaces.Models for more details on the API.

DistMult

class torchkge.models.bilinear.DistMultModel(emb_dim, n_entities, n_relations)[source]

Implementation of DistMult model detailed in 2014 paper by Yang et al.. This class inherits from the torchkge.models.interfaces.BilinearModel interface. It then has its attributes as well.

References

Bishan Yang, Wen-tau Yih, Xiaodong He, Jianfeng Gao, and Li Deng. Embedding Entities and Relations for Learning and Inference in Knowledge Bases. arXiv :1412.6575 [cs], December 2014.

Parameters

emb_dim (int) – Dimension of embedding space.
n_entities (int) – Number of entities in the current data set.
n_relations (int) – Number of relations in the current data set.

ent_emb

Embeddings of the entities, initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_ent, emb_dim)

rel_emb

Embeddings of the relations, initialized with Xavier uniform distribution.

Type: torch.nn.Embedding, shape: (n_rel, emb_dim)

get_embeddings()[source]

Return the embeddings of entities and relations.

Returns

ent_emb (torch.Tensor, shape: (n_ent, emb_dim), dtype: torch.float) – Embeddings of entities.
rel_emb (torch.Tensor, shape: (n_rel, emb_dim), dtype: torch.float) – Embeddings of relations.

inference_prepare_candidates(h_idx, t_idx, r_idx, entities=True)[source]: Link prediction evaluation helper function. Get entities embeddings and relations embeddings. The output will be fed to the inference_scoring_function method. See torchkge.models.interfaces.Models for more details on the API.

inference_scoring_function(h, t, r)[source]: Link prediction evaluation helper function. See torchkge.models.interfaces.Models for more details on the API.

normalize_parameters()[source]: Normalize the entity embeddings, as explained in original paper. This methods should be called at the end of each training epoch and at the end of training as well.

scoring_function(h_idx, t_idx, r_idx)[source]: Compute the scoring function for the triplets given as argument: \(h^T \cdot diag(r) \cdot t\). See referenced paper for more details on the score. See torchkge.models.interfaces.Models for more details on the API.

HolE

class torchkge.models.bilinear.HolEModel(emb_dim, n_entities, n_relations)[source]

Implementation of HolE model detailed in 2015 paper by Nickel et al.. This class inherits from the torchkge.models.interfaces.BilinearModel interface. It then has its attributes as well.

References

Maximilian Nickel, Lorenzo Rosasco, and Tomaso Poggio. Holographic Embeddings of Knowledge Graphs. arXiv :1510.04935 [cs, stat], October 2015.

Parameters

emb_dim (int) – Dimension of embedding space.
n_entities (int) – Number of entities in the current data set.
n_relations (int) – Number of relations in the current data set.

ent_emb

Embeddings of the entities, initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_ent, emb_dim)

rel_emb

Embeddings of the relations, initialized with Xavier uniform distribution.

Type: torch.nn.Embedding, shape: (n_rel, emb_dim)

get_embeddings()[source]

Return the embeddings of entities and relations.

Returns

ent_emb (torch.Tensor, shape: (n_ent, emb_dim), dtype: torch.float) – Embeddings of entities.
rel_emb (torch.Tensor, shape: (n_rel, emb_dim), dtype: torch.float) – Embeddings of relations.

static get_rolling_matrix(x)[source]

Build a rolling matrix.

Parameters: x (torch.Tensor, shape: (b_size, dim)) –
Returns: mat – Rolling matrix such that mat[i,j] = x[j - i mod(dim)]
Return type: torch.Tensor, shape: (b_size, dim, dim)

inference_prepare_candidates(h_idx, t_idx, r_idx, entities=True)[source]: Link prediction evaluation helper function. Get entities embeddings and relations embeddings. The output will be fed to the inference_scoring_function method. See torchkge.models.interfaces.Models for more details on the API.

inference_scoring_function(h, t, r)[source]: Link prediction evaluation helper function. See torchkge.models.interfaces.Models for more details on the API.

normalize_parameters()[source]: Normalize the entity embeddings, as explained in original paper. This methods should be called at the end of each training epoch and at the end of training as well.

scoring_function(h_idx, t_idx, r_idx)[source]: Compute the scoring function for the triplets given as argument: \(h^T \cdot M_r \cdot t\) where \(M_r\) is the rolling matrix built from the relation embedding r. See referenced paper for more details on the score. See torchkge.models.interfaces.Models for more details on the API.

ComplEx

class torchkge.models.bilinear.ComplExModel(emb_dim, n_entities, n_relations)[source]

Implementation of ComplEx model detailed in 2016 paper by Trouillon et al.. This class inherits from the torchkge.models.interfaces.BilinearModel interface. It then has its attributes as well.

References

Théo Trouillon, Johannes Welbl, Sebastian Riedel, Éric Gaussier, and Guillaume Bouchard. Complex Embeddings for Simple Link Prediction. arXiv :1606.06357 [cs, stat], June 2016.

Parameters

emb_dim (int) – Dimension of embedding space.
n_entities (int) – Number of entities in the current data set.
n_relations (int) – Number of relations in the current data set.

re_ent_emb

Real part of the entities complex embeddings. Initialized with Xavier uniform distribution.

Type: torch.nn.Embedding, shape: (n_ent, emb_dim)

im_ent_emb

Imaginary part of the entities complex embeddings. Initialized with Xavier uniform distribution.

Type: torch.nn.Embedding, shape: (n_ent, emb_dim)

re_rel_emb

Real part of the relations complex embeddings. Initialized with Xavier uniform distribution.

Type: torch.nn.Embedding, shape: (n_rel, emb_dim)

im_rel_emb

Imaginary part of the relations complex embeddings. Initialized with Xavier uniform distribution.

Type: torch.nn.Embedding, shape: (n_rel, emb_dim)

get_embeddings()[source]

Return the embeddings of entities and relations.

Returns

re_ent_emb (torch.Tensor, shape: (n_ent, emb_dim), dtype: torch.float) – Real part of embeddings of entities.
im_ent_emb (torch.Tensor, shape: (n_ent, emb_dim), dtype: torch.float) – Imaginary part of embeddings of entities.
re_rel_emb (torch.Tensor, shape: (n_rel, emb_dim), dtype: torch.float) – Real part of embeddings of relations.
im_rel_emb (torch.Tensor, shape: (n_rel, emb_dim), dtype: torch.float) – Imaginary part of embeddings of relations.

inference_prepare_candidates(h_idx, t_idx, r_idx, entities=True)[source]: Link prediction evaluation helper function. Get entities embeddings and relations embeddings. The output will be fed to the inference_scoring_function method. See torchkge.models.interfaces.Models for more details on the API.

inference_scoring_function(h, t, r)[source]: Link prediction evaluation helper function. See torchkge.models.interfaces.Models for more details one the API.

normalize_parameters()[source]: According to original paper, the embeddings should not be normalized.

scoring_function(h_idx, t_idx, r_idx)[source]: Compute the real part of the Hermitian product \(\Re(h^T \cdot diag(r) \cdot \bar{t})\) for each sample of the batch. See referenced paper for more details on the score. See torchkge.models.interfaces.Models for more details on the API.

ANALOGY

class torchkge.models.bilinear.AnalogyModel(emb_dim, n_entities, n_relations, scalar_share=0.5)[source]

Implementation of ANALOGY model detailed in 2017 paper by Liu et al.. According to their remark in the implementation details, the number of scalars on the diagonal of each relation-specific matrix is by default set to be half the embedding dimension. This class inherits from the torchkge.models.interfaces.BilinearModel interface. It then has its attributes as well.

References

Hanxiao Liu, Yuexin Wu, and Yiming Yang. Analogical Inference for Multi-Relational Embeddings. arXiv :1705.02426 [cs], May 2017.

Parameters

emb_dim (int) – Dimension of embedding space.
n_entities (int) – Number of entities in the current data set.
n_relations (int) – Number of relations in the current data set.
scalar_share (float) – Share of the diagonal elements of the relation-specific matrices to be scalars. By default it is set to half according to the original paper.

scalar_dim

Number of diagonal elements of the relation-specific matrices to be scalars. By default it is set to half the embedding dimension according to the original paper.

Type: int

complex_dim

Number of 2x2 matrices on the diagonals of relation-specific matrices.

Type: int

sc_ent_emb

Part of the entities embeddings associated to the scalar part of the relation specific matrices. Initialized with Xavier uniform distribution.

Type: torch.nn.Embedding, shape: (n_ent, emb_dim)

re_ent_emb

Real part of the entities complex embeddings. Initialized with Xavier uniform distribution. As explained in the authors’ paper, almost diagonal matrices can be seen as complex matrices.

Type: torch.nn.Embedding, shape: (n_ent, emb_dim)

im_ent_emb

Imaginary part of the entities complex embeddings. Initialized with Xavier uniform distribution. As explained in the authors’ paper, almost diagonal matrices can be seen as complex matrices.

Type: torch.nn.Embedding, shape: (n_ent, emb_dim)

sc_rel_emb

Part of the entities embeddings associated to the scalar part of the relation specific matrices. Initialized with Xavier uniform distribution.

Type: torch.nn.Embedding, shape: (n_rel, emb_dim)

re_rel_emb

Real part of the relations complex embeddings. Initialized with Xavier uniform distribution. As explained in the authors’ paper, almost diagonal matrices can be seen as complex matrices.

Type: torch.nn.Embedding, shape: (n_rel, emb_dim)

im_rel_emb

Imaginary part of the relations complex embeddings. Initialized with Xavier uniform distribution. As explained in the authors’ paper, almost diagonal matrices can be seen as complex matrices.

Type: torch.nn.Embedding, shape: (n_rel, emb_dim)

get_embeddings()[source]

Return the embeddings of entities and relations.

Returns

sc_ent_emb (torch.Tensor, shape: (n_ent, emb_dim), dtype: torch.float) – Scalar part of embeddings of entities.
re_ent_emb (torch.Tensor, shape: (n_ent, emb_dim), dtype: torch.float) – Real part of embeddings of entities.
im_ent_emb (torch.Tensor, shape: (n_ent, emb_dim), dtype: torch.float) – Imaginary part of embeddings of entities.
sc_rel_emb (torch.Tensor, shape: (n_rel, emb_dim), dtype: torch.float) – Scalar part of embeddings of relations.
re_rel_emb (torch.Tensor, shape: (n_rel, emb_dim), dtype: torch.float) – Real part of embeddings of relations.
im_rel_emb (torch.Tensor, shape: (n_rel, emb_dim), dtype: torch.float) – Imaginary part of embeddings of relations.

inference_prepare_candidates(h_idx, t_idx, r_idx, entities=True)[source]: Link prediction evaluation helper function. Get entities embeddings and relations embeddings. The output will be fed to the inference_scoring_function method. See torchkge.models.interfaces.Models for more details on the API.

inference_scoring_function(h, t, r)[source]: Link prediction evaluation helper function. See torchkge.models.interfaces.Models for more details one the API.

normalize_parameters()[source]: According to original paper, the embeddings should not be normalized.

scoring_function(h_idx, t_idx, r_idx)[source]: Compute the scoring function for the triplets given as argument: \(h_{sc}^T \cdot diag(r_{sc}) \cdot t_{sc} + \Re(h_{compl} \cdot diag(r_{compl} \cdot t_{compl}))\). See referenced paper for more details on the score. See torchkge.models.interfaces.Models for more details on the API.

Deep Models

ConvKB

class torchkge.models.deep.ConvKBModel(emb_dim, n_filters, n_entities, n_relations)[source]

Implementation of ConvKB model detailed in 2018 paper by Nguyen et al.. This class inherits from the torchkge.models.interfaces.Model interface. It then has its attributes as well.

References

Nguyen, D. Q., Nguyen, T. D., Nguyen, D. Q., and Phung, D. A Novel Embed- ding Model for Knowledge Base Completion Based on Convolutional Neural Network. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (2018), vol. 2, pp. 327–333.

Parameters

emb_dim (int) – Dimension of embedding space.
n_filters (int) – Number of filters used for convolution.
n_entities (int) – Number of entities in the current data set.
n_relations (int) – Number of relations in the current data set.

ent_emb

Embeddings of the entities, initialized with Xavier uniform distribution and then normalized.

Type: torch.nn.Embedding, shape: (n_ent, emb_dim)

rel_emb

Embeddings of the relations, initialized with Xavier uniform distribution.

Type: torch.nn.Embedding, shape: (n_rel, emb_dim)

get_embeddings()[source]

Return the embeddings of entities and relations.

Returns

ent_emb (torch.Tensor, shape: (n_ent, emb_dim), dtype: torch.float) – Embeddings of entities.
rel_emb (torch.Tensor, shape: (n_rel, emb_dim), dtype: torch.float) – Embeddings of relations.

inference_prepare_candidates(h_idx, t_idx, r_idx, entities=True)[source]: Link prediction evaluation helper function. Get entities embeddings and relations embeddings. The output will be fed to the inference_scoring_function method. See torchkge.models.interfaces.Models for more details on the API.

inference_scoring_function(h, t, r)[source]: Link prediction evaluation helper function. See torchkge.models.interfaces.Models for more details on the API.

normalize_parameters()[source]: Normalize the entity embeddings, as explained in original paper. This methods should be called at the end of each training epoch and at the end of training as well.

scoring_function(h_idx, t_idx, r_idx)[source]: Compute the scoring function for the triplets given as argument: by applying convolutions to the concatenation of the embeddings. See referenced paper for more details on the score. See torchkge.models.interfaces.Models for more details on the API.