GNNModel

CompGCN

class neuralkg.model.GNNModel.CompGCN.CompGCN(args)[source]

Bases: torch.nn.modules.module.Module

Composition-based multi-relational graph convolutional networks (CompGCN),

which jointly embeds both nodes and relations in a relational graph.

args

Model configuration parameters.

init_model()[source]

Initialize the CompGCN model and embeddings

Parameters

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

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

  • GraphCov – The comp graph convolution layers.

  • conv1 – The convolution layer.

  • fc – The full connection layer.

  • bn0 – The batch Normalization layer.

  • bn1 – The batch Normalization layer.

  • bn2 – The batch Normalization layer.

  • inp_drop – The dropout layer.

  • hid_drop – The dropout layer.

  • feg_drop – The dropout layer.

forward(graph, relation, norm, triples)[source]

The functions used in the training phase

Parameters

  • graph – The knowledge graph recorded in dgl.graph()

  • relation – The relation id sampled in triples

  • norm – The edge norm in graph

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

Returns

The score of triples.

Return type

score

get_score(batch, mode)[source]

The functions used in the testing phase

Parameters

  • batch – A batch of data.

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

Returns

The score of triples.

Return type

score

DistMult(head_emb, rela_emb)[source]

Calculating the score of triples with DistMult model.

concat(ent_embed, rel_embed)[source]
ConvE(sub_emb, rel_emb, all_ent)[source]

Calculating the score of triples with ConvE model.

training: bool
class neuralkg.model.GNNModel.CompGCN.CompGCNCov(in_channels, out_channels, act=<function CompGCNCov.<lambda>>, bias=True, drop_rate=0.0, opn='corr')[source]

Bases: torch.nn.modules.module.Module

The comp graph convolution layers, similar to https://github.com/malllabiisc/CompGCN

get_param(shape)[source]
message_func(edges: dgl.udf.EdgeBatch)[source]
reduce_func(nodes: dgl.udf.NodeBatch)[source]
comp(h, edge_data)[source]
forward(g: dgl.convert.graph, x, rel_repr, edge_type, edge_norm)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

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

training: bool

KBAT

class neuralkg.model.GNNModel.KBAT.KBAT(args)[source]

Bases: torch.nn.modules.module.Module

Learning Attention-based Embeddings for Relation Prediction in Knowledge Graphs (KBAT),

which introduces the attention to aggregate the neighbor node representation.

args

Model configuration parameters.

init_GAT_emb()[source]

Initialize the GAT model and embeddings

Parameters

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

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

  • entity_embeddings – The final embedding used in ConvKB.

  • relation_embeddings – The final embedding used in ConvKB.

  • attentions – The graph attention layers.

  • out_att – The graph attention layers.

init_ConvKB_emb()[source]

Initialize the ConvKB model.

Parameters

  • conv_layer – The convolution layer.

  • dropout – The dropout layer.

  • ReLU – Relu activation function.

  • fc_layer – The full connection layer.

forward(triples, mode, adj_matrix=None, n_hop=None)[source]

The functions used in the training and testing phase

Parameters

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

  • mode – The mode indicates that the model will be used, when it

  • 'GAT' (is) –

  • model (it means graph attetion) –

  • 'ConvKB' (when it is) –

:param : :param it means ConvKB model.:

Returns

The score of triples.

Return type

score

forward_Con(triples, mode)[source]
forward_GAT(triples, adj_matrix, n_hop)[source]
cal_Con_score(head_emb, rela_emb, tail_emb)[source]

Calculating the score of triples with ConvKB model.

Parameters

  • head_emb – The head entity embedding.

  • rela_emb – The relation embedding.

  • tail_emb – The tail entity embedding.

Returns

The score of triples.

Return type

score

cal_GAT_score(head_emb, relation_emb, tail_emb)[source]

Calculating the score of triples with TransE model.

Parameters

  • head_emb – The head entity embedding.

  • rela_emb – The relation embedding.

  • tail_emb – The tail entity embedding.

Returns

The score of triples.

Return type

score

training: bool
class neuralkg.model.GNNModel.KBAT.SpecialSpmmFunctionFinal[source]

Bases: torch.autograd.function.Function

