Extending the Models

You should first read the tutorial on bringing your own interaction module. This tutorial is about how to wrap a custom interaction module with a model module for general reuse and application.

Implementing a model by subclassing pykeen.models.ERModel

The following code block demonstrates how an interaction model can be used to define a full KGEM using the pykeen.models.ERModel base class.

from pykeen.models import ERModel
from pykeen.nn import Embedding, Interaction


class DistMultInteraction(Interaction):
    def forward(self, h, r, t):
        return (h * r * t).sum(dim=-1)


class DistMult(ERModel):
    def __init__(
        self,
        # When defining your class, any hyper-parameters that can be configured should be
        # made as arguments to the __init__() function. When running the pipeline(), these
        # are passed via the ``model_kwargs``.
        embedding_dim: int = 50,
        # All remaining arguments are simply passed through to the parent constructor. If you
        # want access to them, you can name them explicitly. See the pykeen.models.ERModel
        # documentation for a full list
        **kwargs,
    ) -> None:
        # since this is a python class, you can feel free to get creative here. One example of
        # pre-processing is to derive the shape for the relation representation based on the
        # embedding dimension.
        super().__init__(
            # Pass an instance of your interaction function. This is also a place where you can
            # pass hyper-parameters, such as the L_p norm, from the KGEM to the interaction function
            interaction=DistMultInteraction,
            # interaction_kwargs=dict(...),
            # Define the entity representations using a dict. By default, each
            # embedding is a vector. You can use the ``shape`` kwarg to specify higher dimensional
            # tensor shapes.
            entity_representations=Embedding,
            entity_representations_kwargs=dict(
                embedding_dim=embedding_dim,
            ),
            # Define the relation representations the same as the entities
            relation_representations=Embedding,
            relation_representations_kwargs=dict(
                embedding_dim=embedding_dim,
            ),
            # All other arguments are passed through, such as the ``triples_factory``, ``loss``,
            # ``preferred_device``, and others. These are all handled by the pipeline() function
            **kwargs,
        )

The actual implementation of DistMult can be found in pykeen.models.DistMult. Note that it additionally contains configuration for the initializers, constrainers, and regularizers for each of the embeddings as well as class-level defaults for hyper-parameters and hyper-parameter optimization. Modifying these is covered in other tutorials.

Specifying Defaults

If you have a preferred loss function for your model, you can add the loss_default class variable where the value is the loss class.

from typing import ClassVar

from pykeen.models import ERModel
from pykeen.losses import Loss, NSSALoss

class DistMult(ERModel):
    loss_default: ClassVar[Type[Loss]] = NSSALoss
    ...

Now, when using the pipeline, the pykeen.losses.NSSALoss. loss is used by default if none is given. The same kind of modifications can be made to set a default regularizer with regularizer_default.

Specifying Hyper-parameter Optimization Default Ranges

All subclasses of pykeen.models.Model can specify the default ranges or values used during hyper-parameter optimization (HPO). PyKEEN implements a simple dictionary-based configuration that is interpreted by pykeen.hpo.hpo.suggest_kwargs() in the HPO pipeline.

HPO default ranges can be applied to all keyword arguments appearing in the __init__() function of your model by setting a class-level variable called hpo_default.

For example, the embedding_dim can be specified as being on a range between 100 and 150 with the following:

class DistMult(ERModel):
    hpo_default = {
        'embedding_dim': dict(type=int, low=100, high=150)
    }
    ...

A step size can be imposed with q:

class DistMult(ERModel):
    hpo_default = {
        'embedding_dim': dict(type=int, low=100, high=150 q=5)
    }
    ...

An alternative scale can be imposed with scale. Right now, the default is linear, and scale can optionally be set to power_two for integers as in:

class DistMult(ERModel):
    hpo_default = {
        # will uniformly give 16, 32, 64, 128 (left inclusive, right exclusive)
        'hidden_dim': dict(type=int, low=4, high=8, scale='power_two')
    }
    ...

Warning

Alternative scales can not currently be used in combination with step size (q).

There are other possibilities for specifying the type as float, categorical, or as bool.

With float, you can’t use the q option nor set the scale to power_two, but the scale can be set to log (see optuna.distributions.LogUniformDistribution).

hpo_default = {
    # will uniformly give floats on the range of [1.0, 2.0) (exclusive)
    'alpha': dict(type='float', low=1.0, high=2.0),

    # will uniformly give 1.0, 2.0, or 4.0 (exclusive)
    'beta': dict(type='float', low=1.0, high=8.0, scale='log'),
}

With categorical, you can form a dictionary like the following using type='categorical' and giving a choices entry that contains a sequence of either integers, floats, or strings.

hpo_default = {
    'similarity': dict(type='categorical', choices=[...])
}

With bool, you can simply use dict(type=bool) or dict(type='bool').

Note

The HPO rules are subject to change as they are tightly coupled to optuna, which since version 2.0.0 has introduced several new possibilities.

Implementing a model by instantiating pykeen.models.ERModel

Instead of creating a new class, you can also directly use the pykeen.models.ERModel, e.g.

from pykeen.models import ERModel
from pykeen.losses import BCEWithLogitsLoss

model = ERModel(
    triples_factory=...,
    loss="BCEWithLogits",
    interaction="transformer",
    entity_representations_kwargs=dict(embedding_dim=64),
    relation_representations_kwargs=dict(embedding_dim=64),
)

Using a Custom Model with the Pipeline

We can use this new model with all available losses, evaluators, training pipelines, inverse triple modeling, via the pykeen.pipeline.pipeline(), since in addition to the names of models (given as strings), it can also take model classes in the model argument.

from pykeen.pipeline import pipeline

pipeline(
    model=DistMult,
    dataset='Nations',
    loss='NSSA',
)