"""Implementation of TuckEr."""
from collections.abc import Mapping
from typing import Any, ClassVar
from class_resolver import OptionalKwargs
from ..nbase import ERModel
from ...constants import DEFAULT_DROPOUT_HPO_RANGE, DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE
from ...losses import BCEAfterSigmoidLoss, Loss
from ...nn import TuckERInteraction
from ...nn.init import xavier_normal_
from ...typing import FloatTensor, Hint, Initializer
__all__ = [
"TuckER",
]
[docs]
class TuckER(ERModel[FloatTensor, FloatTensor, FloatTensor]):
r"""An implementation of TuckEr from [balazevic2019]_.
It represents entities by $d_e$-dimensional vectors and relations by $d_r$-dimensional vectors, stored in
:class:`~pykeen.nn.representation.Embedding`. The state-ful :class:`~pykeen.nn.modules.TuckERInteraction` is then
used to score triples.
For $E$ entities and $R$ relations, the model has $Ed_e + Rd_r + d_e^2d_r$ effective parameters (ignoring additional
parameters from the :class:`torch.nn.BatchNorm1d` layers in :class:`~pykeen.nn.modules.TuckERInteraction`).
.. seealso::
- Official implementation: https://github.com/ibalazevic/TuckER
- pykg2vec implementation of TuckEr https://github.com/Sujit-O/pykg2vec/blob/master/pykg2vec/core/TuckER.py
---
citation:
author: Balažević
year: 2019
link: https://arxiv.org/abs/1901.09590
github: ibalazevic/TuckER
"""
#: The default strategy for optimizing the model's hyper-parameters
hpo_default: ClassVar[Mapping[str, Any]] = {
"embedding_dim": DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE,
"relation_dim": DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE,
"dropout_0": DEFAULT_DROPOUT_HPO_RANGE,
"dropout_1": DEFAULT_DROPOUT_HPO_RANGE,
"dropout_2": DEFAULT_DROPOUT_HPO_RANGE,
}
#: The default loss function class
loss_default: ClassVar[type[Loss]] = BCEAfterSigmoidLoss # type: ignore[type-abstract]
#: The default parameters for the default loss function class
loss_default_kwargs: ClassVar[Mapping[str, Any]] = {}
def __init__(
self,
*,
embedding_dim: int = 200,
relation_dim: int | None = None,
dropout_0: float = 0.3,
dropout_1: float = 0.4,
dropout_2: float = 0.5,
apply_batch_normalization: bool = True,
entity_initializer: Hint[Initializer] = xavier_normal_,
relation_initializer: Hint[Initializer] = xavier_normal_,
core_tensor_initializer: Hint[Initializer] = None,
core_tensor_initializer_kwargs: OptionalKwargs = None,
**kwargs,
) -> None:
"""
Initialize the model.
:param embedding_dim:
the (entity) embedding dimension
:param relation_dim:
the relation embedding dimension. Defaults to `embedding_dim`.
:param dropout_0:
the first dropout, cf. formula
:param dropout_1:
the second dropout, cf. formula
:param dropout_2:
the third dropout, cf. formula
:param apply_batch_normalization:
whether to apply batch normalization
:param entity_initializer:
the entity representation initializer
:param relation_initializer:
the relation representation initializer
:param core_tensor_initializer:
the core tensor initializer
:param core_tensor_initializer_kwargs:
keyword-based parameters passed to the core tensor initializer
:param kwargs:
additional keyword-based parameters passed to :meth:`ERModel.__init__`
"""
relation_dim = relation_dim or embedding_dim
super().__init__(
interaction=TuckERInteraction,
interaction_kwargs={
"embedding_dim": embedding_dim,
"relation_dim": relation_dim,
"head_dropout": dropout_0, # TODO: rename
"relation_dropout": dropout_1,
"head_relation_dropout": dropout_2,
"apply_batch_normalization": apply_batch_normalization,
"core_initializer": core_tensor_initializer,
"core_initializer_kwargs": core_tensor_initializer_kwargs,
},
entity_representations_kwargs={
"shape": embedding_dim,
"initializer": entity_initializer,
},
relation_representations_kwargs={
"shape": relation_dim,
"initializer": relation_initializer,
},
**kwargs,
)