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",
)