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.
 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
 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 relationspecific subspaces.
Translational Models
Parameters used to train models available in pretrained version :
Dataset 
Dimension 
Optimizer 
Learning Rate 
Batch Size 
Loss 
Margin 
L2 penalization 

TransE 
FB15k 
100 
Adam 
2.1e5 
32768 
Margin 
.651 
1e5 
TransE 
FB15k237 
100 
Adam 
2.1e5 
32768 
Margin 
.651 
1e5 
TransE 
FB15k237 
150 
Adam 
2.7e5 
32768 
Margin 
.648 
1e5 
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 GarciaDuran, Jason Weston, and Oksana Yakhnenko. Translating Embeddings for Modeling Multirelational 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 TwentyEighth 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 relationspecific 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 relationspecific 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.
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 TwentyNinth 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
Relationspecific 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 relationspecific subspace.
 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) – Relationspecific 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.
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
Entityspecific 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
Relationspecific 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 relationspecific subspace.
 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.
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.
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 HansPeter Kriegel. A Threeway Model for Collective Learning on Multirelational 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.
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, Wentau 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.
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[i  j 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.
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 relationspecific 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 MultiRelational 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 relationspecific matrices to be scalars. By default it is set to half according to the original paper.
 scalar_dim
Number of diagonal elements of the relationspecific 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 relationspecific 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.