Special function for only sparse region backpropataion layer, similar to https://arxiv.org/abs/1710.10903

static forward(ctx, edge, edge_w, N, E, out_features)[source]

Performs the operation.

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).

The context can be used to store tensors that can be then retrieved during the backward pass.

static backward(ctx, grad_output)[source]

Defines a formula for differentiating the operation.

This function is to be overridden by all subclasses.

It must accept a context ctx as the first argument, followed by as many outputs did forward() return, and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computated w.r.t. the output.

class neuralkg.model.GNNModel.KBAT.SpecialSpmmFinal[source]

Bases: torch.nn.modules.module.Module

Special spmm final layer, similar to https://arxiv.org/abs/1710.10903.

forward(edge, edge_w, N, E, out_features)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

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

training: bool
class neuralkg.model.GNNModel.KBAT.GraphAttentionLayer(num_nodes, in_features, out_features, nrela_dim, dropout, alpha, concat=True)[source]

Bases: torch.nn.modules.module.Module

Sparse version GAT layer, similar to https://arxiv.org/abs/1710.10903.

forward(input, edge, edge_embed, edge_list_nhop, edge_embed_nhop)[source]

Defines the computation performed at every call.

Should be overridden by all subclasses.

Note

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

training: bool

RGCN

class neuralkg.model.GNNModel.RGCN.RGCN(args)[source]

Bases: torch.nn.modules.module.Module

Modeling Relational Data with Graph Convolutional Networks (RGCN), which use GCN framework to model relation data.

args

Model configuration parameters.

build_model()[source]

Initialize the RGCN model and embeddings

Parameters

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

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

  • RGCN – the relation graph convolution model.

forward(graph, ent, rel, norm, triples, mode='single')[source]

The functions used in the training and testing phase

Parameters

  • graph – The knowledge graph recorded in dgl.graph()

  • ent – The entitiy ids sampled in triples

  • rel – The relation ids sampled in triples

  • norm – The edge norm in graph

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

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

Returns

The score of triples.

Return type

score

get_score(batch, mode)[source]

The functions used in the testing phase

Parameters

  • batch – A batch of data.

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

Returns

The score of triples.

Return type

score

tri2emb(embedding, triples, mode='single')[source]

Get embedding of triples.

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

Parameters

  • embedding (tensor) – This embedding save the entity embeddings.

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

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

Returns

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

Return type

head_emb

build_hidden_layer(idx)[source]

The functions used to initialize the RGCN model

Parameters

  • idx – it`s used to identify rgcn layers. The last rgcn layer should use

  • function. (relu as activation) –

Returns

the relation graph convolution layer

training: bool

XTransE

class neuralkg.model.GNNModel.XTransE.XTransE(args)[source]

Bases: neuralkg.model.KGEModel.model.Model

Explainable Knowledge Graph Embedding for Link Prediction with Lifestyles in e-Commerce (XTransE), which introduces the attention to aggregate the neighbor node representation.

args

Model configuration parameters.

init_emb()[source]

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

Parameters

  • margin – Caculate embedding_range and loss.

  • embedding_range – Uniform distribution range.

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

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

score_func(triples, neighbor=None, mask=None, negs=None, mode='single')[source]

Calculating the score of triples.

Parameters

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

  • neighbor – The neighbors of tail entities.

  • mask – The mask of neighbor nodes

  • negs – Negative samples, defaults to None.

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

Returns

The score of triples.

Return type

score

transe_func(head_emb, rela_emb, tail_emb)[source]

Calculating the score of triples with TransE model.

Parameters

  • head_emb – The head entity embedding.

  • rela_emb – The relation embedding.

  • tail_emb – The tail entity embedding.

Returns

The score of triples.

Return type

score

forward(triples, neighbor=None, mask=None, negs=None, mode='single')[source]

The functions used in the training and testing phase

Parameters

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

  • neighbor – The neighbors of tail entities.

  • mask – The mask of neighbor nodes

  • negs – Negative samples, defaults to None.

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

Returns

The score of triples.

Return type

score

get_score(batch, mode)[source]

The functions used in the testing phase

Parameters

  • batch – A batch of data.

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

Returns

The score of triples.

Return type

score

training: bool