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

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

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

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

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

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

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

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

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

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

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.