扩展模型

您应该首先阅读关于如何引入自己的交互模块的教程。 本教程是关于如何将自定义交互模块与模型模块封装在一起，以便于通用复用和应用。

通过子类化pykeen.models.ERModel实现模型

以下代码块演示了如何使用交互模型来定义一个完整的KGEM，使用pykeen.models.ERModel基类。

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

DistMult 的实际实现可以在 pykeen.models.DistMult 中找到。请注意， 它还包含了每个嵌入的初始化器、约束器和正则化器的配置， 以及超参数和超参数优化的类级默认值。修改这些内容在其他教程中有详细介绍。

指定默认值

如果您为模型有偏好的损失函数，可以添加loss_default类变量，其中值是损失类。

from typing import ClassVar

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

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

现在，当使用管道时，如果没有给出损失函数，默认会使用pykeen.losses.NSSALoss。同样的修改可以通过regularizer_default来设置默认的正则化器。

指定超参数优化默认范围

pykeen.models.Model的所有子类都可以指定在超参数优化（HPO）过程中使用的默认范围或值。PyKEEN实现了一个基于字典的简单配置，该配置由HPO管道中的pykeen.hpo.hpo.suggest_kwargs()解释。

HPO默认范围可以通过设置一个名为hpo_default的类级别变量，应用于出现在模型__init__()函数中的所有关键字参数。

例如，可以使用以下方式将embedding_dim指定在100到150之间的范围内：

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

可以使用 q 来设置步长：

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

可以使用scale来设置一个替代的比例。目前，默认是线性的，scale可以选择设置为power_two，例如对于整数：

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')
    }
    ...

警告

目前，替代尺度不能与步长（q）结合使用。

还有其他可能性可以将type指定为float、categorical或bool。

使用 float 时，你不能使用 q 选项，也不能将比例设置为 power_two， 但可以将比例设置为 log（参见 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'),
}

使用 categorical，你可以通过 type='categorical' 和提供一个包含整数、浮点数或字符串序列的 choices 条目来形成如下所示的字典。

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

使用 bool，你可以简单地使用 dict(type=bool) 或 dict(type='bool')。

注意

HPO规则可能会发生变化，因为它们与optuna紧密耦合，自2.0.0版本以来，它引入了许多新的可能性。

通过实例化pykeen.models.ERModel来实现模型

除了创建一个新类，你也可以直接使用pykeen.models.ERModel，例如。

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

在管道中使用自定义模型

我们可以通过pykeen.pipeline.pipeline()使用这个新模型与所有可用的损失函数、评估器、训练管道、逆三元组建模，因为除了模型名称（以字符串形式给出）之外，它还可以在model参数中接受模型类。

from pykeen.pipeline import pipeline

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