cogdl.models.emb.knowledge_base

Module Contents

Classes

KGEModel

class cogdl.models.emb.knowledge_base.KGEModel(nentity, nrelation, hidden_dim, gamma, double_entity_embedding=False, double_relation_embedding=False)

Bases: cogdl.models.BaseModel

static add_args(parser)

Add model-specific arguments to the parser.

classmethod build_model_from_args(cls, args)

Build a new model instance.

forward(self, sample, mode='single')

Forward function that calculate the score of a batch of triples. In the ‘single’ mode, sample is a batch of triple. In the ‘head-batch’ or ‘tail-batch’ mode, sample consists two part. The first part is usually the positive sample. And the second part is the entities in the negative samples. Because negative samples and positive samples usually share two elements in their triple ((head, relation) or (relation, tail)).

abstract score(self, head, relation, tail, mode)

static train_step(model, optimizer, train_iterator, args)

A single train step. Apply back-propation and return the loss

static test_step(model, test_triples, all_true_triples, args)

Evaluate the model on test or valid datasets