Source code for pykeen.models.unimodal.rotate

# -*- coding: utf-8 -*-

"""Implementation of the RotatE model."""

from typing import Any, ClassVar, Mapping

import torch
from class_resolver import HintOrType, OptionalKwargs

from ..nbase import ERModel
from ...nn.init import init_phases, xavier_uniform_
from ...nn.modules import RotatEInteraction
from ...regularizers import Regularizer
from ...typing import Constrainer, Hint, Initializer
from ...utils import complex_normalize

__all__ = [

[docs]class RotatE(ERModel): r"""An implementation of RotatE from [sun2019]_. RotatE models relations as rotations from head to tail entities in complex space: .. math:: \textbf{e}_t= \textbf{e}_h \odot \textbf{r}_r where $\textbf{e}, \textbf{r} \in \mathbb{C}^{d}$ and the complex elements of $\textbf{r}_r$ are restricted to have a modulus of one ($\|\textbf{r}_r\| = 1$). The interaction model is then defined as: .. math:: f(h,r,t) = -\|\textbf{e}_h \odot \textbf{r}_r - \textbf{e}_t\| which allows to model symmetry, antisymmetry, inversion, and composition. .. seealso:: - Authors' `implementation of RotatE <>`_ --- citation: author: Sun year: 2019 link: github: DeepGraphLearning/KnowledgeGraphEmbedding """ #: The default strategy for optimizing the model's hyper-parameters hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=dict(type=int, low=32, high=1024, q=16), ) def __init__( self, *, embedding_dim: int = 200, entity_initializer: Hint[Initializer] = xavier_uniform_, relation_initializer: Hint[Initializer] = init_phases, relation_constrainer: Hint[Constrainer] = complex_normalize, regularizer: HintOrType[Regularizer] = None, regularizer_kwargs: OptionalKwargs = None, **kwargs, ) -> None: """ Initialize the model. :param embedding_dim: the embedding dimension :param entity_initializer: the entity representation initializer :param relation_initializer: the relation representation initializer :param relation_constrainer: the relation representation constrainer :param regularizer: the regularizer :param regularizer_kwargs: additional keyword-based parameters passed to the regularizer :param kwargs: additional keyword-based parameters passed to :meth:`ERModel.__init__` """ super().__init__( interaction=RotatEInteraction, entity_representations_kwargs=dict( shape=embedding_dim, initializer=entity_initializer, regularizer=regularizer, regularizer_kwargs=regularizer_kwargs, dtype=torch.cfloat, ), relation_representations_kwargs=dict( shape=embedding_dim, initializer=relation_initializer, constrainer=relation_constrainer, dtype=torch.cfloat, ), **kwargs, )