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.

Implement a simple 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 EmbeddingSpecification
from pykeen.nn.modules import DistMultInteraction  # effectively the same as the example above

class DistMult(ERModel):
    def __init__(
        # 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
    ) -> 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.
            # 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
            # Define the entity representations using the EmbeddingSpecification. By default, each
            # embedding is linear. You can use the ``shape`` kwarg to specify higher dimensional
            # tensor shapes.
            # Define the relation representations the same as the entities
            # All other arguments are passed through, such as the ``triples_factory``, ``loss``,
            # ``preferred_device``, and others. These are all handled by the pipeline() function

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.

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

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

model = ERModel(


tutorial on rolling your own more complicated model, like pykeen.nn.modules.NTNInteraction or pykeen.nn.modules.TransDInteraction.


tutorial on using some of the inheriting classes of pykeen.nn.modules.Interaction like pykeen.nn.modules.FunctionalInteraction or pykeen.nn.modules.TranslationalInteraction