From d5d0524a3d26921100d185db1ff2d100a01087a0 Mon Sep 17 00:00:00 2001 From: mali Date: Thu, 6 Aug 2020 11:12:52 +0200 Subject: [PATCH 001/690] Implement interaction function for ERMLP --- src/pykeen/models/unimodal/ermlp.py | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index e8a3e41723..55f169ee2b 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -91,6 +91,20 @@ def _reset_parameters_(self): # noqa: D102 nn.init.zeros_(self.linear2.bias) nn.init.xavier_uniform_(self.linear2.weight, gain=nn.init.calculate_gain('relu')) + def interaction_function( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + ) -> torch.FloatTensor: + """.""" + + # Concatenate them + x_s = torch.cat([h, r, t], dim=-1) + + # Compute scores + return self.mlp(x_s) + def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings h = self.entity_embeddings(hrt_batch[:, 0]) @@ -100,11 +114,7 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: # Embedding Regularization self.regularize_if_necessary(h, r, t) - # Concatenate them - x_s = torch.cat([h, r, t], dim=-1) - - # Compute scores - return self.mlp(x_s) + return self.interaction_function(h, r, t) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings From f7265726ca1cd943c70e75ccbbbcb152e82cc127 Mon Sep 17 00:00:00 2001 From: mali Date: Thu, 6 Aug 2020 12:16:27 +0200 Subject: [PATCH 002/690] Enable usage of user specific-embeddings --- src/pykeen/models/unimodal/ermlp.py | 51 ++++++++++++++++++++--------- 1 file changed, 36 insertions(+), 15 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 55f169ee2b..1e7955ca08 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -105,22 +105,36 @@ def interaction_function( # Compute scores return self.mlp(x_s) - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(hrt_batch[:, 0]) - r = self.relation_embeddings(hrt_batch[:, 1]) - t = self.entity_embeddings(hrt_batch[:, 2]) + def score_hrt( + self, + hrt_batch: torch.LongTensor, + h: Optional[torch.FloatTensor]=None, + r: Optional[torch.FloatTensor]=None, + t: Optional[torch.FloatTensor]=None, + ) -> torch.FloatTensor: # noqa: D102 + if h is None and r is None and t is None: + # Get embeddings + h = self.entity_embeddings(hrt_batch[:, 0]) + r = self.relation_embeddings(hrt_batch[:, 1]) + t = self.entity_embeddings(hrt_batch[:, 2]) # Embedding Regularization self.regularize_if_necessary(h, r, t) return self.interaction_function(h, r, t) - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(hr_batch[:, 0]) - r = self.relation_embeddings(hr_batch[:, 1]) - t = self.entity_embeddings.weight + def score_t( + self, + hr_batch: torch.LongTensor, + h: Optional[torch.FloatTensor] = None, + r: Optional[torch.FloatTensor] = None, + t: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: # noqa: D102 + if h is None and r is None and t is None: + # Get embeddings + h = self.entity_embeddings(hr_batch[:, 0]) + r = self.relation_embeddings(hr_batch[:, 1]) + t = self.entity_embeddings.weight # Embedding Regularization self.regularize_if_necessary(h, r, t) @@ -142,11 +156,18 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 scores = scores.view(-1, self.num_entities) return scores - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings.weight - r = self.relation_embeddings(rt_batch[:, 0]) - t = self.entity_embeddings(rt_batch[:, 1]) + def score_h( + self, + rt_batch: torch.LongTensor, + h: Optional[torch.FloatTensor] = None, + r: Optional[torch.FloatTensor] = None, + t: Optional[torch.FloatTensor] = None, + ) -> torch.FloatTensor: # noqa: D102 + if h is None and r is None and t is None: + # Get embeddings + h = self.entity_embeddings.weight + r = self.relation_embeddings(rt_batch[:, 0]) + t = self.entity_embeddings(rt_batch[:, 1]) # Embedding Regularization self.regularize_if_necessary(h, r, t) From ea677c74ba5147265b70d2cae5b076c519e6bc72 Mon Sep 17 00:00:00 2001 From: mali Date: Thu, 6 Aug 2020 12:18:45 +0200 Subject: [PATCH 003/690] Remove interaction function --- src/pykeen/models/unimodal/ermlp.py | 20 +++++--------------- 1 file changed, 5 insertions(+), 15 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 1e7955ca08..48281aa7fe 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -91,20 +91,6 @@ def _reset_parameters_(self): # noqa: D102 nn.init.zeros_(self.linear2.bias) nn.init.xavier_uniform_(self.linear2.weight, gain=nn.init.calculate_gain('relu')) - def interaction_function( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - ) -> torch.FloatTensor: - """.""" - - # Concatenate them - x_s = torch.cat([h, r, t], dim=-1) - - # Compute scores - return self.mlp(x_s) - def score_hrt( self, hrt_batch: torch.LongTensor, @@ -121,7 +107,11 @@ def score_hrt( # Embedding Regularization self.regularize_if_necessary(h, r, t) - return self.interaction_function(h, r, t) + # Concatenate them + x_s = torch.cat([h, r, t], dim=-1) + + # Compute scores + return self.mlp(x_s) def score_t( self, From 8a2dd3ba61635a9479c5662a76a81aa62a43e8cf Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 6 Aug 2020 14:32:31 +0200 Subject: [PATCH 004/690] Add stub tutorial --- docs/source/index.rst | 1 + docs/source/tutorial/bring_your_own_embeddings.rst | 2 ++ 2 files changed, 3 insertions(+) create mode 100644 docs/source/tutorial/bring_your_own_embeddings.rst diff --git a/docs/source/index.rst b/docs/source/index.rst index 9f0242f3bc..d9f87abcb3 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -12,6 +12,7 @@ PyKEEN tutorial/understanding_evaluation tutorial/running_hpo tutorial/using_mlflow + tutorial/bring_your_own_embeddings .. toctree:: :caption: Reference diff --git a/docs/source/tutorial/bring_your_own_embeddings.rst b/docs/source/tutorial/bring_your_own_embeddings.rst new file mode 100644 index 0000000000..472a9f3ee8 --- /dev/null +++ b/docs/source/tutorial/bring_your_own_embeddings.rst @@ -0,0 +1,2 @@ +Bring Your Own Embeddings +========================= From 0d352e8c340336e0f60bec422410f10f59bbab5c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 17 Aug 2020 14:01:37 +0200 Subject: [PATCH 005/690] Add functional form of ER-MLP interaction function --- src/pykeen/models/unimodal/ermlp.py | 32 +++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 48281aa7fe..bf4103f9bd 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -18,6 +18,38 @@ ] +def score( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + mlp: nn.Module, +) -> torch.FloatTensor: + r""" + Evaluate the ER-MLP interaction function. + + .. math:: + + f(h,r,t) = \textbf{w}^{T} g(\textbf{W} [\textbf{h}; \textbf{r}; \textbf{t}]), + + :param h: shape: (b, d) + The head entity embeddings. + :param r: shape: (b, d) + The relation embeddings. + :param t: shape: (b, d) + The tail entity embeddings. + :param mlp: + The MLP. + + :return: shape: (b,) + The scores. + """ + # Concatenate them + x_s = torch.cat([h, r, t], dim=-1) + + # Compute scores + return mlp(x_s) + + class ERMLP(EntityRelationEmbeddingModel): r"""An implementation of ERMLP from [dong2014]_. From 62c55d234ad373a2b37c33fe1968a6426272fa32 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 17 Aug 2020 14:03:16 +0200 Subject: [PATCH 006/690] Use functional form in score(...) --- src/pykeen/models/unimodal/ermlp.py | 19 ++++++------------- 1 file changed, 6 insertions(+), 13 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index bf4103f9bd..0d8ca5c582 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -126,24 +126,17 @@ def _reset_parameters_(self): # noqa: D102 def score_hrt( self, hrt_batch: torch.LongTensor, - h: Optional[torch.FloatTensor]=None, - r: Optional[torch.FloatTensor]=None, - t: Optional[torch.FloatTensor]=None, ) -> torch.FloatTensor: # noqa: D102 - if h is None and r is None and t is None: - # Get embeddings - h = self.entity_embeddings(hrt_batch[:, 0]) - r = self.relation_embeddings(hrt_batch[:, 1]) - t = self.entity_embeddings(hrt_batch[:, 2]) + # Get embeddings + h = self.entity_embeddings(hrt_batch[:, 0]) + r = self.relation_embeddings(hrt_batch[:, 1]) + t = self.entity_embeddings(hrt_batch[:, 2]) # Embedding Regularization self.regularize_if_necessary(h, r, t) - # Concatenate them - x_s = torch.cat([h, r, t], dim=-1) - - # Compute scores - return self.mlp(x_s) + # Apply interaction function + return score(h=h, r=r, t=t, mlp=self.mlp) def score_t( self, From 41ac013b17de4a4879d04a04d0e2085e71c5a473 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 17 Aug 2020 14:44:26 +0200 Subject: [PATCH 007/690] Inline variable --- src/pykeen/models/unimodal/ermlp.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 0d8ca5c582..7cfd6e730a 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -43,11 +43,7 @@ def score( :return: shape: (b,) The scores. """ - # Concatenate them - x_s = torch.cat([h, r, t], dim=-1) - - # Compute scores - return mlp(x_s) + return mlp(torch.cat([h, r, t], dim=-1)) class ERMLP(EntityRelationEmbeddingModel): From 17439778a7f9a19c040a74a4245378b93c19ec82 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 10 Sep 2020 16:55:12 +0200 Subject: [PATCH 008/690] Add draft of generic interaction function --- src/pykeen/models/base.py | 57 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 57 insertions(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 0f2ecddc3e..1d3883abb0 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -829,3 +829,60 @@ def _can_slice(fn) -> bool: class MultimodalModel(EntityRelationEmbeddingModel): """A multimodal KGE model.""" + + +class InteractionFunction: + """Base class for interaction functions.""" + + def score( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + ) -> torch.FloatTensor: + """ + Generic interaction function. + + :param h: shape: (batch_size, num_heads, d_e) + The head representations. + :param r: shape: (batch_size, num_relations, d_r) + The relation representations. + :param t: shape: (batch_size, num_tails, d_e) + The tail representations. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + raise NotImplementedError + + def score_t( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + all_entities: torch.FloatTensor, + ): + """ + Score all tail entities. + + :param h: shape: (batch_size, d_e) + The head representations. + :param r: shape: (batch_size, d_r) + The relation representations. + :param all_entities: shape: (num_entities, d_e) + The tail representations. + + :return: shape: (batch_size, num_entities) + The scores. + """ + # prepare input to generic score function + h = h.unsqueeze(dim=1) + r = r.unsqueeze(dim=1) + t = all_entities.unsqueeze(dim=0) + + # get scores + scores = self.score(h=h, r=r, t=t) + + # prepare output shape + scores = scores.squeeze(dim=2).squeeze(dim=1) + + return scores From f157578ae6504bd5218c3d22f3aa26b4a3060f92 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 10 Sep 2020 16:58:53 +0200 Subject: [PATCH 009/690] Add correct super-class --- src/pykeen/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 1d3883abb0..a307b9f67c 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -831,7 +831,7 @@ class MultimodalModel(EntityRelationEmbeddingModel): """A multimodal KGE model.""" -class InteractionFunction: +class InteractionFunction(nn.Module): """Base class for interaction functions.""" def score( From eff8c53237df85125ce639d884269b6b816d2b80 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 10 Sep 2020 16:59:05 +0200 Subject: [PATCH 010/690] Add missing type annotation --- src/pykeen/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index a307b9f67c..836795e22d 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -860,7 +860,7 @@ def score_t( h: torch.FloatTensor, r: torch.FloatTensor, all_entities: torch.FloatTensor, - ): + ) -> torch.FloatTensor: """ Score all tail entities. From 8f40a4e9bedfccc41cae83f50f49df1223f015f4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 10 Sep 2020 17:00:50 +0200 Subject: [PATCH 011/690] Add DistMultInteractionFunction --- src/pykeen/models/base.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 836795e22d..a25c529913 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -886,3 +886,15 @@ def score_t( scores = scores.squeeze(dim=2).squeeze(dim=1) return scores + + +class DistMultInteractionFunction(InteractionFunction): + """Interaction function of DistMult.""" + + def score( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + ) -> torch.FloatTensor: # noqa: D102 + return torch.einsum('bhd,brd,btd->bhrt', h, r, t) From 9faabd87b21ae15fb43a584a6d5eb36365dd85e7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 16 Sep 2020 09:44:06 +0200 Subject: [PATCH 012/690] Rename score to forward --- src/pykeen/models/base.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index a25c529913..6e9e6b05b8 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -834,7 +834,7 @@ class MultimodalModel(EntityRelationEmbeddingModel): class InteractionFunction(nn.Module): """Base class for interaction functions.""" - def score( + def forward( self, h: torch.FloatTensor, r: torch.FloatTensor, @@ -880,7 +880,7 @@ def score_t( t = all_entities.unsqueeze(dim=0) # get scores - scores = self.score(h=h, r=r, t=t) + scores = self.forward(h=h, r=r, t=t) # prepare output shape scores = scores.squeeze(dim=2).squeeze(dim=1) @@ -891,7 +891,7 @@ def score_t( class DistMultInteractionFunction(InteractionFunction): """Interaction function of DistMult.""" - def score( + def forward( self, h: torch.FloatTensor, r: torch.FloatTensor, From 3afa4f5f318c455db34ae95950a6e1b8f24bf7c4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 16 Sep 2020 09:44:39 +0200 Subject: [PATCH 013/690] Use __call__ instead of explicitly calling forward This ensures that all hooks get called. --- src/pykeen/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 6e9e6b05b8..55d11338d1 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -880,7 +880,7 @@ def score_t( t = all_entities.unsqueeze(dim=0) # get scores - scores = self.forward(h=h, r=r, t=t) + scores = self(h=h, r=r, t=t) # prepare output shape scores = scores.squeeze(dim=2).squeeze(dim=1) From cc4117bddcc463da54965430865451ccf96484b3 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 16 Sep 2020 09:54:07 +0200 Subject: [PATCH 014/690] Add ER-MLP interaction function --- src/pykeen/models/base.py | 43 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 43 insertions(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 55d11338d1..a251bf7ee1 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -898,3 +898,46 @@ def forward( t: torch.FloatTensor, ) -> torch.FloatTensor: # noqa: D102 return torch.einsum('bhd,brd,btd->bhrt', h, r, t) + + +class ERMLPInteractionFunction(InteractionFunction): + """ + Interaction function of ER-MLP. + + .. math :: + f(h, r, t) = W_2 ReLU(W_1 cat(h, r, t) + b_1) + b_2 + """ + + def __init__( + self, + embedding_dim: int, + hidden_dim: int, + ): + """ + Initialize the interaction function. + + :param embedding_dim: + The embedding vector dimension. + :param hidden_dim: + The hidden dimension of the MLP. + """ + super().__init__() + self.head_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) + self.rel_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=True) + self.tail_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) + self.activation = nn.ReLU() + self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + ) -> torch.FloatTensor: # noqa: D102 + h = self.head_to_hidden(h) + r = self.rel_to_hidden(r) + t = self.tail_to_hidden(t) + # TODO: Choosing which to combine first, h/r, h/t or r/t, depending on the shape might further improve + # performance in a 1:n scenario. + x = self.activation(h[:, :, None, None, :] + r[:, None, :, None, :] + t[:, None, None, :, :]) + return self.hidden_to_score(x).squeeze(dim=-1) From 6fa6bdf8e6afdb5102b5619eaec4b3a7afed4aba Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 16 Sep 2020 10:09:59 +0200 Subject: [PATCH 015/690] Use interaction function in DistMult --- src/pykeen/models/base.py | 31 ++++++++++- src/pykeen/models/unimodal/distmult.py | 74 +++++++------------------- 2 files changed, 47 insertions(+), 58 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index a251bf7ee1..2f56261ca9 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -6,7 +6,7 @@ import logging from abc import abstractmethod from collections import defaultdict -from typing import Any, ClassVar, Collection, Dict, Iterable, List, Mapping, Optional, Set, Type, Union +from typing import Any, ClassVar, Collection, Dict, Iterable, List, Mapping, Optional, Set, Tuple, Type, Union import pandas as pd import torch @@ -888,6 +888,29 @@ def score_t( return scores +def _normalize_for_einsum( + x: torch.FloatTensor, + batch_size: int, + symbol: str, +) -> Tuple[str, torch.FloatTensor]: + """ + Normalize tensor for broadcasting along batch-dimension in einsum. + + :param x: + The tensor. + :param batch_size: + The batch_size + :param symbol: + The symbol for the einsum term. + + :return: + A tuple (reshaped_tensor, term). + """ + if x.shape[0] == batch_size: + return f'b{symbol}d', x + return f'{symbol}d', x.squeeze(dim=0) + + class DistMultInteractionFunction(InteractionFunction): """Interaction function of DistMult.""" @@ -897,7 +920,11 @@ def forward( r: torch.FloatTensor, t: torch.FloatTensor, ) -> torch.FloatTensor: # noqa: D102 - return torch.einsum('bhd,brd,btd->bhrt', h, r, t) + batch_size = max(h.shape[0], r.shape[0], t.shape[0]) + h_term, h = _normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') + r_term, r = _normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') + t_term, t = _normalize_for_einsum(x=t, batch_size=batch_size, symbol='t') + return torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', h, r, t) class ERMLPInteractionFunction(InteractionFunction): diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 5d81001b66..03583ad796 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -4,12 +4,11 @@ from typing import Optional -import torch import torch.autograd from torch import nn from torch.nn import functional -from ..base import EntityRelationEmbeddingModel +from ..base import DistMultInteractionFunction, EntityRelationEmbeddingModel from ...losses import Loss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory @@ -18,6 +17,8 @@ 'DistMult', ] +from ...utils import get_embedding_in_canonical_shape + class DistMult(EntityRelationEmbeddingModel): r"""An implementation of DistMult from [yang2014]_. @@ -88,6 +89,7 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) + self.interaction_function = DistMultInteractionFunction() # Finalize initialization self.reset_parameters_() @@ -106,53 +108,18 @@ def post_parameter_update(self) -> None: # noqa: D102 # Normalize embeddings of entities functional.normalize(self.entity_embeddings.weight.data, out=self.entity_embeddings.weight.data) - @staticmethod - def interaction_function( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, + def _score( + self, + h_ind: Optional[torch.LongTensor] = None, + r_ind: Optional[torch.LongTensor] = None, + t_ind: Optional[torch.LongTensor] = None, ) -> torch.FloatTensor: - """Evaluate the interaction function for given embeddings. - - The embeddings have to be in a broadcastable shape. - - WARNING: Does not ensure forward constraints. - - :param h: shape: (..., e) - Head embeddings. - :param r: shape: (..., e) - Relation embeddings. - :param t: shape: (..., e) - Tail embeddings. - - :return: shape: (...) - The scores. - """ - # Bilinear product - # *: Elementwise multiplication - return torch.sum(h * r * t, dim=-1) - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings - h = self.entity_embeddings(hrt_batch[:, 0]) - r = self.relation_embeddings(hrt_batch[:, 1]) - t = self.entity_embeddings(hrt_batch[:, 2]) + h = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=h_ind) + r = get_embedding_in_canonical_shape(embedding=self.relation_embeddings, ind=r_ind) + t = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=t_ind) # Compute score - scores = self.interaction_function(h=h, r=r, t=t).view(-1, 1) - - # Only regularize relation embeddings - self.regularize_if_necessary(r) - - return scores - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(hr_batch[:, 0]).view(-1, 1, self.embedding_dim) - r = self.relation_embeddings(hr_batch[:, 1]).view(-1, 1, self.embedding_dim) - t = self.entity_embeddings.weight.view(1, -1, self.embedding_dim) - - # Rank against all entities scores = self.interaction_function(h=h, r=r, t=t) # Only regularize relation embeddings @@ -160,16 +127,11 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 return scores - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings.weight.view(1, -1, self.embedding_dim) - r = self.relation_embeddings(rt_batch[:, 0]).view(-1, 1, self.embedding_dim) - t = self.entity_embeddings(rt_batch[:, 1]).view(-1, 1, self.embedding_dim) - - # Rank against all entities - scores = self.interaction_function(h=h, r=r, t=t) + def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=hrt_batch[:, 0], r_ind=hrt_batch[:, 1], t_ind=hrt_batch[:, 2]).view(-1, 1) - # Only regularize relation embeddings - self.regularize_if_necessary(r) + def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=hr_batch[:, 0], r_ind=hr_batch[:, 1], t_ind=None).view(-1, self.num_entities) - return scores + def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) From e6a2e96dc0876b9df2aaaea8f4f76de53b23197c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 16 Sep 2020 10:17:33 +0200 Subject: [PATCH 016/690] Add reset_parameters to ER-MLP Since it is a parametric interaction function, it has internal parameters --- src/pykeen/models/base.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 2f56261ca9..7034ba5e7d 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -955,6 +955,11 @@ def __init__( self.activation = nn.ReLU() self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) + def reset_parameters(self): + for mod in self.modules(): + if hasattr(mod, 'reset_parameters'): + mod.reset_parameters() + def forward( self, h: torch.FloatTensor, From 5a80a4f2edba735076e8dfd3a027881ee1978211 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 16 Sep 2020 10:20:20 +0200 Subject: [PATCH 017/690] Fix reset_parameters infinite recursion since self is in self.modules --- src/pykeen/models/base.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 7034ba5e7d..9b89b282cc 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -949,6 +949,10 @@ def __init__( The hidden dimension of the MLP. """ super().__init__() + """The multi-layer perceptron consisting of an input layer with 3 * self.embedding_dim neurons, a hidden layer + with self.embedding_dim neurons and output layer with one neuron. + The input is represented by the concatenation embeddings of the heads, relations and tail embeddings. + """ self.head_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) self.rel_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=True) self.tail_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) @@ -957,6 +961,8 @@ def __init__( def reset_parameters(self): for mod in self.modules(): + if mod is self: + continue if hasattr(mod, 'reset_parameters'): mod.reset_parameters() From aa17492c9b9f6edc63fa89b82e671b3d3acb97bd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 16 Sep 2020 10:21:09 +0200 Subject: [PATCH 018/690] Modify ER-MLP to use InteractionFunction --- src/pykeen/models/unimodal/ermlp.py | 147 +++++----------------------- 1 file changed, 25 insertions(+), 122 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 7cfd6e730a..2821ce0d2e 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -4,48 +4,19 @@ from typing import Optional -import torch import torch.autograd -from torch import nn -from ..base import EntityRelationEmbeddingModel +from ..base import ERMLPInteractionFunction, EntityRelationEmbeddingModel from ...losses import Loss from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...utils import get_embedding_in_canonical_shape __all__ = [ 'ERMLP', ] -def score( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - mlp: nn.Module, -) -> torch.FloatTensor: - r""" - Evaluate the ER-MLP interaction function. - - .. math:: - - f(h,r,t) = \textbf{w}^{T} g(\textbf{W} [\textbf{h}; \textbf{r}; \textbf{t}]), - - :param h: shape: (b, d) - The head entity embeddings. - :param r: shape: (b, d) - The relation embeddings. - :param t: shape: (b, d) - The tail entity embeddings. - :param mlp: - The MLP. - - :return: shape: (b,) - The scores. - """ - return mlp(torch.cat([h, r, t], dim=-1)) - - class ERMLP(EntityRelationEmbeddingModel): r"""An implementation of ERMLP from [dong2014]_. @@ -89,20 +60,11 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) - if hidden_dim is None: hidden_dim = embedding_dim - self.hidden_dim = hidden_dim - """The multi-layer perceptron consisting of an input layer with 3 * self.embedding_dim neurons, a hidden layer - with self.embedding_dim neurons and output layer with one neuron. - The input is represented by the concatenation embeddings of the heads, relations and tail embeddings. - """ - self.linear1 = nn.Linear(3 * self.embedding_dim, self.hidden_dim) - self.linear2 = nn.Linear(self.hidden_dim, 1) - self.mlp = nn.Sequential( - self.linear1, - nn.ReLU(), - self.linear2, + self.interaction_function = ERMLPInteractionFunction( + embedding_dim=embedding_dim, + hidden_dim=hidden_dim, ) # Finalize initialization @@ -112,91 +74,32 @@ def _reset_parameters_(self): # noqa: D102 # The authors do not specify which initialization was used. Hence, we use the pytorch default. self.entity_embeddings.reset_parameters() self.relation_embeddings.reset_parameters() + self.interaction_function.reset_parameters() - # weight initialization - nn.init.zeros_(self.linear1.bias) - nn.init.xavier_uniform_(self.linear1.weight) - nn.init.zeros_(self.linear2.bias) - nn.init.xavier_uniform_(self.linear2.weight, gain=nn.init.calculate_gain('relu')) - - def score_hrt( + def _score( self, - hrt_batch: torch.LongTensor, - ) -> torch.FloatTensor: # noqa: D102 + h_ind: Optional[torch.LongTensor] = None, + r_ind: Optional[torch.LongTensor] = None, + t_ind: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: # Get embeddings - h = self.entity_embeddings(hrt_batch[:, 0]) - r = self.relation_embeddings(hrt_batch[:, 1]) - t = self.entity_embeddings(hrt_batch[:, 2]) + h = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=h_ind) + r = get_embedding_in_canonical_shape(embedding=self.relation_embeddings, ind=r_ind) + t = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=t_ind) - # Embedding Regularization - self.regularize_if_necessary(h, r, t) + # Compute score + scores = self.interaction_function(h=h, r=r, t=t) - # Apply interaction function - return score(h=h, r=r, t=t, mlp=self.mlp) + # Only regularize relation embeddings + self.regularize_if_necessary(r) - def score_t( - self, - hr_batch: torch.LongTensor, - h: Optional[torch.FloatTensor] = None, - r: Optional[torch.FloatTensor] = None, - t: Optional[torch.FloatTensor] = None, - ) -> torch.FloatTensor: # noqa: D102 - if h is None and r is None and t is None: - # Get embeddings - h = self.entity_embeddings(hr_batch[:, 0]) - r = self.relation_embeddings(hr_batch[:, 1]) - t = self.entity_embeddings.weight - - # Embedding Regularization - self.regularize_if_necessary(h, r, t) - - # First layer can be unrolled - layers = self.mlp.children() - first_layer = next(layers) - w = first_layer.weight - i = 2 * self.embedding_dim - w_hr = w[None, :, :i] @ torch.cat([h, r], dim=-1).unsqueeze(-1) - w_t = w[None, :, i:] @ t.unsqueeze(-1) - b = first_layer.bias - scores = (b[None, None, :] + w_hr[:, None, :, 0]) + w_t[None, :, :, 0] - - # Send scores through rest of the network - scores = scores.view(-1, self.hidden_dim) - for remaining_layer in layers: - scores = remaining_layer(scores) - scores = scores.view(-1, self.num_entities) return scores - def score_h( - self, - rt_batch: torch.LongTensor, - h: Optional[torch.FloatTensor] = None, - r: Optional[torch.FloatTensor] = None, - t: Optional[torch.FloatTensor] = None, - ) -> torch.FloatTensor: # noqa: D102 - if h is None and r is None and t is None: - # Get embeddings - h = self.entity_embeddings.weight - r = self.relation_embeddings(rt_batch[:, 0]) - t = self.entity_embeddings(rt_batch[:, 1]) - - # Embedding Regularization - self.regularize_if_necessary(h, r, t) - - # First layer can be unrolled - layers = self.mlp.children() - first_layer = next(layers) - w = first_layer.weight - i = self.embedding_dim - w_h = w[None, :, :i] @ h.unsqueeze(-1) - w_rt = w[None, :, i:] @ torch.cat([r, t], dim=-1).unsqueeze(-1) - b = first_layer.bias - scores = (b[None, None, :] + w_rt[:, None, :, 0]) + w_h[None, :, :, 0] - - # Send scores through rest of the network - scores = scores.view(-1, self.hidden_dim) - for remaining_layer in layers: - scores = remaining_layer(scores) - scores = scores.view(-1, self.num_entities) + def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=hrt_batch[:, 0], r_ind=hrt_batch[:, 1], t_ind=hrt_batch[:, 2]).view(-1, 1) - return scores + def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=hr_batch[:, 0], r_ind=hr_batch[:, 1], t_ind=None).view(-1, self.num_entities) + + def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) From ce2964c81cbb93ab2ee54bba5766a7a7e0101d95 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 16 Sep 2020 10:22:59 +0200 Subject: [PATCH 019/690] Move reset_parameters to base class InteractionFunction --- src/pykeen/models/base.py | 15 ++++++++------- 1 file changed, 8 insertions(+), 7 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 9b89b282cc..c962e0667b 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -887,6 +887,14 @@ def score_t( return scores + def reset_parameters(self): + """Reset parameters the interaction function may have.""" + for mod in self.modules(): + if mod is self: + continue + if hasattr(mod, 'reset_parameters'): + mod.reset_parameters() + def _normalize_for_einsum( x: torch.FloatTensor, @@ -959,13 +967,6 @@ def __init__( self.activation = nn.ReLU() self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) - def reset_parameters(self): - for mod in self.modules(): - if mod is self: - continue - if hasattr(mod, 'reset_parameters'): - mod.reset_parameters() - def forward( self, h: torch.FloatTensor, From 85a8f6b6c31761fe85bf926ea2024d137595c618 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 15 Oct 2020 08:53:31 +0200 Subject: [PATCH 020/690] Move interaction functions to the respective modules --- src/pykeen/models/base.py | 86 -------------------------- src/pykeen/models/unimodal/distmult.py | 44 ++++++++++++- src/pykeen/models/unimodal/ermlp.py | 50 ++++++++++++++- 3 files changed, 91 insertions(+), 89 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 26a4d900d5..e1c134c26a 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1204,89 +1204,3 @@ def reset_parameters(self): continue if hasattr(mod, 'reset_parameters'): mod.reset_parameters() - - -def _normalize_for_einsum( - x: torch.FloatTensor, - batch_size: int, - symbol: str, -) -> Tuple[str, torch.FloatTensor]: - """ - Normalize tensor for broadcasting along batch-dimension in einsum. - - :param x: - The tensor. - :param batch_size: - The batch_size - :param symbol: - The symbol for the einsum term. - - :return: - A tuple (reshaped_tensor, term). - """ - if x.shape[0] == batch_size: - return f'b{symbol}d', x - return f'{symbol}d', x.squeeze(dim=0) - - -class DistMultInteractionFunction(InteractionFunction): - """Interaction function of DistMult.""" - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa: D102 - batch_size = max(h.shape[0], r.shape[0], t.shape[0]) - h_term, h = _normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') - r_term, r = _normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') - t_term, t = _normalize_for_einsum(x=t, batch_size=batch_size, symbol='t') - return torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', h, r, t) - - -class ERMLPInteractionFunction(InteractionFunction): - """ - Interaction function of ER-MLP. - - .. math :: - f(h, r, t) = W_2 ReLU(W_1 cat(h, r, t) + b_1) + b_2 - """ - - def __init__( - self, - embedding_dim: int, - hidden_dim: int, - ): - """ - Initialize the interaction function. - - :param embedding_dim: - The embedding vector dimension. - :param hidden_dim: - The hidden dimension of the MLP. - """ - super().__init__() - """The multi-layer perceptron consisting of an input layer with 3 * self.embedding_dim neurons, a hidden layer - with self.embedding_dim neurons and output layer with one neuron. - The input is represented by the concatenation embeddings of the heads, relations and tail embeddings. - """ - self.head_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) - self.rel_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=True) - self.tail_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) - self.activation = nn.ReLU() - self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa: D102 - h = self.head_to_hidden(h) - r = self.rel_to_hidden(r) - t = self.tail_to_hidden(t) - # TODO: Choosing which to combine first, h/r, h/t or r/t, depending on the shape might further improve - # performance in a 1:n scenario. - x = self.activation(h[:, :, None, None, :] + r[:, None, :, None, :] + t[:, None, None, :, :]) - return self.hidden_to_score(x).squeeze(dim=-1) diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 03583ad796..a5c3d64e92 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -2,13 +2,14 @@ """Implementation of DistMult.""" -from typing import Optional +from typing import Optional, Tuple +import torch import torch.autograd from torch import nn from torch.nn import functional -from ..base import DistMultInteractionFunction, EntityRelationEmbeddingModel +from ..base import EntityRelationEmbeddingModel, InteractionFunction from ...losses import Loss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory @@ -135,3 +136,42 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) + + +def _normalize_for_einsum( + x: torch.FloatTensor, + batch_size: int, + symbol: str, +) -> Tuple[str, torch.FloatTensor]: + """ + Normalize tensor for broadcasting along batch-dimension in einsum. + + :param x: + The tensor. + :param batch_size: + The batch_size + :param symbol: + The symbol for the einsum term. + + :return: + A tuple (reshaped_tensor, term). + """ + if x.shape[0] == batch_size: + return f'b{symbol}d', x + return f'{symbol}d', x.squeeze(dim=0) + + +class DistMultInteractionFunction(InteractionFunction): + """Interaction function of DistMult.""" + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + ) -> torch.FloatTensor: # noqa: D102 + batch_size = max(h.shape[0], r.shape[0], t.shape[0]) + h_term, h = _normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') + r_term, r = _normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') + t_term, t = _normalize_for_einsum(x=t, batch_size=batch_size, symbol='t') + return torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', h, r, t) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 2821ce0d2e..03d4e0181a 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -5,8 +5,9 @@ from typing import Optional import torch.autograd +from torch import nn -from ..base import ERMLPInteractionFunction, EntityRelationEmbeddingModel +from ..base import EntityRelationEmbeddingModel, InteractionFunction from ...losses import Loss from ...regularizers import Regularizer from ...triples import TriplesFactory @@ -103,3 +104,50 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) + + +class ERMLPInteractionFunction(InteractionFunction): + """ + Interaction function of ER-MLP. + + .. math :: + f(h, r, t) = W_2 ReLU(W_1 cat(h, r, t) + b_1) + b_2 + """ + + def __init__( + self, + embedding_dim: int, + hidden_dim: int, + ): + """ + Initialize the interaction function. + + :param embedding_dim: + The embedding vector dimension. + :param hidden_dim: + The hidden dimension of the MLP. + """ + super().__init__() + """The multi-layer perceptron consisting of an input layer with 3 * self.embedding_dim neurons, a hidden layer + with self.embedding_dim neurons and output layer with one neuron. + The input is represented by the concatenation embeddings of the heads, relations and tail embeddings. + """ + self.head_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) + self.rel_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=True) + self.tail_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) + self.activation = nn.ReLU() + self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + ) -> torch.FloatTensor: # noqa: D102 + h = self.head_to_hidden(h) + r = self.rel_to_hidden(r) + t = self.tail_to_hidden(t) + # TODO: Choosing which to combine first, h/r, h/t or r/t, depending on the shape might further improve + # performance in a 1:n scenario. + x = self.activation(h[:, :, None, None, :] + r[:, None, :, None, :] + t[:, None, None, :, :]) + return self.hidden_to_score(x).squeeze(dim=-1) From 3f1e7e2467a490b5852f93d83de7e24a87908658 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 15 Oct 2020 08:53:51 +0200 Subject: [PATCH 021/690] Fix import order --- src/pykeen/models/unimodal/distmult.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index a5c3d64e92..c08446e842 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -13,13 +13,12 @@ from ...losses import Loss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory +from ...utils import get_embedding_in_canonical_shape __all__ = [ 'DistMult', ] -from ...utils import get_embedding_in_canonical_shape - class DistMult(EntityRelationEmbeddingModel): r"""An implementation of DistMult from [yang2014]_. From 724119f2b60f95d7aa9383f0b2dd9c438b49b1cc Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 15 Oct 2020 08:54:15 +0200 Subject: [PATCH 022/690] Add interaction function to __all__ --- src/pykeen/models/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index e1c134c26a..0039c933a3 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -26,6 +26,7 @@ 'Model', 'EntityEmbeddingModel', 'EntityRelationEmbeddingModel', + 'InteractionFunction', 'MultimodalModel', ] From 6b66aecf72727b1b6b060ec621503513fee1e371 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 15 Oct 2020 08:56:33 +0200 Subject: [PATCH 023/690] Change code order --- src/pykeen/models/unimodal/distmult.py | 81 +++++++++++++------------- 1 file changed, 40 insertions(+), 41 deletions(-) diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index c08446e842..07dc615b01 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -2,9 +2,8 @@ """Implementation of DistMult.""" -from typing import Optional, Tuple +from typing import Optional -import torch import torch.autograd from torch import nn from torch.nn import functional @@ -20,6 +19,45 @@ ] +def _normalize_for_einsum( + x: torch.FloatTensor, + batch_size: int, + symbol: str, +) -> Tuple[str, torch.FloatTensor]: + """ + Normalize tensor for broadcasting along batch-dimension in einsum. + + :param x: + The tensor. + :param batch_size: + The batch_size + :param symbol: + The symbol for the einsum term. + + :return: + A tuple (reshaped_tensor, term). + """ + if x.shape[0] == batch_size: + return f'b{symbol}d', x + return f'{symbol}d', x.squeeze(dim=0) + + +class DistMultInteractionFunction(InteractionFunction): + """Interaction function of DistMult.""" + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + ) -> torch.FloatTensor: # noqa: D102 + batch_size = max(h.shape[0], r.shape[0], t.shape[0]) + h_term, h = _normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') + r_term, r = _normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') + t_term, t = _normalize_for_einsum(x=t, batch_size=batch_size, symbol='t') + return torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', h, r, t) + + class DistMult(EntityRelationEmbeddingModel): r"""An implementation of DistMult from [yang2014]_. @@ -135,42 +173,3 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) - - -def _normalize_for_einsum( - x: torch.FloatTensor, - batch_size: int, - symbol: str, -) -> Tuple[str, torch.FloatTensor]: - """ - Normalize tensor for broadcasting along batch-dimension in einsum. - - :param x: - The tensor. - :param batch_size: - The batch_size - :param symbol: - The symbol for the einsum term. - - :return: - A tuple (reshaped_tensor, term). - """ - if x.shape[0] == batch_size: - return f'b{symbol}d', x - return f'{symbol}d', x.squeeze(dim=0) - - -class DistMultInteractionFunction(InteractionFunction): - """Interaction function of DistMult.""" - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa: D102 - batch_size = max(h.shape[0], r.shape[0], t.shape[0]) - h_term, h = _normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') - r_term, r = _normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') - t_term, t = _normalize_for_einsum(x=t, batch_size=batch_size, symbol='t') - return torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', h, r, t) From 25e91d5f958c0d37a500226e567efd09134f544c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 15 Oct 2020 08:56:42 +0200 Subject: [PATCH 024/690] Export interaction function --- src/pykeen/models/unimodal/distmult.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 07dc615b01..966acf638e 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -16,6 +16,7 @@ __all__ = [ 'DistMult', + 'DistMultInteractionFunction', ] From ab5c80f187af115c2e51bddec3f07d17e167b488 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 15 Oct 2020 08:58:17 +0200 Subject: [PATCH 025/690] Move utility method --- src/pykeen/models/base.py | 23 +++++++++++++++++++ src/pykeen/models/unimodal/distmult.py | 31 ++++---------------------- 2 files changed, 27 insertions(+), 27 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 0039c933a3..7a6ae4960b 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1205,3 +1205,26 @@ def reset_parameters(self): continue if hasattr(mod, 'reset_parameters'): mod.reset_parameters() + + +def normalize_for_einsum( + x: torch.FloatTensor, + batch_size: int, + symbol: str, +) -> Tuple[str, torch.FloatTensor]: + """ + Normalize tensor for broadcasting along batch-dimension in einsum. + + :param x: + The tensor. + :param batch_size: + The batch_size + :param symbol: + The symbol for the einsum term. + + :return: + A tuple (reshaped_tensor, term). + """ + if x.shape[0] == batch_size: + return f'b{symbol}d', x + return f'{symbol}d', x.squeeze(dim=0) \ No newline at end of file diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 966acf638e..0ed16458cb 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -8,7 +8,7 @@ from torch import nn from torch.nn import functional -from ..base import EntityRelationEmbeddingModel, InteractionFunction +from ..base import EntityRelationEmbeddingModel, InteractionFunction, normalize_for_einsum from ...losses import Loss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory @@ -20,29 +20,6 @@ ] -def _normalize_for_einsum( - x: torch.FloatTensor, - batch_size: int, - symbol: str, -) -> Tuple[str, torch.FloatTensor]: - """ - Normalize tensor for broadcasting along batch-dimension in einsum. - - :param x: - The tensor. - :param batch_size: - The batch_size - :param symbol: - The symbol for the einsum term. - - :return: - A tuple (reshaped_tensor, term). - """ - if x.shape[0] == batch_size: - return f'b{symbol}d', x - return f'{symbol}d', x.squeeze(dim=0) - - class DistMultInteractionFunction(InteractionFunction): """Interaction function of DistMult.""" @@ -53,9 +30,9 @@ def forward( t: torch.FloatTensor, ) -> torch.FloatTensor: # noqa: D102 batch_size = max(h.shape[0], r.shape[0], t.shape[0]) - h_term, h = _normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') - r_term, r = _normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') - t_term, t = _normalize_for_einsum(x=t, batch_size=batch_size, symbol='t') + h_term, h = normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') + r_term, r = normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') + t_term, t = normalize_for_einsum(x=t, batch_size=batch_size, symbol='t') return torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', h, r, t) From a1591d64339314d6c3c2a4d25211ae4118ea3d07 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 15 Oct 2020 09:04:07 +0200 Subject: [PATCH 026/690] Extract complex interaction function --- src/pykeen/models/unimodal/complex.py | 94 ++++++++++++++------------- 1 file changed, 48 insertions(+), 46 deletions(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 83ed91db97..f8fbfad78d 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -7,7 +7,7 @@ import torch import torch.nn as nn -from ..base import EntityRelationEmbeddingModel +from ..base import EntityRelationEmbeddingModel, InteractionFunction, normalize_for_einsum from ...losses import Loss, SoftplusLoss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory @@ -15,9 +15,35 @@ __all__ = [ 'ComplEx', + 'ComplexInteractionFunction', ] +class ComplexInteractionFunction(InteractionFunction): + """Interaction function of Complex.""" + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + ) -> torch.FloatTensor: # noqa: D102 + batch_size = max(h.shape[0], r.shape[0], t.shape[0]) + h_term, h = normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') + r_term, r = normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') + t_term, t = normalize_for_einsum(x=t, batch_size=batch_size, symbol='t') + (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] + return sum( + torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', hh, rr, tt) + for hh, rr, tt in [ + (h_re, r_re, t_re), + (h_re, r_im, t_im), + (h_im, r_re, t_im), + (h_im, r_im, t_re), + ] + ) + + class ComplEx(EntityRelationEmbeddingModel): r"""An implementation of ComplEx [trouillon2016]_. @@ -103,6 +129,8 @@ def __init__( regularizer=regularizer, ) + self.interaction_function = ComplexInteractionFunction() + # Finalize initialization self.reset_parameters_() @@ -112,56 +140,30 @@ def _reset_parameters_(self): # noqa: D102 nn.init.normal_(tensor=self.entity_embeddings.weight, mean=0., std=1.) nn.init.normal_(tensor=self.relation_embeddings.weight, mean=0., std=1.) - @staticmethod - def interaction_function( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, + def _score( + self, + h_ind: Optional[torch.LongTensor] = None, + r_ind: Optional[torch.LongTensor] = None, + t_ind: Optional[torch.LongTensor] = None, ) -> torch.FloatTensor: - """Evaluate the interaction function of ComplEx for given embeddings. - - The embeddings have to be in a broadcastable shape. + # Get embeddings + h = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=h_ind) + r = get_embedding_in_canonical_shape(embedding=self.relation_embeddings, ind=r_ind) + t = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=t_ind) - :param h: - Head embeddings. - :param r: - Relation embeddings. - :param t: - Tail embeddings. + # Compute score + scores = self.interaction_function(h=h, r=r, t=t) - :return: shape: (...) - The scores. - """ - # split into real and imaginary part - (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] + # Only regularize relation embeddings + self.regularize_if_necessary(r) - # ComplEx space bilinear product - # *: Elementwise multiplication - return sum( - (hh * rr * tt).sum(dim=-1) - for hh, rr, tt in [ - (h_re, r_re, t_re), - (h_re, r_im, t_im), - (h_im, r_re, t_im), - (h_im, r_im, t_re), - ] - ) + return scores def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # get embeddings - h, r, t = [ - get_embedding_in_canonical_shape(embedding=e, ind=ind) - for e, ind in [ - (self.entity_embeddings, hrt_batch[:, 0]), - (self.relation_embeddings, hrt_batch[:, 1]), - (self.entity_embeddings, hrt_batch[:, 2]), - ] - ] - - # Compute scores - scores = self.interaction_function(h=h, r=r, t=t) + return self._score(h_ind=hrt_batch[:, 0], r_ind=hrt_batch[:, 1], t_ind=hrt_batch[:, 2]).view(-1, 1) - # Regularization - self.regularize_if_necessary(h, r, t) + def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=hr_batch[:, 0], r_ind=hr_batch[:, 1], t_ind=None).view(-1, self.num_entities) - return scores + def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) From 4111a95028775342ca757a20bd171e6534062fdb Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 15 Oct 2020 09:04:31 +0200 Subject: [PATCH 027/690] Add missing newline --- src/pykeen/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 7a6ae4960b..1528712626 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1227,4 +1227,4 @@ def normalize_for_einsum( """ if x.shape[0] == batch_size: return f'b{symbol}d', x - return f'{symbol}d', x.squeeze(dim=0) \ No newline at end of file + return f'{symbol}d', x.squeeze(dim=0) From 025d2a5def6ca9dfcfa64f87925d87c679d0153b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 15 Oct 2020 09:10:04 +0200 Subject: [PATCH 028/690] Extract common baseclass for distmult and complex --- src/pykeen/models/unimodal/complex.py | 112 +++++++++++++++++-------- src/pykeen/models/unimodal/distmult.py | 38 +-------- 2 files changed, 82 insertions(+), 68 deletions(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index f8fbfad78d..248d2adaf7 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -19,6 +19,82 @@ ] +class SimpleVectorEntityRelationEmbeddingModel(EntityRelationEmbeddingModel): + def __init__( + self, + triples_factory: TriplesFactory, + interaction_function: InteractionFunction, + embedding_dim: int = 200, + automatic_memory_optimization: Optional[bool] = None, + loss: Optional[Loss] = None, + preferred_device: Optional[str] = None, + random_seed: Optional[int] = None, + regularizer: Optional[Regularizer] = None, + ) -> None: + """Initialize ComplEx. + + :param triples_factory: TriplesFactory + The triple factory connected to the model. + :param interaction_function: + The interaction function used to compute scores. + :param embedding_dim: + The embedding dimensionality of the entity embeddings. + :param automatic_memory_optimization: bool + Whether to automatically optimize the sub-batch size during training and batch size during evaluation with + regards to the hardware at hand. + :param loss: OptionalLoss (optional) + The loss to use. Defaults to SoftplusLoss. + :param preferred_device: str (optional) + The default device where to model is located. + :param random_seed: int (optional) + An optional random seed to set before the initialization of weights. + :param regularizer: BaseRegularizer + The regularizer to use. + """ + super().__init__( + triples_factory=triples_factory, + embedding_dim=embedding_dim, + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + random_seed=random_seed, + regularizer=regularizer, + ) + + self.interaction_function = interaction_function + + # Finalize initialization + self.reset_parameters_() + + def _score( + self, + h_ind: Optional[torch.LongTensor] = None, + r_ind: Optional[torch.LongTensor] = None, + t_ind: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: + # Get embeddings + h = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=h_ind) + r = get_embedding_in_canonical_shape(embedding=self.relation_embeddings, ind=r_ind) + t = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=t_ind) + + # Compute score + scores = self.interaction_function(h=h, r=r, t=t) + + # Only regularize relation embeddings + self.regularize_if_necessary(r) + + return scores + + def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=hrt_batch[:, 0], r_ind=hrt_batch[:, 1], t_ind=hrt_batch[:, 2]).view(-1, 1) + + def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=hr_batch[:, 0], r_ind=hr_batch[:, 1], t_ind=None).view(-1, self.num_entities) + + def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) + + class ComplexInteractionFunction(InteractionFunction): """Interaction function of Complex.""" @@ -44,7 +120,7 @@ def forward( ) -class ComplEx(EntityRelationEmbeddingModel): +class ComplEx(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of ComplEx [trouillon2016]_. ComplEx is an extension of :class:`pykeen.models.DistMult` that uses complex valued representations for the @@ -121,6 +197,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, + interaction_function=ComplexInteractionFunction(), embedding_dim=2 * embedding_dim, # complex embeddings automatic_memory_optimization=automatic_memory_optimization, loss=loss, @@ -129,41 +206,8 @@ def __init__( regularizer=regularizer, ) - self.interaction_function = ComplexInteractionFunction() - - # Finalize initialization - self.reset_parameters_() - def _reset_parameters_(self): # noqa: D102 # initialize with entity and relation embeddings with standard normal distribution, cf. # https://github.com/ttrouill/complex/blob/dc4eb93408d9a5288c986695b58488ac80b1cc17/efe/models.py#L481-L487 nn.init.normal_(tensor=self.entity_embeddings.weight, mean=0., std=1.) nn.init.normal_(tensor=self.relation_embeddings.weight, mean=0., std=1.) - - def _score( - self, - h_ind: Optional[torch.LongTensor] = None, - r_ind: Optional[torch.LongTensor] = None, - t_ind: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: - # Get embeddings - h = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=h_ind) - r = get_embedding_in_canonical_shape(embedding=self.relation_embeddings, ind=r_ind) - t = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=t_ind) - - # Compute score - scores = self.interaction_function(h=h, r=r, t=t) - - # Only regularize relation embeddings - self.regularize_if_necessary(r) - - return scores - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=hrt_batch[:, 0], r_ind=hrt_batch[:, 1], t_ind=hrt_batch[:, 2]).view(-1, 1) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=hr_batch[:, 0], r_ind=hr_batch[:, 1], t_ind=None).view(-1, self.num_entities) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 0ed16458cb..4eef0cb674 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -8,11 +8,11 @@ from torch import nn from torch.nn import functional -from ..base import EntityRelationEmbeddingModel, InteractionFunction, normalize_for_einsum +from .complex import SimpleVectorEntityRelationEmbeddingModel +from ..base import InteractionFunction, normalize_for_einsum from ...losses import Loss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory -from ...utils import get_embedding_in_canonical_shape __all__ = [ 'DistMult', @@ -36,7 +36,7 @@ def forward( return torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', h, r, t) -class DistMult(EntityRelationEmbeddingModel): +class DistMult(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of DistMult from [yang2014]_. This model simplifies RESCAL by restricting matrices representing relations as diagonal matrices. @@ -98,6 +98,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, + interaction_function=DistMultInteractionFunction(), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, @@ -105,9 +106,6 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) - self.interaction_function = DistMultInteractionFunction() - # Finalize initialization - self.reset_parameters_() def _reset_parameters_(self): # noqa: D102 # xavier uniform, cf. @@ -123,31 +121,3 @@ def post_parameter_update(self) -> None: # noqa: D102 # Normalize embeddings of entities functional.normalize(self.entity_embeddings.weight.data, out=self.entity_embeddings.weight.data) - - def _score( - self, - h_ind: Optional[torch.LongTensor] = None, - r_ind: Optional[torch.LongTensor] = None, - t_ind: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: - # Get embeddings - h = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=h_ind) - r = get_embedding_in_canonical_shape(embedding=self.relation_embeddings, ind=r_ind) - t = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=t_ind) - - # Compute score - scores = self.interaction_function(h=h, r=r, t=t) - - # Only regularize relation embeddings - self.regularize_if_necessary(r) - - return scores - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=hrt_batch[:, 0], r_ind=hrt_batch[:, 1], t_ind=hrt_batch[:, 2]).view(-1, 1) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=hr_batch[:, 0], r_ind=hr_batch[:, 1], t_ind=None).view(-1, self.num_entities) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) From dbcf02171b2834bc19d60b9287e67190994f693a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 15 Oct 2020 09:11:56 +0200 Subject: [PATCH 029/690] Fix docstring --- src/pykeen/models/unimodal/complex.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 248d2adaf7..ddd0dd129e 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -20,6 +20,8 @@ class SimpleVectorEntityRelationEmbeddingModel(EntityRelationEmbeddingModel): + """A base class for embedding models which store a single vector for each entity and relation.""" + def __init__( self, triples_factory: TriplesFactory, @@ -31,7 +33,7 @@ def __init__( random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: - """Initialize ComplEx. + """Initialize embedding model. :param triples_factory: TriplesFactory The triple factory connected to the model. @@ -43,7 +45,7 @@ def __init__( Whether to automatically optimize the sub-batch size during training and batch size during evaluation with regards to the hardware at hand. :param loss: OptionalLoss (optional) - The loss to use. Defaults to SoftplusLoss. + The loss to use. :param preferred_device: str (optional) The default device where to model is located. :param random_seed: int (optional) From 52d7cda8c4ab926a4e2152813126ce0871d1d0a1 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 15 Oct 2020 09:13:33 +0200 Subject: [PATCH 030/690] Use base class for ER-MLP --- src/pykeen/models/unimodal/ermlp.py | 147 +++++++++++----------------- 1 file changed, 59 insertions(+), 88 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 03d4e0181a..e193439970 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -4,21 +4,69 @@ from typing import Optional -import torch.autograd +import torch from torch import nn -from ..base import EntityRelationEmbeddingModel, InteractionFunction +from .complex import SimpleVectorEntityRelationEmbeddingModel +from ..base import InteractionFunction from ...losses import Loss from ...regularizers import Regularizer from ...triples import TriplesFactory -from ...utils import get_embedding_in_canonical_shape __all__ = [ 'ERMLP', + 'ERMLPInteractionFunction', ] -class ERMLP(EntityRelationEmbeddingModel): +class ERMLPInteractionFunction(InteractionFunction): + """ + Interaction function of ER-MLP. + + .. math :: + f(h, r, t) = W_2 ReLU(W_1 cat(h, r, t) + b_1) + b_2 + """ + + def __init__( + self, + embedding_dim: int, + hidden_dim: int, + ): + """ + Initialize the interaction function. + + :param embedding_dim: + The embedding vector dimension. + :param hidden_dim: + The hidden dimension of the MLP. + """ + super().__init__() + """The multi-layer perceptron consisting of an input layer with 3 * self.embedding_dim neurons, a hidden layer + with self.embedding_dim neurons and output layer with one neuron. + The input is represented by the concatenation embeddings of the heads, relations and tail embeddings. + """ + self.head_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) + self.rel_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=True) + self.tail_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) + self.activation = nn.ReLU() + self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + ) -> torch.FloatTensor: # noqa: D102 + h = self.head_to_hidden(h) + r = self.rel_to_hidden(r) + t = self.tail_to_hidden(t) + # TODO: Choosing which to combine first, h/r, h/t or r/t, depending on the shape might further improve + # performance in a 1:n scenario. + x = self.activation(h[:, :, None, None, :] + r[:, None, :, None, :] + t[:, None, None, :, :]) + return self.hidden_to_score(x).squeeze(dim=-1) + + +class ERMLP(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of ERMLP from [dong2014]_. ERMLP is a multi-layer perceptron based approach that uses a single hidden layer and represents entities and @@ -52,8 +100,15 @@ def __init__( regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model.""" + if hidden_dim is None: + hidden_dim = embedding_dim + interaction_function = ERMLPInteractionFunction( + embedding_dim=embedding_dim, + hidden_dim=hidden_dim, + ) super().__init__( triples_factory=triples_factory, + interaction_function=interaction_function, embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, @@ -61,93 +116,9 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) - if hidden_dim is None: - hidden_dim = embedding_dim - self.interaction_function = ERMLPInteractionFunction( - embedding_dim=embedding_dim, - hidden_dim=hidden_dim, - ) - - # Finalize initialization - self.reset_parameters_() def _reset_parameters_(self): # noqa: D102 # The authors do not specify which initialization was used. Hence, we use the pytorch default. self.entity_embeddings.reset_parameters() self.relation_embeddings.reset_parameters() self.interaction_function.reset_parameters() - - def _score( - self, - h_ind: Optional[torch.LongTensor] = None, - r_ind: Optional[torch.LongTensor] = None, - t_ind: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: - # Get embeddings - h = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=h_ind) - r = get_embedding_in_canonical_shape(embedding=self.relation_embeddings, ind=r_ind) - t = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=t_ind) - - # Compute score - scores = self.interaction_function(h=h, r=r, t=t) - - # Only regularize relation embeddings - self.regularize_if_necessary(r) - - return scores - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=hrt_batch[:, 0], r_ind=hrt_batch[:, 1], t_ind=hrt_batch[:, 2]).view(-1, 1) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=hr_batch[:, 0], r_ind=hr_batch[:, 1], t_ind=None).view(-1, self.num_entities) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) - - -class ERMLPInteractionFunction(InteractionFunction): - """ - Interaction function of ER-MLP. - - .. math :: - f(h, r, t) = W_2 ReLU(W_1 cat(h, r, t) + b_1) + b_2 - """ - - def __init__( - self, - embedding_dim: int, - hidden_dim: int, - ): - """ - Initialize the interaction function. - - :param embedding_dim: - The embedding vector dimension. - :param hidden_dim: - The hidden dimension of the MLP. - """ - super().__init__() - """The multi-layer perceptron consisting of an input layer with 3 * self.embedding_dim neurons, a hidden layer - with self.embedding_dim neurons and output layer with one neuron. - The input is represented by the concatenation embeddings of the heads, relations and tail embeddings. - """ - self.head_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) - self.rel_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=True) - self.tail_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) - self.activation = nn.ReLU() - self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa: D102 - h = self.head_to_hidden(h) - r = self.rel_to_hidden(r) - t = self.tail_to_hidden(t) - # TODO: Choosing which to combine first, h/r, h/t or r/t, depending on the shape might further improve - # performance in a 1:n scenario. - x = self.activation(h[:, :, None, None, :] + r[:, None, :, None, :] + t[:, None, None, :, :]) - return self.hidden_to_score(x).squeeze(dim=-1) From c3706ccb2ca44e39332bfc616c7a14cf74e90a8a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 15 Oct 2020 09:22:17 +0200 Subject: [PATCH 031/690] Fix ER-MLP initialization --- src/pykeen/models/unimodal/ermlp.py | 22 +++++++++++++++++++++- 1 file changed, 21 insertions(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index e193439970..84bdd67510 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -1,7 +1,7 @@ # -*- coding: utf-8 -*- """Implementation of ERMLP.""" - +import math from typing import Optional import torch @@ -65,6 +65,26 @@ def forward( x = self.activation(h[:, :, None, None, :] + r[:, None, :, None, :] + t[:, None, None, :, :]) return self.hidden_to_score(x).squeeze(dim=-1) + def reset_parameters(self): + # Initialize biases with zero + nn.init.zeros_(self.rel_to_hidden.bias) + nn.init.zeros_(self.hidden_to_score.bias) + # In the original formulation, + # W_2 sigma(W_1 cat([h, r, t]) + b_1) + b_2 + # W_1 would be initialized with nn.init.xavier_uniform, i.e. with a samples from uniform(-a, a) with + # a = math.sqrt(3.0) * gain * math.sqrt(2.0 / float(fan_in + fan_out)) + # we have: + # fan_out = hidden_dim + # fan_in = 3 * embedding_dim + bound = math.sqrt(3.0) * 1 * math.sqrt(2.0 / float(sum(self.head_to_hidden.weight.shape))) + for mod in [ + self.head_to_hidden, + self.rel_to_hidden, + self.tail_to_hidden, + ]: + nn.init.uniform_(mod.weight, -bound, bound) + nn.init.xavier_uniform_(self.hidden_to_score.weight, gain=nn.init.calculate_gain('relu')) + class ERMLP(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of ERMLP from [dong2014]_. From 78d56a456da6c288b42a9e5745b629df90af0d25 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 15 Oct 2020 09:28:12 +0200 Subject: [PATCH 032/690] Move base class to base --- src/pykeen/models/base.py | 80 ++++++++++++++++++++++++- src/pykeen/models/unimodal/complex.py | 82 +------------------------- src/pykeen/models/unimodal/distmult.py | 3 +- src/pykeen/models/unimodal/ermlp.py | 3 +- 4 files changed, 83 insertions(+), 85 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 1528712626..73c7c6d5c8 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -19,7 +19,7 @@ from ..tqdmw import tqdm from ..triples import TriplesFactory from ..typing import MappedTriples -from ..utils import NoRandomSeedNecessary, get_embedding, resolve_device, set_random_seed +from ..utils import NoRandomSeedNecessary, get_embedding, get_embedding_in_canonical_shape, resolve_device, set_random_seed from ..version import get_version __all__ = [ @@ -1228,3 +1228,81 @@ def normalize_for_einsum( if x.shape[0] == batch_size: return f'b{symbol}d', x return f'{symbol}d', x.squeeze(dim=0) + + +class SimpleVectorEntityRelationEmbeddingModel(EntityRelationEmbeddingModel): + """A base class for embedding models which store a single vector for each entity and relation.""" + + def __init__( + self, + triples_factory: TriplesFactory, + interaction_function: InteractionFunction, + embedding_dim: int = 200, + automatic_memory_optimization: Optional[bool] = None, + loss: Optional[Loss] = None, + preferred_device: Optional[str] = None, + random_seed: Optional[int] = None, + regularizer: Optional[Regularizer] = None, + ) -> None: + """Initialize embedding model. + + :param triples_factory: TriplesFactory + The triple factory connected to the model. + :param interaction_function: + The interaction function used to compute scores. + :param embedding_dim: + The embedding dimensionality of the entity embeddings. + :param automatic_memory_optimization: bool + Whether to automatically optimize the sub-batch size during training and batch size during evaluation with + regards to the hardware at hand. + :param loss: OptionalLoss (optional) + The loss to use. + :param preferred_device: str (optional) + The default device where to model is located. + :param random_seed: int (optional) + An optional random seed to set before the initialization of weights. + :param regularizer: BaseRegularizer + The regularizer to use. + """ + super().__init__( + triples_factory=triples_factory, + embedding_dim=embedding_dim, + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + random_seed=random_seed, + regularizer=regularizer, + ) + + self.interaction_function = interaction_function + + # Finalize initialization + self.reset_parameters_() + + def _score( + self, + h_ind: Optional[torch.LongTensor] = None, + r_ind: Optional[torch.LongTensor] = None, + t_ind: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: + # Get embeddings + h = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=h_ind) + r = get_embedding_in_canonical_shape(embedding=self.relation_embeddings, ind=r_ind) + t = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=t_ind) + + # Compute score + scores = self.interaction_function(h=h, r=r, t=t) + + # Only regularize relation embeddings + self.regularize_if_necessary(r) + + return scores + + def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=hrt_batch[:, 0], r_ind=hrt_batch[:, 1], t_ind=hrt_batch[:, 2]).view(-1, 1) + + def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=hr_batch[:, 0], r_ind=hr_batch[:, 1], t_ind=None).view(-1, self.num_entities) + + def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index ddd0dd129e..a3df4a7dc3 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -7,11 +7,11 @@ import torch import torch.nn as nn -from ..base import EntityRelationEmbeddingModel, InteractionFunction, normalize_for_einsum +from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel, normalize_for_einsum from ...losses import Loss, SoftplusLoss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory -from ...utils import get_embedding_in_canonical_shape, split_complex +from ...utils import split_complex __all__ = [ 'ComplEx', @@ -19,84 +19,6 @@ ] -class SimpleVectorEntityRelationEmbeddingModel(EntityRelationEmbeddingModel): - """A base class for embedding models which store a single vector for each entity and relation.""" - - def __init__( - self, - triples_factory: TriplesFactory, - interaction_function: InteractionFunction, - embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, - loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, - random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, - ) -> None: - """Initialize embedding model. - - :param triples_factory: TriplesFactory - The triple factory connected to the model. - :param interaction_function: - The interaction function used to compute scores. - :param embedding_dim: - The embedding dimensionality of the entity embeddings. - :param automatic_memory_optimization: bool - Whether to automatically optimize the sub-batch size during training and batch size during evaluation with - regards to the hardware at hand. - :param loss: OptionalLoss (optional) - The loss to use. - :param preferred_device: str (optional) - The default device where to model is located. - :param random_seed: int (optional) - An optional random seed to set before the initialization of weights. - :param regularizer: BaseRegularizer - The regularizer to use. - """ - super().__init__( - triples_factory=triples_factory, - embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, - regularizer=regularizer, - ) - - self.interaction_function = interaction_function - - # Finalize initialization - self.reset_parameters_() - - def _score( - self, - h_ind: Optional[torch.LongTensor] = None, - r_ind: Optional[torch.LongTensor] = None, - t_ind: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: - # Get embeddings - h = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=h_ind) - r = get_embedding_in_canonical_shape(embedding=self.relation_embeddings, ind=r_ind) - t = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=t_ind) - - # Compute score - scores = self.interaction_function(h=h, r=r, t=t) - - # Only regularize relation embeddings - self.regularize_if_necessary(r) - - return scores - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=hrt_batch[:, 0], r_ind=hrt_batch[:, 1], t_ind=hrt_batch[:, 2]).view(-1, 1) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=hr_batch[:, 0], r_ind=hr_batch[:, 1], t_ind=None).view(-1, self.num_entities) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) - - class ComplexInteractionFunction(InteractionFunction): """Interaction function of Complex.""" diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 4eef0cb674..042b58f8b3 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -8,8 +8,7 @@ from torch import nn from torch.nn import functional -from .complex import SimpleVectorEntityRelationEmbeddingModel -from ..base import InteractionFunction, normalize_for_einsum +from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel, normalize_for_einsum from ...losses import Loss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 84bdd67510..bbbf5a39da 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -7,8 +7,7 @@ import torch from torch import nn -from .complex import SimpleVectorEntityRelationEmbeddingModel -from ..base import InteractionFunction +from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss from ...regularizers import Regularizer from ...triples import TriplesFactory From 7d4d24ebf2a524d051c5c898ba992aa32fb59924 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 17 Oct 2020 11:22:29 +0200 Subject: [PATCH 033/690] Pass flake --- src/pykeen/models/base.py | 8 +++++--- src/pykeen/models/unimodal/ermlp.py | 2 +- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 73c7c6d5c8..dff7713ca8 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -19,7 +19,10 @@ from ..tqdmw import tqdm from ..triples import TriplesFactory from ..typing import MappedTriples -from ..utils import NoRandomSeedNecessary, get_embedding, get_embedding_in_canonical_shape, resolve_device, set_random_seed +from ..utils import ( + NoRandomSeedNecessary, get_embedding, get_embedding_in_canonical_shape, resolve_device, + set_random_seed, +) from ..version import get_version __all__ = [ @@ -1151,8 +1154,7 @@ def forward( r: torch.FloatTensor, t: torch.FloatTensor, ) -> torch.FloatTensor: - """ - Generic interaction function. + """Score the given triples. :param h: shape: (batch_size, num_heads, d_e) The head representations. diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index bbbf5a39da..98465d5c9b 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -64,7 +64,7 @@ def forward( x = self.activation(h[:, :, None, None, :] + r[:, None, :, None, :] + t[:, None, None, :, :]) return self.hidden_to_score(x).squeeze(dim=-1) - def reset_parameters(self): + def reset_parameters(self): # noqa: D102 # Initialize biases with zero nn.init.zeros_(self.rel_to_hidden.bias) nn.init.zeros_(self.hidden_to_score.bias) From 855ec462fcc1010ca5755818071beae47749469a Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 5 Nov 2020 14:20:31 +0100 Subject: [PATCH 034/690] Remove unnecessary docs --- docs/source/index.rst | 1 - docs/source/tutorial/bring_your_own_embeddings.rst | 2 -- 2 files changed, 3 deletions(-) delete mode 100644 docs/source/tutorial/bring_your_own_embeddings.rst diff --git a/docs/source/index.rst b/docs/source/index.rst index c86be16b18..15c1caeaa0 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -13,7 +13,6 @@ PyKEEN tutorial/understanding_evaluation tutorial/running_hpo tutorial/using_mlflow - tutorial/bring_your_own_embeddings tutorial/using_wandb tutorial/making_predictions tutorial/performance diff --git a/docs/source/tutorial/bring_your_own_embeddings.rst b/docs/source/tutorial/bring_your_own_embeddings.rst deleted file mode 100644 index 472a9f3ee8..0000000000 --- a/docs/source/tutorial/bring_your_own_embeddings.rst +++ /dev/null @@ -1,2 +0,0 @@ -Bring Your Own Embeddings -========================= From 0a8757fa99354b1fd59470f2c7f722dfbb131107 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 5 Nov 2020 14:40:18 +0100 Subject: [PATCH 035/690] Change instantiation of interaction models --- src/pykeen/models/base.py | 4 ++++ src/pykeen/models/unimodal/complex.py | 4 +++- src/pykeen/models/unimodal/distmult.py | 4 +++- src/pykeen/models/unimodal/ermlp.py | 14 +++++++++----- 4 files changed, 19 insertions(+), 7 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 4e576faef0..8cb3a71fad 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1183,6 +1183,10 @@ class MultimodalModel(EntityRelationEmbeddingModel): class InteractionFunction(nn.Module): """Base class for interaction functions.""" + @classmethod + def from_model(cls, module: nn.Module): + return cls() + def forward( self, h: torch.FloatTensor, diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index a3df4a7dc3..37745dfb6b 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -119,9 +119,11 @@ def __init__( :param regularizer: BaseRegularizer The regularizer to use. """ + interaction_function = ComplexInteractionFunction.from_model(self) + super().__init__( triples_factory=triples_factory, - interaction_function=ComplexInteractionFunction(), + interaction_function=interaction_function, embedding_dim=2 * embedding_dim, # complex embeddings automatic_memory_optimization=automatic_memory_optimization, loss=loss, diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 042b58f8b3..5b36a5321a 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -95,9 +95,11 @@ def __init__( :param embedding_dim: The entity embedding dimension $d$. Is usually $d \in [50, 300]$. """ + interaction_function = DistMultInteractionFunction.from_model(self) + super().__init__( triples_factory=triples_factory, - interaction_function=DistMultInteractionFunction(), + interaction_function=interaction_function, embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 98465d5c9b..eeb2242322 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -50,6 +50,10 @@ def __init__( self.activation = nn.ReLU() self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) + @classmethod + def from_model(cls, module: 'ERMLP'): + return cls(embedding_dim=module.embedding_dim, hidden_dim=module.hidden_dim) + def forward( self, h: torch.FloatTensor, @@ -119,12 +123,12 @@ def __init__( regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model.""" + self.embedding_dim = embedding_dim if hidden_dim is None: - hidden_dim = embedding_dim - interaction_function = ERMLPInteractionFunction( - embedding_dim=embedding_dim, - hidden_dim=hidden_dim, - ) + self.hidden_dim = embedding_dim + + interaction_function = ERMLPInteractionFunction.from_model(self) + super().__init__( triples_factory=triples_factory, interaction_function=interaction_function, From f5d7d3a608637e8151c0c1702ff0e3bda943cc2e Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 5 Nov 2020 14:40:58 +0100 Subject: [PATCH 036/690] Update base.py --- src/pykeen/models/base.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 8cb3a71fad..98e12e8613 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1320,6 +1320,11 @@ def __init__( # Finalize initialization self.reset_parameters_() + def __init_subclass__(cls, auto_reset_parameters: bool = True, **kwargs): # noqa: D105 + _track_hyperparameters(cls) + if auto_reset_parameters: + _add_post_reset_parameters(cls) + def _score( self, h_ind: Optional[torch.LongTensor] = None, From f4938865d59a903c55ea90aa51f94a1ce97de4d7 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 5 Nov 2020 14:46:49 +0100 Subject: [PATCH 037/690] Flake8 --- src/pykeen/models/base.py | 29 ++++++-------------------- src/pykeen/models/unimodal/complex.py | 4 ++-- src/pykeen/models/unimodal/distmult.py | 3 ++- src/pykeen/models/unimodal/ermlp.py | 3 ++- src/pykeen/utils.py | 24 +++++++++++++++++++++ 5 files changed, 36 insertions(+), 27 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 98e12e8613..b5b6fcef17 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -30,6 +30,7 @@ 'Model', 'EntityEmbeddingModel', 'EntityRelationEmbeddingModel', + 'SimpleVectorEntityRelationEmbeddingModel', 'InteractionFunction', 'MultimodalModel', ] @@ -1185,6 +1186,11 @@ class InteractionFunction(nn.Module): @classmethod def from_model(cls, module: nn.Module): + """Instantiate the interaction function. + + Override this function if the interaction function needs to share some of the + parameters from the model. + """ return cls() def forward( @@ -1248,29 +1254,6 @@ def reset_parameters(self): mod.reset_parameters() -def normalize_for_einsum( - x: torch.FloatTensor, - batch_size: int, - symbol: str, -) -> Tuple[str, torch.FloatTensor]: - """ - Normalize tensor for broadcasting along batch-dimension in einsum. - - :param x: - The tensor. - :param batch_size: - The batch_size - :param symbol: - The symbol for the einsum term. - - :return: - A tuple (reshaped_tensor, term). - """ - if x.shape[0] == batch_size: - return f'b{symbol}d', x - return f'{symbol}d', x.squeeze(dim=0) - - class SimpleVectorEntityRelationEmbeddingModel(EntityRelationEmbeddingModel): """A base class for embedding models which store a single vector for each entity and relation.""" diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 37745dfb6b..c674fc9ff5 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -7,11 +7,11 @@ import torch import torch.nn as nn -from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel, normalize_for_einsum +from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss, SoftplusLoss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory -from ...utils import split_complex +from ...utils import normalize_for_einsum, split_complex __all__ = [ 'ComplEx', diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 5b36a5321a..4cad01e58e 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -8,10 +8,11 @@ from torch import nn from torch.nn import functional -from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel, normalize_for_einsum +from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory +from ...utils import normalize_for_einsum __all__ = [ 'DistMult', diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index eeb2242322..2f896da038 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- """Implementation of ERMLP.""" + import math from typing import Optional @@ -51,7 +52,7 @@ def __init__( self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) @classmethod - def from_model(cls, module: 'ERMLP'): + def from_model(cls, module: 'ERMLP'): # noqa:D102 return cls(embedding_dim=module.embedding_dim, hidden_dim=module.hidden_dim) def forward( diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index ed400a4b4d..47d9cf560f 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -41,6 +41,7 @@ 'NoRandomSeedNecessary', 'Result', 'fix_dataclass_init_docs', + 'normalize_for_einsum', ] logger = logging.getLogger(__name__) @@ -474,3 +475,26 @@ def random_non_negative_int() -> int: """Generate a random positive integer.""" sq = np.random.SeedSequence(np.random.randint(0, np.iinfo(np.int_).max)) return int(sq.generate_state(1)[0]) + + +def normalize_for_einsum( + x: torch.FloatTensor, + batch_size: int, + symbol: str, +) -> Tuple[str, torch.FloatTensor]: + """ + Normalize tensor for broadcasting along batch-dimension in einsum. + + :param x: + The tensor. + :param batch_size: + The batch_size + :param symbol: + The symbol for the einsum term. + + :return: + A tuple (reshaped_tensor, term). + """ + if x.shape[0] == batch_size: + return f'b{symbol}d', x + return f'{symbol}d', x.squeeze(dim=0) From 73eb7b51f1210087f23c3d2fb3ea37538931b890 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 5 Nov 2020 15:32:58 +0100 Subject: [PATCH 038/690] Update docs --- src/pykeen/models/__init__.py | 10 ++++++++-- src/pykeen/models/base.py | 2 +- 2 files changed, 9 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index 00f5d242aa..0c71eeab80 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -8,7 +8,10 @@ from typing import Mapping, Set, Type, Union -from .base import EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, MultimodalModel +from .base import ( + EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, MultimodalModel, + SimpleVectorEntityRelationEmbeddingModel, +) from .multimodal import ComplExLiteral, DistMultLiteral from .unimodal import ( ComplEx, @@ -73,7 +76,10 @@ def _recur(c): _MODELS: Set[Type[Model]] = { cls for cls in _recur(Model) - if cls not in {Model, MultimodalModel, EntityRelationEmbeddingModel, EntityEmbeddingModel} + if cls not in { + Model, MultimodalModel, EntityRelationEmbeddingModel, + EntityEmbeddingModel, SimpleVectorEntityRelationEmbeddingModel, + } } #: A mapping of models' names to their implementations diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index b5b6fcef17..42716ffa73 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1185,7 +1185,7 @@ class InteractionFunction(nn.Module): """Base class for interaction functions.""" @classmethod - def from_model(cls, module: nn.Module): + def from_model(cls, module: nn.Module) -> 'InteractionFunction': """Instantiate the interaction function. Override this function if the interaction function needs to share some of the From 0c5d4937f2c098bf86c39206f33405b3c905c53a Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 7 Nov 2020 18:16:40 +0100 Subject: [PATCH 039/690] Update base class --- src/pykeen/models/base.py | 34 +++++++++++++++++++++++++++++----- 1 file changed, 29 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 6fc105dc28..818c906447 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -20,7 +20,7 @@ from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory from ..typing import Constrainer, Initializer, MappedTriples, Normalizer -from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed +from ..utils import NoRandomSeedNecessary, get_embedding_in_canonical_shape, resolve_device, set_random_seed __all__ = [ 'Model', @@ -1285,6 +1285,18 @@ def __init__( preferred_device: Optional[str] = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, + entity_initializer: Optional[Initializer] = None, + entity_initializer_kwargs: Optional[Mapping[str, Any]] = None, + entity_normalizer: Optional[Normalizer] = None, + entity_normalizer_kwargs: Optional[Mapping[str, Any]] = None, + entity_constrainer: Optional[Constrainer] = None, + entity_constrainer_kwargs: Optional[Mapping[str, Any]] = None, + relation_initializer: Optional[Initializer] = None, + relation_initializer_kwargs: Optional[Mapping[str, Any]] = None, + relation_normalizer: Optional[Normalizer] = None, + relation_normalizer_kwargs: Optional[Mapping[str, Any]] = None, + relation_constrainer: Optional[Constrainer] = None, + relation_constrainer_kwargs: Optional[Mapping[str, Any]] = None, ) -> None: """Initialize embedding model. @@ -1314,18 +1326,30 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, + entity_initializer=entity_initializer, + entity_initializer_kwargs=entity_initializer_kwargs, + entity_normalizer=entity_normalizer, + entity_normalizer_kwargs=entity_normalizer_kwargs, + entity_constrainer=entity_constrainer, + entity_constrainer_kwargs=entity_constrainer_kwargs, + relation_initializer=relation_initializer, + relation_initializer_kwargs=relation_initializer_kwargs, + relation_normalizer=relation_normalizer, + relation_normalizer_kwargs=relation_normalizer_kwargs, + relation_constrainer=relation_constrainer, + relation_constrainer_kwargs=relation_constrainer_kwargs, ) - self.interaction_function = interaction_function - # Finalize initialization - self.reset_parameters_() - def __init_subclass__(cls, auto_reset_parameters: bool = True, **kwargs): # noqa: D105 _track_hyperparameters(cls) if auto_reset_parameters: _add_post_reset_parameters(cls) + def _reset_parameters_(self): + super()._reset_parameters_() + self.interaction_function.reset_parameters() + def _score( self, h_ind: Optional[torch.LongTensor] = None, From 8f1ca223981151008570dc937b909098dc4e9e70 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 7 Nov 2020 18:21:26 +0100 Subject: [PATCH 040/690] Remove redundant reset parameters function this is automatically taken care of --- src/pykeen/models/unimodal/ermlp.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 1be489897c..74fdbbbc36 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -139,9 +139,3 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) - - def _reset_parameters_(self): # noqa: D102 - # The authors do not specify which initialization was used. Hence, we use the pytorch default. - self.entity_embeddings.reset_parameters() - self.relation_embeddings.reset_parameters() - self.interaction_function.reset_parameters() From 057f58cb103904a63c6a5e00de49fcd94cdf53db Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 7 Nov 2020 18:21:43 +0100 Subject: [PATCH 041/690] Remove from_class from interaction model I don't like how it looks --- src/pykeen/models/base.py | 9 --------- src/pykeen/models/unimodal/complex.py | 2 +- src/pykeen/models/unimodal/distmult.py | 2 +- src/pykeen/models/unimodal/ermlp.py | 12 ++++-------- 4 files changed, 6 insertions(+), 19 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 818c906447..446adebbd5 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1202,15 +1202,6 @@ class MultimodalModel(EntityRelationEmbeddingModel): class InteractionFunction(nn.Module): """Base class for interaction functions.""" - @classmethod - def from_model(cls, module: nn.Module) -> 'InteractionFunction': - """Instantiate the interaction function. - - Override this function if the interaction function needs to share some of the - parameters from the model. - """ - return cls() - def forward( self, h: torch.FloatTensor, diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index efe5b38a50..5183ed8b46 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -119,7 +119,7 @@ def __init__( :param regularizer: BaseRegularizer The regularizer to use. """ - interaction_function = ComplexInteractionFunction.from_model(self) + interaction_function = ComplexInteractionFunction() super().__init__( triples_factory=triples_factory, diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 3d95db704e..9409308a82 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -97,7 +97,7 @@ def __init__( :param embedding_dim: The entity embedding dimension $d$. Is usually $d \in [50, 300]$. """ - interaction_function = DistMultInteractionFunction.from_model(self) + interaction_function = DistMultInteractionFunction() super().__init__( triples_factory=triples_factory, diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 74fdbbbc36..3065f1f45e 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -51,10 +51,6 @@ def __init__( self.activation = nn.ReLU() self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) - @classmethod - def from_model(cls, module: 'ERMLP'): # noqa:D102 - return cls(embedding_dim=module.embedding_dim, hidden_dim=module.hidden_dim) - def forward( self, h: torch.FloatTensor, @@ -124,10 +120,10 @@ def __init__( regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model.""" - if hidden_dim is None: - self.hidden_dim = embedding_dim - - interaction_function = ERMLPInteractionFunction.from_model(self) + interaction_function = ERMLPInteractionFunction( + embedding_dim=embedding_dim, + hidden_dim=hidden_dim, + ) super().__init__( triples_factory=triples_factory, From 68472db526f44acdad37c2b2a7a5a429a46af2e5 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 7 Nov 2020 19:19:05 +0100 Subject: [PATCH 042/690] Fix typo --- src/pykeen/models/unimodal/complex.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 5183ed8b46..711d445128 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -20,7 +20,7 @@ class ComplexInteractionFunction(InteractionFunction): - """Interaction function of Complex.""" + """Interaction function of CompleE.""" def forward( self, From 3f35c86651d7de1e297b681d08754c9f8a03a2f0 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 7 Nov 2020 19:22:30 +0100 Subject: [PATCH 043/690] Upgrade HolE --- src/pykeen/models/unimodal/hole.py | 110 +++++++++-------------------- 1 file changed, 34 insertions(+), 76 deletions(-) diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 01c4a905ea..db39e9f9b8 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -6,8 +6,9 @@ import torch import torch.autograd +from torch.fft import irfft, rfft -from ..base import EntityRelationEmbeddingModel +from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss from ...nn.init import xavier_uniform_ from ...regularizers import Regularizer @@ -19,7 +20,35 @@ ] -class HolE(EntityRelationEmbeddingModel): +class HolEInteractionFunction(InteractionFunction): + """Interaction function for HolE.""" + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + ) -> torch.FloatTensor: # noqa: D102 + # Circular correlation of entity embeddings + a_fft = rfft(h, signal_ndim=1, onesided=True) + b_fft = rfft(t, signal_ndim=1, onesided=True) + + # complex conjugate, a_fft.shape = (batch_size, num_entities, d', 2) + a_fft[:, :, :, 1] *= -1 + + # Hadamard product in frequency domain + p_fft = a_fft * b_fft + + # inverse real FFT, shape: (batch_size, num_entities, d) + composite = irfft(p_fft, signal_ndim=1, onesided=True, signal_sizes=(h.shape[-1],)) + + # inner product with relation embedding + scores = torch.sum(r * composite, dim=-1, keepdim=False) + + return scores + + +class HolE(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of HolE [nickel2016]_. Holographic embeddings (HolE) make use of the circular correlation operator to compute interactions between @@ -63,10 +92,13 @@ def __init__( regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model.""" + interaction_function = HolEInteractionFunction() + super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, loss=loss, + interaction_function=interaction_function, automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, @@ -77,77 +109,3 @@ def __init__( entity_constrainer=clamp_norm, entity_constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ) - - @staticmethod - def interaction_function( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - ) -> torch.FloatTensor: - """Evaluate the interaction function for given embeddings. - - The embeddings have to be in a broadcastable shape. - - :param h: shape: (batch_size, num_entities, d) - Head embeddings. - :param r: shape: (batch_size, num_entities, d) - Relation embeddings. - :param t: shape: (batch_size, num_entities, d) - Tail embeddings. - - :return: shape: (batch_size, num_entities) - The scores. - """ - # Circular correlation of entity embeddings - a_fft = torch.rfft(h, signal_ndim=1, onesided=True) - b_fft = torch.rfft(t, signal_ndim=1, onesided=True) - - # complex conjugate, a_fft.shape = (batch_size, num_entities, d', 2) - a_fft[:, :, :, 1] *= -1 - - # Hadamard product in frequency domain - p_fft = a_fft * b_fft - - # inverse real FFT, shape: (batch_size, num_entities, d) - composite = torch.irfft(p_fft, signal_ndim=1, onesided=True, signal_sizes=(h.shape[-1],)) - - # inner product with relation embedding - scores = torch.sum(r * composite, dim=-1, keepdim=False) - - return scores - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=hrt_batch[:, 0]).unsqueeze(dim=1) - r = self.relation_embeddings(indices=hrt_batch[:, 1]).unsqueeze(dim=1) - t = self.entity_embeddings(indices=hrt_batch[:, 2]).unsqueeze(dim=1) - - # Embedding Regularization - self.regularize_if_necessary(h, r, t) - - scores = self.interaction_function(h=h, r=r, t=t).view(-1, 1) - - return scores - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=hr_batch[:, 0]).unsqueeze(dim=1) - r = self.relation_embeddings(indices=hr_batch[:, 1]).unsqueeze(dim=1) - t = self.entity_embeddings(indices=None).unsqueeze(dim=0) - - # Embedding Regularization - self.regularize_if_necessary(h, r, t) - - scores = self.interaction_function(h=h, r=r, t=t) - - return scores - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=None).unsqueeze(dim=0) - r = self.relation_embeddings(indices=rt_batch[:, 0]).unsqueeze(dim=1) - t = self.entity_embeddings(indices=rt_batch[:, 1]).unsqueeze(dim=1) - - # Embedding Regularization - self.regularize_if_necessary(h, r, t) - - scores = self.interaction_function(h=h, r=r, t=t) - - return scores From cf7d3c1fb3d265165138b246c8b7da84631205bc Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 7 Nov 2020 21:14:22 +0100 Subject: [PATCH 044/690] Fix imports --- src/pykeen/models/unimodal/hole.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index db39e9f9b8..92226ea84d 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -6,7 +6,6 @@ import torch import torch.autograd -from torch.fft import irfft, rfft from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss @@ -30,8 +29,8 @@ def forward( t: torch.FloatTensor, ) -> torch.FloatTensor: # noqa: D102 # Circular correlation of entity embeddings - a_fft = rfft(h, signal_ndim=1, onesided=True) - b_fft = rfft(t, signal_ndim=1, onesided=True) + a_fft = torch.rfft(h, signal_ndim=1, onesided=True) + b_fft = torch.rfft(t, signal_ndim=1, onesided=True) # complex conjugate, a_fft.shape = (batch_size, num_entities, d', 2) a_fft[:, :, :, 1] *= -1 @@ -40,7 +39,7 @@ def forward( p_fft = a_fft * b_fft # inverse real FFT, shape: (batch_size, num_entities, d) - composite = irfft(p_fft, signal_ndim=1, onesided=True, signal_sizes=(h.shape[-1],)) + composite = torch.irfft(p_fft, signal_ndim=1, onesided=True, signal_sizes=(h.shape[-1],)) # inner product with relation embedding scores = torch.sum(r * composite, dim=-1, keepdim=False) From aca40b7d974be1ee26673aa86e022ffe2b040682 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 7 Nov 2020 22:07:28 +0100 Subject: [PATCH 045/690] Fix typos [skip ci] --- src/pykeen/models/unimodal/complex.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 711d445128..1d1ab23bd1 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -15,12 +15,12 @@ __all__ = [ 'ComplEx', - 'ComplexInteractionFunction', + 'ComplExInteractionFunction', ] -class ComplexInteractionFunction(InteractionFunction): - """Interaction function of CompleE.""" +class ComplExInteractionFunction(InteractionFunction): + """Interaction function of ComplEx.""" def forward( self, @@ -119,7 +119,7 @@ def __init__( :param regularizer: BaseRegularizer The regularizer to use. """ - interaction_function = ComplexInteractionFunction() + interaction_function = ComplExInteractionFunction() super().__init__( triples_factory=triples_factory, From ae6116dd9c6922203a877f3590266e7b29bfc3f9 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 9 Nov 2020 11:30:17 +0100 Subject: [PATCH 046/690] Add kwargs --- src/pykeen/models/base.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 446adebbd5..8381100b63 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1207,6 +1207,7 @@ def forward( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, + **kwargs, ) -> torch.FloatTensor: """Score the given triples. @@ -1216,6 +1217,8 @@ def forward( The relation representations. :param t: shape: (batch_size, num_tails, d_e) The tail representations. + :param kwargs: + Additional key-word based arguments. :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. @@ -1227,6 +1230,7 @@ def score_t( h: torch.FloatTensor, r: torch.FloatTensor, all_entities: torch.FloatTensor, + **kwargs, ) -> torch.FloatTensor: """ Score all tail entities. @@ -1237,6 +1241,8 @@ def score_t( The relation representations. :param all_entities: shape: (num_entities, d_e) The tail representations. + :param kwargs: + Additional key-word based arguments. :return: shape: (batch_size, num_entities) The scores. @@ -1247,7 +1253,7 @@ def score_t( t = all_entities.unsqueeze(dim=0) # get scores - scores = self(h=h, r=r, t=t) + scores = self(h=h, r=r, t=t, **kwargs) # prepare output shape scores = scores.squeeze(dim=2).squeeze(dim=1) From 3360271bc13ea1607a33f60a54ccd21e13ce25dd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 9 Nov 2020 11:32:26 +0100 Subject: [PATCH 047/690] Add kwargs to subclasses --- src/pykeen/models/unimodal/complex.py | 3 +++ src/pykeen/models/unimodal/distmult.py | 3 +++ src/pykeen/models/unimodal/ermlp.py | 3 +++ src/pykeen/models/unimodal/hole.py | 3 +++ 4 files changed, 12 insertions(+) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 1d1ab23bd1..c42a7f0c27 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -27,7 +27,10 @@ def forward( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, + **kwargs, ) -> torch.FloatTensor: # noqa: D102 + if len(kwargs) > 0: + raise ValueError(f"{self.__class__.__name__} does not take the following kwargs: {kwargs}") batch_size = max(h.shape[0], r.shape[0], t.shape[0]) h_term, h = normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') r_term, r = normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 9409308a82..e9a7da5cb9 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -29,7 +29,10 @@ def forward( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, + **kwargs, ) -> torch.FloatTensor: # noqa: D102 + if len(kwargs) > 0: + raise ValueError(f"{self.__class__.__name__} does not take the following kwargs: {kwargs}") batch_size = max(h.shape[0], r.shape[0], t.shape[0]) h_term, h = normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') r_term, r = normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 3065f1f45e..edc34b1b3b 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -56,7 +56,10 @@ def forward( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, + **kwargs, ) -> torch.FloatTensor: # noqa: D102 + if len(kwargs) > 0: + raise ValueError(f"{self.__class__.__name__} does not take the following kwargs: {kwargs}") h = self.head_to_hidden(h) r = self.rel_to_hidden(r) t = self.tail_to_hidden(t) diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 92226ea84d..cccd33bfda 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -27,7 +27,10 @@ def forward( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, + **kwargs, ) -> torch.FloatTensor: # noqa: D102 + if len(kwargs) > 0: + raise ValueError(f"{self.__class__.__name__} does not take the following kwargs: {kwargs}") # Circular correlation of entity embeddings a_fft = torch.rfft(h, signal_ndim=1, onesided=True) b_fft = torch.rfft(t, signal_ndim=1, onesided=True) From 115a272669bf99b5649a1f6135305e00a5e9407c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 9 Nov 2020 11:55:01 +0100 Subject: [PATCH 048/690] Add implementation of all score_* methods --- src/pykeen/models/base.py | 154 ++++++++++++++++++++++++- src/pykeen/models/unimodal/complex.py | 3 +- src/pykeen/models/unimodal/conv_e.py | 6 +- src/pykeen/models/unimodal/distmult.py | 3 +- src/pykeen/models/unimodal/ermlp.py | 3 +- src/pykeen/models/unimodal/hole.py | 3 +- 6 files changed, 157 insertions(+), 15 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 8381100b63..25ff738cbb 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1202,6 +1202,12 @@ class MultimodalModel(EntityRelationEmbeddingModel): class InteractionFunction(nn.Module): """Base class for interaction functions.""" + BATCH_DIM: int = 0 + NUM_DIM: int = 1 + HEAD_DIM: int = 1 + RELATION_DIM: int = 2 + TAIL_DIM: int = 3 + def forward( self, h: torch.FloatTensor, @@ -1225,6 +1231,145 @@ def forward( """ raise NotImplementedError + def _check_for_empty_kwargs(self, kwargs: Mapping[str, Any]) -> None: + """Check that kwargs is empty.""" + if len(kwargs) > 0: + raise ValueError(f"{self.__class__.__name__} does not take the following kwargs: {kwargs}") + + @staticmethod + def _add_dim(*x: torch.FloatTensor, dim: int) -> Sequence[torch.FloatTensor]: + """ + Add a dimension to tensors. + + :param x: shape: (d1, ..., dk) + The tensor. + + :return: shape: (1, d1, ..., dk) + The tensor with batch dimension. + """ + out = [xx.unsqueeze(dim=dim) for xx in x] + if len(x) > 1: + return out + return out[0] + + @staticmethod + def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: + """ + Remove dimensions from a tensor. + + :param x: + The tensor. + :param dims: + The dimensions to remove. + + :return: + The squeezed tensor. + """ + # normalize dimensions + dims = [d if d >= 0 else len(x.shape) + d for d in dims] + if len(set(dims)) != len(dims): + raise ValueError(f"Duplicate dimensions: {dims}") + assert all(0 <= d < len(x.shape) for d in dims) + for dim in reversed(sorted(dims)): + x = x.squeeze(dim=dim) + return x + + def score_hrt( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + """ + Score a batch of triples.. + + :param h: shape: (batch_size, d_e) + The head representations. + :param r: shape: (batch_size, d_r) + The relation representations. + :param t: shape: (batch_size, d_e) + The tail representations. + :param kwargs: + Additional key-word based arguments. + + :return: shape: (batch_size, 1) + The scores. + """ + # prepare input to generic score function + h, r, t = self._add_dim(h, r, t, dim=self.NUM_DIM) + + # get scores + scores = self(h=h, r=r, t=t, **kwargs) + + # prepare output shape, (batch_size, num_heads, num_relations, num_tails) -> (batch_size, 1) + return self._remove_dim(scores, self.HEAD_DIM, self.RELATION_DIM, self.TAIL_DIM).unsqueeze(dim=-1) + + def score_h( + self, + all_entities: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + """ + Score all head entities. + + :param all_entities: shape: (num_entities, d_e) + The head representations. + :param r: shape: (batch_size, d_r) + The relation representations. + :param t: shape: (batch_size, d_e) + The tail representations. + :param kwargs: + Additional key-word based arguments. + + :return: shape: (batch_size, num_entities) + The scores. + """ + # TODO: What about unsqueezing for additional e.g. head arguments + # prepare input to generic score function + r, t = self._add_dim(r, t, dim=self.NUM_DIM) + h = self._add_dim(all_entities, dim=self.BATCH_DIM) + + # get scores + scores = self(h=h, r=r, t=t, **kwargs) + + # prepare output shape, (batch_size, num_heads, num_relations, num_tails) -> (batch_size, num_heads) + return self._remove_dim(scores, self.RELATION_DIM, self.TAIL_DIM) + + def score_r( + self, + h: torch.FloatTensor, + all_relations: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + """ + Score all relations. + + :param h: shape: (batch_size, d_e) + The head representations. + :param all_relations: shape: (batch_size, d_r) + The relation representations. + :param t: shape: (num_entities, d_e) + The tail representations. + :param kwargs: + Additional key-word based arguments. + + :return: shape: (batch_size, num_entities) + The scores. + """ + # prepare input to generic score function + h, t = self._add_dim(h, t, dim=self.NUM_DIM) + r = self._add_dim(all_relations, dim=self.BATCH_DIM) + + # get scores + scores = self(h=h, r=r, t=t, **kwargs) + + # prepare output shape + return self._remove_dim(scores, self.HEAD_DIM, self.TAIL_DIM) + def score_t( self, h: torch.FloatTensor, @@ -1248,17 +1393,14 @@ def score_t( The scores. """ # prepare input to generic score function - h = h.unsqueeze(dim=1) - r = r.unsqueeze(dim=1) - t = all_entities.unsqueeze(dim=0) + h, r = self._add_dim(h, r, dim=self.NUM_DIM) + t = self._add_dim(all_entities, dim=self.BATCH_DIM) # get scores scores = self(h=h, r=r, t=t, **kwargs) # prepare output shape - scores = scores.squeeze(dim=2).squeeze(dim=1) - - return scores + return self._remove_dim(scores, self.HEAD_DIM, self.RELATION_DIM) def reset_parameters(self): """Reset parameters the interaction function may have.""" diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index c42a7f0c27..c39ce0e388 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -29,8 +29,7 @@ def forward( t: torch.FloatTensor, **kwargs, ) -> torch.FloatTensor: # noqa: D102 - if len(kwargs) > 0: - raise ValueError(f"{self.__class__.__name__} does not take the following kwargs: {kwargs}") + self._check_for_empty_kwargs(kwargs) batch_size = max(h.shape[0], r.shape[0], t.shape[0]) h_term, h = normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') r_term, r = normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 1fae5d2a0d..6b10b2073b 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -11,7 +11,7 @@ from torch import nn from torch.nn import functional as F # noqa: N812 -from ..base import EntityRelationEmbeddingModel +from ..base import EntityRelationEmbeddingModel, InteractionFunction from ...losses import BCEAfterSigmoidLoss, Loss from ...nn import Embedding from ...nn.init import xavier_normal_ @@ -75,6 +75,10 @@ def _calculate_missing_shape_information( return input_channels, width, height +class ConvEInteractionFunction(InteractionFunction): + """ConvE interaction function.""" + + class ConvE(EntityRelationEmbeddingModel): r"""An implementation of ConvE from [dettmers2018]_. diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index e9a7da5cb9..0cb64dd42b 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -31,8 +31,7 @@ def forward( t: torch.FloatTensor, **kwargs, ) -> torch.FloatTensor: # noqa: D102 - if len(kwargs) > 0: - raise ValueError(f"{self.__class__.__name__} does not take the following kwargs: {kwargs}") + self._check_for_empty_kwargs(kwargs) batch_size = max(h.shape[0], r.shape[0], t.shape[0]) h_term, h = normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') r_term, r = normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index edc34b1b3b..ce9a8a0a1e 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -58,8 +58,7 @@ def forward( t: torch.FloatTensor, **kwargs, ) -> torch.FloatTensor: # noqa: D102 - if len(kwargs) > 0: - raise ValueError(f"{self.__class__.__name__} does not take the following kwargs: {kwargs}") + self._check_for_empty_kwargs(kwargs) h = self.head_to_hidden(h) r = self.rel_to_hidden(r) t = self.tail_to_hidden(t) diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index cccd33bfda..cf07f5c6bf 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -29,8 +29,7 @@ def forward( t: torch.FloatTensor, **kwargs, ) -> torch.FloatTensor: # noqa: D102 - if len(kwargs) > 0: - raise ValueError(f"{self.__class__.__name__} does not take the following kwargs: {kwargs}") + self._check_for_empty_kwargs(kwargs) # Circular correlation of entity embeddings a_fft = torch.rfft(h, signal_ndim=1, onesided=True) b_fft = torch.rfft(t, signal_ndim=1, onesided=True) From 2255cb5737a3c24c76573779502816f9cb5c0153 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 13:06:27 +0100 Subject: [PATCH 049/690] Add index-based interaction function intermediary --- src/pykeen/models/base.py | 167 ++++++++++++++++++++++++++++++++++---- tests/test_models.py | 7 +- 2 files changed, 158 insertions(+), 16 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index be0c066933..fd94aee2b1 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -20,14 +20,16 @@ from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory from ..typing import Constrainer, Initializer, MappedTriples, Normalizer -from ..utils import NoRandomSeedNecessary, get_embedding_in_canonical_shape, resolve_device, set_random_seed +from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed __all__ = [ 'Model', 'EntityEmbeddingModel', 'EntityRelationEmbeddingModel', + 'GeneralVectorEntityRelationEmbeddingModel', 'SimpleVectorEntityRelationEmbeddingModel', 'InteractionFunction', + 'IndexFunction', 'MultimodalModel', ] @@ -1263,13 +1265,73 @@ def reset_parameters(self): mod.reset_parameters() -class SimpleVectorEntityRelationEmbeddingModel(EntityRelationEmbeddingModel): +class IndexFunction(nn.Module): + """A function that handles looking up embeddings by index.""" + + def forward( + self, + model: Model, + h_indices: Optional[torch.LongTensor] = None, + r_indices: Optional[torch.LongTensor] = None, + t_indices: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: + """Get the scores for the given indices. + + :param model: + The KGEM so lookup to attributes like entity embeddings and relation embeddings is possible + :param h_indices: shape: (batch_size,) + The indices for head entities. If None, score against all. + :param r_indices: shape: (batch_size,) + The indices for relations. If None, score against all. + :param t_indices: shape: (batch_size,) + The indices for tail entities. If None, score against all. + + :return: The scores, shape: (batch_size, num_entities) + """ + raise NotImplementedError + + def reset_parameters(self): + """Reset parameters the interaction function may have.""" + + +class InteractionIndexFunction(IndexFunction): + """Wrap a :class:`InteractionFunction` with index-based lookup of entity and relation embeddings.""" + + def __init__(self, interaction_function: InteractionFunction): + super().__init__() + self.interaction_function = interaction_function + + def forward( + self, + model: Model, + h_indices: Optional[torch.LongTensor] = None, + r_indices: Optional[torch.LongTensor] = None, + t_indices: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: # noqa: D102 + # Get embeddings + h = model.entity_embeddings.get_in_canonical_shape(indices=h_indices) + r = model.relation_embeddings.get_in_canonical_shape(indices=r_indices) + t = model.entity_embeddings.get_in_canonical_shape(indices=t_indices) + + # Compute score + scores = self.interaction_function(h=h, r=r, t=t) + + # Only regularize relation embeddings + model.regularize_if_necessary(r) + + return scores + + def reset_parameters(self): # noqa: D102 + self.interaction_function.reset_parameters() + + +class GeneralVectorEntityRelationEmbeddingModel(EntityRelationEmbeddingModel): """A base class for embedding models which store a single vector for each entity and relation.""" def __init__( self, triples_factory: TriplesFactory, - interaction_function: InteractionFunction, + index_function: IndexFunction, embedding_dim: int = 200, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, @@ -1330,7 +1392,7 @@ def __init__( relation_constrainer=relation_constrainer, relation_constrainer_kwargs=relation_constrainer_kwargs, ) - self.interaction_function = interaction_function + self.index_function = index_function def __init_subclass__(cls, auto_reset_parameters: bool = True, **kwargs): # noqa: D105 _track_hyperparameters(cls) @@ -1339,7 +1401,7 @@ def __init_subclass__(cls, auto_reset_parameters: bool = True, **kwargs): # noq def _reset_parameters_(self): super()._reset_parameters_() - self.interaction_function.reset_parameters() + self.index_function.reset_parameters() def _score( self, @@ -1347,18 +1409,18 @@ def _score( r_ind: Optional[torch.LongTensor] = None, t_ind: Optional[torch.LongTensor] = None, ) -> torch.FloatTensor: - # Get embeddings - h = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=h_ind) - r = get_embedding_in_canonical_shape(embedding=self.relation_embeddings, ind=r_ind) - t = get_embedding_in_canonical_shape(embedding=self.entity_embeddings, ind=t_ind) - - # Compute score - scores = self.interaction_function(h=h, r=r, t=t) + """Evaluate the given triples. - # Only regularize relation embeddings - self.regularize_if_necessary(r) + :param h_ind: shape: (batch_size,) + The indices for head entities. If None, score against all. + :param r_ind: shape: (batch_size,) + The indices for relations. If None, score against all. + :param t_ind: shape: (batch_size,) + The indices for tail entities. If None, score against all. - return scores + :return: The scores, shape: (batch_size, num_entities) + """ + return self.index_function(model=self, h_ind=h_ind, r_ind=r_ind, t_ind=t_ind) def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self._score(h_ind=hrt_batch[:, 0], r_ind=hrt_batch[:, 1], t_ind=hrt_batch[:, 2]).view(-1, 1) @@ -1368,3 +1430,78 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) + + def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + raise NotImplementedError + + +class SimpleVectorEntityRelationEmbeddingModel(GeneralVectorEntityRelationEmbeddingModel): + """A base class for embedding models which store a single vector for each entity and relation.""" + + def __init__( + self, + triples_factory: TriplesFactory, + interaction_function: InteractionFunction, + embedding_dim: int = 200, + automatic_memory_optimization: Optional[bool] = None, + loss: Optional[Loss] = None, + preferred_device: Optional[str] = None, + random_seed: Optional[int] = None, + regularizer: Optional[Regularizer] = None, + entity_initializer: Optional[Initializer] = None, + entity_initializer_kwargs: Optional[Mapping[str, Any]] = None, + entity_normalizer: Optional[Normalizer] = None, + entity_normalizer_kwargs: Optional[Mapping[str, Any]] = None, + entity_constrainer: Optional[Constrainer] = None, + entity_constrainer_kwargs: Optional[Mapping[str, Any]] = None, + relation_initializer: Optional[Initializer] = None, + relation_initializer_kwargs: Optional[Mapping[str, Any]] = None, + relation_normalizer: Optional[Normalizer] = None, + relation_normalizer_kwargs: Optional[Mapping[str, Any]] = None, + relation_constrainer: Optional[Constrainer] = None, + relation_constrainer_kwargs: Optional[Mapping[str, Any]] = None, + ) -> None: + """Initialize embedding model. + + :param triples_factory: TriplesFactory + The triple factory connected to the model. + :param interaction_function: + The interaction function used to compute scores. + :param embedding_dim: + The embedding dimensionality of the entity embeddings. + :param automatic_memory_optimization: bool + Whether to automatically optimize the sub-batch size during training and batch size during evaluation with + regards to the hardware at hand. + :param loss: OptionalLoss (optional) + The loss to use. + :param preferred_device: str (optional) + The default device where to model is located. + :param random_seed: int (optional) + An optional random seed to set before the initialization of weights. + :param regularizer: BaseRegularizer + The regularizer to use. + """ + index_function = InteractionIndexFunction(interaction_function=interaction_function) + + super().__init__( + triples_factory=triples_factory, + embedding_dim=embedding_dim, + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + index_function=index_function, + random_seed=random_seed, + regularizer=regularizer, + entity_initializer=entity_initializer, + entity_initializer_kwargs=entity_initializer_kwargs, + entity_normalizer=entity_normalizer, + entity_normalizer_kwargs=entity_normalizer_kwargs, + entity_constrainer=entity_constrainer, + entity_constrainer_kwargs=entity_constrainer_kwargs, + relation_initializer=relation_initializer, + relation_initializer_kwargs=relation_initializer_kwargs, + relation_normalizer=relation_normalizer, + relation_normalizer_kwargs=relation_normalizer_kwargs, + relation_constrainer=relation_constrainer, + relation_constrainer_kwargs=relation_constrainer_kwargs, + ) diff --git a/tests/test_models.py b/tests/test_models.py index 3bd777e426..aaf5891d4f 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -26,8 +26,10 @@ from pykeen.models.base import ( EntityEmbeddingModel, EntityRelationEmbeddingModel, + GeneralVectorEntityRelationEmbeddingModel, Model, MultimodalModel, + SimpleVectorEntityRelationEmbeddingModel, _extend_batch, get_novelty_mask, ) @@ -49,10 +51,11 @@ MultimodalModel.__name__, EntityEmbeddingModel.__name__, EntityRelationEmbeddingModel.__name__, + GeneralVectorEntityRelationEmbeddingModel.__name__, + SimpleVectorEntityRelationEmbeddingModel.__name__, 'MockModel', 'models', 'get_model_cls', - 'SimpleInteractionModel', } for cls in MultimodalModel.__subclasses__(): SKIP_MODULES.add(cls.__name__) @@ -1260,6 +1263,8 @@ def test_abstract(self): self.assertTrue(Model._is_abstract()) self.assertTrue(EntityEmbeddingModel._is_abstract()) self.assertTrue(EntityRelationEmbeddingModel._is_abstract()) + self.assertTrue(GeneralVectorEntityRelationEmbeddingModel._is_abstract()) + self.assertTrue(SimpleVectorEntityRelationEmbeddingModel._is_abstract()) for model_cls in _MODELS: if issubclass(model_cls, MultimodalModel): continue From 3f794bb670fcea9afeab640bb12bed7194f376aa Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 9 Nov 2020 13:08:53 +0100 Subject: [PATCH 050/690] Extract ConvE interaction function --- src/pykeen/models/unimodal/conv_e.py | 395 +++++++++++++-------------- 1 file changed, 186 insertions(+), 209 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 6b10b2073b..47781d3e00 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -4,12 +4,10 @@ import logging import math -import sys from typing import Optional, Tuple, Type import torch from torch import nn -from torch.nn import functional as F # noqa: N812 from ..base import EntityRelationEmbeddingModel, InteractionFunction from ...losses import BCEAfterSigmoidLoss, Loss @@ -75,9 +73,168 @@ def _calculate_missing_shape_information( return input_channels, width, height +def _add_cuda_warning(func): + def wrapped(*args, **kwargs): + try: + return func(*args, **kwargs) + except RuntimeError as e: + if not is_cudnn_error(e): + raise e + raise RuntimeError( + '\nThis code crash might have been caused by a CUDA bug, see ' + 'https://github.com/allenai/allennlp/issues/2888, ' + 'which causes the code to crash during evaluation mode.\n' + 'To avoid this error, the batch size has to be reduced.' + ) from e + + return wrapped + + class ConvEInteractionFunction(InteractionFunction): """ConvE interaction function.""" + #: If batch normalization is enabled, this is: num_features – C from an expected input of size (N,C,L) + bn0: Optional[torch.nn.BatchNorm2d] + #: If batch normalization is enabled, this is: num_features – C from an expected input of size (N,C,H,W) + bn1: Optional[torch.nn.BatchNorm2d] + bn2: Optional[torch.nn.BatchNorm2d] + + def __init__( + self, + input_channels: Optional[int] = None, + output_channels: int = 32, + embedding_height: Optional[int] = None, + embedding_width: Optional[int] = None, + kernel_height: int = 3, + kernel_width: int = 3, + input_dropout: float = 0.2, + output_dropout: float = 0.3, + feature_map_dropout: float = 0.2, + embedding_dim: int = 200, + apply_batch_normalization: bool = True, + ): + super().__init__() + + # Automatic calculation of remaining dimensions + logger.info(f'Resolving {input_channels} * {embedding_width} * {embedding_height} = {embedding_dim}.') + if embedding_dim is None: + embedding_dim = input_channels * embedding_width * embedding_height + + # Parameter need to fulfil: + # input_channels * embedding_height * embedding_width = embedding_dim + input_channels, embedding_width, embedding_height = _calculate_missing_shape_information( + embedding_dim=embedding_dim, + input_channels=input_channels, + width=embedding_width, + height=embedding_height, + ) + logger.info(f'Resolved to {input_channels} * {embedding_width} * {embedding_height} = {embedding_dim}.') + self.embedding_height = embedding_height + self.embedding_width = embedding_width + self.input_channels = input_channels + + if self.input_channels * self.embedding_height * self.embedding_width != self.embedding_dim: + raise ValueError( + f'Product of input channels ({self.input_channels}), height ({self.embedding_height}), and width ' + f'({self.embedding_width}) does not equal target embedding dimension ({self.embedding_dim})', + ) + + self.inp_drop = nn.Dropout(input_dropout) + self.hidden_drop = nn.Dropout(output_dropout) + self.feature_map_drop = nn.Dropout2d(feature_map_dropout) + + self.conv1 = torch.nn.Conv2d( + in_channels=self.input_channels, + out_channels=output_channels, + kernel_size=(kernel_height, kernel_width), + stride=1, + padding=0, + bias=True, + ) + + self.apply_batch_normalization = apply_batch_normalization + if self.apply_batch_normalization: + self.bn0 = nn.BatchNorm2d(self.input_channels) + self.bn1 = nn.BatchNorm2d(output_channels) + self.bn2 = nn.BatchNorm1d(self.embedding_dim) + else: + self.bn0 = None + self.bn1 = None + self.bn2 = None + num_in_features = ( + output_channels + * (2 * self.embedding_height - kernel_height + 1) + * (self.embedding_width - kernel_width + 1) + ) + self.fc = nn.Linear(num_in_features, self.embedding_dim) + self.activation = nn.ReLU() + + @_add_cuda_warning + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + # get tail bias term + if "t_bias" not in kwargs: + raise TypeError(f"{self.__class__.__name__}.forward expects keyword argument 't_bias'.") + t_bias: torch.FloatTensor = kwargs.pop("t_bias") + self._check_for_empty_kwargs(kwargs) + + # bind sizes + batch_size, num_heads = h.shape[:2] + num_relations = r.shape[1] + + # repeat if necessary + h = h.unsqueeze(dim=2).repeat(1, 1, num_relations, 1) + r = r.unsqueeze(dim=1).repeat(1, num_heads, 1, 1) + + # resize and concat head and relation, batch_size', num_input_channels, 2*height, width + # with batch_size' = batch_size * num_heads * num_relations + x = torch.cat([ + h.view(-1, self.input_channels, self.embedding_height, self.embedding_width), + r.view(-1, self.input_channels, self.embedding_height, self.embedding_width), + ], dim=2) + + # batch_size, num_input_channels, 2*height, width + if self.apply_batch_normalization: + x = self.bn0(x) + + # batch_size, num_input_channels, 2*height, width + x = self.inp_drop(x) + + # (N,C_out,H_out,W_out) + x = self.conv1(x) + + if self.apply_batch_normalization: + x = self.bn1(x) + + x = self.activation(x) + x = self.feature_map_drop(x) + + # batch_size', num_output_channels * (2 * height - kernel_height + 1) * (width - kernel_width + 1) + x = x.view(batch_size, -1) + x = self.fc(x) + x = self.hidden_drop(x) + + if self.apply_batch_normalization: + x = self.bn2(x) + x = self.activation(x) + + # reshape: (batch_size', ) + x = x.view(batch_size, num_heads, num_relations, 1, self.embedding_dim) + + # For efficient calculation, each of the convolved [h, r] rows has only to be multiplied with one t row + # output_shape: (batch_size, num_heads, num_relations, num_tails) + x = x @ t.view(batch_size, 1, 1, -1, self.embedding_dim).transpose(-1, -2) + + # add bias term + x = x + t_bias[:, None, None, :] + + return x + class ConvE(EntityRelationEmbeddingModel): r"""An implementation of ConvE from [dettmers2018]_. @@ -161,12 +318,6 @@ class ConvE(EntityRelationEmbeddingModel): #: The default parameters for the default loss function class loss_default_kwargs = {} - #: If batch normalization is enabled, this is: num_features – C from an expected input of size (N,C,L) - bn0: Optional[torch.nn.BatchNorm2d] - #: If batch normalization is enabled, this is: num_features – C from an expected input of size (N,C,H,W) - bn1: Optional[torch.nn.BatchNorm2d] - bn2: Optional[torch.nn.BatchNorm2d] - def __init__( self, triples_factory: TriplesFactory, @@ -216,219 +367,45 @@ def __init__( initializer=nn.init.zeros_, ) - # Automatic calculation of remaining dimensions - logger.info(f'Resolving {input_channels} * {embedding_width} * {embedding_height} = {embedding_dim}.') - if embedding_dim is None: - embedding_dim = input_channels * embedding_width * embedding_height - - # Parameter need to fulfil: - # input_channels * embedding_height * embedding_width = embedding_dim - input_channels, embedding_width, embedding_height = _calculate_missing_shape_information( - embedding_dim=embedding_dim, + self.interaction_function = ConvEInteractionFunction( input_channels=input_channels, - width=embedding_width, - height=embedding_height, - ) - logger.info(f'Resolved to {input_channels} * {embedding_width} * {embedding_height} = {embedding_dim}.') - self.embedding_height = embedding_height - self.embedding_width = embedding_width - self.input_channels = input_channels - - if self.input_channels * self.embedding_height * self.embedding_width != self.embedding_dim: - raise ValueError( - f'Product of input channels ({self.input_channels}), height ({self.embedding_height}), and width ' - f'({self.embedding_width}) does not equal target embedding dimension ({self.embedding_dim})', - ) - - self.inp_drop = nn.Dropout(input_dropout) - self.hidden_drop = nn.Dropout(output_dropout) - self.feature_map_drop = nn.Dropout2d(feature_map_dropout) - - self.conv1 = torch.nn.Conv2d( - in_channels=self.input_channels, - out_channels=output_channels, - kernel_size=(kernel_height, kernel_width), - stride=1, - padding=0, - bias=True, - ) - - self.apply_batch_normalization = apply_batch_normalization - if self.apply_batch_normalization: - self.bn0 = nn.BatchNorm2d(self.input_channels) - self.bn1 = nn.BatchNorm2d(output_channels) - self.bn2 = nn.BatchNorm1d(self.embedding_dim) - else: - self.bn0 = None - self.bn1 = None - self.bn2 = None - num_in_features = ( - output_channels - * (2 * self.embedding_height - kernel_height + 1) - * (self.embedding_width - kernel_width + 1) + output_channels=output_channels, + embedding_height=embedding_height, + embedding_width=embedding_width, + kernel_height=kernel_height, + kernel_width=kernel_width, + input_dropout=input_dropout, + output_dropout=output_dropout, + feature_map_dropout=feature_map_dropout, + embedding_dim=embedding_dim, + apply_batch_normalization=apply_batch_normalization, ) - self.fc = nn.Linear(num_in_features, self.embedding_dim) def _reset_parameters_(self): # noqa: D102 super()._reset_parameters_() - self.bias_term.reset_parameters() - - # weights - for module in [ - self.conv1, - self.bn0, - self.bn1, - self.bn2, - self.fc, - ]: - if module is None: - continue - module.reset_parameters() - - def _convolve_entity_relation(self, h: torch.LongTensor, r: torch.LongTensor) -> torch.FloatTensor: - """Perform the main calculations of the ConvE model.""" - batch_size = h.shape[0] - - # batch_size, num_input_channels, 2*height, width - x = torch.cat([h, r], dim=2) - - try: - # batch_size, num_input_channels, 2*height, width - if self.apply_batch_normalization: - x = self.bn0(x) - - # batch_size, num_input_channels, 2*height, width - x = self.inp_drop(x) - # (N,C_out,H_out,W_out) - x = self.conv1(x) - - if self.apply_batch_normalization: - x = self.bn1(x) - x = F.relu(x) - x = self.feature_map_drop(x) - # batch_size, num_output_channels * (2 * height - kernel_height + 1) * (width - kernel_width + 1) - x = x.view(batch_size, -1) - x = self.fc(x) - x = self.hidden_drop(x) - - if self.apply_batch_normalization: - x = self.bn2(x) - x = F.relu(x) - except RuntimeError as e: - if not is_cudnn_error(e): - raise e - logger.warning( - '\nThis code crash might have been caused by a CUDA bug, see ' - 'https://github.com/allenai/allennlp/issues/2888, ' - 'which causes the code to crash during evaluation mode.\n' - 'To avoid this error, the batch size has to be reduced.\n' - f'The original error message: \n{e.args[0]}', - ) - sys.exit(1) - - return x + self.interaction_function.reset_parameters() def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=hrt_batch[:, 0]).view( - -1, - self.input_channels, - self.embedding_height, - self.embedding_width, - ) - r = self.relation_embeddings(indices=hrt_batch[:, 1]).view( - -1, - self.input_channels, - self.embedding_height, - self.embedding_width, - ) + h = self.entity_embeddings(indices=hrt_batch[:, 0]) + r = self.relation_embeddings(indices=hrt_batch[:, 1]) t = self.entity_embeddings(indices=hrt_batch[:, 2]) - - # Embedding Regularization + t_bias = self.bias_term(indices=hrt_batch[:, 2]) self.regularize_if_necessary(h, r, t) - - x = self._convolve_entity_relation(h, r) - - # For efficient calculation, each of the convolved [h, r] rows has only to be multiplied with one t row - x = (x.view(-1, self.embedding_dim) * t).sum(dim=1, keepdim=True) - - """ - In ConvE the bias term add the end is added for each tail item. In the sLCWA training approach we only have - one tail item for each head and relation. Accordingly the relevant bias for each tail item and triple has to be - looked up. - """ - x = x + self.bias_term(indices=hrt_batch[:, 2]) - # The application of the sigmoid during training is automatically handled by the default loss. - - return x + return self.interaction_function.score_hrt(h=h, r=r, t=t, t_bias=t_bias) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=hr_batch[:, 0]).view( - -1, - self.input_channels, - self.embedding_height, - self.embedding_width, - ) - r = self.relation_embeddings(indices=hr_batch[:, 1]).view( - -1, - self.input_channels, - self.embedding_height, - self.embedding_width, - ) - t = self.entity_embeddings(indices=None).transpose(1, 0) - - # Embedding Regularization - self.regularize_if_necessary(h, r, t) - - x = self._convolve_entity_relation(h, r) - - x = x @ t - x = x + self.bias_term(indices=None).t() - # The application of the sigmoid during training is automatically handled by the default loss. - - return x + h = self.entity_embeddings(indices=hr_batch[:, 0]) + r = self.relation_embeddings(indices=hr_batch[:, 1]) + all_entities = self.entity_embeddings(indices=None) + t_bias = self.bias_term(indices=None) + self.regularize_if_necessary(h, r, all_entities) + return self.interaction_function.score_t(h=h, r=r, all_entities=all_entities, t_bias=t_bias) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - rt_batch_size = rt_batch.shape[0] - h = self.entity_embeddings(indices=None) - r = self.relation_embeddings(indices=rt_batch[:, 0]).view( - -1, - self.input_channels, - self.embedding_height, - self.embedding_width, - ) + all_entities = self.entity_embeddings(indices=None) + r = self.relation_embeddings(indices=rt_batch[:, 0]) t = self.entity_embeddings(indices=rt_batch[:, 1]) - - # Embedding Regularization - self.regularize_if_necessary(h, r, t) - - ''' - Every head has to be convolved with every relation in the rt_batch. Hence we repeat the - relation _num_entities_ times and the head _rt_batch_size_ times. - ''' - r = r.repeat(h.shape[0], 1, 1, 1) - # Code to repeat each item successively instead of the entire tensor - h = h.unsqueeze(1).repeat(1, rt_batch_size, 1).view( - -1, - self.input_channels, - self.embedding_height, - self.embedding_width, - ) - - x = self._convolve_entity_relation(h, r) - - ''' - For efficient computation, each convolved [h, r] pair has only to be multiplied with the corresponding t - embedding found in the rt_batch with [r, t] pairs. - ''' - x = (x.view(self.num_entities, rt_batch_size, self.embedding_dim) * t[None, :, :]).sum(2).transpose(1, 0) - - """ - In ConvE the bias term at the end is added for each tail item. In the score_h function, each row holds - the same tail for many different heads, meaning that these items have to be looked up for each tail of each row - and only then can be added correctly. - """ - x = x + self.bias_term(indices=rt_batch[:, 1]) - # The application of the sigmoid during training is automatically handled by the default loss. - - return x + t_bias = self.bias_term(indices=rt_batch[:, 1]) + self.regularize_if_necessary(all_entities, r, t) + return self.interaction_function.score_h(all_entities=all_entities, r=r, t=t, t_bias=t_bias) From d718f31733e228cde6c7e073acc91740568a0773 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 9 Nov 2020 13:09:57 +0100 Subject: [PATCH 051/690] Fix missing field --- src/pykeen/models/unimodal/conv_e.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 47781d3e00..3b14487729 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -129,6 +129,7 @@ def __init__( height=embedding_height, ) logger.info(f'Resolved to {input_channels} * {embedding_width} * {embedding_height} = {embedding_dim}.') + self.embedding_dim = embedding_dim self.embedding_height = embedding_height self.embedding_width = embedding_width self.input_channels = input_channels From 8cb8381adcdfedb3af3ed38633c76f04162c2af4 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 13:11:24 +0100 Subject: [PATCH 052/690] Streamline class lookup --- src/pykeen/models/__init__.py | 21 ++++++++------------- 1 file changed, 8 insertions(+), 13 deletions(-) diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index 0c71eeab80..52aebd75fd 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -6,11 +6,12 @@ score value is model-dependent, and usually it cannot be directly interpreted as a probability. """ # noqa: D205, D400 +import inspect from typing import Mapping, Set, Type, Union from .base import ( EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, MultimodalModel, - SimpleVectorEntityRelationEmbeddingModel, + SimpleVectorEntityRelationEmbeddingModel, GeneralVectorEntityRelationEmbeddingModel ) from .multimodal import ComplExLiteral, DistMultLiteral from .unimodal import ( @@ -67,20 +68,14 @@ ] -def _recur(c): - for sc in c.__subclasses__(): - yield sc - yield from _recur(sc) +def _concrete_subclasses(cls): + for subcls in cls.__subclasses__(): + if not inspect.isabstract(cls): + yield subcls + yield from _concrete_subclasses(subcls) -_MODELS: Set[Type[Model]] = { - cls - for cls in _recur(Model) - if cls not in { - Model, MultimodalModel, EntityRelationEmbeddingModel, - EntityEmbeddingModel, SimpleVectorEntityRelationEmbeddingModel, - } -} +_MODELS: Set[Type[Model]] = set(_concrete_subclasses(Model)) #: A mapping of models' names to their implementations models: Mapping[str, Type[Model]] = { From 1211e0741aedde21c00bc43b2f3ee1c698d3db92 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 13:13:40 +0100 Subject: [PATCH 053/690] Update base.py --- src/pykeen/models/base.py | 10 +++------- 1 file changed, 3 insertions(+), 7 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index fd94aee2b1..466416cb14 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1394,11 +1394,6 @@ def __init__( ) self.index_function = index_function - def __init_subclass__(cls, auto_reset_parameters: bool = True, **kwargs): # noqa: D105 - _track_hyperparameters(cls) - if auto_reset_parameters: - _add_post_reset_parameters(cls) - def _reset_parameters_(self): super()._reset_parameters_() self.index_function.reset_parameters() @@ -1431,8 +1426,9 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) - def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - raise NotImplementedError + # TODO + # def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + # raise NotImplementedError class SimpleVectorEntityRelationEmbeddingModel(GeneralVectorEntityRelationEmbeddingModel): From 8d2ecae5cf98ee34a2076942459eaf1bf67c6099 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 9 Nov 2020 13:22:45 +0100 Subject: [PATCH 054/690] Add shape verification --- src/pykeen/models/base.py | 35 +++++++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 25ff738cbb..1e136bec69 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1274,6 +1274,29 @@ def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: x = x.squeeze(dim=dim) return x + def _check_shape( + self, + *x: Union[str, torch.Tensor], + ) -> bool: + assert len(x) % 2 == 0 + tensors = x[::2] + shapes = x[1::2] + assert all(torch.is_tensor(t) for t in tensors) + assert all(isinstance(s, str) for s in shapes) + dims = dict() + errors = [] + for tensor, shape in zip(tensors, shapes): + if tensor.ndimension() != len(shape): + errors.append(f"Invalid number of dimensions: {tensor.shape} vs. {shape}") + continue + for dim, name in zip(tensor.shape, shape): + exp_dim = dims.get(name) + if exp_dim is not None and exp_dim != dim: + errors.append(f"{name}: {dim} vs. {exp_dim}") + if len(errors) > 0: + raise ValueError("Shape verification failed:\n" + '\n'.join(errors)) + return True + def score_hrt( self, h: torch.FloatTensor, @@ -1296,6 +1319,9 @@ def score_hrt( :return: shape: (batch_size, 1) The scores. """ + # check shape + assert self._check_shape(h, "be", r, "br", t, "be") + # prepare input to generic score function h, r, t = self._add_dim(h, r, t, dim=self.NUM_DIM) @@ -1327,6 +1353,9 @@ def score_h( :return: shape: (batch_size, num_entities) The scores. """ + # check shape + assert self._check_shape(all_entities, "ne", r, "br", t, "be") + # TODO: What about unsqueezing for additional e.g. head arguments # prepare input to generic score function r, t = self._add_dim(r, t, dim=self.NUM_DIM) @@ -1360,6 +1389,9 @@ def score_r( :return: shape: (batch_size, num_entities) The scores. """ + # check shape + assert self._check_shape(all_relations, "nr", h, "be", t, "be") + # prepare input to generic score function h, t = self._add_dim(h, t, dim=self.NUM_DIM) r = self._add_dim(all_relations, dim=self.BATCH_DIM) @@ -1392,6 +1424,9 @@ def score_t( :return: shape: (batch_size, num_entities) The scores. """ + # check shape + assert self._check_shape(all_entities, "ne", r, "br", h, "be") + # prepare input to generic score function h, r = self._add_dim(h, r, dim=self.NUM_DIM) t = self._add_dim(all_entities, dim=self.BATCH_DIM) From bc47c8a5a44b3778ff2b7b84c53752e32e0df000 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 9 Nov 2020 13:28:55 +0100 Subject: [PATCH 055/690] some fixes for ConvE --- src/pykeen/models/unimodal/conv_e.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 3b14487729..42d11808fd 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -224,12 +224,13 @@ def forward( x = self.bn2(x) x = self.activation(x) - # reshape: (batch_size', ) + # reshape: (batch_size', embedding_dim) x = x.view(batch_size, num_heads, num_relations, 1, self.embedding_dim) # For efficient calculation, each of the convolved [h, r] rows has only to be multiplied with one t row # output_shape: (batch_size, num_heads, num_relations, num_tails) - x = x @ t.view(batch_size, 1, 1, -1, self.embedding_dim).transpose(-1, -2) + t = t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]).transpose(-1, -2) + x = (x @ t).squeeze(dim=-2) # add bias term x = x + t_bias[:, None, None, :] From f3498226738bdff3e733278f46b749b3b8801387 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 9 Nov 2020 13:34:51 +0100 Subject: [PATCH 056/690] More ConvE fixes --- src/pykeen/models/unimodal/conv_e.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 42d11808fd..72c64b0033 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -162,12 +162,12 @@ def __init__( self.bn0 = None self.bn1 = None self.bn2 = None - num_in_features = ( + self.num_in_features = ( output_channels * (2 * self.embedding_height - kernel_height + 1) * (self.embedding_width - kernel_width + 1) ) - self.fc = nn.Linear(num_in_features, self.embedding_dim) + self.fc = nn.Linear(self.num_in_features, self.embedding_dim) self.activation = nn.ReLU() @_add_cuda_warning @@ -185,12 +185,14 @@ def forward( self._check_for_empty_kwargs(kwargs) # bind sizes - batch_size, num_heads = h.shape[:2] + batch_size = max(x.shape[0] for x in (h, r, t)) + num_heads = h.shape[1] num_relations = r.shape[1] + num_tails = t.shape[1] # repeat if necessary - h = h.unsqueeze(dim=2).repeat(1, 1, num_relations, 1) - r = r.unsqueeze(dim=1).repeat(1, num_heads, 1, 1) + h = h.unsqueeze(dim=2).repeat(1 if h.shape[0] == batch_size else batch_size, 1, num_relations, 1) + r = r.unsqueeze(dim=1).repeat(1 if r.shape[0] == batch_size else batch_size, num_heads, 1, 1) # resize and concat head and relation, batch_size', num_input_channels, 2*height, width # with batch_size' = batch_size * num_heads * num_relations @@ -216,7 +218,7 @@ def forward( x = self.feature_map_drop(x) # batch_size', num_output_channels * (2 * height - kernel_height + 1) * (width - kernel_width + 1) - x = x.view(batch_size, -1) + x = x.view(-1, self.num_in_features) x = self.fc(x) x = self.hidden_drop(x) @@ -229,11 +231,11 @@ def forward( # For efficient calculation, each of the convolved [h, r] rows has only to be multiplied with one t row # output_shape: (batch_size, num_heads, num_relations, num_tails) - t = t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]).transpose(-1, -2) + t = t.view(t.shape[0], 1, 1, num_tails, self.embedding_dim).transpose(-1, -2) x = (x @ t).squeeze(dim=-2) # add bias term - x = x + t_bias[:, None, None, :] + x = x + t_bias.view(t.shape[0], 1, 1, num_tails) return x From 85dc528a0a04161af0a7b0e64b28e1e8079b2adc Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 13:55:02 +0100 Subject: [PATCH 057/690] Fix typos --- src/pykeen/models/base.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 1c555fc3c2..6d6b14d0bd 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1583,31 +1583,31 @@ def _reset_parameters_(self): def _score( self, - h_ind: Optional[torch.LongTensor] = None, - r_ind: Optional[torch.LongTensor] = None, - t_ind: Optional[torch.LongTensor] = None, + h_indices: Optional[torch.LongTensor] = None, + r_indices: Optional[torch.LongTensor] = None, + t_indices: Optional[torch.LongTensor] = None, ) -> torch.FloatTensor: """Evaluate the given triples. - :param h_ind: shape: (batch_size,) + :param h_indices: shape: (batch_size,) The indices for head entities. If None, score against all. - :param r_ind: shape: (batch_size,) + :param r_indices: shape: (batch_size,) The indices for relations. If None, score against all. - :param t_ind: shape: (batch_size,) + :param t_indices: shape: (batch_size,) The indices for tail entities. If None, score against all. :return: The scores, shape: (batch_size, num_entities) """ - return self.index_function(model=self, h_ind=h_ind, r_ind=r_ind, t_ind=t_ind) + return self.index_function(model=self, h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=hrt_batch[:, 0], r_ind=hrt_batch[:, 1], t_ind=hrt_batch[:, 2]).view(-1, 1) + return self._score(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]).view(-1, 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=hr_batch[:, 0], r_ind=hr_batch[:, 1], t_ind=None).view(-1, self.num_entities) + return self._score(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], t_indices=None).view(-1, self.num_entities) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_ind=None, r_ind=rt_batch[:, 0], t_ind=rt_batch[:, 1]).view(-1, self.num_entities) + return self._score(h_indices=None, r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]).view(-1, self.num_entities) # TODO # def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 From 72d0210f729ddd381dccd67b3cd0938d9a5f6cfe Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 13:55:15 +0100 Subject: [PATCH 058/690] Add missing relation dim --- src/pykeen/models/base.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 6d6b14d0bd..3b65a2e777 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1516,6 +1516,7 @@ def __init__( triples_factory: TriplesFactory, index_function: IndexFunction, embedding_dim: int = 200, + relation_dim: Optional[int] = None, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: Optional[str] = None, @@ -1557,6 +1558,7 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, + relation_dim=relation_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, @@ -1622,6 +1624,7 @@ def __init__( triples_factory: TriplesFactory, interaction_function: InteractionFunction, embedding_dim: int = 200, + relation_dim: Optional[int] = None, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: Optional[str] = None, @@ -1665,6 +1668,7 @@ def __init__( super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, + relation_dim=relation_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, From a31febea2555c4de4f90d7259bf36866b3478ff4 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 13:55:22 +0100 Subject: [PATCH 059/690] Fix imports --- src/pykeen/models/__init__.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index 52aebd75fd..7acfc4de3f 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -9,9 +9,9 @@ import inspect from typing import Mapping, Set, Type, Union -from .base import ( - EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, MultimodalModel, - SimpleVectorEntityRelationEmbeddingModel, GeneralVectorEntityRelationEmbeddingModel +from .base import ( # noqa:F401 + EntityEmbeddingModel, EntityRelationEmbeddingModel, GeneralVectorEntityRelationEmbeddingModel, Model, + MultimodalModel, SimpleVectorEntityRelationEmbeddingModel, ) from .multimodal import ComplExLiteral, DistMultLiteral from .unimodal import ( From 6493f73f7942d9f9fa3e6a254bc5369ed3dd369c Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 13:55:26 +0100 Subject: [PATCH 060/690] Update docs --- src/pykeen/models/base.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 3b65a2e777..edb16795e3 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1537,22 +1537,22 @@ def __init__( ) -> None: """Initialize embedding model. - :param triples_factory: TriplesFactory + :param triples_factory: The triple factory connected to the model. - :param interaction_function: - The interaction function used to compute scores. + :param index_function: + The index-based interaction function used to compute scores. :param embedding_dim: The embedding dimensionality of the entity embeddings. :param automatic_memory_optimization: bool Whether to automatically optimize the sub-batch size during training and batch size during evaluation with regards to the hardware at hand. - :param loss: OptionalLoss (optional) + :param loss: The loss to use. - :param preferred_device: str (optional) + :param preferred_device: The default device where to model is located. - :param random_seed: int (optional) + :param random_seed: An optional random seed to set before the initialization of weights. - :param regularizer: BaseRegularizer + :param regularizer: The regularizer to use. """ super().__init__( @@ -1645,22 +1645,22 @@ def __init__( ) -> None: """Initialize embedding model. - :param triples_factory: TriplesFactory + :param triples_factory: The triple factory connected to the model. :param interaction_function: - The interaction function used to compute scores. + The embedding-based interaction function used to compute scores. :param embedding_dim: The embedding dimensionality of the entity embeddings. - :param automatic_memory_optimization: bool + :param automatic_memory_optimization: Whether to automatically optimize the sub-batch size during training and batch size during evaluation with regards to the hardware at hand. - :param loss: OptionalLoss (optional) + :param loss: The loss to use. - :param preferred_device: str (optional) + :param preferred_device: The default device where to model is located. - :param random_seed: int (optional) + :param random_seed: An optional random seed to set before the initialization of weights. - :param regularizer: BaseRegularizer + :param regularizer: The regularizer to use. """ index_function = InteractionIndexFunction(interaction_function=interaction_function) From 27579b9b2689fc7717b26496c17a3dd9356757e2 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 13:56:02 +0100 Subject: [PATCH 061/690] Reimplement TransD --- src/pykeen/models/unimodal/trans_d.py | 180 ++++++++++++-------------- src/pykeen/utils.py | 4 + 2 files changed, 87 insertions(+), 97 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 3fcdc97cbe..507d3eaa70 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -7,13 +7,14 @@ import torch import torch.autograd -from ..base import EntityRelationEmbeddingModel +from .. import Model +from ..base import GeneralVectorEntityRelationEmbeddingModel, IndexFunction, InteractionFunction from ...losses import Loss from ...nn import Embedding from ...nn.init import xavier_normal_ from ...regularizers import Regularizer from ...triples import TriplesFactory -from ...utils import clamp_norm +from ...utils import DeviceHint, clamp_norm, resolve_device __all__ = [ 'TransD', @@ -67,7 +68,76 @@ def _project_entity( return e_bot -class TransD(EntityRelationEmbeddingModel): +class TransDInteractionFunction(InteractionFunction): + def __init__(self, p: int = 2, power: int = 2): + # Very similar to TransE, could be generalized + super().__init__() + self.p = p + self.power = power + + def forward(self, h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor) -> torch.FloatTensor: + return -torch.norm(h + r - t, dim=-1, p=self.p) ** self.power + + +class TransDIndexFunction(IndexFunction): + """The index-based interaction function for TransD.""" + + def __init__( + self, + num_entities: int, + num_relations: int, + embedding_dim: int, + relation_dim: int, + device: DeviceHint, + interaction_function: Optional[InteractionFunction] = None, + ): + super().__init__() + device = resolve_device(device) + self.entity_projections = Embedding.init_with_device( + num_embeddings=num_entities, + embedding_dim=embedding_dim, + device=device, + initializer=xavier_normal_, + ) + self.relation_projections = Embedding.init_with_device( + num_embeddings=num_relations, + embedding_dim=relation_dim, + device=device, + initializer=xavier_normal_, + ) + if interaction_function is None: + interaction_function = TransDInteractionFunction() + self.interaction_function = interaction_function + + def reset_parameters(self): # noqa: D102 + self.entity_projections.reset_parameters() + self.relation_projections.reset_parameters() + self.interaction_function.reset_parameters() + + def forward( + self, + model: Model, + h_indices: Optional[torch.LongTensor] = None, + r_indices: Optional[torch.LongTensor] = None, + t_indices: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: # noqa: D102 + h = model.entity_embeddings.get_in_canonical_shape(indices=h_indices) + h_p = model.entity_projections.get_in_canonical_shape(indices=h_indices) + + r = model.relation_embeddings.get_in_canonical_shape(indices=r_indices) + r_p = model.relation_projections.get_in_canonical_shape(indices=r_indices) + + t = model.entity_embeddings.get_in_canonical_shape(indices=t_indices) + t_p = model.entity_projections.get_in_canonical_shape(indices=t_indices) + + # Project entities + h_bot = _project_entity(e=h, e_p=h_p, r=r, r_p=r_p) + t_bot = _project_entity(e=t, e_p=t_p, r=r, r_p=r_p) + + return self.interaction_function(h=h_bot, r=r, t=t_bot) + + +class TransD(GeneralVectorEntityRelationEmbeddingModel): r"""An implementation of TransD from [ji2015]_. TransD is an extension of :class:`pykeen.models.TransR` that, like TransR, considers entities and relations @@ -112,12 +182,21 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, relation_dim: int = 30, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: + index_function = TransDIndexFunction( + num_entities=triples_factory.num_entities, + num_relations=triples_factory.num_relations, + embedding_dim=embedding_dim, + relation_dim=relation_dim, + device=preferred_device, + ) + super().__init__( triples_factory=triples_factory, + index_function=index_function, embedding_dim=embedding_dim, relation_dim=relation_dim, automatic_memory_optimization=automatic_memory_optimization, @@ -132,96 +211,3 @@ def __init__( relation_constrainer=clamp_norm, relation_constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ) - - self.entity_projections = Embedding.init_with_device( - num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, - device=self.device, - initializer=xavier_normal_, - ) - self.relation_projections = Embedding.init_with_device( - num_embeddings=triples_factory.num_relations, - embedding_dim=relation_dim, - device=self.device, - initializer=xavier_normal_, - ) - - def _reset_parameters_(self): # noqa: D102 - super()._reset_parameters_() - self.entity_projections.reset_parameters() - self.relation_projections.reset_parameters() - - @staticmethod - def interaction_function( - h: torch.FloatTensor, - h_p: torch.FloatTensor, - r: torch.FloatTensor, - r_p: torch.FloatTensor, - t: torch.FloatTensor, - t_p: torch.FloatTensor, - ) -> torch.FloatTensor: - """Evaluate the interaction function for given embeddings. - - The embeddings have to be in a broadcastable shape. - - :param h: shape: (batch_size, num_entities, d_e) - Head embeddings. - :param h_p: shape: (batch_size, num_entities, d_e) - Head projections. - :param r: shape: (batch_size, num_entities, d_r) - Relation embeddings. - :param r_p: shape: (batch_size, num_entities, d_r) - Relation projections. - :param t: shape: (batch_size, num_entities, d_e) - Tail embeddings. - :param t_p: shape: (batch_size, num_entities, d_e) - Tail projections. - - :return: shape: (batch_size, num_entities) - The scores. - """ - # Project entities - h_bot = _project_entity(e=h, e_p=h_p, r=r, r_p=r_p) - t_bot = _project_entity(e=t, e_p=t_p, r=r, r_p=r_p) - - # score = -||h_bot + r - t_bot||_2^2 - return -torch.norm(h_bot + r - t_bot, dim=-1, p=2) ** 2 - - def _score( - self, - h_indices: Optional[torch.LongTensor] = None, - r_indices: Optional[torch.LongTensor] = None, - t_indices: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: - """ - Evaluate the interaction function. - - :param h_indices: shape: (batch_size,) - The indices of head entities. If None, score against all. - :param r_indices: shape: (batch_size,) - The indices of relations. If None, score against all. - :param t_indices: shape: (batch_size,) - The indices of tail entities. If None, score against all. - - :return: The scores, shape: (batch_size, num_entities) - """ - # Head - h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - h_p = self.entity_projections.get_in_canonical_shape(indices=h_indices) - - r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) - r_p = self.relation_projections.get_in_canonical_shape(indices=r_indices) - - t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - t_p = self.entity_projections.get_in_canonical_shape(indices=t_indices) - - return self.interaction_function(h=h, h_p=h_p, r=r, r_p=r_p, t=t, t_p=t_p) - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], t_indices=None) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=None, r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 81ef763db6..dcfd878fe4 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -26,6 +26,7 @@ 'random_non_negative_int', 'real_part', 'resolve_device', + 'DeviceHint', 'slice_triples', 'slice_doubles', 'split_complex', @@ -74,6 +75,9 @@ def l2_regularization( return regularization_term +DeviceHint = Union[None, str, torch.device] + + def resolve_device(device: Union[None, str, torch.device] = None) -> torch.device: """Resolve a torch.device given a desired device (string).""" if device is None or device == 'gpu': From 30c35a0176f8c14f9990eb1cb099f0ba1ff90f35 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 13:56:52 +0100 Subject: [PATCH 062/690] Cleanup --- src/pykeen/models/base.py | 8 ++++---- src/pykeen/models/unimodal/conv_e.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index edb16795e3..e229d296fc 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1583,7 +1583,7 @@ def _reset_parameters_(self): super()._reset_parameters_() self.index_function.reset_parameters() - def _score( + def forward( self, h_indices: Optional[torch.LongTensor] = None, r_indices: Optional[torch.LongTensor] = None, @@ -1603,13 +1603,13 @@ def _score( return self.index_function(model=self, h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]).view(-1, 1) + return self(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]).view(-1, 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], t_indices=None).view(-1, self.num_entities) + return self(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], t_indices=None).view(-1, self.num_entities) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=None, r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]).view(-1, self.num_entities) + return self(h_indices=None, r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]).view(-1, self.num_entities) # TODO # def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 72c64b0033..3b15e65bf4 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -84,7 +84,7 @@ def wrapped(*args, **kwargs): '\nThis code crash might have been caused by a CUDA bug, see ' 'https://github.com/allenai/allennlp/issues/2888, ' 'which causes the code to crash during evaluation mode.\n' - 'To avoid this error, the batch size has to be reduced.' + 'To avoid this error, the batch size has to be reduced.', ) from e return wrapped From 14cb050f3cb2e6c709d7434197e42b2166b5a678 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 9 Nov 2020 13:56:57 +0100 Subject: [PATCH 063/690] Extract ConvKB interaction function --- src/pykeen/models/unimodal/conv_kb.py | 123 ++++++++++++++++++-------- tests/test_models.py | 4 +- 2 files changed, 89 insertions(+), 38 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 625c1828da..8c460f20f3 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -9,7 +9,7 @@ import torch.autograd from torch import nn -from ..base import EntityRelationEmbeddingModel +from ..base import EntityRelationEmbeddingModel, InteractionFunction from ...losses import Loss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory @@ -21,6 +21,84 @@ logger = logging.getLogger(__name__) +class ConvKBInteractionFunction(InteractionFunction): + """Interaction function of ConvKB.""" + + def __init__( + self, + hidden_dropout_rate: float = 0., + embedding_dim: int = 200, + num_filters: int = 400, + ): + super().__init__() + self.embedding_dim = embedding_dim + self.num_filters = num_filters + + # The interaction model + # self.conv = nn.Conv2d(in_channels=1, out_channels=num_filters, kernel_size=(1, 3), bias=True) + # decompose convolution for faster computation in 1-n case + self.conv_head = nn.Parameter(torch.empty(num_filters)) + self.conv_rel = nn.Parameter(torch.empty(num_filters)) + self.conv_tail = nn.Parameter(torch.empty(num_filters)) + self.conv_bias = nn.Parameter(torch.empty(num_filters)) + self.relu = nn.ReLU() + self.hidden_dropout = nn.Dropout(p=hidden_dropout_rate) + self.linear = nn.Linear(embedding_dim * num_filters, 1, bias=True) + + def reset_parameters(self): # noqa: D102 + # Use Xavier initialization for weight; bias to zero + nn.init.xavier_uniform_(self.linear.weight, gain=nn.init.calculate_gain('relu')) + nn.init.zeros_(self.linear.bias) + + # Initialize all filters to [0.1, 0.1, -0.1], + # c.f. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L34-L36 + # nn.init.constant_(self.conv.weight[..., :2], 0.1) + # nn.init.constant_(self.conv.weight[..., 2], -0.1) + nn.init.constant_(self.conv_head, 0.1) + nn.init.constant_(self.conv_rel, 0.1) + nn.init.constant_(self.conv_tail, 0.1) + nn.init.zeros_(self.conv_bias) + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + # bind sizes + batch_size = max(x.shape[0] for x in (h, r, t)) + num_heads = h.shape[1] + num_relations = r.shape[1] + num_tails = t.shape[1] + + # compute conv(stack(h, r, t)) + # h.shape: (b, nh, d), conv_head.shape: (o), out.shape: (b, nh, d, o) + x = self.conv_bias.view( + 1, 1, 1, 1, 1, self.num_filters + ) + h.view( + h.shape[0], h.shape[1], 1, 1, self.embedding_dim, 1 + ) * self.conv_head.view( + 1, 1, 1, 1, 1, self.num_filters + ) + r.view( + r.shape[0], 1, r.shape[1], 1, self.embedding_dim, 1 + ) * self.conv_rel.view( + 1, 1, 1, 1, 1, self.num_filters + ) + t.view( + t.shape[0], 1, 1, t.shape[1], self.embedding_dim, 1 + ) * self.conv_tail.view( + 1, 1, 1, 1, 1, self.num_filters + ) + + x = self.relu(x) + + # Apply dropout, cf. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L54-L56 + x = self.hidden_dropout(x) + + # Linear layer for final scores + return self.linear(x.view(-1, self.embedding_dim * self.num_filters)).view(batch_size, num_heads, num_relations, num_tails) + + class ConvKB(EntityRelationEmbeddingModel): r"""An implementation of ConvKB from [nguyen2018]_. @@ -95,51 +173,24 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) - - self.num_filters = num_filters - - # The interaction model - self.conv = nn.Conv2d(in_channels=1, out_channels=num_filters, kernel_size=(1, 3), bias=True) - self.relu = nn.ReLU() - self.hidden_dropout = nn.Dropout(p=hidden_dropout_rate) - self.linear = nn.Linear(embedding_dim * num_filters, 1, bias=True) + self.interaction_function = ConvKBInteractionFunction( + hidden_dropout_rate=hidden_dropout_rate, + embedding_dim=embedding_dim, + num_filters=num_filters, + ) def _reset_parameters_(self): # noqa: D102 # embeddings logger.warning('To be consistent with the paper, initialize entity and relation embeddings from TransE.') super()._reset_parameters_() - - # Use Xavier initialization for weight; bias to zero - nn.init.xavier_uniform_(self.linear.weight, gain=nn.init.calculate_gain('relu')) - nn.init.zeros_(self.linear.bias) - - # Initialize all filters to [0.1, 0.1, -0.1], - # c.f. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L34-L36 - nn.init.constant_(self.conv.weight[..., :2], 0.1) - nn.init.constant_(self.conv.weight[..., 2], -0.1) - nn.init.zeros_(self.conv.bias) + self.interaction_function.reset_parameters() def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 h = self.entity_embeddings(indices=hrt_batch[:, 0]) r = self.relation_embeddings(indices=hrt_batch[:, 1]) t = self.entity_embeddings(indices=hrt_batch[:, 2]) - # Output layer regularization # In the code base only the weights of the output layer are used for regularization # c.f. https://github.com/daiquocnguyen/ConvKB/blob/73a22bfa672f690e217b5c18536647c7cf5667f1/model.py#L60-L66 - self.regularize_if_necessary(self.linear.weight, self.linear.bias) - - # Stack to convolution input - conv_inp = torch.stack([h, r, t], dim=-1).view(-1, 1, self.embedding_dim, 3) - - # Convolution - conv_out = self.conv(conv_inp).view(-1, self.embedding_dim * self.num_filters) - hidden = self.relu(conv_out) - - # Apply dropout, cf. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L54-L56 - hidden = self.hidden_dropout(hidden) - - # Linear layer for final scores - scores = self.linear(hidden) - - return scores + self.regularize_if_necessary(self.interaction_function.linear.weight, self.interaction_function.linear.bias) + return self.interaction_function.score_hrt(h=h, r=r, t=t) diff --git a/tests/test_models.py b/tests/test_models.py index aaf5891d4f..26c7c07286 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -559,8 +559,8 @@ class TestConvKB(_ModelTestCase, unittest.TestCase): model_kwargs = { 'num_filters': 2, } - # two bias terms, one conv-filter - num_constant_init = 3 + # two bias terms, one conv-filter, written as 3 parts + num_constant_init = 5 class TestDistMult(_ModelTestCase, unittest.TestCase): From 8216374e141ba1d63876b9c2f6fa9305b6b18b9d Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 14:22:50 +0100 Subject: [PATCH 064/690] Update device hint --- src/pykeen/models/unimodal/trans_d.py | 3 ++- src/pykeen/typing.py | 5 ++++- src/pykeen/utils.py | 8 +++----- 3 files changed, 9 insertions(+), 7 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 507d3eaa70..003c8b62af 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -14,7 +14,8 @@ from ...nn.init import xavier_normal_ from ...regularizers import Regularizer from ...triples import TriplesFactory -from ...utils import DeviceHint, clamp_norm, resolve_device +from ...utils import clamp_norm, resolve_device +from ...typing import DeviceHint __all__ = [ 'TransD', diff --git a/src/pykeen/typing.py b/src/pykeen/typing.py index 8ee4ab04bb..c12e6c78b6 100644 --- a/src/pykeen/typing.py +++ b/src/pykeen/typing.py @@ -2,7 +2,7 @@ """Type hints for PyKEEN.""" -from typing import Callable, Mapping, TypeVar +from typing import Callable, Mapping, TypeVar, Union import numpy as np import torch @@ -13,6 +13,7 @@ 'EntityMapping', 'RelationMapping', 'InteractionFunction', + 'DeviceHint', ] LabeledTriples = np.ndarray @@ -26,3 +27,5 @@ Initializer = Callable[[TensorType], TensorType] Normalizer = Callable[[TensorType], TensorType] Constrainer = Callable[[TensorType], TensorType] + +DeviceHint = Union[None, str, torch.device] diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index dcfd878fe4..dc0a568e5a 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -15,6 +15,8 @@ import torch import torch.nn +from .typing import DeviceHint + __all__ = [ 'compose', 'clamp_norm', @@ -26,7 +28,6 @@ 'random_non_negative_int', 'real_part', 'resolve_device', - 'DeviceHint', 'slice_triples', 'slice_doubles', 'split_complex', @@ -75,10 +76,7 @@ def l2_regularization( return regularization_term -DeviceHint = Union[None, str, torch.device] - - -def resolve_device(device: Union[None, str, torch.device] = None) -> torch.device: +def resolve_device(device: DeviceHint = None) -> torch.device: """Resolve a torch.device given a desired device (string).""" if device is None or device == 'gpu': device = 'cuda' From dd6d186073f1c48c6957f5200052f1d66855a3d2 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 14:31:57 +0100 Subject: [PATCH 065/690] Update trans_d.py --- src/pykeen/models/unimodal/trans_d.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 003c8b62af..56ccf9d0db 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -123,13 +123,13 @@ def forward( t_indices: Optional[torch.LongTensor] = None, ) -> torch.FloatTensor: # noqa: D102 h = model.entity_embeddings.get_in_canonical_shape(indices=h_indices) - h_p = model.entity_projections.get_in_canonical_shape(indices=h_indices) + h_p = self.entity_projections.get_in_canonical_shape(indices=h_indices) r = model.relation_embeddings.get_in_canonical_shape(indices=r_indices) - r_p = model.relation_projections.get_in_canonical_shape(indices=r_indices) + r_p = self.relation_projections.get_in_canonical_shape(indices=r_indices) t = model.entity_embeddings.get_in_canonical_shape(indices=t_indices) - t_p = model.entity_projections.get_in_canonical_shape(indices=t_indices) + t_p = self.entity_projections.get_in_canonical_shape(indices=t_indices) # Project entities h_bot = _project_entity(e=h, e_p=h_p, r=r, r_p=r_p) From 46e9a6d90bfd2a7b5a5c06c9582347bfc29aea88 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 14:36:46 +0100 Subject: [PATCH 066/690] Add TODO @mberr we should do a bit of cleanup in some of these functions to make it a bit more obvious how to abstract. There's also a question here [skip ci] --- src/pykeen/models/unimodal/trans_e.py | 1 + src/pykeen/models/unimodal/trans_r.py | 1 + 2 files changed, 2 insertions(+) diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index a83b248aba..660183586f 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -89,6 +89,7 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: r = self.relation_embeddings(indices=hrt_batch[:, 1]) t = self.entity_embeddings(indices=hrt_batch[:, 2]) + # TODO question @mberr - why is keepdim=True here but the others it isn't? # TODO: Use torch.dist return -torch.norm(h + r - t, dim=-1, p=self.scoring_fct_norm, keepdim=True) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index ce8cca18c3..54dd872e0c 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -158,6 +158,7 @@ def interaction_function( def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings + # TODO switch to self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 0]) h = self.entity_embeddings(indices=hrt_batch[:, 0]).unsqueeze(dim=1) r = self.relation_embeddings(indices=hrt_batch[:, 1]).unsqueeze(dim=1) t = self.entity_embeddings(indices=hrt_batch[:, 2]).unsqueeze(dim=1) From 25e8b8682607fe0424135174af8dc369259d331f Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 14:58:30 +0100 Subject: [PATCH 067/690] pass flake8 --- src/pykeen/models/unimodal/conv_kb.py | 30 +++++++++++++-------------- src/pykeen/models/unimodal/trans_d.py | 2 +- 2 files changed, 16 insertions(+), 16 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 8c460f20f3..0c3194fb0e 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -74,20 +74,18 @@ def forward( # compute conv(stack(h, r, t)) # h.shape: (b, nh, d), conv_head.shape: (o), out.shape: (b, nh, d, o) - x = self.conv_bias.view( - 1, 1, 1, 1, 1, self.num_filters - ) + h.view( - h.shape[0], h.shape[1], 1, 1, self.embedding_dim, 1 - ) * self.conv_head.view( - 1, 1, 1, 1, 1, self.num_filters - ) + r.view( - r.shape[0], 1, r.shape[1], 1, self.embedding_dim, 1 - ) * self.conv_rel.view( - 1, 1, 1, 1, 1, self.num_filters - ) + t.view( - t.shape[0], 1, 1, t.shape[1], self.embedding_dim, 1 - ) * self.conv_tail.view( - 1, 1, 1, 1, 1, self.num_filters + x = ( + self.conv_bias.view(1, 1, 1, 1, 1, self.num_filters) + + ( + h.view(h.shape[0], h.shape[1], 1, 1, self.embedding_dim, 1) + * self.conv_head.view(1, 1, 1, 1, 1, self.num_filters) + ) + ( + r.view(r.shape[0], 1, r.shape[1], 1, self.embedding_dim, 1) + * self.conv_rel.view(1, 1, 1, 1, 1, self.num_filters) + ) + ( + t.view(t.shape[0], 1, 1, t.shape[1], self.embedding_dim, 1) + * self.conv_tail.view(1, 1, 1, 1, 1, self.num_filters) + ) ) x = self.relu(x) @@ -96,7 +94,9 @@ def forward( x = self.hidden_dropout(x) # Linear layer for final scores - return self.linear(x.view(-1, self.embedding_dim * self.num_filters)).view(batch_size, num_heads, num_relations, num_tails) + return self.linear( + x.view(-1, self.embedding_dim * self.num_filters), + ).view(batch_size, num_heads, num_relations, num_tails) class ConvKB(EntityRelationEmbeddingModel): diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 56ccf9d0db..ba62a01476 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -14,8 +14,8 @@ from ...nn.init import xavier_normal_ from ...regularizers import Regularizer from ...triples import TriplesFactory -from ...utils import clamp_norm, resolve_device from ...typing import DeviceHint +from ...utils import clamp_norm, resolve_device __all__ = [ 'TransD', From 6b45f78a7ce674e9f4d1b3e90930b371512ae39e Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 15:33:46 +0100 Subject: [PATCH 068/690] Merge --- .travis.yml | 3 ++ src/pykeen/evaluation/evaluator.py | 2 +- src/pykeen/models/__init__.py | 2 +- src/pykeen/models/base.py | 10 +++---- .../models/multimodal/complex_literal.py | 6 +++- .../models/multimodal/distmult_literal.py | 6 +++- src/pykeen/models/unimodal/complex.py | 3 +- src/pykeen/models/unimodal/conv_e.py | 3 +- src/pykeen/models/unimodal/conv_kb.py | 3 +- src/pykeen/models/unimodal/distmult.py | 3 +- src/pykeen/models/unimodal/ermlp.py | 3 +- src/pykeen/models/unimodal/ermlpe.py | 3 +- src/pykeen/models/unimodal/hole.py | 3 +- src/pykeen/models/unimodal/kg2e.py | 3 +- src/pykeen/models/unimodal/ntn.py | 3 +- src/pykeen/models/unimodal/proj_e.py | 3 +- src/pykeen/models/unimodal/rescal.py | 3 +- src/pykeen/models/unimodal/rgcn.py | 3 +- src/pykeen/models/unimodal/rotate.py | 3 +- src/pykeen/models/unimodal/simple.py | 3 +- .../models/unimodal/structured_embedding.py | 3 +- src/pykeen/models/unimodal/trans_e.py | 3 +- src/pykeen/models/unimodal/trans_h.py | 3 +- src/pykeen/models/unimodal/trans_r.py | 4 ++- src/pykeen/models/unimodal/tucker.py | 3 +- .../models/unimodal/unstructured_model.py | 3 +- src/pykeen/tqdmw.py | 28 ------------------- src/pykeen/training/training_loop.py | 2 +- src/pykeen/triples/leakage.py | 2 +- src/pykeen/triples/triples_factory.py | 2 +- tox.ini | 6 ++++ 31 files changed, 70 insertions(+), 60 deletions(-) delete mode 100644 src/pykeen/tqdmw.py diff --git a/.travis.yml b/.travis.yml index ebb7ebef36..f786e42416 100644 --- a/.travis.yml +++ b/.travis.yml @@ -15,6 +15,8 @@ jobs: env: TOXENV=flake8 - name: "Check documentation with darglint" env: TOXENV=darglint + # - name: "Check typing with MyPy" + # env: TOXENV=mypy # docs stage - name: "Check RST documentation style" stage: docs @@ -32,6 +34,7 @@ jobs: env: TOXENV=integration allow_failures: - env: TOXENV=darglint + # - env: TOXENV=mypy install: - sh -c 'if [ "$TOXENV" = "py" ]; then pip install tox codecov coverage; else pip install tox; fi' - sh -c 'if [ "$TOXENV" = "docs" ]; then sudo apt-get install graphviz; fi' diff --git a/src/pykeen/evaluation/evaluator.py b/src/pykeen/evaluation/evaluator.py index 59e9efb29e..6af4860259 100644 --- a/src/pykeen/evaluation/evaluator.py +++ b/src/pykeen/evaluation/evaluator.py @@ -13,9 +13,9 @@ import torch from dataclasses_json import dataclass_json +from tqdm.autonotebook import tqdm from ..models.base import Model -from ..tqdmw import tqdm from ..triples.triples_factory import get_unique_entity_ids_from_triples_tensor from ..typing import MappedTriples from ..utils import is_cuda_oom_error, is_cudnn_error, normalize_string, split_list_in_batches_iter diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index 7acfc4de3f..c81b8bd2f0 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -70,7 +70,7 @@ def _concrete_subclasses(cls): for subcls in cls.__subclasses__(): - if not inspect.isabstract(cls): + if not inspect.isabstract(subcls): yield subcls yield from _concrete_subclasses(subcls) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index e229d296fc..4fd38fb5c6 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -19,7 +19,7 @@ from ..nn import Embedding from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory -from ..typing import Constrainer, Initializer, MappedTriples, Normalizer +from ..typing import Constrainer, DeviceHint, Initializer, MappedTriples, Normalizer from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed __all__ = [ @@ -236,7 +236,7 @@ def __init__( loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, automatic_memory_optimization: Optional[bool] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: @@ -367,7 +367,7 @@ def num_relations(self) -> int: # noqa: D401 """The number of unique relation types in the knowledge graph.""" return self.triples_factory.num_relations - def _set_device(self, device: Union[None, str, torch.device] = None) -> None: + def _set_device(self, device: DeviceHint = None) -> None: """Set the Torch device to use.""" self.device = resolve_device(device=device) @@ -1046,7 +1046,7 @@ def __init__( loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, automatic_memory_optimization: Optional[bool] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, entity_initializer: Optional[Initializer] = None, @@ -1110,7 +1110,7 @@ def __init__( loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, automatic_memory_optimization: Optional[bool] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, entity_initializer: Optional[Initializer] = None, diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index f37d6013a9..9e7d585ea2 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -11,6 +11,7 @@ from ..base import MultimodalModel from ...losses import BCEWithLogitsLoss, Loss from ...triples import TriplesNumericLiteralsFactory +from ...typing import DeviceHint from ...utils import slice_doubles @@ -39,7 +40,7 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, input_dropout: float = 0.2, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, ) -> None: """Initialize the model.""" @@ -95,6 +96,9 @@ def _apply_g_function(self, real_embs, img_embs, literals): img = self.img_non_lin_transf(torch.cat([img_embs, literals], 1)) return real, img + def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa:D102 + raise NotImplementedError + def score_t(self, doubles: torch.Tensor) -> torch.Tensor: """Forward pass using right side (tail) prediction for training with the LCWA.""" batch_heads, batch_relations = slice_doubles(doubles) diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index dd7621fde5..67090a0a5c 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -11,6 +11,7 @@ from ..base import MultimodalModel from ...losses import Loss from ...triples import TriplesNumericLiteralsFactory +from ...typing import DeviceHint from ...utils import slice_triples @@ -33,7 +34,7 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, input_dropout: float = 0.0, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, ) -> None: super().__init__( @@ -111,6 +112,9 @@ def _apply_g_function(self, entity_embeddings, literals): """ return self.linear_transformation(torch.cat([entity_embeddings, literals], dim=1)) + def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa:D102 + raise NotImplementedError + def score_t(self, hr_batch: torch.Tensor) -> torch.Tensor: """Forward pass using right side (tail) prediction for training with the LCWA.""" heads, relations, tails = slice_triples(hr_batch) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index c39ce0e388..d2a7e16aa8 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -11,6 +11,7 @@ from ...losses import Loss, SoftplusLoss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint from ...utils import normalize_for_einsum, split_complex __all__ = [ @@ -99,7 +100,7 @@ def __init__( embedding_dim: int = 200, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 3b15e65bf4..679253e66d 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -15,6 +15,7 @@ from ...nn.init import xavier_normal_ from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint from ...utils import is_cudnn_error __all__ = [ @@ -337,7 +338,7 @@ def __init__( embedding_dim: int = 200, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, apply_batch_normalization: bool = True, diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 0c3194fb0e..977cf2d086 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -13,6 +13,7 @@ from ...losses import Loss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint __all__ = [ 'ConvKB', @@ -155,7 +156,7 @@ def __init__( embedding_dim: int = 200, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, num_filters: int = 400, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 0cb64dd42b..be4aa08475 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -12,6 +12,7 @@ from ...losses import Loss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint from ...utils import compose from ...utils import normalize_for_einsum @@ -91,7 +92,7 @@ def __init__( embedding_dim: int = 50, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index ce9a8a0a1e..fcdca0837d 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -12,6 +12,7 @@ from ...losses import Loss from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint __all__ = [ 'ERMLP', @@ -116,7 +117,7 @@ def __init__( embedding_dim: int = 50, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, hidden_dim: Optional[int] = None, regularizer: Optional[Regularizer] = None, diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index e1a2fe7f64..05976b4bd6 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -11,6 +11,7 @@ from ...losses import BCEAfterSigmoidLoss, Loss from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint __all__ = [ 'ERMLPE', @@ -61,7 +62,7 @@ def __init__( embedding_dim: int = 200, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index cf07f5c6bf..29103fda72 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -12,6 +12,7 @@ from ...nn.init import xavier_uniform_ from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint from ...utils import clamp_norm __all__ = [ @@ -88,7 +89,7 @@ def __init__( embedding_dim: int = 200, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 1cb7d26143..846f887acc 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -13,6 +13,7 @@ from ...nn import Embedding from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint from ...utils import clamp_norm __all__ = [ @@ -62,7 +63,7 @@ def __init__( embedding_dim: int = 50, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, dist_similarity: Optional[str] = None, c_min: float = 0.05, diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index eef836e2d3..30c28d3b75 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -11,6 +11,7 @@ from ...losses import Loss from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint __all__ = [ 'NTN', @@ -57,7 +58,7 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, num_slices: int = 4, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, non_linearity: Optional[nn.Module] = None, regularizer: Optional[Regularizer] = None, diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index 8d8dbee1eb..98e39e5a59 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -14,6 +14,7 @@ from ...nn.init import xavier_uniform_ from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint __all__ = [ 'ProjE', @@ -60,7 +61,7 @@ def __init__( embedding_dim: int = 50, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, inner_non_linearity: Optional[nn.Module] = None, regularizer: Optional[Regularizer] = None, diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index 44f8348f8c..d158b1e5cf 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -10,6 +10,7 @@ from ...losses import Loss from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint __all__ = [ 'RESCAL', @@ -55,7 +56,7 @@ def __init__( embedding_dim: int = 50, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index aab8170a74..79bd5a61f7 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -16,6 +16,7 @@ from ...losses import Loss from ...nn import Embedding, RepresentationModule from ...triples import TriplesFactory +from ...typing import DeviceHint __all__ = [ 'RGCN', @@ -488,7 +489,7 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, num_bases_or_blocks: int = 5, num_layers: int = 2, diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index 63718d11ba..e1d3501447 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -14,6 +14,7 @@ from ...nn.init import xavier_uniform_ from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint __all__ = [ 'RotatE', @@ -86,7 +87,7 @@ def __init__( embedding_dim: int = 200, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index ec2d72a41d..9684334355 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -11,6 +11,7 @@ from ...nn import Embedding from ...regularizers import PowerSumRegularizer, Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint __all__ = [ 'SimplE', @@ -67,7 +68,7 @@ def __init__( embedding_dim: int = 200, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, clamp_score: Optional[Union[float, Tuple[float, float]]] = None, diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index e48adb632c..1401504172 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -17,6 +17,7 @@ from ...nn.init import xavier_uniform_ from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint from ...utils import compose __all__ = [ @@ -53,7 +54,7 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 660183586f..b653ed66a3 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -13,6 +13,7 @@ from ...nn.init import xavier_uniform_ from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint from ...utils import compose __all__ = [ @@ -53,7 +54,7 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index bbd6671e92..b233615a82 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -12,6 +12,7 @@ from ...nn import Embedding from ...regularizers import Regularizer, TransHRegularizer from ...triples import TriplesFactory +from ...typing import DeviceHint __all__ = [ 'TransH', @@ -68,7 +69,7 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 54dd872e0c..17cf2d706d 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- """Implementation of TransR.""" + from functools import partial from typing import Optional @@ -15,6 +16,7 @@ from ...nn.init import xavier_uniform_ from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint from ...utils import clamp_norm, compose __all__ = [ @@ -78,7 +80,7 @@ def __init__( relation_dim: int = 30, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index c8a88bf6da..8dd9af3cfd 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -13,6 +13,7 @@ from ...nn.init import xavier_normal_ from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint __all__ = [ 'TuckER', @@ -82,7 +83,7 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, relation_dim: Optional[int] = None, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, dropout_0: float = 0.3, dropout_1: float = 0.4, diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index fd9b73a52d..ef3265c7c0 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -12,6 +12,7 @@ from ...nn.init import xavier_normal_ from ...regularizers import Regularizer from ...triples import TriplesFactory +from ...typing import DeviceHint __all__ = [ 'UnstructuredModel', @@ -49,7 +50,7 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, + preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: diff --git a/src/pykeen/tqdmw.py b/src/pykeen/tqdmw.py deleted file mode 100644 index 0c18987c1b..0000000000 --- a/src/pykeen/tqdmw.py +++ /dev/null @@ -1,28 +0,0 @@ -# -*- coding: utf-8 -*- - -"""A wrapper around tqdm to make it automatically play nice with Jupyter Notebook.""" - -__all__ = [ - 'tqdm', - 'trange', -] - - -def is_notebook() -> bool: - """Check if we're running in a Jupyter notebook.""" - try: - shell = get_ipython().__class__.__name__ - if shell == 'ZMQInteractiveShell': - return True # Jupyter notebook or qtconsole - elif shell == 'TerminalInteractiveShell': - return False # Terminal running IPython - else: - return False # Other type (?) - except NameError: - return False # Probably standard Python interpreter - - -if is_notebook(): - from tqdm.notebook import tqdm, trange -else: - from tqdm import tqdm, trange diff --git a/src/pykeen/training/training_loop.py b/src/pykeen/training/training_loop.py index 32ec1f97c1..f59f46b3ee 100644 --- a/src/pykeen/training/training_loop.py +++ b/src/pykeen/training/training_loop.py @@ -10,11 +10,11 @@ import torch from torch.optim.optimizer import Optimizer from torch.utils.data import DataLoader +from tqdm.autonotebook import tqdm, trange from ..losses import Loss from ..models.base import Model from ..stoppers import Stopper -from ..tqdmw import tqdm, trange from ..trackers import ResultTracker from ..training.schlichtkrull_sampler import GraphSampler from ..triples import Instances, TriplesFactory diff --git a/src/pykeen/triples/leakage.py b/src/pykeen/triples/leakage.py index 01ba1d87af..ef1ec48f23 100644 --- a/src/pykeen/triples/leakage.py +++ b/src/pykeen/triples/leakage.py @@ -18,9 +18,9 @@ import numpy as np from tabulate import tabulate +from tqdm.autonotebook import tqdm from .triples_factory import TriplesFactory, create_entity_mapping, create_relation_mapping -from ..tqdmw import tqdm from ..typing import LabeledTriples __all__ = [ diff --git a/src/pykeen/triples/triples_factory.py b/src/pykeen/triples/triples_factory.py index 68231daae8..b5f3b61521 100644 --- a/src/pykeen/triples/triples_factory.py +++ b/src/pykeen/triples/triples_factory.py @@ -11,10 +11,10 @@ import numpy as np import pandas as pd import torch +from tqdm.autonotebook import tqdm from .instances import LCWAInstances, SLCWAInstances from .utils import load_triples -from ..tqdmw import tqdm from ..typing import EntityMapping, LabeledTriples, MappedTriples, RelationMapping from ..utils import compact_mapping, invert_mapping, random_non_negative_int, slice_triples diff --git a/tox.ini b/tox.ini index bd1180aebd..d3e54783cf 100644 --- a/tox.ini +++ b/tox.ini @@ -106,6 +106,12 @@ commands = add-trailing-comma --py36-plus $(find src/pykeen/ -name *.py) description = Check all python files do not have mistaken trailing commas +[testenv:mypy] +deps = mypy +skip_install = true +commands = mypy --ignore-missing-imports src/pykeen/ +description = Run the mypy tool to check static typing on the project. + [testenv:pyroma] deps = pygments From a3273e8c3a25f2c0366005bdfc24d71a1095ef1e Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 15:43:07 +0100 Subject: [PATCH 069/690] Make device resolvable in Embedding --- src/pykeen/models/unimodal/trans_d.py | 3 +-- src/pykeen/nn/emb.py | 9 +++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index ba62a01476..bf23e4cdd7 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -15,7 +15,7 @@ from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint -from ...utils import clamp_norm, resolve_device +from ...utils import clamp_norm __all__ = [ 'TransD', @@ -93,7 +93,6 @@ def __init__( interaction_function: Optional[InteractionFunction] = None, ): super().__init__() - device = resolve_device(device) self.entity_projections = Embedding.init_with_device( num_embeddings=num_entities, embedding_dim=embedding_dim, diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index c535528938..f7c5862489 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -9,13 +9,14 @@ import torch.nn from torch import nn +from ..typing import Constrainer, DeviceHint, Initializer, Normalizer +from ..utils import resolve_device + __all__ = [ 'RepresentationModule', 'Embedding', ] -from pykeen.typing import Constrainer, Initializer, Normalizer - class RepresentationModule(nn.Module): """A base class for obtaining representations for entities/relations.""" @@ -108,7 +109,7 @@ def init_with_device( cls, num_embeddings: int, embedding_dim: int, - device: torch.device, + device: DeviceHint, initializer: Optional[Initializer] = None, initializer_kwargs: Optional[Mapping[str, Any]] = None, normalizer: Optional[Normalizer] = None, @@ -138,7 +139,7 @@ def init_with_device( normalizer_kwargs=normalizer_kwargs, constrainer=constrainer, constrainer_kwargs=constrainer_kwargs, - ).to(device=device) + ).to(device=resolve_device(device)) @property def num_embeddings(self) -> int: # noqa: D401 From c992fba73c1af15c45803f0153d90bba7e2dd883 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 16:26:13 +0100 Subject: [PATCH 070/690] Fix init problem --- src/pykeen/models/base.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 4fd38fb5c6..571c84af58 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1508,7 +1508,7 @@ def reset_parameters(self): # noqa: D102 self.interaction_function.reset_parameters() -class GeneralVectorEntityRelationEmbeddingModel(EntityRelationEmbeddingModel): +class GeneralVectorEntityRelationEmbeddingModel(EntityRelationEmbeddingModel, reset_parameters_post_init=False): """A base class for embedding models which store a single vector for each entity and relation.""" def __init__( @@ -1616,7 +1616,9 @@ def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 # raise NotImplementedError -class SimpleVectorEntityRelationEmbeddingModel(GeneralVectorEntityRelationEmbeddingModel): +class SimpleVectorEntityRelationEmbeddingModel( + GeneralVectorEntityRelationEmbeddingModel, reset_parameters_post_init=False, +): """A base class for embedding models which store a single vector for each entity and relation.""" def __init__( From 98a6d635bd15c58471217dc05cdcb9d90c9817ca Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 16:32:04 +0100 Subject: [PATCH 071/690] Update test_models.py --- tests/test_models.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 26c7c07286..b5c9302ce9 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1263,8 +1263,6 @@ def test_abstract(self): self.assertTrue(Model._is_abstract()) self.assertTrue(EntityEmbeddingModel._is_abstract()) self.assertTrue(EntityRelationEmbeddingModel._is_abstract()) - self.assertTrue(GeneralVectorEntityRelationEmbeddingModel._is_abstract()) - self.assertTrue(SimpleVectorEntityRelationEmbeddingModel._is_abstract()) for model_cls in _MODELS: if issubclass(model_cls, MultimodalModel): continue From 06fd38fa8fd9833f803984e6cf8e36b0f8aaf868 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 16:33:23 +0100 Subject: [PATCH 072/690] Update test_models.py [skip ci] --- tests/test_models.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_models.py b/tests/test_models.py index b5c9302ce9..876a4229d0 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -879,6 +879,7 @@ def test_score_hrt_manual(self): self.model.relation_projections = relation_projection_embeddings # Compute Scores + # FIXME batch is wrong size? batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 1]], dtype=torch.long) scores = self.model.score_hrt(hrt_batch=batch) self.assertEqual(scores.shape[0], 2) From 90f1ab68fd0d4f9ccf6a85381a155d29e7f7d1ec Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 9 Nov 2020 17:46:36 +0100 Subject: [PATCH 073/690] Update __init__.py --- src/pykeen/models/__init__.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index c81b8bd2f0..e3e4b0abbc 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -6,7 +6,6 @@ score value is model-dependent, and usually it cannot be directly interpreted as a probability. """ # noqa: D205, D400 -import inspect from typing import Mapping, Set, Type, Union from .base import ( # noqa:F401 @@ -68,9 +67,15 @@ ] -def _concrete_subclasses(cls): +_CONCRETE_BASES = { + GeneralVectorEntityRelationEmbeddingModel, + SimpleVectorEntityRelationEmbeddingModel, +} + + +def _concrete_subclasses(cls: Type[Model]): for subcls in cls.__subclasses__(): - if not inspect.isabstract(subcls): + if not subcls._is_abstract() and subcls not in _CONCRETE_BASES: yield subcls yield from _concrete_subclasses(subcls) From a73559358e271502a33c743e7d92864b1b441352 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 11:28:04 +0100 Subject: [PATCH 074/690] Extract functional form of distmult interaction --- src/pykeen/models/unimodal/distmult.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index be4aa08475..96dd4f8db4 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -10,11 +10,11 @@ from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss +from ...nn import functional as pykeen_functional from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint from ...utils import compose -from ...utils import normalize_for_einsum __all__ = [ 'DistMult', @@ -33,11 +33,7 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa: D102 self._check_for_empty_kwargs(kwargs) - batch_size = max(h.shape[0], r.shape[0], t.shape[0]) - h_term, h = normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') - r_term, r = normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') - t_term, t = normalize_for_einsum(x=t, batch_size=batch_size, symbol='t') - return torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', h, r, t) + return pykeen_functional.distmult_interaction(h=h, r=r, t=t) class DistMult(SimpleVectorEntityRelationEmbeddingModel): From 0ec6f8080c1a75d0643cda11ba0952e14cda89c2 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 11:28:19 +0100 Subject: [PATCH 075/690] Add missing file --- src/pykeen/nn/functional.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) create mode 100644 src/pykeen/nn/functional.py diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py new file mode 100644 index 0000000000..d3760adeca --- /dev/null +++ b/src/pykeen/nn/functional.py @@ -0,0 +1,16 @@ +"""Functional forms of interaction methods.""" +import torch + +from ..utils import normalize_for_einsum + + +def distmult_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + batch_size = max(h.shape[0], r.shape[0], t.shape[0]) + h_term, h = normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') + r_term, r = normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') + t_term, t = normalize_for_einsum(x=t, batch_size=batch_size, symbol='t') + return torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', h, r, t) From 336d502eb5fcc6e70726b12c6610bfa2cad34368 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 11:31:44 +0100 Subject: [PATCH 076/690] Extract functional complex interaction function --- src/pykeen/models/unimodal/complex.py | 17 ++---------- src/pykeen/nn/functional.py | 40 +++++++++++++++++++++++++-- 2 files changed, 39 insertions(+), 18 deletions(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index d2a7e16aa8..978914f75a 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -9,10 +9,10 @@ from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss, SoftplusLoss +from ...nn import functional as pykeen_functional from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint -from ...utils import normalize_for_einsum, split_complex __all__ = [ 'ComplEx', @@ -31,20 +31,7 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa: D102 self._check_for_empty_kwargs(kwargs) - batch_size = max(h.shape[0], r.shape[0], t.shape[0]) - h_term, h = normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') - r_term, r = normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') - t_term, t = normalize_for_einsum(x=t, batch_size=batch_size, symbol='t') - (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] - return sum( - torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', hh, rr, tt) - for hh, rr, tt in [ - (h_re, r_re, t_re), - (h_re, r_im, t_im), - (h_im, r_re, t_im), - (h_im, r_im, t_re), - ] - ) + return pykeen_functional.complex_interaction(h=h, r=r, t=t) class ComplEx(SimpleVectorEntityRelationEmbeddingModel): diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index d3760adeca..eec3ef6ae7 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -1,16 +1,50 @@ """Functional forms of interaction methods.""" +from typing import Tuple + import torch -from ..utils import normalize_for_einsum +from ..utils import normalize_for_einsum, split_complex +__all__ = [ + "complex_interaction", + "distmult_interaction", +] -def distmult_interaction( + +def _normalize_terms_for_einsum( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, -) -> torch.FloatTensor: +) -> Tuple[torch.FloatTensor, str, torch.FloatTensor, str, torch.FloatTensor, str]: batch_size = max(h.shape[0], r.shape[0], t.shape[0]) h_term, h = normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') r_term, r = normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') t_term, t = normalize_for_einsum(x=t, batch_size=batch_size, symbol='t') + return h, h_term, r, r_term, t, t_term + + +def distmult_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + h, h_term, r, r_term, t, t_term = _normalize_terms_for_einsum(h, r, t) return torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', h, r, t) + + +def complex_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + h, h_term, r, r_term, t, t_term = _normalize_terms_for_einsum(h, r, t) + (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] + return sum( + torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', hh, rr, tt) + for hh, rr, tt in [ + (h_re, r_re, t_re), + (h_re, r_im, t_im), + (h_im, r_re, t_im), + (h_im, r_im, t_re), + ] + ) From ee3f48b59f59cd0891a85beac77f2519fe937be4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 11:38:21 +0100 Subject: [PATCH 077/690] Add functional conve interaction function --- src/pykeen/models/unimodal/conv_e.py | 96 ++++++------------------- src/pykeen/nn/functional.py | 100 ++++++++++++++++++++++++++- 2 files changed, 119 insertions(+), 77 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 679253e66d..fa2c13a6b3 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -11,12 +11,11 @@ from ..base import EntityRelationEmbeddingModel, InteractionFunction from ...losses import BCEAfterSigmoidLoss, Loss -from ...nn import Embedding +from ...nn import Embedding, functional as pykeen_functional from ...nn.init import xavier_normal_ from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint -from ...utils import is_cudnn_error __all__ = [ 'ConvE', @@ -74,23 +73,6 @@ def _calculate_missing_shape_information( return input_channels, width, height -def _add_cuda_warning(func): - def wrapped(*args, **kwargs): - try: - return func(*args, **kwargs) - except RuntimeError as e: - if not is_cudnn_error(e): - raise e - raise RuntimeError( - '\nThis code crash might have been caused by a CUDA bug, see ' - 'https://github.com/allenai/allennlp/issues/2888, ' - 'which causes the code to crash during evaluation mode.\n' - 'To avoid this error, the batch size has to be reduced.', - ) from e - - return wrapped - - class ConvEInteractionFunction(InteractionFunction): """ConvE interaction function.""" @@ -171,7 +153,6 @@ def __init__( self.fc = nn.Linear(self.num_in_features, self.embedding_dim) self.activation = nn.ReLU() - @_add_cuda_warning def forward( self, h: torch.FloatTensor, @@ -184,61 +165,26 @@ def forward( raise TypeError(f"{self.__class__.__name__}.forward expects keyword argument 't_bias'.") t_bias: torch.FloatTensor = kwargs.pop("t_bias") self._check_for_empty_kwargs(kwargs) - - # bind sizes - batch_size = max(x.shape[0] for x in (h, r, t)) - num_heads = h.shape[1] - num_relations = r.shape[1] - num_tails = t.shape[1] - - # repeat if necessary - h = h.unsqueeze(dim=2).repeat(1 if h.shape[0] == batch_size else batch_size, 1, num_relations, 1) - r = r.unsqueeze(dim=1).repeat(1 if r.shape[0] == batch_size else batch_size, num_heads, 1, 1) - - # resize and concat head and relation, batch_size', num_input_channels, 2*height, width - # with batch_size' = batch_size * num_heads * num_relations - x = torch.cat([ - h.view(-1, self.input_channels, self.embedding_height, self.embedding_width), - r.view(-1, self.input_channels, self.embedding_height, self.embedding_width), - ], dim=2) - - # batch_size, num_input_channels, 2*height, width - if self.apply_batch_normalization: - x = self.bn0(x) - - # batch_size, num_input_channels, 2*height, width - x = self.inp_drop(x) - - # (N,C_out,H_out,W_out) - x = self.conv1(x) - - if self.apply_batch_normalization: - x = self.bn1(x) - - x = self.activation(x) - x = self.feature_map_drop(x) - - # batch_size', num_output_channels * (2 * height - kernel_height + 1) * (width - kernel_width + 1) - x = x.view(-1, self.num_in_features) - x = self.fc(x) - x = self.hidden_drop(x) - - if self.apply_batch_normalization: - x = self.bn2(x) - x = self.activation(x) - - # reshape: (batch_size', embedding_dim) - x = x.view(batch_size, num_heads, num_relations, 1, self.embedding_dim) - - # For efficient calculation, each of the convolved [h, r] rows has only to be multiplied with one t row - # output_shape: (batch_size, num_heads, num_relations, num_tails) - t = t.view(t.shape[0], 1, 1, num_tails, self.embedding_dim).transpose(-1, -2) - x = (x @ t).squeeze(dim=-2) - - # add bias term - x = x + t_bias.view(t.shape[0], 1, 1, num_tails) - - return x + return pykeen_functional.conve_interaction( + h=h, + r=r, + t=t, + t_bias=t_bias, + input_channels=self.input_channels, + embedding_height=self.embedding_height, + embedding_width=self.embedding_width, + num_in_features=self.num_in_features, + embedding_dim=self.embedding_dim, + bn0=self.bn0, + bn1=self.bn1, + bn2=self.bn2, + inp_drop=self.inp_drop, + feature_map_drop=self.feature_map_drop, + hidden_drop=self.hidden_drop, + conv1=self.conv1, + activation=self.activation, + fc=self.fc, + ) class ConvE(EntityRelationEmbeddingModel): diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index eec3ef6ae7..0b31a6c524 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -1,12 +1,14 @@ """Functional forms of interaction methods.""" -from typing import Tuple +from typing import Optional, Tuple import torch +from torch import nn -from ..utils import normalize_for_einsum, split_complex +from ..utils import is_cudnn_error, normalize_for_einsum, split_complex __all__ = [ "complex_interaction", + "conve_interaction", "distmult_interaction", ] @@ -23,6 +25,100 @@ def _normalize_terms_for_einsum( return h, h_term, r, r_term, t, t_term +def _add_cuda_warning(func): + def wrapped(*args, **kwargs): + try: + return func(*args, **kwargs) + except RuntimeError as e: + if not is_cudnn_error(e): + raise e + raise RuntimeError( + '\nThis code crash might have been caused by a CUDA bug, see ' + 'https://github.com/allenai/allennlp/issues/2888, ' + 'which causes the code to crash during evaluation mode.\n' + 'To avoid this error, the batch size has to be reduced.', + ) from e + + return wrapped + + +@_add_cuda_warning +def conve_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + t_bias: torch.FloatTensor, + input_channels: int, + embedding_height: int, + embedding_width: int, + num_in_features: int, + embedding_dim: int, + bn0: Optional[nn.BatchNorm1d], + bn1: Optional[nn.BatchNorm1d], + bn2: Optional[nn.BatchNorm1d], + inp_drop: nn.Dropout, + feature_map_drop: nn.Dropout2d, + hidden_drop: nn.Dropout, + conv1: nn.Conv2d, + activation: nn.Module, + fc: nn.Linear, +) -> torch.FloatTensor: + # bind sizes + batch_size = max(x.shape[0] for x in (h, r, t)) + num_heads = h.shape[1] + num_relations = r.shape[1] + num_tails = t.shape[1] + + # repeat if necessary + h = h.unsqueeze(dim=2).repeat(1 if h.shape[0] == batch_size else batch_size, 1, num_relations, 1) + r = r.unsqueeze(dim=1).repeat(1 if r.shape[0] == batch_size else batch_size, num_heads, 1, 1) + + # resize and concat head and relation, batch_size', num_input_channels, 2*height, width + # with batch_size' = batch_size * num_heads * num_relations + x = torch.cat([ + h.view(-1, input_channels, embedding_height, embedding_width), + r.view(-1, input_channels, embedding_height, embedding_width), + ], dim=2) + + # batch_size, num_input_channels, 2*height, width + if bn0 is not None: + x = bn0(x) + + # batch_size, num_input_channels, 2*height, width + x = inp_drop(x) + + # (N,C_out,H_out,W_out) + x = conv1(x) + + if bn1 is not None: + x = bn1(x) + + x = activation(x) + x = feature_map_drop(x) + + # batch_size', num_output_channels * (2 * height - kernel_height + 1) * (width - kernel_width + 1) + x = x.view(-1, num_in_features) + x = fc(x) + x = hidden_drop(x) + + if bn2 is not None: + x = bn2(x) + x = activation(x) + + # reshape: (batch_size', embedding_dim) + x = x.view(batch_size, num_heads, num_relations, 1, embedding_dim) + + # For efficient calculation, each of the convolved [h, r] rows has only to be multiplied with one t row + # output_shape: (batch_size, num_heads, num_relations, num_tails) + t = t.view(t.shape[0], 1, 1, num_tails, embedding_dim).transpose(-1, -2) + x = (x @ t).squeeze(dim=-2) + + # add bias term + x = x + t_bias.view(t.shape[0], 1, 1, num_tails) + + return x + + def distmult_interaction( h: torch.FloatTensor, r: torch.FloatTensor, From f0c20d63c956296099b4e0d7c4739431cd3abbcb Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 11:38:45 +0100 Subject: [PATCH 078/690] Export is_cudnn_error --- src/pykeen/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index dc0a568e5a..f85019ca25 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -23,6 +23,7 @@ 'compact_mapping', 'imag_part', 'invert_mapping', + 'is_cudnn_error', 'l2_regularization', 'is_cuda_oom_error', 'random_non_negative_int', From a0226896a40b966b917c1025259e8a210680e1f0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 11:54:55 +0100 Subject: [PATCH 079/690] Move utility method --- src/pykeen/models/base.py | 33 +++++---------------------------- src/pykeen/utils.py | 36 ++++++++++++++++++++++++++++++++++++ 2 files changed, 41 insertions(+), 28 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 571c84af58..ef1ee1f3e7 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -20,7 +20,7 @@ from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory from ..typing import Constrainer, DeviceHint, Initializer, MappedTriples, Normalizer -from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed +from ..utils import NoRandomSeedNecessary, check_shapes, resolve_device, set_random_seed __all__ = [ 'Model', @@ -1276,29 +1276,6 @@ def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: x = x.squeeze(dim=dim) return x - def _check_shape( - self, - *x: Union[str, torch.Tensor], - ) -> bool: - assert len(x) % 2 == 0 - tensors = x[::2] - shapes = x[1::2] - assert all(torch.is_tensor(t) for t in tensors) - assert all(isinstance(s, str) for s in shapes) - dims = dict() - errors = [] - for tensor, shape in zip(tensors, shapes): - if tensor.ndimension() != len(shape): - errors.append(f"Invalid number of dimensions: {tensor.shape} vs. {shape}") - continue - for dim, name in zip(tensor.shape, shape): - exp_dim = dims.get(name) - if exp_dim is not None and exp_dim != dim: - errors.append(f"{name}: {dim} vs. {exp_dim}") - if len(errors) > 0: - raise ValueError("Shape verification failed:\n" + '\n'.join(errors)) - return True - def score_hrt( self, h: torch.FloatTensor, @@ -1322,7 +1299,7 @@ def score_hrt( The scores. """ # check shape - assert self._check_shape(h, "be", r, "br", t, "be") + assert check_shapes((h, "be"), (r, "br"), (t, "be")) # prepare input to generic score function h, r, t = self._add_dim(h, r, t, dim=self.NUM_DIM) @@ -1356,7 +1333,7 @@ def score_h( The scores. """ # check shape - assert self._check_shape(all_entities, "ne", r, "br", t, "be") + assert check_shapes((all_entities, "ne"), (r, "br"), (t, "be")) # TODO: What about unsqueezing for additional e.g. head arguments # prepare input to generic score function @@ -1392,7 +1369,7 @@ def score_r( The scores. """ # check shape - assert self._check_shape(all_relations, "nr", h, "be", t, "be") + assert check_shapes((all_relations, "nr"), (h, "be"), (t, "be")) # prepare input to generic score function h, t = self._add_dim(h, t, dim=self.NUM_DIM) @@ -1427,7 +1404,7 @@ def score_t( The scores. """ # check shape - assert self._check_shape(all_entities, "ne", r, "br", h, "be") + assert check_shapes(all_entities, "ne", r, "br", h, "be") # prepare input to generic score function h, r = self._add_dim(h, r, dim=self.NUM_DIM) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index f85019ca25..5cbc6a9876 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -19,6 +19,7 @@ __all__ = [ 'compose', + 'check_shapes', 'clamp_norm', 'compact_mapping', 'imag_part', @@ -455,3 +456,38 @@ def normalize_for_einsum( if x.shape[0] == batch_size: return f'b{symbol}d', x return f'{symbol}d', x.squeeze(dim=0) + + +def check_shapes( + *x: Tuple[torch.Tensor, str], + raise_or_error: bool = True, +) -> bool: + """ + Verify that a sequence of tensors are of matching shapes. + + :param x: + A tuple (tensor, shape), where tensor is a tensor, and shape is a string, where each character corresponds to + a (named) dimension. If the shapes of different tensors share a character, the corresponding dimensions are + expected to be of equal size. + :param raise_or_error: + Whether to raise an exception in case of a mismatch. + + :return: + Whether the shapes matched. + + :raises ValueError: + If the shapes mismatch and raise_on_error is True. + """ + dims = dict() + errors = [] + for tensor, shape in x: + if tensor.ndimension() != len(shape): + errors.append(f"Invalid number of dimensions: {tensor.shape} vs. {shape}") + continue + for dim, name in zip(tensor.shape, shape): + exp_dim = dims.get(name) + if exp_dim is not None and exp_dim != dim: + errors.append(f"{name}: {dim} vs. {exp_dim}") + if raise_or_error and len(errors) > 0: + raise ValueError("Shape verification failed:\n" + '\n'.join(errors)) + return len(errors) > 0 From 356af149fadec99e5d412f33b37a04635fe0d2e0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 12:52:14 +0100 Subject: [PATCH 080/690] Simplify code --- src/pykeen/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index ef1ee1f3e7..572337af51 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1272,7 +1272,7 @@ def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: if len(set(dims)) != len(dims): raise ValueError(f"Duplicate dimensions: {dims}") assert all(0 <= d < len(x.shape) for d in dims) - for dim in reversed(sorted(dims)): + for dim in sorted(dims, reverse=True): x = x.squeeze(dim=dim) return x From 71ebea547a2c6b2cb64c3a8594f0f97c4c6b59dd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 13:10:02 +0100 Subject: [PATCH 081/690] extract ConvKB functional form --- src/pykeen/models/unimodal/conv_kb.py | 58 +++++++-------------------- src/pykeen/nn/functional.py | 48 ++++++++++++++++++++++ 2 files changed, 62 insertions(+), 44 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 977cf2d086..916245048e 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -5,12 +5,12 @@ import logging from typing import Optional -import torch import torch.autograd from torch import nn from ..base import EntityRelationEmbeddingModel, InteractionFunction from ...losses import Loss +from ...nn import functional as pykeen_functional from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -36,13 +36,8 @@ def __init__( self.num_filters = num_filters # The interaction model - # self.conv = nn.Conv2d(in_channels=1, out_channels=num_filters, kernel_size=(1, 3), bias=True) - # decompose convolution for faster computation in 1-n case - self.conv_head = nn.Parameter(torch.empty(num_filters)) - self.conv_rel = nn.Parameter(torch.empty(num_filters)) - self.conv_tail = nn.Parameter(torch.empty(num_filters)) - self.conv_bias = nn.Parameter(torch.empty(num_filters)) - self.relu = nn.ReLU() + self.conv = nn.Conv2d(in_channels=1, out_channels=num_filters, kernel_size=(1, 3), bias=True) + self.activation = nn.ReLU() self.hidden_dropout = nn.Dropout(p=hidden_dropout_rate) self.linear = nn.Linear(embedding_dim * num_filters, 1, bias=True) @@ -53,12 +48,9 @@ def reset_parameters(self): # noqa: D102 # Initialize all filters to [0.1, 0.1, -0.1], # c.f. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L34-L36 - # nn.init.constant_(self.conv.weight[..., :2], 0.1) - # nn.init.constant_(self.conv.weight[..., 2], -0.1) - nn.init.constant_(self.conv_head, 0.1) - nn.init.constant_(self.conv_rel, 0.1) - nn.init.constant_(self.conv_tail, 0.1) - nn.init.zeros_(self.conv_bias) + nn.init.constant_(self.conv.weight[..., :2], 0.1) + nn.init.constant_(self.conv.weight[..., 2], -0.1) + nn.init.zeros_(self.conv.bias) def forward( self, @@ -67,38 +59,16 @@ def forward( t: torch.FloatTensor, **kwargs, ) -> torch.FloatTensor: # noqa: D102 - # bind sizes - batch_size = max(x.shape[0] for x in (h, r, t)) - num_heads = h.shape[1] - num_relations = r.shape[1] - num_tails = t.shape[1] - - # compute conv(stack(h, r, t)) - # h.shape: (b, nh, d), conv_head.shape: (o), out.shape: (b, nh, d, o) - x = ( - self.conv_bias.view(1, 1, 1, 1, 1, self.num_filters) - + ( - h.view(h.shape[0], h.shape[1], 1, 1, self.embedding_dim, 1) - * self.conv_head.view(1, 1, 1, 1, 1, self.num_filters) - ) + ( - r.view(r.shape[0], 1, r.shape[1], 1, self.embedding_dim, 1) - * self.conv_rel.view(1, 1, 1, 1, 1, self.num_filters) - ) + ( - t.view(t.shape[0], 1, 1, t.shape[1], self.embedding_dim, 1) - * self.conv_tail.view(1, 1, 1, 1, 1, self.num_filters) - ) + return pykeen_functional.convkb_interaction( + h=h, + r=r, + t=t, + conv=self.conv, + activation=self.activation, + hidden_dropout=self.hidden_dropout, + linear=self.linear, ) - x = self.relu(x) - - # Apply dropout, cf. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L54-L56 - x = self.hidden_dropout(x) - - # Linear layer for final scores - return self.linear( - x.view(-1, self.embedding_dim * self.num_filters), - ).view(batch_size, num_heads, num_relations, num_tails) - class ConvKB(EntityRelationEmbeddingModel): r"""An implementation of ConvKB from [nguyen2018]_. diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 0b31a6c524..19829f17fe 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -144,3 +144,51 @@ def complex_interaction( (h_im, r_im, t_re), ] ) + +def convkb_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + conv: nn.Conv2d, + activation: nn.Module, + hidden_dropout: nn.Dropout, + linear: nn.Linear, +) -> torch.FloatTensor: + # bind sizes + batch_size = max(x.shape[0] for x in (h, r, t)) + num_heads = h.shape[1] + num_relations = r.shape[1] + num_tails = t.shape[1] + + # decompose convolution for faster computation in 1-n case + num_filters = conv.weight.shape[0] + assert conv.weight.shape == (num_filters, 1, 1, 3) + embedding_dim = h.shape[-1] + + # compute conv(stack(h, r, t)) + conv_head, conv_rel, conv_tail = conv.weight[:, 0, 0, :].t() + conv_bias = conv.bias + # h.shape: (b, nh, d), conv_head.shape: (o), out.shape: (b, nh, d, o) + x = ( + conv_bias.view(1, 1, 1, 1, 1, num_filters) + + ( + h.view(h.shape[0], h.shape[1], 1, 1, embedding_dim, 1) + * conv_head.view(1, 1, 1, 1, 1, num_filters) + ) + ( + r.view(r.shape[0], 1, r.shape[1], 1, embedding_dim, 1) + * conv_rel.view(1, 1, 1, 1, 1, num_filters) + ) + ( + t.view(t.shape[0], 1, 1, t.shape[1], embedding_dim, 1) + * conv_tail.view(1, 1, 1, 1, 1, num_filters) + ) + ) + + x = activation(x) + + # Apply dropout, cf. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L54-L56 + x = hidden_dropout(x) + + # Linear layer for final scores + return linear( + x.view(-1, embedding_dim * num_filters), + ).view(batch_size, num_heads, num_relations, num_tails) \ No newline at end of file From dc9a99b94ec957ad179c5a5fa7838732db1bc9e0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 13:10:14 +0100 Subject: [PATCH 082/690] Reformat --- src/pykeen/nn/functional.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 19829f17fe..d214144031 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -145,6 +145,7 @@ def complex_interaction( ] ) + def convkb_interaction( h: torch.FloatTensor, r: torch.FloatTensor, @@ -191,4 +192,4 @@ def convkb_interaction( # Linear layer for final scores return linear( x.view(-1, embedding_dim * num_filters), - ).view(batch_size, num_heads, num_relations, num_tails) \ No newline at end of file + ).view(batch_size, num_heads, num_relations, num_tails) From cbaa40d93ece5ad9bdbe0750108576e42e97aed1 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 13:10:46 +0100 Subject: [PATCH 083/690] Export convkb_interaction --- src/pykeen/nn/functional.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index d214144031..60e4c5625d 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -9,6 +9,7 @@ __all__ = [ "complex_interaction", "conve_interaction", + "convkb_interaction", "distmult_interaction", ] From 54ce36abfaead5bd45805d7cb34f5ef91f8d9be5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 13:13:57 +0100 Subject: [PATCH 084/690] Add docstring to ComplEx interaction --- src/pykeen/nn/functional.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 60e4c5625d..aebcbadada 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -134,6 +134,19 @@ def complex_interaction( r: torch.FloatTensor, t: torch.FloatTensor, ) -> torch.FloatTensor: + """ + Evaluate the ComplEx interaction function. + + :param h: shape: (batch_size, num_heads, 2*dim) + The complex head representations. + :param r: shape: (batch_size, num_relations, 2*dim) + The complex relation representations. + :param t: shape: (batch_size, num_tails, 2*dim) + The complex tail representations. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ h, h_term, r, r_term, t, t_term = _normalize_terms_for_einsum(h, r, t) (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] return sum( From 91f6b327802e34a0a507f18d89d8797a3c4f07d2 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 13:15:47 +0100 Subject: [PATCH 085/690] Add docstring to ConvKB --- src/pykeen/nn/functional.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index aebcbadada..64fdae4b75 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -169,6 +169,27 @@ def convkb_interaction( hidden_dropout: nn.Dropout, linear: nn.Linear, ) -> torch.FloatTensor: + """ + Evaluate the ConvKB interaction function. + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r: shape: (batch_size, num_relations, dim) + The relation representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + :param conv: + The 3x1 convolution. + :param activation: + The activation function. + :param hidden_dropout: + The dropout layer applied to the hidden activations. + :param linear: + The final linear layer. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ # bind sizes batch_size = max(x.shape[0] for x in (h, r, t)) num_heads = h.shape[1] From 35fae474d1ee2888b168016e909a519478d79ad7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 13:17:52 +0100 Subject: [PATCH 086/690] Add formula to ConvKB interaction function --- src/pykeen/nn/functional.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 64fdae4b75..203ffadf59 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -169,9 +169,12 @@ def convkb_interaction( hidden_dropout: nn.Dropout, linear: nn.Linear, ) -> torch.FloatTensor: - """ + r""" Evaluate the ConvKB interaction function. + .. math:: + W_L drop(act(W_C \ast ([h; r; t]) + b_C)) + b_L + :param h: shape: (batch_size, num_heads, dim) The head representations. :param r: shape: (batch_size, num_relations, dim) From cd08726fdd0458d032c7c257ef7cee69a6da99bb Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 15:12:24 +0100 Subject: [PATCH 087/690] Extract ER-MLP functional form --- src/pykeen/models/unimodal/ermlp.py | 39 ++++++------------ src/pykeen/nn/functional.py | 63 +++++++++++++++++++++-------- 2 files changed, 60 insertions(+), 42 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index fcdca0837d..213154566d 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -2,7 +2,6 @@ """Implementation of ERMLP.""" -import math from typing import Optional import torch @@ -10,6 +9,7 @@ from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss +from ...nn import functional as pykeen_functional from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -46,9 +46,7 @@ def __init__( with self.embedding_dim neurons and output layer with one neuron. The input is represented by the concatenation embeddings of the heads, relations and tail embeddings. """ - self.head_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) - self.rel_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=True) - self.tail_to_hidden = nn.Linear(in_features=embedding_dim, out_features=hidden_dim, bias=False) + self.hidden = nn.Linear(in_features=3 * embedding_dim, out_features=hidden_dim, bias=True) self.activation = nn.ReLU() self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) @@ -60,33 +58,22 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa: D102 self._check_for_empty_kwargs(kwargs) - h = self.head_to_hidden(h) - r = self.rel_to_hidden(r) - t = self.tail_to_hidden(t) - # TODO: Choosing which to combine first, h/r, h/t or r/t, depending on the shape might further improve - # performance in a 1:n scenario. - x = self.activation(h[:, :, None, None, :] + r[:, None, :, None, :] + t[:, None, None, :, :]) - return self.hidden_to_score(x).squeeze(dim=-1) + return pykeen_functional.ermlp_interaction( + h=h, + r=r, + t=t, + hidden=self.hidden, + activation=self.activation, + final=self.hidden_to_score, + ) def reset_parameters(self): # noqa: D102 # Initialize biases with zero - nn.init.zeros_(self.rel_to_hidden.bias) + nn.init.zeros_(self.hidden.bias) nn.init.zeros_(self.hidden_to_score.bias) # In the original formulation, - # W_2 sigma(W_1 cat([h, r, t]) + b_1) + b_2 - # W_1 would be initialized with nn.init.xavier_uniform, i.e. with a samples from uniform(-a, a) with - # a = math.sqrt(3.0) * gain * math.sqrt(2.0 / float(fan_in + fan_out)) - # we have: - # fan_out = hidden_dim - # fan_in = 3 * embedding_dim - bound = math.sqrt(3.0) * 1 * math.sqrt(2.0 / float(sum(self.head_to_hidden.weight.shape))) - for mod in [ - self.head_to_hidden, - self.rel_to_hidden, - self.tail_to_hidden, - ]: - nn.init.uniform_(mod.weight, -bound, bound) - nn.init.xavier_uniform_(self.hidden_to_score.weight, gain=nn.init.calculate_gain('relu')) + nn.init.xavier_uniform_(self.hidden.weight) + nn.init.xavier_uniform_(self.hidden_to_score.weight, gain=nn.init.calculate_gain(self.activation.__class__.__name__.lower())) class ERMLP(SimpleVectorEntityRelationEmbeddingModel): diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 203ffadf59..1b6b51c8ac 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -206,23 +206,12 @@ def convkb_interaction( # compute conv(stack(h, r, t)) conv_head, conv_rel, conv_tail = conv.weight[:, 0, 0, :].t() - conv_bias = conv.bias + conv_bias = conv.bias.view(1, 1, 1, 1, 1, num_filters) # h.shape: (b, nh, d), conv_head.shape: (o), out.shape: (b, nh, d, o) - x = ( - conv_bias.view(1, 1, 1, 1, 1, num_filters) - + ( - h.view(h.shape[0], h.shape[1], 1, 1, embedding_dim, 1) - * conv_head.view(1, 1, 1, 1, 1, num_filters) - ) + ( - r.view(r.shape[0], 1, r.shape[1], 1, embedding_dim, 1) - * conv_rel.view(1, 1, 1, 1, 1, num_filters) - ) + ( - t.view(t.shape[0], 1, 1, t.shape[1], embedding_dim, 1) - * conv_tail.view(1, 1, 1, 1, 1, num_filters) - ) - ) - - x = activation(x) + h = (h.view(h.shape[0], h.shape[1], 1, 1, embedding_dim, 1) * conv_head.view(1, 1, 1, 1, 1, num_filters)) + r = (r.view(r.shape[0], 1, r.shape[1], 1, embedding_dim, 1) * conv_rel.view(1, 1, 1, 1, 1, num_filters)) + t = (t.view(t.shape[0], 1, 1, t.shape[1], embedding_dim, 1) * conv_tail.view(1, 1, 1, 1, 1, num_filters)) + x = activation(conv_bias + h + r + t) # Apply dropout, cf. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L54-L56 x = hidden_dropout(x) @@ -231,3 +220,45 @@ def convkb_interaction( return linear( x.view(-1, embedding_dim * num_filters), ).view(batch_size, num_heads, num_relations, num_tails) + + +def ermlp_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + hidden: nn.Linear, + activation: nn.Module, + final: nn.Linear, +) -> torch.FloatTensor: + r""" + Evaluate the ER-MLP interaction function. + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r: shape: (batch_size, num_relations, dim) + The relation representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + :param hidden: + The first linear layer. + :param activation: + The activation function of the hidden layer. + :param final: + The second linear layer. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + num_heads, num_relations, num_tails = [x.shape[1] for x in (h, r, t)] + hidden_dim, embedding_dim = hidden.weight.shape + assert embedding_dim % 3 == 0 + embedding_dim = embedding_dim // 3 + # split, shape: (embedding_dim, hidden_dim) + head_to_hidden, rel_to_hidden, tail_to_hidden = hidden.weight.t().split(embedding_dim) + bias = hidden.bias.view(1, 1, 1, 1, -1) + h = h.view(-1, num_heads, 1, 1, embedding_dim) @ head_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) + r = r.view(-1, 1, num_relations, 1, embedding_dim) @ rel_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) + t = t.view(-1, 1, 1, num_tails, embedding_dim) @ tail_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) + # TODO: Choosing which to combine first, h/r, h/t or r/t, depending on the shape might further improve + # performance in a 1:n scenario. + return final(activation(bias + h + r + t)).squeeze(dim=-1) From deb2851256ec08b3b7d3396ff381db82e0c5866a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 15:20:06 +0100 Subject: [PATCH 088/690] Add docstring --- src/pykeen/nn/functional.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 1b6b51c8ac..01dd735e78 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -125,6 +125,19 @@ def distmult_interaction( r: torch.FloatTensor, t: torch.FloatTensor, ) -> torch.FloatTensor: + """ + DistMult interaction function. + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r: shape: (batch_size, num_relations, dim) + The relation representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ h, h_term, r, r_term, t, t_term = _normalize_terms_for_einsum(h, r, t) return torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', h, r, t) From ca3abaddc1f4389f26610a7423ef117f74cd3ae8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 15:28:00 +0100 Subject: [PATCH 089/690] Add docstring to ConvE functional interaction function --- src/pykeen/models/unimodal/conv_e.py | 1 - src/pykeen/nn/functional.py | 45 ++++++++++++++++++++++++++-- 2 files changed, 43 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index fa2c13a6b3..2a1ac19edf 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -174,7 +174,6 @@ def forward( embedding_height=self.embedding_height, embedding_width=self.embedding_width, num_in_features=self.num_in_features, - embedding_dim=self.embedding_dim, bn0=self.bn0, bn1=self.bn1, bn2=self.bn2, diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 01dd735e78..2b290df79b 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -53,7 +53,6 @@ def conve_interaction( embedding_height: int, embedding_width: int, num_in_features: int, - embedding_dim: int, bn0: Optional[nn.BatchNorm1d], bn1: Optional[nn.BatchNorm1d], bn2: Optional[nn.BatchNorm1d], @@ -64,11 +63,53 @@ def conve_interaction( activation: nn.Module, fc: nn.Linear, ) -> torch.FloatTensor: + """ + Evaluate the ConvE interaction function. + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r: shape: (batch_size, num_relations, dim) + The relation representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + :param t_bias: shape: (batch_size, num_tails, dim) + The tail entity bias. + :param input_channels: + The number of input channels. + :param embedding_height: + The height of the reshaped embedding. + :param embedding_width: + The width of the reshaped embedding. + :param num_in_features: + The number of output features of the final layer (calculated with kernel and embedding dimensions). + :param bn0: + The first batch normalization layer. + :param bn1: + The second batch normalization layer. + :param bn2: + The third batch normalization layer. + :param inp_drop: + The input dropout layer. + :param feature_map_drop: + The feature map dropout layer. + :param hidden_drop: + The hidden dropout layer. + :param conv1: + The convolution layer. + :param activation: + The activation function. + :param fc: + The final fully connected layer. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ # bind sizes batch_size = max(x.shape[0] for x in (h, r, t)) num_heads = h.shape[1] num_relations = r.shape[1] num_tails = t.shape[1] + embedding_dim = h.shape[-1] # repeat if necessary h = h.unsqueeze(dim=2).repeat(1 if h.shape[0] == batch_size else batch_size, 1, num_relations, 1) @@ -126,7 +167,7 @@ def distmult_interaction( t: torch.FloatTensor, ) -> torch.FloatTensor: """ - DistMult interaction function. + Evaluate the DistMult interaction function. :param h: shape: (batch_size, num_heads, dim) The head representations. From a615ba9eae633ae8f096ef8770b2a09f5a15c439 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 15:29:14 +0100 Subject: [PATCH 090/690] Add TODO --- src/pykeen/nn/functional.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 2b290df79b..ddb1ed6617 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -179,6 +179,7 @@ def distmult_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ + # TODO: check if einsum is still very slow. h, h_term, r, r_term, t, t_term = _normalize_terms_for_einsum(h, r, t) return torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', h, r, t) @@ -203,6 +204,7 @@ def complex_interaction( """ h, h_term, r, r_term, t, t_term = _normalize_terms_for_einsum(h, r, t) (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] + # TODO: check if einsum is still very slow. return sum( torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', hh, rr, tt) for hh, rr, tt in [ From ad2d1d2382a4d31d5f6ab05639b80ed2ea15428e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 15:31:41 +0100 Subject: [PATCH 091/690] Export ermlp_interaction --- src/pykeen/nn/functional.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index ddb1ed6617..40137ed209 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -11,6 +11,7 @@ "conve_interaction", "convkb_interaction", "distmult_interaction", + "ermlp_interaction", ] From 36324eaf9bb42834254f6d6d95b07a4d58d4d0ce Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 15:52:32 +0100 Subject: [PATCH 092/690] fix check_shape usage --- src/pykeen/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 572337af51..225b320ba8 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1404,7 +1404,7 @@ def score_t( The scores. """ # check shape - assert check_shapes(all_entities, "ne", r, "br", h, "be") + assert check_shapes((all_entities, "ne"), (r, "br"), (h, "be")) # prepare input to generic score function h, r = self._add_dim(h, r, dim=self.NUM_DIM) From 24d569406b47542c643071abbbd1389f435055c5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 15:53:07 +0100 Subject: [PATCH 093/690] Fix check_shape --- src/pykeen/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 5cbc6a9876..cc6b8c80f6 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -490,4 +490,4 @@ def check_shapes( errors.append(f"{name}: {dim} vs. {exp_dim}") if raise_or_error and len(errors) > 0: raise ValueError("Shape verification failed:\n" + '\n'.join(errors)) - return len(errors) > 0 + return len(errors) == 0 From af59237b1ed550ca54f048e5f4ff03a3664fd9c5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 15:53:21 +0100 Subject: [PATCH 094/690] Extract broadcast cat; improve ConvE efficiency for score_t --- src/pykeen/nn/functional.py | 93 +++++++++++++++++++++++++++++++++---- 1 file changed, 83 insertions(+), 10 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 40137ed209..dd3d6c665a 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -112,16 +112,13 @@ def conve_interaction( num_tails = t.shape[1] embedding_dim = h.shape[-1] - # repeat if necessary - h = h.unsqueeze(dim=2).repeat(1 if h.shape[0] == batch_size else batch_size, 1, num_relations, 1) - r = r.unsqueeze(dim=1).repeat(1 if r.shape[0] == batch_size else batch_size, num_heads, 1, 1) - - # resize and concat head and relation, batch_size', num_input_channels, 2*height, width + # repeat if necessary, and concat head and relation, batch_size', num_input_channels, 2*height, width # with batch_size' = batch_size * num_heads * num_relations - x = torch.cat([ - h.view(-1, input_channels, embedding_height, embedding_width), - r.view(-1, input_channels, embedding_height, embedding_width), - ], dim=2) + h = h.unsqueeze(dim=2) + h = h.view(*h.shape[:-1], input_channels, embedding_height, embedding_width) + r = r.unsqueeze(dim=1) + r = r.view(*r.shape[:-1], input_channels, embedding_height, embedding_width) + x = broadcast_cat(h, r, dim=2).view(-1, input_channels, 2 * embedding_height, embedding_width) # batch_size, num_input_channels, 2*height, width if bn0 is not None: @@ -149,7 +146,7 @@ def conve_interaction( x = activation(x) # reshape: (batch_size', embedding_dim) - x = x.view(batch_size, num_heads, num_relations, 1, embedding_dim) + x = x.view(-1, num_heads, num_relations, 1, embedding_dim) # For efficient calculation, each of the convolved [h, r] rows has only to be multiplied with one t row # output_shape: (batch_size, num_heads, num_relations, num_tails) @@ -162,6 +159,40 @@ def conve_interaction( return x +def broadcast_cat( + x: torch.FloatTensor, + y: torch.FloatTensor, + dim: int, +) -> torch.FloatTensor: + """ + Concatenate with broadcasting. + + :param x: + The first tensor. + :param y: + The second tensor. + :param dim: + The concat dimension. + + :return: + """ + if x.ndimension() != y.ndimension(): + raise ValueError + x_rep, y_rep = [], [] + for d, (xd, yd) in enumerate(zip(x.shape, y.shape)): + xr = yr = 1 + if d != dim and xd != yd: + if xd == 1: + xr = yd + elif yd == 1: + yr = xd + else: + raise ValueError + x_rep.append(xr) + y_rep.append(yr) + return torch.cat([x.repeat(*x_rep), y.repeat(*y_rep)], dim=dim) + + def distmult_interaction( h: torch.FloatTensor, r: torch.FloatTensor, @@ -319,3 +350,45 @@ def ermlp_interaction( # TODO: Choosing which to combine first, h/r, h/t or r/t, depending on the shape might further improve # performance in a 1:n scenario. return final(activation(bias + h + r + t)).squeeze(dim=-1) + + +def ermlp_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + hidden: nn.Linear, + activation: nn.Module, + final: nn.Linear, +) -> torch.FloatTensor: + r""" + Evaluate the ER-MLPE interaction function. + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r: shape: (batch_size, num_relations, dim) + The relation representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + :param hidden: + The first linear layer. + :param activation: + The activation function of the hidden layer. + :param final: + The second linear layer. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + # Concatenate them + x_s = torch.cat([h, r], dim=-1) + x_s = self.input_dropout(x_s) + + # Predict t embedding + x_t = self.mlp(x_s) + + # compare with all t's + # For efficient calculation, each of the calculated [h, r] rows has only to be multiplied with one t row + x = (x_t.view(-1, self.embedding_dim) * t).sum(dim=1, keepdim=True) + # The application of the sigmoid during training is automatically handled by the default loss. + + return x From a3b18baa43b1fb69e2b0d84036ee8811ba8bf2d3 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 16:01:09 +0100 Subject: [PATCH 095/690] Move broadcast cat --- src/pykeen/utils.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index cc6b8c80f6..84efae9448 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -491,3 +491,37 @@ def check_shapes( if raise_or_error and len(errors) > 0: raise ValueError("Shape verification failed:\n" + '\n'.join(errors)) return len(errors) == 0 + + +def broadcast_cat( + x: torch.FloatTensor, + y: torch.FloatTensor, + dim: int, +) -> torch.FloatTensor: + """ + Concatenate with broadcasting. + + :param x: + The first tensor. + :param y: + The second tensor. + :param dim: + The concat dimension. + + :return: + """ + if x.ndimension() != y.ndimension(): + raise ValueError + x_rep, y_rep = [], [] + for d, (xd, yd) in enumerate(zip(x.shape, y.shape)): + xr = yr = 1 + if d != dim and xd != yd: + if xd == 1: + xr = yd + elif yd == 1: + yr = xd + else: + raise ValueError + x_rep.append(xr) + y_rep.append(yr) + return torch.cat([x.repeat(*x_rep), y.repeat(*y_rep)], dim=dim) \ No newline at end of file From b048ce668f8005b987f583da12f9182df36c1951 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 16:01:18 +0100 Subject: [PATCH 096/690] Move ER-MLP-E function to functional --- src/pykeen/models/unimodal/ermlpe.py | 83 +++++++++------------------- src/pykeen/nn/functional.py | 71 +++++------------------- 2 files changed, 42 insertions(+), 112 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index 05976b4bd6..78c59e4445 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -9,6 +9,7 @@ from ..base import EntityRelationEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss +from ...nn import functional as pykeen_functional from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -96,13 +97,11 @@ def __init__( def _reset_parameters_(self): # noqa: D102 super()._reset_parameters_() - for module in [ - self.linear1, - self.linear2, - self.bn1, - self.bn2, - ]: - module.reset_parameters() + for module in self.modules(): + if module is self: + continue + if hasattr(module, "reset_parameters"): + module.reset_parameters() def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings @@ -113,19 +112,13 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: # Embedding Regularization self.regularize_if_necessary(h, r, t) - # Concatenate them - x_s = torch.cat([h, r], dim=-1) - x_s = self.input_dropout(x_s) - - # Predict t embedding - x_t = self.mlp(x_s) - - # compare with all t's - # For efficient calculation, each of the calculated [h, r] rows has only to be multiplied with one t row - x = (x_t.view(-1, self.embedding_dim) * t).sum(dim=1, keepdim=True) - # The application of the sigmoid during training is automatically handled by the default loss. - - return x + return pykeen_functional.ermlpe_interaction( + h=h, + r=r, + t=t, + input_dropout=self.input_dropout, + mlp=self.mlp, + ).view(hrt_batch.shape[0], 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 h = self.entity_embeddings(indices=hr_batch[:, 0]).view(-1, self.embedding_dim) @@ -135,17 +128,13 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 # Embedding Regularization self.regularize_if_necessary(h, r, t) - # Concatenate them - x_s = torch.cat([h, r], dim=-1) - x_s = self.input_dropout(x_s) - - # Predict t embedding - x_t = self.mlp(x_s) - - x = x_t @ t - # The application of the sigmoid during training is automatically handled by the default loss. - - return x + return pykeen_functional.ermlpe_interaction( + h=h, + r=r, + t=t, + input_dropout=self.input_dropout, + mlp=self.mlp, + ).view(hr_batch.shape[0], self.num_entities) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 h = self.entity_embeddings(indices=None) @@ -155,28 +144,10 @@ def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 # Embedding Regularization self.regularize_if_necessary(h, r, t) - rt_batch_size = t.shape[0] - - # Extend each rt_batch of "r" with shape [rt_batch_size, dim] to [rt_batch_size, dim * num_entities] - r = torch.repeat_interleave(r, self.num_entities, dim=0) - # Extend each h with shape [num_entities, dim] to [rt_batch_size * num_entities, dim] - # h = torch.repeat_interleave(h, rt_batch_size, dim=0) - h = h.repeat(rt_batch_size, 1) - - # Extend t - t = t.repeat_interleave(self.num_entities, dim=0) - - # Concatenate them - x_s = torch.cat([h, r], dim=-1) - x_s = self.input_dropout(x_s) - - # Predict t embedding - x_t = self.mlp(x_s) - - # For efficient calculation, each of the calculated [h, r] rows has only to be multiplied with one t row - x = (x_t.view(-1, self.embedding_dim) * t).sum(dim=1, keepdim=True) - # The results have to be realigned with the expected output of the score_h function - x = x.view(rt_batch_size, self.num_entities) - # The application of the sigmoid during training is automatically handled by the default loss. - - return x + return pykeen_functional.ermlpe_interaction( + h=h, + r=r, + t=t, + input_dropout=self.input_dropout, + mlp=self.mlp, + ).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index dd3d6c665a..02bc552c4f 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -4,7 +4,7 @@ import torch from torch import nn -from ..utils import is_cudnn_error, normalize_for_einsum, split_complex +from ..utils import broadcast_cat, is_cudnn_error, normalize_for_einsum, split_complex __all__ = [ "complex_interaction", @@ -12,6 +12,7 @@ "convkb_interaction", "distmult_interaction", "ermlp_interaction", + "ermlpe_interaction", ] @@ -159,40 +160,6 @@ def conve_interaction( return x -def broadcast_cat( - x: torch.FloatTensor, - y: torch.FloatTensor, - dim: int, -) -> torch.FloatTensor: - """ - Concatenate with broadcasting. - - :param x: - The first tensor. - :param y: - The second tensor. - :param dim: - The concat dimension. - - :return: - """ - if x.ndimension() != y.ndimension(): - raise ValueError - x_rep, y_rep = [], [] - for d, (xd, yd) in enumerate(zip(x.shape, y.shape)): - xr = yr = 1 - if d != dim and xd != yd: - if xd == 1: - xr = yd - elif yd == 1: - yr = xd - else: - raise ValueError - x_rep.append(xr) - y_rep.append(yr) - return torch.cat([x.repeat(*x_rep), y.repeat(*y_rep)], dim=dim) - - def distmult_interaction( h: torch.FloatTensor, r: torch.FloatTensor, @@ -352,13 +319,12 @@ def ermlp_interaction( return final(activation(bias + h + r + t)).squeeze(dim=-1) -def ermlp_interaction( +def ermlpe_interaction( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, - hidden: nn.Linear, - activation: nn.Module, - final: nn.Linear, + input_dropout: nn.Dropout, + mlp: nn.Module, ) -> torch.FloatTensor: r""" Evaluate the ER-MLPE interaction function. @@ -369,26 +335,19 @@ def ermlp_interaction( The relation representations. :param t: shape: (batch_size, num_tails, dim) The tail representations. - :param hidden: - The first linear layer. - :param activation: - The activation function of the hidden layer. - :param final: - The second linear layer. + :param input_dropout: + The input dropout layer. + :param mlp: + The MLP. :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # Concatenate them - x_s = torch.cat([h, r], dim=-1) - x_s = self.input_dropout(x_s) + # repeat if necessary, and concat head and relation, (batch_size, num_heads, num_relations, 2 * embedding_dim) + x = broadcast_cat(h.unsqueeze(dim=2), r.unsqueeze(dim=1), dim=-1) + x = input_dropout(x) - # Predict t embedding - x_t = self.mlp(x_s) + # Predict t embedding, shape: (batch_size, num_heads, num_relations, embedding_dim) + x = mlp(x) - # compare with all t's - # For efficient calculation, each of the calculated [h, r] rows has only to be multiplied with one t row - x = (x_t.view(-1, self.embedding_dim) * t).sum(dim=1, keepdim=True) - # The application of the sigmoid during training is automatically handled by the default loss. - - return x + return (x.unsqueeze(dim=-2) @ t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]).transpose(-2, -1)).squeeze(dim=-1) From fd71d2a3e5c300919b8d2b62cae8718a1c91c1aa Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 12 Nov 2020 16:10:51 +0100 Subject: [PATCH 097/690] Add ER-MLP-E interaction function (non-functional form) --- src/pykeen/models/unimodal/ermlpe.py | 112 +++++++++++++++------------ src/pykeen/nn/functional.py | 4 - 2 files changed, 61 insertions(+), 55 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index 78c59e4445..ed45b551d4 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -7,7 +7,7 @@ import torch from torch import nn -from ..base import EntityRelationEmbeddingModel +from ..base import EntityRelationEmbeddingModel, InteractionFunction from ...losses import BCEAfterSigmoidLoss, Loss from ...nn import functional as pykeen_functional from ...regularizers import Regularizer @@ -19,6 +19,40 @@ ] +class ERMLPEInteractionFunction(InteractionFunction): + """Interaction function of ER-MLP.""" + + def __init__( + self, + hidden_dim: int = 300, + input_dropout: float = 0.2, + hidden_dropout: float = 0.3, + embedding_dim: int = 200, + ): + super().__init__() + self.mlp = nn.Sequential( + nn.Dropout(input_dropout), + nn.Linear(2 * embedding_dim, hidden_dim), + nn.Dropout(hidden_dropout), + nn.BatchNorm1d(hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, embedding_dim), + nn.Dropout(hidden_dropout), + nn.BatchNorm1d(embedding_dim), + nn.ReLU(), + ) + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + self._check_for_empty_kwargs(kwargs=kwargs) + return pykeen_functional.ermlpe_interaction(h=h, r=r, t=t, mlp=self.mlp) + + class ERMLPE(EntityRelationEmbeddingModel): r"""An extension of ERMLP proposed by [sharifzadeh2019]_. @@ -76,78 +110,54 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) - self.hidden_dim = hidden_dim - self.input_dropout = input_dropout - - self.linear1 = nn.Linear(2 * self.embedding_dim, self.hidden_dim) - self.linear2 = nn.Linear(self.hidden_dim, self.embedding_dim) - self.input_dropout = nn.Dropout(self.input_dropout) - self.bn1 = nn.BatchNorm1d(self.hidden_dim) - self.bn2 = nn.BatchNorm1d(self.embedding_dim) - self.mlp = nn.Sequential( - self.linear1, - nn.Dropout(hidden_dropout), - self.bn1, - nn.ReLU(), - self.linear2, - nn.Dropout(hidden_dropout), - self.bn2, - nn.ReLU(), + self.interaction_function = ERMLPEInteractionFunction( + hidden_dim=hidden_dim, + input_dropout=input_dropout, + hidden_dropout=hidden_dropout, + embedding_dim=embedding_dim, ) def _reset_parameters_(self): # noqa: D102 super()._reset_parameters_() - for module in self.modules(): - if module is self: - continue - if hasattr(module, "reset_parameters"): - module.reset_parameters() + self.interaction_function.reset_parameters() def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings - h = self.entity_embeddings(indices=hrt_batch[:, 0]).view(-1, self.embedding_dim) - r = self.relation_embeddings(indices=hrt_batch[:, 1]).view(-1, self.embedding_dim) + h = self.entity_embeddings(indices=hrt_batch[:, 0]) + r = self.relation_embeddings(indices=hrt_batch[:, 1]) t = self.entity_embeddings(indices=hrt_batch[:, 2]) # Embedding Regularization self.regularize_if_necessary(h, r, t) - return pykeen_functional.ermlpe_interaction( - h=h, - r=r, - t=t, - input_dropout=self.input_dropout, - mlp=self.mlp, - ).view(hrt_batch.shape[0], 1) + return self.interaction_function.score_hrt(h=h, r=r, t=t) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=hr_batch[:, 0]).view(-1, self.embedding_dim) - r = self.relation_embeddings(indices=hr_batch[:, 1]).view(-1, self.embedding_dim) - t = self.entity_embeddings(indices=None).transpose(1, 0) + h = self.entity_embeddings(indices=hr_batch[:, 0]) + r = self.relation_embeddings(indices=hr_batch[:, 1]) + t = self.entity_embeddings(indices=None) + + # Embedding Regularization + self.regularize_if_necessary(h, r, t) + + return self.interaction_function.score_t(h=h, r=r, all_entities=t) + + def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + h = self.entity_embeddings(indices=ht_batch[:, 0]) + r = self.relation_embeddings(indices=None) + t = self.entity_embeddings(indices=ht_batch[:, 1]) # Embedding Regularization self.regularize_if_necessary(h, r, t) - return pykeen_functional.ermlpe_interaction( - h=h, - r=r, - t=t, - input_dropout=self.input_dropout, - mlp=self.mlp, - ).view(hr_batch.shape[0], self.num_entities) + return self.interaction_function.score_r(h=h, all_relations=r, t=t) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 h = self.entity_embeddings(indices=None) - r = self.relation_embeddings(indices=rt_batch[:, 0]).view(-1, self.embedding_dim) - t = self.entity_embeddings(indices=rt_batch[:, 1]).view(-1, self.embedding_dim) + r = self.relation_embeddings(indices=rt_batch[:, 0]) + t = self.entity_embeddings(indices=rt_batch[:, 1]) # Embedding Regularization self.regularize_if_necessary(h, r, t) - return pykeen_functional.ermlpe_interaction( - h=h, - r=r, - t=t, - input_dropout=self.input_dropout, - mlp=self.mlp, - ).view(rt_batch.shape[0], self.num_entities) + return self.interaction_function.score_h(all_entities=h, r=r, t=t) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 02bc552c4f..5adf3b7778 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -323,7 +323,6 @@ def ermlpe_interaction( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, - input_dropout: nn.Dropout, mlp: nn.Module, ) -> torch.FloatTensor: r""" @@ -335,8 +334,6 @@ def ermlpe_interaction( The relation representations. :param t: shape: (batch_size, num_tails, dim) The tail representations. - :param input_dropout: - The input dropout layer. :param mlp: The MLP. @@ -345,7 +342,6 @@ def ermlpe_interaction( """ # repeat if necessary, and concat head and relation, (batch_size, num_heads, num_relations, 2 * embedding_dim) x = broadcast_cat(h.unsqueeze(dim=2), r.unsqueeze(dim=1), dim=-1) - x = input_dropout(x) # Predict t embedding, shape: (batch_size, num_heads, num_relations, embedding_dim) x = mlp(x) From 9d13db859861b3d22c180cc417d06b364ec0b340 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 13:35:26 +0100 Subject: [PATCH 098/690] Update utils.py --- src/pykeen/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 84efae9448..8356640605 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -524,4 +524,4 @@ def broadcast_cat( raise ValueError x_rep.append(xr) y_rep.append(yr) - return torch.cat([x.repeat(*x_rep), y.repeat(*y_rep)], dim=dim) \ No newline at end of file + return torch.cat([x.repeat(*x_rep), y.repeat(*y_rep)], dim=dim) From e28a8f7bfc61bba82752749297b2cfeab5199793 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 13:35:34 +0100 Subject: [PATCH 099/690] Make HolE functional --- src/pykeen/models/unimodal/hole.py | 20 ++------------------ src/pykeen/nn/functional.py | 27 +++++++++++++++++++++++++++ 2 files changed, 29 insertions(+), 18 deletions(-) diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 29103fda72..c636396a1a 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -4,11 +4,11 @@ from typing import Optional -import torch import torch.autograd from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss +from ...nn.functional import hole_interaction from ...nn.init import xavier_uniform_ from ...regularizers import Regularizer from ...triples import TriplesFactory @@ -31,23 +31,7 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa: D102 self._check_for_empty_kwargs(kwargs) - # Circular correlation of entity embeddings - a_fft = torch.rfft(h, signal_ndim=1, onesided=True) - b_fft = torch.rfft(t, signal_ndim=1, onesided=True) - - # complex conjugate, a_fft.shape = (batch_size, num_entities, d', 2) - a_fft[:, :, :, 1] *= -1 - - # Hadamard product in frequency domain - p_fft = a_fft * b_fft - - # inverse real FFT, shape: (batch_size, num_entities, d) - composite = torch.irfft(p_fft, signal_ndim=1, onesided=True, signal_sizes=(h.shape[-1],)) - - # inner product with relation embedding - scores = torch.sum(r * composite, dim=-1, keepdim=False) - - return scores + return hole_interaction(h=h, r=r, t=t) class HolE(SimpleVectorEntityRelationEmbeddingModel): diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 5adf3b7778..3a99d17adf 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -1,4 +1,7 @@ +# -*- coding: utf-8 -*- + """Functional forms of interaction methods.""" + from typing import Optional, Tuple import torch @@ -13,6 +16,7 @@ "distmult_interaction", "ermlp_interaction", "ermlpe_interaction", + 'hole_interaction', ] @@ -347,3 +351,26 @@ def ermlpe_interaction( x = mlp(x) return (x.unsqueeze(dim=-2) @ t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]).transpose(-2, -1)).squeeze(dim=-1) + + +def hole_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: # noqa: D102 + """Evaluate the HolE interaction function.""" + # Circular correlation of entity embeddings + a_fft = torch.rfft(h, signal_ndim=1, onesided=True) + b_fft = torch.rfft(t, signal_ndim=1, onesided=True) + + # complex conjugate, a_fft.shape = (batch_size, num_entities, d', 2) + a_fft[:, :, :, 1] *= -1 + + # Hadamard product in frequency domain + p_fft = a_fft * b_fft + + # inverse real FFT, shape: (batch_size, num_entities, d) + composite = torch.irfft(p_fft, signal_ndim=1, onesided=True, signal_sizes=(h.shape[-1],)) + + # inner product with relation embedding + return torch.sum(r * composite, dim=-1, keepdim=False) From e4ea2c6b1ddaeb9ee0f12cec0620057a7180c01f Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 13:35:54 +0100 Subject: [PATCH 100/690] Update ermlp.py --- src/pykeen/models/unimodal/ermlp.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 213154566d..44166ac913 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -73,7 +73,10 @@ def reset_parameters(self): # noqa: D102 nn.init.zeros_(self.hidden_to_score.bias) # In the original formulation, nn.init.xavier_uniform_(self.hidden.weight) - nn.init.xavier_uniform_(self.hidden_to_score.weight, gain=nn.init.calculate_gain(self.activation.__class__.__name__.lower())) + nn.init.xavier_uniform_( + self.hidden_to_score.weight, + gain=nn.init.calculate_gain(self.activation.__class__.__name__.lower()), + ) class ERMLP(SimpleVectorEntityRelationEmbeddingModel): From 0b9fcdd8bada8932b7a49f78cdc89f510fab05a5 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 13:36:30 +0100 Subject: [PATCH 101/690] Remove unused batch size @mberr this is probably important somewhere, but not sure where --- src/pykeen/nn/functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 3a99d17adf..62972b5e01 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -111,7 +111,7 @@ def conve_interaction( The scores. """ # bind sizes - batch_size = max(x.shape[0] for x in (h, r, t)) + # batch_size = max(x.shape[0] for x in (h, r, t)) num_heads = h.shape[1] num_relations = r.shape[1] num_tails = t.shape[1] From 1efd69f797c6827e3c9ef6206c04eaf8b21dd245 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 13:43:35 +0100 Subject: [PATCH 102/690] Switch RotatE to functional interface --- src/pykeen/models/unimodal/rotate.py | 60 ++++++++-------------------- src/pykeen/nn/functional.py | 43 ++++++++++++++++++++ 2 files changed, 60 insertions(+), 43 deletions(-) diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index e1d3501447..3cc745d263 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -9,8 +9,9 @@ import torch.autograd from torch.nn import functional -from ..base import EntityRelationEmbeddingModel +from ..base import EntityRelationEmbeddingModel, InteractionFunction from ...losses import Loss +from ...nn import functional as pykeen_functional from ...nn.init import xavier_uniform_ from ...regularizers import Regularizer from ...triples import TriplesFactory @@ -51,6 +52,20 @@ def complex_normalize(x: torch.Tensor) -> torch.Tensor: return x +class RotatEInteraction(InteractionFunction): + """Interaction function of RotatE.""" + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + self._check_for_empty_kwargs(kwargs) + return pykeen_functional.rotate_interaction(h=h, r=r, t=t) + + class RotatE(EntityRelationEmbeddingModel): r"""An implementation of RotatE from [sun2019]_. @@ -91,6 +106,7 @@ def __init__( random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: + self.interaction_function = RotatEInteraction() super().__init__( triples_factory=triples_factory, embedding_dim=2 * embedding_dim, @@ -105,48 +121,6 @@ def __init__( ) self.real_embedding_dim = embedding_dim - @staticmethod - def interaction_function( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - ) -> torch.FloatTensor: - """Evaluate the interaction function of ComplEx for given embeddings. - - The embeddings have to be in a broadcastable shape. - - WARNING: No forward constraints are applied. - - :param h: shape: (..., e, 2) - Head embeddings. Last dimension corresponds to (real, imag). - :param r: shape: (..., e, 2) - Relation embeddings. Last dimension corresponds to (real, imag). - :param t: shape: (..., e, 2) - Tail embeddings. Last dimension corresponds to (real, imag). - - :return: shape: (...) - The scores. - """ - # Decompose into real and imaginary part - h_re = h[..., 0] - h_im = h[..., 1] - r_re = r[..., 0] - r_im = r[..., 1] - - # Rotate (=Hadamard product in complex space). - rot_h = torch.stack( - [ - h_re * r_re - h_im * r_im, - h_re * r_im + h_im * r_re, - ], - dim=-1, - ) - # Workaround until https://github.com/pytorch/pytorch/issues/30704 is fixed - diff = rot_h - t - scores = -torch.norm(diff.view(diff.shape[:-2] + (-1,)), dim=-1) - - return scores - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings h = self.entity_embeddings(indices=hrt_batch[:, 0]).view(-1, self.real_embedding_dim, 2) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 62972b5e01..c874e947dc 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -17,6 +17,7 @@ "ermlp_interaction", "ermlpe_interaction", 'hole_interaction', + 'rotate_interaction', ] @@ -374,3 +375,45 @@ def hole_interaction( # inner product with relation embedding return torch.sum(r * composite, dim=-1, keepdim=False) + + +def rotate_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Evaluate the interaction function of RotatE for given embeddings. + + The embeddings have to be in a broadcastable shape. + + WARNING: No forward constraints are applied. + + :param h: shape: (..., e, 2) + Head embeddings. Last dimension corresponds to (real, imag). + :param r: shape: (..., e, 2) + Relation embeddings. Last dimension corresponds to (real, imag). + :param t: shape: (..., e, 2) + Tail embeddings. Last dimension corresponds to (real, imag). + + :return: shape: (...) + The scores. + """ + # Decompose into real and imaginary part + h_re = h[..., 0] + h_im = h[..., 1] + r_re = r[..., 0] + r_im = r[..., 1] + + # Rotate (=Hadamard product in complex space). + rot_h = torch.stack( + [ + h_re * r_re - h_im * r_im, + h_re * r_im + h_im * r_re, + ], + dim=-1, + ) + # Workaround until https://github.com/pytorch/pytorch/issues/30704 is fixed + diff = rot_h - t + scores = -torch.norm(diff.view(diff.shape[:-2] + (-1,)), dim=-1) + + return scores From ede132533b8ac477dbee0d19327920c42dd5ebeb Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 13:57:10 +0100 Subject: [PATCH 103/690] Move stateful interaction functions into own module --- src/pykeen/models/base.py | 228 +--------- src/pykeen/models/unimodal/complex.py | 20 +- src/pykeen/models/unimodal/conv_e.py | 171 +------- src/pykeen/models/unimodal/conv_kb.py | 53 +-- src/pykeen/models/unimodal/distmult.py | 20 +- src/pykeen/models/unimodal/ermlp.py | 67 +-- src/pykeen/models/unimodal/ermlpe.py | 39 +- src/pykeen/models/unimodal/hole.py | 3 +- src/pykeen/models/unimodal/trans_d.py | 3 +- src/pykeen/nn/__init__.py | 2 +- src/pykeen/nn/modules.py | 568 +++++++++++++++++++++++++ 11 files changed, 589 insertions(+), 585 deletions(-) create mode 100644 src/pykeen/nn/modules.py diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 225b320ba8..f186d0f4ab 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -17,10 +17,11 @@ from ..losses import Loss, MarginRankingLoss, NSSALoss from ..nn import Embedding +from ..nn.modules import InteractionFunction from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory from ..typing import Constrainer, DeviceHint, Initializer, MappedTriples, Normalizer -from ..utils import NoRandomSeedNecessary, check_shapes, resolve_device, set_random_seed +from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed __all__ = [ 'Model', @@ -28,7 +29,6 @@ 'EntityRelationEmbeddingModel', 'GeneralVectorEntityRelationEmbeddingModel', 'SimpleVectorEntityRelationEmbeddingModel', - 'InteractionFunction', 'IndexFunction', 'MultimodalModel', ] @@ -1201,230 +1201,6 @@ class MultimodalModel(EntityRelationEmbeddingModel): """A multimodal KGE model.""" -class InteractionFunction(nn.Module): - """Base class for interaction functions.""" - - BATCH_DIM: int = 0 - NUM_DIM: int = 1 - HEAD_DIM: int = 1 - RELATION_DIM: int = 2 - TAIL_DIM: int = 3 - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: - """Score the given triples. - - :param h: shape: (batch_size, num_heads, d_e) - The head representations. - :param r: shape: (batch_size, num_relations, d_r) - The relation representations. - :param t: shape: (batch_size, num_tails, d_e) - The tail representations. - :param kwargs: - Additional key-word based arguments. - - :return: shape: (batch_size, num_heads, num_relations, num_tails) - The scores. - """ - raise NotImplementedError - - def _check_for_empty_kwargs(self, kwargs: Mapping[str, Any]) -> None: - """Check that kwargs is empty.""" - if len(kwargs) > 0: - raise ValueError(f"{self.__class__.__name__} does not take the following kwargs: {kwargs}") - - @staticmethod - def _add_dim(*x: torch.FloatTensor, dim: int) -> Sequence[torch.FloatTensor]: - """ - Add a dimension to tensors. - - :param x: shape: (d1, ..., dk) - The tensor. - - :return: shape: (1, d1, ..., dk) - The tensor with batch dimension. - """ - out = [xx.unsqueeze(dim=dim) for xx in x] - if len(x) > 1: - return out - return out[0] - - @staticmethod - def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: - """ - Remove dimensions from a tensor. - - :param x: - The tensor. - :param dims: - The dimensions to remove. - - :return: - The squeezed tensor. - """ - # normalize dimensions - dims = [d if d >= 0 else len(x.shape) + d for d in dims] - if len(set(dims)) != len(dims): - raise ValueError(f"Duplicate dimensions: {dims}") - assert all(0 <= d < len(x.shape) for d in dims) - for dim in sorted(dims, reverse=True): - x = x.squeeze(dim=dim) - return x - - def score_hrt( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: - """ - Score a batch of triples.. - - :param h: shape: (batch_size, d_e) - The head representations. - :param r: shape: (batch_size, d_r) - The relation representations. - :param t: shape: (batch_size, d_e) - The tail representations. - :param kwargs: - Additional key-word based arguments. - - :return: shape: (batch_size, 1) - The scores. - """ - # check shape - assert check_shapes((h, "be"), (r, "br"), (t, "be")) - - # prepare input to generic score function - h, r, t = self._add_dim(h, r, t, dim=self.NUM_DIM) - - # get scores - scores = self(h=h, r=r, t=t, **kwargs) - - # prepare output shape, (batch_size, num_heads, num_relations, num_tails) -> (batch_size, 1) - return self._remove_dim(scores, self.HEAD_DIM, self.RELATION_DIM, self.TAIL_DIM).unsqueeze(dim=-1) - - def score_h( - self, - all_entities: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: - """ - Score all head entities. - - :param all_entities: shape: (num_entities, d_e) - The head representations. - :param r: shape: (batch_size, d_r) - The relation representations. - :param t: shape: (batch_size, d_e) - The tail representations. - :param kwargs: - Additional key-word based arguments. - - :return: shape: (batch_size, num_entities) - The scores. - """ - # check shape - assert check_shapes((all_entities, "ne"), (r, "br"), (t, "be")) - - # TODO: What about unsqueezing for additional e.g. head arguments - # prepare input to generic score function - r, t = self._add_dim(r, t, dim=self.NUM_DIM) - h = self._add_dim(all_entities, dim=self.BATCH_DIM) - - # get scores - scores = self(h=h, r=r, t=t, **kwargs) - - # prepare output shape, (batch_size, num_heads, num_relations, num_tails) -> (batch_size, num_heads) - return self._remove_dim(scores, self.RELATION_DIM, self.TAIL_DIM) - - def score_r( - self, - h: torch.FloatTensor, - all_relations: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: - """ - Score all relations. - - :param h: shape: (batch_size, d_e) - The head representations. - :param all_relations: shape: (batch_size, d_r) - The relation representations. - :param t: shape: (num_entities, d_e) - The tail representations. - :param kwargs: - Additional key-word based arguments. - - :return: shape: (batch_size, num_entities) - The scores. - """ - # check shape - assert check_shapes((all_relations, "nr"), (h, "be"), (t, "be")) - - # prepare input to generic score function - h, t = self._add_dim(h, t, dim=self.NUM_DIM) - r = self._add_dim(all_relations, dim=self.BATCH_DIM) - - # get scores - scores = self(h=h, r=r, t=t, **kwargs) - - # prepare output shape - return self._remove_dim(scores, self.HEAD_DIM, self.TAIL_DIM) - - def score_t( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - all_entities: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: - """ - Score all tail entities. - - :param h: shape: (batch_size, d_e) - The head representations. - :param r: shape: (batch_size, d_r) - The relation representations. - :param all_entities: shape: (num_entities, d_e) - The tail representations. - :param kwargs: - Additional key-word based arguments. - - :return: shape: (batch_size, num_entities) - The scores. - """ - # check shape - assert check_shapes((all_entities, "ne"), (r, "br"), (h, "be")) - - # prepare input to generic score function - h, r = self._add_dim(h, r, dim=self.NUM_DIM) - t = self._add_dim(all_entities, dim=self.BATCH_DIM) - - # get scores - scores = self(h=h, r=r, t=t, **kwargs) - - # prepare output shape - return self._remove_dim(scores, self.HEAD_DIM, self.RELATION_DIM) - - def reset_parameters(self): - """Reset parameters the interaction function may have.""" - for mod in self.modules(): - if mod is self: - continue - if hasattr(mod, 'reset_parameters'): - mod.reset_parameters() - - class IndexFunction(nn.Module): """A function that handles looking up embeddings by index.""" diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 978914f75a..b0d377128a 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -4,36 +4,20 @@ from typing import Optional -import torch import torch.nn as nn -from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel +from ..base import SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss, SoftplusLoss -from ...nn import functional as pykeen_functional +from ...nn.modules import ComplExInteractionFunction from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint __all__ = [ 'ComplEx', - 'ComplExInteractionFunction', ] -class ComplExInteractionFunction(InteractionFunction): - """Interaction function of ComplEx.""" - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return pykeen_functional.complex_interaction(h=h, r=r, t=t) - - class ComplEx(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of ComplEx [trouillon2016]_. diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 2a1ac19edf..f4c26f4888 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -1,18 +1,17 @@ # -*- coding: utf-8 -*- """Implementation of ConvE.""" - import logging -import math -from typing import Optional, Tuple, Type +from typing import Optional, Type import torch from torch import nn -from ..base import EntityRelationEmbeddingModel, InteractionFunction +from ..base import EntityRelationEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss -from ...nn import Embedding, functional as pykeen_functional +from ...nn import Embedding from ...nn.init import xavier_normal_ +from ...nn.modules import ConvEInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -24,168 +23,6 @@ logger = logging.getLogger(__name__) -def _calculate_missing_shape_information( - embedding_dim: int, - input_channels: Optional[int] = None, - width: Optional[int] = None, - height: Optional[int] = None, -) -> Tuple[int, int, int]: - """ - Automatically calculates missing dimensions for ConvE. - - :param embedding_dim: - :param input_channels: - :param width: - :param height: - - :return: (input_channels, width, height), such that - `embedding_dim = input_channels * width * height` - :raises: - If no factorization could be found. - """ - # Store initial input for error message - original = (input_channels, width, height) - - # All are None - if all(factor is None for factor in [input_channels, width, height]): - input_channels = 1 - result_sqrt = math.floor(math.sqrt(embedding_dim)) - height = max(factor for factor in range(1, result_sqrt + 1) if embedding_dim % factor == 0) - width = embedding_dim // height - - # input_channels is None, and any of height or width is None -> set input_channels=1 - if input_channels is None and any(remaining is None for remaining in [width, height]): - input_channels = 1 - - # input channels is not None, and one of height or width is None - assert len([factor for factor in [input_channels, width, height] if factor is None]) <= 1 - if width is None: - width = embedding_dim // (height * input_channels) - if height is None: - height = embedding_dim // (width * input_channels) - if input_channels is None: - input_channels = embedding_dim // (width * height) - assert not any(factor is None for factor in [input_channels, width, height]) - - if input_channels * width * height != embedding_dim: - raise ValueError(f'Could not resolve {original} to a valid factorization of {embedding_dim}.') - - return input_channels, width, height - - -class ConvEInteractionFunction(InteractionFunction): - """ConvE interaction function.""" - - #: If batch normalization is enabled, this is: num_features – C from an expected input of size (N,C,L) - bn0: Optional[torch.nn.BatchNorm2d] - #: If batch normalization is enabled, this is: num_features – C from an expected input of size (N,C,H,W) - bn1: Optional[torch.nn.BatchNorm2d] - bn2: Optional[torch.nn.BatchNorm2d] - - def __init__( - self, - input_channels: Optional[int] = None, - output_channels: int = 32, - embedding_height: Optional[int] = None, - embedding_width: Optional[int] = None, - kernel_height: int = 3, - kernel_width: int = 3, - input_dropout: float = 0.2, - output_dropout: float = 0.3, - feature_map_dropout: float = 0.2, - embedding_dim: int = 200, - apply_batch_normalization: bool = True, - ): - super().__init__() - - # Automatic calculation of remaining dimensions - logger.info(f'Resolving {input_channels} * {embedding_width} * {embedding_height} = {embedding_dim}.') - if embedding_dim is None: - embedding_dim = input_channels * embedding_width * embedding_height - - # Parameter need to fulfil: - # input_channels * embedding_height * embedding_width = embedding_dim - input_channels, embedding_width, embedding_height = _calculate_missing_shape_information( - embedding_dim=embedding_dim, - input_channels=input_channels, - width=embedding_width, - height=embedding_height, - ) - logger.info(f'Resolved to {input_channels} * {embedding_width} * {embedding_height} = {embedding_dim}.') - self.embedding_dim = embedding_dim - self.embedding_height = embedding_height - self.embedding_width = embedding_width - self.input_channels = input_channels - - if self.input_channels * self.embedding_height * self.embedding_width != self.embedding_dim: - raise ValueError( - f'Product of input channels ({self.input_channels}), height ({self.embedding_height}), and width ' - f'({self.embedding_width}) does not equal target embedding dimension ({self.embedding_dim})', - ) - - self.inp_drop = nn.Dropout(input_dropout) - self.hidden_drop = nn.Dropout(output_dropout) - self.feature_map_drop = nn.Dropout2d(feature_map_dropout) - - self.conv1 = torch.nn.Conv2d( - in_channels=self.input_channels, - out_channels=output_channels, - kernel_size=(kernel_height, kernel_width), - stride=1, - padding=0, - bias=True, - ) - - self.apply_batch_normalization = apply_batch_normalization - if self.apply_batch_normalization: - self.bn0 = nn.BatchNorm2d(self.input_channels) - self.bn1 = nn.BatchNorm2d(output_channels) - self.bn2 = nn.BatchNorm1d(self.embedding_dim) - else: - self.bn0 = None - self.bn1 = None - self.bn2 = None - self.num_in_features = ( - output_channels - * (2 * self.embedding_height - kernel_height + 1) - * (self.embedding_width - kernel_width + 1) - ) - self.fc = nn.Linear(self.num_in_features, self.embedding_dim) - self.activation = nn.ReLU() - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - # get tail bias term - if "t_bias" not in kwargs: - raise TypeError(f"{self.__class__.__name__}.forward expects keyword argument 't_bias'.") - t_bias: torch.FloatTensor = kwargs.pop("t_bias") - self._check_for_empty_kwargs(kwargs) - return pykeen_functional.conve_interaction( - h=h, - r=r, - t=t, - t_bias=t_bias, - input_channels=self.input_channels, - embedding_height=self.embedding_height, - embedding_width=self.embedding_width, - num_in_features=self.num_in_features, - bn0=self.bn0, - bn1=self.bn1, - bn2=self.bn2, - inp_drop=self.inp_drop, - feature_map_drop=self.feature_map_drop, - hidden_drop=self.hidden_drop, - conv1=self.conv1, - activation=self.activation, - fc=self.fc, - ) - - class ConvE(EntityRelationEmbeddingModel): r"""An implementation of ConvE from [dettmers2018]_. diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 916245048e..78d92db2a0 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -6,11 +6,10 @@ from typing import Optional import torch.autograd -from torch import nn -from ..base import EntityRelationEmbeddingModel, InteractionFunction +from ..base import EntityRelationEmbeddingModel from ...losses import Loss -from ...nn import functional as pykeen_functional +from ...nn.modules import ConvKBInteractionFunction from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -22,54 +21,6 @@ logger = logging.getLogger(__name__) -class ConvKBInteractionFunction(InteractionFunction): - """Interaction function of ConvKB.""" - - def __init__( - self, - hidden_dropout_rate: float = 0., - embedding_dim: int = 200, - num_filters: int = 400, - ): - super().__init__() - self.embedding_dim = embedding_dim - self.num_filters = num_filters - - # The interaction model - self.conv = nn.Conv2d(in_channels=1, out_channels=num_filters, kernel_size=(1, 3), bias=True) - self.activation = nn.ReLU() - self.hidden_dropout = nn.Dropout(p=hidden_dropout_rate) - self.linear = nn.Linear(embedding_dim * num_filters, 1, bias=True) - - def reset_parameters(self): # noqa: D102 - # Use Xavier initialization for weight; bias to zero - nn.init.xavier_uniform_(self.linear.weight, gain=nn.init.calculate_gain('relu')) - nn.init.zeros_(self.linear.bias) - - # Initialize all filters to [0.1, 0.1, -0.1], - # c.f. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L34-L36 - nn.init.constant_(self.conv.weight[..., :2], 0.1) - nn.init.constant_(self.conv.weight[..., 2], -0.1) - nn.init.zeros_(self.conv.bias) - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - return pykeen_functional.convkb_interaction( - h=h, - r=r, - t=t, - conv=self.conv, - activation=self.activation, - hidden_dropout=self.hidden_dropout, - linear=self.linear, - ) - - class ConvKB(EntityRelationEmbeddingModel): r"""An implementation of ConvKB from [nguyen2018]_. diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 96dd4f8db4..fb8d40ebdd 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -4,13 +4,12 @@ from typing import Optional -import torch.autograd from torch import nn from torch.nn import functional -from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel +from ..base import SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss -from ...nn import functional as pykeen_functional +from ...nn.modules import DistMultInteractionFunction from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -18,24 +17,9 @@ __all__ = [ 'DistMult', - 'DistMultInteractionFunction', ] -class DistMultInteractionFunction(InteractionFunction): - """Interaction function of DistMult.""" - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return pykeen_functional.distmult_interaction(h=h, r=r, t=t) - - class DistMult(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of DistMult from [yang2014]_. diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 44166ac913..4f7766fa04 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -4,81 +4,18 @@ from typing import Optional -import torch -from torch import nn - -from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel +from ..base import SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss -from ...nn import functional as pykeen_functional +from ...nn.modules import ERMLPInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint __all__ = [ 'ERMLP', - 'ERMLPInteractionFunction', ] -class ERMLPInteractionFunction(InteractionFunction): - """ - Interaction function of ER-MLP. - - .. math :: - f(h, r, t) = W_2 ReLU(W_1 cat(h, r, t) + b_1) + b_2 - """ - - def __init__( - self, - embedding_dim: int, - hidden_dim: int, - ): - """ - Initialize the interaction function. - - :param embedding_dim: - The embedding vector dimension. - :param hidden_dim: - The hidden dimension of the MLP. - """ - super().__init__() - """The multi-layer perceptron consisting of an input layer with 3 * self.embedding_dim neurons, a hidden layer - with self.embedding_dim neurons and output layer with one neuron. - The input is represented by the concatenation embeddings of the heads, relations and tail embeddings. - """ - self.hidden = nn.Linear(in_features=3 * embedding_dim, out_features=hidden_dim, bias=True) - self.activation = nn.ReLU() - self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return pykeen_functional.ermlp_interaction( - h=h, - r=r, - t=t, - hidden=self.hidden, - activation=self.activation, - final=self.hidden_to_score, - ) - - def reset_parameters(self): # noqa: D102 - # Initialize biases with zero - nn.init.zeros_(self.hidden.bias) - nn.init.zeros_(self.hidden_to_score.bias) - # In the original formulation, - nn.init.xavier_uniform_(self.hidden.weight) - nn.init.xavier_uniform_( - self.hidden_to_score.weight, - gain=nn.init.calculate_gain(self.activation.__class__.__name__.lower()), - ) - - class ERMLP(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of ERMLP from [dong2014]_. diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index ed45b551d4..5c8c004e90 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -5,11 +5,10 @@ from typing import Optional, Type import torch -from torch import nn -from ..base import EntityRelationEmbeddingModel, InteractionFunction +from ..base import EntityRelationEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss -from ...nn import functional as pykeen_functional +from ...nn.modules import ERMLPEInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -19,40 +18,6 @@ ] -class ERMLPEInteractionFunction(InteractionFunction): - """Interaction function of ER-MLP.""" - - def __init__( - self, - hidden_dim: int = 300, - input_dropout: float = 0.2, - hidden_dropout: float = 0.3, - embedding_dim: int = 200, - ): - super().__init__() - self.mlp = nn.Sequential( - nn.Dropout(input_dropout), - nn.Linear(2 * embedding_dim, hidden_dim), - nn.Dropout(hidden_dropout), - nn.BatchNorm1d(hidden_dim), - nn.ReLU(), - nn.Linear(hidden_dim, embedding_dim), - nn.Dropout(hidden_dropout), - nn.BatchNorm1d(embedding_dim), - nn.ReLU(), - ) - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs=kwargs) - return pykeen_functional.ermlpe_interaction(h=h, r=r, t=t, mlp=self.mlp) - - class ERMLPE(EntityRelationEmbeddingModel): r"""An extension of ERMLP proposed by [sharifzadeh2019]_. diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index c636396a1a..2fba2533e4 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -6,7 +6,8 @@ import torch.autograd -from ..base import InteractionFunction, SimpleVectorEntityRelationEmbeddingModel +from ..base import SimpleVectorEntityRelationEmbeddingModel +from ...nn.modules import InteractionFunction from ...losses import Loss from ...nn.functional import hole_interaction from ...nn.init import xavier_uniform_ diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index bf23e4cdd7..fed88d133c 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -8,7 +8,8 @@ import torch.autograd from .. import Model -from ..base import GeneralVectorEntityRelationEmbeddingModel, IndexFunction, InteractionFunction +from ..base import GeneralVectorEntityRelationEmbeddingModel, IndexFunction +from ...nn.modules import InteractionFunction from ...losses import Loss from ...nn import Embedding from ...nn.init import xavier_normal_ diff --git a/src/pykeen/nn/__init__.py b/src/pykeen/nn/__init__.py index 39930ddcb0..9c60e02c73 100644 --- a/src/pykeen/nn/__init__.py +++ b/src/pykeen/nn/__init__.py @@ -2,7 +2,7 @@ """PyKEEN internal "nn" module.""" -from . import init +from . import functional, init from .emb import Embedding, RepresentationModule __all__ = [ diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py new file mode 100644 index 0000000000..2f2ee4ecf2 --- /dev/null +++ b/src/pykeen/nn/modules.py @@ -0,0 +1,568 @@ +"""Stateful interaction functions.""" +import logging +import math +from typing import Any, Mapping, Optional, Sequence, Tuple + +import torch +from torch import nn + +from . import functional as F +from ..utils import check_shapes + + +class InteractionFunction(nn.Module): + """Base class for interaction functions.""" + + BATCH_DIM: int = 0 + NUM_DIM: int = 1 + HEAD_DIM: int = 1 + RELATION_DIM: int = 2 + TAIL_DIM: int = 3 + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + """Score the given triples. + + :param h: shape: (batch_size, num_heads, d_e) + The head representations. + :param r: shape: (batch_size, num_relations, d_r) + The relation representations. + :param t: shape: (batch_size, num_tails, d_e) + The tail representations. + :param kwargs: + Additional key-word based arguments. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + raise NotImplementedError + + def _check_for_empty_kwargs(self, kwargs: Mapping[str, Any]) -> None: + """Check that kwargs is empty.""" + if len(kwargs) > 0: + raise ValueError(f"{self.__class__.__name__} does not take the following kwargs: {kwargs}") + + @staticmethod + def _add_dim(*x: torch.FloatTensor, dim: int) -> Sequence[torch.FloatTensor]: + """ + Add a dimension to tensors. + + :param x: shape: (d1, ..., dk) + The tensor. + + :return: shape: (1, d1, ..., dk) + The tensor with batch dimension. + """ + out = [xx.unsqueeze(dim=dim) for xx in x] + if len(x) > 1: + return out + return out[0] + + @staticmethod + def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: + """ + Remove dimensions from a tensor. + + :param x: + The tensor. + :param dims: + The dimensions to remove. + + :return: + The squeezed tensor. + """ + # normalize dimensions + dims = [d if d >= 0 else len(x.shape) + d for d in dims] + if len(set(dims)) != len(dims): + raise ValueError(f"Duplicate dimensions: {dims}") + assert all(0 <= d < len(x.shape) for d in dims) + for dim in sorted(dims, reverse=True): + x = x.squeeze(dim=dim) + return x + + def score_hrt( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + """ + Score a batch of triples.. + + :param h: shape: (batch_size, d_e) + The head representations. + :param r: shape: (batch_size, d_r) + The relation representations. + :param t: shape: (batch_size, d_e) + The tail representations. + :param kwargs: + Additional key-word based arguments. + + :return: shape: (batch_size, 1) + The scores. + """ + # check shape + assert check_shapes((h, "be"), (r, "br"), (t, "be")) + + # prepare input to generic score function + h, r, t = self._add_dim(h, r, t, dim=self.NUM_DIM) + + # get scores + scores = self(h=h, r=r, t=t, **kwargs) + + # prepare output shape, (batch_size, num_heads, num_relations, num_tails) -> (batch_size, 1) + return self._remove_dim(scores, self.HEAD_DIM, self.RELATION_DIM, self.TAIL_DIM).unsqueeze(dim=-1) + + def score_h( + self, + all_entities: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + """ + Score all head entities. + + :param all_entities: shape: (num_entities, d_e) + The head representations. + :param r: shape: (batch_size, d_r) + The relation representations. + :param t: shape: (batch_size, d_e) + The tail representations. + :param kwargs: + Additional key-word based arguments. + + :return: shape: (batch_size, num_entities) + The scores. + """ + # check shape + assert check_shapes((all_entities, "ne"), (r, "br"), (t, "be")) + + # TODO: What about unsqueezing for additional e.g. head arguments + # prepare input to generic score function + r, t = self._add_dim(r, t, dim=self.NUM_DIM) + h = self._add_dim(all_entities, dim=self.BATCH_DIM) + + # get scores + scores = self(h=h, r=r, t=t, **kwargs) + + # prepare output shape, (batch_size, num_heads, num_relations, num_tails) -> (batch_size, num_heads) + return self._remove_dim(scores, self.RELATION_DIM, self.TAIL_DIM) + + def score_r( + self, + h: torch.FloatTensor, + all_relations: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + """ + Score all relations. + + :param h: shape: (batch_size, d_e) + The head representations. + :param all_relations: shape: (batch_size, d_r) + The relation representations. + :param t: shape: (num_entities, d_e) + The tail representations. + :param kwargs: + Additional key-word based arguments. + + :return: shape: (batch_size, num_entities) + The scores. + """ + # check shape + assert check_shapes((all_relations, "nr"), (h, "be"), (t, "be")) + + # prepare input to generic score function + h, t = self._add_dim(h, t, dim=self.NUM_DIM) + r = self._add_dim(all_relations, dim=self.BATCH_DIM) + + # get scores + scores = self(h=h, r=r, t=t, **kwargs) + + # prepare output shape + return self._remove_dim(scores, self.HEAD_DIM, self.TAIL_DIM) + + def score_t( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + all_entities: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + """ + Score all tail entities. + + :param h: shape: (batch_size, d_e) + The head representations. + :param r: shape: (batch_size, d_r) + The relation representations. + :param all_entities: shape: (num_entities, d_e) + The tail representations. + :param kwargs: + Additional key-word based arguments. + + :return: shape: (batch_size, num_entities) + The scores. + """ + # check shape + assert check_shapes((all_entities, "ne"), (r, "br"), (h, "be")) + + # prepare input to generic score function + h, r = self._add_dim(h, r, dim=self.NUM_DIM) + t = self._add_dim(all_entities, dim=self.BATCH_DIM) + + # get scores + scores = self(h=h, r=r, t=t, **kwargs) + + # prepare output shape + return self._remove_dim(scores, self.HEAD_DIM, self.RELATION_DIM) + + def reset_parameters(self): + """Reset parameters the interaction function may have.""" + for mod in self.modules(): + if mod is self: + continue + if hasattr(mod, 'reset_parameters'): + mod.reset_parameters() + + +class ComplExInteractionFunction(InteractionFunction): + """Interaction function of ComplEx.""" + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + self._check_for_empty_kwargs(kwargs) + return F.complex_interaction(h=h, r=r, t=t) + + +logger = logging.getLogger(__name__) + + +def _calculate_missing_shape_information( + embedding_dim: int, + input_channels: Optional[int] = None, + width: Optional[int] = None, + height: Optional[int] = None, +) -> Tuple[int, int, int]: + """ + Automatically calculates missing dimensions for ConvE. + + :param embedding_dim: + :param input_channels: + :param width: + :param height: + + :return: (input_channels, width, height), such that + `embedding_dim = input_channels * width * height` + :raises: + If no factorization could be found. + """ + # Store initial input for error message + original = (input_channels, width, height) + + # All are None + if all(factor is None for factor in [input_channels, width, height]): + input_channels = 1 + result_sqrt = math.floor(math.sqrt(embedding_dim)) + height = max(factor for factor in range(1, result_sqrt + 1) if embedding_dim % factor == 0) + width = embedding_dim // height + + # input_channels is None, and any of height or width is None -> set input_channels=1 + if input_channels is None and any(remaining is None for remaining in [width, height]): + input_channels = 1 + + # input channels is not None, and one of height or width is None + assert len([factor for factor in [input_channels, width, height] if factor is None]) <= 1 + if width is None: + width = embedding_dim // (height * input_channels) + if height is None: + height = embedding_dim // (width * input_channels) + if input_channels is None: + input_channels = embedding_dim // (width * height) + assert not any(factor is None for factor in [input_channels, width, height]) + + if input_channels * width * height != embedding_dim: + raise ValueError(f'Could not resolve {original} to a valid factorization of {embedding_dim}.') + + return input_channels, width, height + + +class ConvEInteractionFunction(InteractionFunction): + """ConvE interaction function.""" + + #: If batch normalization is enabled, this is: num_features – C from an expected input of size (N,C,L) + bn0: Optional[torch.nn.BatchNorm2d] + #: If batch normalization is enabled, this is: num_features – C from an expected input of size (N,C,H,W) + bn1: Optional[torch.nn.BatchNorm2d] + bn2: Optional[torch.nn.BatchNorm2d] + + def __init__( + self, + input_channels: Optional[int] = None, + output_channels: int = 32, + embedding_height: Optional[int] = None, + embedding_width: Optional[int] = None, + kernel_height: int = 3, + kernel_width: int = 3, + input_dropout: float = 0.2, + output_dropout: float = 0.3, + feature_map_dropout: float = 0.2, + embedding_dim: int = 200, + apply_batch_normalization: bool = True, + ): + super().__init__() + + # Automatic calculation of remaining dimensions + logger.info(f'Resolving {input_channels} * {embedding_width} * {embedding_height} = {embedding_dim}.') + if embedding_dim is None: + embedding_dim = input_channels * embedding_width * embedding_height + + # Parameter need to fulfil: + # input_channels * embedding_height * embedding_width = embedding_dim + input_channels, embedding_width, embedding_height = _calculate_missing_shape_information( + embedding_dim=embedding_dim, + input_channels=input_channels, + width=embedding_width, + height=embedding_height, + ) + logger.info(f'Resolved to {input_channels} * {embedding_width} * {embedding_height} = {embedding_dim}.') + self.embedding_dim = embedding_dim + self.embedding_height = embedding_height + self.embedding_width = embedding_width + self.input_channels = input_channels + + if self.input_channels * self.embedding_height * self.embedding_width != self.embedding_dim: + raise ValueError( + f'Product of input channels ({self.input_channels}), height ({self.embedding_height}), and width ' + f'({self.embedding_width}) does not equal target embedding dimension ({self.embedding_dim})', + ) + + self.inp_drop = nn.Dropout(input_dropout) + self.hidden_drop = nn.Dropout(output_dropout) + self.feature_map_drop = nn.Dropout2d(feature_map_dropout) + + self.conv1 = torch.nn.Conv2d( + in_channels=self.input_channels, + out_channels=output_channels, + kernel_size=(kernel_height, kernel_width), + stride=1, + padding=0, + bias=True, + ) + + self.apply_batch_normalization = apply_batch_normalization + if self.apply_batch_normalization: + self.bn0 = nn.BatchNorm2d(self.input_channels) + self.bn1 = nn.BatchNorm2d(output_channels) + self.bn2 = nn.BatchNorm1d(self.embedding_dim) + else: + self.bn0 = None + self.bn1 = None + self.bn2 = None + self.num_in_features = ( + output_channels + * (2 * self.embedding_height - kernel_height + 1) + * (self.embedding_width - kernel_width + 1) + ) + self.fc = nn.Linear(self.num_in_features, self.embedding_dim) + self.activation = nn.ReLU() + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + # get tail bias term + if "t_bias" not in kwargs: + raise TypeError(f"{self.__class__.__name__}.forward expects keyword argument 't_bias'.") + t_bias: torch.FloatTensor = kwargs.pop("t_bias") + self._check_for_empty_kwargs(kwargs) + return F.conve_interaction( + h=h, + r=r, + t=t, + t_bias=t_bias, + input_channels=self.input_channels, + embedding_height=self.embedding_height, + embedding_width=self.embedding_width, + num_in_features=self.num_in_features, + bn0=self.bn0, + bn1=self.bn1, + bn2=self.bn2, + inp_drop=self.inp_drop, + feature_map_drop=self.feature_map_drop, + hidden_drop=self.hidden_drop, + conv1=self.conv1, + activation=self.activation, + fc=self.fc, + ) + + +class ConvKBInteractionFunction(InteractionFunction): + """Interaction function of ConvKB.""" + + def __init__( + self, + hidden_dropout_rate: float = 0., + embedding_dim: int = 200, + num_filters: int = 400, + ): + super().__init__() + self.embedding_dim = embedding_dim + self.num_filters = num_filters + + # The interaction model + self.conv = nn.Conv2d(in_channels=1, out_channels=num_filters, kernel_size=(1, 3), bias=True) + self.activation = nn.ReLU() + self.hidden_dropout = nn.Dropout(p=hidden_dropout_rate) + self.linear = nn.Linear(embedding_dim * num_filters, 1, bias=True) + + def reset_parameters(self): # noqa: D102 + # Use Xavier initialization for weight; bias to zero + nn.init.xavier_uniform_(self.linear.weight, gain=nn.init.calculate_gain('relu')) + nn.init.zeros_(self.linear.bias) + + # Initialize all filters to [0.1, 0.1, -0.1], + # c.f. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L34-L36 + nn.init.constant_(self.conv.weight[..., :2], 0.1) + nn.init.constant_(self.conv.weight[..., 2], -0.1) + nn.init.zeros_(self.conv.bias) + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + return F.convkb_interaction( + h=h, + r=r, + t=t, + conv=self.conv, + activation=self.activation, + hidden_dropout=self.hidden_dropout, + linear=self.linear, + ) + + +class DistMultInteractionFunction(InteractionFunction): + """Interaction function of DistMult.""" + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + self._check_for_empty_kwargs(kwargs) + return F.distmult_interaction(h=h, r=r, t=t) + + +class ERMLPInteractionFunction(InteractionFunction): + """ + Interaction function of ER-MLP. + + .. math :: + f(h, r, t) = W_2 ReLU(W_1 cat(h, r, t) + b_1) + b_2 + """ + + def __init__( + self, + embedding_dim: int, + hidden_dim: int, + ): + """ + Initialize the interaction function. + + :param embedding_dim: + The embedding vector dimension. + :param hidden_dim: + The hidden dimension of the MLP. + """ + super().__init__() + """The multi-layer perceptron consisting of an input layer with 3 * self.embedding_dim neurons, a hidden layer + with self.embedding_dim neurons and output layer with one neuron. + The input is represented by the concatenation embeddings of the heads, relations and tail embeddings. + """ + self.hidden = nn.Linear(in_features=3 * embedding_dim, out_features=hidden_dim, bias=True) + self.activation = nn.ReLU() + self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + self._check_for_empty_kwargs(kwargs) + return F.ermlp_interaction( + h=h, + r=r, + t=t, + hidden=self.hidden, + activation=self.activation, + final=self.hidden_to_score, + ) + + def reset_parameters(self): # noqa: D102 + # Initialize biases with zero + nn.init.zeros_(self.hidden.bias) + nn.init.zeros_(self.hidden_to_score.bias) + # In the original formulation, + nn.init.xavier_uniform_(self.hidden.weight) + nn.init.xavier_uniform_( + self.hidden_to_score.weight, + gain=nn.init.calculate_gain(self.activation.__class__.__name__.lower()), + ) + + +class ERMLPEInteractionFunction(InteractionFunction): + """Interaction function of ER-MLP.""" + + def __init__( + self, + hidden_dim: int = 300, + input_dropout: float = 0.2, + hidden_dropout: float = 0.3, + embedding_dim: int = 200, + ): + super().__init__() + self.mlp = nn.Sequential( + nn.Dropout(input_dropout), + nn.Linear(2 * embedding_dim, hidden_dim), + nn.Dropout(hidden_dropout), + nn.BatchNorm1d(hidden_dim), + nn.ReLU(), + nn.Linear(hidden_dim, embedding_dim), + nn.Dropout(hidden_dropout), + nn.BatchNorm1d(embedding_dim), + nn.ReLU(), + ) + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + self._check_for_empty_kwargs(kwargs=kwargs) + return F.ermlpe_interaction(h=h, r=r, t=t, mlp=self.mlp) From 5e3c411ce8550ae1cf2a0bcc76e5b24a978fe25b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 13:57:44 +0100 Subject: [PATCH 104/690] Move logger to start of file --- src/pykeen/nn/modules.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 2f2ee4ecf2..a6d5caf064 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -9,6 +9,8 @@ from . import functional as F from ..utils import check_shapes +logger = logging.getLogger(__name__) + class InteractionFunction(nn.Module): """Base class for interaction functions.""" @@ -248,9 +250,6 @@ def forward( return F.complex_interaction(h=h, r=r, t=t) -logger = logging.getLogger(__name__) - - def _calculate_missing_shape_information( embedding_dim: int, input_channels: Optional[int] = None, From a0f7ad16c68fd6edd56b8225997c87e55a86bf82 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 14:24:18 +0100 Subject: [PATCH 105/690] Abstract the trans-series @mberr i noticed there were some bugs - a lot of the functions did not respect the norm's l_p values given in the __init__... --- src/pykeen/models/base.py | 27 +++++++++++- src/pykeen/models/unimodal/trans_d.py | 30 ++++++++++--- src/pykeen/models/unimodal/trans_e.py | 12 ++--- src/pykeen/models/unimodal/trans_h.py | 12 ++--- src/pykeen/models/unimodal/trans_r.py | 63 ++++++++++++--------------- src/pykeen/nn/functional.py | 52 +++++++++++++++++++++- 6 files changed, 138 insertions(+), 58 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 225b320ba8..c840dbe3cc 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -16,7 +16,7 @@ from torch import nn from ..losses import Loss, MarginRankingLoss, NSSALoss -from ..nn import Embedding +from ..nn import Embedding, functional as pykeen_functional from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory from ..typing import Constrainer, DeviceHint, Initializer, MappedTriples, Normalizer @@ -29,6 +29,7 @@ 'GeneralVectorEntityRelationEmbeddingModel', 'SimpleVectorEntityRelationEmbeddingModel', 'InteractionFunction', + 'TranslationalInteractionFunction', 'IndexFunction', 'MultimodalModel', ] @@ -1425,6 +1426,30 @@ def reset_parameters(self): mod.reset_parameters() +class TranslationalInteractionFunction(InteractionFunction): + """The translational interaction function shared by the TransE, TransR, TransH, and other Trans models.""" + + def __init__(self, p: int): + """Initialize the translational interaction function. + + :param p: The norm used with :func:`torch.norm`. Typically is 1 or 2. + """ + super().__init__() + self.p = p + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa:D102 + return pykeen_functional.translational_interaction( + h=h, r=r, t=t, + p=self.p, dim=kwargs.get('dim', None), keepdim=kwargs.get('keepdim', False), + ) + + class IndexFunction(nn.Module): """A function that handles looking up embeddings by index.""" diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index bf23e4cdd7..dbfd0d1321 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -8,7 +8,10 @@ import torch.autograd from .. import Model -from ..base import GeneralVectorEntityRelationEmbeddingModel, IndexFunction, InteractionFunction +from ..base import ( + GeneralVectorEntityRelationEmbeddingModel, IndexFunction, InteractionFunction, + TranslationalInteractionFunction, +) from ...losses import Loss from ...nn import Embedding from ...nn.init import xavier_normal_ @@ -18,6 +21,8 @@ from ...utils import clamp_norm __all__ = [ + 'TransDInteractionFunction', + 'TransDIndexFunction', 'TransD', ] @@ -69,15 +74,26 @@ def _project_entity( return e_bot -class TransDInteractionFunction(InteractionFunction): +class TransDInteractionFunction(TranslationalInteractionFunction): + """The interaction function for TransD.""" + def __init__(self, p: int = 2, power: int = 2): - # Very similar to TransE, could be generalized - super().__init__() - self.p = p + """Initialize the TransD interaction function. + + :param p: The norm applied by :func:`torch.norm` + :param power: The power applied after :func:`torch.norm`. + """ + super().__init__(p=p) self.power = power - def forward(self, h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor) -> torch.FloatTensor: - return -torch.norm(h + r - t, dim=-1, p=self.p) ** self.power + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa:D102 + return super().forward(h=h, r=r, t=t, **kwargs) ** self.power class TransDIndexFunction(IndexFunction): diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index b653ed66a3..819c3c33bb 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -4,11 +4,10 @@ from typing import Optional -import torch import torch.autograd from torch.nn import functional -from ..base import EntityRelationEmbeddingModel +from ..base import EntityRelationEmbeddingModel, TranslationalInteractionFunction from ...losses import Loss from ...nn.init import xavier_uniform_ from ...regularizers import Regularizer @@ -67,6 +66,7 @@ def __init__( - OpenKE `implementation of TransE `_ """ + self.interaction_function = TranslationalInteractionFunction(p=scoring_fct_norm) super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, @@ -82,7 +82,6 @@ def __init__( ), entity_constrainer=functional.normalize, ) - self.scoring_fct_norm = scoring_fct_norm def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings @@ -92,7 +91,8 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: # TODO question @mberr - why is keepdim=True here but the others it isn't? # TODO: Use torch.dist - return -torch.norm(h + r - t, dim=-1, p=self.scoring_fct_norm, keepdim=True) + + return self.interaction_function(h=h, r=r, t=t, dim=-1, keepdim=True) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings @@ -101,7 +101,7 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 t = self.entity_embeddings(indices=None) # TODO: Use torch.cdist - return -torch.norm(h[:, None, :] + r[:, None, :] - t[None, :, :], dim=-1, p=self.scoring_fct_norm) + return self.interaction_function(h=h[:, None, :], r=r[:, None, :], t=t[None, :, :], dim=-1, keepdim=False) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings @@ -110,4 +110,4 @@ def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 t = self.entity_embeddings(indices=rt_batch[:, 1]) # TODO: Use torch.cdist - return -torch.norm(h[None, :, :] + r[:, None, :] - t[:, None, :], dim=-1, p=self.scoring_fct_norm) + return self.interaction_function(h=h[None, :, :], r=r[:, None, :], t=t[:, None, :], dim=-1, keepdim=False) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index b233615a82..a1a4737f5a 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -7,7 +7,7 @@ import torch from torch.nn import functional -from ..base import EntityRelationEmbeddingModel +from ..base import EntityRelationEmbeddingModel, TranslationalInteractionFunction from ...losses import Loss from ...nn import Embedding from ...regularizers import Regularizer, TransHRegularizer @@ -67,7 +67,7 @@ def __init__( triples_factory: TriplesFactory, embedding_dim: int = 50, automatic_memory_optimization: Optional[bool] = None, - scoring_fct_norm: int = 1, + scoring_fct_norm: int = 2, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -88,7 +88,7 @@ def __init__( regularizer=regularizer, ) - self.scoring_fct_norm = scoring_fct_norm + self.interaction_function = TranslationalInteractionFunction(p=scoring_fct_norm) # embeddings self.normal_vector_embeddings = Embedding.init_with_device( @@ -132,7 +132,7 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: # Regularization term self.regularize_if_necessary() - return -torch.norm(ph + d_r - pt, p=2, dim=-1, keepdim=True) + return self.interaction_function(h=ph, r=d_r, t=pt, dim=-1, keepdim=True) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings @@ -148,7 +148,7 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 # Regularization term self.regularize_if_necessary() - return -torch.norm(ph[:, None, :] + d_r[:, None, :] - pt, p=2, dim=-1) + return self.interaction_function(h=ph[:, None, :], r=d_r[:, None, :], t=pt, dim=-1, keepdim=False) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings @@ -165,4 +165,4 @@ def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 # Regularization term self.regularize_if_necessary() - return -torch.norm(ph + d_r[:, None, :] - pt[:, None, :], p=2, dim=-1) + return self.interaction_function(h=ph, r=d_r[:, None, :], t=pt[:, None, :], dim=-1, keepdim=False) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 17cf2d706d..228ca91f7f 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -10,9 +10,9 @@ import torch.nn.init from torch.nn import functional -from ..base import EntityRelationEmbeddingModel +from ..base import EntityRelationEmbeddingModel, InteractionFunction from ...losses import Loss -from ...nn import Embedding +from ...nn import Embedding, functional as pykeen_functional from ...nn.init import xavier_uniform_ from ...regularizers import Regularizer from ...triples import TriplesFactory @@ -20,6 +20,7 @@ from ...utils import clamp_norm, compose __all__ = [ + 'TransRInteractionFunction', 'TransR', ] @@ -34,6 +35,28 @@ def _projection_initializer( return torch.nn.init.xavier_uniform_(x.view(num_relations, embedding_dim, relation_dim)).view(x.shape) +class TransRInteractionFunction(InteractionFunction): + """The TransR interaction function.""" + + def __init__(self, p: int): + """Initialize the TransR interaction function. + + :param p: The norm applied to the translation + """ + super().__init__() + self.p = p + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa:D102 + m_r = kwargs.pop('m_r') + return pykeen_functional.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p) + + class TransR(EntityRelationEmbeddingModel): r"""An implementation of TransR from [lin2015]_. @@ -85,6 +108,8 @@ def __init__( regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model.""" + self.interaction_function = TransRInteractionFunction(p=scoring_fct_norm) + super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, @@ -104,7 +129,6 @@ def __init__( relation_constrainer=clamp_norm, relation_constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ) - self.scoring_fct_norm = scoring_fct_norm # TODO: Initialize from TransE @@ -125,39 +149,6 @@ def _reset_parameters_(self): # noqa: D102 super()._reset_parameters_() self.relation_projections.reset_parameters() - @staticmethod - def interaction_function( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - m_r: torch.FloatTensor, - ) -> torch.FloatTensor: - """Evaluate the interaction function for given embeddings. - - The embeddings have to be in a broadcastable shape. - - :param h: shape: (batch_size, num_entities, d_e) - Head embeddings. - :param r: shape: (batch_size, num_entities, d_r) - Relation embeddings. - :param t: shape: (batch_size, num_entities, d_e) - Tail embeddings. - :param m_r: shape: (batch_size, num_entities, d_e, d_r) - The relation specific linear transformations. - - :return: shape: (batch_size, num_entities) - The scores. - """ - # project to relation specific subspace, shape: (b, e, d_r) - h_bot = h @ m_r - t_bot = t @ m_r - # ensure constraints - h_bot = clamp_norm(h_bot, p=2, dim=-1, maxnorm=1.) - t_bot = clamp_norm(t_bot, p=2, dim=-1, maxnorm=1.) - - # evaluate score function, shape: (b, e) - return -torch.norm(h_bot + r - t_bot, dim=-1) ** 2 - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings # TODO switch to self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 0]) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index c874e947dc..8419e0c63d 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -2,12 +2,12 @@ """Functional forms of interaction methods.""" -from typing import Optional, Tuple +from typing import Optional, Tuple, Union import torch from torch import nn -from ..utils import broadcast_cat, is_cudnn_error, normalize_for_einsum, split_complex +from ..utils import broadcast_cat, clamp_norm, is_cudnn_error, normalize_for_einsum, split_complex __all__ = [ "complex_interaction", @@ -18,6 +18,8 @@ "ermlpe_interaction", 'hole_interaction', 'rotate_interaction', + 'translational_interaction', + 'transr_interaction', ] @@ -417,3 +419,49 @@ def rotate_interaction( scores = -torch.norm(diff.view(diff.shape[:-2] + (-1,)), dim=-1) return scores + + +def translational_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + dim: int, + p: Union[int, str] = 'fro', + keepdim: bool = False, +) -> torch.FloatTensor: + """Evaluate the translational interaction.""" + return -torch.norm(h + r - t, dim=dim, p=p, keepdim=keepdim) + + +def transr_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + m_r: torch.FloatTensor, + p: int, +) -> torch.FloatTensor: + """Evaluate the interaction function for given embeddings. + + The embeddings have to be in a broadcastable shape. + + :param h: shape: (batch_size, num_entities, d_e) + Head embeddings. + :param r: shape: (batch_size, num_entities, d_r) + Relation embeddings. + :param t: shape: (batch_size, num_entities, d_e) + Tail embeddings. + :param m_r: shape: (batch_size, num_entities, d_e, d_r) + The relation specific linear transformations. + + :return: shape: (batch_size, num_entities) + The scores. + """ + # project to relation specific subspace, shape: (b, e, d_r) + h_bot = h @ m_r + t_bot = t @ m_r + # ensure constraints + h_bot = clamp_norm(h_bot, p=2, dim=-1, maxnorm=1.) + t_bot = clamp_norm(t_bot, p=2, dim=-1, maxnorm=1.) + + # evaluate score function, shape: (b, e) + return translational_interaction(h=h_bot, r=r, t=t_bot, dim=-1, p=p) ** 2 From 9527a4cce5d0936ba94a3c6de436c591223beb8d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 14:44:38 +0100 Subject: [PATCH 106/690] Merge --- src/pykeen/models/unimodal/trans_d.py | 7 ++----- src/pykeen/nn/modules.py | 24 ++++++++++++++++++++++++ 2 files changed, 26 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index d321cb72fb..db22c21781 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -8,14 +8,11 @@ import torch.autograd from .. import Model -from ..base import ( - GeneralVectorEntityRelationEmbeddingModel, IndexFunction -from ...nn.modules import InteractionFunction, - TranslationalInteractionFunction, -) +from ..base import GeneralVectorEntityRelationEmbeddingModel, IndexFunction from ...losses import Loss from ...nn import Embedding from ...nn.init import xavier_normal_ +from ...nn.modules import InteractionFunction, TranslationalInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index a6d5caf064..bd7ae864f3 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -236,6 +236,30 @@ def reset_parameters(self): mod.reset_parameters() +class TranslationalInteractionFunction(InteractionFunction): + """The translational interaction function shared by the TransE, TransR, TransH, and other Trans models.""" + + def __init__(self, p: int): + """Initialize the translational interaction function. + + :param p: The norm used with :func:`torch.norm`. Typically is 1 or 2. + """ + super().__init__() + self.p = p + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa:D102 + return F.translational_interaction( + h=h, r=r, t=t, + p=self.p, dim=kwargs.get('dim', None), keepdim=kwargs.get('keepdim', False), + ) + + class ComplExInteractionFunction(InteractionFunction): """Interaction function of ComplEx.""" From c85ae5136f10d2f7c51dcfd1a26beb7d8e42c3e0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 14:47:34 +0100 Subject: [PATCH 107/690] Post merge fixes --- src/pykeen/models/unimodal/trans_d.py | 25 +-------------- src/pykeen/models/unimodal/trans_e.py | 3 +- src/pykeen/models/unimodal/trans_h.py | 3 +- src/pykeen/models/unimodal/trans_r.py | 27 ++-------------- src/pykeen/nn/modules.py | 46 ++++++++++++++++++++++++++- 5 files changed, 52 insertions(+), 52 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index db22c21781..1a3416ded6 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -12,14 +12,13 @@ from ...losses import Loss from ...nn import Embedding from ...nn.init import xavier_normal_ -from ...nn.modules import InteractionFunction, TranslationalInteractionFunction +from ...nn.modules import InteractionFunction, TransDInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint from ...utils import clamp_norm __all__ = [ - 'TransDInteractionFunction', 'TransDIndexFunction', 'TransD', ] @@ -72,28 +71,6 @@ def _project_entity( return e_bot -class TransDInteractionFunction(TranslationalInteractionFunction): - """The interaction function for TransD.""" - - def __init__(self, p: int = 2, power: int = 2): - """Initialize the TransD interaction function. - - :param p: The norm applied by :func:`torch.norm` - :param power: The power applied after :func:`torch.norm`. - """ - super().__init__(p=p) - self.power = power - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa:D102 - return super().forward(h=h, r=r, t=t, **kwargs) ** self.power - - class TransDIndexFunction(IndexFunction): """The index-based interaction function for TransD.""" diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 819c3c33bb..05617758a7 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -7,9 +7,10 @@ import torch.autograd from torch.nn import functional -from ..base import EntityRelationEmbeddingModel, TranslationalInteractionFunction +from ..base import EntityRelationEmbeddingModel from ...losses import Loss from ...nn.init import xavier_uniform_ +from ...nn.modules import TranslationalInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index a1a4737f5a..da86535acf 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -7,9 +7,10 @@ import torch from torch.nn import functional -from ..base import EntityRelationEmbeddingModel, TranslationalInteractionFunction +from ..base import EntityRelationEmbeddingModel from ...losses import Loss from ...nn import Embedding +from ...nn.modules import TranslationalInteractionFunction from ...regularizers import Regularizer, TransHRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 228ca91f7f..df14dc59eb 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -10,9 +10,9 @@ import torch.nn.init from torch.nn import functional -from ..base import EntityRelationEmbeddingModel, InteractionFunction +from ..base import EntityRelationEmbeddingModel from ...losses import Loss -from ...nn import Embedding, functional as pykeen_functional +from ...nn import Embedding from ...nn.init import xavier_uniform_ from ...regularizers import Regularizer from ...triples import TriplesFactory @@ -20,7 +20,6 @@ from ...utils import clamp_norm, compose __all__ = [ - 'TransRInteractionFunction', 'TransR', ] @@ -35,28 +34,6 @@ def _projection_initializer( return torch.nn.init.xavier_uniform_(x.view(num_relations, embedding_dim, relation_dim)).view(x.shape) -class TransRInteractionFunction(InteractionFunction): - """The TransR interaction function.""" - - def __init__(self, p: int): - """Initialize the TransR interaction function. - - :param p: The norm applied to the translation - """ - super().__init__() - self.p = p - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa:D102 - m_r = kwargs.pop('m_r') - return pykeen_functional.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p) - - class TransR(EntityRelationEmbeddingModel): r"""An implementation of TransR from [lin2015]_. diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index bd7ae864f3..4aa6caaa83 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -6,7 +6,7 @@ import torch from torch import nn -from . import functional as F +from . import functional as F, functional as pykeen_functional from ..utils import check_shapes logger = logging.getLogger(__name__) @@ -589,3 +589,47 @@ def forward( ) -> torch.FloatTensor: # noqa: D102 self._check_for_empty_kwargs(kwargs=kwargs) return F.ermlpe_interaction(h=h, r=r, t=t, mlp=self.mlp) + + +class TransDInteractionFunction(TranslationalInteractionFunction): + """The interaction function for TransD.""" + + def __init__(self, p: int = 2, power: int = 2): + """Initialize the TransD interaction function. + + :param p: The norm applied by :func:`torch.norm` + :param power: The power applied after :func:`torch.norm`. + """ + super().__init__(p=p) + self.power = power + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa:D102 + return super().forward(h=h, r=r, t=t, **kwargs) ** self.power + + +class TransRInteractionFunction(InteractionFunction): + """The TransR interaction function.""" + + def __init__(self, p: int): + """Initialize the TransR interaction function. + + :param p: The norm applied to the translation + """ + super().__init__() + self.p = p + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa:D102 + m_r = kwargs.pop('m_r') + return pykeen_functional.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p) \ No newline at end of file From 24a108ef107a547323bd61a73d3a54425515daf4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 14:47:58 +0100 Subject: [PATCH 108/690] fix TransR --- src/pykeen/models/unimodal/trans_r.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index df14dc59eb..5cdb5a0bb7 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -14,6 +14,7 @@ from ...losses import Loss from ...nn import Embedding from ...nn.init import xavier_uniform_ +from ...nn.modules import TransRInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint From 9f9e83a4ea6124a7e348fd68349bc99feaabd6a8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 14:48:53 +0100 Subject: [PATCH 109/690] Unify ' vs " --- src/pykeen/nn/functional.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 8419e0c63d..e61bab510d 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -16,10 +16,10 @@ "distmult_interaction", "ermlp_interaction", "ermlpe_interaction", - 'hole_interaction', - 'rotate_interaction', - 'translational_interaction', - 'transr_interaction', + "hole_interaction", + "rotate_interaction", + "translational_interaction", + "transr_interaction", ] From 115177a1b495d8476f91d2566e713408c7dad491 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 14:53:06 +0100 Subject: [PATCH 110/690] Move RotatE interaction function --- src/pykeen/models/unimodal/rotate.py | 18 ++---------------- src/pykeen/nn/modules.py | 16 +++++++++++++++- 2 files changed, 17 insertions(+), 17 deletions(-) diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index 3cc745d263..84e81457eb 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -9,10 +9,10 @@ import torch.autograd from torch.nn import functional -from ..base import EntityRelationEmbeddingModel, InteractionFunction +from ..base import EntityRelationEmbeddingModel from ...losses import Loss -from ...nn import functional as pykeen_functional from ...nn.init import xavier_uniform_ +from ...nn.modules import RotatEInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -52,20 +52,6 @@ def complex_normalize(x: torch.Tensor) -> torch.Tensor: return x -class RotatEInteraction(InteractionFunction): - """Interaction function of RotatE.""" - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return pykeen_functional.rotate_interaction(h=h, r=r, t=t) - - class RotatE(EntityRelationEmbeddingModel): r"""An implementation of RotatE from [sun2019]_. diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 4aa6caaa83..9551d2348c 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -632,4 +632,18 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa:D102 m_r = kwargs.pop('m_r') - return pykeen_functional.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p) \ No newline at end of file + return pykeen_functional.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p) + + +class RotatEInteraction(InteractionFunction): + """Interaction function of RotatE.""" + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + self._check_for_empty_kwargs(kwargs) + return pykeen_functional.rotate_interaction(h=h, r=r, t=t) \ No newline at end of file From dc77ab8120c11b09a15e989397c64bf3e23c3a42 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 14:53:37 +0100 Subject: [PATCH 111/690] Move HolE interaction function --- src/pykeen/models/unimodal/hole.py | 19 +------------------ src/pykeen/nn/modules.py | 17 ++++++++++++++++- 2 files changed, 17 insertions(+), 19 deletions(-) diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 2fba2533e4..a08ac7e9e9 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -4,13 +4,10 @@ from typing import Optional -import torch.autograd - from ..base import SimpleVectorEntityRelationEmbeddingModel -from ...nn.modules import InteractionFunction from ...losses import Loss -from ...nn.functional import hole_interaction from ...nn.init import xavier_uniform_ +from ...nn.modules import HolEInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -21,20 +18,6 @@ ] -class HolEInteractionFunction(InteractionFunction): - """Interaction function for HolE.""" - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return hole_interaction(h=h, r=r, t=t) - - class HolE(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of HolE [nickel2016]_. diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 9551d2348c..cf42a3dfb0 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -7,6 +7,7 @@ from torch import nn from . import functional as F, functional as pykeen_functional +from .functional import hole_interaction from ..utils import check_shapes logger = logging.getLogger(__name__) @@ -646,4 +647,18 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa: D102 self._check_for_empty_kwargs(kwargs) - return pykeen_functional.rotate_interaction(h=h, r=r, t=t) \ No newline at end of file + return pykeen_functional.rotate_interaction(h=h, r=r, t=t) + + +class HolEInteractionFunction(InteractionFunction): + """Interaction function for HolE.""" + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + self._check_for_empty_kwargs(kwargs) + return hole_interaction(h=h, r=r, t=t) \ No newline at end of file From c6b199685003e17a72242c51f150b71a416a66cf Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 15:14:15 +0100 Subject: [PATCH 112/690] Abstract simple interaction function and cleanup flake8 --- src/pykeen/nn/__init__.py | 1 + src/pykeen/nn/modules.py | 68 ++++++++++++++++----------------------- 2 files changed, 29 insertions(+), 40 deletions(-) diff --git a/src/pykeen/nn/__init__.py b/src/pykeen/nn/__init__.py index 9c60e02c73..2b24107f24 100644 --- a/src/pykeen/nn/__init__.py +++ b/src/pykeen/nn/__init__.py @@ -9,4 +9,5 @@ 'Embedding', 'RepresentationModule', 'init', + 'functional', ] diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index cf42a3dfb0..d261f3ce3c 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -1,4 +1,7 @@ +# -*- coding: utf-8 -*- + """Stateful interaction functions.""" + import logging import math from typing import Any, Mapping, Optional, Sequence, Tuple @@ -6,8 +9,8 @@ import torch from torch import nn -from . import functional as F, functional as pykeen_functional -from .functional import hole_interaction +from . import functional as pykeen_functional +from .. import typing as pykeen_typing from ..utils import check_shapes logger = logging.getLogger(__name__) @@ -255,14 +258,17 @@ def forward( t: torch.FloatTensor, **kwargs, ) -> torch.FloatTensor: # noqa:D102 - return F.translational_interaction( + return pykeen_functional.translational_interaction( h=h, r=r, t=t, p=self.p, dim=kwargs.get('dim', None), keepdim=kwargs.get('keepdim', False), ) -class ComplExInteractionFunction(InteractionFunction): - """Interaction function of ComplEx.""" +class _SimpleInteractionFunction(InteractionFunction): + """An intermediate base interaction function for simple stateless functions that only take h,r,t.""" + + #: A reference to the stateless interaction function + interaction_function: pykeen_typing.InteractionFunction def forward( self, @@ -272,7 +278,13 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa: D102 self._check_for_empty_kwargs(kwargs) - return F.complex_interaction(h=h, r=r, t=t) + return self.interaction_function(h, r, t) + + +class ComplExInteractionFunction(_SimpleInteractionFunction): + """Interaction function of ComplEx.""" + + interaction_function = pykeen_functional.complex_interaction def _calculate_missing_shape_information( @@ -416,7 +428,7 @@ def forward( raise TypeError(f"{self.__class__.__name__}.forward expects keyword argument 't_bias'.") t_bias: torch.FloatTensor = kwargs.pop("t_bias") self._check_for_empty_kwargs(kwargs) - return F.conve_interaction( + return pykeen_functional.conve_interaction( h=h, r=r, t=t, @@ -474,7 +486,7 @@ def forward( t: torch.FloatTensor, **kwargs, ) -> torch.FloatTensor: # noqa: D102 - return F.convkb_interaction( + return pykeen_functional.convkb_interaction( h=h, r=r, t=t, @@ -485,18 +497,10 @@ def forward( ) -class DistMultInteractionFunction(InteractionFunction): +class DistMultInteractionFunction(_SimpleInteractionFunction): """Interaction function of DistMult.""" - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return F.distmult_interaction(h=h, r=r, t=t) + interaction_function = pykeen_functional.distmult_interaction class ERMLPInteractionFunction(InteractionFunction): @@ -537,7 +541,7 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa: D102 self._check_for_empty_kwargs(kwargs) - return F.ermlp_interaction( + return pykeen_functional.ermlp_interaction( h=h, r=r, t=t, @@ -589,7 +593,7 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa: D102 self._check_for_empty_kwargs(kwargs=kwargs) - return F.ermlpe_interaction(h=h, r=r, t=t, mlp=self.mlp) + return pykeen_functional.ermlpe_interaction(h=h, r=r, t=t, mlp=self.mlp) class TransDInteractionFunction(TranslationalInteractionFunction): @@ -636,29 +640,13 @@ def forward( return pykeen_functional.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p) -class RotatEInteraction(InteractionFunction): +class RotatEInteraction(_SimpleInteractionFunction): """Interaction function of RotatE.""" - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return pykeen_functional.rotate_interaction(h=h, r=r, t=t) + interaction_function = pykeen_functional.rotate_interaction -class HolEInteractionFunction(InteractionFunction): +class HolEInteractionFunction(_SimpleInteractionFunction): """Interaction function for HolE.""" - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return hole_interaction(h=h, r=r, t=t) \ No newline at end of file + interaction_function = pykeen_functional.hole_interaction From 228077a39b6d1c6da06feae58a0a8aa9601b153e Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 15:17:18 +0100 Subject: [PATCH 113/690] Update trans_e.py --- src/pykeen/models/unimodal/trans_e.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 05617758a7..682a6d5ef1 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -67,7 +67,6 @@ def __init__( - OpenKE `implementation of TransE `_ """ - self.interaction_function = TranslationalInteractionFunction(p=scoring_fct_norm) super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, @@ -83,6 +82,7 @@ def __init__( ), entity_constrainer=functional.normalize, ) + self.interaction_function = TranslationalInteractionFunction(p=scoring_fct_norm) def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings From cad15800847e4028a22cac6c566c272180b1a0fc Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 15:19:52 +0100 Subject: [PATCH 114/690] Update trans_r.py --- src/pykeen/models/unimodal/trans_r.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 5cdb5a0bb7..2cabc5c4f0 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -86,8 +86,6 @@ def __init__( regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model.""" - self.interaction_function = TransRInteractionFunction(p=scoring_fct_norm) - super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, @@ -107,6 +105,7 @@ def __init__( relation_constrainer=clamp_norm, relation_constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ) + self.interaction_function = TransRInteractionFunction(p=scoring_fct_norm) # TODO: Initialize from TransE From dbae2030fd321116cb7d68377272dcb82749e46e Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 15:25:43 +0100 Subject: [PATCH 115/690] Update modules.py --- src/pykeen/nn/modules.py | 57 +++++++++++++++++++++++++--------------- 1 file changed, 36 insertions(+), 21 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index d261f3ce3c..e5cc566070 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -10,7 +10,6 @@ from torch import nn from . import functional as pykeen_functional -from .. import typing as pykeen_typing from ..utils import check_shapes logger = logging.getLogger(__name__) @@ -48,10 +47,11 @@ def forward( """ raise NotImplementedError - def _check_for_empty_kwargs(self, kwargs: Mapping[str, Any]) -> None: + @classmethod + def _check_for_empty_kwargs(cls, kwargs: Mapping[str, Any]) -> None: """Check that kwargs is empty.""" if len(kwargs) > 0: - raise ValueError(f"{self.__class__.__name__} does not take the following kwargs: {kwargs}") + raise ValueError(f"{cls.__name__} does not take the following kwargs: {kwargs}") @staticmethod def _add_dim(*x: torch.FloatTensor, dim: int) -> Sequence[torch.FloatTensor]: @@ -264,11 +264,8 @@ def forward( ) -class _SimpleInteractionFunction(InteractionFunction): - """An intermediate base interaction function for simple stateless functions that only take h,r,t.""" - - #: A reference to the stateless interaction function - interaction_function: pykeen_typing.InteractionFunction +class ComplExInteractionFunction(InteractionFunction): + """Interaction function of ComplEx.""" def forward( self, @@ -278,13 +275,7 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa: D102 self._check_for_empty_kwargs(kwargs) - return self.interaction_function(h, r, t) - - -class ComplExInteractionFunction(_SimpleInteractionFunction): - """Interaction function of ComplEx.""" - - interaction_function = pykeen_functional.complex_interaction + return pykeen_functional.complex_interaction(h=h, r=r, t=t) def _calculate_missing_shape_information( @@ -497,10 +488,18 @@ def forward( ) -class DistMultInteractionFunction(_SimpleInteractionFunction): +class DistMultInteractionFunction(InteractionFunction): """Interaction function of DistMult.""" - interaction_function = pykeen_functional.distmult_interaction + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + self._check_for_empty_kwargs(kwargs) + return pykeen_functional.distmult_interaction(h=h, r=r, t=t) class ERMLPInteractionFunction(InteractionFunction): @@ -640,13 +639,29 @@ def forward( return pykeen_functional.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p) -class RotatEInteraction(_SimpleInteractionFunction): +class RotatEInteraction(InteractionFunction): """Interaction function of RotatE.""" - interaction_function = pykeen_functional.rotate_interaction + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + self._check_for_empty_kwargs(kwargs) + return pykeen_functional.rotate_interaction(h=h, r=r, t=t) -class HolEInteractionFunction(_SimpleInteractionFunction): +class HolEInteractionFunction(InteractionFunction): """Interaction function for HolE.""" - interaction_function = pykeen_functional.hole_interaction + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + self._check_for_empty_kwargs(kwargs) + return pykeen_functional.hole_interaction(h=h, r=r, t=t) From a3d085e610db297db59ee2ff00d252fff37c2f2f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 15:49:05 +0100 Subject: [PATCH 116/690] Extract functional form of KG2E --- src/pykeen/models/unimodal/kg2e.py | 137 ++++------------------- src/pykeen/nn/functional.py | 174 ++++++++++++++++++++++++++++- 2 files changed, 193 insertions(+), 118 deletions(-) diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 846f887acc..d612d078e6 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -10,7 +10,8 @@ from ..base import EntityRelationEmbeddingModel from ...losses import Loss -from ...nn import Embedding +from ...nn import Embedding, functional as pykeen_functional +from ...nn.functional import KG2E_SIMILARITIES from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -92,12 +93,10 @@ def __init__( ) # Similarity function used for distributions - if dist_similarity is None or dist_similarity.upper() == 'KL': - self.similarity = self.kullback_leibler_similarity - elif dist_similarity.upper() == 'EL': - self.similarity = self.expected_likelihood - else: - raise ValueError(f'Unknown distribution similarity: "{dist_similarity}".') + dist_similarity = dist_similarity.upper() + if dist_similarity not in KG2E_SIMILARITIES: + raise ValueError(dist_similarity) + self.similarity = dist_similarity # element-wise covariance bounds self.c_min = c_min @@ -163,118 +162,24 @@ def _score( sigma_t = self.entity_covariances.get_in_canonical_shape(indices=t_indices) # Compute entity distribution - mu_e = mu_h - mu_t - sigma_e = sigma_h + sigma_t - return self.similarity(mu_e=mu_e, mu_r=mu_r, sigma_e=sigma_e, sigma_r=sigma_r) + return pykeen_functional.kg2e_interaction( + h_mean=mu_h, + h_var=sigma_h, + r_mean=mu_r, + r_var=sigma_r, + t_mean=mu_t, + t_var=sigma_t, + ) def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]).view(-1, 1) + return self._score( + h_indices=hrt_batch[:, 0], + r_indices=hrt_batch[:, 1], + t_indices=hrt_batch[:, 2], + ).view(hrt_batch.shape[0], 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1]) + return self._score(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1]).view(hr_batch.shape[0], self.num_entities) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]) - - @staticmethod - def expected_likelihood( - mu_e: torch.FloatTensor, - mu_r: torch.FloatTensor, - sigma_e: torch.FloatTensor, - sigma_r: torch.FloatTensor, - epsilon: float = 1.0e-10, - ) -> torch.FloatTensor: - r"""Compute the similarity based on expected likelihood. - - .. math:: - - D((\mu_e, \Sigma_e), (\mu_r, \Sigma_r))) - = \frac{1}{2} \left( - (\mu_e - \mu_r)^T(\Sigma_e + \Sigma_r)^{-1}(\mu_e - \mu_r) - + \log \det (\Sigma_e + \Sigma_r) + d \log (2 \pi) - \right) - = \frac{1}{2} \left( - \mu^T\Sigma^{-1}\mu - + \log \det \Sigma + d \log (2 \pi) - \right) - - :param mu_e: torch.Tensor, shape: (s_1, ..., s_k, d) - The mean of the first Gaussian. - :param mu_r: torch.Tensor, shape: (s_1, ..., s_k, d) - The mean of the second Gaussian. - :param sigma_e: torch.Tensor, shape: (s_1, ..., s_k, d) - The diagonal covariance matrix of the first Gaussian. - :param sigma_r: torch.Tensor, shape: (s_1, ..., s_k, d) - The diagonal covariance matrix of the second Gaussian. - :param epsilon: float (default=1.0) - Small constant used to avoid numerical issues when dividing. - - :return: torch.Tensor, shape: (s_1, ..., s_k) - The similarity. - """ - d = sigma_e.shape[-1] - sigma = sigma_r + sigma_e - mu = mu_e - mu_r - - #: a = \mu^T\Sigma^{-1}\mu - safe_sigma = torch.clamp_min(sigma, min=epsilon) - sigma_inv = torch.reciprocal(safe_sigma) - a = torch.sum(sigma_inv * mu ** 2, dim=-1) - - #: b = \log \det \Sigma - b = safe_sigma.log().sum(dim=-1) - return a + b + d * _LOG_2_PI - - @staticmethod - def kullback_leibler_similarity( - mu_e: torch.FloatTensor, - mu_r: torch.FloatTensor, - sigma_e: torch.FloatTensor, - sigma_r: torch.FloatTensor, - epsilon: float = 1.0e-10, - ) -> torch.FloatTensor: - r"""Compute the similarity based on KL divergence. - - This is done between two Gaussian distributions given by mean mu_* and diagonal covariance matrix sigma_*. - - .. math:: - - D((\mu_e, \Sigma_e), (\mu_r, \Sigma_r))) - = \frac{1}{2} \left( - tr(\Sigma_r^{-1}\Sigma_e) - + (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) - - \log \frac{det(\Sigma_e)}{det(\Sigma_r)} - k_e - \right) - - Note: The sign of the function has been flipped as opposed to the description in the paper, as the - Kullback Leibler divergence is large if the distributions are dissimilar. - - :param mu_e: torch.Tensor, shape: (s_1, ..., s_k, d) - The mean of the first Gaussian. - :param mu_r: torch.Tensor, shape: (s_1, ..., s_k, d) - The mean of the second Gaussian. - :param sigma_e: torch.Tensor, shape: (s_1, ..., s_k, d) - The diagonal covariance matrix of the first Gaussian. - :param sigma_r: torch.Tensor, shape: (s_1, ..., s_k, d) - The diagonal covariance matrix of the second Gaussian. - :param epsilon: float (default=1.0) - Small constant used to avoid numerical issues when dividing. - - :return: torch.Tensor, shape: (s_1, ..., s_k) - The similarity. - """ - d = mu_e.shape[-1] - safe_sigma_r = torch.clamp_min(sigma_r, min=epsilon) - sigma_r_inv = torch.reciprocal(safe_sigma_r) - - #: a = tr(\Sigma_r^{-1}\Sigma_e) - a = torch.sum(sigma_e * sigma_r_inv, dim=-1) - - #: b = (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) - mu = mu_r - mu_e - b = torch.sum(sigma_r_inv * mu ** 2, dim=-1) - - #: c = \log \frac{det(\Sigma_e)}{det(\Sigma_r)} - # = sum log (sigma_e)_i - sum log (sigma_r)_i - c = sigma_e.clamp_min(min=epsilon).log().sum(dim=-1) - safe_sigma_r.log().sum(dim=-1) - return -0.5 * (a + b - c - d) + return self._score(r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index e61bab510d..96b95aafdd 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -1,8 +1,8 @@ # -*- coding: utf-8 -*- """Functional forms of interaction methods.""" - -from typing import Optional, Tuple, Union +import math +from typing import NamedTuple, Optional, Tuple, Union import torch from torch import nn @@ -465,3 +465,173 @@ def transr_interaction( # evaluate score function, shape: (b, e) return translational_interaction(h=h_bot, r=r, t=t_bot, dim=-1, p=p) ** 2 + + +class GaussianDistribution(NamedTuple): + """A gaussian distribution with diagonal covariance matrix.""" + mean: torch.FloatTensor + diagonal_covariance: torch.FloatTensor + + +def _expected_likelihood( + e: GaussianDistribution, + r: GaussianDistribution, + epsilon: float = 1.0e-10, + exact: bool = True, +) -> torch.FloatTensor: + r"""Compute the similarity based on expected likelihood. + + .. math:: + + D((\mu_e, \Sigma_e), (\mu_r, \Sigma_r))) + = \frac{1}{2} \left( + (\mu_e - \mu_r)^T(\Sigma_e + \Sigma_r)^{-1}(\mu_e - \mu_r) + + \log \det (\Sigma_e + \Sigma_r) + d \log (2 \pi) + \right) + = \frac{1}{2} \left( + \mu^T\Sigma^{-1}\mu + + \log \det \Sigma + d \log (2 \pi) + \right) + + :param e: shape: (batch_size, num_heads, num_tails, d) + The entity Gaussian distribution. + :param r: shape: (batch_size, num_relations, d) + The relation Gaussian distribution. + :param epsilon: float (default=1.0) + Small constant used to avoid numerical issues when dividing. + :param exact: + Whether to return the exact similarity, or leave out constant offsets. + + :return: torch.Tensor, shape: (s_1, ..., s_k) + The similarity. + """ + # subtract, shape: (batch_size, num_heads, num_relations, num_tails, dim) + r_shape = r.mean.shape + r_shape = (r_shape[0], 1, r_shape[1], 1, r_shape[2]) + var = r.diagonal_covariance.view(*r_shape) + e.diagonal_covariance.unsqueeze(dim=2) + mean = e.mean.unsqueeze(dim=2) - r.mean.view(*r_shape) + + #: a = \mu^T\Sigma^{-1}\mu + safe_sigma = torch.clamp_min(var, min=epsilon) + sigma_inv = torch.reciprocal(safe_sigma) + sim = torch.sum(sigma_inv * mean ** 2, dim=-1) + + #: b = \log \det \Sigma + sim = sim + safe_sigma.log().sum(dim=-1) + if exact: + sim = sim + sim.shape[-1] * math.log(2. * math.pi) + return sim + + +def _kullback_leibler_similarity( + e: GaussianDistribution, + r: GaussianDistribution, + epsilon: float = 1.0e-10, + exact: bool = True, +) -> torch.FloatTensor: + r"""Compute the similarity based on KL divergence. + + This is done between two Gaussian distributions given by mean mu_* and diagonal covariance matrix sigma_*. + + .. math:: + + D((\mu_e, \Sigma_e), (\mu_r, \Sigma_r))) + = \frac{1}{2} \left( + tr(\Sigma_r^{-1}\Sigma_e) + + (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) + - \log \frac{det(\Sigma_e)}{det(\Sigma_r)} - k_e + \right) + + Note: The sign of the function has been flipped as opposed to the description in the paper, as the + Kullback Leibler divergence is large if the distributions are dissimilar. + + :param e: shape: (batch_size, num_heads, num_tails, d) + The entity Gaussian distributions, as mean/diagonal covariance pairs. + :param r: shape: (batch_size, num_relations, d) + The relation Gaussian distributions, as mean/diagonal covariance pairs. + :param epsilon: float (default=1.0) + Small constant used to avoid numerical issues when dividing. + :param exact: + Whether to return the exact similarity, or leave out constant offsets. + + :return: torch.Tensor, shape: (s_1, ..., s_k) + The similarity. + """ + # invert covariance, shape: (batch_size, num_relations, d) + safe_sigma_r = torch.clamp_min(r.diagonal_covariance, min=epsilon) + sigma_r_inv = torch.reciprocal(safe_sigma_r) + + #: a = tr(\Sigma_r^{-1}\Sigma_e), (batch_size, num_heads, num_relations, num_tails) + # [(b, h, t, d), (b, r, d) -> (b, 1, r, d) -> (b, 1, d, r)] -> (b, h, t, r) -> (b, h, r, t) + sim = (e.diagonal_covariance @ sigma_r_inv.unsqueeze(dim=1).transpose(-2, -1)).transpose(-2, -1) + + #: b = (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) + r_shape = r.mean.shape + # mu.shape: (b, h, r, t, d) + mu = r.mean.view(r_shape[0], 1, r_shape[1], 1, r_shape[2]) - e.mean.unsqueeze(dim=2) + sim = sim + (mu ** 2 @ sigma_r_inv.view(r_shape[0], 1, r_shape[1], r_shape[2], 1)).squeeze(dim=-1) + + #: c = \log \frac{det(\Sigma_e)}{det(\Sigma_r)} + # = sum log (sigma_e)_i - sum log (sigma_r)_i + # ce.shape: (b, h, t) + ce = e.diagonal_covariance.clamp_min(min=epsilon).log().sum(dim=-1) + # cr.shape: (b, r) + cr = safe_sigma_r.log().sum(dim=-1) + sim = sim + ce.unsqueeze(dim=2) - cr.view(r_shape[0], 1, r_shape[1], 1) + + if exact: + sim = sim - e.mean.shape[-1] + sim = 0.5 * sim + + return sim + + +_KG2E_SIMILARITIES = dict( + KL=_kullback_leibler_similarity, + EL=_expected_likelihood, +) +KG2E_SIMILARITIES = set(_KG2E_SIMILARITIES.keys()) + + +def kg2e_interaction( + h_mean: torch.FloatTensor, + h_var: torch.FloatTensor, + r_mean: torch.FloatTensor, + r_var: torch.FloatTensor, + t_mean: torch.FloatTensor, + t_var: torch.FloatTensor, + similarity: str = "KL", + exact: bool = True, +) -> torch.FloatTensor: + """ + Evaluate the KG2E interaction function. + + :param h_mean: shape: (batch_size, num_heads, d) + The head entity distribution mean. + :param h_var: shape: (batch_size, num_heads, d) + The head entity distribution variance. + :param r_mean: shape: (batch_size, num_relations, d) + The relation distribution mean. + :param r_var: shape: (batch_size, num_relations, d) + The relation distribution variance. + :param t_mean: shape: (batch_size, num_tails, d) + The tail entity distribution mean. + :param t_var: shape: (batch_size, num_tails, d) + The tail entity distribution variance. + :param similarity: + The similarity measures for gaussian distributions. From {"KL", "EL"}. + :param exact: + Whether to leave out constants to accelerate similarity computation. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + if similarity not in KG2E_SIMILARITIES: + raise ValueError(similarity) + similarity = _KG2E_SIMILARITIES[similarity] + # Compute entity distribution + e_mean = h_mean.unsqueeze(dim=2) - t_mean.unsqueeze(dim=1) + e_var = h_var.unsqueeze(dim=2) + t_var.unsqueeze(dim=1) + e = GaussianDistribution(mean=e_mean, diagonal_covariance=e_var) + r = GaussianDistribution(mean=r_mean, diagonal_covariance=r_var) + return similarity(e=e, r=r, exact=exact) From a1ffd091ca56b8ec1beea448a04963bd1e397ab4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 15:49:49 +0100 Subject: [PATCH 117/690] Do not store unneeded fields --- src/pykeen/models/unimodal/kg2e.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index d612d078e6..0cf8d2eb17 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -98,10 +98,6 @@ def __init__( raise ValueError(dist_similarity) self.similarity = dist_similarity - # element-wise covariance bounds - self.c_min = c_min - self.c_max = c_max - # Additional covariance embeddings self.entity_covariances = Embedding.init_with_device( num_embeddings=triples_factory.num_entities, @@ -109,7 +105,7 @@ def __init__( device=self.device, # Ensure positive definite covariances matrices and appropriate size by clamping constrainer=torch.clamp, - constrainer_kwargs=dict(min=self.c_min, max=self.c_max), + constrainer_kwargs=dict(min=c_min, max=c_max), ) self.relation_covariances = Embedding.init_with_device( num_embeddings=triples_factory.num_relations, @@ -117,7 +113,7 @@ def __init__( device=self.device, # Ensure positive definite covariances matrices and appropriate size by clamping constrainer=torch.clamp, - constrainer_kwargs=dict(min=self.c_min, max=self.c_max), + constrainer_kwargs=dict(min=c_min, max=c_max), ) def _reset_parameters_(self): # noqa: D102 From a7c84fb08cce5946e5ae498a55af27297027374f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 15:51:48 +0100 Subject: [PATCH 118/690] Implement score_t for KG2E --- src/pykeen/models/unimodal/kg2e.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 0cf8d2eb17..28663624b5 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -177,5 +177,8 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self._score(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1]).view(hr_batch.shape[0], self.num_entities) + def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self._score(h_indices=ht_batch[:, 0], t_indices=ht_batch[:, 1]).view(ht_batch.shape[0], self.num_entities) + def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self._score(r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]).view(rt_batch.shape[0], self.num_entities) From 2576ce857d2f18fde9ec68f24d895d5efbabb3f2 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 16:25:19 +0100 Subject: [PATCH 119/690] Update base.py --- src/pykeen/models/base.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index f186d0f4ab..227e683356 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1361,13 +1361,12 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], t_indices=None).view(-1, self.num_entities) + def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + return self(h_indices=ht_batch[:, 0], r_indices=None, t_indices=ht_batch[:, 1]) + def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self(h_indices=None, r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]).view(-1, self.num_entities) - # TODO - # def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # raise NotImplementedError - class SimpleVectorEntityRelationEmbeddingModel( GeneralVectorEntityRelationEmbeddingModel, reset_parameters_post_init=False, From 250809fbeb2001c13f23220587ba81f6721221c9 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 16:47:07 +0100 Subject: [PATCH 120/690] Update rotate.py --- src/pykeen/models/unimodal/rotate.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index 84e81457eb..7f78a44d19 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -92,7 +92,6 @@ def __init__( random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: - self.interaction_function = RotatEInteraction() super().__init__( triples_factory=triples_factory, embedding_dim=2 * embedding_dim, @@ -105,6 +104,7 @@ def __init__( relation_initializer=init_phases, relation_constrainer=complex_normalize, ) + self.interaction_function = RotatEInteraction() self.real_embedding_dim = embedding_dim def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 From da968ac89006969573adb80e3c45412dd68628f7 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 16:47:49 +0100 Subject: [PATCH 121/690] Snag parameters from partial --- tests/test_models.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index 876a4229d0..2bfbfc5210 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -669,10 +669,13 @@ def _check_constraints(self): * Entity and relation embeddings have to have at most unit L2 norm. * Covariances have to have values between c_min and c_max """ + low = self.model.entity_covariances.constrainer.keywords['min'] + high = self.model.entity_covariances.constrainer.keywords['max'] + for embedding in (self.model.entity_embeddings, self.model.relation_embeddings): assert all_in_bounds(embedding(indices=None).norm(p=2, dim=-1), high=1., a_tol=_EPSILON) for cov in (self.model.entity_covariances, self.model.relation_covariances): - assert all_in_bounds(cov(indices=None), low=self.model.c_min, high=self.model.c_max) + assert all_in_bounds(cov(indices=None), low=low, high=high) class TestKG2EWithKL(_TestKG2E, unittest.TestCase): From b1faffaf6591ef35aa68e682930cfb811c8f108e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 17:00:25 +0100 Subject: [PATCH 122/690] Extract functional form of NTN --- src/pykeen/models/unimodal/ntn.py | 150 +++++++++--------------------- src/pykeen/nn/emb.py | 16 +++- src/pykeen/nn/functional.py | 78 ++++++++++++++++ tests/test_models.py | 2 +- 4 files changed, 133 insertions(+), 113 deletions(-) diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 30c28d3b75..b31ade79f4 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -9,6 +9,7 @@ from ..base import EntityEmbeddingModel from ...losses import Loss +from ...nn import Embedding, functional as F from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -81,46 +82,36 @@ def __init__( ) self.num_slices = num_slices - self.w = nn.Parameter(data=torch.empty( - triples_factory.num_relations, - num_slices, - embedding_dim, - embedding_dim, - device=self.device, - ), requires_grad=True) - self.vh = nn.Parameter(data=torch.empty( - triples_factory.num_relations, - num_slices, - embedding_dim, - device=self.device, - ), requires_grad=True) - self.vt = nn.Parameter(data=torch.empty( - triples_factory.num_relations, - num_slices, - embedding_dim, - device=self.device, - ), requires_grad=True) - self.b = nn.Parameter(data=torch.empty( - triples_factory.num_relations, - num_slices, - device=self.device, - ), requires_grad=True) - self.u = nn.Parameter(data=torch.empty( - triples_factory.num_relations, - num_slices, - device=self.device, - ), requires_grad=True) + self.w = Embedding( + num_embeddings=triples_factory.num_relations, + embedding_dim=num_slices * self.embedding_dim ** 2, + ) + self.vh = Embedding( + num_embeddings=triples_factory.num_relations, + embedding_dim=num_slices * embedding_dim, + ) + self.vt = Embedding( + num_embeddings=triples_factory.num_relations, + embedding_dim=num_slices * embedding_dim, + ) + self.b = Embedding( + num_embeddings=triples_factory.num_relations, + embedding_dim=num_slices, + ) + self.u = Embedding( + num_embeddings=triples_factory.num_relations, + embedding_dim=num_slices, + ) if non_linearity is None: non_linearity = nn.Tanh() self.non_linearity = non_linearity def _reset_parameters_(self): # noqa: D102 - super()._reset_parameters_() - nn.init.normal_(self.w) - nn.init.normal_(self.vh) - nn.init.normal_(self.vt) - nn.init.normal_(self.b) - nn.init.normal_(self.u) + for module in self.modules(): + if module is self: + continue + if hasattr(module, "reset_parameters"): + module.reset_parameters() def _score( self, @@ -136,17 +127,23 @@ def _score( :param r_indices: shape: (batch_size,) :param t_indices: shape: (batch_size,) - :return: shape: (batch_size, num_entities) + :return: shape: (batch_size, num_heads, num_relations, num_tails) """ - assert r_indices is not None + assert slice_size is None, "not implemented" #: shape: (batch_size, num_entities, d) h_all = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) t_all = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) + w = self.w.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim, self.embedding_dim)) + b = self.b.get_in_canonical_shape(indices=r_indices) + u = self.u.get_in_canonical_shape(indices=r_indices) + vh = self.vh.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim)) + vt = self.vt.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim)) if slice_size is None: - return self._interaction_function(h=h_all, t=t_all, r_indices=r_indices) + return F.ntn_interaction(h=h_all, t=t_all, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity) + # TODO: Not implemented if h_all.shape[1] > t_all.shape[1]: h_was_split = True split_tensor = torch.split(h_all, slice_size, dim=1) @@ -164,82 +161,19 @@ def _score( else: h = constant_tensor t = split - score = self._interaction_function(h=h, t=t, r_indices=r_indices) + score = F.ntn_interaction(h=h_all, t=t_all, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity) scores_arr.append(score) return torch.cat(scores_arr, dim=1) - def _interaction_function( - self, - h: torch.FloatTensor, - t: torch.FloatTensor, - r_indices: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: - #: Prepare h: (b, e, d) -> (b, e, 1, 1, d) - h_for_w = h.unsqueeze(dim=-2).unsqueeze(dim=-2) - - #: Prepare t: (b, e, d) -> (b, e, 1, d, 1) - t_for_w = t.unsqueeze(dim=-2).unsqueeze(dim=-1) - - #: Prepare w: (R, k, d, d) -> (b, k, d, d) -> (b, 1, k, d, d) - w_r = self.w.index_select(dim=0, index=r_indices).unsqueeze(dim=1) - - # h.T @ W @ t, shape: (b, e, k, 1, 1) - hwt = (h_for_w @ w_r @ t_for_w) - - #: reduce (b, e, k, 1, 1) -> (b, e, k) - hwt = hwt.squeeze(dim=-1).squeeze(dim=-1) - - #: Prepare vh: (R, k, d) -> (b, k, d) -> (b, 1, k, d) - vh_r = self.vh.index_select(dim=0, index=r_indices).unsqueeze(dim=1) - - #: Prepare h: (b, e, d) -> (b, e, d, 1) - h_for_v = h.unsqueeze(dim=-1) - - # V_h @ h, shape: (b, e, k, 1) - vhh = vh_r @ h_for_v - - #: reduce (b, e, k, 1) -> (b, e, k) - vhh = vhh.squeeze(dim=-1) - - #: Prepare vt: (R, k, d) -> (b, k, d) -> (b, 1, k, d) - vt_r = self.vt.index_select(dim=0, index=r_indices).unsqueeze(dim=1) - - #: Prepare t: (b, e, d) -> (b, e, d, 1) - t_for_v = t.unsqueeze(dim=-1) - - # V_t @ t, shape: (b, e, k, 1) - vtt = vt_r @ t_for_v - - #: reduce (b, e, k, 1) -> (b, e, k) - vtt = vtt.squeeze(dim=-1) - - #: Prepare b: (R, k) -> (b, k) -> (b, 1, k) - b = self.b.index_select(dim=0, index=r_indices).unsqueeze(dim=1) - - # a = f(h.T @ W @ t + Vh @ h + Vt @ t + b), shape: (b, e, k) - pre_act = hwt + vhh + vtt + b - act = self.non_linearity(pre_act) - - # prepare u: (R, k) -> (b, k) -> (b, 1, k, 1) - u = self.u.index_select(dim=0, index=r_indices).unsqueeze(dim=1).unsqueeze(dim=-1) - - # prepare act: (b, e, k) -> (b, e, 1, k) - act = act.unsqueeze(dim=-2) - - # compute score, shape: (b, e, 1, 1) - score = act @ u - - # reduce - score = score.squeeze(dim=-1).squeeze(dim=-1) - - return score - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]) + return self._score(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]).view(hrt_batch.shape[0], 1) def score_t(self, hr_batch: torch.LongTensor, slice_size: int = None) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], slice_size=slice_size) + return self._score(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], slice_size=slice_size).view(hr_batch.shape[0], self.num_entities) + + def score_r(self, ht_batch: torch.LongTensor, slice_size: int = None) -> torch.FloatTensor: # noqa: D102 + return self._score(h_indices=ht_batch[:, 0], t_indices=ht_batch[:, 1], slice_size=slice_size).view(ht_batch.shape[0], self.num_relations) def score_h(self, rt_batch: torch.LongTensor, slice_size: int = None) -> torch.FloatTensor: # noqa: D102 - return self._score(r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1], slice_size=slice_size) + return self._score(r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1], slice_size=slice_size).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index f7c5862489..44446cc11e 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -3,7 +3,7 @@ """Embedding modules.""" import functools -from typing import Any, Mapping, Optional +from typing import Any, Mapping, Optional, Sequence import torch import torch.nn @@ -175,14 +175,22 @@ def forward( def get_in_canonical_shape( self, indices: Optional[torch.LongTensor] = None, + reshape_dim: Optional[Sequence[int]] = None, ) -> torch.FloatTensor: """Get embedding in canonical shape. - :param indices: The indices. If None, return all embeddings. + :param indices: + The indices. If None, return all embeddings. + :param reshape_dim: + Optionally reshape the last dimension. :return: shape: (batch_size, num_embeddings, d) """ x = self(indices=indices) if indices is None: - return x.unsqueeze(dim=0) - return x.unsqueeze(dim=1) + x = x.unsqueeze(dim=0) + else: + x = x.unsqueeze(dim=1) + if reshape_dim is not None: + x = x.view(*x.shape[:-1], *reshape_dim) + return x diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 96b95aafdd..26d9073a6b 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -635,3 +635,81 @@ def kg2e_interaction( e = GaussianDistribution(mean=e_mean, diagonal_covariance=e_var) r = GaussianDistribution(mean=r_mean, diagonal_covariance=r_var) return similarity(e=e, r=r, exact=exact) + + +def _extended_einsum( + eq: str, + *tensors, +) -> torch.FloatTensor: + """Drop dimensions of size 1 to allow broadcasting.""" + lhs, rhs = eq.split("->") + mod_ops, mod_t = [], [] + for op, t in zip(lhs.split(","), tensors): + mod_op = "" + assert len(op) == len(t.shape) + for i, c in reversed(list(enumerate(op))): + if t.shape[i] == 1: + t = t.squeeze(dim=i) + else: + mod_op = c + mod_op + mod_ops.append(mod_op) + mod_t.append(t) + m_lhs = ",".join(mod_ops) + r_keep_dims = set("".join(mod_ops)) + m_rhs = "".join(c for c in rhs if c in r_keep_dims) + m_eq = f"{m_lhs}->{m_rhs}" + mod_r = torch.einsum(m_eq, *mod_t) + # unsqueeze + for i, c in enumerate(rhs): + if c not in r_keep_dims: + mod_r = mod_r.unsqueeze(dim=i) + return mod_r + + +def ntn_interaction( + h: torch.FloatTensor, + t: torch.FloatTensor, + w: torch.FloatTensor, + b: torch.FloatTensor, + u: torch.FloatTensor, + vh: torch.FloatTensor, + vt: torch.FloatTensor, + activation: nn.Module, +) -> torch.FloatTensor: + r""" + Evaluate the NTN interaction function. + + .. math:: + + f(h,r,t) = u_r^T act(h W_r t + V_r h + V_r' t + b_r) + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param vh: shape: (batch_size, num_relations, k, dim) + The head transformation matrix V_h. + :param vt: shape: (batch_size, num_relations, k, dim) + The tail transformation matrix V_h. + :param w: shape: (batch_size, num_relations, k, dim, dim) + The relation specific transformation matrix W_r. + :param b: shape: (batch_size, num_relations, k) + The relation specific offset b_r. + :param u: shape: (batch_size, num_relations, k) + The relation specific final linear transformation b_r. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + :param activation: + The activation function. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + # TODO: check efficiency of einsum + # save sizes + num_heads, num_relations, num_tails = [xx.shape[1] for xx in (h, b, t)] + k = b.shape[-1] + x = _extended_einsum("bhd,brkde,bte->bhrtk", h, w, t) + x = x + _extended_einsum("brkd,bhd->bhk", vh, h).view(-1, num_heads, 1, 1, k) + x = x + _extended_einsum("brkd,btd->btk", vt, t).view(-1, 1, 1, num_tails, k) + x = activation(x) + x = _extended_einsum("bhrtk,brk->bhrt", x, u) + return x diff --git a/tests/test_models.py b/tests/test_models.py index 876a4229d0..9dea615176 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -699,7 +699,7 @@ class _BaseNTNTest(_ModelTestCase, unittest.TestCase): def test_can_slice(self): """Test that the slicing properties are calculated correctly.""" self.assertTrue(self.model.can_slice_h) - self.assertFalse(self.model.can_slice_r) + self.assertTrue(self.model.can_slice_r) self.assertTrue(self.model.can_slice_t) From dc7c839ffd4c3659a82d48e2a7b2c2d6c26190c0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 17:00:39 +0100 Subject: [PATCH 123/690] Fix typo --- src/pykeen/models/unimodal/ntn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index b31ade79f4..5666b332ae 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -161,7 +161,7 @@ def _score( else: h = constant_tensor t = split - score = F.ntn_interaction(h=h_all, t=t_all, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity) + score = F.ntn_interaction(h=h, t=t, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity) scores_arr.append(score) return torch.cat(scores_arr, dim=1) From a5c2b89f37ab4676e93d57f16fb3c639fd0adb1d Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 17:02:06 +0100 Subject: [PATCH 124/690] Update code style --- src/pykeen/models/unimodal/kg2e.py | 8 ++++---- src/pykeen/nn/functional.py | 1 + 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 28663624b5..8329e3ef40 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -133,7 +133,7 @@ def post_parameter_update(self) -> None: # noqa: D102 ): cov.post_parameter_update() - def _score( + def forward( self, h_indices: Optional[torch.LongTensor] = None, r_indices: Optional[torch.LongTensor] = None, @@ -168,17 +168,17 @@ def _score( ) def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score( + return self( h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2], ).view(hrt_batch.shape[0], 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1]).view(hr_batch.shape[0], self.num_entities) + return self(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1]).view(hr_batch.shape[0], self.num_entities) def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self._score(h_indices=ht_batch[:, 0], t_indices=ht_batch[:, 1]).view(ht_batch.shape[0], self.num_entities) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]).view(rt_batch.shape[0], self.num_entities) + return self(r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 26d9073a6b..47a8a0431b 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -469,6 +469,7 @@ def transr_interaction( class GaussianDistribution(NamedTuple): """A gaussian distribution with diagonal covariance matrix.""" + mean: torch.FloatTensor diagonal_covariance: torch.FloatTensor From c0270491e94b45fe58a81f97685676b9b289f294 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 17:04:53 +0100 Subject: [PATCH 125/690] Fix scoring and style --- src/pykeen/models/unimodal/kg2e.py | 2 +- src/pykeen/models/unimodal/ntn.py | 36 +++++++++++++++++++++------- src/pykeen/models/unimodal/simple.py | 8 +++---- 3 files changed, 33 insertions(+), 13 deletions(-) diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 8329e3ef40..639ba356ef 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -178,7 +178,7 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 return self(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1]).view(hr_batch.shape[0], self.num_entities) def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=ht_batch[:, 0], t_indices=ht_batch[:, 1]).view(ht_batch.shape[0], self.num_entities) + return self(h_indices=ht_batch[:, 0], t_indices=ht_batch[:, 1]).view(ht_batch.shape[0], self.num_entities) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self(r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 5666b332ae..e1d354ee34 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -9,7 +9,7 @@ from ..base import EntityEmbeddingModel from ...losses import Loss -from ...nn import Embedding, functional as F +from ...nn import Embedding, functional as pykeen_functional from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -113,7 +113,7 @@ def _reset_parameters_(self): # noqa: D102 if hasattr(module, "reset_parameters"): module.reset_parameters() - def _score( + def forward( self, h_indices: Optional[torch.LongTensor] = None, r_indices: Optional[torch.LongTensor] = None, @@ -134,14 +134,18 @@ def _score( #: shape: (batch_size, num_entities, d) h_all = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) t_all = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - w = self.w.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim, self.embedding_dim)) + w = self.w.get_in_canonical_shape(indices=r_indices, + reshape_dim=(self.num_slices, self.embedding_dim, self.embedding_dim)) b = self.b.get_in_canonical_shape(indices=r_indices) u = self.u.get_in_canonical_shape(indices=r_indices) vh = self.vh.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim)) vt = self.vt.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim)) if slice_size is None: - return F.ntn_interaction(h=h_all, t=t_all, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity) + return pykeen_functional.ntn_interaction( + h=h_all, t=t_all, + w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity, + ) # TODO: Not implemented if h_all.shape[1] > t_all.shape[1]: @@ -167,13 +171,29 @@ def _score( return torch.cat(scores_arr, dim=1) def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]).view(hrt_batch.shape[0], 1) + return self( + h_indices=hrt_batch[:, 0], + r_indices=hrt_batch[:, 1], + t_indices=hrt_batch[:, 2], + ).view(hrt_batch.shape[0], 1) def score_t(self, hr_batch: torch.LongTensor, slice_size: int = None) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], slice_size=slice_size).view(hr_batch.shape[0], self.num_entities) + return self( + h_indices=hr_batch[:, 0], + r_indices=hr_batch[:, 1], + slice_size=slice_size, + ).view(hr_batch.shape[0], self.num_entities) def score_r(self, ht_batch: torch.LongTensor, slice_size: int = None) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=ht_batch[:, 0], t_indices=ht_batch[:, 1], slice_size=slice_size).view(ht_batch.shape[0], self.num_relations) + return self( + h_indices=ht_batch[:, 0], + t_indices=ht_batch[:, 1], + slice_size=slice_size, + ).view(ht_batch.shape[0], self.num_relations) def score_h(self, rt_batch: torch.LongTensor, slice_size: int = None) -> torch.FloatTensor: # noqa: D102 - return self._score(r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1], slice_size=slice_size).view(rt_batch.shape[0], self.num_entities) + return self.forward( + r_indices=rt_batch[:, 0], + t_indices=rt_batch[:, 1], + slice_size=slice_size, + ).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index 9684334355..a37dcbf07e 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -107,7 +107,7 @@ def _reset_parameters_(self): # noqa: D102 ]: emb.reset_parameters() - def _score( + def forward( self, h_indices: Optional[torch.LongTensor], r_indices: Optional[torch.LongTensor], @@ -140,10 +140,10 @@ def _score( return scores def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]).view(-1, 1) + return self(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]).view(-1, 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], t_indices=None) + return self(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], t_indices=None) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._score(h_indices=None, r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]) + return self(h_indices=None, r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]) From 05e32e45dcc5cdf7d20495bec678c47e852c579c Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 17:05:10 +0100 Subject: [PATCH 126/690] Update ntn.py --- src/pykeen/models/unimodal/ntn.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index e1d354ee34..7759df466b 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -165,7 +165,10 @@ def forward( else: h = constant_tensor t = split - score = F.ntn_interaction(h=h, t=t, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity) + score = pykeen_functional.ntn_interaction( + h=h, t=t, + w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity, + ) scores_arr.append(score) return torch.cat(scores_arr, dim=1) From 97a7065569ab15bd5b7d536c05eba7ea59bfca6c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 17:11:23 +0100 Subject: [PATCH 127/690] Update TransE interaction function --- src/pykeen/models/unimodal/trans_e.py | 32 ++++++++++----------------- src/pykeen/nn/functional.py | 28 ++++++++++++++++++----- src/pykeen/nn/modules.py | 6 ++--- 3 files changed, 37 insertions(+), 29 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 682a6d5ef1..dd6ed29a8a 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -86,29 +86,21 @@ def __init__( def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings - h = self.entity_embeddings(indices=hrt_batch[:, 0]) - r = self.relation_embeddings(indices=hrt_batch[:, 1]) - t = self.entity_embeddings(indices=hrt_batch[:, 2]) - - # TODO question @mberr - why is keepdim=True here but the others it isn't? - # TODO: Use torch.dist - - return self.interaction_function(h=h, r=r, t=t, dim=-1, keepdim=True) + h = self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 0]) + r = self.relation_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 1]) + t = self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 2]) + return self.interaction_function(h=h, r=r, t=t).view(-1, 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings - h = self.entity_embeddings(indices=hr_batch[:, 0]) - r = self.relation_embeddings(indices=hr_batch[:, 1]) - t = self.entity_embeddings(indices=None) - - # TODO: Use torch.cdist - return self.interaction_function(h=h[:, None, :], r=r[:, None, :], t=t[None, :, :], dim=-1, keepdim=False) + h = self.entity_embeddings.get_in_canonical_shape(indices=hr_batch[:, 0]) + r = self.relation_embeddings.get_in_canonical_shape(indices=hr_batch[:, 1]) + t = self.entity_embeddings.get_in_canonical_shape(indices=None) + return self.interaction_function(h=h, r=r, t=t).view(hr_batch.shape[0], self.num_entities) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings - h = self.entity_embeddings(indices=None) - r = self.relation_embeddings(indices=rt_batch[:, 0]) - t = self.entity_embeddings(indices=rt_batch[:, 1]) - - # TODO: Use torch.cdist - return self.interaction_function(h=h[None, :, :], r=r[:, None, :], t=t[:, None, :], dim=-1, keepdim=False) + h = self.entity_embeddings.get_in_canonical_shape(indices=None) + r = self.relation_embeddings.get_in_canonical_shape(indices=rt_batch[:, 0]) + t = self.entity_embeddings.get_in_canonical_shape(indices=rt_batch[:, 1]) + return self.interaction_function(h=h, r=r, t=t).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 26d9073a6b..fd22e73951 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -425,12 +425,30 @@ def translational_interaction( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, - dim: int, - p: Union[int, str] = 'fro', - keepdim: bool = False, + p: Union[int, str] = 2, ) -> torch.FloatTensor: - """Evaluate the translational interaction.""" - return -torch.norm(h + r - t, dim=dim, p=p, keepdim=keepdim) + """ + Evaluate the ConvE interaction function. + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r: shape: (batch_size, num_relations, dim) + The relation representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + :param p: + The p for the norm. cf. torch.norm. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + num_heads, num_relations, num_tails = [xx.shape[1] for xx in (h, r, t)] + dim = h.shape[-1] + return -( + h.view(-1, num_heads, 1, 1, dim) + + r.view(-1, 1, num_relations, 1, dim) - + t.view(-1, 1, 1, num_tails, dim) + ).norm(p=p, dim=-1) def transr_interaction( diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index e5cc566070..892a708389 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -258,10 +258,8 @@ def forward( t: torch.FloatTensor, **kwargs, ) -> torch.FloatTensor: # noqa:D102 - return pykeen_functional.translational_interaction( - h=h, r=r, t=t, - p=self.p, dim=kwargs.get('dim', None), keepdim=kwargs.get('keepdim', False), - ) + self._check_for_empty_kwargs(kwargs=kwargs) + return pykeen_functional.translational_interaction(h=h, r=r, t=t, p=self.p) class ComplExInteractionFunction(InteractionFunction): From 34b884f4456a75401723f4d2d3fb120a3d9fce03 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 17:16:43 +0100 Subject: [PATCH 128/690] Improve stability of TransR functional --- src/pykeen/nn/functional.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 6d761a9746..5e50a6dc99 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -426,6 +426,7 @@ def translational_interaction( r: torch.FloatTensor, t: torch.FloatTensor, p: Union[int, str] = 2, + no_root: bool = False, ) -> torch.FloatTensor: """ Evaluate the ConvE interaction function. @@ -438,17 +439,19 @@ def translational_interaction( The tail representations. :param p: The p for the norm. cf. torch.norm. + :param no_root: + Whether to return |x-y|_p^p, cf. https://github.com/pytorch/pytorch/issues/28119 :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ num_heads, num_relations, num_tails = [xx.shape[1] for xx in (h, r, t)] dim = h.shape[-1] - return -( - h.view(-1, num_heads, 1, 1, dim) + - r.view(-1, 1, num_relations, 1, dim) - - t.view(-1, 1, 1, num_tails, dim) - ).norm(p=p, dim=-1) + d = (h.view(-1, num_heads, 1, 1, dim) + r.view(-1, 1, num_relations, 1, dim) - t.view(-1, 1, 1, num_tails, dim)) + if no_root: + return -(d ** p).sum(dim=-1) + else: + return -d.norm(p=p, dim=-1) def transr_interaction( @@ -482,7 +485,7 @@ def transr_interaction( t_bot = clamp_norm(t_bot, p=2, dim=-1, maxnorm=1.) # evaluate score function, shape: (b, e) - return translational_interaction(h=h_bot, r=r, t=t_bot, dim=-1, p=p) ** 2 + return translational_interaction(h=h_bot, r=r, t=t_bot, p=p, no_root=True) class GaussianDistribution(NamedTuple): From 90fb78e3aa54e54f120e85e6e22c79c081031938 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 17:37:03 +0100 Subject: [PATCH 129/690] Adjust TransD --- src/pykeen/models/unimodal/trans_r.py | 6 +- src/pykeen/nn/functional.py | 92 ++++++++++++++++++++------- src/pykeen/nn/modules.py | 3 +- 3 files changed, 73 insertions(+), 28 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 2cabc5c4f0..00d83347c0 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -134,7 +134,7 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: t = self.entity_embeddings(indices=hrt_batch[:, 2]).unsqueeze(dim=1) m_r = self.relation_projections(indices=hrt_batch[:, 1]).view(-1, self.embedding_dim, self.relation_dim) - return self.interaction_function(h=h, r=r, t=t, m_r=m_r).view(-1, 1) + return self.interaction_function(h=h, r=r, t=t, m_r=m_r).view(hrt_batch.shape[0], 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings @@ -143,7 +143,7 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 t = self.entity_embeddings(indices=None).unsqueeze(dim=0) m_r = self.relation_projections(indices=hr_batch[:, 1]).view(-1, self.embedding_dim, self.relation_dim) - return self.interaction_function(h=h, r=r, t=t, m_r=m_r) + return self.interaction_function(h=h, r=r, t=t, m_r=m_r).view(hr_batch.shape[0], self.num_entities) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings @@ -152,4 +152,4 @@ def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 t = self.entity_embeddings(indices=rt_batch[:, 1]).unsqueeze(dim=1) m_r = self.relation_projections(indices=rt_batch[:, 0]).view(-1, self.embedding_dim, self.relation_dim) - return self.interaction_function(h=h, r=r, t=t, m_r=m_r) + return self.interaction_function(h=h, r=r, t=t, m_r=m_r).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 5e50a6dc99..5c74bdd8e8 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -2,7 +2,7 @@ """Functional forms of interaction methods.""" import math -from typing import NamedTuple, Optional, Tuple, Union +from typing import NamedTuple, Optional, SupportsFloat, Tuple, Union import torch from torch import nn @@ -421,15 +421,47 @@ def rotate_interaction( return scores +def _translational_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + p: Union[int, str] = 2, + power_norm: bool = False, +) -> torch.FloatTensor: + """ + Evaluate a translational distance interaction function on already broadcasted representations. + + :param h: shape: (batch_size, num_heads, num_relations, num_tails, dim) + The head representations. + :param r: shape: (batch_size, num_heads, num_relations, num_tails, dim) + The relation representations. + :param t: shape: (batch_size, num_heads, num_relations, num_tails, dim) + The tail representations. + :param p: + The p for the norm. cf. torch.norm. + :param power_norm: + Whether to return |x-y|_p^p, cf. https://github.com/pytorch/pytorch/issues/28119 + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + d = (h + r - t) + if power_norm: + assert isinstance(p, SupportsFloat) + return -(d.abs() ** p).sum(dim=-1) + else: + return -d.norm(p=p, dim=-1) + + def translational_interaction( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, p: Union[int, str] = 2, - no_root: bool = False, + power_norm: bool = False, ) -> torch.FloatTensor: """ - Evaluate the ConvE interaction function. + Evaluate a translational distance interaction function. :param h: shape: (batch_size, num_heads, dim) The head representations. @@ -439,19 +471,24 @@ def translational_interaction( The tail representations. :param p: The p for the norm. cf. torch.norm. - :param no_root: + :param power_norm: Whether to return |x-y|_p^p, cf. https://github.com/pytorch/pytorch/issues/28119 :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ + num_heads, num_relations, num_tails, d_e = _extract_sizes(h, r, t)[:4] + h = h.view(-1, num_heads, 1, 1, d_e) + r = r.view(-1, 1, num_relations, 1, d_e) + t = t.view(-1, 1, 1, num_tails, d_e) + return _translational_interaction(h=h, r=r, t=t, p=p, power_norm=power_norm) + + +def _extract_sizes(h, r, t): num_heads, num_relations, num_tails = [xx.shape[1] for xx in (h, r, t)] - dim = h.shape[-1] - d = (h.view(-1, num_heads, 1, 1, dim) + r.view(-1, 1, num_relations, 1, dim) - t.view(-1, 1, 1, num_tails, dim)) - if no_root: - return -(d ** p).sum(dim=-1) - else: - return -d.norm(p=p, dim=-1) + d_e = h.shape[-1] + d_r = r.shape[-1] + return num_heads, num_relations, num_tails, d_e, d_r def transr_interaction( @@ -460,32 +497,39 @@ def transr_interaction( t: torch.FloatTensor, m_r: torch.FloatTensor, p: int, + power_norm: bool = True, ) -> torch.FloatTensor: """Evaluate the interaction function for given embeddings. - The embeddings have to be in a broadcastable shape. - - :param h: shape: (batch_size, num_entities, d_e) + :param h: shape: (batch_size, num_heads, d_e) Head embeddings. - :param r: shape: (batch_size, num_entities, d_r) + :param r: shape: (batch_size, num_relations, d_r) Relation embeddings. - :param t: shape: (batch_size, num_entities, d_e) - Tail embeddings. - :param m_r: shape: (batch_size, num_entities, d_e, d_r) + :param m_r: shape: (batch_size, num_relations, d_e, d_r) The relation specific linear transformations. + :param t: shape: (batch_size, num_tails, d_e) + Tail embeddings. + :param p: + The parameter p for selecting the norm. + :param power_norm: + Whether to return the powered norm instead. - :return: shape: (batch_size, num_entities) + :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # project to relation specific subspace, shape: (b, e, d_r) - h_bot = h @ m_r - t_bot = t @ m_r - # ensure constraints + num_heads, num_relations, num_tails, d_e, d_r = _extract_sizes(h=h, r=r, t=t) + # project to relation specific subspace and ensure constraints + # head, shape: (b, h, r, 1, d_r) + h_bot = h.view(-1, num_heads, 1, 1, d_e) @ m_r.view(-1, 1, num_relations, d_e, d_r) h_bot = clamp_norm(h_bot, p=2, dim=-1, maxnorm=1.) + + # head, shape: (b, 1, r, t, d_r) + t_bot = t.view(-1, 1, 1, num_tails, d_e) @ m_r.view(-1, 1, num_relations, d_e, d_r) t_bot = clamp_norm(t_bot, p=2, dim=-1, maxnorm=1.) - # evaluate score function, shape: (b, e) - return translational_interaction(h=h_bot, r=r, t=t_bot, p=p, no_root=True) + # evaluate score function, shape: (b, h, r, t) + r = r.view(-1, 1, num_relations, 1, d_r) + return _translational_interaction(h=h_bot, r=r, t=t_bot, p=p, power_norm=power_norm) class GaussianDistribution(NamedTuple): diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 892a708389..ecd0184943 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -634,7 +634,8 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa:D102 m_r = kwargs.pop('m_r') - return pykeen_functional.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p) + self._check_for_empty_kwargs(kwargs=kwargs) + return pykeen_functional.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p, power_norm=True) class RotatEInteraction(InteractionFunction): From dc2f21201fa9723d564636e0053a55353f9be60f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 17:39:40 +0100 Subject: [PATCH 130/690] fix manual example --- tests/test_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index ad3560a4f2..a37eefd45e 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1060,7 +1060,7 @@ def test_score_hrt_manual(self): self.assertEqual(scores.shape[1], 1) first_score = scores[0].item() # second_score = scores[1].item() - self.assertAlmostEqual(first_score, -32, delta=0.01) + self.assertAlmostEqual(first_score, -8, delta=1.0e-06) def _check_constraints(self): """Check model constraints. From 55249bad5aefb4ae19fd53882366810a9481681e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 17:43:48 +0100 Subject: [PATCH 131/690] cleanup code --- src/pykeen/nn/functional.py | 23 ++++++++--------------- 1 file changed, 8 insertions(+), 15 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 5c74bdd8e8..c35fc25259 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -185,9 +185,7 @@ def distmult_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # TODO: check if einsum is still very slow. - h, h_term, r, r_term, t, t_term = _normalize_terms_for_einsum(h, r, t) - return torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', h, r, t) + return _extended_einsum("bhd,brd,btd->bhrt", h, r, t) def complex_interaction( @@ -208,11 +206,9 @@ def complex_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - h, h_term, r, r_term, t, t_term = _normalize_terms_for_einsum(h, r, t) (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] - # TODO: check if einsum is still very slow. return sum( - torch.einsum(f'{h_term},{r_term},{t_term}->bhrt', hh, rr, tt) + _extended_einsum("bhd,brd,btd->bhrt", hh, rr, tt) for hh, rr, tt in [ (h_re, r_re, t_re), (h_re, r_im, t_im), @@ -256,15 +252,11 @@ def convkb_interaction( The scores. """ # bind sizes - batch_size = max(x.shape[0] for x in (h, r, t)) - num_heads = h.shape[1] - num_relations = r.shape[1] - num_tails = t.shape[1] + num_heads, num_relations, num_tails, embedding_dim = _extract_sizes(h, r, t)[:4] # decompose convolution for faster computation in 1-n case num_filters = conv.weight.shape[0] assert conv.weight.shape == (num_filters, 1, 1, 3) - embedding_dim = h.shape[-1] # compute conv(stack(h, r, t)) conv_head, conv_rel, conv_tail = conv.weight[:, 0, 0, :].t() @@ -281,7 +273,7 @@ def convkb_interaction( # Linear layer for final scores return linear( x.view(-1, embedding_dim * num_filters), - ).view(batch_size, num_heads, num_relations, num_tails) + ).view(-1, num_heads, num_relations, num_tails) def ermlp_interaction( @@ -311,8 +303,8 @@ def ermlp_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - num_heads, num_relations, num_tails = [x.shape[1] for x in (h, r, t)] - hidden_dim, embedding_dim = hidden.weight.shape + num_heads, num_relations, num_tails, embedding_dim = _extract_sizes(h, r, t)[:4] + hidden_dim = hidden.weight.shape[0] assert embedding_dim % 3 == 0 embedding_dim = embedding_dim // 3 # split, shape: (embedding_dim, hidden_dim) @@ -708,6 +700,7 @@ def _extended_einsum( *tensors, ) -> torch.FloatTensor: """Drop dimensions of size 1 to allow broadcasting.""" + # TODO: check if einsum is still very slow. lhs, rhs = eq.split("->") mod_ops, mod_t = [], [] for op, t in zip(lhs.split(","), tensors): @@ -771,7 +764,7 @@ def ntn_interaction( """ # TODO: check efficiency of einsum # save sizes - num_heads, num_relations, num_tails = [xx.shape[1] for xx in (h, b, t)] + num_heads, num_relations, num_tails = _extract_sizes(h, b, t)[:3] k = b.shape[-1] x = _extended_einsum("bhd,brkde,bte->bhrtk", h, w, t) x = x + _extended_einsum("brkd,bhd->bhk", vh, h).view(-1, num_heads, 1, 1, k) From 9d03e9f67b687395b5e521cad24c7f8c7985347c Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 17:43:55 +0100 Subject: [PATCH 132/690] Add helper functions to TransE @mberr this is a good place where we can think about how many different models share the same kinds of operations like this --- src/pykeen/models/base.py | 17 ++++++++---- src/pykeen/models/unimodal/trans_e.py | 38 +++++++++++++++++---------- src/pykeen/utils.py | 28 ++++++++++++++++++++ 3 files changed, 64 insertions(+), 19 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 227e683356..9e85b837fa 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -21,7 +21,10 @@ from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory from ..typing import Constrainer, DeviceHint, Initializer, MappedTriples, Normalizer -from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed +from ..utils import ( + NoRandomSeedNecessary, get_hr_indices, get_hrt_indices, get_ht_indices, get_rt_indices, + resolve_device, set_random_seed, +) __all__ = [ 'Model', @@ -1356,16 +1359,20 @@ def forward( return self.index_function(model=self, h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]).view(-1, 1) + h_indices, r_indices, t_indices = get_hrt_indices(hrt_batch) + return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(-1, 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], t_indices=None).view(-1, self.num_entities) + h_indices, r_indices, t_indices = get_hr_indices(hr_batch) + return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(-1, self.num_entities) def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self(h_indices=ht_batch[:, 0], r_indices=None, t_indices=ht_batch[:, 1]) + h_indices, r_indices, t_indices = get_ht_indices(ht_batch) + return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self(h_indices=None, r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]).view(-1, self.num_entities) + h_indices, r_indices, t_indices = get_rt_indices(rt_batch) + return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(-1, self.num_entities) class SimpleVectorEntityRelationEmbeddingModel( diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index dd6ed29a8a..2a89a0b9e4 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -2,7 +2,7 @@ """TransE.""" -from typing import Optional +from typing import Optional, Tuple import torch.autograd from torch.nn import functional @@ -14,7 +14,7 @@ from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint -from ...utils import compose +from ...utils import compose, get_hr_indices, get_hrt_indices, get_ht_indices, get_rt_indices __all__ = [ 'TransE', @@ -85,22 +85,32 @@ def __init__( self.interaction_function = TranslationalInteractionFunction(p=scoring_fct_norm) def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 0]) - r = self.relation_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 1]) - t = self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 2]) + h_indices, r_indices, t_indices = get_hrt_indices(hrt_batch) + h, r, t = self._get_hrt(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) return self.interaction_function(h=h, r=r, t=t).view(-1, 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings.get_in_canonical_shape(indices=hr_batch[:, 0]) - r = self.relation_embeddings.get_in_canonical_shape(indices=hr_batch[:, 1]) - t = self.entity_embeddings.get_in_canonical_shape(indices=None) + h_indices, r_indices, t_indices = get_hr_indices(hr_batch) + h, r, t = self._get_hrt(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) return self.interaction_function(h=h, r=r, t=t).view(hr_batch.shape[0], self.num_entities) + def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + h_indices, r_indices, t_indices = get_ht_indices(ht_batch) + h, r, t = self._get_hrt(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) + return self.interaction_function(h=h, r=r, t=t) + def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings.get_in_canonical_shape(indices=None) - r = self.relation_embeddings.get_in_canonical_shape(indices=rt_batch[:, 0]) - t = self.entity_embeddings.get_in_canonical_shape(indices=rt_batch[:, 1]) + h_indices, r_indices, t_indices = get_rt_indices(rt_batch) + h, r, t = self._get_hrt(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) return self.interaction_function(h=h, r=r, t=t).view(rt_batch.shape[0], self.num_entities) + + def _get_hrt( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: + h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) + r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) + t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) + return h, r, t diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 8356640605..45f632b3de 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -45,6 +45,10 @@ 'Result', 'fix_dataclass_init_docs', 'normalize_for_einsum', + 'get_hrt_indices', + 'get_hr_indices', + 'get_ht_indices', + 'get_rt_indices', ] logger = logging.getLogger(__name__) @@ -525,3 +529,27 @@ def broadcast_cat( x_rep.append(xr) y_rep.append(yr) return torch.cat([x.repeat(*x_rep), y.repeat(*y_rep)], dim=dim) + + +# These three following methods could be used throughout PyKEEN to improve +# the implicit documentation of functions + + +def get_hrt_indices(hrt_batch: torch.LongTensor) -> Tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor]: + """Get indices from a head/relation/tail batch.""" + return hrt_batch[:, 0], hrt_batch[:, 1], hrt_batch[:, 2] + + +def get_hr_indices(hr_batch: torch.LongTensor) -> Tuple[torch.LongTensor, torch.LongTensor, None]: + """Get indices from a head/relation batch.""" + return hr_batch[:, 0], hr_batch[:, 1], None + + +def get_ht_indices(ht_batch: torch.LongTensor) -> Tuple[torch.LongTensor, None, torch.LongTensor]: + """Get indices from a head/tail batch.""" + return ht_batch[:, 0], None, ht_batch[:, 1] + + +def get_rt_indices(rt_batch: torch.LongTensor) -> Tuple[None, torch.LongTensor, torch.LongTensor]: + """Get indices from a relation/tail batch.""" + return None, rt_batch[:, 0], rt_batch[:, 1] From 2b0b22aa26fe81c75dbfbcb28bcc9f8627697310 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 17:44:27 +0100 Subject: [PATCH 133/690] Move helper methods --- src/pykeen/nn/functional.py | 86 ++++++++++++++++--------------------- 1 file changed, 37 insertions(+), 49 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index c35fc25259..29532a6e22 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -2,12 +2,12 @@ """Functional forms of interaction methods.""" import math -from typing import NamedTuple, Optional, SupportsFloat, Tuple, Union +from typing import NamedTuple, Optional, SupportsFloat, Union import torch from torch import nn -from ..utils import broadcast_cat, clamp_norm, is_cudnn_error, normalize_for_einsum, split_complex +from ..utils import broadcast_cat, clamp_norm, is_cudnn_error, split_complex __all__ = [ "complex_interaction", @@ -23,16 +23,41 @@ ] -def _normalize_terms_for_einsum( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, -) -> Tuple[torch.FloatTensor, str, torch.FloatTensor, str, torch.FloatTensor, str]: - batch_size = max(h.shape[0], r.shape[0], t.shape[0]) - h_term, h = normalize_for_einsum(x=h, batch_size=batch_size, symbol='h') - r_term, r = normalize_for_einsum(x=r, batch_size=batch_size, symbol='r') - t_term, t = normalize_for_einsum(x=t, batch_size=batch_size, symbol='t') - return h, h_term, r, r_term, t, t_term +def _extract_sizes(h, r, t): + num_heads, num_relations, num_tails = [xx.shape[1] for xx in (h, r, t)] + d_e = h.shape[-1] + d_r = r.shape[-1] + return num_heads, num_relations, num_tails, d_e, d_r + + +def _extended_einsum( + eq: str, + *tensors, +) -> torch.FloatTensor: + """Drop dimensions of size 1 to allow broadcasting.""" + # TODO: check if einsum is still very slow. + lhs, rhs = eq.split("->") + mod_ops, mod_t = [], [] + for op, t in zip(lhs.split(","), tensors): + mod_op = "" + assert len(op) == len(t.shape) + for i, c in reversed(list(enumerate(op))): + if t.shape[i] == 1: + t = t.squeeze(dim=i) + else: + mod_op = c + mod_op + mod_ops.append(mod_op) + mod_t.append(t) + m_lhs = ",".join(mod_ops) + r_keep_dims = set("".join(mod_ops)) + m_rhs = "".join(c for c in rhs if c in r_keep_dims) + m_eq = f"{m_lhs}->{m_rhs}" + mod_r = torch.einsum(m_eq, *mod_t) + # unsqueeze + for i, c in enumerate(rhs): + if c not in r_keep_dims: + mod_r = mod_r.unsqueeze(dim=i) + return mod_r def _add_cuda_warning(func): @@ -476,13 +501,6 @@ def translational_interaction( return _translational_interaction(h=h, r=r, t=t, p=p, power_norm=power_norm) -def _extract_sizes(h, r, t): - num_heads, num_relations, num_tails = [xx.shape[1] for xx in (h, r, t)] - d_e = h.shape[-1] - d_r = r.shape[-1] - return num_heads, num_relations, num_tails, d_e, d_r - - def transr_interaction( h: torch.FloatTensor, r: torch.FloatTensor, @@ -695,36 +713,6 @@ def kg2e_interaction( return similarity(e=e, r=r, exact=exact) -def _extended_einsum( - eq: str, - *tensors, -) -> torch.FloatTensor: - """Drop dimensions of size 1 to allow broadcasting.""" - # TODO: check if einsum is still very slow. - lhs, rhs = eq.split("->") - mod_ops, mod_t = [], [] - for op, t in zip(lhs.split(","), tensors): - mod_op = "" - assert len(op) == len(t.shape) - for i, c in reversed(list(enumerate(op))): - if t.shape[i] == 1: - t = t.squeeze(dim=i) - else: - mod_op = c + mod_op - mod_ops.append(mod_op) - mod_t.append(t) - m_lhs = ",".join(mod_ops) - r_keep_dims = set("".join(mod_ops)) - m_rhs = "".join(c for c in rhs if c in r_keep_dims) - m_eq = f"{m_lhs}->{m_rhs}" - mod_r = torch.einsum(m_eq, *mod_t) - # unsqueeze - for i, c in enumerate(rhs): - if c not in r_keep_dims: - mod_r = mod_r.unsqueeze(dim=i) - return mod_r - - def ntn_interaction( h: torch.FloatTensor, t: torch.FloatTensor, From 861367d73c24475d52cb9588be830540b78e6459 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 17:44:57 +0100 Subject: [PATCH 134/690] Update functional.py --- src/pykeen/nn/functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 5c74bdd8e8..b0b570904a 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -445,7 +445,7 @@ def _translational_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - d = (h + r - t) + d = h + r - t if power_norm: assert isinstance(p, SupportsFloat) return -(d.abs() ** p).sum(dim=-1) From 684f53977a713bbe520a6bbfdc0f3de0f6ac94b1 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 17:44:59 +0100 Subject: [PATCH 135/690] Further simplify --- src/pykeen/nn/functional.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 29532a6e22..74d755f5ca 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -750,10 +750,8 @@ def ntn_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # TODO: check efficiency of einsum # save sizes - num_heads, num_relations, num_tails = _extract_sizes(h, b, t)[:3] - k = b.shape[-1] + num_heads, num_relations, num_tails, _, k = _extract_sizes(h, b, t) x = _extended_einsum("bhd,brkde,bte->bhrtk", h, w, t) x = x + _extended_einsum("brkd,bhd->bhk", vh, h).view(-1, num_heads, 1, 1, k) x = x + _extended_einsum("brkd,btd->btk", vt, t).view(-1, 1, 1, num_tails, k) From d02ebf7498b12299e92a52cdacb63a9c269bdf32 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 17:56:50 +0100 Subject: [PATCH 136/690] Fix HolE interaction --- src/pykeen/nn/functional.py | 33 +++++++++++++++++++++++---------- 1 file changed, 23 insertions(+), 10 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 74d755f5ca..2cdabf76f5 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -5,6 +5,7 @@ from typing import NamedTuple, Optional, SupportsFloat, Union import torch +import torch.fft from torch import nn from ..utils import broadcast_cat, clamp_norm, is_cudnn_error, split_complex @@ -378,22 +379,34 @@ def hole_interaction( r: torch.FloatTensor, t: torch.FloatTensor, ) -> torch.FloatTensor: # noqa: D102 - """Evaluate the HolE interaction function.""" + """ + Evaluate the HolE interaction function. + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r: shape: (batch_size, num_relations, dim) + The relation representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ # Circular correlation of entity embeddings - a_fft = torch.rfft(h, signal_ndim=1, onesided=True) - b_fft = torch.rfft(t, signal_ndim=1, onesided=True) + a_fft = torch.fft.rfft(h, dim=-1) + b_fft = torch.fft.rfft(t, dim=-1) - # complex conjugate, a_fft.shape = (batch_size, num_entities, d', 2) - a_fft[:, :, :, 1] *= -1 + # complex conjugate, shape = (b, h, d) + a_fft = torch.conj(a_fft) - # Hadamard product in frequency domain - p_fft = a_fft * b_fft + # Hadamard product in frequency domain, shape: (b, h, t, d) + p_fft = a_fft.unsqueeze(dim=2) * b_fft.unsqueeze(dim=1) - # inverse real FFT, shape: (batch_size, num_entities, d) - composite = torch.irfft(p_fft, signal_ndim=1, onesided=True, signal_sizes=(h.shape[-1],)) + # inverse real FFT, shape: (b, h, t, d) + composite = torch.fft.irfft(p_fft, n=h.shape[-1], dim=-1) # inner product with relation embedding - return torch.sum(r * composite, dim=-1, keepdim=False) + return _extended_einsum("bhtd,brd->bhrt", composite, r) def rotate_interaction( From 56ef32b5e453aa5323cf63fffafb466f486c0fad Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 18:18:25 +0100 Subject: [PATCH 137/690] Extend RotatE --- src/pykeen/models/unimodal/rotate.py | 61 ++----------------- src/pykeen/nn/functional.py | 91 ++++++++++++++++++---------- 2 files changed, 63 insertions(+), 89 deletions(-) diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index 7f78a44d19..516eff962a 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -9,7 +9,7 @@ import torch.autograd from torch.nn import functional -from ..base import EntityRelationEmbeddingModel +from .. import SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss from ...nn.init import xavier_uniform_ from ...nn.modules import RotatEInteraction @@ -52,7 +52,7 @@ def complex_normalize(x: torch.Tensor) -> torch.Tensor: return x -class RotatE(EntityRelationEmbeddingModel): +class RotatE(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of RotatE from [sun2019]_. RotatE models relations as rotations from head to tail entities in complex space: @@ -92,8 +92,10 @@ def __init__( random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: + # TODO: regularization super().__init__( triples_factory=triples_factory, + interaction_function=RotatEInteraction(), embedding_dim=2 * embedding_dim, loss=loss, automatic_memory_optimization=automatic_memory_optimization, @@ -104,58 +106,3 @@ def __init__( relation_initializer=init_phases, relation_constrainer=complex_normalize, ) - self.interaction_function = RotatEInteraction() - self.real_embedding_dim = embedding_dim - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(indices=hrt_batch[:, 0]).view(-1, self.real_embedding_dim, 2) - r = self.relation_embeddings(indices=hrt_batch[:, 1]).view(-1, self.real_embedding_dim, 2) - t = self.entity_embeddings(indices=hrt_batch[:, 2]).view(-1, self.real_embedding_dim, 2) - - # Compute scores - scores = self.interaction_function(h=h, r=r, t=t).view(-1, 1) - - # Embedding Regularization - self.regularize_if_necessary(h.view(-1, self.embedding_dim), t.view(-1, self.embedding_dim)) - - return scores - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(indices=hr_batch[:, 0]).view(-1, 1, self.real_embedding_dim, 2) - r = self.relation_embeddings(indices=hr_batch[:, 1]).view(-1, 1, self.real_embedding_dim, 2) - - # Rank against all entities - t = self.entity_embeddings(indices=None).view(1, -1, self.real_embedding_dim, 2) - - # Compute scores - scores = self.interaction_function(h=h, r=r, t=t) - - # Embedding Regularization - self.regularize_if_necessary(h.view(-1, self.embedding_dim), t.view(-1, self.embedding_dim)) - - return scores - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - r = self.relation_embeddings(indices=rt_batch[:, 0]).view(-1, 1, self.real_embedding_dim, 2) - t = self.entity_embeddings(indices=rt_batch[:, 1]).view(-1, 1, self.real_embedding_dim, 2) - - # r expresses a rotation in complex plane. - # The inverse rotation is expressed by the complex conjugate of r. - # The score is computed as the distance of the relation-rotated head to the tail. - # Equivalently, we can rotate the tail by the inverse relation, and measure the distance to the head, i.e. - # |h * r - t| = |h - conj(r) * t| - r_inv = torch.stack([r[:, :, :, 0], -r[:, :, :, 1]], dim=-1) - - # Rank against all entities - h = self.entity_embeddings(indices=None).view(1, -1, self.real_embedding_dim, 2) - - # Compute scores - scores = self.interaction_function(h=t, r=r_inv, t=h) - - # Embedding Regularization - self.regularize_if_necessary(h.view(-1, self.embedding_dim), t.view(-1, self.embedding_dim)) - - return scores diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 8adae66bcb..2729444c14 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -409,6 +409,13 @@ def hole_interaction( return _extended_einsum("bhtd,brd->bhrt", composite, r) +def _view_complex( + x: torch.FloatTensor, +) -> torch.Tensor: + real, imag = split_complex(x=x) + return torch.complex(real=real, imag=imag) + + def rotate_interaction( h: torch.FloatTensor, r: torch.FloatTensor, @@ -416,39 +423,64 @@ def rotate_interaction( ) -> torch.FloatTensor: """Evaluate the interaction function of RotatE for given embeddings. - The embeddings have to be in a broadcastable shape. - - WARNING: No forward constraints are applied. - - :param h: shape: (..., e, 2) - Head embeddings. Last dimension corresponds to (real, imag). - :param r: shape: (..., e, 2) - Relation embeddings. Last dimension corresponds to (real, imag). - :param t: shape: (..., e, 2) - Tail embeddings. Last dimension corresponds to (real, imag). + :param h: shape: (batch_size, num_heads, 2*dim) + The head representations. + :param r: shape: (batch_size, num_relations, 2*dim) + The relation representations. + :param t: shape: (batch_size, num_tails, 2*dim) + The tail representations. - :return: shape: (...) + :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # Decompose into real and imaginary part - h_re = h[..., 0] - h_im = h[..., 1] - r_re = r[..., 0] - r_im = r[..., 1] + # # r expresses a rotation in complex plane. + # # The inverse rotation is expressed by the complex conjugate of r. + # # The score is computed as the distance of the relation-rotated head to the tail. + # # Equivalently, we can rotate the tail by the inverse relation, and measure the distance to the head, i.e. + # # |h * r - t| = |h - conj(r) * t| + # r_inv = torch.stack([r[:, :, :, 0], -r[:, :, :, 1]], dim=-1) + h, r, t = [_view_complex(x) for x in (h, r, t)] # Rotate (=Hadamard product in complex space). - rot_h = torch.stack( - [ - h_re * r_re - h_im * r_im, - h_re * r_im + h_im * r_re, - ], - dim=-1, - ) + hr = _extended_einsum("bhd,brd->bhrd", h, r) + # Workaround until https://github.com/pytorch/pytorch/issues/30704 is fixed - diff = rot_h - t - scores = -torch.norm(diff.view(diff.shape[:-2] + (-1,)), dim=-1) + return negative_norm_of_sum( + hr.unsqueeze(dim=3), + t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]), + p=2, + power_norm=False, + ) + + +def negative_norm_of_sum( + *x: torch.FloatTensor, + p: Union[int, str] = 2, + power_norm: bool = False, +) -> torch.FloatTensor: + """ + Evaluate negative norm of a sum of vectors on already broadcasted representations. - return scores + :param x: shape: (batch_size, num_heads, num_relations, num_tails, dim) + The representations. + :param p: + The p for the norm. cf. torch.norm. + :param power_norm: + Whether to return |x-y|_p^p, cf. https://github.com/pytorch/pytorch/issues/28119 + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + d = sum(x) + if power_norm: + assert isinstance(p, SupportsFloat) + return -(d.abs() ** p).sum(dim=-1) + else: + if torch.is_complex(d): + # workaround for complex numbers: manually compute norm + return -(d.abs() ** p).sum(dim=-1) ** (1 / p) + else: + return -d.norm(p=p, dim=-1) def _translational_interaction( @@ -475,12 +507,7 @@ def _translational_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - d = h + r - t - if power_norm: - assert isinstance(p, SupportsFloat) - return -(d.abs() ** p).sum(dim=-1) - else: - return -d.norm(p=p, dim=-1) + return negative_norm_of_sum(h, r, -t, p=p, power_norm=power_norm) def translational_interaction( From d1bf12ab2e3dd04af5f79b32b830458cf1276fda Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 18:23:18 +0100 Subject: [PATCH 138/690] Let ConvKB subclass from simple --- src/pykeen/models/unimodal/conv_kb.py | 27 +++++++-------------------- tests/test_models.py | 4 ++-- 2 files changed, 9 insertions(+), 22 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 78d92db2a0..f4f263f5d4 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -5,9 +5,7 @@ import logging from typing import Optional -import torch.autograd - -from ..base import EntityRelationEmbeddingModel +from ..base import SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss from ...nn.modules import ConvKBInteractionFunction from ...regularizers import LpRegularizer, Regularizer @@ -21,7 +19,7 @@ logger = logging.getLogger(__name__) -class ConvKB(EntityRelationEmbeddingModel): +class ConvKB(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of ConvKB from [nguyen2018]_. ConvKB uses a convolutional neural network (CNN) whose feature maps capture global interactions of the input. @@ -88,6 +86,11 @@ def __init__( """ super().__init__( triples_factory=triples_factory, + interaction_function=ConvKBInteractionFunction( + hidden_dropout_rate=hidden_dropout_rate, + embedding_dim=embedding_dim, + num_filters=num_filters, + ), embedding_dim=embedding_dim, loss=loss, automatic_memory_optimization=automatic_memory_optimization, @@ -95,24 +98,8 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) - self.interaction_function = ConvKBInteractionFunction( - hidden_dropout_rate=hidden_dropout_rate, - embedding_dim=embedding_dim, - num_filters=num_filters, - ) def _reset_parameters_(self): # noqa: D102 # embeddings logger.warning('To be consistent with the paper, initialize entity and relation embeddings from TransE.') super()._reset_parameters_() - self.interaction_function.reset_parameters() - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=hrt_batch[:, 0]) - r = self.relation_embeddings(indices=hrt_batch[:, 1]) - t = self.entity_embeddings(indices=hrt_batch[:, 2]) - # Output layer regularization - # In the code base only the weights of the output layer are used for regularization - # c.f. https://github.com/daiquocnguyen/ConvKB/blob/73a22bfa672f690e217b5c18536647c7cf5667f1/model.py#L60-L66 - self.regularize_if_necessary(self.interaction_function.linear.weight, self.interaction_function.linear.bias) - return self.interaction_function.score_hrt(h=h, r=r, t=t) diff --git a/tests/test_models.py b/tests/test_models.py index a37eefd45e..18f87d5657 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -559,8 +559,8 @@ class TestConvKB(_ModelTestCase, unittest.TestCase): model_kwargs = { 'num_filters': 2, } - # two bias terms, one conv-filter, written as 3 parts - num_constant_init = 5 + # two bias terms, one conv-filter + num_constant_init = 3 class TestDistMult(_ModelTestCase, unittest.TestCase): From b6ce2cc7fa262c69672763d962f69ea79962cbec Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 18:23:39 +0100 Subject: [PATCH 139/690] Inline --- src/pykeen/models/unimodal/ermlp.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 4f7766fa04..d5d7150bec 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -50,14 +50,12 @@ def __init__( regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model.""" - interaction_function = ERMLPInteractionFunction( - embedding_dim=embedding_dim, - hidden_dim=hidden_dim, - ) - super().__init__( triples_factory=triples_factory, - interaction_function=interaction_function, + interaction_function=ERMLPInteractionFunction( + embedding_dim=embedding_dim, + hidden_dim=hidden_dim, + ), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, From a00641963ac183039cbb21aaa08b3456410cd87c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 18:24:20 +0100 Subject: [PATCH 140/690] Let ERMLPE subclass from simple --- src/pykeen/models/unimodal/ermlpe.py | 63 ++++------------------------ 1 file changed, 8 insertions(+), 55 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index 5c8c004e90..882d274d01 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -4,9 +4,7 @@ from typing import Optional, Type -import torch - -from ..base import EntityRelationEmbeddingModel +from ..base import SimpleVectorEntityRelationEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss from ...nn.modules import ERMLPEInteractionFunction from ...regularizers import Regularizer @@ -18,7 +16,7 @@ ] -class ERMLPE(EntityRelationEmbeddingModel): +class ERMLPE(SimpleVectorEntityRelationEmbeddingModel): r"""An extension of ERMLP proposed by [sharifzadeh2019]_. This model uses a neural network-based approach similar to ERMLP and with slight modifications. @@ -68,6 +66,12 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, + interaction_function=ERMLPEInteractionFunction( + hidden_dim=hidden_dim, + input_dropout=input_dropout, + hidden_dropout=hidden_dropout, + embedding_dim=embedding_dim, + ), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, @@ -75,54 +79,3 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) - self.interaction_function = ERMLPEInteractionFunction( - hidden_dim=hidden_dim, - input_dropout=input_dropout, - hidden_dropout=hidden_dropout, - embedding_dim=embedding_dim, - ) - - def _reset_parameters_(self): # noqa: D102 - super()._reset_parameters_() - self.interaction_function.reset_parameters() - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(indices=hrt_batch[:, 0]) - r = self.relation_embeddings(indices=hrt_batch[:, 1]) - t = self.entity_embeddings(indices=hrt_batch[:, 2]) - - # Embedding Regularization - self.regularize_if_necessary(h, r, t) - - return self.interaction_function.score_hrt(h=h, r=r, t=t) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=hr_batch[:, 0]) - r = self.relation_embeddings(indices=hr_batch[:, 1]) - t = self.entity_embeddings(indices=None) - - # Embedding Regularization - self.regularize_if_necessary(h, r, t) - - return self.interaction_function.score_t(h=h, r=r, all_entities=t) - - def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=ht_batch[:, 0]) - r = self.relation_embeddings(indices=None) - t = self.entity_embeddings(indices=ht_batch[:, 1]) - - # Embedding Regularization - self.regularize_if_necessary(h, r, t) - - return self.interaction_function.score_r(h=h, all_relations=r, t=t) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=None) - r = self.relation_embeddings(indices=rt_batch[:, 0]) - t = self.entity_embeddings(indices=rt_batch[:, 1]) - - # Embedding Regularization - self.regularize_if_necessary(h, r, t) - - return self.interaction_function.score_h(all_entities=h, r=r, t=t) From b87d1fd9a2fd462417f863326340f684b1062e9d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 18:40:46 +0100 Subject: [PATCH 141/690] Extract functional interaction function for ProjE --- src/pykeen/models/unimodal/proj_e.py | 39 +++------------------- src/pykeen/nn/functional.py | 48 ++++++++++++++++++++++++++++ 2 files changed, 53 insertions(+), 34 deletions(-) diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index 98e39e5a59..facea44802 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -11,6 +11,7 @@ from ..base import EntityRelationEmbeddingModel from ...losses import Loss +from ...nn import functional as F from ...nn.init import xavier_uniform_ from ...regularizers import Regularizer from ...triples import TriplesFactory @@ -104,39 +105,9 @@ def _reset_parameters_(self): # noqa: D102 def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings - h = self.entity_embeddings(indices=hrt_batch[:, 0]) - r = self.relation_embeddings(indices=hrt_batch[:, 1]) - t = self.entity_embeddings(indices=hrt_batch[:, 2]) + h = self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 0]) + r = self.relation_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 1]) + t = self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 2]) # Compute score - hidden = self.inner_non_linearity(self.d_e[None, :] * h + self.d_r[None, :] * r + self.b_c[None, :]) - scores = torch.sum(hidden * t, dim=-1, keepdim=True) + self.b_p - - return scores - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(indices=hr_batch[:, 0]) - r = self.relation_embeddings(indices=hr_batch[:, 1]) - t = self.entity_embeddings(indices=None) - - # Rank against all entities - hidden = self.inner_non_linearity(self.d_e[None, :] * h + self.d_r[None, :] * r + self.b_c[None, :]) - scores = torch.sum(hidden[:, None, :] * t[None, :, :], dim=-1) + self.b_p - - return scores - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(indices=None) - r = self.relation_embeddings(indices=rt_batch[:, 0]) - t = self.entity_embeddings(indices=rt_batch[:, 1]) - - # Rank against all entities - hidden = self.inner_non_linearity( - self.d_e[None, None, :] * h[None, :, :] - + (self.d_r[None, None, :] * r[:, None, :] + self.b_c[None, None, :]), - ) - scores = torch.sum(hidden * t[:, None, :], dim=-1) + self.b_p - - return scores + return F.proje_interaction(h=h, r=r, t=t, d_e=self.d_e, d_r=self.d_r, b_c=self.b_c, b_p=self.b_p, activation=self.inner_non_linearity).view(-1, 1) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 2729444c14..192538f33c 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -798,3 +798,51 @@ def ntn_interaction( x = activation(x) x = _extended_einsum("bhrtk,brk->bhrt", x, u) return x + + +def proje_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + d_e: torch.FloatTensor, + d_r: torch.FloatTensor, + b_c: torch.FloatTensor, + b_p: torch.FloatTensor, + activation: nn.Module, +) -> torch.FloatTensor: + r""" + Evaluate the ProjE interaction function. + + .. math:: + + f(h, r, t) = g(t z(D_e h + D_r r + b_c) + b_p) + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r: shape: (batch_size, num_relations, dim) + The relation representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + :param d_e: shape: (dim,) + Global entity projection. + :param d_r: shape: (dim,) + Global relation projection. + :param b_c: shape: (dim,) + Global combination bias. + :param b_p: shape: (1,) + Final score bias + :param activation: + The activation function. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + num_heads, num_relations, num_tails, dim = _extract_sizes(h, r, t)[:4] + # global projections + h = h * d_e.view(1, 1, dim) + r = r * d_r.view(1, 1, dim) + # combination, shape: (b, h, r, d) + x = h.unsqueeze(dim=2) + r.unsqueeze(dim=1) + b_c.view(1, 1, 1, dim) + x = activation(x) + # dot product with t, shape: (b, h, r, t) + return (x @ t.unsqueeze(dim=1).transpose(-2, -1)) + b_p From 8640d3231b8be2f5605af878cea2a2e75939f968 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 18:44:55 +0100 Subject: [PATCH 142/690] Extract ProjE interaction function and make it simple --- src/pykeen/models/unimodal/proj_e.py | 46 +++++----------------------- src/pykeen/nn/modules.py | 45 +++++++++++++++++++++++++++ 2 files changed, 52 insertions(+), 39 deletions(-) diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index facea44802..241e9eadb1 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -4,15 +4,12 @@ from typing import Optional -import numpy -import torch -import torch.autograd from torch import nn -from ..base import EntityRelationEmbeddingModel +from .. import SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss -from ...nn import functional as F from ...nn.init import xavier_uniform_ +from ...nn.modules import ProjEInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -22,7 +19,7 @@ ] -class ProjE(EntityRelationEmbeddingModel): +class ProjE(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of ProjE from [shi2017]_. ProjE is a neural network-based approach with a *combination* and a *projection* layer. The interaction model @@ -69,6 +66,10 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, + interaction_function=ProjEInteractionFunction( + embedding_dim=embedding_dim, + inner_non_linearity=inner_non_linearity, + ), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, @@ -78,36 +79,3 @@ def __init__( entity_initializer=xavier_uniform_, relation_initializer=xavier_uniform_, ) - - # Global entity projection - self.d_e = nn.Parameter(torch.empty(self.embedding_dim, device=self.device), requires_grad=True) - - # Global relation projection - self.d_r = nn.Parameter(torch.empty(self.embedding_dim, device=self.device), requires_grad=True) - - # Global combination bias - self.b_c = nn.Parameter(torch.empty(self.embedding_dim, device=self.device), requires_grad=True) - - # Global combination bias - self.b_p = nn.Parameter(torch.empty(1, device=self.device), requires_grad=True) - - if inner_non_linearity is None: - inner_non_linearity = nn.Tanh() - self.inner_non_linearity = inner_non_linearity - - def _reset_parameters_(self): # noqa: D102 - super()._reset_parameters_() - bound = numpy.sqrt(6) / self.embedding_dim - nn.init.uniform_(self.d_e, a=-bound, b=bound) - nn.init.uniform_(self.d_r, a=-bound, b=bound) - nn.init.uniform_(self.b_c, a=-bound, b=bound) - nn.init.uniform_(self.b_p, a=-bound, b=bound) - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 0]) - r = self.relation_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 1]) - t = self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 2]) - - # Compute score - return F.proje_interaction(h=h, r=r, t=t, d_e=self.d_e, d_r=self.d_r, b_c=self.b_c, b_p=self.b_p, activation=self.inner_non_linearity).view(-1, 1) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index ecd0184943..8cfe67ee44 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -664,3 +664,48 @@ def forward( ) -> torch.FloatTensor: # noqa: D102 self._check_for_empty_kwargs(kwargs) return pykeen_functional.hole_interaction(h=h, r=r, t=t) + + +class ProjEInteractionFunction(InteractionFunction): + """Interaction function for ProjE.""" + + def __init__( + self, + embedding_dim: int = 50, + inner_non_linearity: Optional[nn.Module] = None, + ): + super().__init__() + + # Global entity projection + self.d_e = nn.Parameter(torch.empty(embedding_dim), requires_grad=True) + + # Global relation projection + self.d_r = nn.Parameter(torch.empty(embedding_dim), requires_grad=True) + + # Global combination bias + self.b_c = nn.Parameter(torch.empty(embedding_dim), requires_grad=True) + + # Global combination bias + self.b_p = nn.Parameter(torch.empty(1), requires_grad=True) + + if inner_non_linearity is None: + inner_non_linearity = nn.Tanh() + self.inner_non_linearity = inner_non_linearity + + def reset_parameters(self): # noqa: D102 + embedding_dim = self.d_e.shape[0] + bound = math.sqrt(6) / embedding_dim + for p in self.parameters(): + nn.init.uniform_(p, a=-bound, b=bound) + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + self._check_for_empty_kwargs(kwargs=kwargs) + + # Compute score + return pykeen_functional.proje_interaction(h=h, r=r, t=t, d_e=self.d_e, d_r=self.d_r, b_c=self.b_c, b_p=self.b_p, activation=self.inner_non_linearity).view(-1, 1) From 0bf9bc6a46424ad1efc5bd6308eea52902b4f1ac Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 18:45:55 +0100 Subject: [PATCH 143/690] Add RESCAL functional interaction --- src/pykeen/nn/functional.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 192538f33c..61fc694295 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -846,3 +846,24 @@ def proje_interaction( x = activation(x) # dot product with t, shape: (b, h, r, t) return (x @ t.unsqueeze(dim=1).transpose(-2, -1)) + b_p + + +def rescal_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """ + Evaluate the RESCAL interaction function. + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r: shape: (batch_size, num_relations, dim, dim) + The relation representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + return _extended_einsum("bhd,brde,bte->bhrt", h, r, t) From f9997fda7a324e5d0459148db70d3a77ea2820c0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 18:51:31 +0100 Subject: [PATCH 144/690] Add RESCAL interaction function, and extract a baseclass --- src/pykeen/nn/modules.py | 71 +++++++++++++++++----------------------- 1 file changed, 30 insertions(+), 41 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 8cfe67ee44..f903643abf 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -4,7 +4,7 @@ import logging import math -from typing import Any, Mapping, Optional, Sequence, Tuple +from typing import Any, Callable, Mapping, Optional, Sequence, Tuple import torch from torch import nn @@ -240,6 +240,21 @@ def reset_parameters(self): mod.reset_parameters() +class FunctionalInteractionFunction(InteractionFunction): + """Interaction function without state or additional parameters.""" + interaction: Callable[[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], torch.FloatTensor] + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + self._check_for_empty_kwargs(kwargs) + return self.interaction(h, r, t) + + class TranslationalInteractionFunction(InteractionFunction): """The translational interaction function shared by the TransE, TransR, TransH, and other Trans models.""" @@ -262,18 +277,10 @@ def forward( return pykeen_functional.translational_interaction(h=h, r=r, t=t, p=self.p) -class ComplExInteractionFunction(InteractionFunction): +class ComplExInteractionFunction(FunctionalInteractionFunction): """Interaction function of ComplEx.""" - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return pykeen_functional.complex_interaction(h=h, r=r, t=t) + interaction = pykeen_functional.complex_interaction def _calculate_missing_shape_information( @@ -486,18 +493,10 @@ def forward( ) -class DistMultInteractionFunction(InteractionFunction): +class DistMultInteractionFunction(FunctionalInteractionFunction): """Interaction function of DistMult.""" - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return pykeen_functional.distmult_interaction(h=h, r=r, t=t) + interaction = pykeen_functional.distmult_interaction class ERMLPInteractionFunction(InteractionFunction): @@ -638,32 +637,16 @@ def forward( return pykeen_functional.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p, power_norm=True) -class RotatEInteraction(InteractionFunction): +class RotatEInteraction(FunctionalInteractionFunction): """Interaction function of RotatE.""" - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return pykeen_functional.rotate_interaction(h=h, r=r, t=t) + interaction = pykeen_functional.rotate_interaction -class HolEInteractionFunction(InteractionFunction): +class HolEInteractionFunction(FunctionalInteractionFunction): """Interaction function for HolE.""" - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return pykeen_functional.hole_interaction(h=h, r=r, t=t) + interaction = pykeen_functional.hole_interaction class ProjEInteractionFunction(InteractionFunction): @@ -709,3 +692,9 @@ def forward( # Compute score return pykeen_functional.proje_interaction(h=h, r=r, t=t, d_e=self.d_e, d_r=self.d_r, b_c=self.b_c, b_p=self.b_p, activation=self.inner_non_linearity).view(-1, 1) + + +class RESCALInteractionFunction(FunctionalInteractionFunction): + """RESCAL interaction function.""" + + interaction = pykeen_functional.rescal_interaction From fe604d3dea6eee60588872971a195c876ccea8e6 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 18:56:07 +0100 Subject: [PATCH 145/690] Make RESCAL simple --- src/pykeen/models/unimodal/rescal.py | 54 +++------------------------- 1 file changed, 5 insertions(+), 49 deletions(-) diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index d158b1e5cf..1a132f7c30 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -4,10 +4,9 @@ from typing import Optional -import torch - -from ..base import EntityRelationEmbeddingModel +from ..base import SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss +from ...nn.modules import RESCALInteractionFunction from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -17,7 +16,7 @@ ] -class RESCAL(EntityRelationEmbeddingModel): +class RESCAL(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of RESCAL from [nickel2011]_. This model represents relations as matrices and models interactions between latent features. @@ -68,8 +67,10 @@ def __init__( - OpenKE `implementation of RESCAL `_ """ + # TODO: regularization super().__init__( triples_factory=triples_factory, + interaction_function=RESCALInteractionFunction(), embedding_dim=embedding_dim, relation_dim=embedding_dim ** 2, # d x d matrices automatic_memory_optimization=automatic_memory_optimization, @@ -78,48 +79,3 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - # shape: (b, d) - h = self.entity_embeddings(indices=hrt_batch[:, 0]).view(-1, 1, self.embedding_dim) - # shape: (b, d, d) - r = self.relation_embeddings(indices=hrt_batch[:, 1]).view(-1, self.embedding_dim, self.embedding_dim) - # shape: (b, d) - t = self.entity_embeddings(indices=hrt_batch[:, 2]).view(-1, self.embedding_dim, 1) - - # Compute scores - scores = h @ r @ t - - # Regularization - self.regularize_if_necessary(h, r, t) - - return scores[:, :, 0] - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=hr_batch[:, 0]).view(-1, 1, self.embedding_dim) - r = self.relation_embeddings(indices=hr_batch[:, 1]).view(-1, self.embedding_dim, self.embedding_dim) - t = self.entity_embeddings(indices=None).transpose(0, 1).view(1, self.embedding_dim, self.num_entities) - - # Compute scores - scores = h @ r @ t - - # Regularization - self.regularize_if_necessary(h, r, t) - - return scores[:, 0, :] - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - """Forward pass using left side (head) prediction.""" - # Get embeddings - h = self.entity_embeddings(indices=None).view(1, self.num_entities, self.embedding_dim) - r = self.relation_embeddings(indices=rt_batch[:, 0]).view(-1, self.embedding_dim, self.embedding_dim) - t = self.entity_embeddings(indices=rt_batch[:, 1]).view(-1, self.embedding_dim, 1) - - # Compute scores - scores = h @ r @ t - - # Regularization - self.regularize_if_necessary(h, r, t) - - return scores[:, :, 0] From b696021aa96eaecf603495170ad68980675e9e08 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 19:19:12 +0100 Subject: [PATCH 146/690] Refactor simple --- src/pykeen/models/unimodal/simple.py | 41 ++++++++++++++++------------ src/pykeen/nn/functional.py | 39 +++++++++++++++++++++++++- 2 files changed, 62 insertions(+), 18 deletions(-) diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index a37dcbf07e..c933a57399 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -9,6 +9,7 @@ from ..base import EntityRelationEmbeddingModel from ...losses import Loss, SoftplusLoss from ...nn import Embedding +from ...nn.modules import DistMultInteractionFunction from ...regularizers import PowerSumRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -82,6 +83,7 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) + self.interaction_function = DistMultInteractionFunction() # extra embeddings self.tail_entity_embeddings = Embedding.init_with_device( @@ -113,23 +115,13 @@ def forward( r_indices: Optional[torch.LongTensor], t_indices: Optional[torch.LongTensor], ) -> torch.FloatTensor: # noqa: D102 - # forward model - h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) - t = self.tail_entity_embeddings.get_in_canonical_shape(indices=t_indices) - scores = (h * r * t).sum(dim=-1) - - # Regularization - self.regularize_if_necessary(h, r, t) - - # backward model - h = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - r = self.inverse_relation_embeddings.get_in_canonical_shape(indices=r_indices) - t = self.tail_entity_embeddings.get_in_canonical_shape(indices=h_indices) - scores = 0.5 * (scores + (h * r * t).sum(dim=-1)) - - # Regularization - self.regularize_if_necessary(h, r, t) + scores = 0.5 * sum( + self._single_forward(h_ind=h_ind, r_ind=r_ind, t_ind=t_ind, r_emb=r_emb) + for (h_ind, r_ind, t_ind, r_emb) in ( + (h_indices, r_indices, t_indices, self.relation_embeddings), + (t_indices, r_indices, h_indices, self.inverse_relation_embeddings), + ) + ) # Note: In the code in their repository, the score is clamped to [-20, 20]. # That is not mentioned in the paper, so it is omitted here. @@ -139,6 +131,21 @@ def forward( return scores + def _single_forward( + self, + h_ind: torch.LongTensor, + r_ind: torch.LongTensor, + t_ind: torch.LongTensor, + r_emb: Embedding, + ) -> torch.FloatTensor: + # scores + h = self.entity_embeddings.get_in_canonical_shape(h_ind) + r = r_emb.get_in_canonical_shape(r_ind) + t = self.tail_entity_embeddings.get_in_canonical_shape(t_ind) + # Regularization + self.regularize_if_necessary(h, r, t) + return self.interaction_function(h, r, t) + def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]).view(-1, 1) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 61fc694295..bafd6b5d9f 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -2,7 +2,7 @@ """Functional forms of interaction methods.""" import math -from typing import NamedTuple, Optional, SupportsFloat, Union +from typing import NamedTuple, Optional, SupportsFloat, Tuple, Union import torch import torch.fft @@ -867,3 +867,40 @@ def rescal_interaction( The scores. """ return _extended_einsum("bhd,brde,bte->bhrt", h, r, t) + + +def simple_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + h_inv: torch.FloatTensor, + r_inv: torch.FloatTensor, + t_inv: torch.FloatTensor, + clamp: Optional[Tuple[float, float]] = None, +) -> torch.FloatTensor: + """ + Evaluate the SimplE interaction function. + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r: shape: (batch_size, num_relations, dim, dim) + The relation representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + :param h_inv: shape: (batch_size, num_heads, dim) + The inverse head representations. + :param r_inv: shape: (batch_size, num_relations, dim, dim) + The relation representations. + :param t_inv: shape: (batch_size, num_tails, dim) + The tail representations. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + scores = 0.5 * (distmult_interaction(h=h, r=r, t=t) + distmult_interaction(h=h_inv, r=r_inv, t=t_inv)) + # Note: In the code in their repository, the score is clamped to [-20, 20]. + # That is not mentioned in the paper, so it is made optional here. + if clamp: + min_, max_ = clamp + scores = scores.clamp(min=min_, max=max_) + return scores From de165b0b25843c77553b892e53d18aa474ba8bd8 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 19:31:06 +0100 Subject: [PATCH 147/690] Fix flake8 --- src/pykeen/nn/modules.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index f903643abf..96d1ba11fe 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -687,12 +687,14 @@ def forward( r: torch.FloatTensor, t: torch.FloatTensor, **kwargs, - ) -> torch.FloatTensor: + ) -> torch.FloatTensor: # noqa:D102 self._check_for_empty_kwargs(kwargs=kwargs) # Compute score - return pykeen_functional.proje_interaction(h=h, r=r, t=t, d_e=self.d_e, d_r=self.d_r, b_c=self.b_c, b_p=self.b_p, activation=self.inner_non_linearity).view(-1, 1) - + return pykeen_functional.proje_interaction( + h=h, r=r, t=t, + d_e=self.d_e, d_r=self.d_r, b_c=self.b_c, b_p=self.b_p, activation=self.inner_non_linearity, + ).view(-1, 1) class RESCALInteractionFunction(FunctionalInteractionFunction): """RESCAL interaction function.""" From d6b7fe378e859dfb998e4b8cb8dee00b8daa0581 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 19:34:02 +0100 Subject: [PATCH 148/690] Add builder for functional modules --- src/pykeen/nn/modules.py | 61 ++++++++++++++++++---------------------- 1 file changed, 28 insertions(+), 33 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 96d1ba11fe..c495b9c509 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -4,7 +4,7 @@ import logging import math -from typing import Any, Callable, Mapping, Optional, Sequence, Tuple +from typing import Any, Callable, Mapping, Optional, Sequence, Tuple, Type import torch from torch import nn @@ -240,19 +240,25 @@ def reset_parameters(self): mod.reset_parameters() -class FunctionalInteractionFunction(InteractionFunction): - """Interaction function without state or additional parameters.""" - interaction: Callable[[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], torch.FloatTensor] +def _build_module_from_stateless( + f: Callable[[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], torch.FloatTensor], +) -> Type[InteractionFunction]: + """Build a stateless interaction function module with a pre-defined functional interface.""" - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return self.interaction(h, r, t) + class FunctionalInteractionFunction(InteractionFunction): + """Interaction function without state or additional parameters.""" + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: # noqa: D102 + self._check_for_empty_kwargs(kwargs) + return f(h, r, t) + + return FunctionalInteractionFunction class TranslationalInteractionFunction(InteractionFunction): @@ -277,10 +283,8 @@ def forward( return pykeen_functional.translational_interaction(h=h, r=r, t=t, p=self.p) -class ComplExInteractionFunction(FunctionalInteractionFunction): - """Interaction function of ComplEx.""" - - interaction = pykeen_functional.complex_interaction +#: Interaction function of ComplEx +ComplExInteractionFunction = _build_module_from_stateless(pykeen_functional.complex_interaction) def _calculate_missing_shape_information( @@ -493,10 +497,8 @@ def forward( ) -class DistMultInteractionFunction(FunctionalInteractionFunction): - """Interaction function of DistMult.""" - - interaction = pykeen_functional.distmult_interaction +#: Interaction function for HolE +DistMultInteractionFunction = _build_module_from_stateless(pykeen_functional.distmult_interaction) class ERMLPInteractionFunction(InteractionFunction): @@ -637,16 +639,11 @@ def forward( return pykeen_functional.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p, power_norm=True) -class RotatEInteraction(FunctionalInteractionFunction): - """Interaction function of RotatE.""" - - interaction = pykeen_functional.rotate_interaction - - -class HolEInteractionFunction(FunctionalInteractionFunction): - """Interaction function for HolE.""" +#: Interaction function of RotatE. +RotatEInteraction = _build_module_from_stateless(pykeen_functional.rotate_interaction) - interaction = pykeen_functional.hole_interaction +#: Interaction function for HolE. +HolEInteractionFunction = _build_module_from_stateless(pykeen_functional.hole_interaction) class ProjEInteractionFunction(InteractionFunction): @@ -696,7 +693,5 @@ def forward( d_e=self.d_e, d_r=self.d_r, b_c=self.b_c, b_p=self.b_p, activation=self.inner_non_linearity, ).view(-1, 1) -class RESCALInteractionFunction(FunctionalInteractionFunction): - """RESCAL interaction function.""" - interaction = pykeen_functional.rescal_interaction +RESCALInteractionFunction = _build_module_from_stateless(pykeen_functional.rescal_interaction) From 519f8691c9ce28ad56e0579d62c618110f930433 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 19:35:40 +0100 Subject: [PATCH 149/690] Extract interaction function from structured embedding --- .../models/unimodal/structured_embedding.py | 120 +++--------------- src/pykeen/nn/functional.py | 38 ++++++ src/pykeen/nn/modules.py | 33 +++++ 3 files changed, 87 insertions(+), 104 deletions(-) diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index 1401504172..4ddddeae9f 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -6,15 +6,13 @@ from typing import Optional import numpy as np -import torch -import torch.autograd from torch import nn from torch.nn import functional -from ..base import EntityEmbeddingModel +from .. import SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss -from ...nn import Embedding from ...nn.init import xavier_uniform_ +from ...nn.modules import StructuredEmbeddingInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -25,7 +23,7 @@ ] -class StructuredEmbedding(EntityEmbeddingModel): +class StructuredEmbedding(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of the Structured Embedding (SE) published by [bordes2011]_. SE applies role- and relation-specific projection matrices @@ -63,9 +61,21 @@ def __init__( :param embedding_dim: The entity embedding dimension $d$. Is usually $d \in [50, 300]$. :param scoring_fct_norm: The $l_p$ norm. Usually 1 for SE. """ + # Embeddings + init_bound = 6 / np.sqrt(embedding_dim) + # Initialise relation embeddings to unit length + initializer = compose( + functools.partial(nn.init.uniform_, a=-init_bound, b=+init_bound), + functional.normalize, + ) super().__init__( triples_factory=triples_factory, + interaction_function=StructuredEmbeddingInteractionFunction( + p=scoring_fct_norm, + power_norm=False, + ), embedding_dim=embedding_dim, + relation_dim=2 * embedding_dim ** 2, # head and tail projection matrices automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, @@ -73,103 +83,5 @@ def __init__( regularizer=regularizer, entity_initializer=xavier_uniform_, entity_constrainer=functional.normalize, + relation_initializer=initializer, ) - - self.scoring_fct_norm = scoring_fct_norm - - # Embeddings - init_bound = 6 / np.sqrt(self.embedding_dim) - # Initialise relation embeddings to unit length - initializer = compose( - functools.partial(nn.init.uniform_, a=-init_bound, b=+init_bound), - functional.normalize, - ) - self.left_relation_embeddings = Embedding.init_with_device( - num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim ** 2, - device=self.device, - initializer=initializer, - ) - self.right_relation_embeddings = Embedding.init_with_device( - num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim ** 2, - device=self.device, - initializer=initializer, - ) - - def _reset_parameters_(self): # noqa: D102 - super()._reset_parameters_() - self.left_relation_embeddings.reset_parameters() - self.right_relation_embeddings.reset_parameters() - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(indices=hrt_batch[:, 0]).view(-1, self.embedding_dim, 1) - rel_h = self.left_relation_embeddings(indices=hrt_batch[:, 1]).view(-1, self.embedding_dim, self.embedding_dim) - rel_t = self.right_relation_embeddings(indices=hrt_batch[:, 1]).view(-1, self.embedding_dim, self.embedding_dim) - t = self.entity_embeddings(indices=hrt_batch[:, 2]).view(-1, self.embedding_dim, 1) - - # Project entities - proj_h = rel_h @ h - proj_t = rel_t @ t - - scores = -torch.norm(proj_h - proj_t, dim=1, p=self.scoring_fct_norm) - return scores - - def score_t(self, hr_batch: torch.LongTensor, slice_size: int = None) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(indices=hr_batch[:, 0]).view(-1, self.embedding_dim, 1) - rel_h = self.left_relation_embeddings(indices=hr_batch[:, 1]).view(-1, self.embedding_dim, self.embedding_dim) - rel_t = self.right_relation_embeddings(indices=hr_batch[:, 1]) - rel_t = rel_t.view(-1, 1, self.embedding_dim, self.embedding_dim) - t_all = self.entity_embeddings(indices=None).view(1, -1, self.embedding_dim, 1) - - if slice_size is not None: - proj_t_arr = [] - # Project entities - proj_h = rel_h @ h - - for t in torch.split(t_all, slice_size, dim=1): - # Project entities - proj_t = rel_t @ t - proj_t_arr.append(proj_t) - - proj_t = torch.cat(proj_t_arr, dim=1) - - else: - # Project entities - proj_h = rel_h @ h - proj_t = rel_t @ t_all - - scores = -torch.norm(proj_h[:, None, :, 0] - proj_t[:, :, :, 0], dim=-1, p=self.scoring_fct_norm) - - return scores - - def score_h(self, rt_batch: torch.LongTensor, slice_size: int = None) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h_all = self.entity_embeddings(indices=None).view(1, -1, self.embedding_dim, 1) - rel_h = self.left_relation_embeddings(indices=rt_batch[:, 0]) - rel_h = rel_h.view(-1, 1, self.embedding_dim, self.embedding_dim) - rel_t = self.right_relation_embeddings(indices=rt_batch[:, 0]).view(-1, self.embedding_dim, self.embedding_dim) - t = self.entity_embeddings(indices=rt_batch[:, 1]).view(-1, self.embedding_dim, 1) - - if slice_size is not None: - proj_h_arr = [] - - # Project entities - proj_t = rel_t @ t - - for h in torch.split(h_all, slice_size, dim=1): - # Project entities - proj_h = rel_h @ h - proj_h_arr.append(proj_h) - - proj_h = torch.cat(proj_h_arr, dim=1) - else: - # Project entities - proj_h = rel_h @ h_all - proj_t = rel_t @ t - - scores = -torch.norm(proj_h[:, :, :, 0] - proj_t[:, None, :, 0], dim=-1, p=self.scoring_fct_norm) - - return scores diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index bafd6b5d9f..9670cb9226 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -904,3 +904,41 @@ def simple_interaction( min_, max_ = clamp scores = scores.clamp(min=min_, max=max_) return scores + + +def structured_embedding_interaction( + h: torch.FloatTensor, + r_h: torch.FloatTensor, + r_t: torch.FloatTensor, + t: torch.FloatTensor, + p: int, + power_norm: bool = False, +) -> torch.FloatTensor: + r""" + Evaluate the Structured Embedding interaction function. + + .. math :: + f(h, r, t) = \|R_h r - R_t t\| + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r_h: shape: (batch_size, num_relations, dim, rel_dim) + The relation-specific head projection. + :param r_t: shape: (batch_size, num_relations, dim, rel_dim) + The relation-specific tail projection. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + :param p: + The p for the norm. cf. torch.norm. + :param power_norm: + Whether to return the powered norm. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + return negative_norm_of_sum( + _extended_einsum("brde,bhd->bhre", r_h, h).unsqueeze(dim=3), + -_extended_einsum("brde,btd->brte", r_t, t).unsqueeze(dim=1), + p=p, + power_norm=power_norm, + ) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index f903643abf..ba30a0c361 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -698,3 +698,36 @@ class RESCALInteractionFunction(FunctionalInteractionFunction): """RESCAL interaction function.""" interaction = pykeen_functional.rescal_interaction + + +class StructuredEmbeddingInteractionFunction(InteractionFunction): + """Interaction function of Structured Embedding.""" + + def __init__( + self, + p: int, + power_norm: bool = False + ): + super().__init__() + self.p = p + self.power_norm = power_norm + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + dim = h.shape[-1] + rh, rt = r.split(dim, dim=-1) + rh = rh.view(rh.shape[:-1], dim, dim) + rt = rt.view(rt.shape[:-1], dim, dim) + return pykeen_functional.structured_embedding_interaction( + h=h, + r_h=rh, + r_t=rt, + t=t, + p=self.p, + power_norm=self.power_norm, + ) From c7192bbb0e3bbcb5dca54e5b2ce179517b40f2b5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 19:38:37 +0100 Subject: [PATCH 150/690] Fix interaction function for SE --- src/pykeen/nn/modules.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 078dff80a4..b9b731f105 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -717,9 +717,9 @@ def forward( **kwargs, ) -> torch.FloatTensor: dim = h.shape[-1] - rh, rt = r.split(dim, dim=-1) - rh = rh.view(rh.shape[:-1], dim, dim) - rt = rt.view(rt.shape[:-1], dim, dim) + rh, rt = r.split(dim ** 2, dim=-1) + rh = rh.view(*rh.shape[:-1], dim, dim) + rt = rt.view(*rt.shape[:-1], dim, dim) return pykeen_functional.structured_embedding_interaction( h=h, r_h=rh, From 2e81d98ae697aece53c5910f77babcefbb19be73 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 19:40:01 +0100 Subject: [PATCH 151/690] simplify TransE --- src/pykeen/models/unimodal/trans_e.py | 42 ++++----------------------- 1 file changed, 5 insertions(+), 37 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 2a89a0b9e4..a460cf6d5b 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -2,26 +2,25 @@ """TransE.""" -from typing import Optional, Tuple +from typing import Optional -import torch.autograd from torch.nn import functional -from ..base import EntityRelationEmbeddingModel +from .. import SimpleVectorEntityRelationEmbeddingModel from ...losses import Loss from ...nn.init import xavier_uniform_ from ...nn.modules import TranslationalInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint -from ...utils import compose, get_hr_indices, get_hrt_indices, get_ht_indices, get_rt_indices +from ...utils import compose __all__ = [ 'TransE', ] -class TransE(EntityRelationEmbeddingModel): +class TransE(SimpleVectorEntityRelationEmbeddingModel): r"""TransE models relations as a translation from head to tail entities in :math:`\textbf{e}` [bordes2013]_. .. math:: @@ -69,6 +68,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, + interaction_function=TranslationalInteractionFunction(p=scoring_fct_norm), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, @@ -82,35 +82,3 @@ def __init__( ), entity_constrainer=functional.normalize, ) - self.interaction_function = TranslationalInteractionFunction(p=scoring_fct_norm) - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_hrt_indices(hrt_batch) - h, r, t = self._get_hrt(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) - return self.interaction_function(h=h, r=r, t=t).view(-1, 1) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_hr_indices(hr_batch) - h, r, t = self._get_hrt(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) - return self.interaction_function(h=h, r=r, t=t).view(hr_batch.shape[0], self.num_entities) - - def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_ht_indices(ht_batch) - h, r, t = self._get_hrt(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) - return self.interaction_function(h=h, r=r, t=t) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_rt_indices(rt_batch) - h, r, t = self._get_hrt(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) - return self.interaction_function(h=h, r=r, t=t).view(rt_batch.shape[0], self.num_entities) - - def _get_hrt( - self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], - ) -> Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]: - h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) - t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - return h, r, t From 47e42bb91e6e4e472b0ede7d66f2d2624db24f9e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 19:50:25 +0100 Subject: [PATCH 152/690] Extract functional TransH --- src/pykeen/models/unimodal/trans_h.py | 54 ++++++++------------------- src/pykeen/nn/functional.py | 29 ++++++++++++++ 2 files changed, 45 insertions(+), 38 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index da86535acf..6ccafa7fe4 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -9,7 +9,7 @@ from ..base import EntityRelationEmbeddingModel from ...losses import Loss -from ...nn import Embedding +from ...nn import Embedding, functional as F from ...nn.modules import TranslationalInteractionFunction from ...regularizers import Regularizer, TransHRegularizer from ...triples import TriplesFactory @@ -121,49 +121,27 @@ def regularize_if_necessary(self) -> None: def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings - h = self.entity_embeddings(indices=hrt_batch[:, 0]) - d_r = self.relation_embeddings(indices=hrt_batch[:, 1]) - w_r = self.normal_vector_embeddings(indices=hrt_batch[:, 1]) - t = self.entity_embeddings(indices=hrt_batch[:, 2]) - - # Project to hyperplane - ph = h - torch.sum(w_r * h, dim=-1, keepdim=True) * w_r - pt = t - torch.sum(w_r * t, dim=-1, keepdim=True) * w_r - - # Regularization term + h = self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 0]) + d_r = self.relation_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 1]) + w_r = self.normal_vector_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 1]) + t = self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 2]) self.regularize_if_necessary() - - return self.interaction_function(h=ph, r=d_r, t=pt, dim=-1, keepdim=True) + return F.transh_interaction(h, w_r, d_r, t).view(hrt_batch.shape[0], 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings - h = self.entity_embeddings(indices=hr_batch[:, 0]) - d_r = self.relation_embeddings(indices=hr_batch[:, 1]) - w_r = self.normal_vector_embeddings(indices=hr_batch[:, 1]) - t = self.entity_embeddings(indices=None) - - # Project to hyperplane - ph = h - torch.sum(w_r * h, dim=-1, keepdim=True) * w_r - pt = t[None, :, :] - torch.sum(w_r[:, None, :] * t[None, :, :], dim=-1, keepdim=True) * w_r[:, None, :] - - # Regularization term + h = self.entity_embeddings.get_in_canonical_shape(indices=hr_batch[:, 0]) + d_r = self.relation_embeddings.get_in_canonical_shape(indices=hr_batch[:, 1]) + w_r = self.normal_vector_embeddings.get_in_canonical_shape(indices=hr_batch[:, 1]) + t = self.entity_embeddings.get_in_canonical_shape(indices=None) self.regularize_if_necessary() - - return self.interaction_function(h=ph[:, None, :], r=d_r[:, None, :], t=pt, dim=-1, keepdim=False) + return F.transh_interaction(h, w_r, d_r, t).view(hr_batch.shape[0], self.num_entities) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings - h = self.entity_embeddings(indices=None) - rel_id = rt_batch[:, 0] - d_r = self.relation_embeddings(indices=rel_id) - w_r = self.normal_vector_embeddings(indices=rel_id) - t = self.entity_embeddings(indices=rt_batch[:, 1]) - - # Project to hyperplane - ph = h[None, :, :] - torch.sum(w_r[:, None, :] * h[None, :, :], dim=-1, keepdim=True) * w_r[:, None, :] - pt = t - torch.sum(w_r * t, dim=-1, keepdim=True) * w_r - - # Regularization term + h = self.entity_embeddings.get_in_canonical_shape(indices=None) + d_r = self.relation_embeddings.get_in_canonical_shape(indices=rt_batch[:, 0]) + w_r = self.normal_vector_embeddings.get_in_canonical_shape(indices=rt_batch[:, 0]) + t = self.entity_embeddings.get_in_canonical_shape(indices=rt_batch[:, 1]) self.regularize_if_necessary() - - return self.interaction_function(h=ph, r=d_r[:, None, :], t=pt[:, None, :], dim=-1, keepdim=False) + return F.transh_interaction(h, w_r, d_r, t).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 9670cb9226..c9632396e8 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -942,3 +942,32 @@ def structured_embedding_interaction( p=p, power_norm=power_norm, ) + + +def transh_interaction( + h: torch.FloatTensor, + w_r: torch.FloatTensor, + d_r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """ + Evaluate the DistMult interaction function. + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param w_r: shape: (batch_size, num_relations, dim) + The relation normal vector representations. + :param d_r: shape: (batch_size, num_relations, dim) + The relation difference vector representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + # Project to hyperplane + return _translational_interaction( + h=(h.unsqueeze(dim=2) - _extended_einsum("bhd,brd,bre->bhre", h, w_r, w_r)).unsqueeze(dim=3), + r=d_r.view(d_r.shape[0], 1, d_r.shape[1], 1, d_r.shape[2]), + t=(t.unsqueeze(dim=1) - _extended_einsum("btd,brd,bre->brte", t, w_r, w_r)).unsqueeze(dim=1), + ) From eb829c6887b5722595ea3069642ab5354ecb134c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 19:51:51 +0100 Subject: [PATCH 153/690] fix TransH --- src/pykeen/models/unimodal/trans_h.py | 9 +++------ src/pykeen/nn/functional.py | 8 ++++++++ 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 6ccafa7fe4..1b92845cbe 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -10,7 +10,6 @@ from ..base import EntityRelationEmbeddingModel from ...losses import Loss from ...nn import Embedding, functional as F -from ...nn.modules import TranslationalInteractionFunction from ...regularizers import Regularizer, TransHRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -89,8 +88,6 @@ def __init__( regularizer=regularizer, ) - self.interaction_function = TranslationalInteractionFunction(p=scoring_fct_norm) - # embeddings self.normal_vector_embeddings = Embedding.init_with_device( num_embeddings=triples_factory.num_relations, @@ -126,7 +123,7 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: w_r = self.normal_vector_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 1]) t = self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 2]) self.regularize_if_necessary() - return F.transh_interaction(h, w_r, d_r, t).view(hrt_batch.shape[0], 1) + return F.transh_interaction(h, w_r, d_r, t, p=self.scoring_fct_norm).view(hrt_batch.shape[0], 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings @@ -135,7 +132,7 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 w_r = self.normal_vector_embeddings.get_in_canonical_shape(indices=hr_batch[:, 1]) t = self.entity_embeddings.get_in_canonical_shape(indices=None) self.regularize_if_necessary() - return F.transh_interaction(h, w_r, d_r, t).view(hr_batch.shape[0], self.num_entities) + return F.transh_interaction(h, w_r, d_r, t, p=self.scoring_fct_norm).view(hr_batch.shape[0], self.num_entities) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Get embeddings @@ -144,4 +141,4 @@ def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 w_r = self.normal_vector_embeddings.get_in_canonical_shape(indices=rt_batch[:, 0]) t = self.entity_embeddings.get_in_canonical_shape(indices=rt_batch[:, 1]) self.regularize_if_necessary() - return F.transh_interaction(h, w_r, d_r, t).view(rt_batch.shape[0], self.num_entities) + return F.transh_interaction(h, w_r, d_r, t, p=self.scoring_fct_norm).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index c9632396e8..15ba70df2d 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -949,6 +949,8 @@ def transh_interaction( w_r: torch.FloatTensor, d_r: torch.FloatTensor, t: torch.FloatTensor, + p: int, + power_norm: bool = False, ) -> torch.FloatTensor: """ Evaluate the DistMult interaction function. @@ -961,6 +963,10 @@ def transh_interaction( The relation difference vector representations. :param t: shape: (batch_size, num_tails, dim) The tail representations. + :param p: + The p for the norm. cf. torch.norm. + :param power_norm: + Whether to return |x-y|_p^p. :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. @@ -970,4 +976,6 @@ def transh_interaction( h=(h.unsqueeze(dim=2) - _extended_einsum("bhd,brd,bre->bhre", h, w_r, w_r)).unsqueeze(dim=3), r=d_r.view(d_r.shape[0], 1, d_r.shape[1], 1, d_r.shape[2]), t=(t.unsqueeze(dim=1) - _extended_einsum("btd,brd,bre->brte", t, w_r, w_r)).unsqueeze(dim=1), + p=p, + power_norm=power_norm, ) From 74a6d20d58b82867ad927e2efd1fe60ed5e590b8 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 19:52:47 +0100 Subject: [PATCH 154/690] Update views @mberr please check :) --- src/pykeen/models/base.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 9e85b837fa..80435545fc 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1364,7 +1364,9 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 h_indices, r_indices, t_indices = get_hr_indices(hr_batch) - return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(-1, self.num_entities) + return self( + h_indices=h_indices, r_indices=r_indices, t_indices=t_indices, + ).view(hr_batch.shape[0], self.num_entities) def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 h_indices, r_indices, t_indices = get_ht_indices(ht_batch) @@ -1372,7 +1374,9 @@ def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 h_indices, r_indices, t_indices = get_rt_indices(rt_batch) - return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(-1, self.num_entities) + return self( + h_indices=h_indices, r_indices=r_indices, t_indices=t_indices, + ).view(rt_batch.shape[0], self.num_entities) class SimpleVectorEntityRelationEmbeddingModel( From c0d5849702fc51f2ab6dfb73ed977826b083cf05 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 19:53:00 +0100 Subject: [PATCH 155/690] Better exception type --- src/pykeen/nn/functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 15ba70df2d..7fb4488efc 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -743,7 +743,7 @@ def kg2e_interaction( The scores. """ if similarity not in KG2E_SIMILARITIES: - raise ValueError(similarity) + raise KeyError(similarity) similarity = _KG2E_SIMILARITIES[similarity] # Compute entity distribution e_mean = h_mean.unsqueeze(dim=2) - t_mean.unsqueeze(dim=1) From 39671f515f358a77cbe400fa993a1517e6e18d3e Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 19:53:10 +0100 Subject: [PATCH 156/690] Improve type annotation and unpacking --- src/pykeen/nn/functional.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 7fb4488efc..a49ba92d37 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -24,7 +24,7 @@ ] -def _extract_sizes(h, r, t): +def _extract_sizes(h, r, t) -> Tuple[int, int, int, int, int]: num_heads, num_relations, num_tails = [xx.shape[1] for xx in (h, r, t)] d_e = h.shape[-1] d_r = r.shape[-1] @@ -278,7 +278,7 @@ def convkb_interaction( The scores. """ # bind sizes - num_heads, num_relations, num_tails, embedding_dim = _extract_sizes(h, r, t)[:4] + num_heads, num_relations, num_tails, embedding_dim, _ = _extract_sizes(h, r, t) # decompose convolution for faster computation in 1-n case num_filters = conv.weight.shape[0] @@ -329,7 +329,7 @@ def ermlp_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - num_heads, num_relations, num_tails, embedding_dim = _extract_sizes(h, r, t)[:4] + num_heads, num_relations, num_tails, embedding_dim, _ = _extract_sizes(h, r, t) hidden_dim = hidden.weight.shape[0] assert embedding_dim % 3 == 0 embedding_dim = embedding_dim // 3 @@ -534,7 +534,7 @@ def translational_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - num_heads, num_relations, num_tails, d_e = _extract_sizes(h, r, t)[:4] + num_heads, num_relations, num_tails, d_e, _ = _extract_sizes(h, r, t) h = h.view(-1, num_heads, 1, 1, d_e) r = r.view(-1, 1, num_relations, 1, d_e) t = t.view(-1, 1, 1, num_tails, d_e) @@ -837,7 +837,7 @@ def proje_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - num_heads, num_relations, num_tails, dim = _extract_sizes(h, r, t)[:4] + num_heads, num_relations, num_tails, dim, _ = _extract_sizes(h, r, t) # global projections h = h * d_e.view(1, 1, dim) r = r * d_r.view(1, 1, dim) From fed0445664420f23f2bb388578a16d6803c7c566 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 20:01:25 +0100 Subject: [PATCH 157/690] Use self-documenting functions for slicing @mberr do these need to have view() functions added on them? --- src/pykeen/models/unimodal/simple.py | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index c933a57399..7f168385d4 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -13,6 +13,7 @@ from ...regularizers import PowerSumRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint +from ...utils import get_hr_indices, get_hrt_indices, get_ht_indices, get_rt_indices __all__ = [ 'SimplE', @@ -147,10 +148,17 @@ def _single_forward( return self.interaction_function(h, r, t) def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self(h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2]).view(-1, 1) + h_indices, r_indices, t_indices = get_hrt_indices(hrt_batch) + return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(-1, 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], t_indices=None) + h_indices, r_indices, t_indices = get_hr_indices(hr_batch) + return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) + + def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + h_indices, r_indices, t_indices = get_ht_indices(ht_batch) + return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self(h_indices=None, r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]) + h_indices, r_indices, t_indices = get_rt_indices(rt_batch) + return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) From 13c4bcff0e5cd2cd0b446a2f8b1c6659fea2c17a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 20:14:07 +0100 Subject: [PATCH 158/690] Add yet another base class ;-) --- src/pykeen/models/base.py | 75 ++++++++++++++++++++++++++++++++++++++- src/pykeen/nn/emb.py | 46 ++++++++++++++++++++++-- 2 files changed, 118 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 9e85b837fa..962c4b781b 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1,7 +1,6 @@ # -*- coding: utf-8 -*- """Base module for all KGE models.""" - import functools import inspect import itertools as itt @@ -17,6 +16,7 @@ from ..losses import Loss, MarginRankingLoss, NSSALoss from ..nn import Embedding +from ..nn.emb import EmbeddingSpecification from ..nn.modules import InteractionFunction from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory @@ -1449,3 +1449,76 @@ def __init__( relation_constrainer=relation_constrainer, relation_constrainer_kwargs=relation_constrainer_kwargs, ) + + +class AbstractModel(Model): + def __init__( + self, + interaction_function: InteractionFunction, + triples_factory: TriplesFactory, + entity_embeddings: Optional[Mapping[str, EmbeddingSpecification]] = None, + relation_embeddings: Optional[Mapping[str, EmbeddingSpecification]] = None, + loss: Optional[Loss] = None, + predict_with_sigmoid: bool = False, + automatic_memory_optimization: Optional[bool] = None, + preferred_device: DeviceHint = None, + random_seed: Optional[int] = None, + regularizer: Optional[Regularizer] = None, + ): + super().__init__( + triples_factory=triples_factory, + loss=loss, + predict_with_sigmoid=predict_with_sigmoid, + automatic_memory_optimization=automatic_memory_optimization, + preferred_device=preferred_device, + random_seed=random_seed, + regularizer=regularizer, + ) + self.interaction_function = interaction_function + self.entity_embeddings = nn.ModuleDict({ + key: Embedding.from_specification(num=self.num_entities, specification=spec) + for key, spec in entity_embeddings.items() + }) + self.relation_embeddings = nn.ModuleDict({ + key: Embedding.from_specification(num=self.num_relations, specification=spec) + for key, spec in relation_embeddings.items() + }) + + def _get_representations( + self, + indices: Optional[torch.LongTensor], + source: Mapping[str, Embedding], + prefix: str, + ) -> Mapping[str, torch.FloatTensor]: + return { + prefix + "_" + key if len(key) > 0 else prefix: embedding.get_in_canonical_shape(indices=indices, reshape_dim=embedding.shape) + for key, embedding in source + } + + def _score( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: + return self.interaction_function( + **self._get_representations(indices=h_indices, source=self.entity_embeddings, prefix="h"), + **self._get_representations(indices=r_indices, source=self.relation_embeddings, prefix="r"), + **self._get_representations(indices=t_indices, source=self.entity_embeddings, prefix="t"), + ) + + def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + h_indices, r_indices, t_indices = get_hrt_indices(hrt_batch) + return self._score(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(hrt_batch.shape[0], 1) + + def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + h_indices, r_indices, t_indices = get_hr_indices(hr_batch) + return self._score(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(hr_batch.shape[0], self.num_entities) + + def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + h_indices, r_indices, t_indices = get_ht_indices(ht_batch) + return self._score(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(ht_batch.shape[0], self.num_relations) + + def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + h_indices, r_indices, t_indices = get_rt_indices(rt_batch) + return self._score(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index 44446cc11e..7b5d64be27 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -1,10 +1,11 @@ # -*- coding: utf-8 -*- """Embedding modules.""" - +import dataclasses import functools from typing import Any, Mapping, Optional, Sequence +import numpy import torch import torch.nn from torch import nn @@ -42,6 +43,19 @@ def post_parameter_update(self): """Apply constraints which should not be included in gradients.""" +@dataclasses.dataclass +class EmbeddingSpecification: + shape: Sequence[int] + initializer: Initializer + constrainer: Optional[Constrainer] = None + normalizer: Optional[Normalizer] = None + # regularizer: Optional[Regularizer] + initializer_kwargs: Optional[Mapping[str, Any]] = None + normalizer_kwargs: Optional[Mapping[str, Any]] = None + constrainer_kwargs: Optional[Mapping[str, Any]] = None + # regularizer_kwargs: Optional[Mapping[str, Any]] = None + + class Embedding(RepresentationModule): """Trainable embeddings. @@ -52,7 +66,8 @@ class Embedding(RepresentationModule): def __init__( self, num_embeddings: int, - embedding_dim: int, + embedding_dim: Optional[int], + shape: Optional[Sequence[int]] = None, initializer: Optional[Initializer] = None, initializer_kwargs: Optional[Mapping[str, Any]] = None, normalizer: Optional[Normalizer] = None, @@ -99,11 +114,32 @@ def __init__( self.normalizer = functools.partial(normalizer, **normalizer_kwargs) else: self.normalizer = normalizer + if shape is not None: + embedding_dim = numpy.prod(shape) + self.shape = shape self._embeddings = torch.nn.Embedding( num_embeddings=num_embeddings, embedding_dim=embedding_dim, ) + @classmethod + def from_specification( + cls, + num: int, + specification: EmbeddingSpecification, + ) -> "Embedding": + return Embedding( + num_embeddings=num, + embedding_dim=None, + shape=specification.shape, + initializer=specification.initializer, + initializer_kwargs=specification.initializer_kwargs, + normalizer=specification.normalizer, + normalizer_kwargs=specification.normalizer_kwargs, + constrainer=specification.constrainer, + constrainer_kwargs=specification.constrainer_kwargs, + ) + @classmethod def init_with_device( cls, @@ -151,6 +187,12 @@ def embedding_dim(self) -> int: # noqa: D401 """The representation dimension.""" return self._embeddings.embedding_dim + @property + def shape(self) -> Sequence[int]: + if self.shape is None: + raise AssertionError + return self.shape + def reset_parameters(self) -> None: # noqa: D102 # initialize weights in-place self._embeddings.weight.data = self.initializer(self._embeddings.weight.data) From 86d22482da08de8ca29b91e2220df1d28bfd33c4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 20:19:23 +0100 Subject: [PATCH 159/690] Revert "Add yet another base class ;-)" This reverts commit 13c4bcff --- src/pykeen/models/base.py | 75 +-------------------------------------- src/pykeen/nn/emb.py | 46 ++---------------------- 2 files changed, 3 insertions(+), 118 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 02201bd781..80435545fc 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- """Base module for all KGE models.""" + import functools import inspect import itertools as itt @@ -16,7 +17,6 @@ from ..losses import Loss, MarginRankingLoss, NSSALoss from ..nn import Embedding -from ..nn.emb import EmbeddingSpecification from ..nn.modules import InteractionFunction from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory @@ -1453,76 +1453,3 @@ def __init__( relation_constrainer=relation_constrainer, relation_constrainer_kwargs=relation_constrainer_kwargs, ) - - -class AbstractModel(Model): - def __init__( - self, - interaction_function: InteractionFunction, - triples_factory: TriplesFactory, - entity_embeddings: Optional[Mapping[str, EmbeddingSpecification]] = None, - relation_embeddings: Optional[Mapping[str, EmbeddingSpecification]] = None, - loss: Optional[Loss] = None, - predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, - preferred_device: DeviceHint = None, - random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, - ): - super().__init__( - triples_factory=triples_factory, - loss=loss, - predict_with_sigmoid=predict_with_sigmoid, - automatic_memory_optimization=automatic_memory_optimization, - preferred_device=preferred_device, - random_seed=random_seed, - regularizer=regularizer, - ) - self.interaction_function = interaction_function - self.entity_embeddings = nn.ModuleDict({ - key: Embedding.from_specification(num=self.num_entities, specification=spec) - for key, spec in entity_embeddings.items() - }) - self.relation_embeddings = nn.ModuleDict({ - key: Embedding.from_specification(num=self.num_relations, specification=spec) - for key, spec in relation_embeddings.items() - }) - - def _get_representations( - self, - indices: Optional[torch.LongTensor], - source: Mapping[str, Embedding], - prefix: str, - ) -> Mapping[str, torch.FloatTensor]: - return { - prefix + "_" + key if len(key) > 0 else prefix: embedding.get_in_canonical_shape(indices=indices, reshape_dim=embedding.shape) - for key, embedding in source - } - - def _score( - self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], - ) -> torch.FloatTensor: - return self.interaction_function( - **self._get_representations(indices=h_indices, source=self.entity_embeddings, prefix="h"), - **self._get_representations(indices=r_indices, source=self.relation_embeddings, prefix="r"), - **self._get_representations(indices=t_indices, source=self.entity_embeddings, prefix="t"), - ) - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_hrt_indices(hrt_batch) - return self._score(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(hrt_batch.shape[0], 1) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_hr_indices(hr_batch) - return self._score(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(hr_batch.shape[0], self.num_entities) - - def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_ht_indices(ht_batch) - return self._score(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(ht_batch.shape[0], self.num_relations) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_rt_indices(rt_batch) - return self._score(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index 7b5d64be27..44446cc11e 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -1,11 +1,10 @@ # -*- coding: utf-8 -*- """Embedding modules.""" -import dataclasses + import functools from typing import Any, Mapping, Optional, Sequence -import numpy import torch import torch.nn from torch import nn @@ -43,19 +42,6 @@ def post_parameter_update(self): """Apply constraints which should not be included in gradients.""" -@dataclasses.dataclass -class EmbeddingSpecification: - shape: Sequence[int] - initializer: Initializer - constrainer: Optional[Constrainer] = None - normalizer: Optional[Normalizer] = None - # regularizer: Optional[Regularizer] - initializer_kwargs: Optional[Mapping[str, Any]] = None - normalizer_kwargs: Optional[Mapping[str, Any]] = None - constrainer_kwargs: Optional[Mapping[str, Any]] = None - # regularizer_kwargs: Optional[Mapping[str, Any]] = None - - class Embedding(RepresentationModule): """Trainable embeddings. @@ -66,8 +52,7 @@ class Embedding(RepresentationModule): def __init__( self, num_embeddings: int, - embedding_dim: Optional[int], - shape: Optional[Sequence[int]] = None, + embedding_dim: int, initializer: Optional[Initializer] = None, initializer_kwargs: Optional[Mapping[str, Any]] = None, normalizer: Optional[Normalizer] = None, @@ -114,32 +99,11 @@ def __init__( self.normalizer = functools.partial(normalizer, **normalizer_kwargs) else: self.normalizer = normalizer - if shape is not None: - embedding_dim = numpy.prod(shape) - self.shape = shape self._embeddings = torch.nn.Embedding( num_embeddings=num_embeddings, embedding_dim=embedding_dim, ) - @classmethod - def from_specification( - cls, - num: int, - specification: EmbeddingSpecification, - ) -> "Embedding": - return Embedding( - num_embeddings=num, - embedding_dim=None, - shape=specification.shape, - initializer=specification.initializer, - initializer_kwargs=specification.initializer_kwargs, - normalizer=specification.normalizer, - normalizer_kwargs=specification.normalizer_kwargs, - constrainer=specification.constrainer, - constrainer_kwargs=specification.constrainer_kwargs, - ) - @classmethod def init_with_device( cls, @@ -187,12 +151,6 @@ def embedding_dim(self) -> int: # noqa: D401 """The representation dimension.""" return self._embeddings.embedding_dim - @property - def shape(self) -> Sequence[int]: - if self.shape is None: - raise AssertionError - return self.shape - def reset_parameters(self) -> None: # noqa: D102 # initialize weights in-place self._embeddings.weight.data = self.initializer(self._embeddings.weight.data) From 8fabb606a753d5740fd0b312673094c0dadad3b3 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 20:29:30 +0100 Subject: [PATCH 160/690] Add abstract score function, and implementations for score_* --- src/pykeen/models/base.py | 79 ++++++++++++++++++++++----------------- 1 file changed, 45 insertions(+), 34 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 80435545fc..73a2edb3b3 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -928,6 +928,31 @@ def _compute_loss( ) return self.loss(tensor_1, tensor_2) + self.regularizer.term + @abstractmethod + def score( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: + """Forward pass. + + This method takes head, relation and tail indices and calculates the corresponding score. + + All indices which are not None, have to be either 1-element or have the same shape, which is the batch size. + + :param h_indices: + The head indices. None indicates to use all. + :param r_indices: + The relation indices. None indicates to use all. + :param t_indices: + The tail indices. None indicates to use all. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The score for each triple. + """ + raise NotImplementedError + @abstractmethod def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: """Forward pass. @@ -941,7 +966,11 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: :return: shape: (batch_size, 1), dtype: float The score for each triple. """ - raise NotImplementedError + return self.score( + h_indices=hrt_batch[:, 0], + r_indices=hrt_batch[:, 1], + t_indices=hrt_batch[:, 2], + ).view(hrt_batch.shape[0], 1) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: """Forward pass using right side (tail) prediction. @@ -954,17 +983,11 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: :return: shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails. """ - logger.warning( - 'Calculations will fall back to using the score_hrt method, since this model does not have a specific ' - 'score_t function. This might cause the calculations to take longer than necessary.', - ) - # Extend the hr_batch such that each (h, r) pair is combined with all possible tails - hrt_batch = _extend_batch(batch=hr_batch, all_ids=list(self.triples_factory.entity_to_id.values()), dim=2) - # Calculate the scores for each (h, r, t) triple using the generic interaction function - expanded_scores = self.score_hrt(hrt_batch=hrt_batch) - # Reshape the scores to match the pre-defined output shape of the score_t function. - scores = expanded_scores.view(hr_batch.shape[0], -1) - return scores + return self.score( + h_indices=hr_batch[:, 0], + r_indices=hr_batch[:, 1], + t_indices=None, + ).view(hr_batch.shape[0], self.num_entities) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: """Forward pass using left side (head) prediction. @@ -977,17 +1000,11 @@ def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: :return: shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads. """ - logger.warning( - 'Calculations will fall back to using the score_hrt method, since this model does not have a specific ' - 'score_h function. This might cause the calculations to take longer than necessary.', - ) - # Extend the rt_batch such that each (r, t) pair is combined with all possible heads - hrt_batch = _extend_batch(batch=rt_batch, all_ids=list(self.triples_factory.entity_to_id.values()), dim=0) - # Calculate the scores for each (h, r, t) triple using the generic interaction function - expanded_scores = self.score_hrt(hrt_batch=hrt_batch) - # Reshape the scores to match the pre-defined output shape of the score_h function. - scores = expanded_scores.view(rt_batch.shape[0], -1) - return scores + return self.score( + h_indices=None, + r_indices=rt_batch[:, 0], + t_indices=rt_batch[:, 1], + ).view(rt_batch.shape[0], self.num_entities) def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: """Forward pass using middle (relation) prediction. @@ -1000,17 +1017,11 @@ def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: :return: shape: (batch_size, num_relations), dtype: float For each h-t pair, the scores for all possible relations. """ - logger.warning( - 'Calculations will fall back to using the score_hrt method, since this model does not have a specific ' - 'score_r function. This might cause the calculations to take longer than necessary.', - ) - # Extend the ht_batch such that each (h, t) pair is combined with all possible relations - hrt_batch = _extend_batch(batch=ht_batch, all_ids=list(self.triples_factory.relation_to_id.values()), dim=1) - # Calculate the scores for each (h, r, t) triple using the generic interaction function - expanded_scores = self.score_hrt(hrt_batch=hrt_batch) - # Reshape the scores to match the pre-defined output shape of the score_r function. - scores = expanded_scores.view(ht_batch.shape[0], -1) - return scores + return self.score( + h_indices=ht_batch[:, 0], + r_indices=None, + t_indices=ht_batch[:, 1], + ).view(ht_batch.shape[0], self.num_relations) def get_grad_params(self) -> Iterable[nn.Parameter]: """Get the parameters that require gradients.""" From c7406f7e56e6a1e54a1ca1ae1a8d43e0a5c2a9d6 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 20:31:11 +0100 Subject: [PATCH 161/690] Use abstract score method --- src/pykeen/models/unimodal/trans_h.py | 36 +++++++++------------------ 1 file changed, 12 insertions(+), 24 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 1b92845cbe..9013d83e16 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -87,6 +87,7 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) + self.scoring_fct_norm = scoring_fct_norm # embeddings self.normal_vector_embeddings = Embedding.init_with_device( @@ -116,29 +117,16 @@ def regularize_if_necessary(self) -> None: self.relation_embeddings(indices=None), ) - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 0]) - d_r = self.relation_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 1]) - w_r = self.normal_vector_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 1]) - t = self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 2]) - self.regularize_if_necessary() - return F.transh_interaction(h, w_r, d_r, t, p=self.scoring_fct_norm).view(hrt_batch.shape[0], 1) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings.get_in_canonical_shape(indices=hr_batch[:, 0]) - d_r = self.relation_embeddings.get_in_canonical_shape(indices=hr_batch[:, 1]) - w_r = self.normal_vector_embeddings.get_in_canonical_shape(indices=hr_batch[:, 1]) - t = self.entity_embeddings.get_in_canonical_shape(indices=None) - self.regularize_if_necessary() - return F.transh_interaction(h, w_r, d_r, t, p=self.scoring_fct_norm).view(hr_batch.shape[0], self.num_entities) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 + def score( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: # noqa: D102 # Get embeddings - h = self.entity_embeddings.get_in_canonical_shape(indices=None) - d_r = self.relation_embeddings.get_in_canonical_shape(indices=rt_batch[:, 0]) - w_r = self.normal_vector_embeddings.get_in_canonical_shape(indices=rt_batch[:, 0]) - t = self.entity_embeddings.get_in_canonical_shape(indices=rt_batch[:, 1]) + h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) + d_r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) + w_r = self.normal_vector_embeddings.get_in_canonical_shape(indices=r_indices) + t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) self.regularize_if_necessary() - return F.transh_interaction(h, w_r, d_r, t, p=self.scoring_fct_norm).view(rt_batch.shape[0], self.num_entities) + return F.transh_interaction(h, w_r, d_r, t, p=self.scoring_fct_norm) From 43b37e5ecbb625f3825c1e12af260758c7429f67 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 20:33:02 +0100 Subject: [PATCH 162/690] Use generic score method for conve --- src/pykeen/models/unimodal/conv_e.py | 33 ++++++++++------------------ 1 file changed, 11 insertions(+), 22 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index f4c26f4888..71225fbb82 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -173,26 +173,15 @@ def _reset_parameters_(self): # noqa: D102 self.bias_term.reset_parameters() self.interaction_function.reset_parameters() - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=hrt_batch[:, 0]) - r = self.relation_embeddings(indices=hrt_batch[:, 1]) - t = self.entity_embeddings(indices=hrt_batch[:, 2]) - t_bias = self.bias_term(indices=hrt_batch[:, 2]) + def score( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: # noqa: D102 + h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) + r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) + t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) + t_bias = self.bias_term.get_in_canonical_shape(indices=t_indices) self.regularize_if_necessary(h, r, t) - return self.interaction_function.score_hrt(h=h, r=r, t=t, t_bias=t_bias) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=hr_batch[:, 0]) - r = self.relation_embeddings(indices=hr_batch[:, 1]) - all_entities = self.entity_embeddings(indices=None) - t_bias = self.bias_term(indices=None) - self.regularize_if_necessary(h, r, all_entities) - return self.interaction_function.score_t(h=h, r=r, all_entities=all_entities, t_bias=t_bias) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - all_entities = self.entity_embeddings(indices=None) - r = self.relation_embeddings(indices=rt_batch[:, 0]) - t = self.entity_embeddings(indices=rt_batch[:, 1]) - t_bias = self.bias_term(indices=rt_batch[:, 1]) - self.regularize_if_necessary(all_entities, r, t) - return self.interaction_function.score_h(all_entities=all_entities, r=r, t=t, t_bias=t_bias) + return self.interaction_function(h=h, r=r, t=t, t_bias=t_bias) From a598869fb3440c2df8596a2f7954d41c14fab80a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 20:33:46 +0100 Subject: [PATCH 163/690] Rename generic score to forward --- src/pykeen/models/base.py | 10 +++++----- src/pykeen/models/unimodal/conv_e.py | 2 +- src/pykeen/models/unimodal/trans_h.py | 2 +- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 73a2edb3b3..a92d2c7d19 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -929,7 +929,7 @@ def _compute_loss( return self.loss(tensor_1, tensor_2) + self.regularizer.term @abstractmethod - def score( + def forward( self, h_indices: Optional[torch.LongTensor], r_indices: Optional[torch.LongTensor], @@ -966,7 +966,7 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: :return: shape: (batch_size, 1), dtype: float The score for each triple. """ - return self.score( + return self( h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2], @@ -983,7 +983,7 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: :return: shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails. """ - return self.score( + return self( h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], t_indices=None, @@ -1000,7 +1000,7 @@ def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: :return: shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads. """ - return self.score( + return self( h_indices=None, r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1], @@ -1017,7 +1017,7 @@ def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: :return: shape: (batch_size, num_relations), dtype: float For each h-t pair, the scores for all possible relations. """ - return self.score( + return self( h_indices=ht_batch[:, 0], r_indices=None, t_indices=ht_batch[:, 1], diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 71225fbb82..c126194920 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -173,7 +173,7 @@ def _reset_parameters_(self): # noqa: D102 self.bias_term.reset_parameters() self.interaction_function.reset_parameters() - def score( + def forward( self, h_indices: Optional[torch.LongTensor], r_indices: Optional[torch.LongTensor], diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 9013d83e16..a0ece409c7 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -117,7 +117,7 @@ def regularize_if_necessary(self) -> None: self.relation_embeddings(indices=None), ) - def score( + def forward( self, h_indices: Optional[torch.LongTensor], r_indices: Optional[torch.LongTensor], From d59491c944ea5cdd1d06359351b124161b554294 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 20:36:13 +0100 Subject: [PATCH 164/690] Add default implementation for _reset_parameters_ --- src/pykeen/models/base.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index a92d2c7d19..99ab6e75a9 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -351,7 +351,11 @@ def supports_subbatching(self) -> bool: # noqa: D400, D401 @abstractmethod def _reset_parameters_(self): # noqa: D401 """Reset all parameters of the model in-place.""" - raise NotImplementedError + for module in self.modules(): + if module is self: + continue + if hasattr(module, "reset_parameters"): + module.reset_parameters() def reset_parameters_(self) -> 'Model': # noqa: D401 """Reset all parameters of the model and enforce model constraints.""" From 79d7a3d66ff12af37a95336cd9da961a83127b87 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 20:39:31 +0100 Subject: [PATCH 165/690] Use generic method more often --- src/pykeen/models/unimodal/kg2e.py | 16 ----------- src/pykeen/models/unimodal/ntn.py | 28 -------------------- src/pykeen/models/unimodal/simple.py | 17 ------------ src/pykeen/models/unimodal/trans_r.py | 38 ++++++++------------------- 4 files changed, 11 insertions(+), 88 deletions(-) diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 639ba356ef..6174c34a8d 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -166,19 +166,3 @@ def forward( t_mean=mu_t, t_var=sigma_t, ) - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self( - h_indices=hrt_batch[:, 0], - r_indices=hrt_batch[:, 1], - t_indices=hrt_batch[:, 2], - ).view(hrt_batch.shape[0], 1) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self(h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1]).view(hr_batch.shape[0], self.num_entities) - - def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self(h_indices=ht_batch[:, 0], t_indices=ht_batch[:, 1]).view(ht_batch.shape[0], self.num_entities) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self(r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1]).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 7759df466b..f5e64d4eab 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -172,31 +172,3 @@ def forward( scores_arr.append(score) return torch.cat(scores_arr, dim=1) - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self( - h_indices=hrt_batch[:, 0], - r_indices=hrt_batch[:, 1], - t_indices=hrt_batch[:, 2], - ).view(hrt_batch.shape[0], 1) - - def score_t(self, hr_batch: torch.LongTensor, slice_size: int = None) -> torch.FloatTensor: # noqa: D102 - return self( - h_indices=hr_batch[:, 0], - r_indices=hr_batch[:, 1], - slice_size=slice_size, - ).view(hr_batch.shape[0], self.num_entities) - - def score_r(self, ht_batch: torch.LongTensor, slice_size: int = None) -> torch.FloatTensor: # noqa: D102 - return self( - h_indices=ht_batch[:, 0], - t_indices=ht_batch[:, 1], - slice_size=slice_size, - ).view(ht_batch.shape[0], self.num_relations) - - def score_h(self, rt_batch: torch.LongTensor, slice_size: int = None) -> torch.FloatTensor: # noqa: D102 - return self.forward( - r_indices=rt_batch[:, 0], - t_indices=rt_batch[:, 1], - slice_size=slice_size, - ).view(rt_batch.shape[0], self.num_entities) diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index 7f168385d4..35bfc9cc0c 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -13,7 +13,6 @@ from ...regularizers import PowerSumRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint -from ...utils import get_hr_indices, get_hrt_indices, get_ht_indices, get_rt_indices __all__ = [ 'SimplE', @@ -146,19 +145,3 @@ def _single_forward( # Regularization self.regularize_if_necessary(h, r, t) return self.interaction_function(h, r, t) - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_hrt_indices(hrt_batch) - return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(-1, 1) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_hr_indices(hr_batch) - return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) - - def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_ht_indices(ht_batch) - return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_rt_indices(rt_batch) - return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 00d83347c0..19f776a0ad 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -126,30 +126,14 @@ def _reset_parameters_(self): # noqa: D102 super()._reset_parameters_() self.relation_projections.reset_parameters() - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - # TODO switch to self.entity_embeddings.get_in_canonical_shape(indices=hrt_batch[:, 0]) - h = self.entity_embeddings(indices=hrt_batch[:, 0]).unsqueeze(dim=1) - r = self.relation_embeddings(indices=hrt_batch[:, 1]).unsqueeze(dim=1) - t = self.entity_embeddings(indices=hrt_batch[:, 2]).unsqueeze(dim=1) - m_r = self.relation_projections(indices=hrt_batch[:, 1]).view(-1, self.embedding_dim, self.relation_dim) - - return self.interaction_function(h=h, r=r, t=t, m_r=m_r).view(hrt_batch.shape[0], 1) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(indices=hr_batch[:, 0]).unsqueeze(dim=1) - r = self.relation_embeddings(indices=hr_batch[:, 1]).unsqueeze(dim=1) - t = self.entity_embeddings(indices=None).unsqueeze(dim=0) - m_r = self.relation_projections(indices=hr_batch[:, 1]).view(-1, self.embedding_dim, self.relation_dim) - - return self.interaction_function(h=h, r=r, t=t, m_r=m_r).view(hr_batch.shape[0], self.num_entities) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(indices=None).unsqueeze(dim=0) - r = self.relation_embeddings(indices=rt_batch[:, 0]).unsqueeze(dim=1) - t = self.entity_embeddings(indices=rt_batch[:, 1]).unsqueeze(dim=1) - m_r = self.relation_projections(indices=rt_batch[:, 0]).view(-1, self.embedding_dim, self.relation_dim) - - return self.interaction_function(h=h, r=r, t=t, m_r=m_r).view(rt_batch.shape[0], self.num_entities) + def forward( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: # noqa: D102 + h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) + r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) + t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) + m_r = self.relation_projections.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.embedding_dim, self.relation_dim)) + return self.interaction_function(h=h, r=r, t=t, m_r=m_r) From f8c1fadacc185a5570e647d2825adeafdf8891b6 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 20:47:05 +0100 Subject: [PATCH 166/690] Super reset --- src/pykeen/models/base.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 99ab6e75a9..67c9049f0d 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1109,6 +1109,7 @@ def embedding_dim(self) -> int: # noqa:D401 return self.entity_embeddings.embedding_dim def _reset_parameters_(self): # noqa: D102 + super()._reset_parameters_() self.entity_embeddings.reset_parameters() def post_parameter_update(self) -> None: # noqa: D102 @@ -1201,6 +1202,7 @@ def relation_dim(self): # noqa:D401 return self.relation_embeddings.embedding_dim def _reset_parameters_(self): # noqa: D102 + super()._reset_parameters_() self.entity_embeddings.reset_parameters() self.relation_embeddings.reset_parameters() From 1ccdddf65e43be92133123f0b0578719371db0b2 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 20:52:33 +0100 Subject: [PATCH 167/690] Extract functional tucker --- src/pykeen/models/unimodal/tucker.py | 99 +++------------------------- src/pykeen/nn/functional.py | 70 ++++++++++++++++++++ 2 files changed, 79 insertions(+), 90 deletions(-) diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index 8dd9af3cfd..bb9d9279a6 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -10,6 +10,7 @@ from ..base import EntityRelationEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss +from ...nn import functional as F from ...nn.init import xavier_normal_ from ...regularizers import Regularizer from ...triples import TriplesFactory @@ -20,17 +21,6 @@ ] -def _apply_bn_to_tensor( - batch_norm: nn.BatchNorm1d, - tensor: torch.FloatTensor, -) -> torch.FloatTensor: - shape = tensor.shape - tensor = tensor.view(-1, shape[-1]) - tensor = batch_norm(tensor) - tensor = tensor.view(*shape) - return tensor - - class TuckER(EntityRelationEmbeddingModel): r"""An implementation of TuckEr from [balazevic2019]_. @@ -136,84 +126,13 @@ def _reset_parameters_(self): # noqa: D102 # Initialize core tensor, cf. https://github.com/ibalazevic/TuckER/blob/master/model.py#L12 nn.init.uniform_(self.core_tensor, -1., 1.) - def _scoring_function( + def forward( self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], ) -> torch.FloatTensor: - """ - Evaluate the scoring function. - - Compute scoring function W x_1 h x_2 r x_3 t as in the official implementation, i.e. as - - DO(BN(DO(BN(h)) x_1 DO(W x_2 r))) x_3 t - - where BN denotes BatchNorm and DO denotes Dropout - - :param h: shape: (batch_size, 1, embedding_dim) or (1, num_entities, embedding_dim) - :param r: shape: (batch_size, relation_dim) - :param t: shape: (1, num_entities, embedding_dim) or (batch_size, 1, embedding_dim) - :return: shape: (batch_size, num_entities) or (batch_size, 1) - """ - # Abbreviation - w = self.core_tensor - d_e = self.embedding_dim - d_r = self.relation_dim - - # Compute h_n = DO(BN(h)) - if self.apply_batch_normalization: - h = _apply_bn_to_tensor(batch_norm=self.bn_0, tensor=h) - - h = self.input_dropout(h) - - # Compute wr = DO(W x_2 r) - w = w.view(1, d_e, d_r, d_e) - r = r.view(-1, 1, 1, d_r) - wr = r @ w - wr = self.hidden_dropout_1(wr) - - # compute whr = DO(BN(h_n x_1 wr)) - wr = wr.view(-1, d_e, d_e) - whr = (h @ wr) - if self.apply_batch_normalization: - whr = _apply_bn_to_tensor(batch_norm=self.bn_1, tensor=whr) - whr = self.hidden_dropout_2(whr) - - # Compute whr x_3 t - scores = torch.sum(whr * t, dim=-1) - - return scores - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(indices=hrt_batch[:, 0]).unsqueeze(1) - r = self.relation_embeddings(indices=hrt_batch[:, 1]) - t = self.entity_embeddings(indices=hrt_batch[:, 2]).unsqueeze(1) - - # Compute scores - scores = self._scoring_function(h=h, r=r, t=t) - - return scores - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(indices=hr_batch[:, 0]).unsqueeze(1) - r = self.relation_embeddings(indices=hr_batch[:, 1]) - t = self.entity_embeddings(indices=None).unsqueeze(0) - - # Compute scores - scores = self._scoring_function(h=h, r=r, t=t) - - return scores - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(indices=None).unsqueeze(0) - r = self.relation_embeddings(indices=rt_batch[:, 0]) - t = self.entity_embeddings(indices=rt_batch[:, 1]).unsqueeze(1) - - # Compute scores - scores = self._scoring_function(h=h, r=r, t=t) - - return scores + h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) + r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) + t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) + return F.tucker_interaction(h=h, r=r, t=t) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index a49ba92d37..74c46e3f87 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -979,3 +979,73 @@ def transh_interaction( p=p, power_norm=power_norm, ) + + +def _apply_optional_bn_to_tensor( + batch_norm: Optional[nn.BatchNorm1d], + output_dropout: nn.Dropout, + tensor: torch.FloatTensor, +) -> torch.FloatTensor: + if batch_norm is not None: + shape = tensor.shape + tensor = tensor.view(-1, shape[-1]) + tensor = batch_norm(tensor) + tensor = tensor.view(*shape) + tensor = output_dropout(tensor) + return tensor + + +def tucker_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + core_tensor: torch.FloatTensor, + do1: nn.Dropout, + do0: nn.Dropout, + do2: nn.Dropout, + bn1: Optional[nn.BatchNorm1d], + bn2: Optional[nn.BatchNorm1d], +) -> torch.FloatTensor: + """ + Evaluate the TuckEr interaction function. + + Compute scoring function W x_1 h x_2 r x_3 t as in the official implementation, i.e. as + + DO(BN(DO(BN(h)) x_1 DO(W x_2 r))) x_3 t + + where BN denotes BatchNorm and DO denotes Dropout + + :param h: shape: (batch_size, num_heads, d_e) + The head representations. + :param r: shape: (batch_size, num_relations, d_r) + The relation representations. + :param t: shape: (batch_size, num_tails, d_e) + The tail representations. + :param core_tensor: shape: (d_e, d_r, d_e) + The core tensor. + :param do1: + The dropout layer for the head representations. + :param do0: + The first hidden dropout. + :param do2: + The second hidden dropout. + :param bn1: + The first batch normalization layer. + :param bn2: + The second batch normalization layer. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + # Compute wr = DO(W x_2 r) + x = do0(_extended_einsum("ijd,brd->brij", core_tensor, r)) + + # Compute h_n = DO(BN(h)) + h = _apply_optional_bn_to_tensor(batch_norm=bn1, output_dropout=do1, tensor=h) + + # compute whr = DO(BN(h_n x_1 wr)) + x = _extended_einsum("brid,bhd->bhri", h, x) + x = _apply_optional_bn_to_tensor(batch_norm=bn2, tensor=x, output_dropout=do2) + + # Compute whr x_3 t + return _extended_einsum("bhrd,btd->bhrt", x, t) From ff68de6e7605e259d4d34e63b27f006dc67e3e29 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 20:58:01 +0100 Subject: [PATCH 168/690] Extract Tucker interaction function and make it simple --- src/pykeen/models/unimodal/tucker.py | 38 +++++------------- src/pykeen/nn/functional.py | 2 +- src/pykeen/nn/modules.py | 60 ++++++++++++++++++++++++++++ 3 files changed, 72 insertions(+), 28 deletions(-) diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index bb9d9279a6..ca720c4c79 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -4,14 +4,13 @@ from typing import Optional -import torch import torch.autograd -from torch import nn -from ..base import EntityRelationEmbeddingModel +from .. import SimpleVectorEntityRelationEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss from ...nn import functional as F from ...nn.init import xavier_normal_ +from ...nn.modules import TuckerInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -21,7 +20,7 @@ ] -class TuckER(EntityRelationEmbeddingModel): +class TuckER(SimpleVectorEntityRelationEmbeddingModel): r"""An implementation of TuckEr from [balazevic2019]_. TuckER is a linear model that is based on the tensor factorization method Tucker in which a three-mode tensor @@ -92,6 +91,14 @@ def __init__( """ super().__init__( triples_factory=triples_factory, + interaction_function=TuckerInteractionFunction( + embedding_dim=embedding_dim, + relation_dim=relation_dim, + dropout_0=dropout_0, + dropout_1=dropout_1, + dropout_2=dropout_2, + apply_batch_normalization=apply_batch_normalization, + ), embedding_dim=embedding_dim, relation_dim=relation_dim, automatic_memory_optimization=automatic_memory_optimization, @@ -103,29 +110,6 @@ def __init__( relation_initializer=xavier_normal_, ) - # Core tensor - # Note: we use a different dimension permutation as in the official implementation to match the paper. - self.core_tensor = nn.Parameter( - torch.empty(self.embedding_dim, self.relation_dim, self.embedding_dim, device=self.device), - requires_grad=True, - ) - - # Dropout - self.input_dropout = nn.Dropout(dropout_0) - self.hidden_dropout_1 = nn.Dropout(dropout_1) - self.hidden_dropout_2 = nn.Dropout(dropout_2) - - self.apply_batch_normalization = apply_batch_normalization - - if self.apply_batch_normalization: - self.bn_0 = nn.BatchNorm1d(self.embedding_dim) - self.bn_1 = nn.BatchNorm1d(self.embedding_dim) - - def _reset_parameters_(self): # noqa: D102 - super()._reset_parameters_() - # Initialize core tensor, cf. https://github.com/ibalazevic/TuckER/blob/master/model.py#L12 - nn.init.uniform_(self.core_tensor, -1., 1.) - def forward( self, h_indices: Optional[torch.LongTensor], diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 74c46e3f87..e561950b14 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -1000,8 +1000,8 @@ def tucker_interaction( r: torch.FloatTensor, t: torch.FloatTensor, core_tensor: torch.FloatTensor, - do1: nn.Dropout, do0: nn.Dropout, + do1: nn.Dropout, do2: nn.Dropout, bn1: Optional[nn.BatchNorm1d], bn2: Optional[nn.BatchNorm1d], diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index b9b731f105..313568f5b3 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -728,3 +728,63 @@ def forward( p=self.p, power_norm=self.power_norm, ) + + +class TuckerInteractionFunction(InteractionFunction): + """Interaction function of Tucker.""" + + def __init__( + self, + embedding_dim: int = 200, + relation_dim: Optional[int] = None, + dropout_0: float = 0.3, + dropout_1: float = 0.4, + dropout_2: float = 0.5, + apply_batch_normalization: bool = True, + ): + super().__init__() + + if relation_dim is None: + relation_dim = embedding_dim + + # Core tensor + # Note: we use a different dimension permutation as in the official implementation to match the paper. + self.core_tensor = nn.Parameter( + torch.empty(embedding_dim, relation_dim, embedding_dim), + requires_grad=True, + ) + + # Dropout + self.input_dropout = nn.Dropout(dropout_0) + self.hidden_dropout_1 = nn.Dropout(dropout_1) + self.hidden_dropout_2 = nn.Dropout(dropout_2) + + if apply_batch_normalization: + self.bn1 = nn.BatchNorm1d(embedding_dim) + self.bn2 = nn.BatchNorm1d(embedding_dim) + else: + self.bn1 = self.bn2 = None + + def reset_parameters(self): + # Initialize core tensor, cf. https://github.com/ibalazevic/TuckER/blob/master/model.py#L12 + nn.init.uniform_(self.core_tensor, -1., 1.) + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + self._check_for_empty_kwargs(kwargs=kwargs) + return pykeen_functional.tucker_interaction( + h=h, + r=r, + t=t, + core_tensor=self.core_tensor, + do0=self.input_dropout, + do1=self.hidden_dropout_1, + do2=self.hidden_dropout_2, + bn1=self.bn1, + bn2=self.bn2, + ) From f763058bf45ef7effca71da085b46950edde235d Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 20:59:22 +0100 Subject: [PATCH 169/690] Delete lots of unnecessary code all of the parameters in all of the models that needed initialization can now be done completely automatically :) --- src/pykeen/models/base.py | 13 ------------- src/pykeen/models/unimodal/conv_e.py | 5 ----- src/pykeen/models/unimodal/kg2e.py | 9 --------- src/pykeen/models/unimodal/ntn.py | 7 ------- src/pykeen/models/unimodal/rgcn.py | 4 ---- src/pykeen/models/unimodal/simple.py | 8 -------- src/pykeen/models/unimodal/trans_h.py | 5 ----- src/pykeen/models/unimodal/trans_r.py | 4 ---- 8 files changed, 55 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 67c9049f0d..548ded150a 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1108,10 +1108,6 @@ def embedding_dim(self) -> int: # noqa:D401 """The entity embedding dimension.""" return self.entity_embeddings.embedding_dim - def _reset_parameters_(self): # noqa: D102 - super()._reset_parameters_() - self.entity_embeddings.reset_parameters() - def post_parameter_update(self) -> None: # noqa: D102 # make sure to call this first, to reset regularizer state! super().post_parameter_update() @@ -1201,11 +1197,6 @@ def relation_dim(self): # noqa:D401 """The relation embedding dimension.""" return self.relation_embeddings.embedding_dim - def _reset_parameters_(self): # noqa: D102 - super()._reset_parameters_() - self.entity_embeddings.reset_parameters() - self.relation_embeddings.reset_parameters() - def post_parameter_update(self) -> None: # noqa: D102 # make sure to call this first, to reset regularizer state! super().post_parameter_update() @@ -1352,10 +1343,6 @@ def __init__( ) self.index_function = index_function - def _reset_parameters_(self): - super()._reset_parameters_() - self.index_function.reset_parameters() - def forward( self, h_indices: Optional[torch.LongTensor] = None, diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index c126194920..90cd124266 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -168,11 +168,6 @@ def __init__( apply_batch_normalization=apply_batch_normalization, ) - def _reset_parameters_(self): # noqa: D102 - super()._reset_parameters_() - self.bias_term.reset_parameters() - self.interaction_function.reset_parameters() - def forward( self, h_indices: Optional[torch.LongTensor], diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 6174c34a8d..6d9f443722 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -116,15 +116,6 @@ def __init__( constrainer_kwargs=dict(min=c_min, max=c_max), ) - def _reset_parameters_(self): # noqa: D102 - # Constraints are applied through post_parameter_update - super()._reset_parameters_() - for emb in [ - self.entity_covariances, - self.relation_covariances, - ]: - emb.reset_parameters() - def post_parameter_update(self) -> None: # noqa: D102 super().post_parameter_update() for cov in ( diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index f5e64d4eab..43a9e56ef5 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -106,13 +106,6 @@ def __init__( non_linearity = nn.Tanh() self.non_linearity = non_linearity - def _reset_parameters_(self): # noqa: D102 - for module in self.modules(): - if module is self: - continue - if hasattr(module, "reset_parameters"): - module.reset_parameters() - def forward( self, h_indices: Optional[torch.LongTensor] = None, diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index 79bd5a61f7..72e1d905e4 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -546,10 +546,6 @@ def post_parameter_update(self) -> None: # noqa: D102 self.entity_representations.post_parameter_update() self.relation_embeddings.post_parameter_update() - def _reset_parameters_(self): - self.entity_representations.reset_parameters() - self.relation_embeddings.reset_parameters() - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Enrich embeddings h = self.entity_representations(indices=hrt_batch[:, 0]) diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index 35bfc9cc0c..71f6729fd8 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -101,14 +101,6 @@ def __init__( clamp_score = (-clamp_score, clamp_score) self.clamp = clamp_score - def _reset_parameters_(self): # noqa: D102 - super()._reset_parameters_() - for emb in [ - self.tail_entity_embeddings, - self.inverse_relation_embeddings, - ]: - emb.reset_parameters() - def forward( self, h_indices: Optional[torch.LongTensor], diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index a0ece409c7..d0485bd138 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -102,11 +102,6 @@ def post_parameter_update(self) -> None: # noqa: D102 super().post_parameter_update() self.normal_vector_embeddings.post_parameter_update() - def _reset_parameters_(self): # noqa: D102 - super()._reset_parameters_() - self.normal_vector_embeddings.reset_parameters() - # TODO: Add initialization - def regularize_if_necessary(self) -> None: """Update the regularizer's term given some tensors, if regularization is requested.""" # As described in [wang2014], all entities and relations are used to compute the regularization term diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 19f776a0ad..7916d99d7f 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -122,10 +122,6 @@ def __init__( ), ) - def _reset_parameters_(self): # noqa: D102 - super()._reset_parameters_() - self.relation_projections.reset_parameters() - def forward( self, h_indices: Optional[torch.LongTensor], From a5480fd1ba99f37ec3a1afc7b5f9f2a10b1abd9c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:05:28 +0100 Subject: [PATCH 170/690] Delete index interaction function abstraction --- src/pykeen/models/base.py | 176 ++------------------------- src/pykeen/models/unimodal/tucker.py | 14 --- 2 files changed, 11 insertions(+), 179 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 548ded150a..cdffa06cec 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -22,8 +22,7 @@ from ..triples import TriplesFactory from ..typing import Constrainer, DeviceHint, Initializer, MappedTriples, Normalizer from ..utils import ( - NoRandomSeedNecessary, get_hr_indices, get_hrt_indices, get_ht_indices, get_rt_indices, - resolve_device, set_random_seed, + NoRandomSeedNecessary, resolve_device, set_random_seed, ) __all__ = [ @@ -1212,73 +1211,13 @@ class MultimodalModel(EntityRelationEmbeddingModel): """A multimodal KGE model.""" -class IndexFunction(nn.Module): - """A function that handles looking up embeddings by index.""" - - def forward( - self, - model: Model, - h_indices: Optional[torch.LongTensor] = None, - r_indices: Optional[torch.LongTensor] = None, - t_indices: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: - """Get the scores for the given indices. - - :param model: - The KGEM so lookup to attributes like entity embeddings and relation embeddings is possible - :param h_indices: shape: (batch_size,) - The indices for head entities. If None, score against all. - :param r_indices: shape: (batch_size,) - The indices for relations. If None, score against all. - :param t_indices: shape: (batch_size,) - The indices for tail entities. If None, score against all. - - :return: The scores, shape: (batch_size, num_entities) - """ - raise NotImplementedError - - def reset_parameters(self): - """Reset parameters the interaction function may have.""" - - -class InteractionIndexFunction(IndexFunction): - """Wrap a :class:`InteractionFunction` with index-based lookup of entity and relation embeddings.""" - - def __init__(self, interaction_function: InteractionFunction): - super().__init__() - self.interaction_function = interaction_function - - def forward( - self, - model: Model, - h_indices: Optional[torch.LongTensor] = None, - r_indices: Optional[torch.LongTensor] = None, - t_indices: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = model.entity_embeddings.get_in_canonical_shape(indices=h_indices) - r = model.relation_embeddings.get_in_canonical_shape(indices=r_indices) - t = model.entity_embeddings.get_in_canonical_shape(indices=t_indices) - - # Compute score - scores = self.interaction_function(h=h, r=r, t=t) - - # Only regularize relation embeddings - model.regularize_if_necessary(r) - - return scores - - def reset_parameters(self): # noqa: D102 - self.interaction_function.reset_parameters() - - -class GeneralVectorEntityRelationEmbeddingModel(EntityRelationEmbeddingModel, reset_parameters_post_init=False): +class SimpleVectorEntityRelationEmbeddingModel(EntityRelationEmbeddingModel): """A base class for embedding models which store a single vector for each entity and relation.""" def __init__( self, triples_factory: TriplesFactory, - index_function: IndexFunction, + interaction_function: InteractionFunction, embedding_dim: int = 200, relation_dim: Optional[int] = None, automatic_memory_optimization: Optional[bool] = None, @@ -1303,11 +1242,11 @@ def __init__( :param triples_factory: The triple factory connected to the model. - :param index_function: - The index-based interaction function used to compute scores. + :param interaction_function: + The embedding-based interaction function used to compute scores. :param embedding_dim: The embedding dimensionality of the entity embeddings. - :param automatic_memory_optimization: bool + :param automatic_memory_optimization: Whether to automatically optimize the sub-batch size during training and batch size during evaluation with regards to the hardware at hand. :param loss: @@ -1341,7 +1280,7 @@ def __init__( relation_constrainer=relation_constrainer, relation_constrainer_kwargs=relation_constrainer_kwargs, ) - self.index_function = index_function + self.interaction_function = interaction_function def forward( self, @@ -1360,100 +1299,7 @@ def forward( :return: The scores, shape: (batch_size, num_entities) """ - return self.index_function(model=self, h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_hrt_indices(hrt_batch) - return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices).view(-1, 1) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_hr_indices(hr_batch) - return self( - h_indices=h_indices, r_indices=r_indices, t_indices=t_indices, - ).view(hr_batch.shape[0], self.num_entities) - - def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_ht_indices(ht_batch) - return self(h_indices=h_indices, r_indices=r_indices, t_indices=t_indices) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h_indices, r_indices, t_indices = get_rt_indices(rt_batch) - return self( - h_indices=h_indices, r_indices=r_indices, t_indices=t_indices, - ).view(rt_batch.shape[0], self.num_entities) - - -class SimpleVectorEntityRelationEmbeddingModel( - GeneralVectorEntityRelationEmbeddingModel, reset_parameters_post_init=False, -): - """A base class for embedding models which store a single vector for each entity and relation.""" - - def __init__( - self, - triples_factory: TriplesFactory, - interaction_function: InteractionFunction, - embedding_dim: int = 200, - relation_dim: Optional[int] = None, - automatic_memory_optimization: Optional[bool] = None, - loss: Optional[Loss] = None, - preferred_device: Optional[str] = None, - random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, - entity_initializer: Optional[Initializer] = None, - entity_initializer_kwargs: Optional[Mapping[str, Any]] = None, - entity_normalizer: Optional[Normalizer] = None, - entity_normalizer_kwargs: Optional[Mapping[str, Any]] = None, - entity_constrainer: Optional[Constrainer] = None, - entity_constrainer_kwargs: Optional[Mapping[str, Any]] = None, - relation_initializer: Optional[Initializer] = None, - relation_initializer_kwargs: Optional[Mapping[str, Any]] = None, - relation_normalizer: Optional[Normalizer] = None, - relation_normalizer_kwargs: Optional[Mapping[str, Any]] = None, - relation_constrainer: Optional[Constrainer] = None, - relation_constrainer_kwargs: Optional[Mapping[str, Any]] = None, - ) -> None: - """Initialize embedding model. - - :param triples_factory: - The triple factory connected to the model. - :param interaction_function: - The embedding-based interaction function used to compute scores. - :param embedding_dim: - The embedding dimensionality of the entity embeddings. - :param automatic_memory_optimization: - Whether to automatically optimize the sub-batch size during training and batch size during evaluation with - regards to the hardware at hand. - :param loss: - The loss to use. - :param preferred_device: - The default device where to model is located. - :param random_seed: - An optional random seed to set before the initialization of weights. - :param regularizer: - The regularizer to use. - """ - index_function = InteractionIndexFunction(interaction_function=interaction_function) - - super().__init__( - triples_factory=triples_factory, - embedding_dim=embedding_dim, - relation_dim=relation_dim, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - index_function=index_function, - random_seed=random_seed, - regularizer=regularizer, - entity_initializer=entity_initializer, - entity_initializer_kwargs=entity_initializer_kwargs, - entity_normalizer=entity_normalizer, - entity_normalizer_kwargs=entity_normalizer_kwargs, - entity_constrainer=entity_constrainer, - entity_constrainer_kwargs=entity_constrainer_kwargs, - relation_initializer=relation_initializer, - relation_initializer_kwargs=relation_initializer_kwargs, - relation_normalizer=relation_normalizer, - relation_normalizer_kwargs=relation_normalizer_kwargs, - relation_constrainer=relation_constrainer, - relation_constrainer_kwargs=relation_constrainer_kwargs, - ) + h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) + r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) + t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) + return F.tucker_interaction(h=h, r=r, t=t) diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index ca720c4c79..ae257bf99c 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -4,11 +4,8 @@ from typing import Optional -import torch.autograd - from .. import SimpleVectorEntityRelationEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss -from ...nn import functional as F from ...nn.init import xavier_normal_ from ...nn.modules import TuckerInteractionFunction from ...regularizers import Regularizer @@ -109,14 +106,3 @@ def __init__( entity_initializer=xavier_normal_, relation_initializer=xavier_normal_, ) - - def forward( - self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], - ) -> torch.FloatTensor: - h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) - t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - return F.tucker_interaction(h=h, r=r, t=t) From 4725bfc47c8ebc00f98008c240753b770af1ffe9 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:09:06 +0100 Subject: [PATCH 171/690] Remove traces of index interaction function --- src/pykeen/models/__init__.py | 4 +- src/pykeen/models/unimodal/trans_d.py | 104 ++++++++------------------ 2 files changed, 34 insertions(+), 74 deletions(-) diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index e3e4b0abbc..0967136277 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -9,7 +9,7 @@ from typing import Mapping, Set, Type, Union from .base import ( # noqa:F401 - EntityEmbeddingModel, EntityRelationEmbeddingModel, GeneralVectorEntityRelationEmbeddingModel, Model, + EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, MultimodalModel, SimpleVectorEntityRelationEmbeddingModel, ) from .multimodal import ComplExLiteral, DistMultLiteral @@ -66,9 +66,7 @@ 'get_model_cls', ] - _CONCRETE_BASES = { - GeneralVectorEntityRelationEmbeddingModel, SimpleVectorEntityRelationEmbeddingModel, } diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 1a3416ded6..76d238e3fa 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -7,19 +7,16 @@ import torch import torch.autograd -from .. import Model -from ..base import GeneralVectorEntityRelationEmbeddingModel, IndexFunction +from .. import EntityRelationEmbeddingModel from ...losses import Loss from ...nn import Embedding from ...nn.init import xavier_normal_ -from ...nn.modules import InteractionFunction, TransDInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint from ...utils import clamp_norm __all__ = [ - 'TransDIndexFunction', 'TransD', ] @@ -71,64 +68,7 @@ def _project_entity( return e_bot -class TransDIndexFunction(IndexFunction): - """The index-based interaction function for TransD.""" - - def __init__( - self, - num_entities: int, - num_relations: int, - embedding_dim: int, - relation_dim: int, - device: DeviceHint, - interaction_function: Optional[InteractionFunction] = None, - ): - super().__init__() - self.entity_projections = Embedding.init_with_device( - num_embeddings=num_entities, - embedding_dim=embedding_dim, - device=device, - initializer=xavier_normal_, - ) - self.relation_projections = Embedding.init_with_device( - num_embeddings=num_relations, - embedding_dim=relation_dim, - device=device, - initializer=xavier_normal_, - ) - if interaction_function is None: - interaction_function = TransDInteractionFunction() - self.interaction_function = interaction_function - - def reset_parameters(self): # noqa: D102 - self.entity_projections.reset_parameters() - self.relation_projections.reset_parameters() - self.interaction_function.reset_parameters() - - def forward( - self, - model: Model, - h_indices: Optional[torch.LongTensor] = None, - r_indices: Optional[torch.LongTensor] = None, - t_indices: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: # noqa: D102 - h = model.entity_embeddings.get_in_canonical_shape(indices=h_indices) - h_p = self.entity_projections.get_in_canonical_shape(indices=h_indices) - - r = model.relation_embeddings.get_in_canonical_shape(indices=r_indices) - r_p = self.relation_projections.get_in_canonical_shape(indices=r_indices) - - t = model.entity_embeddings.get_in_canonical_shape(indices=t_indices) - t_p = self.entity_projections.get_in_canonical_shape(indices=t_indices) - - # Project entities - h_bot = _project_entity(e=h, e_p=h_p, r=r, r_p=r_p) - t_bot = _project_entity(e=t, e_p=t_p, r=r, r_p=r_p) - - return self.interaction_function(h=h_bot, r=r, t=t_bot) - - -class TransD(GeneralVectorEntityRelationEmbeddingModel): +class TransD(EntityRelationEmbeddingModel): r"""An implementation of TransD from [ji2015]_. TransD is an extension of :class:`pykeen.models.TransR` that, like TransR, considers entities and relations @@ -177,17 +117,8 @@ def __init__( random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: - index_function = TransDIndexFunction( - num_entities=triples_factory.num_entities, - num_relations=triples_factory.num_relations, - embedding_dim=embedding_dim, - relation_dim=relation_dim, - device=preferred_device, - ) - super().__init__( triples_factory=triples_factory, - index_function=index_function, embedding_dim=embedding_dim, relation_dim=relation_dim, automatic_memory_optimization=automatic_memory_optimization, @@ -202,3 +133,34 @@ def __init__( relation_constrainer=clamp_norm, relation_constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ) + self.entity_projections = Embedding( + num_embeddings=self.num_entities, + embedding_dim=self.embedding_dim, + initializer=xavier_normal_, + ) + self.relation_projections = Embedding( + num_embeddings=self.num_relations, + embedding_dim=self.relation_dim, + initializer=xavier_normal_, + ) + + def forward( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: + h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) + h_p = self.entity_projections.get_in_canonical_shape(indices=h_indices) + + r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) + r_p = self.relation_projections.get_in_canonical_shape(indices=r_indices) + + t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) + t_p = self.entity_projections.get_in_canonical_shape(indices=t_indices) + + # Project entities + h_bot = _project_entity(e=h, e_p=h_p, r=r, r_p=r_p) + t_bot = _project_entity(e=t, e_p=t_p, r=r, r_p=r_p) + + return self.interaction_function(h=h_bot, r=r, t=t_bot) From b95c532da5dc2c7a3167b9dabe963f867df72d5f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:09:39 +0100 Subject: [PATCH 172/690] Fix test imports --- tests/test_models.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 18f87d5657..fdf33efbfe 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -26,7 +26,6 @@ from pykeen.models.base import ( EntityEmbeddingModel, EntityRelationEmbeddingModel, - GeneralVectorEntityRelationEmbeddingModel, Model, MultimodalModel, SimpleVectorEntityRelationEmbeddingModel, @@ -51,7 +50,6 @@ MultimodalModel.__name__, EntityEmbeddingModel.__name__, EntityRelationEmbeddingModel.__name__, - GeneralVectorEntityRelationEmbeddingModel.__name__, SimpleVectorEntityRelationEmbeddingModel.__name__, 'MockModel', 'models', From 88c3e9023ad272e9bc63477f5b67564cc598b590 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:10:19 +0100 Subject: [PATCH 173/690] Remove abstractmethod annotation --- src/pykeen/models/base.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index cdffa06cec..d6c7e09c91 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -956,7 +956,6 @@ def forward( """ raise NotImplementedError - @abstractmethod def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: """Forward pass. From cf4086c866aada03d8b3dbca31938363a1b81aed Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:10:40 +0100 Subject: [PATCH 174/690] Remove abstractmethod annotation --- src/pykeen/models/base.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index d6c7e09c91..f78f606665 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -347,7 +347,6 @@ def supports_subbatching(self) -> bool: # noqa: D400, D401 """Does this model support sub-batching?""" return len(self.modules_not_supporting_sub_batching) == 0 - @abstractmethod def _reset_parameters_(self): # noqa: D401 """Reset all parameters of the model in-place.""" for module in self.modules(): From e043eb9d02ac9e3c0e229cbd7f5ce43d8c12d836 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:11:38 +0100 Subject: [PATCH 175/690] Fix artifact --- src/pykeen/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index f78f606665..9f63d37eef 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1300,4 +1300,4 @@ def forward( h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - return F.tucker_interaction(h=h, r=r, t=t) + return self.interaction_function(h=h, r=r, t=t) From c0f569cb0c0fff56b8cf0913641dea9a561c387d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:19:38 +0100 Subject: [PATCH 176/690] Insert recursive reset_parameters code from another project --- src/pykeen/models/base.py | 36 +++++++++++++++++++++++++++++++++--- 1 file changed, 33 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 9f63d37eef..69aa1f3679 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -8,6 +8,7 @@ import logging from abc import ABC, abstractmethod from collections import defaultdict +from operator import itemgetter from typing import Any, ClassVar, Collection, Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Type, Union import numpy as np @@ -349,11 +350,40 @@ def supports_subbatching(self) -> bool: # noqa: D400, D401 def _reset_parameters_(self): # noqa: D401 """Reset all parameters of the model in-place.""" - for module in self.modules(): + # cf. https://github.com/mberr/ea-sota-comparison/blob/6debd076f93a329753d819ff4d01567a23053720/src/kgm/utils/torch_utils.py#L317-L372 + # Make sure that all modules with parameters do have a reset_parameters method. + uninitialized_parameters = set(map(id, self.parameters())) + parents = defaultdict(list) + + # Recursively visit all sub-modules + task_list = [] + for name, module in self.named_modules(): + + # skip self if module is self: continue - if hasattr(module, "reset_parameters"): - module.reset_parameters() + + # Track parents for blaming + for p in module.parameters(): + parents[id(p)].append(module) + + # call reset_parameters if possible + if hasattr(module, 'reset_parameters'): + task_list.append((name.count('.'), module)) + + # initialize from bottom to top + # This ensures that specialized initializations will take priority over the default ones of its components. + for module in map(itemgetter(1), sorted(task_list, reverse=True, key=itemgetter(0))): + module.reset_parameters() + uninitialized_parameters.difference_update(map(id, module.parameters())) + + # emit warning if there where parameters which were not initialised by reset_parameters. + if len(uninitialized_parameters) > 0: + logger.warning('reset_parameters() not found for all modules containing parameters. %d parameters where likely not initialised.', len(uninitialized_parameters)) + + # Additional debug information + for i, p_id in enumerate(uninitialized_parameters, start=1): + logger.debug('[%3d] Parents to blame: %s', i, parents.get(p_id)) def reset_parameters_(self) -> 'Model': # noqa: D401 """Reset all parameters of the model and enforce model constraints.""" From 09b76f49d66240ab0b6166a7d0e94b01633a2264 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:20:15 +0100 Subject: [PATCH 177/690] Move warning --- src/pykeen/models/unimodal/conv_kb.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index f4f263f5d4..46428337b7 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -98,8 +98,4 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) - - def _reset_parameters_(self): # noqa: D102 - # embeddings logger.warning('To be consistent with the paper, initialize entity and relation embeddings from TransE.') - super()._reset_parameters_() From 9183763cec8b2a6c1f9c95deafbb7b35d4f56b48 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:21:18 +0100 Subject: [PATCH 178/690] Fix ERMLP interaction function --- src/pykeen/nn/functional.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index e561950b14..f000dd729b 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -331,8 +331,6 @@ def ermlp_interaction( """ num_heads, num_relations, num_tails, embedding_dim, _ = _extract_sizes(h, r, t) hidden_dim = hidden.weight.shape[0] - assert embedding_dim % 3 == 0 - embedding_dim = embedding_dim // 3 # split, shape: (embedding_dim, hidden_dim) head_to_hidden, rel_to_hidden, tail_to_hidden = hidden.weight.t().split(embedding_dim) bias = hidden.bias.view(1, 1, 1, 1, -1) From 99d860e0211e05eed0dea5cfe4aa099d9b24d6b7 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 21:21:24 +0100 Subject: [PATCH 179/690] Flake8 cleanup --- src/pykeen/models/base.py | 6 +---- src/pykeen/models/unimodal/trans_d.py | 2 +- src/pykeen/models/unimodal/trans_h.py | 4 +-- src/pykeen/models/unimodal/trans_r.py | 5 +++- .../models/unimodal/unstructured_model.py | 27 +++++++++++++++---- src/pykeen/nn/modules.py | 9 +++++-- 6 files changed, 37 insertions(+), 16 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 69aa1f3679..7f4f62c301 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -22,17 +22,13 @@ from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory from ..typing import Constrainer, DeviceHint, Initializer, MappedTriples, Normalizer -from ..utils import ( - NoRandomSeedNecessary, resolve_device, set_random_seed, -) +from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed __all__ = [ 'Model', 'EntityEmbeddingModel', 'EntityRelationEmbeddingModel', - 'GeneralVectorEntityRelationEmbeddingModel', 'SimpleVectorEntityRelationEmbeddingModel', - 'IndexFunction', 'MultimodalModel', ] diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 76d238e3fa..80f30dacc2 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -149,7 +149,7 @@ def forward( h_indices: Optional[torch.LongTensor], r_indices: Optional[torch.LongTensor], t_indices: Optional[torch.LongTensor], - ) -> torch.FloatTensor: + ) -> torch.FloatTensor: # noqa:D102 h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) h_p = self.entity_projections.get_in_canonical_shape(indices=h_indices) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index d0485bd138..7651e88de5 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -9,7 +9,7 @@ from ..base import EntityRelationEmbeddingModel from ...losses import Loss -from ...nn import Embedding, functional as F +from ...nn import Embedding, functional as pkf from ...regularizers import Regularizer, TransHRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -124,4 +124,4 @@ def forward( w_r = self.normal_vector_embeddings.get_in_canonical_shape(indices=r_indices) t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) self.regularize_if_necessary() - return F.transh_interaction(h, w_r, d_r, t, p=self.scoring_fct_norm) + return pkf.transh_interaction(h, w_r, d_r, t, p=self.scoring_fct_norm) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 7916d99d7f..0cca23a3fe 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -131,5 +131,8 @@ def forward( h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - m_r = self.relation_projections.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.embedding_dim, self.relation_dim)) + m_r = self.relation_projections.get_in_canonical_shape( + indices=r_indices, + reshape_dim=(self.embedding_dim, self.relation_dim), + ) return self.interaction_function(h=h, r=r, t=t, m_r=m_r) diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index ef3265c7c0..b40b92db58 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -7,7 +7,7 @@ import torch import torch.autograd -from ..base import EntityEmbeddingModel +from ..base import EntityEmbeddingModel, InteractionFunction from ...losses import Loss from ...nn.init import xavier_normal_ from ...regularizers import Regularizer @@ -19,6 +19,23 @@ ] +class UMFunction(InteractionFunction): + def __init__(self, p: int, power: int = 2): + super().__init__() + self.p = p + self.power = power + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + keepdim = kwargs.pop('keepdim') + return -torch.norm(h - t, dim=-1, p=self.p, keepdim=keepdim) ** self.power + + class UnstructuredModel(EntityEmbeddingModel): r"""An implementation of the Unstructured Model (UM) published by [bordes2014]_. @@ -69,19 +86,19 @@ def __init__( regularizer=regularizer, entity_initializer=xavier_normal_, ) - self.scoring_fct_norm = scoring_fct_norm + self.interaction_function = UMFunction(p=scoring_fct_norm) def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 h = self.entity_embeddings(indices=hrt_batch[:, 0]) t = self.entity_embeddings(indices=hrt_batch[:, 2]) - return -torch.norm(h - t, dim=-1, p=self.scoring_fct_norm, keepdim=True) ** 2 + return self.interaction_function(h=h, r=None, t=t, keepdim=True) def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 h = self.entity_embeddings(indices=hr_batch[:, 0]).view(-1, 1, self.embedding_dim) t = self.entity_embeddings(indices=None).view(1, -1, self.embedding_dim) - return -torch.norm(h - t, dim=-1, p=self.scoring_fct_norm) ** 2 + return self.interaction_function(h=h, r=None, t=t) def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 h = self.entity_embeddings(indices=None).view(1, -1, self.embedding_dim) t = self.entity_embeddings(indices=rt_batch[:, 1]).view(-1, 1, self.embedding_dim) - return -torch.norm(h - t, dim=-1, p=self.scoring_fct_norm) ** 2 + return self.interaction_function(h=h, r=None, t=t) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 313568f5b3..2e132d9f28 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -705,6 +705,11 @@ def __init__( p: int, power_norm: bool = False ): + """Initialize the SE interaction function. + + :param p: The l_p norm + :param power_norm: Should power normalization be applied? + """ super().__init__() self.p = p self.power_norm = power_norm @@ -765,7 +770,7 @@ def __init__( else: self.bn1 = self.bn2 = None - def reset_parameters(self): + def reset_parameters(self): # noqa:D102 # Initialize core tensor, cf. https://github.com/ibalazevic/TuckER/blob/master/model.py#L12 nn.init.uniform_(self.core_tensor, -1., 1.) @@ -775,7 +780,7 @@ def forward( r: torch.FloatTensor, t: torch.FloatTensor, **kwargs, - ) -> torch.FloatTensor: + ) -> torch.FloatTensor: # noqa:D102 self._check_for_empty_kwargs(kwargs=kwargs) return pykeen_functional.tucker_interaction( h=h, From 8075d22537bc3043127367b55108526d01357812 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:23:08 +0100 Subject: [PATCH 180/690] Fix ERMLPE interaction function --- src/pykeen/nn/functional.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index f000dd729b..60f0caf1a5 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -367,7 +367,8 @@ def ermlpe_interaction( x = broadcast_cat(h.unsqueeze(dim=2), r.unsqueeze(dim=1), dim=-1) # Predict t embedding, shape: (batch_size, num_heads, num_relations, embedding_dim) - x = mlp(x) + shape = x.shape + x = mlp(x.view(-1, shape[-1])).view(*shape[:-1], -1) return (x.unsqueeze(dim=-2) @ t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]).transpose(-2, -1)).squeeze(dim=-1) From 256d6595962544e029d1c6e6024c72f82055b146 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 21:23:12 +0100 Subject: [PATCH 181/690] Pass flake8 --- src/pykeen/models/base.py | 8 ++++++-- src/pykeen/nn/modules.py | 13 +++++++++++-- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 7f4f62c301..86fe0f7e36 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -346,7 +346,7 @@ def supports_subbatching(self) -> bool: # noqa: D400, D401 def _reset_parameters_(self): # noqa: D401 """Reset all parameters of the model in-place.""" - # cf. https://github.com/mberr/ea-sota-comparison/blob/6debd076f93a329753d819ff4d01567a23053720/src/kgm/utils/torch_utils.py#L317-L372 + # cf. https://github.com/mberr/ea-sota-comparison/blob/6debd076f93a329753d819ff4d01567a23053720/src/kgm/utils/torch_utils.py#L317-L372 # noqa:E501 # Make sure that all modules with parameters do have a reset_parameters method. uninitialized_parameters = set(map(id, self.parameters())) parents = defaultdict(list) @@ -375,7 +375,11 @@ def _reset_parameters_(self): # noqa: D401 # emit warning if there where parameters which were not initialised by reset_parameters. if len(uninitialized_parameters) > 0: - logger.warning('reset_parameters() not found for all modules containing parameters. %d parameters where likely not initialised.', len(uninitialized_parameters)) + logger.warning( + 'reset_parameters() not found for all modules containing parameters. ' + '%d parameters where likely not initialized.', + len(uninitialized_parameters), + ) # Additional debug information for i, p_id in enumerate(uninitialized_parameters, start=1): diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 2e132d9f28..a8c4c22f56 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -703,7 +703,7 @@ class StructuredEmbeddingInteractionFunction(InteractionFunction): def __init__( self, p: int, - power_norm: bool = False + power_norm: bool = False, ): """Initialize the SE interaction function. @@ -720,7 +720,7 @@ def forward( r: torch.FloatTensor, t: torch.FloatTensor, **kwargs, - ) -> torch.FloatTensor: + ) -> torch.FloatTensor: # noqa:D102 dim = h.shape[-1] rh, rt = r.split(dim ** 2, dim=-1) rh = rh.view(*rh.shape[:-1], dim, dim) @@ -747,6 +747,15 @@ def __init__( dropout_2: float = 0.5, apply_batch_normalization: bool = True, ): + """Initialize the Tucker interaction function. + + :param embedding_dim: + :param relation_dim: + :param dropout_0: + :param dropout_1: + :param dropout_2: + :param apply_batch_normalization: + """ super().__init__() if relation_dim is None: From c4fa6db7482d5b5cba3713d578e9eb545d48e9f1 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 21:24:51 +0100 Subject: [PATCH 182/690] Shorten import --- src/pykeen/models/unimodal/kg2e.py | 4 ++-- src/pykeen/models/unimodal/ntn.py | 6 +++--- src/pykeen/nn/modules.py | 30 +++++++++++++++--------------- 3 files changed, 20 insertions(+), 20 deletions(-) diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 6d9f443722..025c4a09db 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -10,7 +10,7 @@ from ..base import EntityRelationEmbeddingModel from ...losses import Loss -from ...nn import Embedding, functional as pykeen_functional +from ...nn import Embedding, functional as pkf from ...nn.functional import KG2E_SIMILARITIES from ...regularizers import Regularizer from ...triples import TriplesFactory @@ -149,7 +149,7 @@ def forward( sigma_t = self.entity_covariances.get_in_canonical_shape(indices=t_indices) # Compute entity distribution - return pykeen_functional.kg2e_interaction( + return pkf.kg2e_interaction( h_mean=mu_h, h_var=sigma_h, r_mean=mu_r, diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 43a9e56ef5..4cb4f92657 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -9,7 +9,7 @@ from ..base import EntityEmbeddingModel from ...losses import Loss -from ...nn import Embedding, functional as pykeen_functional +from ...nn import Embedding, functional as pkf from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -135,7 +135,7 @@ def forward( vt = self.vt.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim)) if slice_size is None: - return pykeen_functional.ntn_interaction( + return pkf.ntn_interaction( h=h_all, t=t_all, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity, ) @@ -158,7 +158,7 @@ def forward( else: h = constant_tensor t = split - score = pykeen_functional.ntn_interaction( + score = pkf.ntn_interaction( h=h, t=t, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity, ) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index a8c4c22f56..b22f777c13 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -9,7 +9,7 @@ import torch from torch import nn -from . import functional as pykeen_functional +from . import functional as pkf from ..utils import check_shapes logger = logging.getLogger(__name__) @@ -280,11 +280,11 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa:D102 self._check_for_empty_kwargs(kwargs=kwargs) - return pykeen_functional.translational_interaction(h=h, r=r, t=t, p=self.p) + return pkf.translational_interaction(h=h, r=r, t=t, p=self.p) #: Interaction function of ComplEx -ComplExInteractionFunction = _build_module_from_stateless(pykeen_functional.complex_interaction) +ComplExInteractionFunction = _build_module_from_stateless(pkf.complex_interaction) def _calculate_missing_shape_information( @@ -428,7 +428,7 @@ def forward( raise TypeError(f"{self.__class__.__name__}.forward expects keyword argument 't_bias'.") t_bias: torch.FloatTensor = kwargs.pop("t_bias") self._check_for_empty_kwargs(kwargs) - return pykeen_functional.conve_interaction( + return pkf.conve_interaction( h=h, r=r, t=t, @@ -486,7 +486,7 @@ def forward( t: torch.FloatTensor, **kwargs, ) -> torch.FloatTensor: # noqa: D102 - return pykeen_functional.convkb_interaction( + return pkf.convkb_interaction( h=h, r=r, t=t, @@ -498,7 +498,7 @@ def forward( #: Interaction function for HolE -DistMultInteractionFunction = _build_module_from_stateless(pykeen_functional.distmult_interaction) +DistMultInteractionFunction = _build_module_from_stateless(pkf.distmult_interaction) class ERMLPInteractionFunction(InteractionFunction): @@ -539,7 +539,7 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa: D102 self._check_for_empty_kwargs(kwargs) - return pykeen_functional.ermlp_interaction( + return pkf.ermlp_interaction( h=h, r=r, t=t, @@ -591,7 +591,7 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa: D102 self._check_for_empty_kwargs(kwargs=kwargs) - return pykeen_functional.ermlpe_interaction(h=h, r=r, t=t, mlp=self.mlp) + return pkf.ermlpe_interaction(h=h, r=r, t=t, mlp=self.mlp) class TransDInteractionFunction(TranslationalInteractionFunction): @@ -636,14 +636,14 @@ def forward( ) -> torch.FloatTensor: # noqa:D102 m_r = kwargs.pop('m_r') self._check_for_empty_kwargs(kwargs=kwargs) - return pykeen_functional.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p, power_norm=True) + return pkf.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p, power_norm=True) #: Interaction function of RotatE. -RotatEInteraction = _build_module_from_stateless(pykeen_functional.rotate_interaction) +RotatEInteraction = _build_module_from_stateless(pkf.rotate_interaction) #: Interaction function for HolE. -HolEInteractionFunction = _build_module_from_stateless(pykeen_functional.hole_interaction) +HolEInteractionFunction = _build_module_from_stateless(pkf.hole_interaction) class ProjEInteractionFunction(InteractionFunction): @@ -688,13 +688,13 @@ def forward( self._check_for_empty_kwargs(kwargs=kwargs) # Compute score - return pykeen_functional.proje_interaction( + return pkf.proje_interaction( h=h, r=r, t=t, d_e=self.d_e, d_r=self.d_r, b_c=self.b_c, b_p=self.b_p, activation=self.inner_non_linearity, ).view(-1, 1) -RESCALInteractionFunction = _build_module_from_stateless(pykeen_functional.rescal_interaction) +RESCALInteractionFunction = _build_module_from_stateless(pkf.rescal_interaction) class StructuredEmbeddingInteractionFunction(InteractionFunction): @@ -725,7 +725,7 @@ def forward( rh, rt = r.split(dim ** 2, dim=-1) rh = rh.view(*rh.shape[:-1], dim, dim) rt = rt.view(*rt.shape[:-1], dim, dim) - return pykeen_functional.structured_embedding_interaction( + return pkf.structured_embedding_interaction( h=h, r_h=rh, r_t=rt, @@ -791,7 +791,7 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa:D102 self._check_for_empty_kwargs(kwargs=kwargs) - return pykeen_functional.tucker_interaction( + return pkf.tucker_interaction( h=h, r=r, t=t, From 0e29138b5531d15ee69f0e642a8a593dca8e7973 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:25:11 +0100 Subject: [PATCH 183/690] fix check for reset_parameters --- tests/test_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index fdf33efbfe..16c04f3f4d 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -485,7 +485,7 @@ def test_reset_parameters_constructor_call(self): ) except TypeError as error: assert error.args == ("'NoneType' object is not callable",) - mock_method.assert_called_once() + mock_method.assert_called() def test_custom_representations(self): """Tests whether we can provide custom representations.""" From 83bdf1924f26c56e2f5b083c7d1f3eba962a8a83 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:36:08 +0100 Subject: [PATCH 184/690] Fix UM --- .../models/unimodal/unstructured_model.py | 57 +++++++------------ src/pykeen/nn/functional.py | 26 +++++++++ src/pykeen/nn/modules.py | 19 +++++++ 3 files changed, 67 insertions(+), 35 deletions(-) diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index b40b92db58..0cb18c48a3 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -4,12 +4,12 @@ from typing import Optional -import torch import torch.autograd -from ..base import EntityEmbeddingModel, InteractionFunction +from ..base import EntityEmbeddingModel from ...losses import Loss from ...nn.init import xavier_normal_ +from ...nn.modules import UnstructuredModelInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -19,23 +19,6 @@ ] -class UMFunction(InteractionFunction): - def __init__(self, p: int, power: int = 2): - super().__init__() - self.p = p - self.power = power - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: - keepdim = kwargs.pop('keepdim') - return -torch.norm(h - t, dim=-1, p=self.p, keepdim=keepdim) ** self.power - - class UnstructuredModel(EntityEmbeddingModel): r"""An implementation of the Unstructured Model (UM) published by [bordes2014]_. @@ -86,19 +69,23 @@ def __init__( regularizer=regularizer, entity_initializer=xavier_normal_, ) - self.interaction_function = UMFunction(p=scoring_fct_norm) - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=hrt_batch[:, 0]) - t = self.entity_embeddings(indices=hrt_batch[:, 2]) - return self.interaction_function(h=h, r=None, t=t, keepdim=True) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=hr_batch[:, 0]).view(-1, 1, self.embedding_dim) - t = self.entity_embeddings(indices=None).view(1, -1, self.embedding_dim) - return self.interaction_function(h=h, r=None, t=t) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings(indices=None).view(1, -1, self.embedding_dim) - t = self.entity_embeddings(indices=rt_batch[:, 1]).view(-1, 1, self.embedding_dim) - return self.interaction_function(h=h, r=None, t=t) + self.interaction_function = UnstructuredModelInteractionFunction(p=scoring_fct_norm) + + def forward( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: # noqa: D102 + h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) + t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) + scores = self.interaction_function(h=h, r=None, t=t) + # same score for all relations + repeats = [1, 1, 1, 1] + if r_indices is None: + repeats[2] = self.num_relations + else: + relation_batch_size = len(r_indices) + if scores.shape[0] < relation_batch_size: + repeats[0] = relation_batch_size + return scores.repeat(*repeats) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 60f0caf1a5..376ef74428 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -1048,3 +1048,29 @@ def tucker_interaction( # Compute whr x_3 t return _extended_einsum("bhrd,btd->bhrt", x, t) + + +def unstructured_model_interaction( + h: torch.FloatTensor, + t: torch.FloatTensor, + p: int, + power_norm: bool = True, +) -> torch.FloatTensor: + """ + Evaluate the SimplE interaction function. + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + :param p: + The parameter p for selecting the norm. + :param power_norm: + Whether to return the powered norm instead. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + h = h.unsqueeze(dim=2).unsqueeze(dim=3) + t = t.unsqueeze(dim=1).unsqueeze(dim=2) + return negative_norm_of_sum(h, -t, p=p, power_norm=power_norm) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index b22f777c13..993b32fa53 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -802,3 +802,22 @@ def forward( bn1=self.bn1, bn2=self.bn2, ) + + +class UnstructuredModelInteractionFunction(InteractionFunction): + """Interaction function of UnstructuredModel.""" + + def __init__(self, p: int, power_norm: bool = True): + super().__init__() + self.p = p + self.power_norm = power_norm + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + self._check_for_empty_kwargs(kwargs=kwargs) + return pkf.unstructured_model_interaction(h, t, p=self.p, power_norm=self.power_norm) From cf84b9e2308457518b1e9a56be6a29c245823395 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:40:10 +0100 Subject: [PATCH 185/690] Fix RESCAL interaction shape --- src/pykeen/nn/functional.py | 4 ++-- src/pykeen/nn/modules.py | 6 ------ 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 376ef74428..5608b5972e 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -857,7 +857,7 @@ def rescal_interaction( :param h: shape: (batch_size, num_heads, dim) The head representations. - :param r: shape: (batch_size, num_relations, dim, dim) + :param r: shape: (batch_size, num_relations, dim ** 2) The relation representations. :param t: shape: (batch_size, num_tails, dim) The tail representations. @@ -865,7 +865,7 @@ def rescal_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return _extended_einsum("bhd,brde,bte->bhrt", h, r, t) + return _extended_einsum("bhd,brde,bte->bhrt", h, r.view(*r.shape[:-1], h.shape[-1], h.shape[-1]), t) def simple_interaction( diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 993b32fa53..32425c0e13 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -339,12 +339,6 @@ def _calculate_missing_shape_information( class ConvEInteractionFunction(InteractionFunction): """ConvE interaction function.""" - #: If batch normalization is enabled, this is: num_features – C from an expected input of size (N,C,L) - bn0: Optional[torch.nn.BatchNorm2d] - #: If batch normalization is enabled, this is: num_features – C from an expected input of size (N,C,H,W) - bn1: Optional[torch.nn.BatchNorm2d] - bn2: Optional[torch.nn.BatchNorm2d] - def __init__( self, input_channels: Optional[int] = None, From a3f564639d4ec190366e5b16ba9b7a65282320c7 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 13 Nov 2020 21:42:34 +0100 Subject: [PATCH 186/690] Update docs --- docs/source/reference/models.rst | 12 ++++++++++++ src/pykeen/nn/functional.py | 9 +++++---- 2 files changed, 17 insertions(+), 4 deletions(-) diff --git a/docs/source/reference/models.rst b/docs/source/reference/models.rst index f84366148f..8e3baac80c 100644 --- a/docs/source/reference/models.rst +++ b/docs/source/reference/models.rst @@ -20,3 +20,15 @@ Extra Modules ------------- .. automodule:: pykeen.nn :members: + +Interactors +----------- +Functional Interface +~~~~~~~~~~~~~~~~~~~~ +.. automodule:: pykeen.nn.functional + :members: + +Module Interface +~~~~~~~~~~~~~~~~ +.. automodule:: pykeen.nn.modules + :members: diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 5608b5972e..c43da38c3f 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- """Functional forms of interaction methods.""" + import math from typing import NamedTuple, Optional, SupportsFloat, Tuple, Union @@ -465,7 +466,7 @@ def negative_norm_of_sum( :param p: The p for the norm. cf. torch.norm. :param power_norm: - Whether to return |x-y|_p^p, cf. https://github.com/pytorch/pytorch/issues/28119 + Whether to return $|x-y|_p^p$, cf. https://github.com/pytorch/pytorch/issues/28119 :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. @@ -501,7 +502,7 @@ def _translational_interaction( :param p: The p for the norm. cf. torch.norm. :param power_norm: - Whether to return |x-y|_p^p, cf. https://github.com/pytorch/pytorch/issues/28119 + Whether to return $|x-y|_p^p$, cf. https://github.com/pytorch/pytorch/issues/28119 :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. @@ -516,7 +517,7 @@ def translational_interaction( p: Union[int, str] = 2, power_norm: bool = False, ) -> torch.FloatTensor: - """ + r""" Evaluate a translational distance interaction function. :param h: shape: (batch_size, num_heads, dim) @@ -528,7 +529,7 @@ def translational_interaction( :param p: The p for the norm. cf. torch.norm. :param power_norm: - Whether to return |x-y|_p^p, cf. https://github.com/pytorch/pytorch/issues/28119 + Whether to return $|x-y|_p^p$, cf. https://github.com/pytorch/pytorch/issues/28119 :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. From 8f68a458966925535f1e71d947ca3dafa1a0d5dd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:58:10 +0100 Subject: [PATCH 187/690] Fix simple --- src/pykeen/models/unimodal/simple.py | 27 ++++++++------------------- 1 file changed, 8 insertions(+), 19 deletions(-) diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index 71f6729fd8..c4df20cb6a 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -108,10 +108,14 @@ def forward( t_indices: Optional[torch.LongTensor], ) -> torch.FloatTensor: # noqa: D102 scores = 0.5 * sum( - self._single_forward(h_ind=h_ind, r_ind=r_ind, t_ind=t_ind, r_emb=r_emb) - for (h_ind, r_ind, t_ind, r_emb) in ( - (h_indices, r_indices, t_indices, self.relation_embeddings), - (t_indices, r_indices, h_indices, self.inverse_relation_embeddings), + self.interaction_function( + h_source.get_in_canonical_shape(indices=h_indices), + r_source.get_in_canonical_shape(indices=r_indices), + t_source.get_in_canonical_shape(indices=t_indices) + ) + for h_source, r_source, t_source in ( + (self.entity_embeddings, self.relation_embeddings, self.tail_entity_embeddings), + (self.tail_entity_embeddings, self.inverse_relation_embeddings, self.entity_embeddings), ) ) @@ -122,18 +126,3 @@ def forward( scores = scores.clamp(min=min_, max=max_) return scores - - def _single_forward( - self, - h_ind: torch.LongTensor, - r_ind: torch.LongTensor, - t_ind: torch.LongTensor, - r_emb: Embedding, - ) -> torch.FloatTensor: - # scores - h = self.entity_embeddings.get_in_canonical_shape(h_ind) - r = r_emb.get_in_canonical_shape(r_ind) - t = self.tail_entity_embeddings.get_in_canonical_shape(t_ind) - # Regularization - self.regularize_if_necessary(h, r, t) - return self.interaction_function(h, r, t) From fc1455a0c79fc051aeed338cb8e57236db17e10a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 21:59:23 +0100 Subject: [PATCH 188/690] Make the interaction function a parameter of SimplE --- src/pykeen/models/unimodal/simple.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index c4df20cb6a..d26fc8ce62 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -12,7 +12,7 @@ from ...nn.modules import DistMultInteractionFunction from ...regularizers import PowerSumRegularizer, Regularizer from ...triples import TriplesFactory -from ...typing import DeviceHint +from ...typing import DeviceHint, InteractionFunction __all__ = [ 'SimplE', @@ -73,6 +73,7 @@ def __init__( random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, clamp_score: Optional[Union[float, Tuple[float, float]]] = None, + interaction_function: Optional[InteractionFunction] = None, ) -> None: super().__init__( triples_factory=triples_factory, @@ -83,7 +84,9 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) - self.interaction_function = DistMultInteractionFunction() + if interaction_function is None: + interaction_function = DistMultInteractionFunction() + self.interaction_function = interaction_function # extra embeddings self.tail_entity_embeddings = Embedding.init_with_device( From de6fa41b3ec5e7c3b5e84798cdbf73f51cb29171 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 22:05:01 +0100 Subject: [PATCH 189/690] Extract base class from SimplE --- src/pykeen/models/base.py | 86 ++++++++++++++++++++++++++++ src/pykeen/models/unimodal/simple.py | 36 +++--------- 2 files changed, 93 insertions(+), 29 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 86fe0f7e36..1535e0d2fc 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1250,6 +1250,7 @@ def __init__( relation_dim: Optional[int] = None, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, + predict_with_sigmoid: bool = False, preferred_device: Optional[str] = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, @@ -1292,6 +1293,7 @@ def __init__( relation_dim=relation_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, + predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, @@ -1331,3 +1333,87 @@ def forward( r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) return self.interaction_function(h=h, r=r, t=t) + + +class TwoSideERModel(EntityRelationEmbeddingModel): + """A model with two sets of entity and relation embeddings.""" + + def __init__( + self, + triples_factory: TriplesFactory, + interaction_function: InteractionFunction, + embedding_dim: int = 50, + relation_dim: Optional[int] = None, + loss: Optional[Loss] = None, + predict_with_sigmoid: bool = False, + automatic_memory_optimization: Optional[bool] = None, + preferred_device: DeviceHint = None, + random_seed: Optional[int] = None, + regularizer: Optional[Regularizer] = None, + entity_initializer: Optional[Initializer] = None, + entity_initializer_kwargs: Optional[Mapping[str, Any]] = None, + entity_normalizer: Optional[Normalizer] = None, + entity_normalizer_kwargs: Optional[Mapping[str, Any]] = None, + entity_constrainer: Optional[Constrainer] = None, + entity_constrainer_kwargs: Optional[Mapping[str, Any]] = None, + relation_initializer: Optional[Initializer] = None, + relation_initializer_kwargs: Optional[Mapping[str, Any]] = None, + relation_normalizer: Optional[Normalizer] = None, + relation_normalizer_kwargs: Optional[Mapping[str, Any]] = None, + relation_constrainer: Optional[Constrainer] = None, + relation_constrainer_kwargs: Optional[Mapping[str, Any]] = None, + ) -> None: + super().__init__( + triples_factory=triples_factory, + embedding_dim=embedding_dim, + relation_dim=relation_dim, + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + predict_with_sigmoid=predict_with_sigmoid, + preferred_device=preferred_device, + random_seed=random_seed, + regularizer=regularizer, + entity_initializer=entity_initializer, + entity_initializer_kwargs=entity_initializer_kwargs, + entity_normalizer=entity_normalizer, + entity_normalizer_kwargs=entity_normalizer_kwargs, + entity_constrainer=entity_constrainer, + entity_constrainer_kwargs=entity_constrainer_kwargs, + relation_initializer=relation_initializer, + relation_initializer_kwargs=relation_initializer_kwargs, + relation_normalizer=relation_normalizer, + relation_normalizer_kwargs=relation_normalizer_kwargs, + relation_constrainer=relation_constrainer, + relation_constrainer_kwargs=relation_constrainer_kwargs, + ) + self.interaction_function = interaction_function + + # extra embeddings + self.reverse_entity_embeddings = Embedding.init_with_device( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + device=self.device, + ) + self.reverse_relation_embeddings = Embedding.init_with_device( + num_embeddings=triples_factory.num_relations, + embedding_dim=embedding_dim, + device=self.device, + ) + + def forward( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: # noqa: D102 + return 0.5 * sum( + self.interaction_function( + h_source.get_in_canonical_shape(indices=h_indices), + r_source.get_in_canonical_shape(indices=r_indices), + t_source.get_in_canonical_shape(indices=t_indices) + ) + for h_source, r_source, t_source in ( + (self.entity_embeddings, self.relation_embeddings, self.reverse_entity_embeddings), + (self.reverse_entity_embeddings, self.reverse_relation_embeddings, self.entity_embeddings), + ) + ) diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index d26fc8ce62..c51f84dd3a 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -6,9 +6,8 @@ import torch.autograd -from ..base import EntityRelationEmbeddingModel +from ..base import TwoSideERModel from ...losses import Loss, SoftplusLoss -from ...nn import Embedding from ...nn.modules import DistMultInteractionFunction from ...regularizers import PowerSumRegularizer, Regularizer from ...triples import TriplesFactory @@ -19,7 +18,7 @@ ] -class SimplE(EntityRelationEmbeddingModel): +class SimplE(TwoSideERModel): r"""An implementation of SimplE [kazemi2018]_. SimplE is an extension of canonical polyadic (CP), an early tensor factorization approach in which each entity @@ -77,6 +76,7 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, + interaction_function=DistMultInteractionFunction(), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, @@ -84,22 +84,6 @@ def __init__( random_seed=random_seed, regularizer=regularizer, ) - if interaction_function is None: - interaction_function = DistMultInteractionFunction() - self.interaction_function = interaction_function - - # extra embeddings - self.tail_entity_embeddings = Embedding.init_with_device( - num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, - device=self.device, - ) - self.inverse_relation_embeddings = Embedding.init_with_device( - num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim, - device=self.device, - ) - if isinstance(clamp_score, float): clamp_score = (-clamp_score, clamp_score) self.clamp = clamp_score @@ -110,16 +94,10 @@ def forward( r_indices: Optional[torch.LongTensor], t_indices: Optional[torch.LongTensor], ) -> torch.FloatTensor: # noqa: D102 - scores = 0.5 * sum( - self.interaction_function( - h_source.get_in_canonical_shape(indices=h_indices), - r_source.get_in_canonical_shape(indices=r_indices), - t_source.get_in_canonical_shape(indices=t_indices) - ) - for h_source, r_source, t_source in ( - (self.entity_embeddings, self.relation_embeddings, self.tail_entity_embeddings), - (self.tail_entity_embeddings, self.inverse_relation_embeddings, self.entity_embeddings), - ) + scores = super().forward( + h_indices=h_indices, + r_indices=r_indices, + t_indices=t_indices, ) # Note: In the code in their repository, the score is clamped to [-20, 20]. From bc325d8a3c76ed8d9ca6d00604275809111bb398 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 22:15:05 +0100 Subject: [PATCH 190/690] Start extracting another superclass --- src/pykeen/models/__init__.py | 4 +- src/pykeen/models/base.py | 97 +++++++++++++++++-- src/pykeen/models/unimodal/complex.py | 4 +- src/pykeen/models/unimodal/conv_kb.py | 4 +- src/pykeen/models/unimodal/distmult.py | 4 +- src/pykeen/models/unimodal/ermlp.py | 4 +- src/pykeen/models/unimodal/ermlpe.py | 4 +- src/pykeen/models/unimodal/hole.py | 4 +- src/pykeen/models/unimodal/kg2e.py | 53 +++------- src/pykeen/models/unimodal/proj_e.py | 4 +- src/pykeen/models/unimodal/rescal.py | 4 +- src/pykeen/models/unimodal/rotate.py | 4 +- src/pykeen/models/unimodal/simple.py | 4 +- .../models/unimodal/structured_embedding.py | 4 +- src/pykeen/models/unimodal/trans_d.py | 42 +++----- src/pykeen/models/unimodal/trans_e.py | 4 +- src/pykeen/models/unimodal/tucker.py | 4 +- tests/test_models.py | 4 +- 18 files changed, 144 insertions(+), 108 deletions(-) diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index 0967136277..5273620aeb 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -10,7 +10,7 @@ from .base import ( # noqa:F401 EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, - MultimodalModel, SimpleVectorEntityRelationEmbeddingModel, + MultimodalModel, SingleVectorEmbeddingModel, ) from .multimodal import ComplExLiteral, DistMultLiteral from .unimodal import ( @@ -67,7 +67,7 @@ ] _CONCRETE_BASES = { - SimpleVectorEntityRelationEmbeddingModel, + SingleVectorEmbeddingModel, } diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 1535e0d2fc..f97c2ac3d3 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -28,7 +28,7 @@ 'Model', 'EntityEmbeddingModel', 'EntityRelationEmbeddingModel', - 'SimpleVectorEntityRelationEmbeddingModel', + 'SingleVectorEmbeddingModel', 'MultimodalModel', ] @@ -1239,7 +1239,7 @@ class MultimodalModel(EntityRelationEmbeddingModel): """A multimodal KGE model.""" -class SimpleVectorEntityRelationEmbeddingModel(EntityRelationEmbeddingModel): +class SingleVectorEmbeddingModel(EntityRelationEmbeddingModel): """A base class for embedding models which store a single vector for each entity and relation.""" def __init__( @@ -1335,13 +1335,12 @@ def forward( return self.interaction_function(h=h, r=r, t=t) -class TwoSideERModel(EntityRelationEmbeddingModel): - """A model with two sets of entity and relation embeddings.""" +class TwoVectorEmbeddingModel(EntityRelationEmbeddingModel): + """A model with two vectors for each entity and relation.""" def __init__( self, triples_factory: TriplesFactory, - interaction_function: InteractionFunction, embedding_dim: int = 50, relation_dim: Optional[int] = None, loss: Optional[Loss] = None, @@ -1386,20 +1385,98 @@ def __init__( relation_constrainer=relation_constrainer, relation_constrainer_kwargs=relation_constrainer_kwargs, ) - self.interaction_function = interaction_function # extra embeddings - self.reverse_entity_embeddings = Embedding.init_with_device( + self.second_entity_embeddings = Embedding.init_with_device( num_embeddings=triples_factory.num_entities, embedding_dim=embedding_dim, device=self.device, ) - self.reverse_relation_embeddings = Embedding.init_with_device( + self.second_relation_embeddings = Embedding.init_with_device( num_embeddings=triples_factory.num_relations, embedding_dim=embedding_dim, device=self.device, ) + def _forward( + self, + h1: torch.FloatTensor, + h2: torch.FloatTensor, + r1: torch.FloatTensor, + r2: torch.FloatTensor, + t1: torch.FloatTensor, + t2: torch.FloatTensor, + ) -> torch.FloatTensor: + raise NotImplementedError + + def forward( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: # noqa: D102 + h1 = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) + h2 = self.second_entity_embeddings.get_in_canonical_shape(indices=h_indices) + r1 = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) + r2 = self.second_relation_embeddings.get_in_canonical_shape(indices=r_indices) + t1 = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) + t2 = self.second_entity_embeddings.get_in_canonical_shape(indices=t_indices) + return self._forward(h1, h2, r1, r2, t1, t2) + + +class TwoSideEmbeddingModel(TwoVectorEmbeddingModel): + """A model which averages scores for forward and backward model.""" + + def __init__( + self, + triples_factory: TriplesFactory, + interaction_function: InteractionFunction, + embedding_dim: int = 50, + relation_dim: Optional[int] = None, + loss: Optional[Loss] = None, + predict_with_sigmoid: bool = False, + automatic_memory_optimization: Optional[bool] = None, + preferred_device: DeviceHint = None, + random_seed: Optional[int] = None, + regularizer: Optional[Regularizer] = None, + entity_initializer: Optional[Initializer] = None, + entity_initializer_kwargs: Optional[Mapping[str, Any]] = None, + entity_normalizer: Optional[Normalizer] = None, + entity_normalizer_kwargs: Optional[Mapping[str, Any]] = None, + entity_constrainer: Optional[Constrainer] = None, + entity_constrainer_kwargs: Optional[Mapping[str, Any]] = None, + relation_initializer: Optional[Initializer] = None, + relation_initializer_kwargs: Optional[Mapping[str, Any]] = None, + relation_normalizer: Optional[Normalizer] = None, + relation_normalizer_kwargs: Optional[Mapping[str, Any]] = None, + relation_constrainer: Optional[Constrainer] = None, + relation_constrainer_kwargs: Optional[Mapping[str, Any]] = None, + ): + super().__init__( + triples_factory=triples_factory, + embedding_dim=embedding_dim, + relation_dim=relation_dim, + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + predict_with_sigmoid=predict_with_sigmoid, + preferred_device=preferred_device, + random_seed=random_seed, + regularizer=regularizer, + entity_initializer=entity_initializer, + entity_initializer_kwargs=entity_initializer_kwargs, + entity_normalizer=entity_normalizer, + entity_normalizer_kwargs=entity_normalizer_kwargs, + entity_constrainer=entity_constrainer, + entity_constrainer_kwargs=entity_constrainer_kwargs, + relation_initializer=relation_initializer, + relation_initializer_kwargs=relation_initializer_kwargs, + relation_normalizer=relation_normalizer, + relation_normalizer_kwargs=relation_normalizer_kwargs, + relation_constrainer=relation_constrainer, + relation_constrainer_kwargs=relation_constrainer_kwargs, + ) + self.interaction_function = interaction_function + def forward( self, h_indices: Optional[torch.LongTensor], @@ -1413,7 +1490,7 @@ def forward( t_source.get_in_canonical_shape(indices=t_indices) ) for h_source, r_source, t_source in ( - (self.entity_embeddings, self.relation_embeddings, self.reverse_entity_embeddings), - (self.reverse_entity_embeddings, self.reverse_relation_embeddings, self.entity_embeddings), + (self.entity_embeddings, self.relation_embeddings, self.second_entity_embeddings), + (self.second_entity_embeddings, self.second_relation_embeddings, self.entity_embeddings), ) ) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index b0d377128a..8da5cd92fe 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -6,7 +6,7 @@ import torch.nn as nn -from ..base import SimpleVectorEntityRelationEmbeddingModel +from ..base import SingleVectorEmbeddingModel from ...losses import Loss, SoftplusLoss from ...nn.modules import ComplExInteractionFunction from ...regularizers import LpRegularizer, Regularizer @@ -18,7 +18,7 @@ ] -class ComplEx(SimpleVectorEntityRelationEmbeddingModel): +class ComplEx(SingleVectorEmbeddingModel): r"""An implementation of ComplEx [trouillon2016]_. ComplEx is an extension of :class:`pykeen.models.DistMult` that uses complex valued representations for the diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 46428337b7..7825b8207f 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -5,7 +5,7 @@ import logging from typing import Optional -from ..base import SimpleVectorEntityRelationEmbeddingModel +from ..base import SingleVectorEmbeddingModel from ...losses import Loss from ...nn.modules import ConvKBInteractionFunction from ...regularizers import LpRegularizer, Regularizer @@ -19,7 +19,7 @@ logger = logging.getLogger(__name__) -class ConvKB(SimpleVectorEntityRelationEmbeddingModel): +class ConvKB(SingleVectorEmbeddingModel): r"""An implementation of ConvKB from [nguyen2018]_. ConvKB uses a convolutional neural network (CNN) whose feature maps capture global interactions of the input. diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index fb8d40ebdd..b0a9f0bc0e 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -7,7 +7,7 @@ from torch import nn from torch.nn import functional -from ..base import SimpleVectorEntityRelationEmbeddingModel +from ..base import SingleVectorEmbeddingModel from ...losses import Loss from ...nn.modules import DistMultInteractionFunction from ...regularizers import LpRegularizer, Regularizer @@ -20,7 +20,7 @@ ] -class DistMult(SimpleVectorEntityRelationEmbeddingModel): +class DistMult(SingleVectorEmbeddingModel): r"""An implementation of DistMult from [yang2014]_. This model simplifies RESCAL by restricting matrices representing relations as diagonal matrices. diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index d5d7150bec..ac273aed52 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -4,7 +4,7 @@ from typing import Optional -from ..base import SimpleVectorEntityRelationEmbeddingModel +from ..base import SingleVectorEmbeddingModel from ...losses import Loss from ...nn.modules import ERMLPInteractionFunction from ...regularizers import Regularizer @@ -16,7 +16,7 @@ ] -class ERMLP(SimpleVectorEntityRelationEmbeddingModel): +class ERMLP(SingleVectorEmbeddingModel): r"""An implementation of ERMLP from [dong2014]_. ERMLP is a multi-layer perceptron based approach that uses a single hidden layer and represents entities and diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index 882d274d01..863e92a978 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -4,7 +4,7 @@ from typing import Optional, Type -from ..base import SimpleVectorEntityRelationEmbeddingModel +from ..base import SingleVectorEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss from ...nn.modules import ERMLPEInteractionFunction from ...regularizers import Regularizer @@ -16,7 +16,7 @@ ] -class ERMLPE(SimpleVectorEntityRelationEmbeddingModel): +class ERMLPE(SingleVectorEmbeddingModel): r"""An extension of ERMLP proposed by [sharifzadeh2019]_. This model uses a neural network-based approach similar to ERMLP and with slight modifications. diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index a08ac7e9e9..1e9e9ed6e1 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -4,7 +4,7 @@ from typing import Optional -from ..base import SimpleVectorEntityRelationEmbeddingModel +from ..base import SingleVectorEmbeddingModel from ...losses import Loss from ...nn.init import xavier_uniform_ from ...nn.modules import HolEInteractionFunction @@ -18,7 +18,7 @@ ] -class HolE(SimpleVectorEntityRelationEmbeddingModel): +class HolE(SingleVectorEmbeddingModel): r"""An implementation of HolE [nickel2016]_. Holographic embeddings (HolE) make use of the circular correlation operator to compute interactions between diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 025c4a09db..e90ea99e8b 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -8,7 +8,7 @@ import torch import torch.autograd -from ..base import EntityRelationEmbeddingModel +from ..base import TwoVectorEmbeddingModel from ...losses import Loss from ...nn import Embedding, functional as pkf from ...nn.functional import KG2E_SIMILARITIES @@ -24,7 +24,7 @@ _LOG_2_PI = math.log(2. * math.pi) -class KG2E(EntityRelationEmbeddingModel): +class KG2E(TwoVectorEmbeddingModel): r"""An implementation of KG2E from [he2015]_. KG2E aims to explicitly model (un)certainties in entities and relations (e.g. influenced by the number of triples @@ -116,44 +116,21 @@ def __init__( constrainer_kwargs=dict(min=c_min, max=c_max), ) - def post_parameter_update(self) -> None: # noqa: D102 - super().post_parameter_update() - for cov in ( - self.entity_covariances, - self.relation_covariances, - ): - cov.post_parameter_update() - - def forward( + def _forward( self, - h_indices: Optional[torch.LongTensor] = None, - r_indices: Optional[torch.LongTensor] = None, - t_indices: Optional[torch.LongTensor] = None, + h1: torch.FloatTensor, + h2: torch.FloatTensor, + r1: torch.FloatTensor, + r2: torch.FloatTensor, + t1: torch.FloatTensor, + t2: torch.FloatTensor, ) -> torch.FloatTensor: - """ - Compute scores for NTN. - - :param h_indices: shape: (batch_size,) - :param r_indices: shape: (batch_size,) - :param t_indices: shape: (batch_size,) - - :return: shape: (batch_size, num_entities) - """ - # Get embeddings - mu_h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - mu_r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) - mu_t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - - sigma_h = self.entity_covariances.get_in_canonical_shape(indices=h_indices) - sigma_r = self.relation_covariances.get_in_canonical_shape(indices=r_indices) - sigma_t = self.entity_covariances.get_in_canonical_shape(indices=t_indices) - # Compute entity distribution return pkf.kg2e_interaction( - h_mean=mu_h, - h_var=sigma_h, - r_mean=mu_r, - r_var=sigma_r, - t_mean=mu_t, - t_var=sigma_t, + h_mean=h1, + h_var=h2, + r_mean=r1, + r_var=r2, + t_mean=t1, + t_var=t2, ) diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index 241e9eadb1..c2842050bf 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -6,7 +6,7 @@ from torch import nn -from .. import SimpleVectorEntityRelationEmbeddingModel +from .. import SingleVectorEmbeddingModel from ...losses import Loss from ...nn.init import xavier_uniform_ from ...nn.modules import ProjEInteractionFunction @@ -19,7 +19,7 @@ ] -class ProjE(SimpleVectorEntityRelationEmbeddingModel): +class ProjE(SingleVectorEmbeddingModel): r"""An implementation of ProjE from [shi2017]_. ProjE is a neural network-based approach with a *combination* and a *projection* layer. The interaction model diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index 1a132f7c30..6fd639130a 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -4,7 +4,7 @@ from typing import Optional -from ..base import SimpleVectorEntityRelationEmbeddingModel +from ..base import SingleVectorEmbeddingModel from ...losses import Loss from ...nn.modules import RESCALInteractionFunction from ...regularizers import LpRegularizer, Regularizer @@ -16,7 +16,7 @@ ] -class RESCAL(SimpleVectorEntityRelationEmbeddingModel): +class RESCAL(SingleVectorEmbeddingModel): r"""An implementation of RESCAL from [nickel2011]_. This model represents relations as matrices and models interactions between latent features. diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index 516eff962a..e956ad3831 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -9,7 +9,7 @@ import torch.autograd from torch.nn import functional -from .. import SimpleVectorEntityRelationEmbeddingModel +from .. import SingleVectorEmbeddingModel from ...losses import Loss from ...nn.init import xavier_uniform_ from ...nn.modules import RotatEInteraction @@ -52,7 +52,7 @@ def complex_normalize(x: torch.Tensor) -> torch.Tensor: return x -class RotatE(SimpleVectorEntityRelationEmbeddingModel): +class RotatE(SingleVectorEmbeddingModel): r"""An implementation of RotatE from [sun2019]_. RotatE models relations as rotations from head to tail entities in complex space: diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index c51f84dd3a..6bdd9bd9b1 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -6,7 +6,7 @@ import torch.autograd -from ..base import TwoSideERModel +from ..base import TwoSideModel from ...losses import Loss, SoftplusLoss from ...nn.modules import DistMultInteractionFunction from ...regularizers import PowerSumRegularizer, Regularizer @@ -18,7 +18,7 @@ ] -class SimplE(TwoSideERModel): +class SimplE(TwoSideModel): r"""An implementation of SimplE [kazemi2018]_. SimplE is an extension of canonical polyadic (CP), an early tensor factorization approach in which each entity diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index 4ddddeae9f..d04cd503dd 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -9,7 +9,7 @@ from torch import nn from torch.nn import functional -from .. import SimpleVectorEntityRelationEmbeddingModel +from .. import SingleVectorEmbeddingModel from ...losses import Loss from ...nn.init import xavier_uniform_ from ...nn.modules import StructuredEmbeddingInteractionFunction @@ -23,7 +23,7 @@ ] -class StructuredEmbedding(SimpleVectorEntityRelationEmbeddingModel): +class StructuredEmbedding(SingleVectorEmbeddingModel): r"""An implementation of the Structured Embedding (SE) published by [bordes2011]_. SE applies role- and relation-specific projection matrices diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 80f30dacc2..c9a1809721 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -7,9 +7,8 @@ import torch import torch.autograd -from .. import EntityRelationEmbeddingModel +from ..base import TwoVectorEmbeddingModel from ...losses import Loss -from ...nn import Embedding from ...nn.init import xavier_normal_ from ...regularizers import Regularizer from ...triples import TriplesFactory @@ -68,7 +67,7 @@ def _project_entity( return e_bot -class TransD(EntityRelationEmbeddingModel): +class TransD(TwoVectorEmbeddingModel): r"""An implementation of TransD from [ji2015]_. TransD is an extension of :class:`pykeen.models.TransR` that, like TransR, considers entities and relations @@ -133,34 +132,17 @@ def __init__( relation_constrainer=clamp_norm, relation_constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ) - self.entity_projections = Embedding( - num_embeddings=self.num_entities, - embedding_dim=self.embedding_dim, - initializer=xavier_normal_, - ) - self.relation_projections = Embedding( - num_embeddings=self.num_relations, - embedding_dim=self.relation_dim, - initializer=xavier_normal_, - ) - def forward( + def _forward( self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], + h1: torch.FloatTensor, + h2: torch.FloatTensor, + r1: torch.FloatTensor, + r2: torch.FloatTensor, + t1: torch.FloatTensor, + t2: torch.FloatTensor, ) -> torch.FloatTensor: # noqa:D102 - h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - h_p = self.entity_projections.get_in_canonical_shape(indices=h_indices) - - r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) - r_p = self.relation_projections.get_in_canonical_shape(indices=r_indices) - - t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - t_p = self.entity_projections.get_in_canonical_shape(indices=t_indices) - # Project entities - h_bot = _project_entity(e=h, e_p=h_p, r=r, r_p=r_p) - t_bot = _project_entity(e=t, e_p=t_p, r=r, r_p=r_p) - - return self.interaction_function(h=h_bot, r=r, t=t_bot) + h_bot = _project_entity(e=h1, e_p=h2, r=r1, r_p=r2) + t_bot = _project_entity(e=t1, e_p=t2, r=r1, r_p=r2) + return self.interaction_function(h=h_bot, r=r1, t=t_bot) diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index a460cf6d5b..da74ecc353 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -6,7 +6,7 @@ from torch.nn import functional -from .. import SimpleVectorEntityRelationEmbeddingModel +from .. import SingleVectorEmbeddingModel from ...losses import Loss from ...nn.init import xavier_uniform_ from ...nn.modules import TranslationalInteractionFunction @@ -20,7 +20,7 @@ ] -class TransE(SimpleVectorEntityRelationEmbeddingModel): +class TransE(SingleVectorEmbeddingModel): r"""TransE models relations as a translation from head to tail entities in :math:`\textbf{e}` [bordes2013]_. .. math:: diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index ae257bf99c..723469412f 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -4,7 +4,7 @@ from typing import Optional -from .. import SimpleVectorEntityRelationEmbeddingModel +from .. import SingleVectorEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss from ...nn.init import xavier_normal_ from ...nn.modules import TuckerInteractionFunction @@ -17,7 +17,7 @@ ] -class TuckER(SimpleVectorEntityRelationEmbeddingModel): +class TuckER(SingleVectorEmbeddingModel): r"""An implementation of TuckEr from [balazevic2019]_. TuckER is a linear model that is based on the tensor factorization method Tucker in which a three-mode tensor diff --git a/tests/test_models.py b/tests/test_models.py index 16c04f3f4d..c2fef792b8 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -28,7 +28,7 @@ EntityRelationEmbeddingModel, Model, MultimodalModel, - SimpleVectorEntityRelationEmbeddingModel, + SingleVectorEmbeddingModel, _extend_batch, get_novelty_mask, ) @@ -50,7 +50,7 @@ MultimodalModel.__name__, EntityEmbeddingModel.__name__, EntityRelationEmbeddingModel.__name__, - SimpleVectorEntityRelationEmbeddingModel.__name__, + SingleVectorEmbeddingModel.__name__, 'MockModel', 'models', 'get_model_cls', From f957414cc7af90f6b5757a9a445b24f1513f5919 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 22:15:35 +0100 Subject: [PATCH 191/690] Remove line breaks --- src/pykeen/models/unimodal/kg2e.py | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index e90ea99e8b..cd61c30ebf 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -124,13 +124,5 @@ def _forward( r2: torch.FloatTensor, t1: torch.FloatTensor, t2: torch.FloatTensor, - ) -> torch.FloatTensor: - # Compute entity distribution - return pkf.kg2e_interaction( - h_mean=h1, - h_var=h2, - r_mean=r1, - r_var=r2, - t_mean=t1, - t_var=t2, - ) + ) -> torch.FloatTensor: # noqa: D102 + return pkf.kg2e_interaction(h_mean=h1, h_var=h2, r_mean=r1, r_var=r2, t_mean=t1, t_var=t2) From ed01b4fb8c045047418cd63c63cf1753101ce2a2 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 22:31:03 +0100 Subject: [PATCH 192/690] Adapt TwoVector models --- src/pykeen/models/base.py | 109 +++++++------------------- src/pykeen/models/unimodal/kg2e.py | 48 ++++++------ src/pykeen/models/unimodal/trans_d.py | 17 ++-- src/pykeen/nn/emb.py | 58 +++++++++++++- 4 files changed, 119 insertions(+), 113 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index f97c2ac3d3..bfb720ef23 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -18,6 +18,7 @@ from ..losses import Loss, MarginRankingLoss, NSSALoss from ..nn import Embedding +from ..nn.emb import EmbeddingSpecification from ..nn.modules import InteractionFunction from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory @@ -1141,7 +1142,7 @@ def post_parameter_update(self) -> None: # noqa: D102 self.entity_embeddings.post_parameter_update() -class EntityRelationEmbeddingModel(Model): +class EntityRelationEmbeddingModel(Model, ABC): """A base module for KGE models that have different embeddings for entities and relations.""" def __init__( @@ -1155,18 +1156,8 @@ def __init__( preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, - entity_initializer: Optional[Initializer] = None, - entity_initializer_kwargs: Optional[Mapping[str, Any]] = None, - entity_normalizer: Optional[Normalizer] = None, - entity_normalizer_kwargs: Optional[Mapping[str, Any]] = None, - entity_constrainer: Optional[Constrainer] = None, - entity_constrainer_kwargs: Optional[Mapping[str, Any]] = None, - relation_initializer: Optional[Initializer] = None, - relation_initializer_kwargs: Optional[Mapping[str, Any]] = None, - relation_normalizer: Optional[Normalizer] = None, - relation_normalizer_kwargs: Optional[Mapping[str, Any]] = None, - relation_constrainer: Optional[Constrainer] = None, - relation_constrainer_kwargs: Optional[Mapping[str, Any]] = None, + embedding_specification: EmbeddingSpecification = None, + relation_embedding_specification: EmbeddingSpecification = None, ) -> None: """Initialize the entity embedding model. @@ -1186,32 +1177,22 @@ def __init__( regularizer=regularizer, predict_with_sigmoid=predict_with_sigmoid, ) - self.entity_embeddings = Embedding.init_with_device( + self.entity_embeddings = Embedding.from_specification( num_embeddings=triples_factory.num_entities, embedding_dim=embedding_dim, - device=self.device, - initializer=entity_initializer, - initializer_kwargs=entity_initializer_kwargs, - normalizer=entity_normalizer, - normalizer_kwargs=entity_normalizer_kwargs, - constrainer=entity_constrainer, - constrainer_kwargs=entity_constrainer_kwargs, + specification=embedding_specification, ) # Default for relation dimensionality if relation_dim is None: relation_dim = embedding_dim + if relation_embedding_specification is None: + relation_embedding_specification = embedding_specification - self.relation_embeddings = Embedding.init_with_device( + self.relation_embeddings = Embedding.from_specification( num_embeddings=triples_factory.num_relations, embedding_dim=relation_dim, - device=self.device, - initializer=relation_initializer, - initializer_kwargs=relation_initializer_kwargs, - normalizer=relation_normalizer, - normalizer_kwargs=relation_normalizer_kwargs, - constrainer=relation_constrainer, - constrainer_kwargs=relation_constrainer_kwargs, + specification=relation_embedding_specification, ) @property @@ -1254,18 +1235,8 @@ def __init__( preferred_device: Optional[str] = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, - entity_initializer: Optional[Initializer] = None, - entity_initializer_kwargs: Optional[Mapping[str, Any]] = None, - entity_normalizer: Optional[Normalizer] = None, - entity_normalizer_kwargs: Optional[Mapping[str, Any]] = None, - entity_constrainer: Optional[Constrainer] = None, - entity_constrainer_kwargs: Optional[Mapping[str, Any]] = None, - relation_initializer: Optional[Initializer] = None, - relation_initializer_kwargs: Optional[Mapping[str, Any]] = None, - relation_normalizer: Optional[Normalizer] = None, - relation_normalizer_kwargs: Optional[Mapping[str, Any]] = None, - relation_constrainer: Optional[Constrainer] = None, - relation_constrainer_kwargs: Optional[Mapping[str, Any]] = None, + embedding_specification: EmbeddingSpecification = None, + relation_embedding_specification: EmbeddingSpecification = None, ) -> None: """Initialize embedding model. @@ -1297,18 +1268,8 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - entity_initializer=entity_initializer, - entity_initializer_kwargs=entity_initializer_kwargs, - entity_normalizer=entity_normalizer, - entity_normalizer_kwargs=entity_normalizer_kwargs, - entity_constrainer=entity_constrainer, - entity_constrainer_kwargs=entity_constrainer_kwargs, - relation_initializer=relation_initializer, - relation_initializer_kwargs=relation_initializer_kwargs, - relation_normalizer=relation_normalizer, - relation_normalizer_kwargs=relation_normalizer_kwargs, - relation_constrainer=relation_constrainer, - relation_constrainer_kwargs=relation_constrainer_kwargs, + embedding_specification=embedding_specification, + relation_embedding_specification=relation_embedding_specification, ) self.interaction_function = interaction_function @@ -1349,18 +1310,10 @@ def __init__( preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, - entity_initializer: Optional[Initializer] = None, - entity_initializer_kwargs: Optional[Mapping[str, Any]] = None, - entity_normalizer: Optional[Normalizer] = None, - entity_normalizer_kwargs: Optional[Mapping[str, Any]] = None, - entity_constrainer: Optional[Constrainer] = None, - entity_constrainer_kwargs: Optional[Mapping[str, Any]] = None, - relation_initializer: Optional[Initializer] = None, - relation_initializer_kwargs: Optional[Mapping[str, Any]] = None, - relation_normalizer: Optional[Normalizer] = None, - relation_normalizer_kwargs: Optional[Mapping[str, Any]] = None, - relation_constrainer: Optional[Constrainer] = None, - relation_constrainer_kwargs: Optional[Mapping[str, Any]] = None, + embedding_specification: EmbeddingSpecification = None, + relation_embedding_specification: EmbeddingSpecification = None, + second_embedding_specification: EmbeddingSpecification = None, + second_relation_embedding_specification: EmbeddingSpecification = None, ) -> None: super().__init__( triples_factory=triples_factory, @@ -1372,30 +1325,24 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - entity_initializer=entity_initializer, - entity_initializer_kwargs=entity_initializer_kwargs, - entity_normalizer=entity_normalizer, - entity_normalizer_kwargs=entity_normalizer_kwargs, - entity_constrainer=entity_constrainer, - entity_constrainer_kwargs=entity_constrainer_kwargs, - relation_initializer=relation_initializer, - relation_initializer_kwargs=relation_initializer_kwargs, - relation_normalizer=relation_normalizer, - relation_normalizer_kwargs=relation_normalizer_kwargs, - relation_constrainer=relation_constrainer, - relation_constrainer_kwargs=relation_constrainer_kwargs, + embedding_specification=embedding_specification, + relation_embedding_specification=relation_embedding_specification, ) # extra embeddings - self.second_entity_embeddings = Embedding.init_with_device( + if second_embedding_specification is None: + second_embedding_specification = embedding_specification + self.second_entity_embeddings = Embedding.from_specification( num_embeddings=triples_factory.num_entities, embedding_dim=embedding_dim, - device=self.device, + specification=second_embedding_specification, ) - self.second_relation_embeddings = Embedding.init_with_device( + if second_relation_embedding_specification is None: + second_relation_embedding_specification = relation_embedding_specification + self.second_relation_embeddings = Embedding.from_specification( num_embeddings=triples_factory.num_relations, embedding_dim=embedding_dim, - device=self.device, + specification=second_relation_embedding_specification, ) def _forward( diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index cd61c30ebf..faecf77c30 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -10,7 +10,8 @@ from ..base import TwoVectorEmbeddingModel from ...losses import Loss -from ...nn import Embedding, functional as pkf +from ...nn import functional as pkf +from ...nn.emb import EmbeddingSpecification from ...nn.functional import KG2E_SIMILARITIES from ...regularizers import Regularizer from ...triples import TriplesFactory @@ -86,36 +87,31 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - entity_constrainer=clamp_norm, - entity_constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), - relation_constrainer=clamp_norm, - relation_constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + embedding_specification=EmbeddingSpecification( + constrainer=clamp_norm, + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + ), + relation_embedding_specification=EmbeddingSpecification( + constrainer=clamp_norm, + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + ), + # Ensure positive definite covariances matrices and appropriate size by clamping + second_embedding_specification=EmbeddingSpecification( + constrainer=torch.clamp, + constrainer_kwargs=dict(min=c_min, max=c_max), + ), + second_relation_embedding_specification=EmbeddingSpecification( + # Ensure positive definite covariances matrices and appropriate size by clamping + constrainer=torch.clamp, + constrainer_kwargs=dict(min=c_min, max=c_max), + ) ) - # Similarity function used for distributions dist_similarity = dist_similarity.upper() if dist_similarity not in KG2E_SIMILARITIES: raise ValueError(dist_similarity) self.similarity = dist_similarity - # Additional covariance embeddings - self.entity_covariances = Embedding.init_with_device( - num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, - device=self.device, - # Ensure positive definite covariances matrices and appropriate size by clamping - constrainer=torch.clamp, - constrainer_kwargs=dict(min=c_min, max=c_max), - ) - self.relation_covariances = Embedding.init_with_device( - num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim, - device=self.device, - # Ensure positive definite covariances matrices and appropriate size by clamping - constrainer=torch.clamp, - constrainer_kwargs=dict(min=c_min, max=c_max), - ) - def _forward( self, h1: torch.FloatTensor, @@ -125,4 +121,6 @@ def _forward( t1: torch.FloatTensor, t2: torch.FloatTensor, ) -> torch.FloatTensor: # noqa: D102 - return pkf.kg2e_interaction(h_mean=h1, h_var=h2, r_mean=r1, r_var=r2, t_mean=t1, t_var=t2) + return pkf.kg2e_interaction( + h_mean=h1, h_var=h2, r_mean=r1, r_var=r2, t_mean=t1, t_var=t2, similarity=self.similarity, + ) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index c9a1809721..26836c1e49 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -9,6 +9,7 @@ from ..base import TwoVectorEmbeddingModel from ...losses import Loss +from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_normal_ from ...regularizers import Regularizer from ...triples import TriplesFactory @@ -125,12 +126,16 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - entity_initializer=xavier_normal_, - relation_initializer=xavier_normal_, - entity_constrainer=clamp_norm, - entity_constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), - relation_constrainer=clamp_norm, - relation_constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + embedding_specification=EmbeddingSpecification( + initializer=xavier_normal_, + constrainer=clamp_norm, + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + ), + relation_embedding_specification=EmbeddingSpecification( + initializer=xavier_normal_, + constrainer=clamp_norm, + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + ), ) def _forward( diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index 44446cc11e..66a99dd98e 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -1,7 +1,7 @@ # -*- coding: utf-8 -*- """Embedding modules.""" - +import dataclasses import functools from typing import Any, Mapping, Optional, Sequence @@ -42,6 +42,25 @@ def post_parameter_update(self): """Apply constraints which should not be included in gradients.""" +@dataclasses.dataclass +class EmbeddingSpecification: + """An embedding specification.""" + + # embedding_dim: int + # shape: Optional[Sequence[int]] = None + + initializer: Optional[Initializer] = None + initializer_kwargs: Optional[Mapping[str, Any]] = None + + normalizer: Optional[Normalizer] = None + normalizer_kwargs: Optional[Mapping[str, Any]] = None + + constrainer: Optional[Constrainer] = None + constrainer_kwargs: Optional[Mapping[str, Any]] = None + + # regularizer: Optional[Regularizer] = None + + class Embedding(RepresentationModule): """Trainable embeddings. @@ -104,6 +123,43 @@ def __init__( embedding_dim=embedding_dim, ) + @classmethod + def from_specification( + cls, + num_embeddings: int, + embedding_dim: int, + specification: Optional[EmbeddingSpecification], + device: Optional[torch.device] = None, + ) -> "Embedding": + """ + Create an embedding based on an specification. + + :param num_embeddings: + The number of embeddings. + :param specification: + The specification. + :param device: + If given, move to device. + + :return: + An embedding object. + """ + if specification is None: + specification = EmbeddingSpecification() + embedding = Embedding( + num_embeddings=num_embeddings, + embedding_dim=embedding_dim, + initializer=specification.initializer, + initializer_kwargs=specification.initializer_kwargs, + normalizer=specification.normalizer, + normalizer_kwargs=specification.normalizer_kwargs, + constrainer=specification.constrainer, + constrainer_kwargs=specification.constrainer_kwargs, + ) + if device is not None: + embedding = embedding.to(device=device) + return embedding + @classmethod def init_with_device( cls, From d3b5daf9156927d066408986947074ec89cc4dc0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 22:33:47 +0100 Subject: [PATCH 193/690] small fixes --- src/pykeen/models/base.py | 39 ++++++--------------------- src/pykeen/models/unimodal/simple.py | 7 +++-- src/pykeen/models/unimodal/trans_e.py | 15 +++++++---- 3 files changed, 21 insertions(+), 40 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index bfb720ef23..f3133e6a7a 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1186,9 +1186,6 @@ def __init__( # Default for relation dimensionality if relation_dim is None: relation_dim = embedding_dim - if relation_embedding_specification is None: - relation_embedding_specification = embedding_specification - self.relation_embeddings = Embedding.from_specification( num_embeddings=triples_factory.num_relations, embedding_dim=relation_dim, @@ -1330,15 +1327,11 @@ def __init__( ) # extra embeddings - if second_embedding_specification is None: - second_embedding_specification = embedding_specification self.second_entity_embeddings = Embedding.from_specification( num_embeddings=triples_factory.num_entities, embedding_dim=embedding_dim, specification=second_embedding_specification, ) - if second_relation_embedding_specification is None: - second_relation_embedding_specification = relation_embedding_specification self.second_relation_embeddings = Embedding.from_specification( num_embeddings=triples_factory.num_relations, embedding_dim=embedding_dim, @@ -1386,18 +1379,10 @@ def __init__( preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, - entity_initializer: Optional[Initializer] = None, - entity_initializer_kwargs: Optional[Mapping[str, Any]] = None, - entity_normalizer: Optional[Normalizer] = None, - entity_normalizer_kwargs: Optional[Mapping[str, Any]] = None, - entity_constrainer: Optional[Constrainer] = None, - entity_constrainer_kwargs: Optional[Mapping[str, Any]] = None, - relation_initializer: Optional[Initializer] = None, - relation_initializer_kwargs: Optional[Mapping[str, Any]] = None, - relation_normalizer: Optional[Normalizer] = None, - relation_normalizer_kwargs: Optional[Mapping[str, Any]] = None, - relation_constrainer: Optional[Constrainer] = None, - relation_constrainer_kwargs: Optional[Mapping[str, Any]] = None, + embedding_specification: EmbeddingSpecification = None, + relation_embedding_specification: EmbeddingSpecification = None, + second_embedding_specification: EmbeddingSpecification = None, + second_relation_embedding_specification: EmbeddingSpecification = None, ): super().__init__( triples_factory=triples_factory, @@ -1409,18 +1394,10 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - entity_initializer=entity_initializer, - entity_initializer_kwargs=entity_initializer_kwargs, - entity_normalizer=entity_normalizer, - entity_normalizer_kwargs=entity_normalizer_kwargs, - entity_constrainer=entity_constrainer, - entity_constrainer_kwargs=entity_constrainer_kwargs, - relation_initializer=relation_initializer, - relation_initializer_kwargs=relation_initializer_kwargs, - relation_normalizer=relation_normalizer, - relation_normalizer_kwargs=relation_normalizer_kwargs, - relation_constrainer=relation_constrainer, - relation_constrainer_kwargs=relation_constrainer_kwargs, + embedding_specification=embedding_specification, + relation_embedding_specification=relation_embedding_specification, + second_embedding_specification=second_embedding_specification, + second_relation_embedding_specification=second_relation_embedding_specification, ) self.interaction_function = interaction_function diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index 6bdd9bd9b1..ebe0a17253 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -6,19 +6,19 @@ import torch.autograd -from ..base import TwoSideModel +from ..base import TwoSideEmbeddingModel from ...losses import Loss, SoftplusLoss from ...nn.modules import DistMultInteractionFunction from ...regularizers import PowerSumRegularizer, Regularizer from ...triples import TriplesFactory -from ...typing import DeviceHint, InteractionFunction +from ...typing import DeviceHint __all__ = [ 'SimplE', ] -class SimplE(TwoSideModel): +class SimplE(TwoSideEmbeddingModel): r"""An implementation of SimplE [kazemi2018]_. SimplE is an extension of canonical polyadic (CP), an early tensor factorization approach in which each entity @@ -72,7 +72,6 @@ def __init__( random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, clamp_score: Optional[Union[float, Tuple[float, float]]] = None, - interaction_function: Optional[InteractionFunction] = None, ) -> None: super().__init__( triples_factory=triples_factory, diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index da74ecc353..a31e52964a 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -8,6 +8,7 @@ from .. import SingleVectorEmbeddingModel from ...losses import Loss +from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import TranslationalInteractionFunction from ...regularizers import Regularizer @@ -75,10 +76,14 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - entity_initializer=xavier_uniform_, - relation_initializer=compose( - xavier_uniform_, - functional.normalize, + embedding_specification=EmbeddingSpecification( + initializer=xavier_uniform_, + constrainer=functional.normalize, + ), + relation_embedding_specification=EmbeddingSpecification( + initializer=compose( + xavier_uniform_, + functional.normalize, + ), ), - entity_constrainer=functional.normalize, ) From 495892721fb55faf69caf063b1ace5f60e77c63a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 22:41:01 +0100 Subject: [PATCH 194/690] Use specification --- src/pykeen/models/unimodal/complex.py | 9 ++++++-- src/pykeen/models/unimodal/conv_e.py | 9 ++++++-- src/pykeen/models/unimodal/distmult.py | 23 +++++++++++-------- src/pykeen/models/unimodal/hole.py | 13 +++++++---- src/pykeen/models/unimodal/proj_e.py | 9 ++++++-- src/pykeen/models/unimodal/rotate.py | 11 ++++++--- .../models/unimodal/structured_embedding.py | 13 +++++++---- src/pykeen/models/unimodal/trans_h.py | 10 ++++---- src/pykeen/models/unimodal/trans_r.py | 22 +++++++++++------- src/pykeen/models/unimodal/tucker.py | 9 ++++++-- 10 files changed, 88 insertions(+), 40 deletions(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 8da5cd92fe..04c9c81c08 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -8,6 +8,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss, SoftplusLoss +from ...nn.emb import EmbeddingSpecification from ...nn.modules import ComplExInteractionFunction from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory @@ -106,6 +107,10 @@ def __init__( regularizer=regularizer, # initialize with entity and relation embeddings with standard normal distribution, cf. # https://github.com/ttrouill/complex/blob/dc4eb93408d9a5288c986695b58488ac80b1cc17/efe/models.py#L481-L487 - entity_initializer=nn.init.normal_, - relation_initializer=nn.init.normal_, + embedding_specification=EmbeddingSpecification( + initializer=nn.init.normal_, + ), + relation_embedding_specification=EmbeddingSpecification( + initializer=nn.init.normal_, + ), ) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 90cd124266..033abfde42 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -10,6 +10,7 @@ from ..base import EntityRelationEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss from ...nn import Embedding +from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_normal_ from ...nn.modules import ConvEInteractionFunction from ...regularizers import Regularizer @@ -142,8 +143,12 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - entity_initializer=xavier_normal_, - relation_initializer=xavier_normal_, + embedding_specification=EmbeddingSpecification( + initializer=xavier_normal_, + ), + relation_embedding_specification=EmbeddingSpecification( + initializer=xavier_normal_, + ), ) # ConvE uses one bias for each entity diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index b0a9f0bc0e..f70dee9267 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -9,6 +9,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss +from ...nn.emb import EmbeddingSpecification from ...nn.modules import DistMultInteractionFunction from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory @@ -91,14 +92,18 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - # xavier uniform, cf. - # https://github.com/thunlp/OpenKE/blob/adeed2c0d2bef939807ed4f69c1ea4db35fd149b/models/DistMult.py#L16-L17 - entity_initializer=nn.init.xavier_uniform_, - # Constrain entity embeddings to unit length - entity_constrainer=functional.normalize, - # relations are initialized to unit length (but not constraint) - relation_initializer=compose( - nn.init.xavier_uniform_, - functional.normalize, + embedding_specification=EmbeddingSpecification( + # xavier uniform, cf. + # https://github.com/thunlp/OpenKE/blob/adeed2c0d2bef939807ed4f69c1ea4db35fd149b/models/DistMult.py#L16-L17 + initializer=nn.init.xavier_uniform_, + # Constrain entity embeddings to unit length + constrainer=functional.normalize, + ), + relation_embedding_specification=EmbeddingSpecification( + # relations are initialized to unit length (but not constraint) + initializer=compose( + nn.init.xavier_uniform_, + functional.normalize, + ), ), ) diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 1e9e9ed6e1..68cf707114 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -6,6 +6,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss +from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import HolEInteractionFunction from ...regularizers import Regularizer @@ -74,8 +75,12 @@ def __init__( random_seed=random_seed, regularizer=regularizer, # Initialisation, cf. https://github.com/mnick/scikit-kge/blob/master/skge/param.py#L18-L27 - entity_initializer=xavier_uniform_, - relation_initializer=xavier_uniform_, - entity_constrainer=clamp_norm, - entity_constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + embedding_specification=EmbeddingSpecification( + initializer=xavier_uniform_, + constrainer=clamp_norm, + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + ), + relation_embedding_specification=EmbeddingSpecification( + initializer=xavier_uniform_, + ), ) diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index c2842050bf..3adea5c568 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -8,6 +8,7 @@ from .. import SingleVectorEmbeddingModel from ...losses import Loss +from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import ProjEInteractionFunction from ...regularizers import Regularizer @@ -76,6 +77,10 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - entity_initializer=xavier_uniform_, - relation_initializer=xavier_uniform_, + embedding_specification=EmbeddingSpecification( + initializer=xavier_uniform_, + ), + relation_embedding_specification=EmbeddingSpecification( + initializer=xavier_uniform_, + ), ) diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index e956ad3831..a5a35f172a 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -11,6 +11,7 @@ from .. import SingleVectorEmbeddingModel from ...losses import Loss +from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import RotatEInteraction from ...regularizers import Regularizer @@ -102,7 +103,11 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - entity_initializer=xavier_uniform_, - relation_initializer=init_phases, - relation_constrainer=complex_normalize, + embedding_specification=EmbeddingSpecification( + initializer=xavier_uniform_, + ), + relation_embedding_specification=EmbeddingSpecification( + initializer=init_phases, + constrainer=complex_normalize, + ), ) diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index d04cd503dd..51e940ada6 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -11,6 +11,7 @@ from .. import SingleVectorEmbeddingModel from ...losses import Loss +from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import StructuredEmbeddingInteractionFunction from ...regularizers import Regularizer @@ -64,7 +65,7 @@ def __init__( # Embeddings init_bound = 6 / np.sqrt(embedding_dim) # Initialise relation embeddings to unit length - initializer = compose( + relation_initializer = compose( functools.partial(nn.init.uniform_, a=-init_bound, b=+init_bound), functional.normalize, ) @@ -81,7 +82,11 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - entity_initializer=xavier_uniform_, - entity_constrainer=functional.normalize, - relation_initializer=initializer, + embedding_specification=EmbeddingSpecification( + initializer=xavier_uniform_, + constrainer=functional.normalize, + ), + relation_embedding_specification=EmbeddingSpecification( + initializer=relation_initializer, + ), ) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 7651e88de5..29dae9d63c 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -10,6 +10,7 @@ from ..base import EntityRelationEmbeddingModel from ...losses import Loss from ...nn import Embedding, functional as pkf +from ...nn.emb import EmbeddingSpecification from ...regularizers import Regularizer, TransHRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -90,12 +91,13 @@ def __init__( self.scoring_fct_norm = scoring_fct_norm # embeddings - self.normal_vector_embeddings = Embedding.init_with_device( + self.normal_vector_embeddings = Embedding.from_specification( num_embeddings=triples_factory.num_relations, embedding_dim=embedding_dim, - device=self.device, - # Normalise the normal vectors by their l2 norms - constrainer=functional.normalize, + specification=EmbeddingSpecification( + # Normalise the normal vectors by their l2 norms + constrainer=functional.normalize, + ), ) def post_parameter_update(self) -> None: # noqa: D102 diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 0cca23a3fe..01af972d22 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -13,6 +13,7 @@ from ..base import EntityRelationEmbeddingModel from ...losses import Loss from ...nn import Embedding +from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import TransRInteractionFunction from ...regularizers import Regularizer @@ -95,15 +96,20 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - entity_initializer=xavier_uniform_, - entity_constrainer=clamp_norm, - entity_constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), - relation_initializer=compose( - xavier_uniform_, - functional.normalize, + embedding_specification=EmbeddingSpecification( + initializer=xavier_uniform_, + constrainer=clamp_norm, + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), - relation_constrainer=clamp_norm, - relation_constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + relation_embedding_specification=EmbeddingSpecification( + initializer=compose( + xavier_uniform_, + functional.normalize, + ), + constrainer=clamp_norm, + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + ), + ) self.interaction_function = TransRInteractionFunction(p=scoring_fct_norm) diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index 723469412f..718d7caef0 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -6,6 +6,7 @@ from .. import SingleVectorEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss +from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_normal_ from ...nn.modules import TuckerInteractionFunction from ...regularizers import Regularizer @@ -103,6 +104,10 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - entity_initializer=xavier_normal_, - relation_initializer=xavier_normal_, + embedding_specification=EmbeddingSpecification( + initializer=xavier_normal_, + ), + relation_embedding_specification=EmbeddingSpecification( + initializer=xavier_normal_, + ), ) From 618e398f8f61a16e4c1a78fdf7b32c7cb4ca78ff Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 22:44:44 +0100 Subject: [PATCH 195/690] Fix TransD --- src/pykeen/models/unimodal/trans_d.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 26836c1e49..f5b56a395b 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -11,6 +11,7 @@ from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_normal_ +from ...nn.modules import TransDInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -137,6 +138,7 @@ def __init__( constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), ) + self.interaction_function = TransDInteractionFunction() def _forward( self, From d761931938971f0b774cd77e4ba1cdfd2e14e5c3 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 22:46:39 +0100 Subject: [PATCH 196/690] Fix TransD --- src/pykeen/models/unimodal/trans_d.py | 4 ++-- src/pykeen/nn/modules.py | 27 +++------------------------ 2 files changed, 5 insertions(+), 26 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index f5b56a395b..37f5bd7b6f 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -11,7 +11,7 @@ from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_normal_ -from ...nn.modules import TransDInteractionFunction +from ...nn.modules import TranslationalInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -138,7 +138,7 @@ def __init__( constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), ) - self.interaction_function = TransDInteractionFunction() + self.interaction_function = TranslationalInteractionFunction(p=2, power_norm=True) def _forward( self, diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 32425c0e13..64ec5fb839 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -264,13 +264,14 @@ def forward( class TranslationalInteractionFunction(InteractionFunction): """The translational interaction function shared by the TransE, TransR, TransH, and other Trans models.""" - def __init__(self, p: int): + def __init__(self, p: int, power_norm: bool = False): """Initialize the translational interaction function. :param p: The norm used with :func:`torch.norm`. Typically is 1 or 2. """ super().__init__() self.p = p + self.power_norm = power_norm def forward( self, @@ -280,7 +281,7 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa:D102 self._check_for_empty_kwargs(kwargs=kwargs) - return pkf.translational_interaction(h=h, r=r, t=t, p=self.p) + return pkf.translational_interaction(h=h, r=r, t=t, p=self.p, power_norm=self.power_norm) #: Interaction function of ComplEx @@ -588,28 +589,6 @@ def forward( return pkf.ermlpe_interaction(h=h, r=r, t=t, mlp=self.mlp) -class TransDInteractionFunction(TranslationalInteractionFunction): - """The interaction function for TransD.""" - - def __init__(self, p: int = 2, power: int = 2): - """Initialize the TransD interaction function. - - :param p: The norm applied by :func:`torch.norm` - :param power: The power applied after :func:`torch.norm`. - """ - super().__init__(p=p) - self.power = power - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa:D102 - return super().forward(h=h, r=r, t=t, **kwargs) ** self.power - - class TransRInteractionFunction(InteractionFunction): """The TransR interaction function.""" From 768d9ed071edd7a7d5943501cf3bff3cd7542bca Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 23:03:44 +0100 Subject: [PATCH 197/690] Fix broadcast_cat --- src/pykeen/nn/functional.py | 8 ++++---- src/pykeen/utils.py | 2 ++ 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index c43da38c3f..cc7ad57c1e 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -153,7 +153,7 @@ def conve_interaction( h = h.view(*h.shape[:-1], input_channels, embedding_height, embedding_width) r = r.unsqueeze(dim=1) r = r.view(*r.shape[:-1], input_channels, embedding_height, embedding_width) - x = broadcast_cat(h, r, dim=2).view(-1, input_channels, 2 * embedding_height, embedding_width) + x = broadcast_cat(h, r, dim=-2).view(-1, input_channels, 2 * embedding_height, embedding_width) # batch_size, num_input_channels, 2*height, width if bn0 is not None: @@ -180,7 +180,7 @@ def conve_interaction( x = bn2(x) x = activation(x) - # reshape: (batch_size', embedding_dim) + # reshape: (batch_size', embedding_dim) -> (b, h, r, 1, d) x = x.view(-1, num_heads, num_relations, 1, embedding_dim) # For efficient calculation, each of the convolved [h, r] rows has only to be multiplied with one t row @@ -1038,13 +1038,13 @@ def tucker_interaction( The scores. """ # Compute wr = DO(W x_2 r) - x = do0(_extended_einsum("ijd,brd->brij", core_tensor, r)) + x = do0(_extended_einsum("idj,brd->brij", core_tensor, r)) # Compute h_n = DO(BN(h)) h = _apply_optional_bn_to_tensor(batch_norm=bn1, output_dropout=do1, tensor=h) # compute whr = DO(BN(h_n x_1 wr)) - x = _extended_einsum("brid,bhd->bhri", h, x) + x = _extended_einsum("brid,bhd->bhri", x, h) x = _apply_optional_bn_to_tensor(batch_norm=bn2, tensor=x, output_dropout=do2) # Compute whr x_3 t diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 45f632b3de..10c795a091 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -516,6 +516,8 @@ def broadcast_cat( """ if x.ndimension() != y.ndimension(): raise ValueError + if dim < 0: + dim = x.ndimension() + dim x_rep, y_rep = [], [] for d, (xd, yd) in enumerate(zip(x.shape, y.shape)): xr = yr = 1 From a00a302164df11903c9b3ea9c089ca10214ece25 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 23:08:56 +0100 Subject: [PATCH 198/690] Fix tucker --- src/pykeen/nn/functional.py | 2 +- tests/test_models.py | 9 +++++---- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index cc7ad57c1e..da20206cfb 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -988,7 +988,7 @@ def _apply_optional_bn_to_tensor( ) -> torch.FloatTensor: if batch_norm is not None: shape = tensor.shape - tensor = tensor.view(-1, shape[-1]) + tensor = tensor.reshape(-1, shape[-1]) tensor = batch_norm(tensor) tensor = tensor.view(*shape) tensor = output_dropout(tensor) diff --git a/tests/test_models.py b/tests/test_models.py index c2fef792b8..622c857c19 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -29,7 +29,7 @@ Model, MultimodalModel, SingleVectorEmbeddingModel, - _extend_batch, + TwoVectorEmbeddingModel, _extend_batch, get_novelty_mask, ) from pykeen.models.cli import build_cli_from_cls @@ -667,12 +667,13 @@ def _check_constraints(self): * Entity and relation embeddings have to have at most unit L2 norm. * Covariances have to have values between c_min and c_max """ - low = self.model.entity_covariances.constrainer.keywords['min'] - high = self.model.entity_covariances.constrainer.keywords['max'] + self.model: TwoVectorEmbeddingModel + low = self.model.second_entity_embeddings.constrainer.keywords['min'] + high = self.model.second_entity_embeddings.constrainer.keywords['max'] for embedding in (self.model.entity_embeddings, self.model.relation_embeddings): assert all_in_bounds(embedding(indices=None).norm(p=2, dim=-1), high=1., a_tol=_EPSILON) - for cov in (self.model.entity_covariances, self.model.relation_covariances): + for cov in (self.model.second_entity_embeddings, self.model.second_relation_embeddings): assert all_in_bounds(cov(indices=None), low=low, high=high) From 845c5df395de176d600f61832399f5355260c0f8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 23:26:37 +0100 Subject: [PATCH 199/690] Extract TransD functional form --- src/pykeen/models/base.py | 2 +- src/pykeen/models/unimodal/trans_d.py | 56 +--------------- src/pykeen/nn/functional.py | 95 +++++++++++++++++++++++++++ src/pykeen/nn/modules.py | 2 +- tests/test_models.py | 6 +- 5 files changed, 101 insertions(+), 60 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index f3133e6a7a..041930c1f9 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1334,7 +1334,7 @@ def __init__( ) self.second_relation_embeddings = Embedding.from_specification( num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim, + embedding_dim=relation_dim, specification=second_relation_embedding_specification, ) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 37f5bd7b6f..a84e176662 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -4,14 +4,13 @@ from typing import Optional -import torch import torch.autograd from ..base import TwoVectorEmbeddingModel from ...losses import Loss +from ...nn import functional as F from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_normal_ -from ...nn.modules import TranslationalInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -22,53 +21,6 @@ ] -def _project_entity( - e: torch.FloatTensor, - e_p: torch.FloatTensor, - r: torch.FloatTensor, - r_p: torch.FloatTensor, -) -> torch.FloatTensor: - r"""Project entity relation-specific. - - .. math:: - - e_{\bot} = M_{re} e - = (r_p e_p^T + I^{d_r \times d_e}) e - = r_p e_p^T e + I^{d_r \times d_e} e - = r_p (e_p^T e) + e' - - and additionally enforces - - .. math:: - - \|e_{\bot}\|_2 \leq 1 - - :param e: shape: (batch_size, num_entities, d_e) - The entity embedding. - :param e_p: shape: (batch_size, num_entities, d_e) - The entity projection. - :param r: shape: (batch_size, num_entities, d_r) - The relation embedding. - :param r_p: shape: (batch_size, num_entities, d_r) - The relation projection. - - :return: shape: (batch_size, num_entities, d_r) - - """ - # The dimensions affected by e' - change_dim = min(e.shape[-1], r.shape[-1]) - - # Project entities - # r_p (e_p.T e) + e' - e_bot = r_p * torch.sum(e_p * e, dim=-1, keepdim=True) - e_bot[:, :, :change_dim] += e[:, :, :change_dim] - - # Enforce constraints - e_bot = clamp_norm(e_bot, p=2, dim=-1, maxnorm=1) - - return e_bot - - class TransD(TwoVectorEmbeddingModel): r"""An implementation of TransD from [ji2015]_. @@ -138,7 +90,6 @@ def __init__( constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), ) - self.interaction_function = TranslationalInteractionFunction(p=2, power_norm=True) def _forward( self, @@ -149,7 +100,4 @@ def _forward( t1: torch.FloatTensor, t2: torch.FloatTensor, ) -> torch.FloatTensor: # noqa:D102 - # Project entities - h_bot = _project_entity(e=h1, e_p=h2, r=r1, r_p=r2) - t_bot = _project_entity(e=t1, e_p=t2, r=r1, r_p=r2) - return self.interaction_function(h=h_bot, r=r1, t=t_bot) + return F.transd_interaction(h=h1, r=r1, t=t1, h_p=h2, r_p=r2, t_p=t2, p=2, power_norm=True) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index da20206cfb..03b5f66912 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -1075,3 +1075,98 @@ def unstructured_model_interaction( h = h.unsqueeze(dim=2).unsqueeze(dim=3) t = t.unsqueeze(dim=1).unsqueeze(dim=2) return negative_norm_of_sum(h, -t, p=p, power_norm=power_norm) + + +def _project_entity( + e: torch.FloatTensor, + e_p: torch.FloatTensor, + r_p: torch.FloatTensor, +) -> torch.FloatTensor: + r"""Project entity relation-specific. + + .. math:: + + e_{\bot} = M_{re} e + = (r_p e_p^T + I^{d_r \times d_e}) e + = r_p e_p^T e + I^{d_r \times d_e} e + = r_p (e_p^T e) + e' + + and additionally enforces + + .. math:: + + \|e_{\bot}\|_2 \leq 1 + + :param e: shape: (..., d_e) + The entity embedding. + :param e_p: shape: (..., d_e) + The entity projection. + :param r_p: shape: (..., d_r) + The relation projection. + + :return: shape: (..., d_r) + + """ + # The dimensions affected by e' + change_dim = min(e.shape[-1], r_p.shape[-1]) + + # Project entities + # r_p (e_p.T e) + e' + e_bot = r_p * torch.sum(e_p * e, dim=-1, keepdim=True) + e_bot[..., :change_dim] += e[..., :change_dim] + + # Enforce constraints + e_bot = clamp_norm(e_bot, p=2, dim=-1, maxnorm=1) + + return e_bot + + +def transd_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + h_p: torch.FloatTensor, + r_p: torch.FloatTensor, + t_p: torch.FloatTensor, + p: int, + power_norm: bool = False, +) -> torch.FloatTensor: + """ + Evaluate the DistMult interaction function. + + :param h: shape: (batch_size, num_heads, d_e) + The head representations. + :param r: shape: (batch_size, num_relations, d_r) + The relation representations. + :param t: shape: (batch_size, num_tails, d_e) + The tail representations. + :param h_p: shape: (batch_size, num_heads, d_e) + The head projections. + :param r_p: shape: (batch_size, num_relations, d_r) + The relation projections. + :param t_p: shape: (batch_size, num_tails, d_e) + The tail projections. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + # Project entities + # shape: (b, h, r, 1, d_r) + h_bot = _project_entity( + e=h.unsqueeze(dim=2), + e_p=h_p.unsqueeze(dim=2), + r_p=r_p.unsqueeze(dim=1), + ).unsqueeze(dim=-2) + # shape: (b, 1, r, t, d_r) + t_bot = _project_entity( + e=t.unsqueeze(dim=1), + e_p=t_p.unsqueeze(dim=1), + r_p=r_p.unsqueeze(dim=2) + ).unsqueeze(dim=1) + return _translational_interaction( + h=h_bot, + r=r.view(r.shape[0], 1, r.shape[1], 1, r.shape[2]), + t=t_bot, + p=p, + power_norm=power_norm, + ) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 64ec5fb839..3f4f1ab317 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -281,7 +281,7 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa:D102 self._check_for_empty_kwargs(kwargs=kwargs) - return pkf.translational_interaction(h=h, r=r, t=t, p=self.p, power_norm=self.power_norm) + return pkf._translational_interaction(h=h, r=r, t=t, p=self.p, power_norm=self.power_norm) #: Interaction function of ComplEx diff --git a/tests/test_models.py b/tests/test_models.py index 622c857c19..8b410e1807 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -38,8 +38,8 @@ inverse_outdegree_edge_weights, symmetric_edge_weights, ) -from pykeen.models.unimodal.trans_d import _project_entity from pykeen.nn import Embedding, RepresentationModule +from pykeen.nn.functional import _project_entity from pykeen.training import LCWATrainingLoop, SLCWATrainingLoop, TrainingLoop from pykeen.triples import TriplesFactory from pykeen.utils import all_in_bounds, clamp_norm, set_random_seed @@ -974,12 +974,10 @@ def test_project_entity(self): e_p = torch.rand(1, self.model.num_entities, self.embedding_dim, generator=self.generator) # random relation embeddings & projections - r = torch.rand(self.batch_size, 1, self.model.relation_dim, generator=self.generator) - r = clamp_norm(r, maxnorm=1, p=2, dim=-1) r_p = torch.rand(self.batch_size, 1, self.model.relation_dim, generator=self.generator) # project - e_bot = _project_entity(e=e, e_p=e_p, r=r, r_p=r_p) + e_bot = _project_entity(e=e, e_p=e_p, r_p=r_p) # check shape: assert e_bot.shape == (self.batch_size, self.model.num_entities, self.model.relation_dim) From 1f3f1ae276dd1fd1a5e09564de350461cdcbec11 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 23:27:22 +0100 Subject: [PATCH 200/690] Fix docstring --- src/pykeen/nn/functional.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 03b5f66912..160f2ab4ec 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -1146,6 +1146,10 @@ def transd_interaction( The relation projections. :param t_p: shape: (batch_size, num_tails, d_e) The tail projections. + :param p: + The parameter p for selecting the norm. + :param power_norm: + Whether to return the powered norm instead. :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. @@ -1163,10 +1167,5 @@ def transd_interaction( e_p=t_p.unsqueeze(dim=1), r_p=r_p.unsqueeze(dim=2) ).unsqueeze(dim=1) - return _translational_interaction( - h=h_bot, - r=r.view(r.shape[0], 1, r.shape[1], 1, r.shape[2]), - t=t_bot, - p=p, - power_norm=power_norm, - ) + r = r.view(r.shape[0], 1, r.shape[1], 1, r.shape[2]) + return _translational_interaction(h=h_bot, r=r, t=t_bot, p=p, power_norm=power_norm) From 3d9213bc3815a4972deac2094a0762f8e59bbc33 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 23:33:23 +0100 Subject: [PATCH 201/690] Add TransD interaction function --- src/pykeen/nn/modules.py | 73 +++++++++++++++++++--------------------- 1 file changed, 35 insertions(+), 38 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 3f4f1ab317..aaa314a29c 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -4,6 +4,7 @@ import logging import math +from abc import ABC from typing import Any, Callable, Mapping, Optional, Sequence, Tuple, Type import torch @@ -261,18 +262,26 @@ def forward( return FunctionalInteractionFunction -class TranslationalInteractionFunction(InteractionFunction): +class TranslationalInteractionFunction(InteractionFunction, ABC): """The translational interaction function shared by the TransE, TransR, TransH, and other Trans models.""" def __init__(self, p: int, power_norm: bool = False): """Initialize the translational interaction function. - :param p: The norm used with :func:`torch.norm`. Typically is 1 or 2. + :param p: + The norm used with :func:`torch.norm`. Typically is 1 or 2. + :param power_norm: + Whether to use the p-th power of the L_p norm. It has the advantage of being differentiable around 0, + and numerically more stable. """ super().__init__() self.p = p self.power_norm = power_norm + +class TransEInteractionFunction(TranslationalInteractionFunction): + """The TransE interaction function.""" + def forward( self, h: torch.FloatTensor, @@ -589,16 +598,11 @@ def forward( return pkf.ermlpe_interaction(h=h, r=r, t=t, mlp=self.mlp) -class TransRInteractionFunction(InteractionFunction): +class TransRInteractionFunction(TranslationalInteractionFunction): """The TransR interaction function.""" - def __init__(self, p: int): - """Initialize the TransR interaction function. - - :param p: The norm applied to the translation - """ - super().__init__() - self.p = p + def __init__(self, p: int, power_norm: bool = True): + super().__init__(p=p, power_norm=power_norm) def forward( self, @@ -609,7 +613,7 @@ def forward( ) -> torch.FloatTensor: # noqa:D102 m_r = kwargs.pop('m_r') self._check_for_empty_kwargs(kwargs=kwargs) - return pkf.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p, power_norm=True) + return pkf.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p, power_norm=self.power_norm) #: Interaction function of RotatE. @@ -670,23 +674,9 @@ def forward( RESCALInteractionFunction = _build_module_from_stateless(pkf.rescal_interaction) -class StructuredEmbeddingInteractionFunction(InteractionFunction): +class StructuredEmbeddingInteractionFunction(TranslationalInteractionFunction): """Interaction function of Structured Embedding.""" - def __init__( - self, - p: int, - power_norm: bool = False, - ): - """Initialize the SE interaction function. - - :param p: The l_p norm - :param power_norm: Should power normalization be applied? - """ - super().__init__() - self.p = p - self.power_norm = power_norm - def forward( self, h: torch.FloatTensor, @@ -698,14 +688,7 @@ def forward( rh, rt = r.split(dim ** 2, dim=-1) rh = rh.view(*rh.shape[:-1], dim, dim) rt = rt.view(*rt.shape[:-1], dim, dim) - return pkf.structured_embedding_interaction( - h=h, - r_h=rh, - r_t=rt, - t=t, - p=self.p, - power_norm=self.power_norm, - ) + return pkf.structured_embedding_interaction(h=h, r_h=rh, r_t=rt, t=t, p=self.p, power_norm=self.power_norm) class TuckerInteractionFunction(InteractionFunction): @@ -777,13 +760,11 @@ def forward( ) -class UnstructuredModelInteractionFunction(InteractionFunction): +class UnstructuredModelInteractionFunction(TranslationalInteractionFunction): """Interaction function of UnstructuredModel.""" def __init__(self, p: int, power_norm: bool = True): - super().__init__() - self.p = p - self.power_norm = power_norm + super().__init__(p=p, power_norm=power_norm) def forward( self, @@ -794,3 +775,19 @@ def forward( ) -> torch.FloatTensor: self._check_for_empty_kwargs(kwargs=kwargs) return pkf.unstructured_model_interaction(h, t, p=self.p, power_norm=self.power_norm) + + +class TransDInteractionFunction(TranslationalInteractionFunction): + """Interaction function of TransD.""" + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: + h_p = kwargs.pop("h_p") + r_p = kwargs.pop("r_p") + t_p = kwargs.pop("t_p") + return pkf.transd_interaction(h=h, r=r, t=t, h_p=h_p, r_p=r_p, t_p=t_p, p=self.p, power_norm=self.power_norm) From 07d48f9cbf81a76b3c428d887d02dd42b906fc32 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 23:34:07 +0100 Subject: [PATCH 202/690] Use interaction function --- src/pykeen/models/unimodal/trans_d.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index a84e176662..e5a34953ae 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -8,9 +8,9 @@ from ..base import TwoVectorEmbeddingModel from ...losses import Loss -from ...nn import functional as F from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_normal_ +from ...nn.modules import TransDInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -90,6 +90,7 @@ def __init__( constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), ) + self.interaction_function = TransDInteractionFunction(p=2, power_norm=True) def _forward( self, @@ -100,4 +101,4 @@ def _forward( t1: torch.FloatTensor, t2: torch.FloatTensor, ) -> torch.FloatTensor: # noqa:D102 - return F.transd_interaction(h=h1, r=r1, t=t1, h_p=h2, r_p=r2, t_p=t2, p=2, power_norm=True) + return self.interaction_function(h=h1, r=r1, t=t1, h_p=h2, r_p=r2, t_p=t2, p=2, power_norm=True) From e2940d3b1faa08d6b640cc0dea54b05ef5b92571 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 23:35:34 +0100 Subject: [PATCH 203/690] Fix unittest --- tests/test_models.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 8b410e1807..b40bef1d06 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -846,6 +846,8 @@ def _check_constraints(self): def test_score_hrt_manual(self): """Manually test interaction function of TransD.""" + self.model: TwoVectorEmbeddingModel + # entity embeddings weights = torch.as_tensor(data=[[2., 2.], [4., 4.]], dtype=torch.float) entity_embeddings = Embedding( @@ -861,7 +863,7 @@ def test_score_hrt_manual(self): embedding_dim=2, ) entity_projection_embeddings._embeddings.weight.data.copy_(projection_weights) - self.model.entity_projections = entity_projection_embeddings + self.model.second_entity_embeddings = entity_projection_embeddings # relation embeddings relation_weights = torch.as_tensor(data=[[4.], [4.]], dtype=torch.float) @@ -878,7 +880,7 @@ def test_score_hrt_manual(self): embedding_dim=1, ) relation_projection_embeddings._embeddings.weight.data.copy_(relation_projection_weights) - self.model.relation_projections = relation_projection_embeddings + self.model.second_relation_embeddings = relation_projection_embeddings # Compute Scores # FIXME batch is wrong size? @@ -905,7 +907,7 @@ def test_score_hrt_manual(self): embedding_dim=3, ) relation_projection_embeddings._embeddings.weight.data.copy_(relation_projection_weights) - self.model.relation_projections = relation_projection_embeddings + self.model.second_relation_embeddings = relation_projection_embeddings # Compute Scores batch = torch.as_tensor(data=[[0, 0, 0]], dtype=torch.long) @@ -937,7 +939,7 @@ def test_score_hrt_manual(self): embedding_dim=3, ) entity_projection_embeddings._embeddings.weight.data.copy_(projection_weights) - self.model.entity_projections = entity_projection_embeddings + self.model.second_entity_embeddings = entity_projection_embeddings # relation embeddings relation_weights = torch.as_tensor(data=[[3., 3.], [3., 3.]], dtype=torch.float) @@ -954,7 +956,7 @@ def test_score_hrt_manual(self): embedding_dim=2, ) relation_projection_embeddings._embeddings.weight.data.copy_(relation_projection_weights) - self.model.relation_projections = relation_projection_embeddings + self.model.second_relation_embeddings = relation_projection_embeddings # Compute Scores batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 0]], dtype=torch.long) From 2de40f019340b6259200fc5dc77ba0429c638e73 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 23:35:50 +0100 Subject: [PATCH 204/690] Remove obsolete FIXME --- tests/test_models.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index b40bef1d06..fcf150c100 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -883,7 +883,6 @@ def test_score_hrt_manual(self): self.model.second_relation_embeddings = relation_projection_embeddings # Compute Scores - # FIXME batch is wrong size? batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 1]], dtype=torch.long) scores = self.model.score_hrt(hrt_batch=batch) self.assertEqual(scores.shape[0], 2) From 9a54c5b84badbfda768f366c9fd9a015182ffe5f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 23:37:21 +0100 Subject: [PATCH 205/690] Use resolved relation embedding for the second relation embedding --- src/pykeen/models/base.py | 4 ++-- src/pykeen/nn/emb.py | 2 ++ 2 files changed, 4 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 041930c1f9..9e3b7fdd3e 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1329,12 +1329,12 @@ def __init__( # extra embeddings self.second_entity_embeddings = Embedding.from_specification( num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, + embedding_dim=self.embedding_dim, specification=second_embedding_specification, ) self.second_relation_embeddings = Embedding.from_specification( num_embeddings=triples_factory.num_relations, - embedding_dim=relation_dim, + embedding_dim=self.relation_dim, specification=second_relation_embedding_specification, ) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index 66a99dd98e..806889a021 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -136,6 +136,8 @@ def from_specification( :param num_embeddings: The number of embeddings. + :param embedding_dim: + The embedding dimension. :param specification: The specification. :param device: From b46e0e9f36d55b7a4398bbc61fd3617a3121c7fc Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 23:39:52 +0100 Subject: [PATCH 206/690] Make post_parameter update recursive --- src/pykeen/models/base.py | 16 +++++----------- src/pykeen/models/unimodal/rgcn.py | 5 ----- src/pykeen/models/unimodal/trans_h.py | 4 ---- 3 files changed, 5 insertions(+), 20 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 9e3b7fdd3e..19cebf787e 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -865,6 +865,11 @@ def make_labeled_df( def post_parameter_update(self) -> None: """Has to be called after each parameter update.""" self.regularizer.reset() + for module in self.modules(): + if module is self: + continue + if hasattr(module, "post_parameter_update"): + module.post_parameter_update() def regularize_if_necessary(self, *tensors: torch.FloatTensor) -> None: """Update the regularizer's term given some tensors, if regularization is requested. @@ -1136,11 +1141,6 @@ def embedding_dim(self) -> int: # noqa:D401 """The entity embedding dimension.""" return self.entity_embeddings.embedding_dim - def post_parameter_update(self) -> None: # noqa: D102 - # make sure to call this first, to reset regularizer state! - super().post_parameter_update() - self.entity_embeddings.post_parameter_update() - class EntityRelationEmbeddingModel(Model, ABC): """A base module for KGE models that have different embeddings for entities and relations.""" @@ -1202,12 +1202,6 @@ def relation_dim(self): # noqa:D401 """The relation embedding dimension.""" return self.relation_embeddings.embedding_dim - def post_parameter_update(self) -> None: # noqa: D102 - # make sure to call this first, to reset regularizer state! - super().post_parameter_update() - self.entity_embeddings.post_parameter_update() - self.relation_embeddings.post_parameter_update() - def _can_slice(fn) -> bool: return 'slice_size' in inspect.getfullargspec(fn).args diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index 72e1d905e4..508f077eb6 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -541,11 +541,6 @@ def __init__( # TODO: Dummy self.decoder = Decoder() - def post_parameter_update(self) -> None: # noqa: D102 - super().post_parameter_update() - self.entity_representations.post_parameter_update() - self.relation_embeddings.post_parameter_update() - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 # Enrich embeddings h = self.entity_representations(indices=hrt_batch[:, 0]) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 29dae9d63c..b35c78c832 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -100,10 +100,6 @@ def __init__( ), ) - def post_parameter_update(self) -> None: # noqa: D102 - super().post_parameter_update() - self.normal_vector_embeddings.post_parameter_update() - def regularize_if_necessary(self) -> None: """Update the regularizer's term given some tensors, if regularization is requested.""" # As described in [wang2014], all entities and relations are used to compute the regularization term From dd01c72bd7c5c45a6bd6a9ecd76e83dfa90befb0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 13 Nov 2020 23:53:59 +0100 Subject: [PATCH 207/690] Remove unused method --- src/pykeen/utils.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 10c795a091..4d3d37f0b5 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -44,7 +44,6 @@ 'NoRandomSeedNecessary', 'Result', 'fix_dataclass_init_docs', - 'normalize_for_einsum', 'get_hrt_indices', 'get_hr_indices', 'get_ht_indices', @@ -439,29 +438,6 @@ def random_non_negative_int() -> int: return int(sq.generate_state(1)[0]) -def normalize_for_einsum( - x: torch.FloatTensor, - batch_size: int, - symbol: str, -) -> Tuple[str, torch.FloatTensor]: - """ - Normalize tensor for broadcasting along batch-dimension in einsum. - - :param x: - The tensor. - :param batch_size: - The batch_size - :param symbol: - The symbol for the einsum term. - - :return: - A tuple (reshaped_tensor, term). - """ - if x.shape[0] == batch_size: - return f'b{symbol}d', x - return f'{symbol}d', x.squeeze(dim=0) - - def check_shapes( *x: Tuple[torch.Tensor, str], raise_or_error: bool = True, From de28d914f3ed785bc065ee7818116801db637595 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 14 Nov 2020 11:26:53 +0100 Subject: [PATCH 208/690] Pass flake8 --- src/pykeen/models/base.py | 2 +- src/pykeen/models/unimodal/kg2e.py | 2 +- src/pykeen/nn/functional.py | 2 +- src/pykeen/nn/modules.py | 8 ++++---- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 19cebf787e..d0195f0631 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1405,7 +1405,7 @@ def forward( self.interaction_function( h_source.get_in_canonical_shape(indices=h_indices), r_source.get_in_canonical_shape(indices=r_indices), - t_source.get_in_canonical_shape(indices=t_indices) + t_source.get_in_canonical_shape(indices=t_indices), ) for h_source, r_source, t_source in ( (self.entity_embeddings, self.relation_embeddings, self.second_entity_embeddings), diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index faecf77c30..d8861f4ffd 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -104,7 +104,7 @@ def __init__( # Ensure positive definite covariances matrices and appropriate size by clamping constrainer=torch.clamp, constrainer_kwargs=dict(min=c_min, max=c_max), - ) + ), ) # Similarity function used for distributions dist_similarity = dist_similarity.upper() diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 160f2ab4ec..709cb19ec3 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -1165,7 +1165,7 @@ def transd_interaction( t_bot = _project_entity( e=t.unsqueeze(dim=1), e_p=t_p.unsqueeze(dim=1), - r_p=r_p.unsqueeze(dim=2) + r_p=r_p.unsqueeze(dim=2), ).unsqueeze(dim=1) r = r.view(r.shape[0], 1, r.shape[1], 1, r.shape[2]) return _translational_interaction(h=h_bot, r=r, t=t_bot, p=p, power_norm=power_norm) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index aaa314a29c..0646f84763 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -246,7 +246,7 @@ def _build_module_from_stateless( ) -> Type[InteractionFunction]: """Build a stateless interaction function module with a pre-defined functional interface.""" - class FunctionalInteractionFunction(InteractionFunction): + class StatelessInteractionFunction(InteractionFunction): """Interaction function without state or additional parameters.""" def forward( @@ -259,7 +259,7 @@ def forward( self._check_for_empty_kwargs(kwargs) return f(h, r, t) - return FunctionalInteractionFunction + return StatelessInteractionFunction class TranslationalInteractionFunction(InteractionFunction, ABC): @@ -772,7 +772,7 @@ def forward( r: torch.FloatTensor, t: torch.FloatTensor, **kwargs, - ) -> torch.FloatTensor: + ) -> torch.FloatTensor: # noqa:D102 self._check_for_empty_kwargs(kwargs=kwargs) return pkf.unstructured_model_interaction(h, t, p=self.p, power_norm=self.power_norm) @@ -786,7 +786,7 @@ def forward( r: torch.FloatTensor, t: torch.FloatTensor, **kwargs, - ) -> torch.FloatTensor: + ) -> torch.FloatTensor: # noqa:D102 h_p = kwargs.pop("h_p") r_p = kwargs.pop("r_p") t_p = kwargs.pop("t_p") From 85876dbae507269922cf0709fdea091595456662 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 14 Nov 2020 11:49:22 +0100 Subject: [PATCH 209/690] Revert test --- tests/test_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index fcf150c100..8303e9b3c1 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -485,7 +485,7 @@ def test_reset_parameters_constructor_call(self): ) except TypeError as error: assert error.args == ("'NoneType' object is not callable",) - mock_method.assert_called() + mock_method.assert_called_once() def test_custom_representations(self): """Tests whether we can provide custom representations.""" From 5978b73704e51de8bd70dce50ad52805513129fc Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 14 Nov 2020 11:49:30 +0100 Subject: [PATCH 210/690] Update error info --- tests/test_models.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 8303e9b3c1..3e2591ed2b 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -423,7 +423,7 @@ def test_score_h_with_score_hrt_equality(self) -> None: scores_h = self.model.score_h(batch) scores_hrt = super(self.model.__class__, self.model).score_h(batch) except NotImplementedError: - self.fail(msg='Score_h not yet implemented') + self.fail(msg=f'{self.model.__class__.__name__}.score_h() has not yet been implemented') except RuntimeError as e: if str(e) == 'fft: ATen not compiled with MKL support': self.skipTest(str(e)) @@ -444,7 +444,7 @@ def test_score_r_with_score_hrt_equality(self) -> None: scores_r = self.model.score_r(batch) scores_hrt = super(self.model.__class__, self.model).score_r(batch) except NotImplementedError: - self.fail(msg='Score_h not yet implemented') + self.fail(msg=f'{self.model.__class__.__name__}.score_r() has not yet been implemented') except RuntimeError as e: if str(e) == 'fft: ATen not compiled with MKL support': self.skipTest(str(e)) @@ -465,7 +465,7 @@ def test_score_t_with_score_hrt_equality(self) -> None: scores_t = self.model.score_t(batch) scores_hrt = super(self.model.__class__, self.model).score_t(batch) except NotImplementedError: - self.fail(msg='Score_h not yet implemented') + self.fail(msg=f'{self.model.__class__.__name__}.score_t() has not yet been implemented') except RuntimeError as e: if str(e) == 'fft: ATen not compiled with MKL support': self.skipTest(str(e)) From 5661c2aadd77c6736771ddb7f5a4bc71d562365a Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 14 Nov 2020 11:49:39 +0100 Subject: [PATCH 211/690] Get rid of extraneous logged tqdms --- tests/test_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index 3e2591ed2b..b471cc6fba 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -266,7 +266,7 @@ def test_train_lcwa(self) -> None: def _safe_train_loop(self, loop: TrainingLoop, num_epochs, batch_size, sampler): try: - losses = loop.train(num_epochs=num_epochs, batch_size=batch_size, sampler=sampler) + losses = loop.train(num_epochs=num_epochs, batch_size=batch_size, sampler=sampler, use_tqdm=False) except RuntimeError as e: if str(e) == 'fft: ATen not compiled with MKL support': self.skipTest(str(e)) From 1cd1ca962911258da6de47172a79433f26be91f7 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 14 Nov 2020 11:49:47 +0100 Subject: [PATCH 212/690] Fix imports --- tests/test_models.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index b471cc6fba..93afc32a6d 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -29,7 +29,9 @@ Model, MultimodalModel, SingleVectorEmbeddingModel, - TwoVectorEmbeddingModel, _extend_batch, + TwoSideEmbeddingModel, + TwoVectorEmbeddingModel, + _extend_batch, get_novelty_mask, ) from pykeen.models.cli import build_cli_from_cls @@ -51,6 +53,8 @@ EntityEmbeddingModel.__name__, EntityRelationEmbeddingModel.__name__, SingleVectorEmbeddingModel.__name__, + TwoVectorEmbeddingModel.__name__, + TwoSideEmbeddingModel.__name__, 'MockModel', 'models', 'get_model_cls', From eb1b5da1da1e23816df133b1cf6b0aed3788d8a6 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 14 Nov 2020 11:49:50 +0100 Subject: [PATCH 213/690] Update rgcn.py --- src/pykeen/models/unimodal/rgcn.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index 508f077eb6..5ba8c67141 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -541,9 +541,13 @@ def __init__( # TODO: Dummy self.decoder = Decoder() - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Enrich embeddings - h = self.entity_representations(indices=hrt_batch[:, 0]) - t = self.entity_representations(indices=hrt_batch[:, 2]) - r = self.relation_embeddings(indices=hrt_batch[:, 1]) + def forward( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: # noqa: D102 + h = self.entity_representations(indices=h_indices) + r = self.relation_embeddings(indices=r_indices) + t = self.entity_representations(indices=t_indices) return self.decoder(h, r, t).unsqueeze(dim=-1) From 0ada54024914b8e258c968d1fe364df479c31907 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sat, 14 Nov 2020 11:59:15 +0100 Subject: [PATCH 214/690] Reorganize embedding specification code and type hints --- src/pykeen/models/base.py | 24 ++++++++++++------------ src/pykeen/nn/emb.py | 33 +++++++++++++++++++-------------- 2 files changed, 31 insertions(+), 26 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index d0195f0631..a60ef51354 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1156,8 +1156,8 @@ def __init__( preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, - embedding_specification: EmbeddingSpecification = None, - relation_embedding_specification: EmbeddingSpecification = None, + embedding_specification: Optional[EmbeddingSpecification] = None, + relation_embedding_specification: Optional[EmbeddingSpecification] = None, ) -> None: """Initialize the entity embedding model. @@ -1226,8 +1226,8 @@ def __init__( preferred_device: Optional[str] = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, - embedding_specification: EmbeddingSpecification = None, - relation_embedding_specification: EmbeddingSpecification = None, + embedding_specification: Optional[EmbeddingSpecification] = None, + relation_embedding_specification: Optional[EmbeddingSpecification] = None, ) -> None: """Initialize embedding model. @@ -1301,10 +1301,10 @@ def __init__( preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, - embedding_specification: EmbeddingSpecification = None, - relation_embedding_specification: EmbeddingSpecification = None, - second_embedding_specification: EmbeddingSpecification = None, - second_relation_embedding_specification: EmbeddingSpecification = None, + embedding_specification: Optional[EmbeddingSpecification] = None, + relation_embedding_specification: Optional[EmbeddingSpecification] = None, + second_embedding_specification: Optional[EmbeddingSpecification] = None, + second_relation_embedding_specification: Optional[EmbeddingSpecification] = None, ) -> None: super().__init__( triples_factory=triples_factory, @@ -1373,10 +1373,10 @@ def __init__( preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, - embedding_specification: EmbeddingSpecification = None, - relation_embedding_specification: EmbeddingSpecification = None, - second_embedding_specification: EmbeddingSpecification = None, - second_relation_embedding_specification: EmbeddingSpecification = None, + embedding_specification: Optional[EmbeddingSpecification] = None, + relation_embedding_specification: Optional[EmbeddingSpecification] = None, + second_embedding_specification: Optional[EmbeddingSpecification] = None, + second_relation_embedding_specification: Optional[EmbeddingSpecification] = None, ): super().__init__( triples_factory=triples_factory, diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index 806889a021..46a981b162 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -60,6 +60,20 @@ class EmbeddingSpecification: # regularizer: Optional[Regularizer] = None + def make(self, num_embeddings: int, embedding_dim: int, device: DeviceHint) -> 'Embedding': + """Create an embedding with this specification.""" + return Embedding.init_with_device( + num_embeddings=num_embeddings, + embedding_dim=embedding_dim, + initializer=self.initializer, + initializer_kwargs=self.initializer_kwargs, + normalizer=self.normalizer, + normalizer_kwargs=self.normalizer_kwargs, + constrainer=self.constrainer, + constrainer_kwargs=self.constrainer_kwargs, + device=device, + ) + class Embedding(RepresentationModule): """Trainable embeddings. @@ -129,10 +143,9 @@ def from_specification( num_embeddings: int, embedding_dim: int, specification: Optional[EmbeddingSpecification], - device: Optional[torch.device] = None, - ) -> "Embedding": - """ - Create an embedding based on an specification. + device: DeviceHint = None, + ) -> 'Embedding': + """Create an embedding based on a specification. :param num_embeddings: The number of embeddings. @@ -148,19 +161,11 @@ def from_specification( """ if specification is None: specification = EmbeddingSpecification() - embedding = Embedding( + return specification.make( num_embeddings=num_embeddings, embedding_dim=embedding_dim, - initializer=specification.initializer, - initializer_kwargs=specification.initializer_kwargs, - normalizer=specification.normalizer, - normalizer_kwargs=specification.normalizer_kwargs, - constrainer=specification.constrainer, - constrainer_kwargs=specification.constrainer_kwargs, + device=device, ) - if device is not None: - embedding = embedding.to(device=device) - return embedding @classmethod def init_with_device( From 6fd0f5aca2b6871277320f50936826d1c0850d1e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 12:56:09 +0100 Subject: [PATCH 215/690] Fix typo --- src/pykeen/nn/modules.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 0646f84763..640ed68fc7 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -176,7 +176,7 @@ def score_r( The head representations. :param all_relations: shape: (batch_size, d_r) The relation representations. - :param t: shape: (num_entities, d_e) + :param t: shape: (batch_size, d_e) The tail representations. :param kwargs: Additional key-word based arguments. From 8ef901f3cae8febf2ed16bc78d322dbd9e343af1 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 12:57:04 +0100 Subject: [PATCH 216/690] Fix check_shapes --- src/pykeen/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 4d3d37f0b5..72bb7ba30c 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -468,6 +468,7 @@ def check_shapes( exp_dim = dims.get(name) if exp_dim is not None and exp_dim != dim: errors.append(f"{name}: {dim} vs. {exp_dim}") + dims[name] = dim if raise_or_error and len(errors) > 0: raise ValueError("Shape verification failed:\n" + '\n'.join(errors)) return len(errors) == 0 From 35dfbbbe58ddec96669b8b2a7e62824b13359761 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 12:58:27 +0100 Subject: [PATCH 217/690] Fix typo - part 2 --- src/pykeen/nn/modules.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 640ed68fc7..44494420d9 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -174,7 +174,7 @@ def score_r( :param h: shape: (batch_size, d_e) The head representations. - :param all_relations: shape: (batch_size, d_r) + :param all_relations: shape: (num_relations, d_r) The relation representations. :param t: shape: (batch_size, d_e) The tail representations. From f97ace9f7d1f43729198c6253b726834b9e98d07 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 13:09:01 +0100 Subject: [PATCH 218/690] Cleanup functional --- src/pykeen/nn/functional.py | 1051 ++++++++++++++++++----------------- src/pykeen/nn/modules.py | 2 +- 2 files changed, 533 insertions(+), 520 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 709cb19ec3..b0b4f1c04f 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -19,9 +19,19 @@ "ermlp_interaction", "ermlpe_interaction", "hole_interaction", + "kg2e_interaction", + "ntn_interaction", + "proje_interaction", + "rescal_interaction", "rotate_interaction", - "translational_interaction", + "simple_interaction", + "structured_embedding_interaction", + "transd_interaction", + "transe_interaction", + "transh_interaction", "transr_interaction", + "tucker_interaction", + "unstructured_model_interaction", ] @@ -32,6 +42,221 @@ def _extract_sizes(h, r, t) -> Tuple[int, int, int, int, int]: return num_heads, num_relations, num_tails, d_e, d_r +def _negative_norm_of_sum( + *x: torch.FloatTensor, + p: Union[int, str] = 2, + power_norm: bool = False, +) -> torch.FloatTensor: + """ + Evaluate negative norm of a sum of vectors on already broadcasted representations. + + :param x: shape: (batch_size, num_heads, num_relations, num_tails, dim) + The representations. + :param p: + The p for the norm. cf. torch.norm. + :param power_norm: + Whether to return $|x-y|_p^p$, cf. https://github.com/pytorch/pytorch/issues/28119 + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + d = sum(x) + if power_norm: + assert isinstance(p, SupportsFloat) + return -(d.abs() ** p).sum(dim=-1) + else: + if torch.is_complex(d): + # workaround for complex numbers: manually compute norm + return -(d.abs() ** p).sum(dim=-1) ** (1 / p) + else: + return -d.norm(p=p, dim=-1) + + +class GaussianDistribution(NamedTuple): + """A gaussian distribution with diagonal covariance matrix.""" + + mean: torch.FloatTensor + diagonal_covariance: torch.FloatTensor + + +def _apply_optional_bn_to_tensor( + batch_norm: Optional[nn.BatchNorm1d], + output_dropout: nn.Dropout, + tensor: torch.FloatTensor, +) -> torch.FloatTensor: + if batch_norm is not None: + shape = tensor.shape + tensor = tensor.reshape(-1, shape[-1]) + tensor = batch_norm(tensor) + tensor = tensor.view(*shape) + tensor = output_dropout(tensor) + return tensor + + +def _project_entity( + e: torch.FloatTensor, + e_p: torch.FloatTensor, + r_p: torch.FloatTensor, +) -> torch.FloatTensor: + r"""Project entity relation-specific. + + .. math:: + + e_{\bot} = M_{re} e + = (r_p e_p^T + I^{d_r \times d_e}) e + = r_p e_p^T e + I^{d_r \times d_e} e + = r_p (e_p^T e) + e' + + and additionally enforces + + .. math:: + + \|e_{\bot}\|_2 \leq 1 + + :param e: shape: (..., d_e) + The entity embedding. + :param e_p: shape: (..., d_e) + The entity projection. + :param r_p: shape: (..., d_r) + The relation projection. + + :return: shape: (..., d_r) + + """ + # The dimensions affected by e' + change_dim = min(e.shape[-1], r_p.shape[-1]) + + # Project entities + # r_p (e_p.T e) + e' + e_bot = r_p * torch.sum(e_p * e, dim=-1, keepdim=True) + e_bot[..., :change_dim] += e[..., :change_dim] + + # Enforce constraints + e_bot = clamp_norm(e_bot, p=2, dim=-1, maxnorm=1) + + return e_bot + + +def _expected_likelihood( + e: GaussianDistribution, + r: GaussianDistribution, + epsilon: float = 1.0e-10, + exact: bool = True, +) -> torch.FloatTensor: + r"""Compute the similarity based on expected likelihood. + + .. math:: + + D((\mu_e, \Sigma_e), (\mu_r, \Sigma_r))) + = \frac{1}{2} \left( + (\mu_e - \mu_r)^T(\Sigma_e + \Sigma_r)^{-1}(\mu_e - \mu_r) + + \log \det (\Sigma_e + \Sigma_r) + d \log (2 \pi) + \right) + = \frac{1}{2} \left( + \mu^T\Sigma^{-1}\mu + + \log \det \Sigma + d \log (2 \pi) + \right) + + :param e: shape: (batch_size, num_heads, num_tails, d) + The entity Gaussian distribution. + :param r: shape: (batch_size, num_relations, d) + The relation Gaussian distribution. + :param epsilon: float (default=1.0) + Small constant used to avoid numerical issues when dividing. + :param exact: + Whether to return the exact similarity, or leave out constant offsets. + + :return: torch.Tensor, shape: (s_1, ..., s_k) + The similarity. + """ + # subtract, shape: (batch_size, num_heads, num_relations, num_tails, dim) + r_shape = r.mean.shape + r_shape = (r_shape[0], 1, r_shape[1], 1, r_shape[2]) + var = r.diagonal_covariance.view(*r_shape) + e.diagonal_covariance.unsqueeze(dim=2) + mean = e.mean.unsqueeze(dim=2) - r.mean.view(*r_shape) + + #: a = \mu^T\Sigma^{-1}\mu + safe_sigma = torch.clamp_min(var, min=epsilon) + sigma_inv = torch.reciprocal(safe_sigma) + sim = torch.sum(sigma_inv * mean ** 2, dim=-1) + + #: b = \log \det \Sigma + sim = sim + safe_sigma.log().sum(dim=-1) + if exact: + sim = sim + sim.shape[-1] * math.log(2. * math.pi) + return sim + + +def _kullback_leibler_similarity( + e: GaussianDistribution, + r: GaussianDistribution, + epsilon: float = 1.0e-10, + exact: bool = True, +) -> torch.FloatTensor: + r"""Compute the similarity based on KL divergence. + + This is done between two Gaussian distributions given by mean mu_* and diagonal covariance matrix sigma_*. + + .. math:: + + D((\mu_e, \Sigma_e), (\mu_r, \Sigma_r))) + = \frac{1}{2} \left( + tr(\Sigma_r^{-1}\Sigma_e) + + (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) + - \log \frac{det(\Sigma_e)}{det(\Sigma_r)} - k_e + \right) + + Note: The sign of the function has been flipped as opposed to the description in the paper, as the + Kullback Leibler divergence is large if the distributions are dissimilar. + + :param e: shape: (batch_size, num_heads, num_tails, d) + The entity Gaussian distributions, as mean/diagonal covariance pairs. + :param r: shape: (batch_size, num_relations, d) + The relation Gaussian distributions, as mean/diagonal covariance pairs. + :param epsilon: float (default=1.0) + Small constant used to avoid numerical issues when dividing. + :param exact: + Whether to return the exact similarity, or leave out constant offsets. + + :return: torch.Tensor, shape: (s_1, ..., s_k) + The similarity. + """ + # invert covariance, shape: (batch_size, num_relations, d) + safe_sigma_r = torch.clamp_min(r.diagonal_covariance, min=epsilon) + sigma_r_inv = torch.reciprocal(safe_sigma_r) + + #: a = tr(\Sigma_r^{-1}\Sigma_e), (batch_size, num_heads, num_relations, num_tails) + # [(b, h, t, d), (b, r, d) -> (b, 1, r, d) -> (b, 1, d, r)] -> (b, h, t, r) -> (b, h, r, t) + sim = (e.diagonal_covariance @ sigma_r_inv.unsqueeze(dim=1).transpose(-2, -1)).transpose(-2, -1) + + #: b = (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) + r_shape = r.mean.shape + # mu.shape: (b, h, r, t, d) + mu = r.mean.view(r_shape[0], 1, r_shape[1], 1, r_shape[2]) - e.mean.unsqueeze(dim=2) + sim = sim + (mu ** 2 @ sigma_r_inv.view(r_shape[0], 1, r_shape[1], r_shape[2], 1)).squeeze(dim=-1) + + #: c = \log \frac{det(\Sigma_e)}{det(\Sigma_r)} + # = sum log (sigma_e)_i - sum log (sigma_r)_i + # ce.shape: (b, h, t) + ce = e.diagonal_covariance.clamp_min(min=epsilon).log().sum(dim=-1) + # cr.shape: (b, r) + cr = safe_sigma_r.log().sum(dim=-1) + sim = sim + ce.unsqueeze(dim=2) - cr.view(r_shape[0], 1, r_shape[1], 1) + + if exact: + sim = sim - e.mean.shape[-1] + sim = 0.5 * sim + + return sim + + +_KG2E_SIMILARITIES = dict( + KL=_kullback_leibler_similarity, + EL=_expected_likelihood, +) +KG2E_SIMILARITIES = set(_KG2E_SIMILARITIES.keys()) + + def _extended_einsum( eq: str, *tensors, @@ -62,6 +287,40 @@ def _extended_einsum( return mod_r +def _view_complex( + x: torch.FloatTensor, +) -> torch.Tensor: + real, imag = split_complex(x=x) + return torch.complex(real=real, imag=imag) + + +def _translational_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + p: Union[int, str] = 2, + power_norm: bool = False, +) -> torch.FloatTensor: + """ + Evaluate a translational distance interaction function on already broadcasted representations. + + :param h: shape: (batch_size, num_heads, num_relations, num_tails, dim) + The head representations. + :param r: shape: (batch_size, num_heads, num_relations, num_tails, dim) + The relation representations. + :param t: shape: (batch_size, num_heads, num_relations, num_tails, dim) + The tail representations. + :param p: + The p for the norm. cf. torch.norm. + :param power_norm: + Whether to return $|x-y|_p^p$, cf. https://github.com/pytorch/pytorch/issues/28119 + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + return _negative_norm_of_sum(h, r, -t, p=p, power_norm=power_norm) + + def _add_cuda_warning(func): def wrapped(*args, **kwargs): try: @@ -79,6 +338,36 @@ def wrapped(*args, **kwargs): return wrapped +def complex_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """ + Evaluate the ComplEx interaction function. + + :param h: shape: (batch_size, num_heads, 2*dim) + The complex head representations. + :param r: shape: (batch_size, num_relations, 2*dim) + The complex relation representations. + :param t: shape: (batch_size, num_tails, 2*dim) + The complex tail representations. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] + return sum( + _extended_einsum("bhd,brd,btd->bhrt", hh, rr, tt) + for hh, rr, tt in [ + (h_re, r_re, t_re), + (h_re, r_im, t_im), + (h_im, r_re, t_im), + (h_im, r_im, t_re), + ] + ) + + @_add_cuda_warning def conve_interaction( h: torch.FloatTensor, @@ -194,57 +483,6 @@ def conve_interaction( return x -def distmult_interaction( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, -) -> torch.FloatTensor: - """ - Evaluate the DistMult interaction function. - - :param h: shape: (batch_size, num_heads, dim) - The head representations. - :param r: shape: (batch_size, num_relations, dim) - The relation representations. - :param t: shape: (batch_size, num_tails, dim) - The tail representations. - - :return: shape: (batch_size, num_heads, num_relations, num_tails) - The scores. - """ - return _extended_einsum("bhd,brd,btd->bhrt", h, r, t) - - -def complex_interaction( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, -) -> torch.FloatTensor: - """ - Evaluate the ComplEx interaction function. - - :param h: shape: (batch_size, num_heads, 2*dim) - The complex head representations. - :param r: shape: (batch_size, num_relations, 2*dim) - The complex relation representations. - :param t: shape: (batch_size, num_tails, 2*dim) - The complex tail representations. - - :return: shape: (batch_size, num_heads, num_relations, num_tails) - The scores. - """ - (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] - return sum( - _extended_einsum("bhd,brd,btd->bhrt", hh, rr, tt) - for hh, rr, tt in [ - (h_re, r_re, t_re), - (h_re, r_im, t_im), - (h_im, r_re, t_im), - (h_im, r_im, t_re), - ] - ) - - def convkb_interaction( h: torch.FloatTensor, r: torch.FloatTensor, @@ -290,235 +528,50 @@ def convkb_interaction( conv_bias = conv.bias.view(1, 1, 1, 1, 1, num_filters) # h.shape: (b, nh, d), conv_head.shape: (o), out.shape: (b, nh, d, o) h = (h.view(h.shape[0], h.shape[1], 1, 1, embedding_dim, 1) * conv_head.view(1, 1, 1, 1, 1, num_filters)) - r = (r.view(r.shape[0], 1, r.shape[1], 1, embedding_dim, 1) * conv_rel.view(1, 1, 1, 1, 1, num_filters)) - t = (t.view(t.shape[0], 1, 1, t.shape[1], embedding_dim, 1) * conv_tail.view(1, 1, 1, 1, 1, num_filters)) - x = activation(conv_bias + h + r + t) - - # Apply dropout, cf. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L54-L56 - x = hidden_dropout(x) - - # Linear layer for final scores - return linear( - x.view(-1, embedding_dim * num_filters), - ).view(-1, num_heads, num_relations, num_tails) - - -def ermlp_interaction( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - hidden: nn.Linear, - activation: nn.Module, - final: nn.Linear, -) -> torch.FloatTensor: - r""" - Evaluate the ER-MLP interaction function. - - :param h: shape: (batch_size, num_heads, dim) - The head representations. - :param r: shape: (batch_size, num_relations, dim) - The relation representations. - :param t: shape: (batch_size, num_tails, dim) - The tail representations. - :param hidden: - The first linear layer. - :param activation: - The activation function of the hidden layer. - :param final: - The second linear layer. - - :return: shape: (batch_size, num_heads, num_relations, num_tails) - The scores. - """ - num_heads, num_relations, num_tails, embedding_dim, _ = _extract_sizes(h, r, t) - hidden_dim = hidden.weight.shape[0] - # split, shape: (embedding_dim, hidden_dim) - head_to_hidden, rel_to_hidden, tail_to_hidden = hidden.weight.t().split(embedding_dim) - bias = hidden.bias.view(1, 1, 1, 1, -1) - h = h.view(-1, num_heads, 1, 1, embedding_dim) @ head_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) - r = r.view(-1, 1, num_relations, 1, embedding_dim) @ rel_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) - t = t.view(-1, 1, 1, num_tails, embedding_dim) @ tail_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) - # TODO: Choosing which to combine first, h/r, h/t or r/t, depending on the shape might further improve - # performance in a 1:n scenario. - return final(activation(bias + h + r + t)).squeeze(dim=-1) - - -def ermlpe_interaction( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - mlp: nn.Module, -) -> torch.FloatTensor: - r""" - Evaluate the ER-MLPE interaction function. - - :param h: shape: (batch_size, num_heads, dim) - The head representations. - :param r: shape: (batch_size, num_relations, dim) - The relation representations. - :param t: shape: (batch_size, num_tails, dim) - The tail representations. - :param mlp: - The MLP. - - :return: shape: (batch_size, num_heads, num_relations, num_tails) - The scores. - """ - # repeat if necessary, and concat head and relation, (batch_size, num_heads, num_relations, 2 * embedding_dim) - x = broadcast_cat(h.unsqueeze(dim=2), r.unsqueeze(dim=1), dim=-1) - - # Predict t embedding, shape: (batch_size, num_heads, num_relations, embedding_dim) - shape = x.shape - x = mlp(x.view(-1, shape[-1])).view(*shape[:-1], -1) - - return (x.unsqueeze(dim=-2) @ t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]).transpose(-2, -1)).squeeze(dim=-1) - - -def hole_interaction( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, -) -> torch.FloatTensor: # noqa: D102 - """ - Evaluate the HolE interaction function. - - :param h: shape: (batch_size, num_heads, dim) - The head representations. - :param r: shape: (batch_size, num_relations, dim) - The relation representations. - :param t: shape: (batch_size, num_tails, dim) - The tail representations. - - :return: shape: (batch_size, num_heads, num_relations, num_tails) - The scores. - """ - # Circular correlation of entity embeddings - a_fft = torch.fft.rfft(h, dim=-1) - b_fft = torch.fft.rfft(t, dim=-1) - - # complex conjugate, shape = (b, h, d) - a_fft = torch.conj(a_fft) - - # Hadamard product in frequency domain, shape: (b, h, t, d) - p_fft = a_fft.unsqueeze(dim=2) * b_fft.unsqueeze(dim=1) - - # inverse real FFT, shape: (b, h, t, d) - composite = torch.fft.irfft(p_fft, n=h.shape[-1], dim=-1) - - # inner product with relation embedding - return _extended_einsum("bhtd,brd->bhrt", composite, r) - - -def _view_complex( - x: torch.FloatTensor, -) -> torch.Tensor: - real, imag = split_complex(x=x) - return torch.complex(real=real, imag=imag) - - -def rotate_interaction( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, -) -> torch.FloatTensor: - """Evaluate the interaction function of RotatE for given embeddings. - - :param h: shape: (batch_size, num_heads, 2*dim) - The head representations. - :param r: shape: (batch_size, num_relations, 2*dim) - The relation representations. - :param t: shape: (batch_size, num_tails, 2*dim) - The tail representations. - - :return: shape: (batch_size, num_heads, num_relations, num_tails) - The scores. - """ - # # r expresses a rotation in complex plane. - # # The inverse rotation is expressed by the complex conjugate of r. - # # The score is computed as the distance of the relation-rotated head to the tail. - # # Equivalently, we can rotate the tail by the inverse relation, and measure the distance to the head, i.e. - # # |h * r - t| = |h - conj(r) * t| - # r_inv = torch.stack([r[:, :, :, 0], -r[:, :, :, 1]], dim=-1) - h, r, t = [_view_complex(x) for x in (h, r, t)] - - # Rotate (=Hadamard product in complex space). - hr = _extended_einsum("bhd,brd->bhrd", h, r) - - # Workaround until https://github.com/pytorch/pytorch/issues/30704 is fixed - return negative_norm_of_sum( - hr.unsqueeze(dim=3), - t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]), - p=2, - power_norm=False, - ) - - -def negative_norm_of_sum( - *x: torch.FloatTensor, - p: Union[int, str] = 2, - power_norm: bool = False, -) -> torch.FloatTensor: - """ - Evaluate negative norm of a sum of vectors on already broadcasted representations. - - :param x: shape: (batch_size, num_heads, num_relations, num_tails, dim) - The representations. - :param p: - The p for the norm. cf. torch.norm. - :param power_norm: - Whether to return $|x-y|_p^p$, cf. https://github.com/pytorch/pytorch/issues/28119 + r = (r.view(r.shape[0], 1, r.shape[1], 1, embedding_dim, 1) * conv_rel.view(1, 1, 1, 1, 1, num_filters)) + t = (t.view(t.shape[0], 1, 1, t.shape[1], embedding_dim, 1) * conv_tail.view(1, 1, 1, 1, 1, num_filters)) + x = activation(conv_bias + h + r + t) - :return: shape: (batch_size, num_heads, num_relations, num_tails) - The scores. - """ - d = sum(x) - if power_norm: - assert isinstance(p, SupportsFloat) - return -(d.abs() ** p).sum(dim=-1) - else: - if torch.is_complex(d): - # workaround for complex numbers: manually compute norm - return -(d.abs() ** p).sum(dim=-1) ** (1 / p) - else: - return -d.norm(p=p, dim=-1) + # Apply dropout, cf. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L54-L56 + x = hidden_dropout(x) + + # Linear layer for final scores + return linear( + x.view(-1, embedding_dim * num_filters), + ).view(-1, num_heads, num_relations, num_tails) -def _translational_interaction( +def distmult_interaction( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, - p: Union[int, str] = 2, - power_norm: bool = False, ) -> torch.FloatTensor: """ - Evaluate a translational distance interaction function on already broadcasted representations. + Evaluate the DistMult interaction function. - :param h: shape: (batch_size, num_heads, num_relations, num_tails, dim) + :param h: shape: (batch_size, num_heads, dim) The head representations. - :param r: shape: (batch_size, num_heads, num_relations, num_tails, dim) + :param r: shape: (batch_size, num_relations, dim) The relation representations. - :param t: shape: (batch_size, num_heads, num_relations, num_tails, dim) + :param t: shape: (batch_size, num_tails, dim) The tail representations. - :param p: - The p for the norm. cf. torch.norm. - :param power_norm: - Whether to return $|x-y|_p^p$, cf. https://github.com/pytorch/pytorch/issues/28119 :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return negative_norm_of_sum(h, r, -t, p=p, power_norm=power_norm) + return _extended_einsum("bhd,brd,btd->bhrt", h, r, t) -def translational_interaction( +def ermlp_interaction( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, - p: Union[int, str] = 2, - power_norm: bool = False, + hidden: nn.Linear, + activation: nn.Module, + final: nn.Linear, ) -> torch.FloatTensor: r""" - Evaluate a translational distance interaction function. + Evaluate the ER-MLP interaction function. :param h: shape: (batch_size, num_heads, dim) The head representations. @@ -526,187 +579,93 @@ def translational_interaction( The relation representations. :param t: shape: (batch_size, num_tails, dim) The tail representations. - :param p: - The p for the norm. cf. torch.norm. - :param power_norm: - Whether to return $|x-y|_p^p$, cf. https://github.com/pytorch/pytorch/issues/28119 + :param hidden: + The first linear layer. + :param activation: + The activation function of the hidden layer. + :param final: + The second linear layer. :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - num_heads, num_relations, num_tails, d_e, _ = _extract_sizes(h, r, t) - h = h.view(-1, num_heads, 1, 1, d_e) - r = r.view(-1, 1, num_relations, 1, d_e) - t = t.view(-1, 1, 1, num_tails, d_e) - return _translational_interaction(h=h, r=r, t=t, p=p, power_norm=power_norm) + num_heads, num_relations, num_tails, embedding_dim, _ = _extract_sizes(h, r, t) + hidden_dim = hidden.weight.shape[0] + # split, shape: (embedding_dim, hidden_dim) + head_to_hidden, rel_to_hidden, tail_to_hidden = hidden.weight.t().split(embedding_dim) + bias = hidden.bias.view(1, 1, 1, 1, -1) + h = h.view(-1, num_heads, 1, 1, embedding_dim) @ head_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) + r = r.view(-1, 1, num_relations, 1, embedding_dim) @ rel_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) + t = t.view(-1, 1, 1, num_tails, embedding_dim) @ tail_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) + # TODO: Choosing which to combine first, h/r, h/t or r/t, depending on the shape might further improve + # performance in a 1:n scenario. + return final(activation(bias + h + r + t)).squeeze(dim=-1) -def transr_interaction( +def ermlpe_interaction( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, - m_r: torch.FloatTensor, - p: int, - power_norm: bool = True, + mlp: nn.Module, ) -> torch.FloatTensor: - """Evaluate the interaction function for given embeddings. + r""" + Evaluate the ER-MLPE interaction function. - :param h: shape: (batch_size, num_heads, d_e) - Head embeddings. - :param r: shape: (batch_size, num_relations, d_r) - Relation embeddings. - :param m_r: shape: (batch_size, num_relations, d_e, d_r) - The relation specific linear transformations. - :param t: shape: (batch_size, num_tails, d_e) - Tail embeddings. - :param p: - The parameter p for selecting the norm. - :param power_norm: - Whether to return the powered norm instead. + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r: shape: (batch_size, num_relations, dim) + The relation representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + :param mlp: + The MLP. :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - num_heads, num_relations, num_tails, d_e, d_r = _extract_sizes(h=h, r=r, t=t) - # project to relation specific subspace and ensure constraints - # head, shape: (b, h, r, 1, d_r) - h_bot = h.view(-1, num_heads, 1, 1, d_e) @ m_r.view(-1, 1, num_relations, d_e, d_r) - h_bot = clamp_norm(h_bot, p=2, dim=-1, maxnorm=1.) - - # head, shape: (b, 1, r, t, d_r) - t_bot = t.view(-1, 1, 1, num_tails, d_e) @ m_r.view(-1, 1, num_relations, d_e, d_r) - t_bot = clamp_norm(t_bot, p=2, dim=-1, maxnorm=1.) - - # evaluate score function, shape: (b, h, r, t) - r = r.view(-1, 1, num_relations, 1, d_r) - return _translational_interaction(h=h_bot, r=r, t=t_bot, p=p, power_norm=power_norm) - - -class GaussianDistribution(NamedTuple): - """A gaussian distribution with diagonal covariance matrix.""" - - mean: torch.FloatTensor - diagonal_covariance: torch.FloatTensor - - -def _expected_likelihood( - e: GaussianDistribution, - r: GaussianDistribution, - epsilon: float = 1.0e-10, - exact: bool = True, -) -> torch.FloatTensor: - r"""Compute the similarity based on expected likelihood. + # repeat if necessary, and concat head and relation, (batch_size, num_heads, num_relations, 2 * embedding_dim) + x = broadcast_cat(h.unsqueeze(dim=2), r.unsqueeze(dim=1), dim=-1) - .. math:: + # Predict t embedding, shape: (batch_size, num_heads, num_relations, embedding_dim) + shape = x.shape + x = mlp(x.view(-1, shape[-1])).view(*shape[:-1], -1) - D((\mu_e, \Sigma_e), (\mu_r, \Sigma_r))) - = \frac{1}{2} \left( - (\mu_e - \mu_r)^T(\Sigma_e + \Sigma_r)^{-1}(\mu_e - \mu_r) - + \log \det (\Sigma_e + \Sigma_r) + d \log (2 \pi) - \right) - = \frac{1}{2} \left( - \mu^T\Sigma^{-1}\mu - + \log \det \Sigma + d \log (2 \pi) - \right) + return (x.unsqueeze(dim=-2) @ t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]).transpose(-2, -1)).squeeze(dim=-1) - :param e: shape: (batch_size, num_heads, num_tails, d) - The entity Gaussian distribution. - :param r: shape: (batch_size, num_relations, d) - The relation Gaussian distribution. - :param epsilon: float (default=1.0) - Small constant used to avoid numerical issues when dividing. - :param exact: - Whether to return the exact similarity, or leave out constant offsets. - :return: torch.Tensor, shape: (s_1, ..., s_k) - The similarity. +def hole_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: # noqa: D102 """ - # subtract, shape: (batch_size, num_heads, num_relations, num_tails, dim) - r_shape = r.mean.shape - r_shape = (r_shape[0], 1, r_shape[1], 1, r_shape[2]) - var = r.diagonal_covariance.view(*r_shape) + e.diagonal_covariance.unsqueeze(dim=2) - mean = e.mean.unsqueeze(dim=2) - r.mean.view(*r_shape) - - #: a = \mu^T\Sigma^{-1}\mu - safe_sigma = torch.clamp_min(var, min=epsilon) - sigma_inv = torch.reciprocal(safe_sigma) - sim = torch.sum(sigma_inv * mean ** 2, dim=-1) - - #: b = \log \det \Sigma - sim = sim + safe_sigma.log().sum(dim=-1) - if exact: - sim = sim + sim.shape[-1] * math.log(2. * math.pi) - return sim - - -def _kullback_leibler_similarity( - e: GaussianDistribution, - r: GaussianDistribution, - epsilon: float = 1.0e-10, - exact: bool = True, -) -> torch.FloatTensor: - r"""Compute the similarity based on KL divergence. - - This is done between two Gaussian distributions given by mean mu_* and diagonal covariance matrix sigma_*. - - .. math:: - - D((\mu_e, \Sigma_e), (\mu_r, \Sigma_r))) - = \frac{1}{2} \left( - tr(\Sigma_r^{-1}\Sigma_e) - + (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) - - \log \frac{det(\Sigma_e)}{det(\Sigma_r)} - k_e - \right) - - Note: The sign of the function has been flipped as opposed to the description in the paper, as the - Kullback Leibler divergence is large if the distributions are dissimilar. + Evaluate the HolE interaction function. - :param e: shape: (batch_size, num_heads, num_tails, d) - The entity Gaussian distributions, as mean/diagonal covariance pairs. - :param r: shape: (batch_size, num_relations, d) - The relation Gaussian distributions, as mean/diagonal covariance pairs. - :param epsilon: float (default=1.0) - Small constant used to avoid numerical issues when dividing. - :param exact: - Whether to return the exact similarity, or leave out constant offsets. + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r: shape: (batch_size, num_relations, dim) + The relation representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. - :return: torch.Tensor, shape: (s_1, ..., s_k) - The similarity. + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. """ - # invert covariance, shape: (batch_size, num_relations, d) - safe_sigma_r = torch.clamp_min(r.diagonal_covariance, min=epsilon) - sigma_r_inv = torch.reciprocal(safe_sigma_r) - - #: a = tr(\Sigma_r^{-1}\Sigma_e), (batch_size, num_heads, num_relations, num_tails) - # [(b, h, t, d), (b, r, d) -> (b, 1, r, d) -> (b, 1, d, r)] -> (b, h, t, r) -> (b, h, r, t) - sim = (e.diagonal_covariance @ sigma_r_inv.unsqueeze(dim=1).transpose(-2, -1)).transpose(-2, -1) - - #: b = (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) - r_shape = r.mean.shape - # mu.shape: (b, h, r, t, d) - mu = r.mean.view(r_shape[0], 1, r_shape[1], 1, r_shape[2]) - e.mean.unsqueeze(dim=2) - sim = sim + (mu ** 2 @ sigma_r_inv.view(r_shape[0], 1, r_shape[1], r_shape[2], 1)).squeeze(dim=-1) - - #: c = \log \frac{det(\Sigma_e)}{det(\Sigma_r)} - # = sum log (sigma_e)_i - sum log (sigma_r)_i - # ce.shape: (b, h, t) - ce = e.diagonal_covariance.clamp_min(min=epsilon).log().sum(dim=-1) - # cr.shape: (b, r) - cr = safe_sigma_r.log().sum(dim=-1) - sim = sim + ce.unsqueeze(dim=2) - cr.view(r_shape[0], 1, r_shape[1], 1) + # Circular correlation of entity embeddings + a_fft = torch.fft.rfft(h, dim=-1) + b_fft = torch.fft.rfft(t, dim=-1) - if exact: - sim = sim - e.mean.shape[-1] - sim = 0.5 * sim + # complex conjugate, shape = (b, h, d) + a_fft = torch.conj(a_fft) - return sim + # Hadamard product in frequency domain, shape: (b, h, t, d) + p_fft = a_fft.unsqueeze(dim=2) * b_fft.unsqueeze(dim=1) + # inverse real FFT, shape: (b, h, t, d) + composite = torch.fft.irfft(p_fft, n=h.shape[-1], dim=-1) -_KG2E_SIMILARITIES = dict( - KL=_kullback_leibler_similarity, - EL=_expected_likelihood, -) -KG2E_SIMILARITIES = set(_KG2E_SIMILARITIES.keys()) + # inner product with relation embedding + return _extended_einsum("bhtd,brd->bhrt", composite, r) def kg2e_interaction( @@ -869,6 +828,43 @@ def rescal_interaction( return _extended_einsum("bhd,brde,bte->bhrt", h, r.view(*r.shape[:-1], h.shape[-1], h.shape[-1]), t) +def rotate_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Evaluate the interaction function of RotatE for given embeddings. + + :param h: shape: (batch_size, num_heads, 2*dim) + The head representations. + :param r: shape: (batch_size, num_relations, 2*dim) + The relation representations. + :param t: shape: (batch_size, num_tails, 2*dim) + The tail representations. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + # # r expresses a rotation in complex plane. + # # The inverse rotation is expressed by the complex conjugate of r. + # # The score is computed as the distance of the relation-rotated head to the tail. + # # Equivalently, we can rotate the tail by the inverse relation, and measure the distance to the head, i.e. + # # |h * r - t| = |h - conj(r) * t| + # r_inv = torch.stack([r[:, :, :, 0], -r[:, :, :, 1]], dim=-1) + h, r, t = [_view_complex(x) for x in (h, r, t)] + + # Rotate (=Hadamard product in complex space). + hr = _extended_einsum("bhd,brd->bhrd", h, r) + + # Workaround until https://github.com/pytorch/pytorch/issues/30704 is fixed + return _negative_norm_of_sum( + hr.unsqueeze(dim=3), + t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]), + p=2, + power_norm=False, + ) + + def simple_interaction( h: torch.FloatTensor, r: torch.FloatTensor, @@ -936,7 +932,7 @@ def structured_embedding_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return negative_norm_of_sum( + return _negative_norm_of_sum( _extended_einsum("brde,bhd->bhre", r_h, h).unsqueeze(dim=3), -_extended_einsum("brde,btd->brte", r_t, t).unsqueeze(dim=1), p=p, @@ -944,6 +940,90 @@ def structured_embedding_interaction( ) +def transd_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + h_p: torch.FloatTensor, + r_p: torch.FloatTensor, + t_p: torch.FloatTensor, + p: int, + power_norm: bool = False, +) -> torch.FloatTensor: + """ + Evaluate the DistMult interaction function. + + :param h: shape: (batch_size, num_heads, d_e) + The head representations. + :param r: shape: (batch_size, num_relations, d_r) + The relation representations. + :param t: shape: (batch_size, num_tails, d_e) + The tail representations. + :param h_p: shape: (batch_size, num_heads, d_e) + The head projections. + :param r_p: shape: (batch_size, num_relations, d_r) + The relation projections. + :param t_p: shape: (batch_size, num_tails, d_e) + The tail projections. + :param p: + The parameter p for selecting the norm. + :param power_norm: + Whether to return the powered norm instead. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + # Project entities + # shape: (b, h, r, 1, d_r) + h_bot = _project_entity( + e=h.unsqueeze(dim=2), + e_p=h_p.unsqueeze(dim=2), + r_p=r_p.unsqueeze(dim=1), + ).unsqueeze(dim=-2) + # shape: (b, 1, r, t, d_r) + t_bot = _project_entity( + e=t.unsqueeze(dim=1), + e_p=t_p.unsqueeze(dim=1), + r_p=r_p.unsqueeze(dim=2), + ).unsqueeze(dim=1) + r = r.view(r.shape[0], 1, r.shape[1], 1, r.shape[2]) + return _translational_interaction(h=h_bot, r=r, t=t_bot, p=p, power_norm=power_norm) + + +def transe_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + p: Union[int, str] = 2, + power_norm: bool = False, +) -> torch.FloatTensor: + """ + Evaluate the TransE interaction function. + + :param h: shape: (batch_size, num_heads, dim) + The head representations. + :param r: shape: (batch_size, num_relations, dim) + The relation representations. + :param t: shape: (batch_size, num_tails, dim) + The tail representations. + :param p: + The p for the norm. + :param power_norm: + Whether to return the powered norm. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + num_heads, num_relations, num_tails, embedding_dim = _extract_sizes(h, r, t)[:4] + return _translational_interaction( + h=h.view(-1, num_heads, 1, 1, embedding_dim), + r=r.view(-1, 1, num_relations, 1, embedding_dim), + t=t.view(-1, 1, 1, num_tails, embedding_dim), + p=p, + power_norm=power_norm, + ) + + def transh_interaction( h: torch.FloatTensor, w_r: torch.FloatTensor, @@ -981,18 +1061,45 @@ def transh_interaction( ) -def _apply_optional_bn_to_tensor( - batch_norm: Optional[nn.BatchNorm1d], - output_dropout: nn.Dropout, - tensor: torch.FloatTensor, +def transr_interaction( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + m_r: torch.FloatTensor, + p: int, + power_norm: bool = True, ) -> torch.FloatTensor: - if batch_norm is not None: - shape = tensor.shape - tensor = tensor.reshape(-1, shape[-1]) - tensor = batch_norm(tensor) - tensor = tensor.view(*shape) - tensor = output_dropout(tensor) - return tensor + """Evaluate the interaction function for given embeddings. + + :param h: shape: (batch_size, num_heads, d_e) + Head embeddings. + :param r: shape: (batch_size, num_relations, d_r) + Relation embeddings. + :param m_r: shape: (batch_size, num_relations, d_e, d_r) + The relation specific linear transformations. + :param t: shape: (batch_size, num_tails, d_e) + Tail embeddings. + :param p: + The parameter p for selecting the norm. + :param power_norm: + Whether to return the powered norm instead. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + num_heads, num_relations, num_tails, d_e, d_r = _extract_sizes(h=h, r=r, t=t) + # project to relation specific subspace and ensure constraints + # head, shape: (b, h, r, 1, d_r) + h_bot = h.view(-1, num_heads, 1, 1, d_e) @ m_r.view(-1, 1, num_relations, d_e, d_r) + h_bot = clamp_norm(h_bot, p=2, dim=-1, maxnorm=1.) + + # head, shape: (b, 1, r, t, d_r) + t_bot = t.view(-1, 1, 1, num_tails, d_e) @ m_r.view(-1, 1, num_relations, d_e, d_r) + t_bot = clamp_norm(t_bot, p=2, dim=-1, maxnorm=1.) + + # evaluate score function, shape: (b, h, r, t) + r = r.view(-1, 1, num_relations, 1, d_r) + return _translational_interaction(h=h_bot, r=r, t=t_bot, p=p, power_norm=power_norm) def tucker_interaction( @@ -1074,98 +1181,4 @@ def unstructured_model_interaction( """ h = h.unsqueeze(dim=2).unsqueeze(dim=3) t = t.unsqueeze(dim=1).unsqueeze(dim=2) - return negative_norm_of_sum(h, -t, p=p, power_norm=power_norm) - - -def _project_entity( - e: torch.FloatTensor, - e_p: torch.FloatTensor, - r_p: torch.FloatTensor, -) -> torch.FloatTensor: - r"""Project entity relation-specific. - - .. math:: - - e_{\bot} = M_{re} e - = (r_p e_p^T + I^{d_r \times d_e}) e - = r_p e_p^T e + I^{d_r \times d_e} e - = r_p (e_p^T e) + e' - - and additionally enforces - - .. math:: - - \|e_{\bot}\|_2 \leq 1 - - :param e: shape: (..., d_e) - The entity embedding. - :param e_p: shape: (..., d_e) - The entity projection. - :param r_p: shape: (..., d_r) - The relation projection. - - :return: shape: (..., d_r) - - """ - # The dimensions affected by e' - change_dim = min(e.shape[-1], r_p.shape[-1]) - - # Project entities - # r_p (e_p.T e) + e' - e_bot = r_p * torch.sum(e_p * e, dim=-1, keepdim=True) - e_bot[..., :change_dim] += e[..., :change_dim] - - # Enforce constraints - e_bot = clamp_norm(e_bot, p=2, dim=-1, maxnorm=1) - - return e_bot - - -def transd_interaction( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - h_p: torch.FloatTensor, - r_p: torch.FloatTensor, - t_p: torch.FloatTensor, - p: int, - power_norm: bool = False, -) -> torch.FloatTensor: - """ - Evaluate the DistMult interaction function. - - :param h: shape: (batch_size, num_heads, d_e) - The head representations. - :param r: shape: (batch_size, num_relations, d_r) - The relation representations. - :param t: shape: (batch_size, num_tails, d_e) - The tail representations. - :param h_p: shape: (batch_size, num_heads, d_e) - The head projections. - :param r_p: shape: (batch_size, num_relations, d_r) - The relation projections. - :param t_p: shape: (batch_size, num_tails, d_e) - The tail projections. - :param p: - The parameter p for selecting the norm. - :param power_norm: - Whether to return the powered norm instead. - - :return: shape: (batch_size, num_heads, num_relations, num_tails) - The scores. - """ - # Project entities - # shape: (b, h, r, 1, d_r) - h_bot = _project_entity( - e=h.unsqueeze(dim=2), - e_p=h_p.unsqueeze(dim=2), - r_p=r_p.unsqueeze(dim=1), - ).unsqueeze(dim=-2) - # shape: (b, 1, r, t, d_r) - t_bot = _project_entity( - e=t.unsqueeze(dim=1), - e_p=t_p.unsqueeze(dim=1), - r_p=r_p.unsqueeze(dim=2), - ).unsqueeze(dim=1) - r = r.view(r.shape[0], 1, r.shape[1], 1, r.shape[2]) - return _translational_interaction(h=h_bot, r=r, t=t_bot, p=p, power_norm=power_norm) + return _negative_norm_of_sum(h, -t, p=p, power_norm=power_norm) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 44494420d9..6ed1cae0ec 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -290,7 +290,7 @@ def forward( **kwargs, ) -> torch.FloatTensor: # noqa:D102 self._check_for_empty_kwargs(kwargs=kwargs) - return pkf._translational_interaction(h=h, r=r, t=t, p=self.p, power_norm=self.power_norm) + return pkf.transe_interaction(h=h, r=r, t=t, p=self.p, power_norm=self.power_norm) #: Interaction function of ComplEx From 1dd265e99d39d6c71be89b6b0abdbaf9b92f8c1a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 13:09:33 +0100 Subject: [PATCH 219/690] Use TransEInteractionFunction --- src/pykeen/models/unimodal/trans_e.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index a31e52964a..499ec89f22 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -10,7 +10,7 @@ from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ -from ...nn.modules import TranslationalInteractionFunction +from ...nn.modules import TransEInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -69,7 +69,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - interaction_function=TranslationalInteractionFunction(p=scoring_fct_norm), + interaction_function=TransEInteractionFunction(p=scoring_fct_norm, power_norm=False), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, From dc0e051905646b1044481a0fe04fd5abd045aba2 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 13:21:37 +0100 Subject: [PATCH 220/690] Remove **kwargs from interaction function --- src/pykeen/nn/modules.py | 70 ++++++++++++++++++++++++---------------- 1 file changed, 42 insertions(+), 28 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 6ed1cae0ec..8c7ca3ed72 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -5,7 +5,7 @@ import logging import math from abc import ABC -from typing import Any, Callable, Mapping, Optional, Sequence, Tuple, Type +from typing import Any, Callable, Mapping, Optional, Sequence, Tuple, Type, Union import torch from torch import nn @@ -27,21 +27,18 @@ class InteractionFunction(nn.Module): def forward( self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, + h: Sequence[torch.FloatTensor] = tuple(), + r: Sequence[torch.FloatTensor] = tuple(), + t: Sequence[torch.FloatTensor] = tuple(), ) -> torch.FloatTensor: - """Score the given triples. + """Compute broadcasted triple scores given representations for head, relation and tails. - :param h: shape: (batch_size, num_heads, d_e) + :param h: shape: (batch_size, num_heads, *) The head representations. - :param r: shape: (batch_size, num_relations, d_r) + :param r: shape: (batch_size, num_relations, *) The relation representations. - :param t: shape: (batch_size, num_tails, d_e) + :param t: shape: (batch_size, num_tails, *) The tail representations. - :param kwargs: - Additional key-word based arguments. :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. @@ -94,10 +91,9 @@ def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: def score_hrt( self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, + h: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, + r: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, + t: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, ) -> torch.FloatTensor: """ Score a batch of triples.. @@ -121,17 +117,16 @@ def score_hrt( h, r, t = self._add_dim(h, r, t, dim=self.NUM_DIM) # get scores - scores = self(h=h, r=r, t=t, **kwargs) + scores = self(h=h, r=r, t=t) # prepare output shape, (batch_size, num_heads, num_relations, num_tails) -> (batch_size, 1) return self._remove_dim(scores, self.HEAD_DIM, self.RELATION_DIM, self.TAIL_DIM).unsqueeze(dim=-1) def score_h( self, - all_entities: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, + all_entities: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, + r: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, + t: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, ) -> torch.FloatTensor: """ Score all head entities. @@ -164,10 +159,9 @@ def score_h( def score_r( self, - h: torch.FloatTensor, - all_relations: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, + h: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, + all_relations: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, + t: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, ) -> torch.FloatTensor: """ Score all relations. @@ -199,10 +193,9 @@ def score_r( def score_t( self, - h: torch.FloatTensor, - r: torch.FloatTensor, - all_entities: torch.FloatTensor, - **kwargs, + h: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, + r: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, + all_entities: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, ) -> torch.FloatTensor: """ Score all tail entities. @@ -791,3 +784,24 @@ def forward( r_p = kwargs.pop("r_p") t_p = kwargs.pop("t_p") return pkf.transd_interaction(h=h, r=r, t=t, h_p=h_p, r_p=r_p, t_p=t_p, p=self.p, power_norm=self.power_norm) + + +class NTNInteractionFunction(InteractionFunction): + """The interaction function of NTN.""" + + def __init__( + self, + non_linearity: Optional[nn.Module] = None, + ): + super().__init__() + if non_linearity is None: + non_linearity = nn.Tanh() + self.non_linearity = non_linearity + + def forward( + self, + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, + **kwargs, + ) -> torch.FloatTensor: From 35f9a202f6d0712a03ff520df8419d1f79bdab52 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 13:27:58 +0100 Subject: [PATCH 221/690] Add proper typing --- src/pykeen/nn/modules.py | 44 ++++++++++++++++++++++------------------ 1 file changed, 24 insertions(+), 20 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 8c7ca3ed72..d663519743 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -5,7 +5,7 @@ import logging import math from abc import ABC -from typing import Any, Callable, Mapping, Optional, Sequence, Tuple, Type, Union +from typing import Any, Callable, Generic, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union import torch from torch import nn @@ -15,8 +15,11 @@ logger = logging.getLogger(__name__) +EntityRepresentation = TypeVar("EntityRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) +RelationRepresentation = TypeVar("RelationRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) -class InteractionFunction(nn.Module): + +class InteractionFunction(nn.Module, Generic[EntityRepresentation, RelationRepresentation]): """Base class for interaction functions.""" BATCH_DIM: int = 0 @@ -27,9 +30,9 @@ class InteractionFunction(nn.Module): def forward( self, - h: Sequence[torch.FloatTensor] = tuple(), - r: Sequence[torch.FloatTensor] = tuple(), - t: Sequence[torch.FloatTensor] = tuple(), + h: EntityRepresentation = tuple(), + r: RelationRepresentation = tuple(), + t: EntityRepresentation = tuple(), ) -> torch.FloatTensor: """Compute broadcasted triple scores given representations for head, relation and tails. @@ -48,6 +51,7 @@ def forward( @classmethod def _check_for_empty_kwargs(cls, kwargs: Mapping[str, Any]) -> None: """Check that kwargs is empty.""" + # TODO: Deprecated if len(kwargs) > 0: raise ValueError(f"{cls.__name__} does not take the following kwargs: {kwargs}") @@ -91,9 +95,9 @@ def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: def score_hrt( self, - h: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, - r: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, - t: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, + h: EntityRepresentation = tuple(), + r: RelationRepresentation = tuple(), + t: EntityRepresentation = tuple(), ) -> torch.FloatTensor: """ Score a batch of triples.. @@ -124,9 +128,9 @@ def score_hrt( def score_h( self, - all_entities: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, - r: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, - t: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, + all_entities: EntityRepresentation = tuple(), + r: RelationRepresentation = tuple(), + t: EntityRepresentation = tuple(), ) -> torch.FloatTensor: """ Score all head entities. @@ -152,16 +156,16 @@ def score_h( h = self._add_dim(all_entities, dim=self.BATCH_DIM) # get scores - scores = self(h=h, r=r, t=t, **kwargs) + scores = self(h=h, r=r, t=t) # prepare output shape, (batch_size, num_heads, num_relations, num_tails) -> (batch_size, num_heads) return self._remove_dim(scores, self.RELATION_DIM, self.TAIL_DIM) def score_r( self, - h: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, - all_relations: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, - t: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, + h: EntityRepresentation = tuple(), + all_relations: RelationRepresentation = tuple(), + t: EntityRepresentation = tuple(), ) -> torch.FloatTensor: """ Score all relations. @@ -186,16 +190,16 @@ def score_r( r = self._add_dim(all_relations, dim=self.BATCH_DIM) # get scores - scores = self(h=h, r=r, t=t, **kwargs) + scores = self(h=h, r=r, t=t) # prepare output shape return self._remove_dim(scores, self.HEAD_DIM, self.TAIL_DIM) def score_t( self, - h: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, - r: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, - all_entities: Union[None, torch.FloatTensor, Sequence[torch.FloatTensor]] = None, + h: EntityRepresentation = tuple(), + r: RelationRepresentation = tuple(), + all_entities: EntityRepresentation = tuple(), ) -> torch.FloatTensor: """ Score all tail entities. @@ -220,7 +224,7 @@ def score_t( t = self._add_dim(all_entities, dim=self.BATCH_DIM) # get scores - scores = self(h=h, r=r, t=t, **kwargs) + scores = self(h=h, r=r, t=t) # prepare output shape return self._remove_dim(scores, self.HEAD_DIM, self.RELATION_DIM) From 28d1d7c05c76be6ed5dad17db3ef28cbaf25f2f8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 13:29:29 +0100 Subject: [PATCH 222/690] Add field --- src/pykeen/nn/modules.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index d663519743..94686adf7e 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -5,7 +5,7 @@ import logging import math from abc import ABC -from typing import Any, Callable, Generic, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union +from typing import Any, Callable, Generic, Mapping, Optional, Sequence, Tuple, Type, TypeVar import torch from torch import nn @@ -22,12 +22,19 @@ class InteractionFunction(nn.Module, Generic[EntityRepresentation, RelationRepresentation]): """Base class for interaction functions.""" + # Dimensions BATCH_DIM: int = 0 NUM_DIM: int = 1 HEAD_DIM: int = 1 RELATION_DIM: int = 2 TAIL_DIM: int = 3 + #: The symbolic shapes for entity representations + entity_shape: Sequence[str] + + #: The symbolic shapes for relation representations + relation_shape: Sequence[str] + def forward( self, h: EntityRepresentation = tuple(), From cc2579b9467a24347121f4d8489f6d1dfee020b5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 13:36:22 +0100 Subject: [PATCH 223/690] improve shape check --- src/pykeen/nn/modules.py | 30 +++++++++++++++++++++--------- 1 file changed, 21 insertions(+), 9 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 94686adf7e..cc5621d527 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -1,7 +1,7 @@ # -*- coding: utf-8 -*- """Stateful interaction functions.""" - +import itertools import logging import math from abc import ABC @@ -100,6 +100,22 @@ def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: x = x.squeeze(dim=dim) return x + def _check_shapes( + self, + h: EntityRepresentation = tuple(), + r: RelationRepresentation = tuple(), + t: EntityRepresentation = tuple(), + h_prefix: str = "b", + r_prefix: str = "b", + t_prefix: str = "b", + raise_on_error: bool = True, + ) -> bool: + return len(h) == len(self.entity_shape) and len(r) == len(self.relation_shape) and len(t) == len(self.entity_shape) and check_shapes(*itertools.chain( + ((hh, h_prefix + hs) for hh, hs in zip(h, self.entity_shape)), + ((rr, r_prefix + rs) for rr, rs in zip(r, self.relation_shape)), + ((tt, t_prefix + ts) for tt, ts in zip(t, self.entity_shape)), + ), raise_or_error=raise_on_error) + def score_hrt( self, h: EntityRepresentation = tuple(), @@ -121,8 +137,7 @@ def score_hrt( :return: shape: (batch_size, 1) The scores. """ - # check shape - assert check_shapes((h, "be"), (r, "br"), (t, "be")) + assert self._check_shapes(h=h, r=r, t=t) # prepare input to generic score function h, r, t = self._add_dim(h, r, t, dim=self.NUM_DIM) @@ -154,8 +169,7 @@ def score_h( :return: shape: (batch_size, num_entities) The scores. """ - # check shape - assert check_shapes((all_entities, "ne"), (r, "br"), (t, "be")) + assert self._check_shapes(h=all_entities, r=r, t=t, h_prefix="n") # TODO: What about unsqueezing for additional e.g. head arguments # prepare input to generic score function @@ -189,8 +203,7 @@ def score_r( :return: shape: (batch_size, num_entities) The scores. """ - # check shape - assert check_shapes((all_relations, "nr"), (h, "be"), (t, "be")) + assert self._check_shapes(h=h, r=all_relations, t=t, r_prefix="n") # prepare input to generic score function h, t = self._add_dim(h, t, dim=self.NUM_DIM) @@ -223,8 +236,7 @@ def score_t( :return: shape: (batch_size, num_entities) The scores. """ - # check shape - assert check_shapes((all_entities, "ne"), (r, "br"), (h, "be")) + assert self._check_shapes(h=h, r=r, t=all_entities, t_prefix="n") # prepare input to generic score function h, r = self._add_dim(h, r, dim=self.NUM_DIM) From 90489755f5f1f24eb2016b5ae7238283dcfd6d9e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 13:38:08 +0100 Subject: [PATCH 224/690] Improve dimension extension --- src/pykeen/nn/modules.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index cc5621d527..f7a577abd2 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -140,7 +140,9 @@ def score_hrt( assert self._check_shapes(h=h, r=r, t=t) # prepare input to generic score function - h, r, t = self._add_dim(h, r, t, dim=self.NUM_DIM) + h = self._add_dim(h, dim=self.NUM_DIM) + r = self._add_dim(r, dim=self.NUM_DIM) + t = self._add_dim(t, dim=self.NUM_DIM) # get scores scores = self(h=h, r=r, t=t) @@ -171,10 +173,10 @@ def score_h( """ assert self._check_shapes(h=all_entities, r=r, t=t, h_prefix="n") - # TODO: What about unsqueezing for additional e.g. head arguments # prepare input to generic score function - r, t = self._add_dim(r, t, dim=self.NUM_DIM) h = self._add_dim(all_entities, dim=self.BATCH_DIM) + r = self._add_dim(r, dim=self.NUM_DIM) + t = self._add_dim(t, dim=self.NUM_DIM) # get scores scores = self(h=h, r=r, t=t) @@ -206,8 +208,9 @@ def score_r( assert self._check_shapes(h=h, r=all_relations, t=t, r_prefix="n") # prepare input to generic score function - h, t = self._add_dim(h, t, dim=self.NUM_DIM) + h = self._add_dim(h, dim=self.NUM_DIM) r = self._add_dim(all_relations, dim=self.BATCH_DIM) + t = self._add_dim(t, dim=self.NUM_DIM) # get scores scores = self(h=h, r=r, t=t) @@ -239,7 +242,8 @@ def score_t( assert self._check_shapes(h=h, r=r, t=all_entities, t_prefix="n") # prepare input to generic score function - h, r = self._add_dim(h, r, dim=self.NUM_DIM) + h = self._add_dim(h, dim=self.NUM_DIM) + r = self._add_dim(r, dim=self.NUM_DIM) t = self._add_dim(all_entities, dim=self.BATCH_DIM) # get scores From 0eb088a24d9b1b7f690f4b31e707ea37a071559f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 13:38:35 +0100 Subject: [PATCH 225/690] remove kwargs from docstring --- src/pykeen/nn/modules.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index f7a577abd2..1e8b606ffa 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -131,8 +131,6 @@ def score_hrt( The relation representations. :param t: shape: (batch_size, d_e) The tail representations. - :param kwargs: - Additional key-word based arguments. :return: shape: (batch_size, 1) The scores. @@ -165,8 +163,6 @@ def score_h( The relation representations. :param t: shape: (batch_size, d_e) The tail representations. - :param kwargs: - Additional key-word based arguments. :return: shape: (batch_size, num_entities) The scores. @@ -199,8 +195,6 @@ def score_r( The relation representations. :param t: shape: (batch_size, d_e) The tail representations. - :param kwargs: - Additional key-word based arguments. :return: shape: (batch_size, num_entities) The scores. @@ -233,8 +227,6 @@ def score_t( The relation representations. :param all_entities: shape: (num_entities, d_e) The tail representations. - :param kwargs: - Additional key-word based arguments. :return: shape: (batch_size, num_entities) The scores. From f4c02ffd704bf1856af0371db050d96179694315 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 13:41:15 +0100 Subject: [PATCH 226/690] Use new interface for UM --- src/pykeen/nn/modules.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 1e8b606ffa..489a102f8d 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -274,7 +274,7 @@ def forward( return StatelessInteractionFunction -class TranslationalInteractionFunction(InteractionFunction, ABC): +class TranslationalInteractionFunction(InteractionFunction[EntityRepresentation, RelationRepresentation], ABC): """The translational interaction function shared by the TransE, TransR, TransH, and other Trans models.""" def __init__(self, p: int, power_norm: bool = False): @@ -772,20 +772,22 @@ def forward( ) -class UnstructuredModelInteractionFunction(TranslationalInteractionFunction): +class UnstructuredModelInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, None]): """Interaction function of UnstructuredModel.""" + # shapes + entity_shape = ("d",) + relation_shape = tuple() + def __init__(self, p: int, power_norm: bool = True): super().__init__(p=p, power_norm=power_norm) def forward( self, h: torch.FloatTensor, - r: torch.FloatTensor, + r: None, t: torch.FloatTensor, - **kwargs, ) -> torch.FloatTensor: # noqa:D102 - self._check_for_empty_kwargs(kwargs=kwargs) return pkf.unstructured_model_interaction(h, t, p=self.p, power_norm=self.power_norm) From 35d4bece35190418fdbef7abc3079af6089b3b18 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 13:44:14 +0100 Subject: [PATCH 227/690] Update SE interaction --- src/pykeen/nn/modules.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 489a102f8d..631862a2b8 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -686,20 +686,19 @@ def forward( RESCALInteractionFunction = _build_module_from_stateless(pkf.rescal_interaction) -class StructuredEmbeddingInteractionFunction(TranslationalInteractionFunction): +class StructuredEmbeddingInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor]]): """Interaction function of Structured Embedding.""" + entity_shape = ("d",) + relation_shape = ("dd",) + def forward( self, h: torch.FloatTensor, - r: torch.FloatTensor, + r: Tuple[torch.FloatTensor, torch.FloatTensor], t: torch.FloatTensor, - **kwargs, ) -> torch.FloatTensor: # noqa:D102 - dim = h.shape[-1] - rh, rt = r.split(dim ** 2, dim=-1) - rh = rh.view(*rh.shape[:-1], dim, dim) - rt = rt.view(*rt.shape[:-1], dim, dim) + rh, rt = r return pkf.structured_embedding_interaction(h=h, r_h=rh, r_t=rt, t=t, p=self.p, power_norm=self.power_norm) From 715f0ac9d4a6d38879c909ea17a41edd934d5e07 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 13:52:41 +0100 Subject: [PATCH 228/690] Allow different shape for tail entities (for ConvE) --- src/pykeen/nn/modules.py | 62 +++++++++++++++++++++++----------------- 1 file changed, 36 insertions(+), 26 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 631862a2b8..0d90f3d2c0 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -5,7 +5,7 @@ import logging import math from abc import ABC -from typing import Any, Callable, Generic, Mapping, Optional, Sequence, Tuple, Type, TypeVar +from typing import Any, Callable, Generic, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union import torch from torch import nn @@ -15,11 +15,12 @@ logger = logging.getLogger(__name__) -EntityRepresentation = TypeVar("EntityRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) +HeadRepresentation = TypeVar("HeadRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) RelationRepresentation = TypeVar("RelationRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) +TailRepresentation = TypeVar("TailRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) -class InteractionFunction(nn.Module, Generic[EntityRepresentation, RelationRepresentation]): +class InteractionFunction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation]): """Base class for interaction functions.""" # Dimensions @@ -30,16 +31,17 @@ class InteractionFunction(nn.Module, Generic[EntityRepresentation, RelationRepre TAIL_DIM: int = 3 #: The symbolic shapes for entity representations - entity_shape: Sequence[str] + entity_shape: Union[str, Sequence[str]] = "d" + tail_entity_shape: Union[None, str, Sequence[str]] = None #: The symbolic shapes for relation representations - relation_shape: Sequence[str] + relation_shape: Union[str, Sequence[str]] = "d" def forward( self, - h: EntityRepresentation = tuple(), + h: HeadRepresentation = tuple(), r: RelationRepresentation = tuple(), - t: EntityRepresentation = tuple(), + t: TailRepresentation = tuple(), ) -> torch.FloatTensor: """Compute broadcasted triple scores given representations for head, relation and tails. @@ -102,25 +104,36 @@ def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: def _check_shapes( self, - h: EntityRepresentation = tuple(), + h: HeadRepresentation = tuple(), r: RelationRepresentation = tuple(), - t: EntityRepresentation = tuple(), + t: TailRepresentation = tuple(), h_prefix: str = "b", r_prefix: str = "b", t_prefix: str = "b", raise_on_error: bool = True, ) -> bool: - return len(h) == len(self.entity_shape) and len(r) == len(self.relation_shape) and len(t) == len(self.entity_shape) and check_shapes(*itertools.chain( - ((hh, h_prefix + hs) for hh, hs in zip(h, self.entity_shape)), - ((rr, r_prefix + rs) for rr, rs in zip(r, self.relation_shape)), - ((tt, t_prefix + ts) for tt, ts in zip(t, self.entity_shape)), + entity_shape = self.entity_shape + if isinstance(entity_shape, str): + entity_shape = (entity_shape,) + relation_shape = self.relation_shape + if isinstance(relation_shape, str): + relation_shape = (relation_shape,) + tail_entity_shape = self.tail_entity_shape + if tail_entity_shape is None: + tail_entity_shape = entity_shape + if isinstance(tail_entity_shape, str): + tail_entity_shape = (tail_entity_shape,) + return len(h) == len(entity_shape) and len(r) == len(relation_shape) and len(t) == len(tail_entity_shape) and check_shapes(*itertools.chain( + ((hh, h_prefix + hs) for hh, hs in zip(h, entity_shape)), + ((rr, r_prefix + rs) for rr, rs in zip(r, relation_shape)), + ((tt, t_prefix + ts) for tt, ts in zip(t, tail_entity_shape)), ), raise_or_error=raise_on_error) def score_hrt( self, - h: EntityRepresentation = tuple(), + h: HeadRepresentation = tuple(), r: RelationRepresentation = tuple(), - t: EntityRepresentation = tuple(), + t: TailRepresentation = tuple(), ) -> torch.FloatTensor: """ Score a batch of triples.. @@ -150,9 +163,9 @@ def score_hrt( def score_h( self, - all_entities: EntityRepresentation = tuple(), + all_entities: HeadRepresentation = tuple(), r: RelationRepresentation = tuple(), - t: EntityRepresentation = tuple(), + t: TailRepresentation = tuple(), ) -> torch.FloatTensor: """ Score all head entities. @@ -182,9 +195,9 @@ def score_h( def score_r( self, - h: EntityRepresentation = tuple(), + h: HeadRepresentation = tuple(), all_relations: RelationRepresentation = tuple(), - t: EntityRepresentation = tuple(), + t: TailRepresentation = tuple(), ) -> torch.FloatTensor: """ Score all relations. @@ -214,9 +227,9 @@ def score_r( def score_t( self, - h: EntityRepresentation = tuple(), + h: HeadRepresentation = tuple(), r: RelationRepresentation = tuple(), - all_entities: EntityRepresentation = tuple(), + all_entities: TailRepresentation = tuple(), ) -> torch.FloatTensor: """ Score all tail entities. @@ -358,7 +371,7 @@ def _calculate_missing_shape_information( return input_channels, width, height -class ConvEInteractionFunction(InteractionFunction): +class ConvEInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor]): """ConvE interaction function.""" def __init__( @@ -437,7 +450,6 @@ def forward( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, - **kwargs, ) -> torch.FloatTensor: # noqa: D102 # get tail bias term if "t_bias" not in kwargs: @@ -689,8 +701,7 @@ def forward( class StructuredEmbeddingInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor]]): """Interaction function of Structured Embedding.""" - entity_shape = ("d",) - relation_shape = ("dd",) + relation_shape = "dd" def forward( self, @@ -775,7 +786,6 @@ class UnstructuredModelInteractionFunction(TranslationalInteractionFunction[torc """Interaction function of UnstructuredModel.""" # shapes - entity_shape = ("d",) relation_shape = tuple() def __init__(self, p: int, power_norm: bool = True): From ca1cfc25ac028cdc14651da049d1771db507a6bd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 13:53:24 +0100 Subject: [PATCH 229/690] Adapt types for SE & UM --- src/pykeen/nn/modules.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 0d90f3d2c0..d2de4106a5 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -371,7 +371,7 @@ def _calculate_missing_shape_information( return input_channels, width, height -class ConvEInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor]): +class ConvEInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor]]): """ConvE interaction function.""" def __init__( @@ -698,7 +698,7 @@ def forward( RESCALInteractionFunction = _build_module_from_stateless(pkf.rescal_interaction) -class StructuredEmbeddingInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor]]): +class StructuredEmbeddingInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor], torch.FloatTensor]): """Interaction function of Structured Embedding.""" relation_shape = "dd" @@ -782,7 +782,7 @@ def forward( ) -class UnstructuredModelInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, None]): +class UnstructuredModelInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, None, torch.FloatTensor]): """Interaction function of UnstructuredModel.""" # shapes From c02ddc117a2f3b4139538b7132fb312faaef4fae Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:08:44 +0100 Subject: [PATCH 230/690] Adjust interaction function types and shapes --- src/pykeen/models/unimodal/ntn.py | 17 ++-- src/pykeen/nn/modules.py | 129 ++++++++++++++++++++---------- 2 files changed, 94 insertions(+), 52 deletions(-) diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 4cb4f92657..b619b98819 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -86,14 +86,6 @@ def __init__( num_embeddings=triples_factory.num_relations, embedding_dim=num_slices * self.embedding_dim ** 2, ) - self.vh = Embedding( - num_embeddings=triples_factory.num_relations, - embedding_dim=num_slices * embedding_dim, - ) - self.vt = Embedding( - num_embeddings=triples_factory.num_relations, - embedding_dim=num_slices * embedding_dim, - ) self.b = Embedding( num_embeddings=triples_factory.num_relations, embedding_dim=num_slices, @@ -102,6 +94,15 @@ def __init__( num_embeddings=triples_factory.num_relations, embedding_dim=num_slices, ) + self.vh = Embedding( + num_embeddings=triples_factory.num_relations, + embedding_dim=num_slices * embedding_dim, + ) + self.vt = Embedding( + num_embeddings=triples_factory.num_relations, + embedding_dim=num_slices * embedding_dim, + ) + if non_linearity is None: non_linearity = nn.Tanh() self.non_linearity = non_linearity diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index d2de4106a5..b7668b5294 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -266,6 +266,22 @@ def reset_parameters(self): mod.reset_parameters() +class StatelessInteractionFunction(InteractionFunction[HeadRepresentation, RelationRepresentation, TailRepresentation]): + """Interaction function without state.""" + + def __init__(self, f: Callable[[torch.FloatTensor, ...], torch.FloatTensor]): + super().__init__() + self.f = f + + def forward( + self, + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + ) -> torch.FloatTensor: # noqa: D102 + return self.f(*h, *r, *t) + + def _build_module_from_stateless( f: Callable[[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], torch.FloatTensor], ) -> Type[InteractionFunction]: @@ -287,7 +303,7 @@ def forward( return StatelessInteractionFunction -class TranslationalInteractionFunction(InteractionFunction[EntityRepresentation, RelationRepresentation], ABC): +class TranslationalInteractionFunction(InteractionFunction[HeadRepresentation, RelationRepresentation, TailRepresentation], ABC): """The translational interaction function shared by the TransE, TransR, TransH, and other Trans models.""" def __init__(self, p: int, power_norm: bool = False): @@ -304,7 +320,7 @@ def __init__(self, p: int, power_norm: bool = False): self.power_norm = power_norm -class TransEInteractionFunction(TranslationalInteractionFunction): +class TransEInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): """The TransE interaction function.""" def forward( @@ -318,8 +334,11 @@ def forward( return pkf.transe_interaction(h=h, r=r, t=t, p=self.p, power_norm=self.power_norm) -#: Interaction function of ComplEx -ComplExInteractionFunction = _build_module_from_stateless(pkf.complex_interaction) +class ComplExInteractionFunction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): + """Interaction function of ComplEx.""" + + def __init__(self): + super().__init__(f=pkf.complex_interaction) def _calculate_missing_shape_information( @@ -477,7 +496,7 @@ def forward( ) -class ConvKBInteractionFunction(InteractionFunction): +class ConvKBInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): """Interaction function of ConvKB.""" def __init__( @@ -512,7 +531,6 @@ def forward( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, - **kwargs, ) -> torch.FloatTensor: # noqa: D102 return pkf.convkb_interaction( h=h, @@ -525,11 +543,14 @@ def forward( ) -#: Interaction function for HolE -DistMultInteractionFunction = _build_module_from_stateless(pkf.distmult_interaction) +class DistMultInteractionFunction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): + """Interaction function of DistMult.""" + + def __init__(self): + super().__init__(f=pkf.distmult_interaction) -class ERMLPInteractionFunction(InteractionFunction): +class ERMLPInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): """ Interaction function of ER-MLP. @@ -564,9 +585,7 @@ def forward( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, - **kwargs, ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) return pkf.ermlp_interaction( h=h, r=r, @@ -588,7 +607,7 @@ def reset_parameters(self): # noqa: D102 ) -class ERMLPEInteractionFunction(InteractionFunction): +class ERMLPEInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): """Interaction function of ER-MLP.""" def __init__( @@ -616,38 +635,43 @@ def forward( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, - **kwargs, ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs=kwargs) return pkf.ermlpe_interaction(h=h, r=r, t=t, mlp=self.mlp) -class TransRInteractionFunction(TranslationalInteractionFunction): +class TransRInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor], torch.FloatTensor]): """The TransR interaction function.""" + relation_shape = ("e", "de") + def __init__(self, p: int, power_norm: bool = True): super().__init__(p=p, power_norm=power_norm) def forward( self, h: torch.FloatTensor, - r: torch.FloatTensor, + r: Tuple[torch.FloatTensor, torch.FloatTensor], t: torch.FloatTensor, - **kwargs, ) -> torch.FloatTensor: # noqa:D102 - m_r = kwargs.pop('m_r') - self._check_for_empty_kwargs(kwargs=kwargs) + r, m_r = r return pkf.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p, power_norm=self.power_norm) -#: Interaction function of RotatE. -RotatEInteraction = _build_module_from_stateless(pkf.rotate_interaction) +class RotatEInteraction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): + """Interaction function of RotatE.""" + + def __init__(self): + super().__init__(f=pkf.rotate_interaction) -#: Interaction function for HolE. -HolEInteractionFunction = _build_module_from_stateless(pkf.hole_interaction) +class HolEInteraction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): + """Interaction function for HolE.""" -class ProjEInteractionFunction(InteractionFunction): + def __init__(self): + super().__init__(f=pkf.hole_interaction) + + +class ProjEInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): """Interaction function for ProjE.""" def __init__( @@ -684,18 +708,20 @@ def forward( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, - **kwargs, ) -> torch.FloatTensor: # noqa:D102 - self._check_for_empty_kwargs(kwargs=kwargs) - - # Compute score return pkf.proje_interaction( h=h, r=r, t=t, d_e=self.d_e, d_r=self.d_r, b_c=self.b_c, b_p=self.b_p, activation=self.inner_non_linearity, - ).view(-1, 1) + ) + +class RESCALInteractionFunction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): + """Interaction function of RESCAL.""" -RESCALInteractionFunction = _build_module_from_stateless(pkf.rescal_interaction) + relation_shape = "dd" + + def __init__(self): + super().__init__(f=pkf.rescal_interaction) class StructuredEmbeddingInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor], torch.FloatTensor]): @@ -713,7 +739,7 @@ def forward( return pkf.structured_embedding_interaction(h=h, r_h=rh, r_t=rt, t=t, p=self.p, power_norm=self.power_norm) -class TuckerInteractionFunction(InteractionFunction): +class TuckerInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): """Interaction function of Tucker.""" def __init__( @@ -766,9 +792,7 @@ def forward( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, - **kwargs, ) -> torch.FloatTensor: # noqa:D102 - self._check_for_empty_kwargs(kwargs=kwargs) return pkf.tucker_interaction( h=h, r=r, @@ -800,25 +824,41 @@ def forward( return pkf.unstructured_model_interaction(h, t, p=self.p, power_norm=self.power_norm) -class TransDInteractionFunction(TranslationalInteractionFunction): +class TransDInteractionFunction( + TranslationalInteractionFunction[ + Tuple[torch.FloatTensor, torch.FloatTensor], + Tuple[torch.FloatTensor, torch.FloatTensor], + Tuple[torch.FloatTensor, torch.FloatTensor], + ] +): """Interaction function of TransD.""" + entity_shape = ("d", "d") + relation_shape = ("e", "e") + def forward( self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, + h: Tuple[torch.FloatTensor, torch.FloatTensor], + r: Tuple[torch.FloatTensor, torch.FloatTensor], + t: Tuple[torch.FloatTensor, torch.FloatTensor], ) -> torch.FloatTensor: # noqa:D102 - h_p = kwargs.pop("h_p") - r_p = kwargs.pop("r_p") - t_p = kwargs.pop("t_p") + h, h_p = h + r, r_p = r + t, t_p = t return pkf.transd_interaction(h=h, r=r, t=t, h_p=h_p, r_p=r_p, t_p=t_p, p=self.p, power_norm=self.power_norm) -class NTNInteractionFunction(InteractionFunction): +class NTNInteractionFunction( + InteractionFunction[ + torch.FloatTensor, + Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], + torch.FloatTensor, + ] +): """The interaction function of NTN.""" + relation_shape = ("kdd", "k", "k", "kd", "kd") + def __init__( self, non_linearity: Optional[nn.Module] = None, @@ -831,7 +871,8 @@ def __init__( def forward( self, h: torch.FloatTensor, - r: torch.FloatTensor, + r: Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], t: torch.FloatTensor, - **kwargs, ) -> torch.FloatTensor: + w, b, u, vh, vt = r + return pkf.ntn_interaction(h=h, t=t, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity) From fc865e765d2c470cf95a1d35615f9463c700a752 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:09:33 +0100 Subject: [PATCH 231/690] Use NTN interaction function --- src/pykeen/models/unimodal/ntn.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index b619b98819..e8568afce6 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -10,6 +10,7 @@ from ..base import EntityEmbeddingModel from ...losses import Loss from ...nn import Embedding, functional as pkf +from ...nn.modules import NTNInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -102,10 +103,7 @@ def __init__( num_embeddings=triples_factory.num_relations, embedding_dim=num_slices * embedding_dim, ) - - if non_linearity is None: - non_linearity = nn.Tanh() - self.non_linearity = non_linearity + self.interaction = NTNInteractionFunction(non_linearity=non_linearity) def forward( self, @@ -136,7 +134,7 @@ def forward( vt = self.vt.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim)) if slice_size is None: - return pkf.ntn_interaction( + return self.interaction( h=h_all, t=t_all, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity, ) From b5cc5d3adafa6acfbe1f1ed6d09ce213cf7b4819 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:11:05 +0100 Subject: [PATCH 232/690] Remove unused method --- src/pykeen/nn/modules.py | 37 ++----------------------------------- 1 file changed, 2 insertions(+), 35 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index b7668b5294..2f921a1533 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -5,7 +5,7 @@ import logging import math from abc import ABC -from typing import Any, Callable, Generic, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union +from typing import Callable, Generic, Optional, Sequence, Tuple, TypeVar, Union import torch from torch import nn @@ -57,13 +57,6 @@ def forward( """ raise NotImplementedError - @classmethod - def _check_for_empty_kwargs(cls, kwargs: Mapping[str, Any]) -> None: - """Check that kwargs is empty.""" - # TODO: Deprecated - if len(kwargs) > 0: - raise ValueError(f"{cls.__name__} does not take the following kwargs: {kwargs}") - @staticmethod def _add_dim(*x: torch.FloatTensor, dim: int) -> Sequence[torch.FloatTensor]: """ @@ -282,27 +275,6 @@ def forward( return self.f(*h, *r, *t) -def _build_module_from_stateless( - f: Callable[[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], torch.FloatTensor], -) -> Type[InteractionFunction]: - """Build a stateless interaction function module with a pre-defined functional interface.""" - - class StatelessInteractionFunction(InteractionFunction): - """Interaction function without state or additional parameters.""" - - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - **kwargs, - ) -> torch.FloatTensor: # noqa: D102 - self._check_for_empty_kwargs(kwargs) - return f(h, r, t) - - return StatelessInteractionFunction - - class TranslationalInteractionFunction(InteractionFunction[HeadRepresentation, RelationRepresentation, TailRepresentation], ABC): """The translational interaction function shared by the TransE, TransR, TransH, and other Trans models.""" @@ -328,9 +300,7 @@ def forward( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, - **kwargs, ) -> torch.FloatTensor: # noqa:D102 - self._check_for_empty_kwargs(kwargs=kwargs) return pkf.transe_interaction(h=h, r=r, t=t, p=self.p, power_norm=self.power_norm) @@ -471,10 +441,7 @@ def forward( t: torch.FloatTensor, ) -> torch.FloatTensor: # noqa: D102 # get tail bias term - if "t_bias" not in kwargs: - raise TypeError(f"{self.__class__.__name__}.forward expects keyword argument 't_bias'.") - t_bias: torch.FloatTensor = kwargs.pop("t_bias") - self._check_for_empty_kwargs(kwargs) + t, t_bias = t return pkf.conve_interaction( h=h, r=r, From ffaf7875ae17b7998bea0337f75fdb2080cd4522 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:18:09 +0100 Subject: [PATCH 233/690] Fix type annotation for ConvE --- src/pykeen/nn/modules.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 2f921a1533..c706d26a4a 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -438,7 +438,7 @@ def forward( self, h: torch.FloatTensor, r: torch.FloatTensor, - t: torch.FloatTensor, + t: Tuple[torch.FloatTensor, torch.FloatTensor], ) -> torch.FloatTensor: # noqa: D102 # get tail bias term t, t_bias = t From 94e4b97c87744169595d5eda15b2fa664049b811 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:19:59 +0100 Subject: [PATCH 234/690] fix type annotation of StatelessInteractionFunction --- src/pykeen/nn/modules.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index c706d26a4a..64d116affe 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -262,7 +262,7 @@ def reset_parameters(self): class StatelessInteractionFunction(InteractionFunction[HeadRepresentation, RelationRepresentation, TailRepresentation]): """Interaction function without state.""" - def __init__(self, f: Callable[[torch.FloatTensor, ...], torch.FloatTensor]): + def __init__(self, f: Callable[..., torch.FloatTensor]): super().__init__() self.f = f From 394768c2f1d06fea63d5a7356cd9eb0e801eb0d7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:20:28 +0100 Subject: [PATCH 235/690] Fix class names --- src/pykeen/nn/modules.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 64d116affe..f93599c73b 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -624,14 +624,14 @@ def forward( return pkf.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p, power_norm=self.power_norm) -class RotatEInteraction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): +class RotatEInteractionFunction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): """Interaction function of RotatE.""" def __init__(self): super().__init__(f=pkf.rotate_interaction) -class HolEInteraction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): +class HolEInteractionFunction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): """Interaction function for HolE.""" def __init__(self): From 077087fb48635ad365518cf15464b91a3f67c112 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:22:19 +0100 Subject: [PATCH 236/690] Fix RotatE import --- src/pykeen/models/unimodal/rotate.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index a5a35f172a..406428ee2b 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -13,7 +13,7 @@ from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ -from ...nn.modules import RotatEInteraction +from ...nn.modules import RotatEInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -96,7 +96,7 @@ def __init__( # TODO: regularization super().__init__( triples_factory=triples_factory, - interaction_function=RotatEInteraction(), + interaction_function=RotatEInteractionFunction(), embedding_dim=2 * embedding_dim, loss=loss, automatic_memory_optimization=automatic_memory_optimization, From 7786e760a7ca2f38ec65a0bf82bcf3276ff8e2b2 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:23:27 +0100 Subject: [PATCH 237/690] Fix tqdm usage in training_loop.py --- src/pykeen/training/training_loop.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/pykeen/training/training_loop.py b/src/pykeen/training/training_loop.py index f59f46b3ee..8c26f3acb3 100644 --- a/src/pykeen/training/training_loop.py +++ b/src/pykeen/training/training_loop.py @@ -317,7 +317,8 @@ def _train( # noqa: C901 num_training_instances = self.training_instances.num_instances # When size probing, we don't want progress bars - if not only_size_probing and use_tqdm: + use_tqdm = use_tqdm and not only_size_probing + if use_tqdm: # Create progress bar _tqdm_kwargs = dict(desc=f'Training epochs on {self.device}', unit='epoch') if tqdm_kwargs is not None: @@ -407,10 +408,11 @@ def _train( # noqa: C901 result_tracker.log_metrics({'loss': epoch_loss}, step=epoch) # Print loss information to console - epochs.set_postfix({ - 'loss': self.losses_per_epochs[-1], - 'prev_loss': self.losses_per_epochs[-2] if epoch > 2 else float('nan'), - }) + if use_tqdm: + epochs.set_postfix({ + 'loss': self.losses_per_epochs[-1], + 'prev_loss': self.losses_per_epochs[-2] if epoch > 2 else float('nan'), + }) if stopper is not None and stopper.should_evaluate(epoch) and stopper.should_stop(epoch): return self.losses_per_epochs From 43b326c849de3e9f0e739871f13ff41128fc1efb Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:23:39 +0100 Subject: [PATCH 238/690] Fix ConvE --- src/pykeen/models/unimodal/conv_e.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 033abfde42..2338454800 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -184,4 +184,4 @@ def forward( t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) t_bias = self.bias_term.get_in_canonical_shape(indices=t_indices) self.regularize_if_necessary(h, r, t) - return self.interaction_function(h=h, r=r, t=t, t_bias=t_bias) + return self.interaction_function(h=h, r=r, t=(t, t_bias)) From 7022564733a2bf75c1d1be9c99ed13566f33d331 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:26:23 +0100 Subject: [PATCH 239/690] Fix NTN interaction usage --- src/pykeen/models/unimodal/ntn.py | 38 ++++--------------------------- 1 file changed, 4 insertions(+), 34 deletions(-) diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index e8568afce6..954d2f1f46 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -9,7 +9,7 @@ from ..base import EntityEmbeddingModel from ...losses import Loss -from ...nn import Embedding, functional as pkf +from ...nn import Embedding from ...nn.modules import NTNInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory @@ -124,8 +124,8 @@ def forward( assert slice_size is None, "not implemented" #: shape: (batch_size, num_entities, d) - h_all = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - t_all = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) + h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) + t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) w = self.w.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim, self.embedding_dim)) b = self.b.get_in_canonical_shape(indices=r_indices) @@ -133,34 +133,4 @@ def forward( vh = self.vh.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim)) vt = self.vt.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim)) - if slice_size is None: - return self.interaction( - h=h_all, t=t_all, - w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity, - ) - - # TODO: Not implemented - if h_all.shape[1] > t_all.shape[1]: - h_was_split = True - split_tensor = torch.split(h_all, slice_size, dim=1) - constant_tensor = t_all - else: - h_was_split = False - split_tensor = torch.split(t_all, slice_size, dim=1) - constant_tensor = h_all - - scores_arr = [] - for split in split_tensor: - if h_was_split: - h = split - t = constant_tensor - else: - h = constant_tensor - t = split - score = pkf.ntn_interaction( - h=h, t=t, - w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity, - ) - scores_arr.append(score) - - return torch.cat(scores_arr, dim=1) + return self.interaction(h=h, t=t, r=(w, b, u, vh, vt)), From 97b57a5afb7e4d6f55668771f8bd2a6dc8da3a54 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:26:56 +0100 Subject: [PATCH 240/690] Adjust type annotation of SingleVectorEmbeddingModel --- src/pykeen/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index a60ef51354..786f13921b 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1217,7 +1217,7 @@ class SingleVectorEmbeddingModel(EntityRelationEmbeddingModel): def __init__( self, triples_factory: TriplesFactory, - interaction_function: InteractionFunction, + interaction_function: InteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], embedding_dim: int = 200, relation_dim: Optional[int] = None, automatic_memory_optimization: Optional[bool] = None, From a16b5913d1aab2cb54aca69df11196ad717b672e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:29:49 +0100 Subject: [PATCH 241/690] Fix StatelessInteractionFunction with single vector representation --- src/pykeen/nn/modules.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index f93599c73b..fa3ef5fa0e 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -272,6 +272,8 @@ def forward( r: RelationRepresentation, t: TailRepresentation, ) -> torch.FloatTensor: # noqa: D102 + # normalization + h, r, t = [(x,) if torch.is_tensor(x) else x for x in (h, r, t)] return self.f(*h, *r, *t) From 7fccf7060babdb42b0bc4348d1b8de9a98f9ced5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:32:57 +0100 Subject: [PATCH 242/690] Fix typo --- src/pykeen/models/unimodal/ntn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 954d2f1f46..5ddb2f9b05 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -133,4 +133,4 @@ def forward( vh = self.vh.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim)) vt = self.vt.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim)) - return self.interaction(h=h, t=t, r=(w, b, u, vh, vt)), + return self.interaction(h=h, t=t, r=(w, b, u, vh, vt)) From 04aa8ed498a06e6dae1920b0528e0fda84bd7628 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:33:50 +0100 Subject: [PATCH 243/690] Add ABC to mitigate multiple initialization from __init_subclass__ hook --- src/pykeen/models/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 786f13921b..06be6ef242 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1211,7 +1211,7 @@ class MultimodalModel(EntityRelationEmbeddingModel): """A multimodal KGE model.""" -class SingleVectorEmbeddingModel(EntityRelationEmbeddingModel): +class SingleVectorEmbeddingModel(EntityRelationEmbeddingModel, ABC): """A base class for embedding models which store a single vector for each entity and relation.""" def __init__( @@ -1287,7 +1287,7 @@ def forward( return self.interaction_function(h=h, r=r, t=t) -class TwoVectorEmbeddingModel(EntityRelationEmbeddingModel): +class TwoVectorEmbeddingModel(EntityRelationEmbeddingModel, ABC): """A model with two vectors for each entity and relation.""" def __init__( From 10beb04d2bc44e5d0e533c5945ed8c584465229d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:44:55 +0100 Subject: [PATCH 244/690] Fix subclasses of TwoVectorEmbeddingModel --- src/pykeen/models/base.py | 34 +++++++------ src/pykeen/models/unimodal/kg2e.py | 22 ++------- .../models/unimodal/structured_embedding.py | 45 ++++++++++++++--- src/pykeen/models/unimodal/trans_d.py | 2 +- src/pykeen/nn/modules.py | 49 +++++++++++++++++-- 5 files changed, 108 insertions(+), 44 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 06be6ef242..d63ca65f1d 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1293,6 +1293,11 @@ class TwoVectorEmbeddingModel(EntityRelationEmbeddingModel, ABC): def __init__( self, triples_factory: TriplesFactory, + interaction_function: InteractionFunction[ + Tuple[torch.FloatTensor, torch.FloatTensor], + Tuple[torch.FloatTensor, torch.FloatTensor], + Tuple[torch.FloatTensor, torch.FloatTensor], + ], embedding_dim: int = 50, relation_dim: Optional[int] = None, loss: Optional[Loss] = None, @@ -1331,17 +1336,7 @@ def __init__( embedding_dim=self.relation_dim, specification=second_relation_embedding_specification, ) - - def _forward( - self, - h1: torch.FloatTensor, - h2: torch.FloatTensor, - r1: torch.FloatTensor, - r2: torch.FloatTensor, - t1: torch.FloatTensor, - t2: torch.FloatTensor, - ) -> torch.FloatTensor: - raise NotImplementedError + self.interaction_function = interaction_function def forward( self, @@ -1355,10 +1350,10 @@ def forward( r2 = self.second_relation_embeddings.get_in_canonical_shape(indices=r_indices) t1 = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) t2 = self.second_entity_embeddings.get_in_canonical_shape(indices=t_indices) - return self._forward(h1, h2, r1, r2, t1, t2) + return self.interaction_function(h=(h1, h2), r=(r1, r2), t=(t1, t2)) -class TwoSideEmbeddingModel(TwoVectorEmbeddingModel): +class TwoSideEmbeddingModel(EntityRelationEmbeddingModel): """A model which averages scores for forward and backward model.""" def __init__( @@ -1390,8 +1385,17 @@ def __init__( regularizer=regularizer, embedding_specification=embedding_specification, relation_embedding_specification=relation_embedding_specification, - second_embedding_specification=second_embedding_specification, - second_relation_embedding_specification=second_relation_embedding_specification, + ) + # extra embeddings + self.second_entity_embeddings = Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=self.embedding_dim, + specification=second_embedding_specification, + ) + self.second_relation_embeddings = Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=self.relation_dim, + specification=second_relation_embedding_specification, ) self.interaction_function = interaction_function diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index d8861f4ffd..28dde17cca 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -13,6 +13,7 @@ from ...nn import functional as pkf from ...nn.emb import EmbeddingSpecification from ...nn.functional import KG2E_SIMILARITIES +from ...nn.modules import KG2EInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -81,6 +82,9 @@ def __init__( """ super().__init__( triples_factory=triples_factory, + interaction_function=KG2EInteractionFunction( + similarity=dist_similarity, + ), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, @@ -106,21 +110,3 @@ def __init__( constrainer_kwargs=dict(min=c_min, max=c_max), ), ) - # Similarity function used for distributions - dist_similarity = dist_similarity.upper() - if dist_similarity not in KG2E_SIMILARITIES: - raise ValueError(dist_similarity) - self.similarity = dist_similarity - - def _forward( - self, - h1: torch.FloatTensor, - h2: torch.FloatTensor, - r1: torch.FloatTensor, - r2: torch.FloatTensor, - t1: torch.FloatTensor, - t2: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa: D102 - return pkf.kg2e_interaction( - h_mean=h1, h_var=h2, r_mean=r1, r_var=r2, t_mean=t1, t_var=t2, similarity=self.similarity, - ) diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index 51e940ada6..27e156ca5e 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -6,10 +6,11 @@ from typing import Optional import numpy as np +import torch from torch import nn from torch.nn import functional -from .. import SingleVectorEmbeddingModel +from .. import EntityRelationEmbeddingModel from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ @@ -24,7 +25,7 @@ ] -class StructuredEmbedding(SingleVectorEmbeddingModel): +class StructuredEmbedding(EntityRelationEmbeddingModel): r"""An implementation of the Structured Embedding (SE) published by [bordes2011]_. SE applies role- and relation-specific projection matrices @@ -71,12 +72,8 @@ def __init__( ) super().__init__( triples_factory=triples_factory, - interaction_function=StructuredEmbeddingInteractionFunction( - p=scoring_fct_norm, - power_norm=False, - ), embedding_dim=embedding_dim, - relation_dim=2 * embedding_dim ** 2, # head and tail projection matrices + relation_dim=embedding_dim ** 2, # head projection matrices automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, @@ -90,3 +87,37 @@ def __init__( initializer=relation_initializer, ), ) + self.second_relation_embedding = EmbeddingSpecification( + initializer=relation_initializer, + ).make( + num_embeddings=self.num_relations, + embedding_dim=self.relation_dim, + device=self.device, + ) + self.interaction_function = StructuredEmbeddingInteractionFunction( + p=scoring_fct_norm, + power_norm=False, + ), + + def forward( + self, + h_indices: Optional[torch.LongTensor] = None, + r_indices: Optional[torch.LongTensor] = None, + t_indices: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: + """Evaluate the given triples. + + :param h_indices: shape: (batch_size,) + The indices for head entities. If None, score against all. + :param r_indices: shape: (batch_size,) + The indices for relations. If None, score against all. + :param t_indices: shape: (batch_size,) + The indices for tail entities. If None, score against all. + + :return: The scores, shape: (batch_size, num_entities) + """ + h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) + r_h = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) + r_t = self.second_relation_embedding.get_in_canonical_shape(indices=r_indices) + t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) + return self.interaction_function(h=h, r=(r_h, r_t), t=t) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index e5a34953ae..dc3cdada0e 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -72,6 +72,7 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, + interaction_function=TransDInteractionFunction(p=2, power_norm=True), embedding_dim=embedding_dim, relation_dim=relation_dim, automatic_memory_optimization=automatic_memory_optimization, @@ -90,7 +91,6 @@ def __init__( constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), ) - self.interaction_function = TransDInteractionFunction(p=2, power_norm=True) def _forward( self, diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index fa3ef5fa0e..6f27dfe0b1 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -11,6 +11,7 @@ from torch import nn from . import functional as pkf +from .functional import KG2E_SIMILARITIES from ..utils import check_shapes logger = logging.getLogger(__name__) @@ -39,9 +40,9 @@ class InteractionFunction(nn.Module, Generic[HeadRepresentation, RelationReprese def forward( self, - h: HeadRepresentation = tuple(), - r: RelationRepresentation = tuple(), - t: TailRepresentation = tuple(), + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, ) -> torch.FloatTensor: """Compute broadcasted triple scores given representations for head, relation and tails. @@ -845,3 +846,45 @@ def forward( ) -> torch.FloatTensor: w, b, u, vh, vt = r return pkf.ntn_interaction(h=h, t=t, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity) + + +class KG2EInteractionFunction( + InteractionFunction[ + Tuple[torch.FloatTensor, torch.FloatTensor], + Tuple[torch.FloatTensor, torch.FloatTensor], + Tuple[torch.FloatTensor, torch.FloatTensor], + ] +): + """Interaction function of KG2E.""" + + def __init__( + self, + similarity: str = "KL", + exact: bool = True, + ): + super().__init__() + similarity = similarity.upper() + if similarity not in KG2E_SIMILARITIES: + raise ValueError(similarity) + self.similarity = similarity + self.exact = exact + + def forward( + self, + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + ) -> torch.FloatTensor: + h_mean, h_var = h + r_mean, r_var = r + t_mean, t_var = t + return pkf.kg2e_interaction( + h_mean=h_mean, + h_var=h_var, + r_mean=r_mean, + r_var=r_var, + t_mean=t_mean, + t_var=t_var, + similarity=self.similarity, + exact=self.exact, + ) From 1f2c1c2904dbb9772761f097651e67564ae33cfe Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:46:26 +0100 Subject: [PATCH 245/690] Fix TransR --- src/pykeen/models/unimodal/trans_r.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 01af972d22..fc0e387025 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -141,4 +141,4 @@ def forward( indices=r_indices, reshape_dim=(self.embedding_dim, self.relation_dim), ) - return self.interaction_function(h=h, r=r, t=t, m_r=m_r) + return self.interaction_function(h=h, r=(r, m_r), t=t) From 3b8abea91536a4deeeec9976851c263520859cf8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:50:29 +0100 Subject: [PATCH 246/690] Remove typo --- src/pykeen/models/unimodal/structured_embedding.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index 27e156ca5e..e3169c553a 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -97,7 +97,7 @@ def __init__( self.interaction_function = StructuredEmbeddingInteractionFunction( p=scoring_fct_norm, power_norm=False, - ), + ) def forward( self, From 892e80999fb91c2f24d4755a30357f2c6b51076a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:51:26 +0100 Subject: [PATCH 247/690] Add embedding specification --- src/pykeen/nn/emb.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index 46a981b162..d78b027422 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -15,6 +15,7 @@ __all__ = [ 'RepresentationModule', 'Embedding', + 'EmbeddingSpecification', ] From c4227b47d6f3d7481ef2cd8ce11d1b4aaee6b75c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:55:21 +0100 Subject: [PATCH 248/690] Pass representations from outside --- src/pykeen/models/base.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index d63ca65f1d..6edbb5468b 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -17,7 +17,7 @@ from torch import nn from ..losses import Loss, MarginRankingLoss, NSSALoss -from ..nn import Embedding +from ..nn import Embedding, RepresentationModule from ..nn.emb import EmbeddingSpecification from ..nn.modules import InteractionFunction from ..regularizers import NoRegularizer, Regularizer @@ -239,6 +239,9 @@ def __init__( preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, + entity_representations: Optional[Sequence[RepresentationModule]] = None, + relation_representations: Optional[Sequence[RepresentationModule]] = None, + tail_representations: Optional[Sequence[RepresentationModule]] = None, ) -> None: """Initialize the module. @@ -306,6 +309,12 @@ def __init__( # This allows to store the optimized parameters self.automatic_memory_optimization = automatic_memory_optimization + # Important: use ModuleList to ensure that Pytorch correctly handles their devices and parameters + self.entity_representations = nn.ModuleList(entity_representations) + self.relation_representations = nn.ModuleList(relation_representations) + if tail_representations is not None: + self.tail_representations = nn.ModuleList(tail_representations) + @classmethod def _is_abstract(cls) -> bool: return inspect.isabstract(cls) From 2f3f9998953ecd9bca48a832c79ac911eb5e44af Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 14:57:15 +0100 Subject: [PATCH 249/690] Move type annotations to typing --- src/pykeen/nn/modules.py | 7 ++----- src/pykeen/typing.py | 11 ++++++++--- 2 files changed, 10 insertions(+), 8 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 6f27dfe0b1..511ebe8845 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -5,21 +5,18 @@ import logging import math from abc import ABC -from typing import Callable, Generic, Optional, Sequence, Tuple, TypeVar, Union +from typing import Callable, Generic, Optional, Sequence, Tuple, Union import torch from torch import nn from . import functional as pkf from .functional import KG2E_SIMILARITIES +from ..typing import HeadRepresentation, RelationRepresentation, TailRepresentation from ..utils import check_shapes logger = logging.getLogger(__name__) -HeadRepresentation = TypeVar("HeadRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) -RelationRepresentation = TypeVar("RelationRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) -TailRepresentation = TypeVar("TailRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) - class InteractionFunction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation]): """Base class for interaction functions.""" diff --git a/src/pykeen/typing.py b/src/pykeen/typing.py index c12e6c78b6..e363d3f4a6 100644 --- a/src/pykeen/typing.py +++ b/src/pykeen/typing.py @@ -2,7 +2,7 @@ """Type hints for PyKEEN.""" -from typing import Callable, Mapping, TypeVar, Union +from typing import Callable, Mapping, Sequence, TypeVar, Union import numpy as np import torch @@ -12,8 +12,10 @@ 'MappedTriples', 'EntityMapping', 'RelationMapping', - 'InteractionFunction', 'DeviceHint', + 'HeadRepresentation', + 'RelationRepresentation', + 'TailRepresentation', ] LabeledTriples = np.ndarray @@ -23,9 +25,12 @@ # comment: TypeVar expects none, or at least two super-classes TensorType = TypeVar("TensorType", torch.Tensor, torch.FloatTensor) -InteractionFunction = Callable[[TensorType, TensorType, TensorType], TensorType] Initializer = Callable[[TensorType], TensorType] Normalizer = Callable[[TensorType], TensorType] Constrainer = Callable[[TensorType], TensorType] DeviceHint = Union[None, str, torch.device] + +HeadRepresentation = TypeVar("HeadRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) +RelationRepresentation = TypeVar("RelationRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) +TailRepresentation = TypeVar("TailRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) From 9b4ce068452b62d5ac38056c3c3d94aa5c42ffdd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 15:02:49 +0100 Subject: [PATCH 250/690] store entity/relation representations and interaction function in base module --- src/pykeen/models/base.py | 18 ++++++++++++------ src/pykeen/typing.py | 7 ++++--- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 6edbb5468b..2fe259d842 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -9,7 +9,7 @@ from abc import ABC, abstractmethod from collections import defaultdict from operator import itemgetter -from typing import Any, ClassVar, Collection, Dict, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Type, Union +from typing import Any, ClassVar, Collection, Dict, Generic, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Type, Union import numpy as np import pandas as pd @@ -22,7 +22,7 @@ from ..nn.modules import InteractionFunction from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory -from ..typing import Constrainer, DeviceHint, Initializer, MappedTriples, Normalizer +from ..typing import Constrainer, DeviceHint, HeadRepresentation, Initializer, MappedTriples, Normalizer, RelationRepresentation, TailRepresentation from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed __all__ = [ @@ -209,7 +209,7 @@ def _new_init(self, *args, **kwargs): cls.__init__ = _new_init -class Model(nn.Module, ABC): +class Model(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation], ABC): """A base module for all of the KGE models.""" #: A dictionary of hyper-parameters to the models that use them @@ -230,6 +230,12 @@ class Model(nn.Module, ABC): #: The instance of the regularizer regularizer: Regularizer + #: The entity representations + entity_representations: Sequence[RepresentationModule] + + #: The relation representations + relation_representations: Sequence[RepresentationModule] + def __init__( self, triples_factory: TriplesFactory, @@ -241,7 +247,7 @@ def __init__( regularizer: Optional[Regularizer] = None, entity_representations: Optional[Sequence[RepresentationModule]] = None, relation_representations: Optional[Sequence[RepresentationModule]] = None, - tail_representations: Optional[Sequence[RepresentationModule]] = None, + interaction_function: InteractionFunction[HeadRepresentation, RelationRepresentation, TailRepresentation] = None, ) -> None: """Initialize the module. @@ -312,8 +318,8 @@ def __init__( # Important: use ModuleList to ensure that Pytorch correctly handles their devices and parameters self.entity_representations = nn.ModuleList(entity_representations) self.relation_representations = nn.ModuleList(relation_representations) - if tail_representations is not None: - self.tail_representations = nn.ModuleList(tail_representations) + + self.interaction_function = interaction_function @classmethod def _is_abstract(cls) -> bool: diff --git a/src/pykeen/typing.py b/src/pykeen/typing.py index e363d3f4a6..ab21b6eb20 100644 --- a/src/pykeen/typing.py +++ b/src/pykeen/typing.py @@ -31,6 +31,7 @@ DeviceHint = Union[None, str, torch.device] -HeadRepresentation = TypeVar("HeadRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) -RelationRepresentation = TypeVar("RelationRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) -TailRepresentation = TypeVar("TailRepresentation", torch.FloatTensor, Sequence[torch.FloatTensor]) +Representation = torch.FloatTensor +HeadRepresentation = TypeVar("HeadRepresentation", Representation, Sequence[Representation]) +RelationRepresentation = TypeVar("RelationRepresentation", Representation, Sequence[Representation]) +TailRepresentation = TypeVar("TailRepresentation", Representation, Sequence[Representation]) From a0accbc9dc2b944846f18cf743f118874a37545e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 15:11:26 +0100 Subject: [PATCH 251/690] Move implementation of forward to base class --- src/pykeen/models/base.py | 93 +++++++++++++++++++++++++-------------- 1 file changed, 61 insertions(+), 32 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 2fe259d842..68908b40a8 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -245,8 +245,8 @@ def __init__( preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, - entity_representations: Optional[Sequence[RepresentationModule]] = None, - relation_representations: Optional[Sequence[RepresentationModule]] = None, + entity_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, + relation_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, interaction_function: InteractionFunction[HeadRepresentation, RelationRepresentation, TailRepresentation] = None, ) -> None: """Initialize the module. @@ -315,6 +315,11 @@ def __init__( # This allows to store the optimized parameters self.automatic_memory_optimization = automatic_memory_optimization + # normalization + if not isinstance(entity_representations, Sequence): + entity_representations = [entity_representations] + if not isinstance(relation_representations, Sequence): + relation_representations = [relation_representations] # Important: use ModuleList to ensure that Pytorch correctly handles their devices and parameters self.entity_representations = nn.ModuleList(entity_representations) self.relation_representations = nn.ModuleList(relation_representations) @@ -981,7 +986,7 @@ def _compute_loss( ) return self.loss(tensor_1, tensor_2) + self.regularizer.term - @abstractmethod + # @abstractmethod def forward( self, h_indices: Optional[torch.LongTensor], @@ -1004,7 +1009,20 @@ def forward( :return: shape: (batch_size, num_heads, num_relations, num_tails) The score for each triple. """ - raise NotImplementedError + h, r, t = [ + [ + representation.get_in_canonical_shape(indices=indices) + for representation in representations + ] + for indices, representations in ( + (h_indices, self.entity_representations), + (r_indices, self.relation_representations), + (t_indices, self.entity_representations), + ) + ] + # normalization + h, r, t = [x[0] if len(x) == 1 else x for x in (h, r, t)] + return self.interaction_function(h=h, r=r, t=t) def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: """Forward pass. @@ -1226,7 +1244,7 @@ class MultimodalModel(EntityRelationEmbeddingModel): """A multimodal KGE model.""" -class SingleVectorEmbeddingModel(EntityRelationEmbeddingModel, ABC): +class SingleVectorEmbeddingModel(Model, ABC): """A base class for embedding models which store a single vector for each entity and relation.""" def __init__( @@ -1264,42 +1282,53 @@ def __init__( :param regularizer: The regularizer to use. """ + # Default for relation dimensionality + if relation_dim is None: + relation_dim = embedding_dim + self.embedding_dim = embedding_dim + self.relation_dim = relation_dim super().__init__( triples_factory=triples_factory, - embedding_dim=embedding_dim, - relation_dim=relation_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - embedding_specification=embedding_specification, - relation_embedding_specification=relation_embedding_specification, + interaction_function=interaction_function, + entity_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + specification=embedding_specification, + ), + relation_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=relation_dim, + specification=relation_embedding_specification, + ), ) - self.interaction_function = interaction_function - - def forward( - self, - h_indices: Optional[torch.LongTensor] = None, - r_indices: Optional[torch.LongTensor] = None, - t_indices: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: - """Evaluate the given triples. - - :param h_indices: shape: (batch_size,) - The indices for head entities. If None, score against all. - :param r_indices: shape: (batch_size,) - The indices for relations. If None, score against all. - :param t_indices: shape: (batch_size,) - The indices for tail entities. If None, score against all. - - :return: The scores, shape: (batch_size, num_entities) - """ - h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) - t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - return self.interaction_function(h=h, r=r, t=t) + # + # def forward( + # self, + # h_indices: Optional[torch.LongTensor] = None, + # r_indices: Optional[torch.LongTensor] = None, + # t_indices: Optional[torch.LongTensor] = None, + # ) -> torch.FloatTensor: + # """Evaluate the given triples. + # + # :param h_indices: shape: (batch_size,) + # The indices for head entities. If None, score against all. + # :param r_indices: shape: (batch_size,) + # The indices for relations. If None, score against all. + # :param t_indices: shape: (batch_size,) + # The indices for tail entities. If None, score against all. + # + # :return: The scores, shape: (batch_size, num_entities) + # """ + # h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) + # r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) + # t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) + # return self.interaction_function(h=h, r=r, t=t) class TwoVectorEmbeddingModel(EntityRelationEmbeddingModel, ABC): From ca87c654d795b5e141f6ba34a2b0f4b5dbfdb8ad Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 15:18:01 +0100 Subject: [PATCH 252/690] Subclass NTN from Model (since it also has relation representations) --- src/pykeen/models/base.py | 37 ++++++++-------- src/pykeen/models/unimodal/ntn.py | 74 +++++++++++-------------------- 2 files changed, 44 insertions(+), 67 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 68908b40a8..0c6d668fac 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -6,7 +6,7 @@ import inspect import itertools as itt import logging -from abc import ABC, abstractmethod +from abc import ABC from collections import defaultdict from operator import itemgetter from typing import Any, ClassVar, Collection, Dict, Generic, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Type, Union @@ -22,7 +22,7 @@ from ..nn.modules import InteractionFunction from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory -from ..typing import Constrainer, DeviceHint, HeadRepresentation, Initializer, MappedTriples, Normalizer, RelationRepresentation, TailRepresentation +from ..typing import Constrainer, DeviceHint, HeadRepresentation, Initializer, MappedTriples, Normalizer, RelationRepresentation, Representation, TailRepresentation from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed __all__ = [ @@ -326,6 +326,9 @@ def __init__( self.interaction_function = interaction_function + # reset parameters + self.reset_parameters_() + @classmethod def _is_abstract(cls) -> bool: return inspect.isabstract(cls) @@ -1126,6 +1129,7 @@ class EntityEmbeddingModel(Model): def __init__( self, triples_factory: TriplesFactory, + interaction_function: InteractionFunction[Representation, RelationRepresentation, Representation], embedding_dim: int = 50, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, @@ -1148,6 +1152,8 @@ def __init__( .. seealso:: Constructor of the base class :class:`pykeen.models.Model` """ + # TODO: use spec + self.embedding_dim = embedding_dim super().__init__( triples_factory=triples_factory, automatic_memory_optimization=automatic_memory_optimization, @@ -1156,23 +1162,18 @@ def __init__( random_seed=random_seed, regularizer=regularizer, predict_with_sigmoid=predict_with_sigmoid, + entity_representations=Embedding( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + initializer=entity_initializer, + initializer_kwargs=entity_initializer_kwargs, + normalizer=entity_normalizer, + normalizer_kwargs=entity_normalizer_kwargs, + constrainer=entity_constrainer, + constrainer_kwargs=entity_constrainer_kwargs, + ), + interaction_function=interaction_function, ) - self.entity_embeddings = Embedding.init_with_device( - num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, - device=self.device, - initializer=entity_initializer, - initializer_kwargs=entity_initializer_kwargs, - normalizer=entity_normalizer, - normalizer_kwargs=entity_normalizer_kwargs, - constrainer=entity_constrainer, - constrainer_kwargs=entity_constrainer_kwargs, - ) - - @property - def embedding_dim(self) -> int: # noqa:D401 - """The entity embedding dimension.""" - return self.entity_embeddings.embedding_dim class EntityRelationEmbeddingModel(Model, ABC): diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 5ddb2f9b05..7ff26b8598 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -4,10 +4,9 @@ from typing import Optional -import torch from torch import nn -from ..base import EntityEmbeddingModel +from .. import Model from ...losses import Loss from ...nn import Embedding from ...nn.modules import NTNInteractionFunction @@ -20,7 +19,7 @@ ] -class NTN(EntityEmbeddingModel): +class NTN(Model): r"""An implementation of NTN from [socher2013]_. NTN uses a bilinear tensor layer instead of a standard linear neural network layer: @@ -72,65 +71,42 @@ def __init__( :param non_linearity: A non-linear activation function. Defaults to the hyperbolic tangent :class:`torch.nn.Tanh`. """ - super().__init__( - triples_factory=triples_factory, - embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, - regularizer=regularizer, - ) + self.embedding_dim = embedding_dim self.num_slices = num_slices - self.w = Embedding( + w = Embedding( num_embeddings=triples_factory.num_relations, - embedding_dim=num_slices * self.embedding_dim ** 2, + embedding_dim=num_slices * embedding_dim ** 2, ) - self.b = Embedding( + b = Embedding( num_embeddings=triples_factory.num_relations, embedding_dim=num_slices, ) - self.u = Embedding( + u = Embedding( num_embeddings=triples_factory.num_relations, embedding_dim=num_slices, ) - self.vh = Embedding( + vh = Embedding( num_embeddings=triples_factory.num_relations, embedding_dim=num_slices * embedding_dim, ) - self.vt = Embedding( + vt = Embedding( num_embeddings=triples_factory.num_relations, embedding_dim=num_slices * embedding_dim, ) - self.interaction = NTNInteractionFunction(non_linearity=non_linearity) - - def forward( - self, - h_indices: Optional[torch.LongTensor] = None, - r_indices: Optional[torch.LongTensor] = None, - t_indices: Optional[torch.LongTensor] = None, - slice_size: int = None, - ) -> torch.FloatTensor: - """ - Compute scores for NTN. - - :param h_indices: shape: (batch_size,) - :param r_indices: shape: (batch_size,) - :param t_indices: shape: (batch_size,) - - :return: shape: (batch_size, num_heads, num_relations, num_tails) - """ - assert slice_size is None, "not implemented" - - #: shape: (batch_size, num_entities, d) - h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - w = self.w.get_in_canonical_shape(indices=r_indices, - reshape_dim=(self.num_slices, self.embedding_dim, self.embedding_dim)) - b = self.b.get_in_canonical_shape(indices=r_indices) - u = self.u.get_in_canonical_shape(indices=r_indices) - vh = self.vh.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim)) - vt = self.vt.get_in_canonical_shape(indices=r_indices, reshape_dim=(self.num_slices, self.embedding_dim)) - - return self.interaction(h=h, t=t, r=(w, b, u, vh, vt)) + super().__init__( + triples_factory=triples_factory, + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + random_seed=random_seed, + regularizer=regularizer, + interaction_function=NTNInteractionFunction( + non_linearity=non_linearity, + ), + entity_representations=Embedding( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + ), + relation_representations=(w, b, u, vh, vt), + ) From 7b6a2395eec937f3872c1ce30fd73aec9ffe9de6 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 15:20:58 +0100 Subject: [PATCH 253/690] Let UM subclass directly from Model --- src/pykeen/models/base.py | 2 +- .../models/unimodal/unstructured_model.py | 26 +++++++++++++------ 2 files changed, 19 insertions(+), 9 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 0c6d668fac..dc1f025714 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1125,6 +1125,7 @@ def load_state(self, path: str) -> None: class EntityEmbeddingModel(Model): """A base module for most KGE models that have one embedding for entities.""" + # TODO: deprecated def __init__( self, @@ -1152,7 +1153,6 @@ def __init__( .. seealso:: Constructor of the base class :class:`pykeen.models.Model` """ - # TODO: use spec self.embedding_dim = embedding_dim super().__init__( triples_factory=triples_factory, diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index 0cb18c48a3..f0e84a110e 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -6,8 +6,9 @@ import torch.autograd -from ..base import EntityEmbeddingModel +from .. import Model from ...losses import Loss +from ...nn import Embedding from ...nn.init import xavier_normal_ from ...nn.modules import UnstructuredModelInteractionFunction from ...regularizers import Regularizer @@ -19,7 +20,7 @@ ] -class UnstructuredModel(EntityEmbeddingModel): +class UnstructuredModel(Model): r"""An implementation of the Unstructured Model (UM) published by [bordes2014]_. UM computes the distance between head and tail entities then applies the $l_p$ norm. @@ -50,6 +51,7 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, + predict_with_sigmoid: bool = False, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, @@ -59,17 +61,22 @@ def __init__( :param embedding_dim: The entity embedding dimension $d$. Is usually $d \in [50, 300]$. :param scoring_fct_norm: The $l_p$ norm. Usually 1 for UM. """ + self.embedding_dim = embedding_dim super().__init__( triples_factory=triples_factory, - embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, + predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - entity_initializer=xavier_normal_, + interaction_function=UnstructuredModelInteractionFunction(p=scoring_fct_norm), + entity_representations=Embedding( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + initializer=xavier_normal_, + ), ) - self.interaction_function = UnstructuredModelInteractionFunction(p=scoring_fct_norm) def forward( self, @@ -77,9 +84,12 @@ def forward( r_indices: Optional[torch.LongTensor], t_indices: Optional[torch.LongTensor], ) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - scores = self.interaction_function(h=h, r=None, t=t) + scores = super(UnstructuredModel, self).forward( + h_indices=h_indices, + r_indices=r_indices, + t_indices=t_indices, + ) + # TODO: move this logic to superclass # same score for all relations repeats = [1, 1, 1, 1] if r_indices is None: From d0c05190fe3b2cbbbd830849f94cd4ac338e62a8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 15:28:26 +0100 Subject: [PATCH 254/690] simplify test --- tests/test_models.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 93afc32a6d..d4ed24dd2f 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -481,14 +481,11 @@ def test_score_t_with_score_hrt_equality(self) -> None: def test_reset_parameters_constructor_call(self): """Tests whether reset_parameters is called in the constructor.""" with patch.object(self.model_cls, 'reset_parameters_', return_value=None) as mock_method: - try: - self.model_cls( - triples_factory=self.factory, - embedding_dim=self.embedding_dim, - **(self.model_kwargs or {}), - ) - except TypeError as error: - assert error.args == ("'NoneType' object is not callable",) + self.model_cls( + triples_factory=self.factory, + embedding_dim=self.embedding_dim, + **(self.model_kwargs or {}), + ) mock_method.assert_called_once() def test_custom_representations(self): From 3ea3d21eab3bd798ca7f3f6d3ab004e7e8ef51ba Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 15:34:05 +0100 Subject: [PATCH 255/690] Fix comparison --- src/pykeen/models/base.py | 19 ++++++++++++--- .../models/unimodal/unstructured_model.py | 24 ------------------- 2 files changed, 16 insertions(+), 27 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index dc1f025714..2e48035b2a 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -316,9 +316,9 @@ def __init__( self.automatic_memory_optimization = automatic_memory_optimization # normalization - if not isinstance(entity_representations, Sequence): + if entity_representations is not None and not isinstance(entity_representations, Sequence): entity_representations = [entity_representations] - if not isinstance(relation_representations, Sequence): + if relation_representations is not None and not isinstance(relation_representations, Sequence): relation_representations = [relation_representations] # Important: use ModuleList to ensure that Pytorch correctly handles their devices and parameters self.entity_representations = nn.ModuleList(entity_representations) @@ -1025,7 +1025,20 @@ def forward( ] # normalization h, r, t = [x[0] if len(x) == 1 else x for x in (h, r, t)] - return self.interaction_function(h=h, r=r, t=t) + + scores = self.interaction_function(h=h, r=r, t=t) + if len(self.relation_representations) == 0: + # same score for all relations + repeats = [1, 1, 1, 1] + if r_indices is None: + repeats[2] = self.num_relations + else: + relation_batch_size = len(r_indices) + if scores.shape[0] < relation_batch_size: + repeats[0] = relation_batch_size + scores = scores.repeat(*repeats) + + return scores def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: """Forward pass. diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index f0e84a110e..faeb298c81 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -4,8 +4,6 @@ from typing import Optional -import torch.autograd - from .. import Model from ...losses import Loss from ...nn import Embedding @@ -77,25 +75,3 @@ def __init__( initializer=xavier_normal_, ), ) - - def forward( - self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], - ) -> torch.FloatTensor: # noqa: D102 - scores = super(UnstructuredModel, self).forward( - h_indices=h_indices, - r_indices=r_indices, - t_indices=t_indices, - ) - # TODO: move this logic to superclass - # same score for all relations - repeats = [1, 1, 1, 1] - if r_indices is None: - repeats[2] = self.num_relations - else: - relation_batch_size = len(r_indices) - if scores.shape[0] < relation_batch_size: - repeats[0] = relation_batch_size - return scores.repeat(*repeats) From fdfe0bb8b91edf4c6c5f0dfe06a845bc1cc2ba5d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 15:37:47 +0100 Subject: [PATCH 256/690] Fix double initialization --- src/pykeen/models/base.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 2e48035b2a..bd6779cd48 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -333,11 +333,9 @@ def __init__( def _is_abstract(cls) -> bool: return inspect.isabstract(cls) - def __init_subclass__(cls, reset_parameters_post_init: bool = True, **kwargs): # noqa:D105 + def __init_subclass__(cls, **kwargs): # noqa:D105 if not cls._is_abstract(): _track_hyperparameters(cls) - if reset_parameters_post_init: - _add_post_reset_parameters(cls) @property def can_slice_h(self) -> bool: From 66f58eaf048ad48df2fc15c8db05753e7c47ef4d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 15:40:45 +0100 Subject: [PATCH 257/690] Fix _check_constraints --- tests/test_models.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index d4ed24dd2f..f9e292857c 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -572,7 +572,7 @@ def _check_constraints(self): Entity embeddings have to have unit L2 norm. """ - entity_norms = self.model.entity_embeddings(indices=None).norm(p=2, dim=-1) + entity_norms = self.model.entity_representations[0](indices=None).norm(p=2, dim=-1) assert torch.allclose(entity_norms, torch.ones_like(entity_norms)) def _test_score_all_triples(self, k: Optional[int], batch_size: int = 16): @@ -654,7 +654,7 @@ def _check_constraints(self): Entity embeddings have to have at most unit L2 norm. """ - assert all_in_bounds(self.model.entity_embeddings(indices=None).norm(p=2, dim=-1), high=1., a_tol=_EPSILON) + assert all_in_bounds(self.model.entity_representations[0](indices=None).norm(p=2, dim=-1), high=1., a_tol=_EPSILON) class _TestKG2E(_ModelTestCase): @@ -786,7 +786,7 @@ def _check_constraints(self): """ relation_abs = ( self.model - .relation_embeddings(indices=None) + .relation_representations[0](indices=None) .view(self.factory.num_relations, -1, 2) .norm(p=2, dim=-1) ) @@ -998,7 +998,7 @@ def _check_constraints(self): Entity embeddings have to have unit L2 norm. """ - entity_norms = self.model.entity_embeddings(indices=None).norm(p=2, dim=-1) + entity_norms = self.model.entity_representations[0](indices=None).norm(p=2, dim=-1) assert torch.allclose(entity_norms, torch.ones_like(entity_norms)) From 9a7c472d240878c3cb2d1b5818304dd16bc7fdf1 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 15:48:39 +0100 Subject: [PATCH 258/690] Fix ConvE --- src/pykeen/models/base.py | 42 +++++++--------- src/pykeen/models/unimodal/conv_e.py | 75 +++++++++++++++------------- 2 files changed, 59 insertions(+), 58 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index bd6779cd48..3336c27322 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1136,6 +1136,7 @@ def load_state(self, path: str) -> None: class EntityEmbeddingModel(Model): """A base module for most KGE models that have one embedding for entities.""" + # TODO: deprecated def __init__( @@ -1193,6 +1194,7 @@ class EntityRelationEmbeddingModel(Model, ABC): def __init__( self, triples_factory: TriplesFactory, + interaction_function: InteractionFunction[Representation, Representation, Representation], embedding_dim: int = 50, relation_dim: Optional[int] = None, loss: Optional[Loss] = None, @@ -1213,39 +1215,31 @@ def __init__( .. seealso:: Constructor of the base class :class:`pykeen.models.Model` .. seealso:: Constructor of the base class :class:`pykeen.models.EntityEmbeddingModel` """ + self.embedding_dim = embedding_dim + # Default for relation dimensionality + if relation_dim is None: + relation_dim = embedding_dim + self.relation_dim = relation_dim super().__init__( triples_factory=triples_factory, + interaction_function=interaction_function, automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, predict_with_sigmoid=predict_with_sigmoid, + entity_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + specification=embedding_specification, + ), + relation_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=relation_dim, + specification=relation_embedding_specification, + ), ) - self.entity_embeddings = Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, - specification=embedding_specification, - ) - - # Default for relation dimensionality - if relation_dim is None: - relation_dim = embedding_dim - self.relation_embeddings = Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - embedding_dim=relation_dim, - specification=relation_embedding_specification, - ) - - @property - def embedding_dim(self) -> int: # noqa:D401 - """The entity embedding dimension.""" - return self.entity_embeddings.embedding_dim - - @property - def relation_dim(self): # noqa:D401 - """The relation embedding dimension.""" - return self.relation_embeddings.embedding_dim def _can_slice(fn) -> bool: diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 2338454800..75a19b30a0 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -7,7 +7,7 @@ import torch from torch import nn -from ..base import EntityRelationEmbeddingModel +from .. import Model from ...losses import BCEAfterSigmoidLoss, Loss from ...nn import Embedding from ...nn.emb import EmbeddingSpecification @@ -24,7 +24,7 @@ logger = logging.getLogger(__name__) -class ConvE(EntityRelationEmbeddingModel): +class ConvE(Model): r"""An implementation of ConvE from [dettmers2018]_. ConvE is a CNN-based approach. For each triple $(h,r,t)$, the input to ConvE is a matrix @@ -134,54 +134,61 @@ def __init__( 'This can be done by defining the TriplesFactory class with the _create_inverse_triples_ parameter set ' 'to true.', ) - super().__init__( triples_factory=triples_factory, - embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - embedding_specification=EmbeddingSpecification( - initializer=xavier_normal_, + entity_representations=[ + Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + specification=EmbeddingSpecification( + initializer=xavier_normal_, + ), + ), + # ConvE uses one bias for each entity + Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=1, + specification=EmbeddingSpecification( + initializer=nn.init.zeros_, + ), + ) + ], + relation_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=embedding_dim, + specification=EmbeddingSpecification( + initializer=xavier_normal_, + ), ), - relation_embedding_specification=EmbeddingSpecification( - initializer=xavier_normal_, + interaction_function=ConvEInteractionFunction( + input_channels=input_channels, + output_channels=output_channels, + embedding_height=embedding_height, + embedding_width=embedding_width, + kernel_height=kernel_height, + kernel_width=kernel_width, + input_dropout=input_dropout, + output_dropout=output_dropout, + feature_map_dropout=feature_map_dropout, + embedding_dim=embedding_dim, + apply_batch_normalization=apply_batch_normalization, ), ) - # ConvE uses one bias for each entity - self.bias_term = Embedding.init_with_device( - num_embeddings=triples_factory.num_entities, - embedding_dim=1, - device=self.device, - initializer=nn.init.zeros_, - ) - - self.interaction_function = ConvEInteractionFunction( - input_channels=input_channels, - output_channels=output_channels, - embedding_height=embedding_height, - embedding_width=embedding_width, - kernel_height=kernel_height, - kernel_width=kernel_width, - input_dropout=input_dropout, - output_dropout=output_dropout, - feature_map_dropout=feature_map_dropout, - embedding_dim=embedding_dim, - apply_batch_normalization=apply_batch_normalization, - ) - def forward( self, h_indices: Optional[torch.LongTensor], r_indices: Optional[torch.LongTensor], t_indices: Optional[torch.LongTensor], ) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) - t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - t_bias = self.bias_term.get_in_canonical_shape(indices=t_indices) + h = self.entity_representations[0].get_in_canonical_shape(indices=h_indices) + r = self.relation_representations[0].get_in_canonical_shape(indices=r_indices) + t = self.entity_representations[0].get_in_canonical_shape(indices=t_indices) + t_bias = self.entity_representations[1].get_in_canonical_shape(indices=t_indices) self.regularize_if_necessary(h, r, t) return self.interaction_function(h=h, r=r, t=(t, t_bias)) From df0d05ccf559ab2f2928c52db47801e05e79e86c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 15:52:41 +0100 Subject: [PATCH 259/690] Fix class hierarchy for TwoVectorModel --- src/pykeen/models/base.py | 86 ++++++++++++++------------------------- tests/test_models.py | 8 ++-- 2 files changed, 34 insertions(+), 60 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 3336c27322..3c9f7bf0fe 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1313,31 +1313,9 @@ def __init__( specification=relation_embedding_specification, ), ) - # - # def forward( - # self, - # h_indices: Optional[torch.LongTensor] = None, - # r_indices: Optional[torch.LongTensor] = None, - # t_indices: Optional[torch.LongTensor] = None, - # ) -> torch.FloatTensor: - # """Evaluate the given triples. - # - # :param h_indices: shape: (batch_size,) - # The indices for head entities. If None, score against all. - # :param r_indices: shape: (batch_size,) - # The indices for relations. If None, score against all. - # :param t_indices: shape: (batch_size,) - # The indices for tail entities. If None, score against all. - # - # :return: The scores, shape: (batch_size, num_entities) - # """ - # h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - # r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) - # t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - # return self.interaction_function(h=h, r=r, t=t) - - -class TwoVectorEmbeddingModel(EntityRelationEmbeddingModel, ABC): + + +class TwoVectorEmbeddingModel(Model, ABC): """A model with two vectors for each entity and relation.""" def __init__( @@ -1361,46 +1339,42 @@ def __init__( second_embedding_specification: Optional[EmbeddingSpecification] = None, second_relation_embedding_specification: Optional[EmbeddingSpecification] = None, ) -> None: + if relation_dim is None: + relation_dim = embedding_dim super().__init__( triples_factory=triples_factory, - embedding_dim=embedding_dim, - relation_dim=relation_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - embedding_specification=embedding_specification, - relation_embedding_specification=relation_embedding_specification, - ) - - # extra embeddings - self.second_entity_embeddings = Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - embedding_dim=self.embedding_dim, - specification=second_embedding_specification, - ) - self.second_relation_embeddings = Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - embedding_dim=self.relation_dim, - specification=second_relation_embedding_specification, + entity_representations=[ + Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + specification=embedding_specification, + ), + Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + specification=second_embedding_specification, + ) + ], + relation_representations=[ + Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=relation_dim, + specification=relation_embedding_specification, + ), + Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=relation_dim, + specification=second_relation_embedding_specification, + ) + ], + interaction_function=interaction_function, ) - self.interaction_function = interaction_function - - def forward( - self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], - ) -> torch.FloatTensor: # noqa: D102 - h1 = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - h2 = self.second_entity_embeddings.get_in_canonical_shape(indices=h_indices) - r1 = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) - r2 = self.second_relation_embeddings.get_in_canonical_shape(indices=r_indices) - t1 = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - t2 = self.second_entity_embeddings.get_in_canonical_shape(indices=t_indices) - return self.interaction_function(h=(h1, h2), r=(r1, r2), t=(t1, t2)) class TwoSideEmbeddingModel(EntityRelationEmbeddingModel): diff --git a/tests/test_models.py b/tests/test_models.py index f9e292857c..051539dc17 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -669,12 +669,12 @@ def _check_constraints(self): * Covariances have to have values between c_min and c_max """ self.model: TwoVectorEmbeddingModel - low = self.model.second_entity_embeddings.constrainer.keywords['min'] - high = self.model.second_entity_embeddings.constrainer.keywords['max'] + low = self.model.entity_representations[1].constrainer.keywords['min'] + high = self.model.entity_representations[1].constrainer.keywords['max'] - for embedding in (self.model.entity_embeddings, self.model.relation_embeddings): + for embedding in (self.model.entity_representations[0], self.model.relation_representations[0]): assert all_in_bounds(embedding(indices=None).norm(p=2, dim=-1), high=1., a_tol=_EPSILON) - for cov in (self.model.second_entity_embeddings, self.model.second_relation_embeddings): + for cov in (self.model.entity_representations[1], self.model.relation_representations[1]): assert all_in_bounds(cov(indices=None), low=low, high=high) From d6f56e6eb10a672f1c47e3725c341fbe2d97ea5c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 15:56:06 +0100 Subject: [PATCH 260/690] Temporarily comment out TransD manual test --- src/pykeen/models/unimodal/trans_d.py | 13 -- tests/test_models.py | 244 +++++++++++++------------- 2 files changed, 123 insertions(+), 134 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index dc3cdada0e..e09e433437 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -4,8 +4,6 @@ from typing import Optional -import torch.autograd - from ..base import TwoVectorEmbeddingModel from ...losses import Loss from ...nn.emb import EmbeddingSpecification @@ -91,14 +89,3 @@ def __init__( constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), ) - - def _forward( - self, - h1: torch.FloatTensor, - h2: torch.FloatTensor, - r1: torch.FloatTensor, - r2: torch.FloatTensor, - t1: torch.FloatTensor, - t2: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa:D102 - return self.interaction_function(h=h1, r=r1, t=t1, h_p=h2, r_p=r2, t_p=t2, p=2, power_norm=True) diff --git a/tests/test_models.py b/tests/test_models.py index 051539dc17..7aa72a8f79 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -8,6 +8,7 @@ import traceback import unittest from typing import Any, ClassVar, Mapping, Optional, Type +from unittest.case import SkipTest from unittest.mock import patch import numpy @@ -842,131 +843,132 @@ def _check_constraints(self): Entity and relation embeddings have to have at most unit L2 norm. """ - for emb in (self.model.entity_embeddings, self.model.relation_embeddings): + for emb in (self.model.entity_representations[0], self.model.relation_representations[0]): assert all_in_bounds(emb(indices=None).norm(p=2, dim=-1), high=1., a_tol=_EPSILON) def test_score_hrt_manual(self): """Manually test interaction function of TransD.""" - self.model: TwoVectorEmbeddingModel - - # entity embeddings - weights = torch.as_tensor(data=[[2., 2.], [4., 4.]], dtype=torch.float) - entity_embeddings = Embedding( - num_embeddings=2, - embedding_dim=2, - ) - entity_embeddings._embeddings.weight.data.copy_(weights) - self.model.entity_embeddings = entity_embeddings - - projection_weights = torch.as_tensor(data=[[3., 3.], [2., 2.]], dtype=torch.float) - entity_projection_embeddings = Embedding( - num_embeddings=2, - embedding_dim=2, - ) - entity_projection_embeddings._embeddings.weight.data.copy_(projection_weights) - self.model.second_entity_embeddings = entity_projection_embeddings - - # relation embeddings - relation_weights = torch.as_tensor(data=[[4.], [4.]], dtype=torch.float) - relation_embeddings = Embedding( - num_embeddings=2, - embedding_dim=1, - ) - relation_embeddings._embeddings.weight.data.copy_(relation_weights) - self.model.relation_embeddings = relation_embeddings - - relation_projection_weights = torch.as_tensor(data=[[5.], [3.]], dtype=torch.float) - relation_projection_embeddings = Embedding( - num_embeddings=2, - embedding_dim=1, - ) - relation_projection_embeddings._embeddings.weight.data.copy_(relation_projection_weights) - self.model.second_relation_embeddings = relation_projection_embeddings - - # Compute Scores - batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 1]], dtype=torch.long) - scores = self.model.score_hrt(hrt_batch=batch) - self.assertEqual(scores.shape[0], 2) - self.assertEqual(scores.shape[1], 1) - first_score = scores[0].item() - self.assertAlmostEqual(first_score, -16, delta=0.01) - - # Use different dimension for relation embedding: relation_dim > entity_dim - # relation embeddings - relation_weights = torch.as_tensor(data=[[3., 3., 3.], [3., 3., 3.]], dtype=torch.float) - relation_embeddings = Embedding( - num_embeddings=2, - embedding_dim=3, - ) - relation_embeddings._embeddings.weight.data.copy_(relation_weights) - self.model.relation_embeddings = relation_embeddings - - relation_projection_weights = torch.as_tensor(data=[[4., 4., 4.], [4., 4., 4.]], dtype=torch.float) - relation_projection_embeddings = Embedding( - num_embeddings=2, - embedding_dim=3, - ) - relation_projection_embeddings._embeddings.weight.data.copy_(relation_projection_weights) - self.model.second_relation_embeddings = relation_projection_embeddings - - # Compute Scores - batch = torch.as_tensor(data=[[0, 0, 0]], dtype=torch.long) - scores = self.model.score_hrt(hrt_batch=batch) - self.assertAlmostEqual(scores.item(), -27, delta=0.01) - - batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 0]], dtype=torch.long) - scores = self.model.score_hrt(hrt_batch=batch) - self.assertEqual(scores.shape[0], 2) - self.assertEqual(scores.shape[1], 1) - first_score = scores[0].item() - second_score = scores[1].item() - self.assertAlmostEqual(first_score, -27, delta=0.01) - self.assertAlmostEqual(second_score, -27, delta=0.01) - - # Use different dimension for relation embedding: relation_dim < entity_dim - # entity embeddings - weights = torch.as_tensor(data=[[1., 1., 1.], [1., 1., 1.]], dtype=torch.float) - entity_embeddings = Embedding( - num_embeddings=2, - embedding_dim=3, - ) - entity_embeddings._embeddings.weight.data.copy_(weights) - self.model.entity_embeddings = entity_embeddings - - projection_weights = torch.as_tensor(data=[[2., 2., 2.], [2., 2., 2.]], dtype=torch.float) - entity_projection_embeddings = Embedding( - num_embeddings=2, - embedding_dim=3, - ) - entity_projection_embeddings._embeddings.weight.data.copy_(projection_weights) - self.model.second_entity_embeddings = entity_projection_embeddings - - # relation embeddings - relation_weights = torch.as_tensor(data=[[3., 3.], [3., 3.]], dtype=torch.float) - relation_embeddings = Embedding( - num_embeddings=2, - embedding_dim=2, - ) - relation_embeddings._embeddings.weight.data.copy_(relation_weights) - self.model.relation_embeddings = relation_embeddings - - relation_projection_weights = torch.as_tensor(data=[[4., 4.], [4., 4.]], dtype=torch.float) - relation_projection_embeddings = Embedding( - num_embeddings=2, - embedding_dim=2, - ) - relation_projection_embeddings._embeddings.weight.data.copy_(relation_projection_weights) - self.model.second_relation_embeddings = relation_projection_embeddings - - # Compute Scores - batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 0]], dtype=torch.long) - scores = self.model.score_hrt(hrt_batch=batch) - self.assertEqual(scores.shape[0], 2) - self.assertEqual(scores.shape[1], 1) - first_score = scores[0].item() - second_score = scores[1].item() - self.assertAlmostEqual(first_score, -18, delta=0.01) - self.assertAlmostEqual(second_score, -18, delta=0.01) + raise SkipTest("TODO: Move this test to interaction function checks.") + # self.model: TwoVectorEmbeddingModel + # + # # entity embeddings + # weights = torch.as_tensor(data=[[2., 2.], [4., 4.]], dtype=torch.float) + # entity_embeddings = Embedding( + # num_embeddings=2, + # embedding_dim=2, + # ) + # entity_embeddings._embeddings.weight.data.copy_(weights) + # self.model.entity_representations[0] = entity_embeddings + # + # projection_weights = torch.as_tensor(data=[[3., 3.], [2., 2.]], dtype=torch.float) + # entity_projection_embeddings = Embedding( + # num_embeddings=2, + # embedding_dim=2, + # ) + # entity_projection_embeddings._embeddings.weight.data.copy_(projection_weights) + # self.model.second_entity_embeddings = entity_projection_embeddings + # + # # relation embeddings + # relation_weights = torch.as_tensor(data=[[4.], [4.]], dtype=torch.float) + # relation_embeddings = Embedding( + # num_embeddings=2, + # embedding_dim=1, + # ) + # relation_embeddings._embeddings.weight.data.copy_(relation_weights) + # self.model.relation_representations[0] = relation_embeddings + # + # relation_projection_weights = torch.as_tensor(data=[[5.], [3.]], dtype=torch.float) + # relation_projection_embeddings = Embedding( + # num_embeddings=2, + # embedding_dim=1, + # ) + # relation_projection_embeddings._embeddings.weight.data.copy_(relation_projection_weights) + # self.model.second_relation_embeddings = relation_projection_embeddings + # + # # Compute Scores + # batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 1]], dtype=torch.long) + # scores = self.model.score_hrt(hrt_batch=batch) + # self.assertEqual(scores.shape[0], 2) + # self.assertEqual(scores.shape[1], 1) + # first_score = scores[0].item() + # self.assertAlmostEqual(first_score, -16, delta=0.01) + # + # # Use different dimension for relation embedding: relation_dim > entity_dim + # # relation embeddings + # relation_weights = torch.as_tensor(data=[[3., 3., 3.], [3., 3., 3.]], dtype=torch.float) + # relation_embeddings = Embedding( + # num_embeddings=2, + # embedding_dim=3, + # ) + # relation_embeddings._embeddings.weight.data.copy_(relation_weights) + # self.model.relation_representations[0] = relation_embeddings + # + # relation_projection_weights = torch.as_tensor(data=[[4., 4., 4.], [4., 4., 4.]], dtype=torch.float) + # relation_projection_embeddings = Embedding( + # num_embeddings=2, + # embedding_dim=3, + # ) + # relation_projection_embeddings._embeddings.weight.data.copy_(relation_projection_weights) + # self.model.relation_representations[1] = relation_projection_embeddings + # + # # Compute Scores + # batch = torch.as_tensor(data=[[0, 0, 0]], dtype=torch.long) + # scores = self.model.score_hrt(hrt_batch=batch) + # self.assertAlmostEqual(scores.item(), -27, delta=0.01) + # + # batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 0]], dtype=torch.long) + # scores = self.model.score_hrt(hrt_batch=batch) + # self.assertEqual(scores.shape[0], 2) + # self.assertEqual(scores.shape[1], 1) + # first_score = scores[0].item() + # second_score = scores[1].item() + # self.assertAlmostEqual(first_score, -27, delta=0.01) + # self.assertAlmostEqual(second_score, -27, delta=0.01) + # + # # Use different dimension for relation embedding: relation_dim < entity_dim + # # entity embeddings + # weights = torch.as_tensor(data=[[1., 1., 1.], [1., 1., 1.]], dtype=torch.float) + # entity_embeddings = Embedding( + # num_embeddings=2, + # embedding_dim=3, + # ) + # entity_embeddings._embeddings.weight.data.copy_(weights) + # self.model.entity_representations[0] = entity_embeddings + # + # projection_weights = torch.as_tensor(data=[[2., 2., 2.], [2., 2., 2.]], dtype=torch.float) + # entity_projection_embeddings = Embedding( + # num_embeddings=2, + # embedding_dim=3, + # ) + # entity_projection_embeddings._embeddings.weight.data.copy_(projection_weights) + # self.model.second_entity_embeddings = entity_projection_embeddings + # + # # relation embeddings + # relation_weights = torch.as_tensor(data=[[3., 3.], [3., 3.]], dtype=torch.float) + # relation_embeddings = Embedding( + # num_embeddings=2, + # embedding_dim=2, + # ) + # relation_embeddings._embeddings.weight.data.copy_(relation_weights) + # self.model.relation_representations[0] = relation_embeddings + # + # relation_projection_weights = torch.as_tensor(data=[[4., 4.], [4., 4.]], dtype=torch.float) + # relation_projection_embeddings = Embedding( + # num_embeddings=2, + # embedding_dim=2, + # ) + # relation_projection_embeddings._embeddings.weight.data.copy_(relation_projection_weights) + # self.model.relation_representations[1] = relation_projection_embeddings + # + # # Compute Scores + # batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 0]], dtype=torch.long) + # scores = self.model.score_hrt(hrt_batch=batch) + # self.assertEqual(scores.shape[0], 2) + # self.assertEqual(scores.shape[1], 1) + # first_score = scores[0].item() + # second_score = scores[1].item() + # self.assertAlmostEqual(first_score, -18, delta=0.01) + # self.assertAlmostEqual(second_score, -18, delta=0.01) def test_project_entity(self): """Test _project_entity.""" From 3e78825edbf1b0fed6f0b3d03b366258eb8a9933 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 16:15:52 +0100 Subject: [PATCH 261/690] Start to add tests for interaction functions --- tests/test_interactions.py | 94 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 94 insertions(+) create mode 100644 tests/test_interactions.py diff --git a/tests/test_interactions.py b/tests/test_interactions.py new file mode 100644 index 0000000000..75bf66e5bb --- /dev/null +++ b/tests/test_interactions.py @@ -0,0 +1,94 @@ +"""Tests for interaction functions.""" +import unittest +from typing import Any, Generic, Mapping, MutableMapping, Optional, Tuple, Type, TypeVar + +import torch + +from pykeen.nn.modules import InteractionFunction, TransEInteractionFunction +from pykeen.typing import Representation + +T = TypeVar("T") + + +class GenericTests(Generic[T]): + """Generic tests.""" + + cls: Type[T] + kwargs: Optional[Mapping[str, Any]] = None + instance: T + + def setUp(self) -> None: + kwargs = self.kwargs or {} + kwargs = self._pre_instantiation_hook(kwargs=dict(kwargs)) + self.instance = self.cls(**kwargs) + self.post_instantiation_hook() + + def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: + """Perform actions before instantiation, potentially modyfing kwargs.""" + return kwargs + + def post_instantiation_hook(self) -> None: + """Perform actions after instantiation.""" + + +class InteractionTests(GenericTests[InteractionFunction]): + """Generic test for interaction functions.""" + + dim: int = 2 + batch_size: int = 3 + num_relations: int = 5 + num_entities: int = 7 + + def _get_hrt( + self, + *shapes: Tuple[int, ...], + ) -> Tuple[Representation, ...]: + return tuple([torch.rand(*s, requires_grad=True) for s in shapes]) + + def _get_shapes_for_score_(self, dim: int) -> Tuple[Tuple[int, int, int], Tuple[int, int, int], Tuple[int, int, int], Tuple[int, int, int, int]]: + result = [] + num_choices = self.num_entities if dim != 1 else self.num_relations + for i in range(3): + shape = [1, 1, 1, 1] + if i == dim: + shape[i + 1] = num_choices + else: + shape[0] = self.batch_size + result.append(tuple(shape)) + shape = [self.batch_size, 1, 1, 1] + shape[dim + 1] = num_choices + result.append(tuple(shape)) + return tuple(result) + + def test_forward(self): + for hs, rs, ts, exp in [ + # slcwa + [ + (self.batch_size, 1, self.dim), + (self.batch_size, 1, self.dim), + (self.batch_size, 1, self.dim), + (self.batch_size, 1, 1, 1), + ], + # score_h + self._get_shapes_for_score_(dim=0), + # score_r + self._get_shapes_for_score_(dim=1), + # score_t + self._get_shapes_for_score_(dim=2), + ]: + h, r, t = self._get_hrt(hs, rs, ts) + scores = self.instance.forward(h=h, r=r, t=t) + assert torch.is_tensor(scores) + assert scores.dtype == torch.float32 + assert scores.ndimension() == 4 + assert scores.shape == exp + assert scores.requires_grad + + +class TransETests(InteractionTests, unittest.TestCase): + """Tests for TransE interaction function.""" + + cls = TransEInteractionFunction + kwargs = dict( + p=2, + ) From 57e3deb3925ef153e3730c845b5c56fe7ea0c1f4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 16:32:39 +0100 Subject: [PATCH 262/690] Small fixes and better testing --- src/pykeen/nn/modules.py | 53 +++++++++++++------ tests/test_interactions.py | 106 +++++++++++++++++++++++++------------ 2 files changed, 109 insertions(+), 50 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 511ebe8845..2cd5f67ab1 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -12,7 +12,7 @@ from . import functional as pkf from .functional import KG2E_SIMILARITIES -from ..typing import HeadRepresentation, RelationRepresentation, TailRepresentation +from ..typing import HeadRepresentation, RelationRepresentation, Representation, TailRepresentation from ..utils import check_shapes logger = logging.getLogger(__name__) @@ -114,12 +114,27 @@ def _check_shapes( tail_entity_shape = entity_shape if isinstance(tail_entity_shape, str): tail_entity_shape = (tail_entity_shape,) - return len(h) == len(entity_shape) and len(r) == len(relation_shape) and len(t) == len(tail_entity_shape) and check_shapes(*itertools.chain( + if len(h) != len(entity_shape): + if raise_on_error: + raise ValueError + return False + if len(r) != len(relation_shape): + if raise_on_error: + raise ValueError + return False + if len(t) != len(tail_entity_shape): + if raise_on_error: + raise ValueError + return False + return check_shapes(*itertools.chain( ((hh, h_prefix + hs) for hh, hs in zip(h, entity_shape)), ((rr, r_prefix + rs) for rr, rs in zip(r, relation_shape)), ((tt, t_prefix + ts) for tt, ts in zip(t, tail_entity_shape)), ), raise_or_error=raise_on_error) + def _ensure_tuple(self, *x: Union[Representation, Sequence[Representation]]) -> Tuple[Sequence[Representation], ...]: + return tuple(xx if isinstance(xx, Sequence) else (xx,) for xx in x) + def score_hrt( self, h: HeadRepresentation = tuple(), @@ -139,12 +154,13 @@ def score_hrt( :return: shape: (batch_size, 1) The scores. """ + h, r, t = self._ensure_tuple(h, r, t) assert self._check_shapes(h=h, r=r, t=t) # prepare input to generic score function - h = self._add_dim(h, dim=self.NUM_DIM) - r = self._add_dim(r, dim=self.NUM_DIM) - t = self._add_dim(t, dim=self.NUM_DIM) + h = self._add_dim(*h, dim=self.NUM_DIM) + r = self._add_dim(*r, dim=self.NUM_DIM) + t = self._add_dim(*t, dim=self.NUM_DIM) # get scores scores = self(h=h, r=r, t=t) @@ -171,12 +187,13 @@ def score_h( :return: shape: (batch_size, num_entities) The scores. """ - assert self._check_shapes(h=all_entities, r=r, t=t, h_prefix="n") + h, r, t = self._ensure_tuple(all_entities, r, t) + assert self._check_shapes(h=h, r=r, t=t, h_prefix="n") # prepare input to generic score function - h = self._add_dim(all_entities, dim=self.BATCH_DIM) - r = self._add_dim(r, dim=self.NUM_DIM) - t = self._add_dim(t, dim=self.NUM_DIM) + h = self._add_dim(*h, dim=self.BATCH_DIM) + r = self._add_dim(*r, dim=self.NUM_DIM) + t = self._add_dim(*t, dim=self.NUM_DIM) # get scores scores = self(h=h, r=r, t=t) @@ -203,12 +220,13 @@ def score_r( :return: shape: (batch_size, num_entities) The scores. """ - assert self._check_shapes(h=h, r=all_relations, t=t, r_prefix="n") + h, r, t = self._ensure_tuple(h, all_relations, t) + assert self._check_shapes(h=h, r=r, t=t, r_prefix="n") # prepare input to generic score function - h = self._add_dim(h, dim=self.NUM_DIM) - r = self._add_dim(all_relations, dim=self.BATCH_DIM) - t = self._add_dim(t, dim=self.NUM_DIM) + h = self._add_dim(*h, dim=self.NUM_DIM) + r = self._add_dim(*r, dim=self.BATCH_DIM) + t = self._add_dim(*t, dim=self.NUM_DIM) # get scores scores = self(h=h, r=r, t=t) @@ -235,12 +253,13 @@ def score_t( :return: shape: (batch_size, num_entities) The scores. """ - assert self._check_shapes(h=h, r=r, t=all_entities, t_prefix="n") + h, r, t = self._ensure_tuple(h, r, all_entities) + assert self._check_shapes(h=h, r=r, t=t, t_prefix="n") # prepare input to generic score function - h = self._add_dim(h, dim=self.NUM_DIM) - r = self._add_dim(r, dim=self.NUM_DIM) - t = self._add_dim(all_entities, dim=self.BATCH_DIM) + h = self._add_dim(*h, dim=self.NUM_DIM) + r = self._add_dim(*r, dim=self.NUM_DIM) + t = self._add_dim(*t, dim=self.BATCH_DIM) # get scores scores = self(h=h, r=r, t=t) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 75bf66e5bb..8fecdfd05a 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -4,7 +4,7 @@ import torch -from pykeen.nn.modules import InteractionFunction, TransEInteractionFunction +from pykeen.nn.modules import DistMultInteractionFunction, InteractionFunction, TransEInteractionFunction from pykeen.typing import Representation T = TypeVar("T") @@ -45,44 +45,84 @@ def _get_hrt( ) -> Tuple[Representation, ...]: return tuple([torch.rand(*s, requires_grad=True) for s in shapes]) - def _get_shapes_for_score_(self, dim: int) -> Tuple[Tuple[int, int, int], Tuple[int, int, int], Tuple[int, int, int], Tuple[int, int, int, int]]: - result = [] - num_choices = self.num_entities if dim != 1 else self.num_relations - for i in range(3): - shape = [1, 1, 1, 1] - if i == dim: - shape[i + 1] = num_choices - else: - shape[0] = self.batch_size - result.append(tuple(shape)) - shape = [self.batch_size, 1, 1, 1] - shape[dim + 1] = num_choices - result.append(tuple(shape)) - return tuple(result) + def _check_scores(self, scores: torch.FloatTensor, exp_shape: Tuple[int, ...]): + """Check shape, dtype and gradients of scores.""" + assert torch.is_tensor(scores) + assert scores.dtype == torch.float32 + assert scores.ndimension() == len(exp_shape) + assert scores.shape == exp_shape + assert scores.requires_grad + + def test_score_hrt(self): + """Test score_hrt.""" + h, r, t = self._get_hrt( + (self.batch_size, self.dim), + (self.batch_size, self.dim), + (self.batch_size, self.dim), + ) + scores = self.instance.score_hrt(h=h, r=r, t=t) + self._check_scores(scores=scores, exp_shape=(self.batch_size, 1)) + + def test_score_h(self): + """Test score_h.""" + h, r, t = self._get_hrt( + (self.num_entities, self.dim), + (self.batch_size, self.dim), + (self.batch_size, self.dim), + ) + scores = self.instance.score_h(all_entities=h, r=r, t=t) + self._check_scores(scores=scores, exp_shape=(self.batch_size, self.num_entities)) + + def test_score_r(self): + """Test score_r.""" + h, r, t = self._get_hrt( + (self.batch_size, self.dim), + (self.num_relations, self.dim), + (self.batch_size, self.dim), + ) + scores = self.instance.score_r(h=h, all_relations=r, t=t) + self._check_scores(scores=scores, exp_shape=(self.batch_size, self.num_relations)) + + def test_score_t(self): + """Test score_t.""" + h, r, t = self._get_hrt( + (self.batch_size, self.dim), + (self.batch_size, self.dim), + (self.num_entities, self.dim), + ) + scores = self.instance.score_t(h=h, r=r, all_entities=t) + self._check_scores(scores=scores, exp_shape=(self.batch_size, self.num_entities)) def test_forward(self): - for hs, rs, ts, exp in [ - # slcwa + """Test forward.""" + for hs, rs, ts in [ [ (self.batch_size, 1, self.dim), - (self.batch_size, 1, self.dim), - (self.batch_size, 1, self.dim), - (self.batch_size, 1, 1, 1), + (1, self.num_relations, self.dim), + (self.batch_size, self.num_entities, self.dim), + ], + [ + (1, 1, self.dim), + (1, self.num_relations, self.dim), + (self.batch_size, self.num_entities, self.dim), + ], + [ + (1, self.num_entities, self.dim), + (1, self.num_relations, self.dim), + (1, self.num_entities, self.dim), ], - # score_h - self._get_shapes_for_score_(dim=0), - # score_r - self._get_shapes_for_score_(dim=1), - # score_t - self._get_shapes_for_score_(dim=2), ]: - h, r, t = self._get_hrt(hs, rs, ts) - scores = self.instance.forward(h=h, r=r, t=t) - assert torch.is_tensor(scores) - assert scores.dtype == torch.float32 - assert scores.ndimension() == 4 - assert scores.shape == exp - assert scores.requires_grad + with self.subTest(f"forward({hs}, {rs}, {ts})"): + expected_shape = (max(hs[0], rs[0], ts[0]), hs[1], rs[1], ts[1]) + h, r, t = self._get_hrt(hs, rs, ts) + scores = self.instance(h=h, r=r, t=t) + self._check_scores(scores=scores, exp_shape=expected_shape) + + +class DistMultTests(InteractionTests, unittest.TestCase): + """Tests for DistMult interaction function.""" + + cls = DistMultInteractionFunction class TransETests(InteractionTests, unittest.TestCase): From f41f088c32f1b4fe77f5f7b62494766b534dae80 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 16:40:56 +0100 Subject: [PATCH 263/690] More tests --- tests/test_interactions.py | 25 +++++++++++++++++++++++++ 1 file changed, 25 insertions(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 8fecdfd05a..2418a77a53 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -1,6 +1,7 @@ """Tests for interaction functions.""" import unittest from typing import Any, Generic, Mapping, MutableMapping, Optional, Tuple, Type, TypeVar +from unittest.case import SkipTest import torch @@ -52,6 +53,10 @@ def _check_scores(self, scores: torch.FloatTensor, exp_shape: Tuple[int, ...]): assert scores.ndimension() == len(exp_shape) assert scores.shape == exp_shape assert scores.requires_grad + self._additional_score_checks(scores) + + def _additional_score_checks(self, scores): + """Additional checks for scores.""" def test_score_hrt(self): """Test score_hrt.""" @@ -118,12 +123,26 @@ def test_forward(self): scores = self.instance(h=h, r=r, t=t) self._check_scores(scores=scores, exp_shape=expected_shape) + def test_scores(self): + """Test individual scores.""" + for i in range(10): + h, r, t = self._get_hrt((1, 1, self.dim), (1, 1, self.dim), (1, 1, self.dim)) + scores = self.instance(h=h, r=r, t=t) + exp_score = self._exp_score(h, r, t).item() + assert scores.item() == exp_score + + def _exp_score(self, h, r, t) -> torch.FloatTensor: + raise SkipTest() + class DistMultTests(InteractionTests, unittest.TestCase): """Tests for DistMult interaction function.""" cls = DistMultInteractionFunction + def _exp_score(self, h, r, t) -> torch.FloatTensor: + return (h * r * t).sum(dim=-1) + class TransETests(InteractionTests, unittest.TestCase): """Tests for TransE interaction function.""" @@ -132,3 +151,9 @@ class TransETests(InteractionTests, unittest.TestCase): kwargs = dict( p=2, ) + + def _additional_score_checks(self, scores): + assert (scores <= 0).all() + + def _exp_score(self, h, r, t) -> torch.FloatTensor: + return -(h + r - t).norm(p=2, dim=-1) From 367087bf4361e68409c965da276e30b3b2dea38a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 17:05:44 +0100 Subject: [PATCH 264/690] Fix ConvE interaction tests --- src/pykeen/nn/modules.py | 2 + tests/test_interactions.py | 156 ++++++++++++++++++++++++++++++------- 2 files changed, 128 insertions(+), 30 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 2cd5f67ab1..f08e693daa 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -382,6 +382,8 @@ def _calculate_missing_shape_information( class ConvEInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor]]): """ConvE interaction function.""" + tail_entity_shape = ("d", "k") # with k=1 + def __init__( self, input_channels: Optional[int] = None, diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 2418a77a53..efcd52a705 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -1,11 +1,11 @@ """Tests for interaction functions.""" import unittest -from typing import Any, Generic, Mapping, MutableMapping, Optional, Tuple, Type, TypeVar +from typing import Any, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union from unittest.case import SkipTest import torch -from pykeen.nn.modules import DistMultInteractionFunction, InteractionFunction, TransEInteractionFunction +import pykeen.nn.modules from pykeen.typing import Representation T = TypeVar("T") @@ -32,7 +32,7 @@ def post_instantiation_hook(self) -> None: """Perform actions after instantiation.""" -class InteractionTests(GenericTests[InteractionFunction]): +class InteractionTests(GenericTests[pykeen.nn.modules.InteractionFunction]): """Generic test for interaction functions.""" dim: int = 2 @@ -43,8 +43,13 @@ class InteractionTests(GenericTests[InteractionFunction]): def _get_hrt( self, *shapes: Tuple[int, ...], - ) -> Tuple[Representation, ...]: - return tuple([torch.rand(*s, requires_grad=True) for s in shapes]) + **kwargs + ) -> Tuple[Union[Representation, Sequence[Representation]], ...]: + kwargs.setdefault("d", self.dim) + return tuple( + torch.rand(*s, *(kwargs.get(ss) for ss in ms), requires_grad=True) + for s, ms in zip(shapes, [self.cls.entity_shape, self.cls.relation_shape, self.cls.entity_shape]) + ) def _check_scores(self, scores: torch.FloatTensor, exp_shape: Tuple[int, ...]): """Check shape, dtype and gradients of scores.""" @@ -61,9 +66,9 @@ def _additional_score_checks(self, scores): def test_score_hrt(self): """Test score_hrt.""" h, r, t = self._get_hrt( - (self.batch_size, self.dim), - (self.batch_size, self.dim), - (self.batch_size, self.dim), + (self.batch_size,), + (self.batch_size,), + (self.batch_size,), ) scores = self.instance.score_hrt(h=h, r=r, t=t) self._check_scores(scores=scores, exp_shape=(self.batch_size, 1)) @@ -71,9 +76,9 @@ def test_score_hrt(self): def test_score_h(self): """Test score_h.""" h, r, t = self._get_hrt( - (self.num_entities, self.dim), - (self.batch_size, self.dim), - (self.batch_size, self.dim), + (self.num_entities,), + (self.batch_size,), + (self.batch_size,), ) scores = self.instance.score_h(all_entities=h, r=r, t=t) self._check_scores(scores=scores, exp_shape=(self.batch_size, self.num_entities)) @@ -81,9 +86,9 @@ def test_score_h(self): def test_score_r(self): """Test score_r.""" h, r, t = self._get_hrt( - (self.batch_size, self.dim), - (self.num_relations, self.dim), - (self.batch_size, self.dim), + (self.batch_size,), + (self.num_relations,), + (self.batch_size,), ) scores = self.instance.score_r(h=h, all_relations=r, t=t) self._check_scores(scores=scores, exp_shape=(self.batch_size, self.num_relations)) @@ -91,9 +96,9 @@ def test_score_r(self): def test_score_t(self): """Test score_t.""" h, r, t = self._get_hrt( - (self.batch_size, self.dim), - (self.batch_size, self.dim), - (self.num_entities, self.dim), + (self.batch_size,), + (self.batch_size,), + (self.num_entities,), ) scores = self.instance.score_t(h=h, r=r, all_entities=t) self._check_scores(scores=scores, exp_shape=(self.batch_size, self.num_entities)) @@ -102,19 +107,19 @@ def test_forward(self): """Test forward.""" for hs, rs, ts in [ [ - (self.batch_size, 1, self.dim), - (1, self.num_relations, self.dim), - (self.batch_size, self.num_entities, self.dim), + (self.batch_size, 1), + (1, self.num_relations), + (self.batch_size, self.num_entities), ], [ - (1, 1, self.dim), - (1, self.num_relations, self.dim), - (self.batch_size, self.num_entities, self.dim), + (1, 1), + (1, self.num_relations), + (self.batch_size, self.num_entities), ], [ - (1, self.num_entities, self.dim), - (1, self.num_relations, self.dim), - (1, self.num_entities, self.dim), + (1, self.num_entities), + (1, self.num_relations), + (1, self.num_entities), ], ]: with self.subTest(f"forward({hs}, {rs}, {ts})"): @@ -125,29 +130,120 @@ def test_forward(self): def test_scores(self): """Test individual scores.""" + self.instance.eval() for i in range(10): - h, r, t = self._get_hrt((1, 1, self.dim), (1, 1, self.dim), (1, 1, self.dim)) + h, r, t = self._get_hrt((1, 1), (1, 1), (1, 1)) scores = self.instance(h=h, r=r, t=t) exp_score = self._exp_score(h, r, t).item() assert scores.item() == exp_score def _exp_score(self, h, r, t) -> torch.FloatTensor: - raise SkipTest() + raise SkipTest("No score check implemented.") + + +class ComplExTests(InteractionTests, unittest.TestCase): + """Tests for ComplEx interaction function.""" + + cls = pykeen.nn.modules.ComplExInteractionFunction + + +class ConvETests(InteractionTests, unittest.TestCase): + """Tests for ConvE interaction function.""" + + cls = pykeen.nn.modules.ConvEInteractionFunction + kwargs = dict( + embedding_height=1, + embedding_width=2, + kernel_height=2, + kernel_width=1, + embedding_dim=InteractionTests.dim, + ) + + def _get_hrt( + self, + *shapes: Tuple[int, ...], + **kwargs + ) -> Tuple[Union[Representation, Sequence[Representation]], ...]: + h, r, t = super()._get_hrt(*shapes, **kwargs) + t_bias = torch.rand_like(t[..., 0, None]) + return h, r, (t, t_bias) + + +class ConvKBTests(InteractionTests, unittest.TestCase): + """Tests for ConvKB interaction function.""" + + cls = pykeen.nn.modules.ConvKBInteractionFunction + kwargs = dict( + embedding_dim=InteractionTests.dim, + num_filters=2 * InteractionTests.dim - 1, + ) class DistMultTests(InteractionTests, unittest.TestCase): """Tests for DistMult interaction function.""" - cls = DistMultInteractionFunction + cls = pykeen.nn.modules.DistMultInteractionFunction def _exp_score(self, h, r, t) -> torch.FloatTensor: return (h * r * t).sum(dim=-1) +class ERMLPTests(InteractionTests, unittest.TestCase): + """Tests for ERMLP interaction function.""" + + cls = pykeen.nn.modules.ERMLPInteractionFunction + kwargs = dict( + embedding_dim=InteractionTests.dim, + hidden_dim=2 * InteractionTests.dim - 1, + ) + + +class ERMLPETests(InteractionTests, unittest.TestCase): + """Tests for ERMLP-E interaction function.""" + + cls = pykeen.nn.modules.ERMLPInteractionFunction + kwargs = dict( + embedding_dim=InteractionTests.dim, + hidden_dim=2 * InteractionTests.dim - 1, + ) + + +class HolETests(InteractionTests, unittest.TestCase): + """Tests for HolE interaction function.""" + + cls = pykeen.nn.modules.HolEInteractionFunction + + +class NTNTests(InteractionTests, unittest.TestCase): + """Tests for NTN interaction function.""" + + cls = pykeen.nn.modules.NTNInteractionFunction + + num_slices: int = 2 + + def _get_hrt( + self, + *shapes: Tuple[int, ...], + ) -> Tuple[Union[Representation, Sequence[Representation]], ...]: + h, r, t = super()._get_hrt(*shapes) + r_shape = r.shape[:-1] + r = [ + torch.rand(*r_shape, *shape) + for shape in [ + (self.num_slices, self.dim, self.dim), + (self.num_slices,), + (self.num_slices,), + (self.num_slices, self.dim), + (self.num_slices, self.dim), + ] + ] + return h, r, t + + class TransETests(InteractionTests, unittest.TestCase): """Tests for TransE interaction function.""" - cls = TransEInteractionFunction + cls = pykeen.nn.modules.TransEInteractionFunction kwargs = dict( p=2, ) From fe0a21fa9ba3d73550aae5dee7d1e5a5e40e0a4e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 17:13:07 +0100 Subject: [PATCH 265/690] Fix tests --- tests/test_interactions.py | 34 ++++++++++++---------------------- 1 file changed, 12 insertions(+), 22 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index efcd52a705..68a9cc7ef1 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -40,16 +40,21 @@ class InteractionTests(GenericTests[pykeen.nn.modules.InteractionFunction]): num_relations: int = 5 num_entities: int = 7 + shape_kwargs = dict() + def _get_hrt( self, *shapes: Tuple[int, ...], - **kwargs ) -> Tuple[Union[Representation, Sequence[Representation]], ...]: - kwargs.setdefault("d", self.dim) - return tuple( - torch.rand(*s, *(kwargs.get(ss) for ss in ms), requires_grad=True) + self.shape_kwargs.setdefault("d", self.dim) + result = tuple( + tuple( + torch.rand(*s, *(self.shape_kwargs[sss] for sss in mss), requires_grad=True) + for mss in ms + ) for s, ms in zip(shapes, [self.cls.entity_shape, self.cls.relation_shape, self.cls.entity_shape]) ) + return [x[0] if len(x) == 1 else x for x in result] def _check_scores(self, scores: torch.FloatTensor, exp_shape: Tuple[int, ...]): """Check shape, dtype and gradients of scores.""" @@ -220,24 +225,9 @@ class NTNTests(InteractionTests, unittest.TestCase): cls = pykeen.nn.modules.NTNInteractionFunction num_slices: int = 2 - - def _get_hrt( - self, - *shapes: Tuple[int, ...], - ) -> Tuple[Union[Representation, Sequence[Representation]], ...]: - h, r, t = super()._get_hrt(*shapes) - r_shape = r.shape[:-1] - r = [ - torch.rand(*r_shape, *shape) - for shape in [ - (self.num_slices, self.dim, self.dim), - (self.num_slices,), - (self.num_slices,), - (self.num_slices, self.dim), - (self.num_slices, self.dim), - ] - ] - return h, r, t + shape_kwargs = dict( + k=2, + ) class TransETests(InteractionTests, unittest.TestCase): From a701affc760402ce48949dfdcd9bc975c9b90b7a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 17:39:05 +0100 Subject: [PATCH 266/690] Some fixes --- src/pykeen/nn/functional.py | 4 +- src/pykeen/nn/modules.py | 36 +++++++----- tests/test_interactions.py | 114 +++++++++++++++++++++++++++++++++--- 3 files changed, 127 insertions(+), 27 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index b0b4f1c04f..5b95d4e397 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -817,7 +817,7 @@ def rescal_interaction( :param h: shape: (batch_size, num_heads, dim) The head representations. - :param r: shape: (batch_size, num_relations, dim ** 2) + :param r: shape: (batch_size, num_relations, dim, dim) The relation representations. :param t: shape: (batch_size, num_tails, dim) The tail representations. @@ -825,7 +825,7 @@ def rescal_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return _extended_einsum("bhd,brde,bte->bhrt", h, r.view(*r.shape[:-1], h.shape[-1], h.shape[-1]), t) + return _extended_einsum("bhd,brde,bte->bhrt", h, r, t) def rotate_interaction( diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index f08e693daa..5dbd5df1a4 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -18,6 +18,10 @@ logger = logging.getLogger(__name__) +def _ensure_tuple(*x: Union[Representation, Sequence[Representation]]) -> Tuple[Sequence[Representation], ...]: + return tuple(xx if isinstance(xx, Sequence) else (xx,) for xx in x) + + class InteractionFunction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation]): """Base class for interaction functions.""" @@ -29,11 +33,11 @@ class InteractionFunction(nn.Module, Generic[HeadRepresentation, RelationReprese TAIL_DIM: int = 3 #: The symbolic shapes for entity representations - entity_shape: Union[str, Sequence[str]] = "d" - tail_entity_shape: Union[None, str, Sequence[str]] = None + entity_shape: Tuple[str, ...] = ("d",) + tail_entity_shape: Union[None, Tuple[str, ...]] = None #: The symbolic shapes for relation representations - relation_shape: Union[str, Sequence[str]] = "d" + relation_shape: Tuple[str, ...] = ("d",) def forward( self, @@ -67,9 +71,9 @@ def _add_dim(*x: torch.FloatTensor, dim: int) -> Sequence[torch.FloatTensor]: The tensor with batch dimension. """ out = [xx.unsqueeze(dim=dim) for xx in x] - if len(x) > 1: - return out - return out[0] + if len(x) == 1: + return out[0] + return out @staticmethod def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: @@ -132,9 +136,6 @@ def _check_shapes( ((tt, t_prefix + ts) for tt, ts in zip(t, tail_entity_shape)), ), raise_or_error=raise_on_error) - def _ensure_tuple(self, *x: Union[Representation, Sequence[Representation]]) -> Tuple[Sequence[Representation], ...]: - return tuple(xx if isinstance(xx, Sequence) else (xx,) for xx in x) - def score_hrt( self, h: HeadRepresentation = tuple(), @@ -154,7 +155,7 @@ def score_hrt( :return: shape: (batch_size, 1) The scores. """ - h, r, t = self._ensure_tuple(h, r, t) + h, r, t = _ensure_tuple(h, r, t) assert self._check_shapes(h=h, r=r, t=t) # prepare input to generic score function @@ -187,7 +188,7 @@ def score_h( :return: shape: (batch_size, num_entities) The scores. """ - h, r, t = self._ensure_tuple(all_entities, r, t) + h, r, t = _ensure_tuple(all_entities, r, t) assert self._check_shapes(h=h, r=r, t=t, h_prefix="n") # prepare input to generic score function @@ -220,7 +221,7 @@ def score_r( :return: shape: (batch_size, num_entities) The scores. """ - h, r, t = self._ensure_tuple(h, all_relations, t) + h, r, t = _ensure_tuple(h, all_relations, t) assert self._check_shapes(h=h, r=r, t=t, r_prefix="n") # prepare input to generic score function @@ -253,7 +254,7 @@ def score_t( :return: shape: (batch_size, num_entities) The scores. """ - h, r, t = self._ensure_tuple(h, r, all_entities) + h, r, t = _ensure_tuple(h, r, all_entities) assert self._check_shapes(h=h, r=r, t=t, t_prefix="n") # prepare input to generic score function @@ -290,7 +291,7 @@ def forward( t: TailRepresentation, ) -> torch.FloatTensor: # noqa: D102 # normalization - h, r, t = [(x,) if torch.is_tensor(x) else x for x in (h, r, t)] + h, r, t = _ensure_tuple(h, r, t) return self.f(*h, *r, *t) @@ -706,7 +707,7 @@ def forward( class RESCALInteractionFunction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): """Interaction function of RESCAL.""" - relation_shape = "dd" + relation_shape = ("dd",) def __init__(self): super().__init__(f=pkf.rescal_interaction) @@ -715,7 +716,7 @@ def __init__(self): class StructuredEmbeddingInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor], torch.FloatTensor]): """Interaction function of Structured Embedding.""" - relation_shape = "dd" + relation_shape = ("dd", "dd") def forward( self, @@ -875,6 +876,9 @@ class KG2EInteractionFunction( ): """Interaction function of KG2E.""" + entity_shape = ("d", "d") + relation_shape = ("d", "d") + def __init__( self, similarity: str = "KL", diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 68a9cc7ef1..792b138016 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -32,6 +32,10 @@ def post_instantiation_hook(self) -> None: """Perform actions after instantiation.""" +def _unpack_singletons(*xs: Tuple) -> Sequence[Tuple]: + return [x[0] if len(x) == 1 else x for x in xs] + + class InteractionTests(GenericTests[pykeen.nn.modules.InteractionFunction]): """Generic test for interaction functions.""" @@ -49,12 +53,15 @@ def _get_hrt( self.shape_kwargs.setdefault("d", self.dim) result = tuple( tuple( - torch.rand(*s, *(self.shape_kwargs[sss] for sss in mss), requires_grad=True) - for mss in ms + torch.rand(*prefix_shape, *(self.shape_kwargs[dim] for dim in weight_shape), requires_grad=True) + for weight_shape in weight_shapes + ) + for prefix_shape, weight_shapes in zip( + shapes, + [self.cls.entity_shape, self.cls.relation_shape, self.cls.entity_shape], ) - for s, ms in zip(shapes, [self.cls.entity_shape, self.cls.relation_shape, self.cls.entity_shape]) ) - return [x[0] if len(x) == 1 else x for x in result] + return tuple(_unpack_singletons(*result)) def _check_scores(self, scores: torch.FloatTensor, exp_shape: Tuple[int, ...]): """Check shape, dtype and gradients of scores.""" @@ -96,7 +103,11 @@ def test_score_r(self): (self.batch_size,), ) scores = self.instance.score_r(h=h, all_relations=r, t=t) - self._check_scores(scores=scores, exp_shape=(self.batch_size, self.num_relations)) + if len(self.cls.relation_shape) == 0: + exp_shape = (self.batch_size, 1) + else: + exp_shape = (self.batch_size, self.num_relations) + self._check_scores(scores=scores, exp_shape=exp_shape) def test_score_t(self): """Test score_t.""" @@ -127,8 +138,13 @@ def test_forward(self): (1, self.num_entities), ], ]: + if any(isinstance(m, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d)) for m in self.instance.modules()): + continue with self.subTest(f"forward({hs}, {rs}, {ts})"): - expected_shape = (max(hs[0], rs[0], ts[0]), hs[1], rs[1], ts[1]) + expected_shape = [max(hs[0], rs[0], ts[0]), hs[1], rs[1], ts[1]] + if len(self.cls.relation_shape) == 0: + expected_shape[2] = 1 + expected_shape = tuple(expected_shape) h, r, t = self._get_hrt(hs, rs, ts) scores = self.instance(h=h, r=r, t=t) self._check_scores(scores=scores, exp_shape=expected_shape) @@ -230,10 +246,45 @@ class NTNTests(InteractionTests, unittest.TestCase): ) -class TransETests(InteractionTests, unittest.TestCase): - """Tests for TransE interaction function.""" +class ProjETests(InteractionTests, unittest.TestCase): + """Tests for ProjE interaction function.""" + + cls = pykeen.nn.modules.ProjEInteractionFunction + kwargs = dict( + embedding_dim=InteractionTests.dim, + ) + + +class RESCALTests(InteractionTests, unittest.TestCase): + """Tests for RESCAL interaction function.""" + + cls = pykeen.nn.modules.RESCALInteractionFunction + + +class KG2ETests(InteractionTests, unittest.TestCase): + """Tests for KG2E interaction function.""" + + cls = pykeen.nn.modules.KG2EInteractionFunction + + +class TuckerTests(InteractionTests, unittest.TestCase): + """Tests for Tucker interaction function.""" + + cls = pykeen.nn.modules.TuckerInteractionFunction + kwargs = dict( + embedding_dim=InteractionTests.dim, + ) + + +class RotatETests(InteractionTests, unittest.TestCase): + """Tests for RotatE interaction function.""" + + cls = pykeen.nn.modules.RotatEInteractionFunction + + +class TranslationalInteractionTests(InteractionTests): + """Common tests for translational interaction.""" - cls = pykeen.nn.modules.TransEInteractionFunction kwargs = dict( p=2, ) @@ -241,5 +292,50 @@ class TransETests(InteractionTests, unittest.TestCase): def _additional_score_checks(self, scores): assert (scores <= 0).all() + +class TransDTests(TranslationalInteractionTests, unittest.TestCase): + """Tests for TransD interaction function.""" + + cls = pykeen.nn.modules.TransDInteractionFunction + shape_kwargs = dict( + e=3, + ) + + +class TransETests(TranslationalInteractionTests, unittest.TestCase): + """Tests for TransE interaction function.""" + + cls = pykeen.nn.modules.TransEInteractionFunction + def _exp_score(self, h, r, t) -> torch.FloatTensor: return -(h + r - t).norm(p=2, dim=-1) + + +# class TransHTests(TranslationalInteractionTests, unittest.TestCase): +# """Tests for TransH interaction function.""" +# +# cls = pykeen.nn.modules.TransHInteractionFunction +# # shape_kwargs = dict( +# # e=3, +# # ) + + +class TransRTests(TranslationalInteractionTests, unittest.TestCase): + """Tests for TransR interaction function.""" + + cls = pykeen.nn.modules.TransRInteractionFunction + shape_kwargs = dict( + e=3, + ) + + +class SETests(TranslationalInteractionTests, unittest.TestCase): + """Tests for SE interaction function.""" + + cls = pykeen.nn.modules.StructuredEmbeddingInteractionFunction + + +class UMTests(TranslationalInteractionTests, unittest.TestCase): + """Tests for UM interaction function.""" + + cls = pykeen.nn.modules.UnstructuredModelInteractionFunction From 1a54a9cc132233262a92d6747350e6b15da0bfe6 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 17:44:41 +0100 Subject: [PATCH 267/690] Subclass TransH from Model instead of EntityRelationModel --- src/pykeen/models/unimodal/trans_h.py | 69 ++++++++++++--------------- src/pykeen/nn/modules.py | 19 ++++++++ 2 files changed, 49 insertions(+), 39 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index b35c78c832..fdcf6a0663 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -4,13 +4,13 @@ from typing import Optional -import torch from torch.nn import functional -from ..base import EntityRelationEmbeddingModel +from .. import Model from ...losses import Loss -from ...nn import Embedding, functional as pkf +from ...nn import Embedding from ...nn.emb import EmbeddingSpecification +from ...nn.modules import TransHInteractionFunction from ...regularizers import Regularizer, TransHRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -20,7 +20,7 @@ ] -class TransH(EntityRelationEmbeddingModel): +class TransH(Model): r"""An implementation of TransH [wang2014]_. This model extends :class:`pykeen.models.TransE` by applying the translation from head to tail entity in a @@ -70,6 +70,7 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 2, loss: Optional[Loss] = None, + predict_with_sigmoid: bool = False, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, @@ -79,47 +80,37 @@ def __init__( :param embedding_dim: The entity embedding dimension $d$. Is usually $d \in [50, 300]$. :param scoring_fct_norm: The :math:`l_p` norm applied in the interaction function. Is usually ``1`` or ``2.``. """ + embedding_dim = embedding_dim super().__init__( triples_factory=triples_factory, - embedding_dim=embedding_dim, + interaction_function=TransHInteractionFunction( + p=scoring_fct_norm, + power_norm=False, + ), automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - ) - self.scoring_fct_norm = scoring_fct_norm - - # embeddings - self.normal_vector_embeddings = Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim, - specification=EmbeddingSpecification( - # Normalise the normal vectors by their l2 norms - constrainer=functional.normalize, + predict_with_sigmoid=predict_with_sigmoid, + entity_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + specification=None, ), + relation_representations=[ + Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=embedding_dim, + specification=None, + ), + Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=embedding_dim, + specification=EmbeddingSpecification( + # Normalise the normal vectors by their l2 norms + constrainer=functional.normalize, + ), + ), + ], ) - - def regularize_if_necessary(self) -> None: - """Update the regularizer's term given some tensors, if regularization is requested.""" - # As described in [wang2014], all entities and relations are used to compute the regularization term - # which enforces the defined soft constraints. - super().regularize_if_necessary( - self.entity_embeddings(indices=None), - self.normal_vector_embeddings(indices=None), # FIXME - self.relation_embeddings(indices=None), - ) - - def forward( - self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], - ) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - d_r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) - w_r = self.normal_vector_embeddings.get_in_canonical_shape(indices=r_indices) - t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - self.regularize_if_necessary() - return pkf.transh_interaction(h, w_r, d_r, t, p=self.scoring_fct_norm) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 5dbd5df1a4..3d10775e5d 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -910,3 +910,22 @@ def forward( similarity=self.similarity, exact=self.exact, ) + + +class TransHInteractionFunction( + TranslationalInteractionFunction[ + torch.FloatTensor, + Tuple[torch.FloatTensor, torch.FloatTensor], + torch.FloatTensor, + ] +): + """Interaction function of TransH.""" + + def forward( + self, + h: torch.FloatTensor, + r: Tuple[torch.FloatTensor, torch.FloatTensor], + t: torch.FloatTensor, + ) -> torch.FloatTensor: # noqa: D102 + w_r, d_r = r + return pkf.transh_interaction(h, w_r, d_r, t, p=self.scoring_fct_norm) From 00c44fce276f5aa89c1610f71647f28557017695 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 17:45:36 +0100 Subject: [PATCH 268/690] Add TransH interaction test and fix --- src/pykeen/nn/modules.py | 4 +++- tests/test_interactions.py | 11 ++++------- 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 3d10775e5d..c0be04b5d2 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -921,6 +921,8 @@ class TransHInteractionFunction( ): """Interaction function of TransH.""" + relation_shape = ("d", "d") + def forward( self, h: torch.FloatTensor, @@ -928,4 +930,4 @@ def forward( t: torch.FloatTensor, ) -> torch.FloatTensor: # noqa: D102 w_r, d_r = r - return pkf.transh_interaction(h, w_r, d_r, t, p=self.scoring_fct_norm) + return pkf.transh_interaction(h, w_r, d_r, t, p=self.p, power_norm=self.power_norm) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 792b138016..169815e63c 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -311,13 +311,10 @@ def _exp_score(self, h, r, t) -> torch.FloatTensor: return -(h + r - t).norm(p=2, dim=-1) -# class TransHTests(TranslationalInteractionTests, unittest.TestCase): -# """Tests for TransH interaction function.""" -# -# cls = pykeen.nn.modules.TransHInteractionFunction -# # shape_kwargs = dict( -# # e=3, -# # ) +class TransHTests(TranslationalInteractionTests, unittest.TestCase): + """Tests for TransH interaction function.""" + + cls = pykeen.nn.modules.TransHInteractionFunction class TransRTests(TranslationalInteractionTests, unittest.TestCase): From 3eb308eff7cc2bced3eeb5023767149da7f95577 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 17:48:50 +0100 Subject: [PATCH 269/690] Do not subclass two-side model from entity-relation model --- src/pykeen/models/base.py | 45 ++++++++++++++++++++++++--------------- 1 file changed, 28 insertions(+), 17 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 3c9f7bf0fe..79eca3348d 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1377,7 +1377,7 @@ def __init__( ) -class TwoSideEmbeddingModel(EntityRelationEmbeddingModel): +class TwoSideEmbeddingModel(Model): """A model which averages scores for forward and backward model.""" def __init__( @@ -1397,31 +1397,42 @@ def __init__( second_embedding_specification: Optional[EmbeddingSpecification] = None, second_relation_embedding_specification: Optional[EmbeddingSpecification] = None, ): + if relation_dim is None: + relation_dim = embedding_dim super().__init__( triples_factory=triples_factory, - embedding_dim=embedding_dim, - relation_dim=relation_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - embedding_specification=embedding_specification, - relation_embedding_specification=relation_embedding_specification, - ) - # extra embeddings - self.second_entity_embeddings = Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - embedding_dim=self.embedding_dim, - specification=second_embedding_specification, - ) - self.second_relation_embeddings = Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - embedding_dim=self.relation_dim, - specification=second_relation_embedding_specification, + entity_representations=[ + Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + specification=embedding_specification, + ), + Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + specification=second_embedding_specification, + ), + ], + relation_representations=[ + Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=relation_dim, + specification=relation_embedding_specification, + ), + Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=relation_dim, + specification=second_relation_embedding_specification, + ), + ], + interaction_function=interaction_function, ) - self.interaction_function = interaction_function def forward( self, From 8faa1178fef4463b377f287104cd2b79dbaca80f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 17:56:17 +0100 Subject: [PATCH 270/690] Extract base class for models with two relation embeddings --- src/pykeen/models/base.py | 56 +++++++++++++++++++ .../models/unimodal/structured_embedding.py | 48 ++++------------ src/pykeen/models/unimodal/trans_h.py | 31 +++------- src/pykeen/models/unimodal/trans_r.py | 50 +++++------------ 4 files changed, 89 insertions(+), 96 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 79eca3348d..22118e1974 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1315,6 +1315,62 @@ def __init__( ) +class DoubleRelationEmbeddingModel(Model, ABC): + """A model with one vector for each entity and two vectors for each relation.""" + + def __init__( + self, + triples_factory: TriplesFactory, + interaction_function: InteractionFunction[ + torch.FloatTensor, + Tuple[torch.FloatTensor, torch.FloatTensor], + torch.FloatTensor, + ], + embedding_dim: int = 50, + relation_dim: Optional[int] = None, + loss: Optional[Loss] = None, + predict_with_sigmoid: bool = False, + automatic_memory_optimization: Optional[bool] = None, + preferred_device: DeviceHint = None, + random_seed: Optional[int] = None, + regularizer: Optional[Regularizer] = None, + embedding_specification: Optional[EmbeddingSpecification] = None, + relation_embedding_specification: Optional[EmbeddingSpecification] = None, + second_relation_embedding_specification: Optional[EmbeddingSpecification] = None, + ) -> None: + if relation_dim is None: + relation_dim = embedding_dim + super().__init__( + triples_factory=triples_factory, + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + predict_with_sigmoid=predict_with_sigmoid, + preferred_device=preferred_device, + random_seed=random_seed, + regularizer=regularizer, + entity_representations=[ + Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + specification=embedding_specification, + ), + ], + relation_representations=[ + Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=relation_dim, + specification=relation_embedding_specification, + ), + Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=relation_dim, + specification=second_relation_embedding_specification, + ) + ], + interaction_function=interaction_function, + ) + + class TwoVectorEmbeddingModel(Model, ABC): """A model with two vectors for each entity and relation.""" diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index e3169c553a..241ed983e4 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -6,11 +6,10 @@ from typing import Optional import numpy as np -import torch from torch import nn from torch.nn import functional -from .. import EntityRelationEmbeddingModel +from ..base import DoubleRelationEmbeddingModel from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ @@ -25,7 +24,7 @@ ] -class StructuredEmbedding(EntityRelationEmbeddingModel): +class StructuredEmbedding(DoubleRelationEmbeddingModel): r"""An implementation of the Structured Embedding (SE) published by [bordes2011]_. SE applies role- and relation-specific projection matrices @@ -72,8 +71,12 @@ def __init__( ) super().__init__( triples_factory=triples_factory, + interaction_function=StructuredEmbeddingInteractionFunction( + p=scoring_fct_norm, + power_norm=False, + ), embedding_dim=embedding_dim, - relation_dim=embedding_dim ** 2, # head projection matrices + relation_dim=embedding_dim ** 2, automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, @@ -86,38 +89,7 @@ def __init__( relation_embedding_specification=EmbeddingSpecification( initializer=relation_initializer, ), + second_relation_embedding_specification=EmbeddingSpecification( + initializer=relation_initializer, + ) ) - self.second_relation_embedding = EmbeddingSpecification( - initializer=relation_initializer, - ).make( - num_embeddings=self.num_relations, - embedding_dim=self.relation_dim, - device=self.device, - ) - self.interaction_function = StructuredEmbeddingInteractionFunction( - p=scoring_fct_norm, - power_norm=False, - ) - - def forward( - self, - h_indices: Optional[torch.LongTensor] = None, - r_indices: Optional[torch.LongTensor] = None, - t_indices: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: - """Evaluate the given triples. - - :param h_indices: shape: (batch_size,) - The indices for head entities. If None, score against all. - :param r_indices: shape: (batch_size,) - The indices for relations. If None, score against all. - :param t_indices: shape: (batch_size,) - The indices for tail entities. If None, score against all. - - :return: The scores, shape: (batch_size, num_entities) - """ - h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - r_h = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) - r_t = self.second_relation_embedding.get_in_canonical_shape(indices=r_indices) - t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - return self.interaction_function(h=h, r=(r_h, r_t), t=t) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index fdcf6a0663..1d69fba0d6 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -6,9 +6,8 @@ from torch.nn import functional -from .. import Model +from ..base import DoubleRelationEmbeddingModel from ...losses import Loss -from ...nn import Embedding from ...nn.emb import EmbeddingSpecification from ...nn.modules import TransHInteractionFunction from ...regularizers import Regularizer, TransHRegularizer @@ -20,7 +19,7 @@ ] -class TransH(Model): +class TransH(DoubleRelationEmbeddingModel): r"""An implementation of TransH [wang2014]_. This model extends :class:`pykeen.models.TransE` by applying the translation from head to tail entity in a @@ -80,37 +79,23 @@ def __init__( :param embedding_dim: The entity embedding dimension $d$. Is usually $d \in [50, 300]$. :param scoring_fct_norm: The :math:`l_p` norm applied in the interaction function. Is usually ``1`` or ``2.``. """ - embedding_dim = embedding_dim super().__init__( triples_factory=triples_factory, interaction_function=TransHInteractionFunction( p=scoring_fct_norm, power_norm=False, ), + embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, predict_with_sigmoid=predict_with_sigmoid, - entity_representations=Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, - specification=None, + embedding_specification=None, + relation_embedding_specification=None, + second_relation_embedding_specification=EmbeddingSpecification( + # Normalise the normal vectors by their l2 norms + constrainer=functional.normalize, ), - relation_representations=[ - Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim, - specification=None, - ), - Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim, - specification=EmbeddingSpecification( - # Normalise the normal vectors by their l2 norms - constrainer=functional.normalize, - ), - ), - ], ) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index fc0e387025..489491e7fe 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -1,7 +1,7 @@ # -*- coding: utf-8 -*- """Implementation of TransR.""" - +import logging from functools import partial from typing import Optional @@ -10,9 +10,8 @@ import torch.nn.init from torch.nn import functional -from ..base import EntityRelationEmbeddingModel +from ..base import DoubleRelationEmbeddingModel from ...losses import Loss -from ...nn import Embedding from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import TransRInteractionFunction @@ -36,7 +35,7 @@ def _projection_initializer( return torch.nn.init.xavier_uniform_(x.view(num_relations, embedding_dim, relation_dim)).view(x.shape) -class TransR(EntityRelationEmbeddingModel): +class TransR(DoubleRelationEmbeddingModel): r"""An implementation of TransR from [lin2015]_. TransR is an extension of :class:`pykeen.models.TransH` that explicitly considers entities and relations as @@ -109,36 +108,17 @@ def __init__( constrainer=clamp_norm, constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), - - ) - self.interaction_function = TransRInteractionFunction(p=scoring_fct_norm) - - # TODO: Initialize from TransE - - # embeddings - self.relation_projections = Embedding.init_with_device( - num_embeddings=triples_factory.num_relations, - embedding_dim=relation_dim * embedding_dim, - device=self.device, - initializer=partial( - _projection_initializer, - num_relations=self.num_relations, - embedding_dim=self.embedding_dim, - relation_dim=self.relation_dim, + # Relation projections + second_relation_embedding_specification=EmbeddingSpecification( + initializer=partial( + _projection_initializer, + num_relations=triples_factory.num_relations, + embedding_dim=embedding_dim, + relation_dim=relation_dim, + ), + ), + interaction_function=TransRInteractionFunction( + p=scoring_fct_norm, ), ) - - def forward( - self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], - ) -> torch.FloatTensor: # noqa: D102 - h = self.entity_embeddings.get_in_canonical_shape(indices=h_indices) - r = self.relation_embeddings.get_in_canonical_shape(indices=r_indices) - t = self.entity_embeddings.get_in_canonical_shape(indices=t_indices) - m_r = self.relation_projections.get_in_canonical_shape( - indices=r_indices, - reshape_dim=(self.embedding_dim, self.relation_dim), - ) - return self.interaction_function(h=h, r=(r, m_r), t=t) + logging.warning("Initialize from TransE") From 55a79787bb97bf7c73628d5c87824c5c39309a41 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 18:02:06 +0100 Subject: [PATCH 271/690] Add shape to Embedding --- src/pykeen/nn/emb.py | 69 ++++++++++++---------------------------- src/pykeen/nn/modules.py | 2 +- 2 files changed, 22 insertions(+), 49 deletions(-) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index d78b027422..d47ca66406 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -5,12 +5,12 @@ import functools from typing import Any, Mapping, Optional, Sequence +import numpy import torch import torch.nn from torch import nn -from ..typing import Constrainer, DeviceHint, Initializer, Normalizer -from ..utils import resolve_device +from ..typing import Constrainer, Initializer, Normalizer __all__ = [ 'RepresentationModule', @@ -61,18 +61,23 @@ class EmbeddingSpecification: # regularizer: Optional[Regularizer] = None - def make(self, num_embeddings: int, embedding_dim: int, device: DeviceHint) -> 'Embedding': + def make( + self, + num_embeddings: int, + embedding_dim: int, + shape: Optional[Sequence[int]], + ) -> 'Embedding': """Create an embedding with this specification.""" - return Embedding.init_with_device( + return Embedding( num_embeddings=num_embeddings, embedding_dim=embedding_dim, + shape=shape, initializer=self.initializer, initializer_kwargs=self.initializer_kwargs, normalizer=self.normalizer, normalizer_kwargs=self.normalizer_kwargs, constrainer=self.constrainer, constrainer_kwargs=self.constrainer_kwargs, - device=device, ) @@ -87,6 +92,7 @@ def __init__( self, num_embeddings: int, embedding_dim: int, + shape: Optional[Sequence[int]] = None, initializer: Optional[Initializer] = None, initializer_kwargs: Optional[Mapping[str, Any]] = None, normalizer: Optional[Normalizer] = None, @@ -118,6 +124,11 @@ def __init__( Additional keyword arguments passed to the constrainer """ super().__init__() + if shape is not None: + embedding_dim = numpy.prod(shape) + else: + shape = (embedding_dim,) + self.shape = shape if initializer is None: initializer = nn.init.normal_ @@ -143,8 +154,8 @@ def from_specification( cls, num_embeddings: int, embedding_dim: int, - specification: Optional[EmbeddingSpecification], - device: DeviceHint = None, + shape: Optional[Sequence[int]] = None, + specification: Optional[EmbeddingSpecification] = None, ) -> 'Embedding': """Create an embedding based on a specification. @@ -154,9 +165,6 @@ def from_specification( The embedding dimension. :param specification: The specification. - :param device: - If given, move to device. - :return: An embedding object. """ @@ -165,46 +173,9 @@ def from_specification( return specification.make( num_embeddings=num_embeddings, embedding_dim=embedding_dim, - device=device, + shape=shape, ) - @classmethod - def init_with_device( - cls, - num_embeddings: int, - embedding_dim: int, - device: DeviceHint, - initializer: Optional[Initializer] = None, - initializer_kwargs: Optional[Mapping[str, Any]] = None, - normalizer: Optional[Normalizer] = None, - normalizer_kwargs: Optional[Mapping[str, Any]] = None, - constrainer: Optional[Constrainer] = None, - constrainer_kwargs: Optional[Mapping[str, Any]] = None, - ) -> 'Embedding': # noqa:E501 - """Create an embedding object on the given device by wrapping :func:`__init__`. - - This method is a hotfix for not being able to pass a device during initialization of - :class:`torch.nn.Embedding`. Instead the weight is always initialized on CPU and has - to be moved to GPU afterwards. - - .. seealso:: - - https://developer.nvidia.com/gpugems/gpugems3/part-vi-gpu-computing/chapter-37-efficient-random-number-generation-and-application - - :return: - The embedding. - """ - return cls( - num_embeddings=num_embeddings, - embedding_dim=embedding_dim, - initializer=initializer, - initializer_kwargs=initializer_kwargs, - normalizer=normalizer, - normalizer_kwargs=normalizer_kwargs, - constrainer=constrainer, - constrainer_kwargs=constrainer_kwargs, - ).to(device=resolve_device(device)) - @property def num_embeddings(self) -> int: # noqa: D401 """The total number of representations (i.e. the maximum ID).""" @@ -255,6 +226,8 @@ def get_in_canonical_shape( x = x.unsqueeze(dim=0) else: x = x.unsqueeze(dim=1) + if len(self.shape) > 1 and reshape_dim is None: + reshape_dim = self.shape if reshape_dim is not None: x = x.view(*x.shape[:-1], *reshape_dim) return x diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index c0be04b5d2..78963f807c 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -19,7 +19,7 @@ def _ensure_tuple(*x: Union[Representation, Sequence[Representation]]) -> Tuple[Sequence[Representation], ...]: - return tuple(xx if isinstance(xx, Sequence) else (xx,) for xx in x) + return c class InteractionFunction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation]): From 33d1645511768dd467dab3f9b860d60238e406da Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 18:03:55 +0100 Subject: [PATCH 272/690] Use shape for NTN matrix embeddings --- src/pykeen/models/unimodal/ntn.py | 9 +++------ src/pykeen/nn/emb.py | 4 +++- 2 files changed, 6 insertions(+), 7 deletions(-) diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 7ff26b8598..2d7ee08739 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -71,12 +71,9 @@ def __init__( :param non_linearity: A non-linear activation function. Defaults to the hyperbolic tangent :class:`torch.nn.Tanh`. """ - self.embedding_dim = embedding_dim - self.num_slices = num_slices - w = Embedding( num_embeddings=triples_factory.num_relations, - embedding_dim=num_slices * embedding_dim ** 2, + shape=(num_slices, embedding_dim, embedding_dim) ) b = Embedding( num_embeddings=triples_factory.num_relations, @@ -88,11 +85,11 @@ def __init__( ) vh = Embedding( num_embeddings=triples_factory.num_relations, - embedding_dim=num_slices * embedding_dim, + shape=(num_slices, embedding_dim), ) vt = Embedding( num_embeddings=triples_factory.num_relations, - embedding_dim=num_slices * embedding_dim, + shape=(num_slices, embedding_dim), ) super().__init__( triples_factory=triples_factory, diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index d47ca66406..c9a579a169 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -91,7 +91,7 @@ class Embedding(RepresentationModule): def __init__( self, num_embeddings: int, - embedding_dim: int, + embedding_dim: Optional[int] = None, shape: Optional[Sequence[int]] = None, initializer: Optional[Initializer] = None, initializer_kwargs: Optional[Mapping[str, Any]] = None, @@ -124,6 +124,8 @@ def __init__( Additional keyword arguments passed to the constrainer """ super().__init__() + if shape is None and embedding_dim is None: + raise ValueError if shape is not None: embedding_dim = numpy.prod(shape) else: From 474edfac611e6df935e4cdf58baa89e884cdbb6a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 18:09:37 +0100 Subject: [PATCH 273/690] Specify shape for more than 1-d embeddings --- src/pykeen/models/base.py | 30 +++++++++++++------ src/pykeen/models/unimodal/rescal.py | 2 +- .../models/unimodal/structured_embedding.py | 2 +- src/pykeen/nn/emb.py | 14 +++++---- 4 files changed, 31 insertions(+), 17 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 22118e1974..e339dc9ea2 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1258,7 +1258,7 @@ def __init__( triples_factory: TriplesFactory, interaction_function: InteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], embedding_dim: int = 200, - relation_dim: Optional[int] = None, + relation_dim: Union[None, int, Sequence[int]] = None, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, @@ -1291,8 +1291,6 @@ def __init__( # Default for relation dimensionality if relation_dim is None: relation_dim = embedding_dim - self.embedding_dim = embedding_dim - self.relation_dim = relation_dim super().__init__( triples_factory=triples_factory, automatic_memory_optimization=automatic_memory_optimization, @@ -1304,16 +1302,30 @@ def __init__( interaction_function=interaction_function, entity_representations=Embedding.from_specification( num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, + shape=embedding_dim, specification=embedding_specification, ), relation_representations=Embedding.from_specification( num_embeddings=triples_factory.num_relations, - embedding_dim=relation_dim, + shape=relation_dim, specification=relation_embedding_specification, ), ) + @property + def embedding_dim(self) -> int: + """The entity embedding dim.""" + embedding = self.entity_representations[0] + assert isinstance(embedding, Embedding) + return embedding.embedding_dim + + @property + def relation_dim(self) -> int: + """The relation embedding dim.""" + embedding = self.relation_representations[0] + assert isinstance(embedding, Embedding) + return embedding.embedding_dim + class DoubleRelationEmbeddingModel(Model, ABC): """A model with one vector for each entity and two vectors for each relation.""" @@ -1327,7 +1339,7 @@ def __init__( torch.FloatTensor, ], embedding_dim: int = 50, - relation_dim: Optional[int] = None, + relation_dim: Union[None, int, Sequence[int]] = None, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, automatic_memory_optimization: Optional[bool] = None, @@ -1351,19 +1363,19 @@ def __init__( entity_representations=[ Embedding.from_specification( num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, + shape=embedding_dim, specification=embedding_specification, ), ], relation_representations=[ Embedding.from_specification( num_embeddings=triples_factory.num_relations, - embedding_dim=relation_dim, + shape=relation_dim, specification=relation_embedding_specification, ), Embedding.from_specification( num_embeddings=triples_factory.num_relations, - embedding_dim=relation_dim, + shape=relation_dim, specification=second_relation_embedding_specification, ) ], diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index 6fd639130a..6fbd574692 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -72,7 +72,7 @@ def __init__( triples_factory=triples_factory, interaction_function=RESCALInteractionFunction(), embedding_dim=embedding_dim, - relation_dim=embedding_dim ** 2, # d x d matrices + relation_dim=(embedding_dim, embedding_dim), automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index 241ed983e4..f2af434609 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -76,7 +76,7 @@ def __init__( power_norm=False, ), embedding_dim=embedding_dim, - relation_dim=embedding_dim ** 2, + relation_dim=(embedding_dim, embedding_dim), automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index c9a579a169..b3f0b8e853 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -3,7 +3,7 @@ """Embedding modules.""" import dataclasses import functools -from typing import Any, Mapping, Optional, Sequence +from typing import Any, Mapping, Optional, Sequence, Union import numpy import torch @@ -64,8 +64,8 @@ class EmbeddingSpecification: def make( self, num_embeddings: int, - embedding_dim: int, - shape: Optional[Sequence[int]], + embedding_dim: Optional[int], + shape: Optional[Union[int, Sequence[int]]] ) -> 'Embedding': """Create an embedding with this specification.""" return Embedding( @@ -92,7 +92,7 @@ def __init__( self, num_embeddings: int, embedding_dim: Optional[int] = None, - shape: Optional[Sequence[int]] = None, + shape: Optional[Union[int, Sequence[int]]] = None, initializer: Optional[Initializer] = None, initializer_kwargs: Optional[Mapping[str, Any]] = None, normalizer: Optional[Normalizer] = None, @@ -127,6 +127,8 @@ def __init__( if shape is None and embedding_dim is None: raise ValueError if shape is not None: + if not isinstance(shape, Sequence): + shape = (shape,) embedding_dim = numpy.prod(shape) else: shape = (embedding_dim,) @@ -155,8 +157,8 @@ def __init__( def from_specification( cls, num_embeddings: int, - embedding_dim: int, - shape: Optional[Sequence[int]] = None, + embedding_dim: Optional[int] = None, + shape: Optional[Union[int, Sequence[int]]] = None, specification: Optional[EmbeddingSpecification] = None, ) -> 'Embedding': """Create an embedding based on a specification. From f87706f74c148d5fd20c14001eab681491554639 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 18:10:43 +0100 Subject: [PATCH 274/690] Fix typo --- src/pykeen/nn/modules.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 78963f807c..c0be04b5d2 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -19,7 +19,7 @@ def _ensure_tuple(*x: Union[Representation, Sequence[Representation]]) -> Tuple[Sequence[Representation], ...]: - return c + return tuple(xx if isinstance(xx, Sequence) else (xx,) for xx in x) class InteractionFunction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation]): From e31264b943077dc61220d335f307087ce9569739 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 18:20:41 +0100 Subject: [PATCH 275/690] Make TransR directly subclass from Model --- src/pykeen/models/unimodal/trans_r.py | 70 ++++++++++++--------------- src/pykeen/nn/emb.py | 4 +- tests/test_models.py | 10 ++-- 3 files changed, 41 insertions(+), 43 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 489491e7fe..89f9876cbe 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -2,16 +2,13 @@ """Implementation of TransR.""" import logging -from functools import partial from typing import Optional -import torch -import torch.autograd -import torch.nn.init from torch.nn import functional -from ..base import DoubleRelationEmbeddingModel +from .. import Model from ...losses import Loss +from ...nn import Embedding from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import TransRInteractionFunction @@ -25,17 +22,7 @@ ] -def _projection_initializer( - x: torch.FloatTensor, - num_relations: int, - embedding_dim: int, - relation_dim: int, -) -> torch.FloatTensor: - """Initialize by Glorot.""" - return torch.nn.init.xavier_uniform_(x.view(num_relations, embedding_dim, relation_dim)).view(x.shape) - - -class TransR(DoubleRelationEmbeddingModel): +class TransR(Model): r"""An implementation of TransR from [lin2015]_. TransR is an extension of :class:`pykeen.models.TransH` that explicitly considers entities and relations as @@ -88,35 +75,42 @@ def __init__( """Initialize the model.""" super().__init__( triples_factory=triples_factory, - embedding_dim=embedding_dim, - relation_dim=relation_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - embedding_specification=EmbeddingSpecification( - initializer=xavier_uniform_, - constrainer=clamp_norm, - constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + entity_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + shape=embedding_dim, + specification=EmbeddingSpecification( + initializer=xavier_uniform_, + constrainer=clamp_norm, + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + ) ), - relation_embedding_specification=EmbeddingSpecification( - initializer=compose( - xavier_uniform_, - functional.normalize, + relation_representations=[ + Embedding.from_specification( + triples_factory.num_relations, + shape=relation_dim, + specification=EmbeddingSpecification( + initializer=compose( + xavier_uniform_, + functional.normalize, + ), + constrainer=clamp_norm, + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + ), ), - constrainer=clamp_norm, - constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), - ), - # Relation projections - second_relation_embedding_specification=EmbeddingSpecification( - initializer=partial( - _projection_initializer, - num_relations=triples_factory.num_relations, - embedding_dim=embedding_dim, - relation_dim=relation_dim, - ), - ), + # Relation projections + Embedding.from_specification( + triples_factory.num_relations, + shape=(relation_dim, embedding_dim), + specification=EmbeddingSpecification( + initializer=xavier_uniform_, + ), + ) + ], interaction_function=TransRInteractionFunction( p=scoring_fct_norm, ), diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index b3f0b8e853..b19883d29b 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -192,7 +192,9 @@ def embedding_dim(self) -> int: # noqa: D401 def reset_parameters(self) -> None: # noqa: D102 # initialize weights in-place - self._embeddings.weight.data = self.initializer(self._embeddings.weight.data) + self._embeddings.weight.data = self.initializer( + self._embeddings.weight.data.view(self.num_embeddings, *self.shape) + ).view(self.num_embeddings, self.embedding_dim) def post_parameter_update(self): # noqa: D102 # apply constraints in-place diff --git a/tests/test_models.py b/tests/test_models.py index 7aa72a8f79..bbaa750e64 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -25,7 +25,7 @@ from pykeen.datasets.nations import NATIONS_TEST_PATH, NATIONS_TRAIN_PATH, Nations from pykeen.models import _MODELS from pykeen.models.base import ( - EntityEmbeddingModel, + DoubleRelationEmbeddingModel, EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, MultimodalModel, @@ -810,7 +810,7 @@ def _check_constraints(self): Entity embeddings have to have unit L2 norm. """ - norms = self.model.entity_embeddings(indices=None).norm(p=2, dim=-1) + norms = self.model.entity_representations[0](indices=None).norm(p=2, dim=-1) assert torch.allclose(norms, torch.ones_like(norms)) @@ -1014,7 +1014,8 @@ def _check_constraints(self): Entity embeddings have to have unit L2 norm. """ - entity_norms = self.model.normal_vector_embeddings(indices=None).norm(p=2, dim=-1) + self.model: DoubleRelationEmbeddingModel + entity_norms = self.model.relation_representations[1](indices=None).norm(p=2, dim=-1) assert torch.allclose(entity_norms, torch.ones_like(entity_norms)) @@ -1028,6 +1029,7 @@ class TestTransR(_DistanceModelTestCase, unittest.TestCase): def test_score_hrt_manual(self): """Manually test interaction function of TransR.""" + # TODO: Move to interaction tests. # entity embeddings weights = torch.as_tensor(data=[[2., 2.], [3., 3.]], dtype=torch.float) entity_embeddings = Embedding( @@ -1068,7 +1070,7 @@ def _check_constraints(self): Entity and relation embeddings have to have at most unit L2 norm. """ - for emb in (self.model.entity_embeddings, self.model.relation_embeddings): + for emb in (self.model.entity_representations[0], self.model.relation_representations[0]): assert all_in_bounds(emb(indices=None).norm(p=2, dim=-1), high=1., a_tol=1.0e-06) From 6c1011809642b58a5ff1065731f191bc4f73f00d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 18:28:32 +0100 Subject: [PATCH 276/690] Move manual TransR score test to interaction function tests --- tests/test_interactions.py | 11 +++++++++++ tests/test_models.py | 40 +------------------------------------- 2 files changed, 12 insertions(+), 39 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 169815e63c..387492412b 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -325,6 +325,17 @@ class TransRTests(TranslationalInteractionTests, unittest.TestCase): e=3, ) + def test_manual(self): + """A manual test of the score function.""" + # Compute Scores + h = torch.as_tensor(data=[2, 2], dtype=torch.float32).view(1, 2) + r = torch.as_tensor(data=[4, 4], dtype=torch.float32).view(1, 2) + m_r = torch.as_tensor(data=[5, 5, 6, 6], dtype=torch.float32).view(1, 2, 2) + t = torch.as_tensor(data=[2, 2], dtype=torch.float32).view(1, 2) + scores = self.instance.score_hrt(h=h, r=(r, m_r), t=t) + first_score = scores[0].item() + self.assertAlmostEqual(first_score, -32, delta=1.0e-04) + class SETests(TranslationalInteractionTests, unittest.TestCase): """Tests for SE interaction function.""" diff --git a/tests/test_models.py b/tests/test_models.py index bbaa750e64..9a0ad111c0 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -41,7 +41,7 @@ inverse_outdegree_edge_weights, symmetric_edge_weights, ) -from pykeen.nn import Embedding, RepresentationModule +from pykeen.nn import RepresentationModule from pykeen.nn.functional import _project_entity from pykeen.training import LCWATrainingLoop, SLCWATrainingLoop, TrainingLoop from pykeen.triples import TriplesFactory @@ -1027,44 +1027,6 @@ class TestTransR(_DistanceModelTestCase, unittest.TestCase): 'relation_dim': 4, } - def test_score_hrt_manual(self): - """Manually test interaction function of TransR.""" - # TODO: Move to interaction tests. - # entity embeddings - weights = torch.as_tensor(data=[[2., 2.], [3., 3.]], dtype=torch.float) - entity_embeddings = Embedding( - num_embeddings=2, - embedding_dim=2, - ) - entity_embeddings._embeddings.weight.data.copy_(weights) - self.model.entity_embeddings = entity_embeddings - - # relation embeddings - relation_weights = torch.as_tensor(data=[[4., 4], [5., 5.]], dtype=torch.float) - relation_embeddings = Embedding( - num_embeddings=2, - embedding_dim=2, - ) - relation_embeddings._embeddings.weight.data.copy_(relation_weights) - self.model.relation_embeddings = relation_embeddings - - relation_projection_weights = torch.as_tensor(data=[[5., 5., 6., 6.], [7., 7., 8., 8.]], dtype=torch.float) - relation_projection_embeddings = Embedding( - num_embeddings=2, - embedding_dim=4, - ) - relation_projection_embeddings._embeddings.weight.data.copy_(relation_projection_weights) - self.model.relation_projections = relation_projection_embeddings - - # Compute Scores - batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 1]], dtype=torch.long) - scores = self.model.score_hrt(hrt_batch=batch) - self.assertEqual(scores.shape[0], 2) - self.assertEqual(scores.shape[1], 1) - first_score = scores[0].item() - # second_score = scores[1].item() - self.assertAlmostEqual(first_score, -8, delta=1.0e-06) - def _check_constraints(self): """Check model constraints. From 24a72cc2cf0bd2405f3706eddf8ac3f51eca0ccc Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 18:37:43 +0100 Subject: [PATCH 277/690] Move manual score tests for TransD to interaction fucntion tests --- src/pykeen/nn/functional.py | 2 +- src/pykeen/nn/modules.py | 3 + tests/test_interactions.py | 31 ++++++++- tests/test_models.py | 125 ------------------------------------ 4 files changed, 34 insertions(+), 127 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 5b95d4e397..a106123ee4 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -951,7 +951,7 @@ def transd_interaction( power_norm: bool = False, ) -> torch.FloatTensor: """ - Evaluate the DistMult interaction function. + Evaluate the TransD interaction function. :param h: shape: (batch_size, num_heads, d_e) The head representations. diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index c0be04b5d2..724b6b051e 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -825,6 +825,9 @@ class TransDInteractionFunction( entity_shape = ("d", "d") relation_shape = ("e", "e") + def __init__(self, p: int = 2, power_norm: bool = True): + super().__init__(p=p, power_norm=power_norm) + def forward( self, h: Tuple[torch.FloatTensor, torch.FloatTensor], diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 387492412b..463f99b730 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -301,6 +301,35 @@ class TransDTests(TranslationalInteractionTests, unittest.TestCase): e=3, ) + def test_manual_small_relation_dim(self): + """Manually test the value of the interaction function.""" + # entity embeddings + h = t = torch.as_tensor(data=[2., 2.], dtype=torch.float).view(1, 2) + h_p = t_p = torch.as_tensor(data=[3., 3.], dtype=torch.float).view(1, 2) + + # relation embeddings + r = torch.as_tensor(data=[4.], dtype=torch.float).view(1, 1) + r_p = torch.as_tensor(data=[5.], dtype=torch.float).view(1, 1) + + # Compute Scores + scores = self.instance.score_hrt(h=(h, h_p), r=(r, r_p), t=(t, t_p)) + first_score = scores[0].item() + self.assertAlmostEqual(first_score, -16, delta=0.01) + + def test_manual_big_relation_dim(self): + """Manually test the value of the interaction function.""" + # entity embeddings + h = t = torch.as_tensor(data=[2., 2.], dtype=torch.float).view(1, 2) + h_p = t_p = torch.as_tensor(data=[3., 3.], dtype=torch.float).view(1, 2) + + # relation embeddings + r = torch.as_tensor(data=[3., 3., 3.], dtype=torch.float).view(1, 3) + r_p = torch.as_tensor(data=[4., 4., 4.], dtype=torch.float).view(1, 3) + + # Compute Scores + scores = self.instance.score_hrt(h=(h, h_p), r=(r, r_p), t=(t, t_p)) + self.assertAlmostEqual(scores.item(), -27, delta=0.01) + class TransETests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransE interaction function.""" @@ -326,7 +355,7 @@ class TransRTests(TranslationalInteractionTests, unittest.TestCase): ) def test_manual(self): - """A manual test of the score function.""" + """Manually test the value of the interaction function.""" # Compute Scores h = torch.as_tensor(data=[2, 2], dtype=torch.float32).view(1, 2) r = torch.as_tensor(data=[4, 4], dtype=torch.float32).view(1, 2) diff --git a/tests/test_models.py b/tests/test_models.py index 9a0ad111c0..65cd4cc904 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -8,7 +8,6 @@ import traceback import unittest from typing import Any, ClassVar, Mapping, Optional, Type -from unittest.case import SkipTest from unittest.mock import patch import numpy @@ -846,130 +845,6 @@ def _check_constraints(self): for emb in (self.model.entity_representations[0], self.model.relation_representations[0]): assert all_in_bounds(emb(indices=None).norm(p=2, dim=-1), high=1., a_tol=_EPSILON) - def test_score_hrt_manual(self): - """Manually test interaction function of TransD.""" - raise SkipTest("TODO: Move this test to interaction function checks.") - # self.model: TwoVectorEmbeddingModel - # - # # entity embeddings - # weights = torch.as_tensor(data=[[2., 2.], [4., 4.]], dtype=torch.float) - # entity_embeddings = Embedding( - # num_embeddings=2, - # embedding_dim=2, - # ) - # entity_embeddings._embeddings.weight.data.copy_(weights) - # self.model.entity_representations[0] = entity_embeddings - # - # projection_weights = torch.as_tensor(data=[[3., 3.], [2., 2.]], dtype=torch.float) - # entity_projection_embeddings = Embedding( - # num_embeddings=2, - # embedding_dim=2, - # ) - # entity_projection_embeddings._embeddings.weight.data.copy_(projection_weights) - # self.model.second_entity_embeddings = entity_projection_embeddings - # - # # relation embeddings - # relation_weights = torch.as_tensor(data=[[4.], [4.]], dtype=torch.float) - # relation_embeddings = Embedding( - # num_embeddings=2, - # embedding_dim=1, - # ) - # relation_embeddings._embeddings.weight.data.copy_(relation_weights) - # self.model.relation_representations[0] = relation_embeddings - # - # relation_projection_weights = torch.as_tensor(data=[[5.], [3.]], dtype=torch.float) - # relation_projection_embeddings = Embedding( - # num_embeddings=2, - # embedding_dim=1, - # ) - # relation_projection_embeddings._embeddings.weight.data.copy_(relation_projection_weights) - # self.model.second_relation_embeddings = relation_projection_embeddings - # - # # Compute Scores - # batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 1]], dtype=torch.long) - # scores = self.model.score_hrt(hrt_batch=batch) - # self.assertEqual(scores.shape[0], 2) - # self.assertEqual(scores.shape[1], 1) - # first_score = scores[0].item() - # self.assertAlmostEqual(first_score, -16, delta=0.01) - # - # # Use different dimension for relation embedding: relation_dim > entity_dim - # # relation embeddings - # relation_weights = torch.as_tensor(data=[[3., 3., 3.], [3., 3., 3.]], dtype=torch.float) - # relation_embeddings = Embedding( - # num_embeddings=2, - # embedding_dim=3, - # ) - # relation_embeddings._embeddings.weight.data.copy_(relation_weights) - # self.model.relation_representations[0] = relation_embeddings - # - # relation_projection_weights = torch.as_tensor(data=[[4., 4., 4.], [4., 4., 4.]], dtype=torch.float) - # relation_projection_embeddings = Embedding( - # num_embeddings=2, - # embedding_dim=3, - # ) - # relation_projection_embeddings._embeddings.weight.data.copy_(relation_projection_weights) - # self.model.relation_representations[1] = relation_projection_embeddings - # - # # Compute Scores - # batch = torch.as_tensor(data=[[0, 0, 0]], dtype=torch.long) - # scores = self.model.score_hrt(hrt_batch=batch) - # self.assertAlmostEqual(scores.item(), -27, delta=0.01) - # - # batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 0]], dtype=torch.long) - # scores = self.model.score_hrt(hrt_batch=batch) - # self.assertEqual(scores.shape[0], 2) - # self.assertEqual(scores.shape[1], 1) - # first_score = scores[0].item() - # second_score = scores[1].item() - # self.assertAlmostEqual(first_score, -27, delta=0.01) - # self.assertAlmostEqual(second_score, -27, delta=0.01) - # - # # Use different dimension for relation embedding: relation_dim < entity_dim - # # entity embeddings - # weights = torch.as_tensor(data=[[1., 1., 1.], [1., 1., 1.]], dtype=torch.float) - # entity_embeddings = Embedding( - # num_embeddings=2, - # embedding_dim=3, - # ) - # entity_embeddings._embeddings.weight.data.copy_(weights) - # self.model.entity_representations[0] = entity_embeddings - # - # projection_weights = torch.as_tensor(data=[[2., 2., 2.], [2., 2., 2.]], dtype=torch.float) - # entity_projection_embeddings = Embedding( - # num_embeddings=2, - # embedding_dim=3, - # ) - # entity_projection_embeddings._embeddings.weight.data.copy_(projection_weights) - # self.model.second_entity_embeddings = entity_projection_embeddings - # - # # relation embeddings - # relation_weights = torch.as_tensor(data=[[3., 3.], [3., 3.]], dtype=torch.float) - # relation_embeddings = Embedding( - # num_embeddings=2, - # embedding_dim=2, - # ) - # relation_embeddings._embeddings.weight.data.copy_(relation_weights) - # self.model.relation_representations[0] = relation_embeddings - # - # relation_projection_weights = torch.as_tensor(data=[[4., 4.], [4., 4.]], dtype=torch.float) - # relation_projection_embeddings = Embedding( - # num_embeddings=2, - # embedding_dim=2, - # ) - # relation_projection_embeddings._embeddings.weight.data.copy_(relation_projection_weights) - # self.model.relation_representations[1] = relation_projection_embeddings - # - # # Compute Scores - # batch = torch.as_tensor(data=[[0, 0, 0], [0, 0, 0]], dtype=torch.long) - # scores = self.model.score_hrt(hrt_batch=batch) - # self.assertEqual(scores.shape[0], 2) - # self.assertEqual(scores.shape[1], 1) - # first_score = scores[0].item() - # second_score = scores[1].item() - # self.assertAlmostEqual(first_score, -18, delta=0.01) - # self.assertAlmostEqual(second_score, -18, delta=0.01) - def test_project_entity(self): """Test _project_entity.""" # random entity embeddings & projections From 7ebd8cd22d228ad19da98737857b937f1766d81b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 18:40:19 +0100 Subject: [PATCH 278/690] cleanup imports --- src/pykeen/models/unimodal/kg2e.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 28dde17cca..53b8d3e38c 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -10,9 +10,7 @@ from ..base import TwoVectorEmbeddingModel from ...losses import Loss -from ...nn import functional as pkf from ...nn.emb import EmbeddingSpecification -from ...nn.functional import KG2E_SIMILARITIES from ...nn.modules import KG2EInteractionFunction from ...regularizers import Regularizer from ...triples import TriplesFactory From e6c3d9defb34fda26269cd0b6cdcb3b636e01527 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 18:41:26 +0100 Subject: [PATCH 279/690] Export broadcast_cat --- src/pykeen/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 72bb7ba30c..24d91ce411 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -18,6 +18,7 @@ from .typing import DeviceHint __all__ = [ + 'broadcast_cat', 'compose', 'check_shapes', 'clamp_norm', From a879bbc1b035bb8d8c5760c8e7222b4d1499344b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 18:42:37 +0100 Subject: [PATCH 280/690] Add deprecation notice to EntityEmbeddingModel and EntityRelationEmbeddingModel --- src/pykeen/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index e339dc9ea2..ddf20d383c 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1136,7 +1136,6 @@ def load_state(self, path: str) -> None: class EntityEmbeddingModel(Model): """A base module for most KGE models that have one embedding for entities.""" - # TODO: deprecated def __init__( @@ -1190,6 +1189,7 @@ def __init__( class EntityRelationEmbeddingModel(Model, ABC): """A base module for KGE models that have different embeddings for entities and relations.""" + # TODO: Deprecated. def __init__( self, From 1401612d2f3801a91a09da57d28952e7a1009739 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 19:27:59 +0100 Subject: [PATCH 281/690] Add slicing for all models --- src/pykeen/nn/modules.py | 126 ++++++++++++++++++++++--------------- tests/test_interactions.py | 47 ++++++++++++-- 2 files changed, 116 insertions(+), 57 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 724b6b051e..a9490d43a8 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -22,6 +22,10 @@ def _ensure_tuple(*x: Union[Representation, Sequence[Representation]]) -> Tuple[ return tuple(xx if isinstance(xx, Sequence) else (xx,) for xx in x) +def _unpack_singletons(*xs: Tuple) -> Sequence[Tuple]: + return [x[0] if len(x) == 1 else x for x in xs] + + class InteractionFunction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation]): """Base class for interaction functions.""" @@ -136,6 +140,69 @@ def _check_shapes( ((tt, t_prefix + ts) for tt, ts in zip(t, tail_entity_shape)), ), raise_or_error=raise_on_error) + def _score( + self, + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + h_prefix: str = "b", + r_prefix: str = "b", + t_prefix: str = "b", + slice_size: Optional[int] = None, + ) -> torch.FloatTensor: + assert {h_prefix, r_prefix, t_prefix}.issubset(list("bn")) + # at most one of h_prefix, r_prefix, t_prefix equals n + slice_dim = [dim for dim, prefix in zip("hrt", (h_prefix, r_prefix, t_prefix)) if prefix == "n"] + if len(slice_dim) == 1: + slice_dim = slice_dim[0] + else: + slice_dim = None + h, r, t = _ensure_tuple(h, r, t) + assert self._check_shapes(h=h, r=r, t=t, h_prefix=h_prefix, r_prefix=r_prefix, t_prefix=t_prefix) + + # prepare input to generic score function: bh*, br*, bt* + h = self._add_dim(*h, dim=self.BATCH_DIM if h_prefix == "n" else self.NUM_DIM) + r = self._add_dim(*r, dim=self.BATCH_DIM if r_prefix == "n" else self.NUM_DIM) + t = self._add_dim(*t, dim=self.BATCH_DIM if t_prefix == "n" else self.NUM_DIM) + + # get scores + if slice_size is None: + scores = self(h=h, r=r, t=t) + else: + assert slice_dim is not None + scores = [] + if slice_dim == "h": + cat_dim = self.HEAD_DIM + for h_batch in zip(*(hh.split(slice_size, dim=1) for hh in _ensure_tuple(h)[0])): + if len(h_batch) == 1: + h_batch = h_batch[0] + scores.append(self(h=h_batch, r=r, t=t)) + elif slice_dim == "r": + cat_dim = self.RELATION_DIM + for r_batch in zip(*(rr.split(slice_size, dim=1) for rr in _ensure_tuple(r)[0])): + if len(r_batch) == 1: + r_batch = r_batch[0] + scores.append(self(h=h, r=r_batch, t=t)) + elif slice_dim == "t": + cat_dim = self.TAIL_DIM + for t_batch in zip(*(tt.split(slice_size, dim=1) for tt in _ensure_tuple(t)[0])): + if len(t_batch) == 1: + t_batch = t_batch[0] + scores.append(self(h=h, r=r, t=t_batch)) + else: + raise ValueError(slice_dim) + scores = torch.cat(scores, dim=cat_dim) + remove_dims = [ + dim + for dim, prefix in zip( + (self.HEAD_DIM, self.RELATION_DIM, self.TAIL_DIM), + (h_prefix, r_prefix, t_prefix), + ) + if prefix == "b" + ] + # prepare output shape + return self._remove_dim(scores, *remove_dims) + def score_hrt( self, h: HeadRepresentation = tuple(), @@ -155,25 +222,14 @@ def score_hrt( :return: shape: (batch_size, 1) The scores. """ - h, r, t = _ensure_tuple(h, r, t) - assert self._check_shapes(h=h, r=r, t=t) - - # prepare input to generic score function - h = self._add_dim(*h, dim=self.NUM_DIM) - r = self._add_dim(*r, dim=self.NUM_DIM) - t = self._add_dim(*t, dim=self.NUM_DIM) - - # get scores - scores = self(h=h, r=r, t=t) - - # prepare output shape, (batch_size, num_heads, num_relations, num_tails) -> (batch_size, 1) - return self._remove_dim(scores, self.HEAD_DIM, self.RELATION_DIM, self.TAIL_DIM).unsqueeze(dim=-1) + return self._score(h=h, r=r, t=t).unsqueeze(dim=-1) def score_h( self, all_entities: HeadRepresentation = tuple(), r: RelationRepresentation = tuple(), t: TailRepresentation = tuple(), + slice_size: Optional[int] = None, ) -> torch.FloatTensor: """ Score all head entities. @@ -188,25 +244,14 @@ def score_h( :return: shape: (batch_size, num_entities) The scores. """ - h, r, t = _ensure_tuple(all_entities, r, t) - assert self._check_shapes(h=h, r=r, t=t, h_prefix="n") - - # prepare input to generic score function - h = self._add_dim(*h, dim=self.BATCH_DIM) - r = self._add_dim(*r, dim=self.NUM_DIM) - t = self._add_dim(*t, dim=self.NUM_DIM) - - # get scores - scores = self(h=h, r=r, t=t) - - # prepare output shape, (batch_size, num_heads, num_relations, num_tails) -> (batch_size, num_heads) - return self._remove_dim(scores, self.RELATION_DIM, self.TAIL_DIM) + return self._score(h=all_entities, r=r, t=t, h_prefix="n", slice_size=slice_size) def score_r( self, h: HeadRepresentation = tuple(), all_relations: RelationRepresentation = tuple(), t: TailRepresentation = tuple(), + slice_size: Optional[int] = None, ) -> torch.FloatTensor: """ Score all relations. @@ -221,25 +266,14 @@ def score_r( :return: shape: (batch_size, num_entities) The scores. """ - h, r, t = _ensure_tuple(h, all_relations, t) - assert self._check_shapes(h=h, r=r, t=t, r_prefix="n") - - # prepare input to generic score function - h = self._add_dim(*h, dim=self.NUM_DIM) - r = self._add_dim(*r, dim=self.BATCH_DIM) - t = self._add_dim(*t, dim=self.NUM_DIM) - - # get scores - scores = self(h=h, r=r, t=t) - - # prepare output shape - return self._remove_dim(scores, self.HEAD_DIM, self.TAIL_DIM) + return self._score(h=h, r=all_relations, t=t, r_prefix="n", slice_size=slice_size) def score_t( self, h: HeadRepresentation = tuple(), r: RelationRepresentation = tuple(), all_entities: TailRepresentation = tuple(), + slice_size: Optional[int] = None, ) -> torch.FloatTensor: """ Score all tail entities. @@ -254,19 +288,7 @@ def score_t( :return: shape: (batch_size, num_entities) The scores. """ - h, r, t = _ensure_tuple(h, r, all_entities) - assert self._check_shapes(h=h, r=r, t=t, t_prefix="n") - - # prepare input to generic score function - h = self._add_dim(*h, dim=self.NUM_DIM) - r = self._add_dim(*r, dim=self.NUM_DIM) - t = self._add_dim(*t, dim=self.BATCH_DIM) - - # get scores - scores = self(h=h, r=r, t=t) - - # prepare output shape - return self._remove_dim(scores, self.HEAD_DIM, self.RELATION_DIM) + return self._score(h=h, r=r, t=all_entities, t_prefix="n", slice_size=slice_size) def reset_parameters(self): """Reset parameters the interaction function may have.""" diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 463f99b730..0b71f189a8 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -32,10 +32,6 @@ def post_instantiation_hook(self) -> None: """Perform actions after instantiation.""" -def _unpack_singletons(*xs: Tuple) -> Sequence[Tuple]: - return [x[0] if len(x) == 1 else x for x in xs] - - class InteractionTests(GenericTests[pykeen.nn.modules.InteractionFunction]): """Generic test for interaction functions.""" @@ -61,7 +57,7 @@ def _get_hrt( [self.cls.entity_shape, self.cls.relation_shape, self.cls.entity_shape], ) ) - return tuple(_unpack_singletons(*result)) + return tuple(pykeen.nn.modules._unpack_singletons(*result)) def _check_scores(self, scores: torch.FloatTensor, exp_shape: Tuple[int, ...]): """Check shape, dtype and gradients of scores.""" @@ -95,6 +91,19 @@ def test_score_h(self): scores = self.instance.score_h(all_entities=h, r=r, t=t) self._check_scores(scores=scores, exp_shape=(self.batch_size, self.num_entities)) + def test_score_h_slicing(self): + """Test score_h with slicing.""" + #: The equivalence for models with batch norm only holds in evaluation mode + self.instance.eval() + h, r, t = self._get_hrt( + (self.num_entities,), + (self.batch_size,), + (self.batch_size,), + ) + scores = self.instance.score_h(all_entities=h, r=r, t=t, slice_size=self.num_entities // 2 + 1) + scores_no_slice = self.instance.score_h(all_entities=h, r=r, t=t, slice_size=None) + assert torch.allclose(scores, scores_no_slice) + def test_score_r(self): """Test score_r.""" h, r, t = self._get_hrt( @@ -109,6 +118,21 @@ def test_score_r(self): exp_shape = (self.batch_size, self.num_relations) self._check_scores(scores=scores, exp_shape=exp_shape) + def test_score_r_slicing(self): + """Test score_r with slicing.""" + if len(self.cls.relation_shape) == 0: + raise SkipTest("No use in slicing relations for models without relation information.") + #: The equivalence for models with batch norm only holds in evaluation mode + self.instance.eval() + h, r, t = self._get_hrt( + (self.batch_size,), + (self.num_relations,), + (self.batch_size,), + ) + scores = self.instance.score_r(h=h, all_relations=r, t=t, slice_size=self.num_relations // 2 + 1) + scores_no_slice = self.instance.score_r(h=h, all_relations=r, t=t, slice_size=None) + assert torch.allclose(scores, scores_no_slice) + def test_score_t(self): """Test score_t.""" h, r, t = self._get_hrt( @@ -119,6 +143,19 @@ def test_score_t(self): scores = self.instance.score_t(h=h, r=r, all_entities=t) self._check_scores(scores=scores, exp_shape=(self.batch_size, self.num_entities)) + def test_score_t_slicing(self): + """Test score_t with slicing.""" + #: The equivalence for models with batch norm only holds in evaluation mode + self.instance.eval() + h, r, t = self._get_hrt( + (self.batch_size,), + (self.batch_size,), + (self.num_entities,), + ) + scores = self.instance.score_t(h=h, r=r, all_entities=t, slice_size=self.num_entities // 2 + 1) + scores_no_slice = self.instance.score_t(h=h, r=r, all_entities=t, slice_size=None) + assert torch.allclose(scores, scores_no_slice) + def test_forward(self): """Test forward.""" for hs, rs, ts in [ From 47d601f482a919b89ba15002755aaaa0a7e79dc8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 19:40:28 +0100 Subject: [PATCH 282/690] Add generic testtest and fix ER-MLPE --- src/pykeen/nn/functional.py | 11 ++++++++--- src/pykeen/utils.py | 12 ++++++++++++ tests/test_interactions.py | 32 ++++++++++++++++++++++++++++++-- 3 files changed, 50 insertions(+), 5 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index a106123ee4..184a8af4c5 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -626,11 +626,16 @@ def ermlpe_interaction( # repeat if necessary, and concat head and relation, (batch_size, num_heads, num_relations, 2 * embedding_dim) x = broadcast_cat(h.unsqueeze(dim=2), r.unsqueeze(dim=1), dim=-1) - # Predict t embedding, shape: (batch_size, num_heads, num_relations, embedding_dim) + # Predict t embedding, shape: (b, h, r, 1, d) shape = x.shape - x = mlp(x.view(-1, shape[-1])).view(*shape[:-1], -1) + x = mlp(x.view(-1, shape[-1])).view(*shape[:-1], -1).unsqueeze(dim=-2) - return (x.unsqueeze(dim=-2) @ t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]).transpose(-2, -1)).squeeze(dim=-1) + # transpose t, (b, 1, 1, d, t) + t = t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]).transpose(-2, -1) + + # dot product, (b, h, r, 1, t) + x = (x @ t).squeeze(dim=-2) + return x def hole_interaction( diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 24d91ce411..c4c84ee5b5 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -533,3 +533,15 @@ def get_ht_indices(ht_batch: torch.LongTensor) -> Tuple[torch.LongTensor, None, def get_rt_indices(rt_batch: torch.LongTensor) -> Tuple[None, torch.LongTensor, torch.LongTensor]: """Get indices from a relation/tail batch.""" return None, rt_batch[:, 0], rt_batch[:, 1] + + +def get_subclasses(cls: Type[X]) -> Iterable[Type[X]]: + """ + Get all subclasses. + + Credit to: https://stackoverflow.com/a/33607093. + + """ + for subclass in cls.__subclasses__(): + yield from get_subclasses(subclass) + yield subclass diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 0b71f189a8..436efff9f2 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -1,12 +1,14 @@ """Tests for interaction functions.""" import unittest -from typing import Any, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union +from typing import Any, Collection, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union from unittest.case import SkipTest import torch import pykeen.nn.modules +from pykeen.nn.modules import InteractionFunction, StatelessInteractionFunction, TranslationalInteractionFunction from pykeen.typing import Representation +from pykeen.utils import get_subclasses T = TypeVar("T") @@ -32,6 +34,21 @@ def post_instantiation_hook(self) -> None: """Perform actions after instantiation.""" +class TestsTest(Generic[T]): + """A generic test for tests.""" + + base_cls: Type[T] + base_test: Type[GenericTests[T]] + skip_cls: Collection[T] = tuple() + + def test_testing(self): + """Check that there is a test for all subclasses.""" + to_test = set(get_subclasses(self.base_cls)).difference(self.skip_cls) + tested = (test_cls.cls for test_cls in get_subclasses(self.base_test) if hasattr(test_cls, "cls")) + not_tested = to_test.difference(tested) + assert not_tested == set() + + class InteractionTests(GenericTests[pykeen.nn.modules.InteractionFunction]): """Generic test for interaction functions.""" @@ -259,7 +276,7 @@ class ERMLPTests(InteractionTests, unittest.TestCase): class ERMLPETests(InteractionTests, unittest.TestCase): """Tests for ERMLP-E interaction function.""" - cls = pykeen.nn.modules.ERMLPInteractionFunction + cls = pykeen.nn.modules.ERMLPEInteractionFunction kwargs = dict( embedding_dim=InteractionTests.dim, hidden_dim=2 * InteractionTests.dim - 1, @@ -413,3 +430,14 @@ class UMTests(TranslationalInteractionTests, unittest.TestCase): """Tests for UM interaction function.""" cls = pykeen.nn.modules.UnstructuredModelInteractionFunction + + +class InteractionTestsTest(TestsTest[InteractionFunction], unittest.TestCase): + """Test for tests for all interaction functions.""" + + base_cls = InteractionFunction + base_test = InteractionTests + skip_cls = { + TranslationalInteractionFunction, + StatelessInteractionFunction, + } From 965f77ed2c279a97bae79b06e84843d5bd32b97f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 19:45:19 +0100 Subject: [PATCH 283/690] Add __all__ --- src/pykeen/nn/modules.py | 21 +++++++++++++++++++++ 1 file changed, 21 insertions(+) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index a9490d43a8..6b114ba623 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -17,6 +17,27 @@ logger = logging.getLogger(__name__) +__all__ = [ + "ComplExInteractionFunction", + "ConvEInteractionFunction", + "ConvKBInteractionFunction", + "DistMultInteractionFunction", + "ERMLPInteractionFunction", + "ERMLPEInteractionFunction", + "HolEInteractionFunction", + "InteractionFunction", + "KG2EInteractionFunction", + "NTNInteractionFunction", + "ProjEInteractionFunction", + "RESCALInteractionFunction", + "RotatEInteractionFunction", + "StructuredEmbeddingInteractionFunction", + "TransDInteractionFunction", + "TransEInteractionFunction", + "TransHInteractionFunction", + "TransRInteractionFunction", +] + def _ensure_tuple(*x: Union[Representation, Sequence[Representation]]) -> Tuple[Sequence[Representation], ...]: return tuple(xx if isinstance(xx, Sequence) else (xx,) for xx in x) From 13a930d9087ed782f341fa7d8c1376584ea7a82d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sat, 14 Nov 2020 19:45:38 +0100 Subject: [PATCH 284/690] Add Representation to __all__ --- src/pykeen/typing.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/typing.py b/src/pykeen/typing.py index ab21b6eb20..a05b1e8e9f 100644 --- a/src/pykeen/typing.py +++ b/src/pykeen/typing.py @@ -15,6 +15,7 @@ 'DeviceHint', 'HeadRepresentation', 'RelationRepresentation', + 'Representation', 'TailRepresentation', ] From 317649058d262b99a486b20978c31a18bd75a48b Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sun, 15 Nov 2020 19:03:14 +0100 Subject: [PATCH 285/690] Pass flake8 and start mypy --- .travis.yml | 5 +- src/pykeen/models/base.py | 42 ++-- src/pykeen/models/unimodal/complex.py | 4 +- src/pykeen/models/unimodal/conv_e.py | 6 +- src/pykeen/models/unimodal/conv_kb.py | 4 +- src/pykeen/models/unimodal/distmult.py | 4 +- src/pykeen/models/unimodal/ermlp.py | 4 +- src/pykeen/models/unimodal/ermlpe.py | 4 +- src/pykeen/models/unimodal/hole.py | 4 +- src/pykeen/models/unimodal/ntn.py | 6 +- src/pykeen/models/unimodal/proj_e.py | 4 +- src/pykeen/models/unimodal/rescal.py | 4 +- src/pykeen/models/unimodal/simple.py | 4 +- .../models/unimodal/structured_embedding.py | 6 +- src/pykeen/models/unimodal/trans_d.py | 4 +- src/pykeen/models/unimodal/trans_h.py | 4 +- src/pykeen/models/unimodal/trans_r.py | 8 +- src/pykeen/models/unimodal/tucker.py | 4 +- .../models/unimodal/unstructured_model.py | 4 +- src/pykeen/nn/emb.py | 49 +++-- src/pykeen/nn/functional.py | 14 +- src/pykeen/nn/modules.py | 181 ++++++++++-------- src/pykeen/typing.py | 9 +- src/pykeen/utils.py | 19 +- tests/test_interactions.py | 58 +++--- tests/test_models.py | 6 +- tox.ini | 9 +- 27 files changed, 262 insertions(+), 208 deletions(-) diff --git a/.travis.yml b/.travis.yml index f786e42416..7a0b6dcc20 100644 --- a/.travis.yml +++ b/.travis.yml @@ -15,8 +15,8 @@ jobs: env: TOXENV=flake8 - name: "Check documentation with darglint" env: TOXENV=darglint - # - name: "Check typing with MyPy" - # env: TOXENV=mypy + - name: "Check typing with MyPy" + env: TOXENV=mypy # docs stage - name: "Check RST documentation style" stage: docs @@ -34,7 +34,6 @@ jobs: env: TOXENV=integration allow_failures: - env: TOXENV=darglint - # - env: TOXENV=mypy install: - sh -c 'if [ "$TOXENV" = "py" ]; then pip install tox codecov coverage; else pip install tox; fi' - sh -c 'if [ "$TOXENV" = "docs" ]; then sudo apt-get install graphviz; fi' diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index ddf20d383c..6ad5f8c02a 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -9,7 +9,10 @@ from abc import ABC from collections import defaultdict from operator import itemgetter -from typing import Any, ClassVar, Collection, Dict, Generic, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, Type, Union +from typing import ( + Any, ClassVar, Collection, Dict, Generic, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, + Type, Union, +) import numpy as np import pandas as pd @@ -19,10 +22,13 @@ from ..losses import Loss, MarginRankingLoss, NSSALoss from ..nn import Embedding, RepresentationModule from ..nn.emb import EmbeddingSpecification -from ..nn.modules import InteractionFunction +from ..nn.modules import Interaction from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory -from ..typing import Constrainer, DeviceHint, HeadRepresentation, Initializer, MappedTriples, Normalizer, RelationRepresentation, Representation, TailRepresentation +from ..typing import ( + Constrainer, DeviceHint, HeadRepresentation, Initializer, MappedTriples, Normalizer, + RelationRepresentation, Representation, TailRepresentation, +) from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed __all__ = [ @@ -217,12 +223,14 @@ class Model(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailR #: The default strategy for optimizing the model's hyper-parameters hpo_default: ClassVar[Mapping[str, Any]] + #: The default loss function class loss_default: ClassVar[Type[Loss]] = MarginRankingLoss #: The default parameters for the default loss function class loss_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = dict(margin=1.0, reduction='mean') #: The instance of the loss loss: Loss + #: The default regularizer class regularizer_default: ClassVar[Type[Regularizer]] = NoRegularizer #: The default parameters for the default regularizer class @@ -247,7 +255,8 @@ def __init__( regularizer: Optional[Regularizer] = None, entity_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, relation_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, - interaction_function: InteractionFunction[HeadRepresentation, RelationRepresentation, TailRepresentation] = None, + interaction_function: Interaction[ + HeadRepresentation, RelationRepresentation, TailRepresentation] = None, ) -> None: """Initialize the module. @@ -1136,12 +1145,13 @@ def load_state(self, path: str) -> None: class EntityEmbeddingModel(Model): """A base module for most KGE models that have one embedding for entities.""" + # TODO: deprecated def __init__( self, triples_factory: TriplesFactory, - interaction_function: InteractionFunction[Representation, RelationRepresentation, Representation], + interaction_function: Interaction[Representation, RelationRepresentation, Representation], embedding_dim: int = 50, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, @@ -1155,7 +1165,6 @@ def __init__( entity_normalizer_kwargs: Optional[Mapping[str, Any]] = None, entity_constrainer: Optional[Constrainer] = None, entity_constrainer_kwargs: Optional[Mapping[str, Any]] = None, - ) -> None: """Initialize the entity embedding model. @@ -1189,12 +1198,13 @@ def __init__( class EntityRelationEmbeddingModel(Model, ABC): """A base module for KGE models that have different embeddings for entities and relations.""" + # TODO: Deprecated. def __init__( self, triples_factory: TriplesFactory, - interaction_function: InteractionFunction[Representation, Representation, Representation], + interaction_function: Interaction[Representation, Representation, Representation], embedding_dim: int = 50, relation_dim: Optional[int] = None, loss: Optional[Loss] = None, @@ -1256,7 +1266,7 @@ class SingleVectorEmbeddingModel(Model, ABC): def __init__( self, triples_factory: TriplesFactory, - interaction_function: InteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], + interaction_function: Interaction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], embedding_dim: int = 200, relation_dim: Union[None, int, Sequence[int]] = None, automatic_memory_optimization: Optional[bool] = None, @@ -1313,14 +1323,14 @@ def __init__( ) @property - def embedding_dim(self) -> int: + def embedding_dim(self) -> int: # noqa:D401 """The entity embedding dim.""" embedding = self.entity_representations[0] assert isinstance(embedding, Embedding) return embedding.embedding_dim @property - def relation_dim(self) -> int: + def relation_dim(self) -> int: # noqa:D401 """The relation embedding dim.""" embedding = self.relation_representations[0] assert isinstance(embedding, Embedding) @@ -1333,7 +1343,7 @@ class DoubleRelationEmbeddingModel(Model, ABC): def __init__( self, triples_factory: TriplesFactory, - interaction_function: InteractionFunction[ + interaction_function: Interaction[ torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor], torch.FloatTensor, @@ -1377,7 +1387,7 @@ def __init__( num_embeddings=triples_factory.num_relations, shape=relation_dim, specification=second_relation_embedding_specification, - ) + ), ], interaction_function=interaction_function, ) @@ -1389,7 +1399,7 @@ class TwoVectorEmbeddingModel(Model, ABC): def __init__( self, triples_factory: TriplesFactory, - interaction_function: InteractionFunction[ + interaction_function: Interaction[ Tuple[torch.FloatTensor, torch.FloatTensor], Tuple[torch.FloatTensor, torch.FloatTensor], Tuple[torch.FloatTensor, torch.FloatTensor], @@ -1427,7 +1437,7 @@ def __init__( num_embeddings=triples_factory.num_entities, embedding_dim=embedding_dim, specification=second_embedding_specification, - ) + ), ], relation_representations=[ Embedding.from_specification( @@ -1439,7 +1449,7 @@ def __init__( num_embeddings=triples_factory.num_relations, embedding_dim=relation_dim, specification=second_relation_embedding_specification, - ) + ), ], interaction_function=interaction_function, ) @@ -1451,7 +1461,7 @@ class TwoSideEmbeddingModel(Model): def __init__( self, triples_factory: TriplesFactory, - interaction_function: InteractionFunction, + interaction_function: Interaction, embedding_dim: int = 50, relation_dim: Optional[int] = None, loss: Optional[Loss] = None, diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 04c9c81c08..ac57238775 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -9,7 +9,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss, SoftplusLoss from ...nn.emb import EmbeddingSpecification -from ...nn.modules import ComplExInteractionFunction +from ...nn.modules import ComplExInteraction from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -94,7 +94,7 @@ def __init__( :param regularizer: BaseRegularizer The regularizer to use. """ - interaction_function = ComplExInteractionFunction() + interaction_function = ComplExInteraction() super().__init__( triples_factory=triples_factory, diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 75a19b30a0..8785b2bf53 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -12,7 +12,7 @@ from ...nn import Embedding from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_normal_ -from ...nn.modules import ConvEInteractionFunction +from ...nn.modules import ConvEInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -156,7 +156,7 @@ def __init__( specification=EmbeddingSpecification( initializer=nn.init.zeros_, ), - ) + ), ], relation_representations=Embedding.from_specification( num_embeddings=triples_factory.num_relations, @@ -165,7 +165,7 @@ def __init__( initializer=xavier_normal_, ), ), - interaction_function=ConvEInteractionFunction( + interaction_function=ConvEInteraction( input_channels=input_channels, output_channels=output_channels, embedding_height=embedding_height, diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 7825b8207f..6e5256a787 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -7,7 +7,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss -from ...nn.modules import ConvKBInteractionFunction +from ...nn.modules import ConvKBInteraction from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -86,7 +86,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - interaction_function=ConvKBInteractionFunction( + interaction_function=ConvKBInteraction( hidden_dropout_rate=hidden_dropout_rate, embedding_dim=embedding_dim, num_filters=num_filters, diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index f70dee9267..e305ce2987 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -10,7 +10,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss from ...nn.emb import EmbeddingSpecification -from ...nn.modules import DistMultInteractionFunction +from ...nn.modules import DistMultInteraction from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -81,7 +81,7 @@ def __init__( :param embedding_dim: The entity embedding dimension $d$. Is usually $d \in [50, 300]$. """ - interaction_function = DistMultInteractionFunction() + interaction_function = DistMultInteraction() super().__init__( triples_factory=triples_factory, diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index ac273aed52..0c4b2f6578 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -6,7 +6,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss -from ...nn.modules import ERMLPInteractionFunction +from ...nn.modules import ERMLPInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -52,7 +52,7 @@ def __init__( """Initialize the model.""" super().__init__( triples_factory=triples_factory, - interaction_function=ERMLPInteractionFunction( + interaction_function=ERMLPInteraction( embedding_dim=embedding_dim, hidden_dim=hidden_dim, ), diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index 863e92a978..66f3e5066c 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -6,7 +6,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss -from ...nn.modules import ERMLPEInteractionFunction +from ...nn.modules import ERMLPEInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -66,7 +66,7 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, - interaction_function=ERMLPEInteractionFunction( + interaction_function=ERMLPEInteraction( hidden_dim=hidden_dim, input_dropout=input_dropout, hidden_dropout=hidden_dropout, diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 68cf707114..01f9ab9e7c 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -8,7 +8,7 @@ from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ -from ...nn.modules import HolEInteractionFunction +from ...nn.modules import HolEInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -63,7 +63,7 @@ def __init__( regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model.""" - interaction_function = HolEInteractionFunction() + interaction_function = HolEInteraction() super().__init__( triples_factory=triples_factory, diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 2d7ee08739..4f78ad982a 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -9,7 +9,7 @@ from .. import Model from ...losses import Loss from ...nn import Embedding -from ...nn.modules import NTNInteractionFunction +from ...nn.modules import NTNInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -73,7 +73,7 @@ def __init__( """ w = Embedding( num_embeddings=triples_factory.num_relations, - shape=(num_slices, embedding_dim, embedding_dim) + shape=(num_slices, embedding_dim, embedding_dim), ) b = Embedding( num_embeddings=triples_factory.num_relations, @@ -98,7 +98,7 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - interaction_function=NTNInteractionFunction( + interaction_function=NTNInteraction( non_linearity=non_linearity, ), entity_representations=Embedding( diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index 3adea5c568..dc4549d14f 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -10,7 +10,7 @@ from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ -from ...nn.modules import ProjEInteractionFunction +from ...nn.modules import ProjEInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -67,7 +67,7 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, - interaction_function=ProjEInteractionFunction( + interaction_function=ProjEInteraction( embedding_dim=embedding_dim, inner_non_linearity=inner_non_linearity, ), diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index 6fbd574692..dc755e8ce2 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -6,7 +6,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss -from ...nn.modules import RESCALInteractionFunction +from ...nn.modules import RESCALInteraction from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -70,7 +70,7 @@ def __init__( # TODO: regularization super().__init__( triples_factory=triples_factory, - interaction_function=RESCALInteractionFunction(), + interaction_function=RESCALInteraction(), embedding_dim=embedding_dim, relation_dim=(embedding_dim, embedding_dim), automatic_memory_optimization=automatic_memory_optimization, diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index ebe0a17253..84b3a602b4 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -8,7 +8,7 @@ from ..base import TwoSideEmbeddingModel from ...losses import Loss, SoftplusLoss -from ...nn.modules import DistMultInteractionFunction +from ...nn.modules import DistMultInteraction from ...regularizers import PowerSumRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -75,7 +75,7 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, - interaction_function=DistMultInteractionFunction(), + interaction_function=DistMultInteraction(), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index f2af434609..7b6a3cf02b 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -13,7 +13,7 @@ from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ -from ...nn.modules import StructuredEmbeddingInteractionFunction +from ...nn.modules import StructuredEmbeddingInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -71,7 +71,7 @@ def __init__( ) super().__init__( triples_factory=triples_factory, - interaction_function=StructuredEmbeddingInteractionFunction( + interaction_function=StructuredEmbeddingInteraction( p=scoring_fct_norm, power_norm=False, ), @@ -91,5 +91,5 @@ def __init__( ), second_relation_embedding_specification=EmbeddingSpecification( initializer=relation_initializer, - ) + ), ) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index e09e433437..bd62e10fc0 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -8,7 +8,7 @@ from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_normal_ -from ...nn.modules import TransDInteractionFunction +from ...nn.modules import TransDInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -70,7 +70,7 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, - interaction_function=TransDInteractionFunction(p=2, power_norm=True), + interaction_function=TransDInteraction(p=2, power_norm=True), embedding_dim=embedding_dim, relation_dim=relation_dim, automatic_memory_optimization=automatic_memory_optimization, diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 1d69fba0d6..2bdea2ebfb 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -9,7 +9,7 @@ from ..base import DoubleRelationEmbeddingModel from ...losses import Loss from ...nn.emb import EmbeddingSpecification -from ...nn.modules import TransHInteractionFunction +from ...nn.modules import TransHInteraction from ...regularizers import Regularizer, TransHRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -81,7 +81,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - interaction_function=TransHInteractionFunction( + interaction_function=TransHInteraction( p=scoring_fct_norm, power_norm=False, ), diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 89f9876cbe..4390a31a55 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -11,7 +11,7 @@ from ...nn import Embedding from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ -from ...nn.modules import TransRInteractionFunction +from ...nn.modules import TransRInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -87,7 +87,7 @@ def __init__( initializer=xavier_uniform_, constrainer=clamp_norm, constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), - ) + ), ), relation_representations=[ Embedding.from_specification( @@ -109,9 +109,9 @@ def __init__( specification=EmbeddingSpecification( initializer=xavier_uniform_, ), - ) + ), ], - interaction_function=TransRInteractionFunction( + interaction_function=TransRInteraction( p=scoring_fct_norm, ), ) diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index 718d7caef0..c8f6bbd3d8 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -8,7 +8,7 @@ from ...losses import BCEAfterSigmoidLoss, Loss from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_normal_ -from ...nn.modules import TuckerInteractionFunction +from ...nn.modules import TuckerInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -89,7 +89,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - interaction_function=TuckerInteractionFunction( + interaction_function=TuckerInteraction( embedding_dim=embedding_dim, relation_dim=relation_dim, dropout_0=dropout_0, diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index faeb298c81..734b395369 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -8,7 +8,7 @@ from ...losses import Loss from ...nn import Embedding from ...nn.init import xavier_normal_ -from ...nn.modules import UnstructuredModelInteractionFunction +from ...nn.modules import UnstructuredModelInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -68,7 +68,7 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - interaction_function=UnstructuredModelInteractionFunction(p=scoring_fct_norm), + interaction_function=UnstructuredModelInteraction(p=scoring_fct_norm), entity_representations=Embedding( num_embeddings=triples_factory.num_entities, embedding_dim=embedding_dim, diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index b19883d29b..fc92fce828 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -1,8 +1,10 @@ # -*- coding: utf-8 -*- """Embedding modules.""" + import dataclasses import functools +import logging from typing import Any, Mapping, Optional, Sequence, Union import numpy @@ -18,14 +20,13 @@ 'EmbeddingSpecification', ] +logger = logging.getLogger(__name__) + class RepresentationModule(nn.Module): """A base class for obtaining representations for entities/relations.""" - def forward( - self, - indices: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: + def forward(self, indices: Optional[torch.LongTensor] = None) -> torch.FloatTensor: """Get representations for indices. :param indices: shape: (m,) @@ -65,7 +66,7 @@ def make( self, num_embeddings: int, embedding_dim: Optional[int], - shape: Optional[Union[int, Sequence[int]]] + shape: Optional[Union[int, Sequence[int]]], ) -> 'Embedding': """Create an embedding with this specification.""" return Embedding( @@ -88,11 +89,15 @@ class Embedding(RepresentationModule): can be used throughout PyKEEN as a more fully featured drop-in replacement. """ + shape: Sequence[int] + normalizer: Optional[Normalizer] + constrainer: Optional[Constrainer] + def __init__( self, num_embeddings: int, embedding_dim: Optional[int] = None, - shape: Optional[Union[int, Sequence[int]]] = None, + shape: Union[None, int, Sequence[int]] = None, initializer: Optional[Initializer] = None, initializer_kwargs: Optional[Mapping[str, Any]] = None, normalizer: Optional[Normalizer] = None, @@ -125,29 +130,33 @@ def __init__( """ super().__init__() if shape is None and embedding_dim is None: - raise ValueError - if shape is not None: - if not isinstance(shape, Sequence): - shape = (shape,) + raise ValueError('Missing both shape and embedding_dim') + elif shape is not None: + if isinstance(shape, int): + self.shape = (shape,) + else: + self.shape = shape embedding_dim = numpy.prod(shape) else: - shape = (embedding_dim,) - self.shape = shape + assert embedding_dim is not None + self.shape = (embedding_dim,) if initializer is None: initializer = nn.init.normal_ if initializer_kwargs: - self.initializer = functools.partial(initializer, **initializer_kwargs) - else: - self.initializer = initializer - if constrainer_kwargs: - self.constrainer = functools.partial(constrainer, **constrainer_kwargs) + initializer = functools.partial(initializer, **initializer_kwargs) + self.initializer = initializer + + if constrainer is not None and constrainer_kwargs: + self.constrainer: Constrainer = functools.partial(constrainer, **constrainer_kwargs) else: self.constrainer = constrainer - if normalizer_kwargs: - self.normalizer = functools.partial(normalizer, **normalizer_kwargs) + + if normalizer is not None and normalizer_kwargs: + self.normalizer: Normalizer = functools.partial(normalizer, **normalizer_kwargs) else: self.normalizer = normalizer + self._embeddings = torch.nn.Embedding( num_embeddings=num_embeddings, embedding_dim=embedding_dim, @@ -193,7 +202,7 @@ def embedding_dim(self) -> int: # noqa: D401 def reset_parameters(self) -> None: # noqa: D102 # initialize weights in-place self._embeddings.weight.data = self.initializer( - self._embeddings.weight.data.view(self.num_embeddings, *self.shape) + self._embeddings.weight.data.view(self.num_embeddings, *self.shape), ).view(self.num_embeddings, self.embedding_dim) def post_parameter_update(self): # noqa: D102 diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 184a8af4c5..3bc8bd9bd4 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -44,7 +44,7 @@ def _extract_sizes(h, r, t) -> Tuple[int, int, int, int, int]: def _negative_norm_of_sum( *x: torch.FloatTensor, - p: Union[int, str] = 2, + p: Union[str, int] = 2, power_norm: bool = False, ) -> torch.FloatTensor: """ @@ -60,12 +60,13 @@ def _negative_norm_of_sum( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - d = sum(x) + d: torch.FloatTensor = sum(x) if power_norm: assert isinstance(p, SupportsFloat) return -(d.abs() ** p).sum(dim=-1) else: if torch.is_complex(d): + assert isinstance(p, SupportsFloat) # workaround for complex numbers: manually compute norm return -(d.abs() ** p).sum(dim=-1) ** (1 / p) else: @@ -250,11 +251,10 @@ def _kullback_leibler_similarity( return sim -_KG2E_SIMILARITIES = dict( +KG2E_SIMILARITIES = dict( KL=_kullback_leibler_similarity, EL=_expected_likelihood, ) -KG2E_SIMILARITIES = set(_KG2E_SIMILARITIES.keys()) def _extended_einsum( @@ -706,15 +706,13 @@ def kg2e_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - if similarity not in KG2E_SIMILARITIES: - raise KeyError(similarity) - similarity = _KG2E_SIMILARITIES[similarity] + similarity_fn = KG2E_SIMILARITIES[similarity] # Compute entity distribution e_mean = h_mean.unsqueeze(dim=2) - t_mean.unsqueeze(dim=1) e_var = h_var.unsqueeze(dim=2) + t_var.unsqueeze(dim=1) e = GaussianDistribution(mean=e_mean, diagonal_covariance=e_var) r = GaussianDistribution(mean=r_mean, diagonal_covariance=r_var) - return similarity(e=e, r=r, exact=exact) + return similarity_fn(e=e, r=r, exact=exact) def ntn_interaction( diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 6b114ba623..0755f3d976 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -1,14 +1,15 @@ # -*- coding: utf-8 -*- """Stateful interaction functions.""" + import itertools import logging import math from abc import ABC -from typing import Callable, Generic, Optional, Sequence, Tuple, Union +from typing import Callable, Generic, List, Optional, Sequence, Tuple, Union import torch -from torch import nn +from torch import FloatTensor, nn from . import functional as pkf from .functional import KG2E_SIMILARITIES @@ -18,24 +19,24 @@ logger = logging.getLogger(__name__) __all__ = [ - "ComplExInteractionFunction", - "ConvEInteractionFunction", - "ConvKBInteractionFunction", - "DistMultInteractionFunction", - "ERMLPInteractionFunction", - "ERMLPEInteractionFunction", - "HolEInteractionFunction", - "InteractionFunction", - "KG2EInteractionFunction", - "NTNInteractionFunction", - "ProjEInteractionFunction", - "RESCALInteractionFunction", - "RotatEInteractionFunction", - "StructuredEmbeddingInteractionFunction", - "TransDInteractionFunction", - "TransEInteractionFunction", - "TransHInteractionFunction", - "TransRInteractionFunction", + "ComplExInteraction", + "ConvEInteraction", + "ConvKBInteraction", + "DistMultInteraction", + "ERMLPInteraction", + "ERMLPEInteraction", + "HolEInteraction", + "Interaction", + "KG2EInteraction", + "NTNInteraction", + "ProjEInteraction", + "RESCALInteraction", + "RotatEInteraction", + "StructuredEmbeddingInteraction", + "TransDInteraction", + "TransEInteraction", + "TransHInteraction", + "TransRInteraction", ] @@ -47,7 +48,7 @@ def _unpack_singletons(*xs: Tuple) -> Sequence[Tuple]: return [x[0] if len(x) == 1 else x for x in xs] -class InteractionFunction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation]): +class Interaction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation]): """Base class for interaction functions.""" # Dimensions @@ -58,11 +59,12 @@ class InteractionFunction(nn.Module, Generic[HeadRepresentation, RelationReprese TAIL_DIM: int = 3 #: The symbolic shapes for entity representations - entity_shape: Tuple[str, ...] = ("d",) - tail_entity_shape: Union[None, Tuple[str, ...]] = None + entity_shape: Sequence[str] = ("d",) + # TODO add docstring + tail_entity_shape: Optional[Sequence[str]] = None #: The symbolic shapes for relation representations - relation_shape: Tuple[str, ...] = ("d",) + relation_shape: Sequence[str] = ("d",) def forward( self, @@ -114,7 +116,7 @@ def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: The squeezed tensor. """ # normalize dimensions - dims = [d if d >= 0 else len(x.shape) + d for d in dims] + dims = tuple(d if d >= 0 else len(x.shape) + d for d in dims) if len(set(dims)) != len(dims): raise ValueError(f"Duplicate dimensions: {dims}") assert all(0 <= d < len(x.shape) for d in dims) @@ -130,7 +132,7 @@ def _check_shapes( h_prefix: str = "b", r_prefix: str = "b", t_prefix: str = "b", - raise_on_error: bool = True, + raise_on_errors: bool = True, ) -> bool: entity_shape = self.entity_shape if isinstance(entity_shape, str): @@ -144,22 +146,27 @@ def _check_shapes( if isinstance(tail_entity_shape, str): tail_entity_shape = (tail_entity_shape,) if len(h) != len(entity_shape): - if raise_on_error: + if raise_on_errors: raise ValueError return False if len(r) != len(relation_shape): - if raise_on_error: + if raise_on_errors: raise ValueError return False if len(t) != len(tail_entity_shape): - if raise_on_error: + if raise_on_errors: raise ValueError return False - return check_shapes(*itertools.chain( - ((hh, h_prefix + hs) for hh, hs in zip(h, entity_shape)), - ((rr, r_prefix + rs) for rr, rs in zip(r, relation_shape)), - ((tt, t_prefix + ts) for tt, ts in zip(t, tail_entity_shape)), - ), raise_or_error=raise_on_error) + + # TODO make helper function + unit test + a = ((hh, h_prefix + hs) for hh, hs in zip(h, entity_shape)) # type: ignore + b = ((rr, r_prefix + rs) for rr, rs in zip(r, relation_shape)) # type: ignore + c = ((tt, t_prefix + ts) for tt, ts in zip(t, tail_entity_shape)) # type: ignore + + return check_shapes( + *itertools.chain(a, b, c), + raise_on_errors=raise_on_errors, + ) def _score( self, @@ -173,11 +180,12 @@ def _score( ) -> torch.FloatTensor: assert {h_prefix, r_prefix, t_prefix}.issubset(list("bn")) # at most one of h_prefix, r_prefix, t_prefix equals n - slice_dim = [dim for dim, prefix in zip("hrt", (h_prefix, r_prefix, t_prefix)) if prefix == "n"] - if len(slice_dim) == 1: - slice_dim = slice_dim[0] - else: - slice_dim = None + slice_dims: List[str] = [ + dim + for dim, prefix in zip("hrt", (h_prefix, r_prefix, t_prefix)) + if prefix == "n" + ] + slice_dim: Optional[str] = slice_dims[0] if len(slice_dims) == 1 else None h, r, t = _ensure_tuple(h, r, t) assert self._check_shapes(h=h, r=r, t=t, h_prefix=h_prefix, r_prefix=r_prefix, t_prefix=t_prefix) @@ -320,7 +328,7 @@ def reset_parameters(self): mod.reset_parameters() -class StatelessInteractionFunction(InteractionFunction[HeadRepresentation, RelationRepresentation, TailRepresentation]): +class StatelessInteraction(Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation]): """Interaction function without state.""" def __init__(self, f: Callable[..., torch.FloatTensor]): @@ -338,7 +346,10 @@ def forward( return self.f(*h, *r, *t) -class TranslationalInteractionFunction(InteractionFunction[HeadRepresentation, RelationRepresentation, TailRepresentation], ABC): +class TranslationalInteraction( + Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], + ABC, +): """The translational interaction function shared by the TransE, TransR, TransH, and other Trans models.""" def __init__(self, p: int, power_norm: bool = False): @@ -355,7 +366,7 @@ def __init__(self, p: int, power_norm: bool = False): self.power_norm = power_norm -class TransEInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): +class TransEInteraction(TranslationalInteraction[FloatTensor, FloatTensor, FloatTensor]): """The TransE interaction function.""" def forward( @@ -367,7 +378,7 @@ def forward( return pkf.transe_interaction(h=h, r=r, t=t, p=self.p, power_norm=self.power_norm) -class ComplExInteractionFunction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): +class ComplExInteraction(StatelessInteraction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function of ComplEx.""" def __init__(self): @@ -408,14 +419,14 @@ def _calculate_missing_shape_information( input_channels = 1 # input channels is not None, and one of height or width is None - assert len([factor for factor in [input_channels, width, height] if factor is None]) <= 1 - if width is None: + if width is None and height is not None and input_channels is not None: width = embedding_dim // (height * input_channels) - if height is None: + elif height is None and width is not None and input_channels is not None: height = embedding_dim // (width * input_channels) - if input_channels is None: + elif input_channels is None and height is not None and width is not None: input_channels = embedding_dim // (width * height) - assert not any(factor is None for factor in [input_channels, width, height]) + else: + raise ValueError('More than one of width, height, and input_channels was None') if input_channels * width * height != embedding_dim: raise ValueError(f'Could not resolve {original} to a valid factorization of {embedding_dim}.') @@ -423,7 +434,9 @@ def _calculate_missing_shape_information( return input_channels, width, height -class ConvEInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor]]): +class ConvEInteraction( + Interaction[torch.FloatTensor, torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor]], +): """ConvE interaction function.""" tail_entity_shape = ("d", "k") # with k=1 @@ -528,7 +541,7 @@ def forward( ) -class ConvKBInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): +class ConvKBInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function of ConvKB.""" def __init__( @@ -575,14 +588,14 @@ def forward( ) -class DistMultInteractionFunction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): +class DistMultInteraction(StatelessInteraction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function of DistMult.""" def __init__(self): super().__init__(f=pkf.distmult_interaction) -class ERMLPInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): +class ERMLPInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): """ Interaction function of ER-MLP. @@ -639,7 +652,7 @@ def reset_parameters(self): # noqa: D102 ) -class ERMLPEInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): +class ERMLPEInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function of ER-MLP.""" def __init__( @@ -671,7 +684,13 @@ def forward( return pkf.ermlpe_interaction(h=h, r=r, t=t, mlp=self.mlp) -class TransRInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor], torch.FloatTensor]): +class TransRInteraction( + TranslationalInteraction[ + torch.FloatTensor, + Tuple[torch.FloatTensor, torch.FloatTensor], + torch.FloatTensor, + ], +): """The TransR interaction function.""" relation_shape = ("e", "de") @@ -689,21 +708,21 @@ def forward( return pkf.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p, power_norm=self.power_norm) -class RotatEInteractionFunction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): +class RotatEInteraction(StatelessInteraction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function of RotatE.""" def __init__(self): super().__init__(f=pkf.rotate_interaction) -class HolEInteractionFunction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): +class HolEInteraction(StatelessInteraction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function for HolE.""" def __init__(self): super().__init__(f=pkf.hole_interaction) -class ProjEInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): +class ProjEInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function for ProjE.""" def __init__( @@ -747,7 +766,7 @@ def forward( ) -class RESCALInteractionFunction(StatelessInteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): +class RESCALInteraction(StatelessInteraction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function of RESCAL.""" relation_shape = ("dd",) @@ -756,7 +775,13 @@ def __init__(self): super().__init__(f=pkf.rescal_interaction) -class StructuredEmbeddingInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor], torch.FloatTensor]): +class StructuredEmbeddingInteraction( + TranslationalInteraction[ + torch.FloatTensor, + Tuple[torch.FloatTensor, torch.FloatTensor], + torch.FloatTensor, + ], +): """Interaction function of Structured Embedding.""" relation_shape = ("dd", "dd") @@ -771,7 +796,7 @@ def forward( return pkf.structured_embedding_interaction(h=h, r_h=rh, r_t=rt, t=t, p=self.p, power_norm=self.power_norm) -class TuckerInteractionFunction(InteractionFunction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): +class TuckerInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function of Tucker.""" def __init__( @@ -838,11 +863,13 @@ def forward( ) -class UnstructuredModelInteractionFunction(TranslationalInteractionFunction[torch.FloatTensor, None, torch.FloatTensor]): +class UnstructuredModelInteraction( + TranslationalInteraction[torch.FloatTensor, None, torch.FloatTensor], +): """Interaction function of UnstructuredModel.""" # shapes - relation_shape = tuple() + relation_shape: Sequence[str] = tuple() def __init__(self, p: int, power_norm: bool = True): super().__init__(p=p, power_norm=power_norm) @@ -856,12 +883,12 @@ def forward( return pkf.unstructured_model_interaction(h, t, p=self.p, power_norm=self.power_norm) -class TransDInteractionFunction( - TranslationalInteractionFunction[ +class TransDInteraction( + TranslationalInteraction[ Tuple[torch.FloatTensor, torch.FloatTensor], Tuple[torch.FloatTensor, torch.FloatTensor], Tuple[torch.FloatTensor, torch.FloatTensor], - ] + ], ): """Interaction function of TransD.""" @@ -883,12 +910,12 @@ def forward( return pkf.transd_interaction(h=h, r=r, t=t, h_p=h_p, r_p=r_p, t_p=t_p, p=self.p, power_norm=self.power_norm) -class NTNInteractionFunction( - InteractionFunction[ +class NTNInteraction( + Interaction[ torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], torch.FloatTensor, - ] + ], ): """The interaction function of NTN.""" @@ -908,17 +935,17 @@ def forward( h: torch.FloatTensor, r: Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], t: torch.FloatTensor, - ) -> torch.FloatTensor: + ) -> torch.FloatTensor: # noqa:D102 w, b, u, vh, vt = r return pkf.ntn_interaction(h=h, t=t, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity) -class KG2EInteractionFunction( - InteractionFunction[ +class KG2EInteraction( + Interaction[ Tuple[torch.FloatTensor, torch.FloatTensor], Tuple[torch.FloatTensor, torch.FloatTensor], Tuple[torch.FloatTensor, torch.FloatTensor], - ] + ], ): """Interaction function of KG2E.""" @@ -942,7 +969,7 @@ def forward( h: HeadRepresentation, r: RelationRepresentation, t: TailRepresentation, - ) -> torch.FloatTensor: + ) -> torch.FloatTensor: # noqa:D102 h_mean, h_var = h r_mean, r_var = r t_mean, t_var = t @@ -958,13 +985,7 @@ def forward( ) -class TransHInteractionFunction( - TranslationalInteractionFunction[ - torch.FloatTensor, - Tuple[torch.FloatTensor, torch.FloatTensor], - torch.FloatTensor, - ] -): +class TransHInteraction(TranslationalInteraction[FloatTensor, Tuple[FloatTensor, FloatTensor], FloatTensor]): """Interaction function of TransH.""" relation_shape = ("d", "d") diff --git a/src/pykeen/typing.py b/src/pykeen/typing.py index a05b1e8e9f..eef6133202 100644 --- a/src/pykeen/typing.py +++ b/src/pykeen/typing.py @@ -12,6 +12,9 @@ 'MappedTriples', 'EntityMapping', 'RelationMapping', + 'Initializer', + 'Normalizer', + 'Constrainer', 'DeviceHint', 'HeadRepresentation', 'RelationRepresentation', @@ -33,6 +36,6 @@ DeviceHint = Union[None, str, torch.device] Representation = torch.FloatTensor -HeadRepresentation = TypeVar("HeadRepresentation", Representation, Sequence[Representation]) -RelationRepresentation = TypeVar("RelationRepresentation", Representation, Sequence[Representation]) -TailRepresentation = TypeVar("TailRepresentation", Representation, Sequence[Representation]) +HeadRepresentation = TypeVar("HeadRepresentation", Representation, Sequence[Representation]) # type: ignore +RelationRepresentation = TypeVar("RelationRepresentation", Representation, Sequence[Representation]) # type: ignore +TailRepresentation = TypeVar("TailRepresentation", Representation, Sequence[Representation]) # type: ignore diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index c4c84ee5b5..f8a1537f28 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -251,13 +251,12 @@ def chained_op(x: torch.Tensor): return chained_op -def set_random_seed(seed: int): +def set_random_seed(seed: int) -> Tuple[None, torch._C.Generator, None]: """Set the random seed on numpy, torch, and python.""" - return ( - np.random.seed(seed=seed), - torch.manual_seed(seed=seed), - random.seed(seed), - ) + np.random.seed(seed=seed) + generator = torch.manual_seed(seed=seed) + random.seed(seed) + return None, generator, None class NoRandomSeedNecessary: @@ -441,7 +440,7 @@ def random_non_negative_int() -> int: def check_shapes( *x: Tuple[torch.Tensor, str], - raise_or_error: bool = True, + raise_on_errors: bool = True, ) -> bool: """ Verify that a sequence of tensors are of matching shapes. @@ -450,7 +449,7 @@ def check_shapes( A tuple (tensor, shape), where tensor is a tensor, and shape is a string, where each character corresponds to a (named) dimension. If the shapes of different tensors share a character, the corresponding dimensions are expected to be of equal size. - :param raise_or_error: + :param raise_on_errors: Whether to raise an exception in case of a mismatch. :return: @@ -459,7 +458,7 @@ def check_shapes( :raises ValueError: If the shapes mismatch and raise_on_error is True. """ - dims = dict() + dims: Dict[str, Tuple[int, ...]] = dict() errors = [] for tensor, shape in x: if tensor.ndimension() != len(shape): @@ -470,7 +469,7 @@ def check_shapes( if exp_dim is not None and exp_dim != dim: errors.append(f"{name}: {dim} vs. {exp_dim}") dims[name] = dim - if raise_or_error and len(errors) > 0: + if raise_on_errors and errors: raise ValueError("Shape verification failed:\n" + '\n'.join(errors)) return len(errors) == 0 diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 436efff9f2..4827017d32 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -1,4 +1,7 @@ +# -*- coding: utf-8 -*- + """Tests for interaction functions.""" + import unittest from typing import Any, Collection, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union from unittest.case import SkipTest @@ -6,7 +9,7 @@ import torch import pykeen.nn.modules -from pykeen.nn.modules import InteractionFunction, StatelessInteractionFunction, TranslationalInteractionFunction +from pykeen.nn.modules import Interaction, StatelessInteraction, TranslationalInteraction from pykeen.typing import Representation from pykeen.utils import get_subclasses @@ -21,6 +24,7 @@ class GenericTests(Generic[T]): instance: T def setUp(self) -> None: + """Set up the generic testing method.""" kwargs = self.kwargs or {} kwargs = self._pre_instantiation_hook(kwargs=dict(kwargs)) self.instance = self.cls(**kwargs) @@ -49,7 +53,7 @@ def test_testing(self): assert not_tested == set() -class InteractionTests(GenericTests[pykeen.nn.modules.InteractionFunction]): +class InteractionTests(GenericTests[pykeen.nn.modules.Interaction]): """Generic test for interaction functions.""" dim: int = 2 @@ -206,7 +210,7 @@ def test_forward(self): def test_scores(self): """Test individual scores.""" self.instance.eval() - for i in range(10): + for _ in range(10): h, r, t = self._get_hrt((1, 1), (1, 1), (1, 1)) scores = self.instance(h=h, r=r, t=t) exp_score = self._exp_score(h, r, t).item() @@ -219,13 +223,13 @@ def _exp_score(self, h, r, t) -> torch.FloatTensor: class ComplExTests(InteractionTests, unittest.TestCase): """Tests for ComplEx interaction function.""" - cls = pykeen.nn.modules.ComplExInteractionFunction + cls = pykeen.nn.modules.ComplExInteraction class ConvETests(InteractionTests, unittest.TestCase): """Tests for ConvE interaction function.""" - cls = pykeen.nn.modules.ConvEInteractionFunction + cls = pykeen.nn.modules.ConvEInteraction kwargs = dict( embedding_height=1, embedding_width=2, @@ -237,7 +241,7 @@ class ConvETests(InteractionTests, unittest.TestCase): def _get_hrt( self, *shapes: Tuple[int, ...], - **kwargs + **kwargs, ) -> Tuple[Union[Representation, Sequence[Representation]], ...]: h, r, t = super()._get_hrt(*shapes, **kwargs) t_bias = torch.rand_like(t[..., 0, None]) @@ -247,7 +251,7 @@ def _get_hrt( class ConvKBTests(InteractionTests, unittest.TestCase): """Tests for ConvKB interaction function.""" - cls = pykeen.nn.modules.ConvKBInteractionFunction + cls = pykeen.nn.modules.ConvKBInteraction kwargs = dict( embedding_dim=InteractionTests.dim, num_filters=2 * InteractionTests.dim - 1, @@ -257,7 +261,7 @@ class ConvKBTests(InteractionTests, unittest.TestCase): class DistMultTests(InteractionTests, unittest.TestCase): """Tests for DistMult interaction function.""" - cls = pykeen.nn.modules.DistMultInteractionFunction + cls = pykeen.nn.modules.DistMultInteraction def _exp_score(self, h, r, t) -> torch.FloatTensor: return (h * r * t).sum(dim=-1) @@ -266,7 +270,7 @@ def _exp_score(self, h, r, t) -> torch.FloatTensor: class ERMLPTests(InteractionTests, unittest.TestCase): """Tests for ERMLP interaction function.""" - cls = pykeen.nn.modules.ERMLPInteractionFunction + cls = pykeen.nn.modules.ERMLPInteraction kwargs = dict( embedding_dim=InteractionTests.dim, hidden_dim=2 * InteractionTests.dim - 1, @@ -276,7 +280,7 @@ class ERMLPTests(InteractionTests, unittest.TestCase): class ERMLPETests(InteractionTests, unittest.TestCase): """Tests for ERMLP-E interaction function.""" - cls = pykeen.nn.modules.ERMLPEInteractionFunction + cls = pykeen.nn.modules.ERMLPEInteraction kwargs = dict( embedding_dim=InteractionTests.dim, hidden_dim=2 * InteractionTests.dim - 1, @@ -286,13 +290,13 @@ class ERMLPETests(InteractionTests, unittest.TestCase): class HolETests(InteractionTests, unittest.TestCase): """Tests for HolE interaction function.""" - cls = pykeen.nn.modules.HolEInteractionFunction + cls = pykeen.nn.modules.HolEInteraction class NTNTests(InteractionTests, unittest.TestCase): """Tests for NTN interaction function.""" - cls = pykeen.nn.modules.NTNInteractionFunction + cls = pykeen.nn.modules.NTNInteraction num_slices: int = 2 shape_kwargs = dict( @@ -303,7 +307,7 @@ class NTNTests(InteractionTests, unittest.TestCase): class ProjETests(InteractionTests, unittest.TestCase): """Tests for ProjE interaction function.""" - cls = pykeen.nn.modules.ProjEInteractionFunction + cls = pykeen.nn.modules.ProjEInteraction kwargs = dict( embedding_dim=InteractionTests.dim, ) @@ -312,19 +316,19 @@ class ProjETests(InteractionTests, unittest.TestCase): class RESCALTests(InteractionTests, unittest.TestCase): """Tests for RESCAL interaction function.""" - cls = pykeen.nn.modules.RESCALInteractionFunction + cls = pykeen.nn.modules.RESCALInteraction class KG2ETests(InteractionTests, unittest.TestCase): """Tests for KG2E interaction function.""" - cls = pykeen.nn.modules.KG2EInteractionFunction + cls = pykeen.nn.modules.KG2EInteraction class TuckerTests(InteractionTests, unittest.TestCase): """Tests for Tucker interaction function.""" - cls = pykeen.nn.modules.TuckerInteractionFunction + cls = pykeen.nn.modules.TuckerInteraction kwargs = dict( embedding_dim=InteractionTests.dim, ) @@ -333,7 +337,7 @@ class TuckerTests(InteractionTests, unittest.TestCase): class RotatETests(InteractionTests, unittest.TestCase): """Tests for RotatE interaction function.""" - cls = pykeen.nn.modules.RotatEInteractionFunction + cls = pykeen.nn.modules.RotatEInteraction class TranslationalInteractionTests(InteractionTests): @@ -350,7 +354,7 @@ def _additional_score_checks(self, scores): class TransDTests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransD interaction function.""" - cls = pykeen.nn.modules.TransDInteractionFunction + cls = pykeen.nn.modules.TransDInteraction shape_kwargs = dict( e=3, ) @@ -388,7 +392,7 @@ def test_manual_big_relation_dim(self): class TransETests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransE interaction function.""" - cls = pykeen.nn.modules.TransEInteractionFunction + cls = pykeen.nn.modules.TransEInteraction def _exp_score(self, h, r, t) -> torch.FloatTensor: return -(h + r - t).norm(p=2, dim=-1) @@ -397,13 +401,13 @@ def _exp_score(self, h, r, t) -> torch.FloatTensor: class TransHTests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransH interaction function.""" - cls = pykeen.nn.modules.TransHInteractionFunction + cls = pykeen.nn.modules.TransHInteraction class TransRTests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransR interaction function.""" - cls = pykeen.nn.modules.TransRInteractionFunction + cls = pykeen.nn.modules.TransRInteraction shape_kwargs = dict( e=3, ) @@ -423,21 +427,21 @@ def test_manual(self): class SETests(TranslationalInteractionTests, unittest.TestCase): """Tests for SE interaction function.""" - cls = pykeen.nn.modules.StructuredEmbeddingInteractionFunction + cls = pykeen.nn.modules.StructuredEmbeddingInteraction class UMTests(TranslationalInteractionTests, unittest.TestCase): """Tests for UM interaction function.""" - cls = pykeen.nn.modules.UnstructuredModelInteractionFunction + cls = pykeen.nn.modules.UnstructuredModelInteraction -class InteractionTestsTest(TestsTest[InteractionFunction], unittest.TestCase): +class InteractionTestsTest(TestsTest[Interaction], unittest.TestCase): """Test for tests for all interaction functions.""" - base_cls = InteractionFunction + base_cls = Interaction base_test = InteractionTests skip_cls = { - TranslationalInteractionFunction, - StatelessInteractionFunction, + TranslationalInteraction, + StatelessInteraction, } diff --git a/tests/test_models.py b/tests/test_models.py index 65cd4cc904..5905b3f6a7 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -654,7 +654,11 @@ def _check_constraints(self): Entity embeddings have to have at most unit L2 norm. """ - assert all_in_bounds(self.model.entity_representations[0](indices=None).norm(p=2, dim=-1), high=1., a_tol=_EPSILON) + assert all_in_bounds( + self.model.entity_representations[0](indices=None).norm(p=2, dim=-1), + high=1., + a_tol=_EPSILON, + ) class _TestKG2E(_ModelTestCase): diff --git a/tox.ini b/tox.ini index d3e54783cf..ef61f1d1b1 100644 --- a/tox.ini +++ b/tox.ini @@ -35,6 +35,7 @@ whitelist_externals = /bin/cat /bin/cp /bin/mkdir + /bin/rm /usr/bin/cat /usr/bin/cp /usr/bin/mkdir @@ -109,7 +110,12 @@ description = Check all python files do not have mistaken trailing commas [testenv:mypy] deps = mypy skip_install = true -commands = mypy --ignore-missing-imports src/pykeen/ +commands = + mypy --ignore-missing-imports \ + src/pykeen/typing.py \ + src/pykeen/utils.py \ + src/pykeen/nn + description = Run the mypy tool to check static typing on the project. [testenv:pyroma] @@ -143,6 +149,7 @@ extras = commands = mkdir -p {envtmpdir} cp -r source {envtmpdir}/source + rm -rf source/api sphinx-build -W -b html -d {envtmpdir}/build/doctrees {envtmpdir}/source {envtmpdir}/build/html sphinx-build -W -b coverage -d {envtmpdir}/build/doctrees {envtmpdir}/source {envtmpdir}/build/coverage cat {envtmpdir}/build/coverage/c.txt From a14eb5e02f2ebbeb2e0628520e305373b16b3a14 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sun, 15 Nov 2020 19:56:50 +0100 Subject: [PATCH 286/690] More typing --- src/pykeen/models/__init__.py | 10 ++++- src/pykeen/models/base.py | 40 ++++++++++--------- .../models/multimodal/complex_literal.py | 4 +- src/pykeen/models/unimodal/conv_e.py | 7 ++-- src/pykeen/models/unimodal/ermlp.py | 5 ++- src/pykeen/models/unimodal/ermlpe.py | 6 +-- src/pykeen/models/unimodal/hole.py | 2 +- src/pykeen/models/unimodal/kg2e.py | 8 ++-- src/pykeen/models/unimodal/rgcn.py | 2 +- src/pykeen/models/unimodal/rotate.py | 4 +- src/pykeen/models/unimodal/simple.py | 4 +- src/pykeen/models/unimodal/trans_d.py | 4 +- src/pykeen/models/unimodal/trans_e.py | 4 +- src/pykeen/models/unimodal/trans_r.py | 4 +- src/pykeen/models/unimodal/tucker.py | 4 +- src/pykeen/nn/functional.py | 14 +++---- src/pykeen/nn/modules.py | 16 ++++---- src/pykeen/regularizers.py | 12 +++--- src/pykeen/triples/triples_factory.py | 17 +++++--- .../triples_numeric_literals_factory.py | 12 ++++-- tox.ini | 8 +++- 21 files changed, 108 insertions(+), 79 deletions(-) diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index 5273620aeb..d2ab204f80 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -9,8 +9,8 @@ from typing import Mapping, Set, Type, Union from .base import ( # noqa:F401 - EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, - MultimodalModel, SingleVectorEmbeddingModel, + DoubleRelationEmbeddingModel, EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, MultimodalModel, + SingleVectorEmbeddingModel, TwoSideEmbeddingModel, TwoVectorEmbeddingModel, ) from .multimodal import ComplExLiteral, DistMultLiteral from .unimodal import ( @@ -68,6 +68,12 @@ _CONCRETE_BASES = { SingleVectorEmbeddingModel, + DoubleRelationEmbeddingModel, + TwoSideEmbeddingModel, + EntityEmbeddingModel, + EntityRelationEmbeddingModel, + TwoVectorEmbeddingModel, + MultimodalModel, } diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 6ad5f8c02a..d50409d92a 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -212,7 +212,8 @@ def _new_init(self, *args, **kwargs): _original_init(self, *args, **kwargs) self.reset_parameters_() - cls.__init__ = _new_init + # sorry mypy, but this kind of evil must be permitted. + cls.__init__ = _new_init # type: ignore class Model(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation], ABC): @@ -231,8 +232,9 @@ class Model(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailR #: The instance of the loss loss: Loss + # FIXME problem with typing (remove type: ignore) #: The default regularizer class - regularizer_default: ClassVar[Type[Regularizer]] = NoRegularizer + regularizer_default: ClassVar[Type[Regularizer]] = NoRegularizer # type: ignore #: The default parameters for the default regularizer class regularizer_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = None #: The instance of the regularizer @@ -255,8 +257,7 @@ def __init__( regularizer: Optional[Regularizer] = None, entity_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, relation_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, - interaction_function: Interaction[ - HeadRepresentation, RelationRepresentation, TailRepresentation] = None, + interaction_function: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation] = None, ) -> None: """Initialize the module. @@ -295,12 +296,12 @@ def __init__( # Loss if loss is None: - self.loss = self.loss_default(**self.loss_default_kwargs) + self.loss = self.loss_default(**(self.loss_default_kwargs or {})) else: self.loss = loss # TODO: Check loss functions that require 1 and -1 as label but only - self.is_mr_loss = isinstance(self.loss, MarginRankingLoss) + self.is_mr_loss: bool = isinstance(self.loss, MarginRankingLoss) # Regularizer if regularizer is None: @@ -310,7 +311,7 @@ def __init__( ) self.regularizer = regularizer - self.is_nssa_loss = isinstance(self.loss, NSSALoss) + self.is_nssa_loss: bool = isinstance(self.loss, NSSALoss) # The triples factory facilitates access to the dataset. self.triples_factory = triples_factory @@ -495,10 +496,10 @@ def predict_scores_all_tails( """ # Enforce evaluation mode self.eval() - if slice_size is None: - scores = self.score_t(hr_batch) + if slice_size is not None and self.can_slice_t: + scores = self.score_t(hr_batch, slice_size=slice_size) # type: ignore else: - scores = self.score_t(hr_batch, slice_size=slice_size) + scores = self.score_t(hr_batch) if self.predict_with_sigmoid: scores = torch.sigmoid(scores) return scores @@ -634,10 +635,10 @@ def predict_scores_all_relations( """ # Enforce evaluation mode self.eval() - if slice_size is None: - scores = self.score_r(ht_batch) + if slice_size is not None and self.can_slice_r: + scores = self.score_r(ht_batch, slice_size=slice_size) # type: ignore else: - scores = self.score_r(ht_batch, slice_size=slice_size) + scores = self.score_r(ht_batch) if self.predict_with_sigmoid: scores = torch.sigmoid(scores) return scores @@ -670,10 +671,10 @@ def predict_scores_all_heads( for a (tail, inverse_relation) pair. ''' if not self.triples_factory.create_inverse_triples: - if slice_size is None: - scores = self.score_h(rt_batch) + if slice_size is not None and self.can_slice_h: + scores = self.score_h(rt_batch, slice_size=slice_size) # type: ignore else: - scores = self.score_h(rt_batch, slice_size=slice_size) + scores = self.score_h(rt_batch) if self.predict_with_sigmoid: scores = torch.sigmoid(scores) return scores @@ -695,10 +696,11 @@ def predict_scores_all_heads( # The score_t function requires (entity, relation) pairs instead of (relation, entity) pairs rt_batch_cloned = rt_batch_cloned.flip(1) - if slice_size is None: - scores = self.score_t(rt_batch_cloned) + if slice_size is not None and self.can_slice_t: + scores = self.score_t(rt_batch_cloned, slice_size=slice_size) # type: ignore else: - scores = self.score_t(rt_batch_cloned, slice_size=slice_size) + scores = self.score_t(rt_batch_cloned) + if self.predict_with_sigmoid: scores = torch.sigmoid(scores) return scores diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index 9e7d585ea2..53a325ac91 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -2,7 +2,7 @@ """Implementation of the ComplexLiteral model based on the local closed world assumption (LCWA) training approach.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional import torch import torch.nn as nn @@ -31,7 +31,7 @@ class ComplExLiteral(MultimodalModel): #: The default loss function class loss_default = BCEWithLogitsLoss #: The default parameters for the default loss function class - loss_default_kwargs = {} + loss_default_kwargs: ClassVar[Mapping[str, Any]] = {} def __init__( self, diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 8785b2bf53..3b4a24131a 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -1,8 +1,9 @@ # -*- coding: utf-8 -*- """Implementation of ConvE.""" + import logging -from typing import Optional, Type +from typing import Any, ClassVar, Mapping, Optional, Type import torch from torch import nn @@ -102,9 +103,9 @@ class ConvE(Model): feature_map_dropout=dict(type=float, low=0.0, high=1.0), ) #: The default loss function class - loss_default: Type[Loss] = BCEAfterSigmoidLoss + loss_default: ClassVar[Type[Loss]] = BCEAfterSigmoidLoss #: The default parameters for the default loss function class - loss_default_kwargs = {} + loss_default_kwargs: ClassVar[Mapping[str, Any]] = {} def __init__( self, diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 0c4b2f6578..b326e35e05 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -49,7 +49,10 @@ def __init__( hidden_dim: Optional[int] = None, regularizer: Optional[Regularizer] = None, ) -> None: - """Initialize the model.""" + """Initialize ERMLP.""" + if hidden_dim is None: + hidden_dim = embedding_dim + super().__init__( triples_factory=triples_factory, interaction_function=ERMLPInteraction( diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index 66f3e5066c..0f4c163901 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -2,7 +2,7 @@ """An implementation of the extension to ERMLP.""" -from typing import Optional, Type +from typing import Any, ClassVar, Mapping, Optional, Type from ..base import SingleVectorEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss @@ -47,9 +47,9 @@ class ERMLPE(SingleVectorEmbeddingModel): hidden_dropout=dict(type=float, low=0.0, high=0.8, q=0.1), ) #: The default loss function class - loss_default: Type[Loss] = BCEAfterSigmoidLoss + loss_default: ClassVar[Type[Loss]] = BCEAfterSigmoidLoss #: The default parameters for the default loss function class - loss_default_kwargs = {} + loss_default_kwargs: ClassVar[Mapping[str, Any]] = {} def __init__( self, diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 01f9ab9e7c..611b942a97 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -77,7 +77,7 @@ def __init__( # Initialisation, cf. https://github.com/mnick/scikit-kge/blob/master/skge/param.py#L18-L27 embedding_specification=EmbeddingSpecification( initializer=xavier_uniform_, - constrainer=clamp_norm, + constrainer=clamp_norm, # type: ignore constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), relation_embedding_specification=EmbeddingSpecification( diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 53b8d3e38c..3e93c150ba 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -11,7 +11,7 @@ from ..base import TwoVectorEmbeddingModel from ...losses import Loss from ...nn.emb import EmbeddingSpecification -from ...nn.modules import KG2EInteractionFunction +from ...nn.modules import KG2EInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -80,7 +80,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - interaction_function=KG2EInteractionFunction( + interaction_function=KG2EInteraction( similarity=dist_similarity, ), embedding_dim=embedding_dim, @@ -90,11 +90,11 @@ def __init__( random_seed=random_seed, regularizer=regularizer, embedding_specification=EmbeddingSpecification( - constrainer=clamp_norm, + constrainer=clamp_norm, # type: ignore constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), relation_embedding_specification=EmbeddingSpecification( - constrainer=clamp_norm, + constrainer=clamp_norm, # type: ignore constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), # Ensure positive definite covariances matrices and appropriate size by clamping diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index 5ba8c67141..69d901ec77 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -340,7 +340,7 @@ def forward( # Normalize messages by relation-specific in-degree if self.edge_weighting is not None: - m_r *= self.edge_weighting(source=sources_r, target=targets_r).unsqueeze(dim=-1) + m_r *= self.edge_weighting(sources_r, targets_r).unsqueeze(dim=-1) # Aggregate messages in target new_x.index_add_(dim=0, index=targets_r, source=m_r) diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index 406428ee2b..a5a35f172a 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -13,7 +13,7 @@ from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ -from ...nn.modules import RotatEInteractionFunction +from ...nn.modules import RotatEInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -96,7 +96,7 @@ def __init__( # TODO: regularization super().__init__( triples_factory=triples_factory, - interaction_function=RotatEInteractionFunction(), + interaction_function=RotatEInteraction(), embedding_dim=2 * embedding_dim, loss=loss, automatic_memory_optimization=automatic_memory_optimization, diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index 84b3a602b4..e0f0978cac 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -2,7 +2,7 @@ """Implementation of SimplE.""" -from typing import Optional, Tuple, Union +from typing import Any, ClassVar, Mapping, Optional, Tuple, Union import torch.autograd @@ -50,7 +50,7 @@ class SimplE(TwoSideEmbeddingModel): #: The default loss function class loss_default = SoftplusLoss #: The default parameters for the default loss function class - loss_default_kwargs = {} + loss_default_kwargs: ClassVar[Mapping[str, Any]] = {} #: The regularizer used by [trouillon2016]_ for SimplE #: In the paper, they use weight of 0.1, and do not normalize the #: regularization term by the number of elements, which is 200. diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index bd62e10fc0..69a27e42fb 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -80,12 +80,12 @@ def __init__( regularizer=regularizer, embedding_specification=EmbeddingSpecification( initializer=xavier_normal_, - constrainer=clamp_norm, + constrainer=clamp_norm, # type: ignore constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), relation_embedding_specification=EmbeddingSpecification( initializer=xavier_normal_, - constrainer=clamp_norm, + constrainer=clamp_norm, # type: ignore constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), ) diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 499ec89f22..197fc2676b 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -10,7 +10,7 @@ from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ -from ...nn.modules import TransEInteractionFunction +from ...nn.modules import TransEInteraction from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -69,7 +69,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - interaction_function=TransEInteractionFunction(p=scoring_fct_norm, power_norm=False), + interaction_function=TransEInteraction(p=scoring_fct_norm, power_norm=False), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 4390a31a55..1a1661461f 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -85,7 +85,7 @@ def __init__( shape=embedding_dim, specification=EmbeddingSpecification( initializer=xavier_uniform_, - constrainer=clamp_norm, + constrainer=clamp_norm, # type: ignore constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), ), @@ -98,7 +98,7 @@ def __init__( xavier_uniform_, functional.normalize, ), - constrainer=clamp_norm, + constrainer=clamp_norm, # type: ignore constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), ), diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index c8f6bbd3d8..b7390a9271 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -2,7 +2,7 @@ """Implementation of TuckEr.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional from .. import SingleVectorEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss @@ -61,7 +61,7 @@ class TuckER(SingleVectorEmbeddingModel): #: The default loss function class loss_default = BCEAfterSigmoidLoss #: The default parameters for the default loss function class - loss_default_kwargs = {} + loss_default_kwargs: ClassVar[Mapping[str, Any]] = {} def __init__( self, diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 3bc8bd9bd4..bf8df457fe 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -196,7 +196,7 @@ def _kullback_leibler_similarity( ) -> torch.FloatTensor: r"""Compute the similarity based on KL divergence. - This is done between two Gaussian distributions given by mean mu_* and diagonal covariance matrix sigma_*. + This is done between two Gaussian distributions given by mean `mu_*` and diagonal covariance matrix `sigma_*`. .. math:: @@ -251,10 +251,10 @@ def _kullback_leibler_similarity( return sim -KG2E_SIMILARITIES = dict( - KL=_kullback_leibler_similarity, - EL=_expected_likelihood, -) +KG2E_SIMILARITIES = { + 'KL': _kullback_leibler_similarity, + 'EL': _expected_likelihood, +} def _extended_einsum( @@ -346,7 +346,7 @@ def complex_interaction( """ Evaluate the ComplEx interaction function. - :param h: shape: (batch_size, num_heads, 2*dim) + :param h: shape: (batch_size, num_heads, `2*dim`) The complex head representations. :param r: shape: (batch_size, num_relations, 2*dim) The complex relation representations. @@ -1049,7 +1049,7 @@ def transh_interaction( :param p: The p for the norm. cf. torch.norm. :param power_norm: - Whether to return |x-y|_p^p. + Whether to return $|x-y|_p^p$. :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 0755f3d976..b707321629 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -12,7 +12,6 @@ from torch import FloatTensor, nn from . import functional as pkf -from .functional import KG2E_SIMILARITIES from ..typing import HeadRepresentation, RelationRepresentation, Representation, TailRepresentation from ..utils import check_shapes @@ -74,11 +73,11 @@ def forward( ) -> torch.FloatTensor: """Compute broadcasted triple scores given representations for head, relation and tails. - :param h: shape: (batch_size, num_heads, *) + :param h: shape: (batch_size, num_heads, ``*``) The head representations. - :param r: shape: (batch_size, num_relations, *) + :param r: shape: (batch_size, num_relations, ``*``) The relation representations. - :param t: shape: (batch_size, num_tails, *) + :param t: shape: (batch_size, num_tails, ``*``) The tail representations. :return: shape: (batch_size, num_heads, num_relations, num_tails) @@ -951,16 +950,17 @@ class KG2EInteraction( entity_shape = ("d", "d") relation_shape = ("d", "d") + similarity: str + exact: bool def __init__( self, - similarity: str = "KL", + similarity: Optional[str] = None, exact: bool = True, ): super().__init__() - similarity = similarity.upper() - if similarity not in KG2E_SIMILARITIES: - raise ValueError(similarity) + if similarity is None: + similarity = 'KL' self.similarity = similarity self.exact = exact diff --git a/src/pykeen/regularizers.py b/src/pykeen/regularizers.py index c79e9538da..c746c7142d 100644 --- a/src/pykeen/regularizers.py +++ b/src/pykeen/regularizers.py @@ -2,7 +2,7 @@ """Regularization in PyKEEN.""" -from abc import abstractmethod +from abc import ABC, abstractmethod from typing import Any, ClassVar, Collection, Iterable, Mapping, Optional, Type, Union import torch @@ -24,7 +24,7 @@ _REGULARIZER_SUFFIX = 'Regularizer' -class Regularizer(nn.Module): +class Regularizer(nn.Module, ABC): """A base class for all regularizers.""" #: The overall regularization weight @@ -76,7 +76,7 @@ def update(self, *tensors: torch.FloatTensor) -> None: """Update the regularization term based on passed tensors.""" if self.apply_only_once and self.updated: return - self.regularization_term = self.regularization_term + sum(self.forward(x=x) for x in tensors) + self.regularization_term = self.regularization_term + sum(self(x) for x in tensors) self.updated = True @property @@ -92,7 +92,7 @@ class NoRegularizer(Regularizer): """ #: The default strategy for optimizing the regularizer's hyper-parameters - hpo_default = {} + hpo_default: ClassVar[Mapping[str, Any]] = {} def update(self, *tensors: torch.FloatTensor) -> None: # noqa: D102 # no need to compute anything @@ -251,7 +251,7 @@ def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102 _REGULARIZERS: Collection[Type[Regularizer]] = { - NoRegularizer, + NoRegularizer, # type: ignore LpRegularizer, PowerSumRegularizer, CombinedRegularizer, @@ -269,7 +269,7 @@ def get_regularizer_cls(query: Union[None, str, Type[Regularizer]]) -> Type[Regu """Get the regularizer class.""" return get_cls( query, - base=Regularizer, + base=Regularizer, # type: ignore lookup_dict=regularizers, default=NoRegularizer, suffix=_REGULARIZER_SUFFIX, diff --git a/src/pykeen/triples/triples_factory.py b/src/pykeen/triples/triples_factory.py index b5f3b61521..c8ac719cf1 100644 --- a/src/pykeen/triples/triples_factory.py +++ b/src/pykeen/triples/triples_factory.py @@ -5,6 +5,7 @@ import logging import os import re +import typing from collections import Counter, defaultdict from typing import Collection, Dict, Iterable, List, Mapping, Optional, Sequence, Set, TextIO, Tuple, Union @@ -283,9 +284,9 @@ def __init__( # Generate relation mapping if necessary if relation_to_id is None: if self.create_inverse_triples: - relation_to_id = create_relation_mapping( - set(self.relation_to_inverse.keys()).union(set(self.relation_to_inverse.values())), - ) + _keys = set(self.relation_to_inverse.keys()) if self.relation_to_inverse else set() + _values = set(self.relation_to_inverse.values()) if self.relation_to_inverse else set() + relation_to_id = create_relation_mapping(_keys.union(_values)) else: relation_to_id = create_relation_mapping(unique_relations) if compact_id: @@ -328,6 +329,8 @@ def get_inverse_relation_id(self, relation: str) -> int: """Get the inverse relation identifier for the given relation.""" if not self.create_inverse_triples: raise ValueError('Can not get inverse triple, they have not been created.') + if self.relation_to_inverse is None: + raise ValueError inverse_relation = self.relation_to_inverse[relation] return self.relation_to_id[inverse_relation] @@ -481,7 +484,7 @@ def get_most_frequent_relations(self, n: Union[int, float]) -> Set[str]: elif not isinstance(n, int): raise TypeError('n must be either an integer or a float') - counter = Counter(self.triples[:, 1]) + counter: typing.Counter[str] = Counter(self.triples[:, 1]) return { relation for relation, _ in counter.most_common(n) @@ -566,7 +569,7 @@ def _word_cloud(self, *, text: List[str], top: int): def tensor_to_df( self, tensor: torch.LongTensor, - **kwargs: Union[torch.Tensor, np.ndarray, Sequence], + **kwargs: Union[torch.Tensor, np.ndarray, Sequence], # FIXME fix the type annotation ) -> pd.DataFrame: """Take a tensor of triples and make a pandas dataframe with labels. @@ -605,7 +608,7 @@ def tensor_to_df( for key, values in kwargs.items(): # convert PyTorch tensors to numpy if torch.is_tensor(values): - values = values.cpu().numpy() + values = values.cpu().numpy() # type: ignore data[key] = values # convert to dataframe @@ -637,6 +640,8 @@ def new_with_restriction( 'relations as well.', str(self), ) + if self.relation_to_inverse is None: + raise ValueError relations = list(relations) + list(map(self.relation_to_inverse.__getitem__, relations)) keep_mask = None diff --git a/src/pykeen/triples/triples_numeric_literals_factory.py b/src/pykeen/triples/triples_numeric_literals_factory.py index 5af8772489..761a7755e8 100644 --- a/src/pykeen/triples/triples_numeric_literals_factory.py +++ b/src/pykeen/triples/triples_numeric_literals_factory.py @@ -47,6 +47,9 @@ def create_matrix_of_literals( class TriplesNumericLiteralsFactory(TriplesFactory): """Create multi-modal instances given the path to triples.""" + numeric_literals: np.ndarray + literals_to_id: Dict[str, int] + def __init__( self, *, @@ -79,9 +82,6 @@ def __init__( self.path_to_numeric_triples = '' self.numeric_triples = numeric_triples - self.numeric_literals = None - self.literals_to_id = None - self._create_numeric_literals() def __repr__(self): # noqa: D105 @@ -104,6 +104,9 @@ def create_slcwa_instances(self) -> MultimodalSLCWAInstances: if self.numeric_literals is None: self._create_numeric_literals() + assert self.numeric_literals is not None + assert self.literals_to_id is not None + return MultimodalSLCWAInstances( mapped_triples=slcwa_instances.mapped_triples, entity_to_id=slcwa_instances.entity_to_id, @@ -119,6 +122,9 @@ def create_lcwa_instances(self, use_tqdm: Optional[bool] = None) -> MultimodalLC if self.numeric_literals is None: self._create_numeric_literals() + assert self.numeric_literals is not None + assert self.literals_to_id is not None + return MultimodalLCWAInstances( mapped_triples=lcwa_instances.mapped_triples, entity_to_id=lcwa_instances.entity_to_id, diff --git a/tox.ini b/tox.ini index ef61f1d1b1..d5e033448c 100644 --- a/tox.ini +++ b/tox.ini @@ -12,6 +12,7 @@ envlist = flake8 darglint pyroma + mypy # documentation linters/checkers readme doc8 @@ -114,7 +115,12 @@ commands = mypy --ignore-missing-imports \ src/pykeen/typing.py \ src/pykeen/utils.py \ - src/pykeen/nn + src/pykeen/nn \ + src/pykeen/regularizers.py \ + src/pykeen/losses.py \ + src/pykeen/trackers.py \ + src/pykeen/models/base.py \ + src/pykeen/triples/triples_factory.py description = Run the mypy tool to check static typing on the project. From 79deecb560799277ba5a89dedddc42889b3a0d82 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sun, 15 Nov 2020 20:09:36 +0100 Subject: [PATCH 287/690] More type annotation fixes --- src/pykeen/models/base.py | 4 +-- .../models/multimodal/complex_literal.py | 14 ++++---- .../models/multimodal/distmult_literal.py | 6 ++++ src/pykeen/models/unimodal/rgcn.py | 32 ++++++++++--------- 4 files changed, 31 insertions(+), 25 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index d50409d92a..cd8e4ad48b 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -112,7 +112,7 @@ def get_novelty_mask( """ other_cols = sorted(set(range(mapped_triples.shape[1])).difference({col})) other_col_ids = torch.tensor(data=other_col_ids, dtype=torch.long, device=mapped_triples.device) - filter_mask = (mapped_triples[:, other_cols] == other_col_ids[None, :]).all(dim=-1) + filter_mask = (mapped_triples[:, other_cols] == other_col_ids[None, :]).all(dim=-1) # type: ignore known_ids = mapped_triples[filter_mask, col].unique().cpu().numpy() return np.isin(element=query_ids, test_elements=known_ids, assume_unique=True, invert=True) @@ -249,6 +249,7 @@ class Model(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailR def __init__( self, triples_factory: TriplesFactory, + interaction_function: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, automatic_memory_optimization: Optional[bool] = None, @@ -257,7 +258,6 @@ def __init__( regularizer: Optional[Regularizer] = None, entity_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, relation_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, - interaction_function: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation] = None, ) -> None: """Initialize the module. diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index 53a325ac91..58b18d6841 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -10,10 +10,15 @@ from ..base import MultimodalModel from ...losses import BCEWithLogitsLoss, Loss +from ...nn.modules import ComplExInteraction from ...triples import TriplesNumericLiteralsFactory from ...typing import DeviceHint from ...utils import slice_doubles +__all__ = [ + 'ComplExLiteral', +] + # TODO: Check entire build of the model class ComplExLiteral(MultimodalModel): @@ -46,6 +51,7 @@ def __init__( """Initialize the model.""" super().__init__( triples_factory=triples_factory, + interaction_function=ComplExInteraction(), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, @@ -53,11 +59,6 @@ def __init__( random_seed=random_seed, ) - self.entity_embs_real = None - self.entity_embs_img = None - self.relation_embs_real = None - self.relation_embs_img = None - # Literal # num_ent x num_lit numeric_literals = triples_factory.numeric_literals @@ -79,9 +80,6 @@ def __init__( self.inp_drop = torch.nn.Dropout(input_dropout) - self._init_embeddings() - - def _init_embeddings(self): self.entity_embs_real = nn.Embedding(self.num_entities, self.embedding_dim, padding_idx=0) self.entity_embs_img = nn.Embedding(self.num_entities, self.embedding_dim, padding_idx=0) self.relation_embs_real = nn.Embedding(self.num_relations, self.embedding_dim, padding_idx=0) diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 67090a0a5c..bd0cd0a6ee 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -10,10 +10,15 @@ from ..base import MultimodalModel from ...losses import Loss +from ...nn.modules import DistMultInteraction from ...triples import TriplesNumericLiteralsFactory from ...typing import DeviceHint from ...utils import slice_triples +__all__ = [ + 'DistMultLiteral', +] + # TODO: Check entire build of the model class DistMultLiteral(MultimodalModel): @@ -38,6 +43,7 @@ def __init__( random_seed: Optional[int] = None, ) -> None: super().__init__( + interaction_function=DistMultInteraction(), triples_factory=triples_factory, embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index 69d901ec77..09b2821ed9 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -509,15 +509,8 @@ def __init__( ): if triples_factory.create_inverse_triples: raise ValueError('R-GCN handles edges in an undirected manner.') - super().__init__( - triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - predict_with_sigmoid=predict_with_sigmoid, - preferred_device=preferred_device, - random_seed=random_seed, - ) - self.entity_representations = RGCNRepresentations( + + entity_representations = RGCNRepresentations( triples_factory=triples_factory, embedding_dim=embedding_dim, num_bases_or_blocks=num_bases_or_blocks, @@ -534,12 +527,21 @@ def __init__( buffer_messages=buffer_messages, base_representations=None, ) - self.relation_embeddings = Embedding( + relation_representations = Embedding( num_embeddings=triples_factory.num_relations, embedding_dim=embedding_dim, ) - # TODO: Dummy - self.decoder = Decoder() + super().__init__( + triples_factory=triples_factory, + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + predict_with_sigmoid=predict_with_sigmoid, + preferred_device=preferred_device, + random_seed=random_seed, + interaction_function=Decoder(), + entity_representations=entity_representations, + relation_representations=relation_representations, + ) def forward( self, @@ -547,7 +549,7 @@ def forward( r_indices: Optional[torch.LongTensor], t_indices: Optional[torch.LongTensor], ) -> torch.FloatTensor: # noqa: D102 - h = self.entity_representations(indices=h_indices) - r = self.relation_embeddings(indices=r_indices) - t = self.entity_representations(indices=t_indices) + h = self.entity_representations[0](indices=h_indices) + r = self.relation_representations[0](indices=r_indices) + t = self.entity_representations[0](indices=t_indices) return self.decoder(h, r, t).unsqueeze(dim=-1) From 55de2686ea452bba374d29cc3da83ccac4376c8d Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Sun, 15 Nov 2020 20:15:32 +0100 Subject: [PATCH 288/690] Update base.py --- src/pykeen/models/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index cd8e4ad48b..8729f799e0 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1527,7 +1527,7 @@ def forward( t_source.get_in_canonical_shape(indices=t_indices), ) for h_source, r_source, t_source in ( - (self.entity_embeddings, self.relation_embeddings, self.second_entity_embeddings), - (self.second_entity_embeddings, self.second_relation_embeddings, self.entity_embeddings), + (self.entity_representations[0], self.relation_representations[0], self.entity_representations[1]), + (self.entity_representations[1], self.relation_representations[1], self.entity_representations[0]), ) ) From 9f3980646aad7966d0c4fb9d7024cdbb97159d1a Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 01:52:48 +0100 Subject: [PATCH 289/690] More cleanup --- src/pykeen/models/base.py | 34 +++++++++---------- .../models/multimodal/complex_literal.py | 2 +- .../models/multimodal/distmult_literal.py | 2 +- src/pykeen/models/unimodal/complex.py | 4 +-- src/pykeen/models/unimodal/conv_e.py | 4 +-- src/pykeen/models/unimodal/conv_kb.py | 2 +- src/pykeen/models/unimodal/distmult.py | 4 +-- src/pykeen/models/unimodal/ermlp.py | 2 +- src/pykeen/models/unimodal/ermlpe.py | 2 +- src/pykeen/models/unimodal/hole.py | 4 +-- src/pykeen/models/unimodal/kg2e.py | 2 +- src/pykeen/models/unimodal/ntn.py | 2 +- src/pykeen/models/unimodal/proj_e.py | 2 +- src/pykeen/models/unimodal/rescal.py | 2 +- src/pykeen/models/unimodal/rgcn.py | 9 ++--- src/pykeen/models/unimodal/rotate.py | 2 +- src/pykeen/models/unimodal/simple.py | 2 +- .../models/unimodal/structured_embedding.py | 2 +- src/pykeen/models/unimodal/trans_d.py | 2 +- src/pykeen/models/unimodal/trans_e.py | 2 +- src/pykeen/models/unimodal/trans_h.py | 2 +- src/pykeen/models/unimodal/trans_r.py | 2 +- src/pykeen/models/unimodal/tucker.py | 2 +- .../models/unimodal/unstructured_model.py | 2 +- tests/test_early_stopping.py | 9 ++--- tests/test_evaluators.py | 9 ++--- tests/test_interactions.py | 2 +- tests/test_model_mode.py | 26 ++++++-------- tests/test_models.py | 1 + 29 files changed, 64 insertions(+), 78 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 8729f799e0..d2ee15c181 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -249,7 +249,7 @@ class Model(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailR def __init__( self, triples_factory: TriplesFactory, - interaction_function: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], + interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, automatic_memory_optimization: Optional[bool] = None, @@ -334,7 +334,7 @@ def __init__( self.entity_representations = nn.ModuleList(entity_representations) self.relation_representations = nn.ModuleList(relation_representations) - self.interaction_function = interaction_function + self.interaction = interaction # reset parameters self.reset_parameters_() @@ -1035,7 +1035,7 @@ def forward( # normalization h, r, t = [x[0] if len(x) == 1 else x for x in (h, r, t)] - scores = self.interaction_function(h=h, r=r, t=t) + scores = self.interaction(h=h, r=r, t=t) if len(self.relation_representations) == 0: # same score for all relations repeats = [1, 1, 1, 1] @@ -1153,7 +1153,7 @@ class EntityEmbeddingModel(Model): def __init__( self, triples_factory: TriplesFactory, - interaction_function: Interaction[Representation, RelationRepresentation, Representation], + interaction: Interaction[Representation, RelationRepresentation, Representation], embedding_dim: int = 50, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, @@ -1178,6 +1178,7 @@ def __init__( self.embedding_dim = embedding_dim super().__init__( triples_factory=triples_factory, + interaction=interaction, automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, @@ -1194,7 +1195,6 @@ def __init__( constrainer=entity_constrainer, constrainer_kwargs=entity_constrainer_kwargs, ), - interaction_function=interaction_function, ) @@ -1206,7 +1206,7 @@ class EntityRelationEmbeddingModel(Model, ABC): def __init__( self, triples_factory: TriplesFactory, - interaction_function: Interaction[Representation, Representation, Representation], + interaction: Interaction[Representation, Representation, Representation], embedding_dim: int = 50, relation_dim: Optional[int] = None, loss: Optional[Loss] = None, @@ -1234,7 +1234,7 @@ def __init__( self.relation_dim = relation_dim super().__init__( triples_factory=triples_factory, - interaction_function=interaction_function, + interaction=interaction, automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, @@ -1268,7 +1268,7 @@ class SingleVectorEmbeddingModel(Model, ABC): def __init__( self, triples_factory: TriplesFactory, - interaction_function: Interaction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], + interaction: Interaction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], embedding_dim: int = 200, relation_dim: Union[None, int, Sequence[int]] = None, automatic_memory_optimization: Optional[bool] = None, @@ -1284,7 +1284,7 @@ def __init__( :param triples_factory: The triple factory connected to the model. - :param interaction_function: + :param interaction: The embedding-based interaction function used to compute scores. :param embedding_dim: The embedding dimensionality of the entity embeddings. @@ -1311,7 +1311,7 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - interaction_function=interaction_function, + interaction=interaction, entity_representations=Embedding.from_specification( num_embeddings=triples_factory.num_entities, shape=embedding_dim, @@ -1345,7 +1345,7 @@ class DoubleRelationEmbeddingModel(Model, ABC): def __init__( self, triples_factory: TriplesFactory, - interaction_function: Interaction[ + interaction: Interaction[ torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor], torch.FloatTensor, @@ -1366,6 +1366,7 @@ def __init__( relation_dim = embedding_dim super().__init__( triples_factory=triples_factory, + interaction=interaction, automatic_memory_optimization=automatic_memory_optimization, loss=loss, predict_with_sigmoid=predict_with_sigmoid, @@ -1391,7 +1392,6 @@ def __init__( specification=second_relation_embedding_specification, ), ], - interaction_function=interaction_function, ) @@ -1401,7 +1401,7 @@ class TwoVectorEmbeddingModel(Model, ABC): def __init__( self, triples_factory: TriplesFactory, - interaction_function: Interaction[ + interaction: Interaction[ Tuple[torch.FloatTensor, torch.FloatTensor], Tuple[torch.FloatTensor, torch.FloatTensor], Tuple[torch.FloatTensor, torch.FloatTensor], @@ -1453,7 +1453,7 @@ def __init__( specification=second_relation_embedding_specification, ), ], - interaction_function=interaction_function, + interaction=interaction, ) @@ -1463,7 +1463,7 @@ class TwoSideEmbeddingModel(Model): def __init__( self, triples_factory: TriplesFactory, - interaction_function: Interaction, + interaction: Interaction, embedding_dim: int = 50, relation_dim: Optional[int] = None, loss: Optional[Loss] = None, @@ -1511,7 +1511,7 @@ def __init__( specification=second_relation_embedding_specification, ), ], - interaction_function=interaction_function, + interaction=interaction, ) def forward( @@ -1521,7 +1521,7 @@ def forward( t_indices: Optional[torch.LongTensor], ) -> torch.FloatTensor: # noqa: D102 return 0.5 * sum( - self.interaction_function( + self.interaction( h_source.get_in_canonical_shape(indices=h_indices), r_source.get_in_canonical_shape(indices=r_indices), t_source.get_in_canonical_shape(indices=t_indices), diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index 58b18d6841..bb2588ad3f 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -51,7 +51,7 @@ def __init__( """Initialize the model.""" super().__init__( triples_factory=triples_factory, - interaction_function=ComplExInteraction(), + interaction=ComplExInteraction(), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index bd0cd0a6ee..8586af010d 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -43,7 +43,7 @@ def __init__( random_seed: Optional[int] = None, ) -> None: super().__init__( - interaction_function=DistMultInteraction(), + interaction=DistMultInteraction(), triples_factory=triples_factory, embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index ac57238775..98c6f70a18 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -94,11 +94,9 @@ def __init__( :param regularizer: BaseRegularizer The regularizer to use. """ - interaction_function = ComplExInteraction() - super().__init__( triples_factory=triples_factory, - interaction_function=interaction_function, + interaction=ComplExInteraction(), embedding_dim=2 * embedding_dim, # complex embeddings automatic_memory_optimization=automatic_memory_optimization, loss=loss, diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 3b4a24131a..ad44399946 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -166,7 +166,7 @@ def __init__( initializer=xavier_normal_, ), ), - interaction_function=ConvEInteraction( + interaction=ConvEInteraction( input_channels=input_channels, output_channels=output_channels, embedding_height=embedding_height, @@ -192,4 +192,4 @@ def forward( t = self.entity_representations[0].get_in_canonical_shape(indices=t_indices) t_bias = self.entity_representations[1].get_in_canonical_shape(indices=t_indices) self.regularize_if_necessary(h, r, t) - return self.interaction_function(h=h, r=r, t=(t, t_bias)) + return self.interaction(h=h, r=r, t=(t, t_bias)) diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 6e5256a787..0fb2441d50 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -86,7 +86,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - interaction_function=ConvKBInteraction( + interaction=ConvKBInteraction( hidden_dropout_rate=hidden_dropout_rate, embedding_dim=embedding_dim, num_filters=num_filters, diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index e305ce2987..c5a7389845 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -81,11 +81,9 @@ def __init__( :param embedding_dim: The entity embedding dimension $d$. Is usually $d \in [50, 300]$. """ - interaction_function = DistMultInteraction() - super().__init__( triples_factory=triples_factory, - interaction_function=interaction_function, + interaction=DistMultInteraction(), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index b326e35e05..b85f34aa51 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -55,7 +55,7 @@ def __init__( super().__init__( triples_factory=triples_factory, - interaction_function=ERMLPInteraction( + interaction=ERMLPInteraction( embedding_dim=embedding_dim, hidden_dim=hidden_dim, ), diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index 0f4c163901..1386ff4eb2 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -66,7 +66,7 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, - interaction_function=ERMLPEInteraction( + interaction=ERMLPEInteraction( hidden_dim=hidden_dim, input_dropout=input_dropout, hidden_dropout=hidden_dropout, diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 611b942a97..509f240dbe 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -63,13 +63,11 @@ def __init__( regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model.""" - interaction_function = HolEInteraction() - super().__init__( triples_factory=triples_factory, embedding_dim=embedding_dim, loss=loss, - interaction_function=interaction_function, + interaction=HolEInteraction(), automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 3e93c150ba..4e73f17f92 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -80,7 +80,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - interaction_function=KG2EInteraction( + interaction=KG2EInteraction( similarity=dist_similarity, ), embedding_dim=embedding_dim, diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 4f78ad982a..a096108ccc 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -98,7 +98,7 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - interaction_function=NTNInteraction( + interaction=NTNInteraction( non_linearity=non_linearity, ), entity_representations=Embedding( diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index dc4549d14f..2445fcd6a3 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -67,7 +67,7 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, - interaction_function=ProjEInteraction( + interaction=ProjEInteraction( embedding_dim=embedding_dim, inner_non_linearity=inner_non_linearity, ), diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index dc755e8ce2..9b9fb28627 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -70,7 +70,7 @@ def __init__( # TODO: regularization super().__init__( triples_factory=triples_factory, - interaction_function=RESCALInteraction(), + interaction=RESCALInteraction(), embedding_dim=embedding_dim, relation_dim=(embedding_dim, embedding_dim), automatic_memory_optimization=automatic_memory_optimization, diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index 09b2821ed9..df6619630f 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -15,6 +15,7 @@ from ..base import Model from ...losses import Loss from ...nn import Embedding, RepresentationModule +from ...nn.modules import DistMultInteraction from ...triples import TriplesFactory from ...typing import DeviceHint @@ -418,12 +419,6 @@ def reset_parameters(self): act.reset_parameters() -class Decoder(nn.Module): - # TODO: Replace this by interaction function, once https://github.com/pykeen/pykeen/pull/107 is merged. - def forward(self, h, r, t): - return (h * r * t).sum(dim=-1) - - class RGCN(Model): """An implementation of R-GCN from [schlichtkrull2018]_. @@ -538,7 +533,7 @@ def __init__( predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, random_seed=random_seed, - interaction_function=Decoder(), + interaction=DistMultInteraction(), entity_representations=entity_representations, relation_representations=relation_representations, ) diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index a5a35f172a..7404b40c8d 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -96,7 +96,7 @@ def __init__( # TODO: regularization super().__init__( triples_factory=triples_factory, - interaction_function=RotatEInteraction(), + interaction=RotatEInteraction(), embedding_dim=2 * embedding_dim, loss=loss, automatic_memory_optimization=automatic_memory_optimization, diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index e0f0978cac..bbb7bcc8b2 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -75,7 +75,7 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, - interaction_function=DistMultInteraction(), + interaction=DistMultInteraction(), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index 7b6a3cf02b..0334f68dc0 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -71,7 +71,7 @@ def __init__( ) super().__init__( triples_factory=triples_factory, - interaction_function=StructuredEmbeddingInteraction( + interaction=StructuredEmbeddingInteraction( p=scoring_fct_norm, power_norm=False, ), diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 69a27e42fb..a1c90834e5 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -70,7 +70,7 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, - interaction_function=TransDInteraction(p=2, power_norm=True), + interaction=TransDInteraction(p=2, power_norm=True), embedding_dim=embedding_dim, relation_dim=relation_dim, automatic_memory_optimization=automatic_memory_optimization, diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 197fc2676b..3d5333ac5d 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -69,7 +69,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - interaction_function=TransEInteraction(p=scoring_fct_norm, power_norm=False), + interaction=TransEInteraction(p=scoring_fct_norm, power_norm=False), embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 2bdea2ebfb..afa289e771 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -81,7 +81,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - interaction_function=TransHInteraction( + interaction=TransHInteraction( p=scoring_fct_norm, power_norm=False, ), diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 1a1661461f..8a2fa628b5 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -111,7 +111,7 @@ def __init__( ), ), ], - interaction_function=TransRInteraction( + interaction=TransRInteraction( p=scoring_fct_norm, ), ) diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index b7390a9271..0bd9fa158f 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -89,7 +89,7 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - interaction_function=TuckerInteraction( + interaction=TuckerInteraction( embedding_dim=embedding_dim, relation_dim=relation_dim, dropout_0=dropout_0, diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index 734b395369..d91a410811 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -68,7 +68,7 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, regularizer=regularizer, - interaction_function=UnstructuredModelInteraction(p=scoring_fct_norm), + interaction=UnstructuredModelInteraction(p=scoring_fct_norm), entity_representations=Embedding( num_embeddings=triples_factory.num_entities, embedding_dim=embedding_dim, diff --git a/tests/test_early_stopping.py b/tests/test_early_stopping.py index dc4946371d..587e250530 100644 --- a/tests/test_early_stopping.py +++ b/tests/test_early_stopping.py @@ -106,7 +106,11 @@ class MockModel(EntityRelationEmbeddingModel): """A mock model returning fake scores.""" def __init__(self, triples_factory: TriplesFactory, automatic_memory_optimization: bool): - super().__init__(triples_factory=triples_factory, automatic_memory_optimization=automatic_memory_optimization) + super().__init__( + triples_factory=triples_factory, + automatic_memory_optimization=automatic_memory_optimization, + interaction=None, + ) num_entities = self.num_entities self.scores = torch.arange(num_entities, dtype=torch.float) @@ -126,9 +130,6 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self._generate_fake_scores(batch=rt_batch) - def reset_parameters_(self) -> Model: # noqa: D102 - pass # Not needed for unittest - class LogCallWrapper: """An object which wraps functions and checks whether they have been called.""" diff --git a/tests/test_evaluators.py b/tests/test_evaluators.py index 64b277118d..54107db0c2 100644 --- a/tests/test_evaluators.py +++ b/tests/test_evaluators.py @@ -432,7 +432,11 @@ class DummyModel(EntityRelationEmbeddingModel): """A dummy model returning fake scores.""" def __init__(self, triples_factory: TriplesFactory, automatic_memory_optimization: bool): - super().__init__(triples_factory=triples_factory, automatic_memory_optimization=automatic_memory_optimization) + super().__init__( + triples_factory=triples_factory, + automatic_memory_optimization=automatic_memory_optimization, + interaction=..., + ) num_entities = self.num_entities self.scores = torch.arange(num_entities, dtype=torch.float) @@ -452,9 +456,6 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D10 def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 return self._generate_fake_scores(batch=rt_batch) - def reset_parameters_(self) -> Model: # noqa: D102 - pass # Not needed for unittest - class TestEvaluationStructure(unittest.TestCase): """Tests for testing the correct structure of the evaluation procedure.""" diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 4827017d32..35e5735b88 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -50,7 +50,7 @@ def test_testing(self): to_test = set(get_subclasses(self.base_cls)).difference(self.skip_cls) tested = (test_cls.cls for test_cls in get_subclasses(self.base_test) if hasattr(test_cls, "cls")) not_tested = to_test.difference(tested) - assert not_tested == set() + assert not not_tested, not_tested class InteractionTests(GenericTests[pykeen.nn.modules.Interaction]): diff --git a/tests/test_model_mode.py b/tests/test_model_mode.py index 5e8e9e8702..34ff0df60b 100644 --- a/tests/test_model_mode.py +++ b/tests/test_model_mode.py @@ -6,11 +6,11 @@ from dataclasses import dataclass import torch -from torch import nn from pykeen.datasets import Nations from pykeen.models import TransE -from pykeen.models.base import EntityRelationEmbeddingModel, Model +from pykeen.models.base import EntityRelationEmbeddingModel +from pykeen.nn.modules import Interaction from pykeen.triples import TriplesFactory from pykeen.utils import resolve_device @@ -158,24 +158,18 @@ def test_alignment_of_score_r_fall_back(self) -> None: assert all(scores_r_function == scores_hrt_function) -class SimpleInteractionModel(EntityRelationEmbeddingModel): - """A model with a simple interaction function for testing the base model.""" +class SimpleInteraction(Interaction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): + """A simple interaction as the sum of the vectors.""" - def __init__(self, triples_factory: TriplesFactory): - super().__init__(triples_factory=triples_factory) - self.entity_embeddings = nn.Embedding(self.num_entities, self.embedding_dim) - self.relation_embeddings = nn.Embedding(self.num_relations, self.embedding_dim) + def forward(self, h, r, t) -> torch.FloatTensor: # noqa:D102 + return torch.sum(h + r + t, dim=1) - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - # Get embeddings - h = self.entity_embeddings(hrt_batch[:, 0]) - r = self.relation_embeddings(hrt_batch[:, 1]) - t = self.entity_embeddings(hrt_batch[:, 2]) - return torch.sum(h + r + t, dim=1) +class SimpleInteractionModel(EntityRelationEmbeddingModel): + """A model with a simple interaction function for testing the base model.""" - def reset_parameters_(self) -> Model: # noqa: D102 - pass # Not needed for unittest + def __init__(self, *args, **kwargs): + super().__init__(*args, interaction=SimpleInteraction(), **kwargs) @dataclass diff --git a/tests/test_models.py b/tests/test_models.py index 5905b3f6a7..4f29550c30 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -50,6 +50,7 @@ Model.__name__, 'DummyModel', MultimodalModel.__name__, + DoubleRelationEmbeddingModel.__name__, EntityEmbeddingModel.__name__, EntityRelationEmbeddingModel.__name__, SingleVectorEmbeddingModel.__name__, From d2a685d2c02119a83dd27caa78f6a1df6d06c316 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 03:08:21 +0100 Subject: [PATCH 290/690] Add additional abstraction layer into the Model class See the docstring for explanation --- src/pykeen/models/__init__.py | 11 +- src/pykeen/models/base.py | 255 ++++++++++++------ src/pykeen/models/unimodal/conv_e.py | 4 +- src/pykeen/models/unimodal/ntn.py | 4 +- src/pykeen/models/unimodal/rgcn.py | 4 +- src/pykeen/models/unimodal/trans_r.py | 4 +- .../models/unimodal/unstructured_model.py | 4 +- src/pykeen/nn/__init__.py | 3 + src/pykeen/nn/modules.py | 41 +-- tests/test_models.py | 66 +++-- 10 files changed, 252 insertions(+), 144 deletions(-) diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index d2ab204f80..25d58a71fb 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -9,7 +9,7 @@ from typing import Mapping, Set, Type, Union from .base import ( # noqa:F401 - DoubleRelationEmbeddingModel, EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, MultimodalModel, + DoubleRelationEmbeddingModel, ERModel, EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, MultimodalModel, SingleVectorEmbeddingModel, TwoSideEmbeddingModel, TwoVectorEmbeddingModel, ) from .multimodal import ComplExLiteral, DistMultLiteral @@ -66,7 +66,8 @@ 'get_model_cls', ] -_CONCRETE_BASES = { +_BASE_MODELS = { + ERModel, SingleVectorEmbeddingModel, DoubleRelationEmbeddingModel, TwoSideEmbeddingModel, @@ -79,12 +80,12 @@ def _concrete_subclasses(cls: Type[Model]): for subcls in cls.__subclasses__(): - if not subcls._is_abstract() and subcls not in _CONCRETE_BASES: + if not subcls._is_abstract and subcls not in _BASE_MODELS: yield subcls yield from _concrete_subclasses(subcls) -_MODELS: Set[Type[Model]] = set(_concrete_subclasses(Model)) +_MODELS: Set[Type[Model]] = set(_concrete_subclasses(Model)) # type: ignore #: A mapping of models' names to their implementations models: Mapping[str, Type[Model]] = { @@ -101,6 +102,6 @@ def get_model_cls(query: Union[str, Type[Model]]) -> Type[Model]: """ return get_cls( query, - base=Model, + base=Model, # type: ignore lookup_dict=models, ) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index d2ee15c181..85cd2227f3 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -6,7 +6,7 @@ import inspect import itertools as itt import logging -from abc import ABC +from abc import ABC, abstractmethod from collections import defaultdict from operator import itemgetter from typing import ( @@ -33,9 +33,13 @@ __all__ = [ 'Model', + 'ERModel', 'EntityEmbeddingModel', 'EntityRelationEmbeddingModel', 'SingleVectorEmbeddingModel', + 'DoubleRelationEmbeddingModel', + 'TwoVectorEmbeddingModel', + 'TwoSideEmbeddingModel', 'MultimodalModel', ] @@ -193,6 +197,10 @@ def _process_remove_known(df: pd.DataFrame, remove_known: bool, testing: Optiona return df +def _can_slice(fn) -> bool: + return 'slice_size' in inspect.getfullargspec(fn).args + + def _track_hyperparameters(cls: Type['Model']) -> None: """Initialize the subclass while keeping track of hyper-parameters.""" # Keep track of the hyper-parameters that are used across all @@ -216,12 +224,29 @@ def _new_init(self, *args, **kwargs): cls.__init__ = _new_init # type: ignore -class Model(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation], ABC): - """A base module for all of the KGE models.""" +class Model(nn.Module, ABC): + """An abstract class for knowledge graph embedding models (KGEMs). + + The only function that needs to be implemented for a given subclass is + :meth:`Model.forward`. The job of the :meth:`Model.forward` function, as + opposed to the completely general :meth:`torch.nn.Module.forward` is + to take indices for the head, relation, and tails' respective representation(s) + and to determine a score. + + Subclasses of Model can decide however they want on how to store entities' and + relations' representations, how they want to be looked up, and how they should + be scored. The :class:`GeneralModel` provides a commonly useful implementation + which allows for the specification of one or more entity representations and + one or more relation representations in the form of :class:`pykeen.nn.Embedding` + as well as a matching instance of a :class:`pykeen.nn.Interaction`. + """ #: A dictionary of hyper-parameters to the models that use them _hyperparameter_usage: ClassVar[Dict[str, Set[str]]] = defaultdict(set) + #: Keep track of if this is a base model + _is_abstract: ClassVar[bool] + #: The default strategy for optimizing the model's hyper-parameters hpo_default: ClassVar[Mapping[str, Any]] @@ -240,24 +265,15 @@ class Model(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailR #: The instance of the regularizer regularizer: Regularizer - #: The entity representations - entity_representations: Sequence[RepresentationModule] - - #: The relation representations - relation_representations: Sequence[RepresentationModule] - def __init__( self, triples_factory: TriplesFactory, - interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, regularizer: Optional[Regularizer] = None, - entity_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, - relation_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, ) -> None: """Initialize the module. @@ -325,27 +341,11 @@ def __init__( # This allows to store the optimized parameters self.automatic_memory_optimization = automatic_memory_optimization - # normalization - if entity_representations is not None and not isinstance(entity_representations, Sequence): - entity_representations = [entity_representations] - if relation_representations is not None and not isinstance(relation_representations, Sequence): - relation_representations = [relation_representations] - # Important: use ModuleList to ensure that Pytorch correctly handles their devices and parameters - self.entity_representations = nn.ModuleList(entity_representations) - self.relation_representations = nn.ModuleList(relation_representations) - - self.interaction = interaction - - # reset parameters - self.reset_parameters_() - - @classmethod - def _is_abstract(cls) -> bool: - return inspect.isabstract(cls) - - def __init_subclass__(cls, **kwargs): # noqa:D105 - if not cls._is_abstract(): + def __init_subclass__(cls, autoreset: bool = True, **kwargs): # noqa:D105 + cls._is_abstract = not autoreset + if not cls._is_abstract: _track_hyperparameters(cls) + _add_post_reset_parameters(cls) @property def can_slice_h(self) -> bool: @@ -998,7 +998,7 @@ def _compute_loss( ) return self.loss(tensor_1, tensor_2) + self.regularizer.term - # @abstractmethod + @abstractmethod def forward( self, h_indices: Optional[torch.LongTensor], @@ -1021,33 +1021,7 @@ def forward( :return: shape: (batch_size, num_heads, num_relations, num_tails) The score for each triple. """ - h, r, t = [ - [ - representation.get_in_canonical_shape(indices=indices) - for representation in representations - ] - for indices, representations in ( - (h_indices, self.entity_representations), - (r_indices, self.relation_representations), - (t_indices, self.entity_representations), - ) - ] - # normalization - h, r, t = [x[0] if len(x) == 1 else x for x in (h, r, t)] - - scores = self.interaction(h=h, r=r, t=t) - if len(self.relation_representations) == 0: - # same score for all relations - repeats = [1, 1, 1, 1] - if r_indices is None: - repeats[2] = self.num_relations - else: - relation_batch_size = len(r_indices) - if scores.shape[0] < relation_batch_size: - repeats[0] = relation_batch_size - scores = scores.repeat(*repeats) - - return scores + raise NotImplementedError def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: """Forward pass. @@ -1056,8 +1030,7 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: :param hrt_batch: shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples. - :raises NotImplementedError: - If the method was not implemented for this class. + :return: shape: (batch_size, 1), dtype: float The score for each triple. """ @@ -1145,7 +1118,122 @@ def load_state(self, path: str) -> None: self.load_state_dict(torch.load(path, map_location=self.device)) -class EntityEmbeddingModel(Model): +class ERModel(Model, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation], autoreset=False): + """A commonly useful base for KGEMs using embeddings and interaction modules.""" + + #: The entity representations + entity_representations: Sequence[RepresentationModule] + + #: The relation representations + relation_representations: Sequence[RepresentationModule] + + def __init__( + self, + triples_factory: TriplesFactory, + interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], + loss: Optional[Loss] = None, + predict_with_sigmoid: bool = False, + automatic_memory_optimization: Optional[bool] = None, + preferred_device: DeviceHint = None, + random_seed: Optional[int] = None, + regularizer: Optional[Regularizer] = None, + entity_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, + relation_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, + ) -> None: + """Initialize the module. + + :param triples_factory: + The triples factory facilitates access to the dataset. + :param loss: + The loss to use. If None is given, use the loss default specific to the model subclass. + :param predict_with_sigmoid: + Whether to apply sigmoid onto the scores when predicting scores. Applying sigmoid at prediction time may + lead to exactly equal scores for certain triples with very high, or very low score. When not trained with + applying sigmoid (or using BCEWithLogitsLoss), the scores are not calibrated to perform well with sigmoid. + :param automatic_memory_optimization: + If set to `True`, the model derives the maximum possible batch sizes for the scoring of triples during + evaluation and also training (if no batch size was given). This allows to fully utilize the hardware at hand + and achieves the fastest calculations possible. + :param preferred_device: + The preferred device for model training and inference. + :param random_seed: + A random seed to use for initialising the model's weights. **Should** be set when aiming at reproducibility. + :param regularizer: + A regularizer to use for training. + """ + super().__init__( + triples_factory=triples_factory, + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + random_seed=random_seed, + regularizer=regularizer, + predict_with_sigmoid=predict_with_sigmoid, + ) + + # normalization + if entity_representations is not None and not isinstance(entity_representations, Sequence): + entity_representations = [entity_representations] + if relation_representations is not None and not isinstance(relation_representations, Sequence): + relation_representations = [relation_representations] + # Important: use ModuleList to ensure that Pytorch correctly handles their devices and parameters + self.entity_representations = nn.ModuleList(entity_representations) + self.relation_representations = nn.ModuleList(relation_representations) + + self.interaction = interaction + + def forward( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: + """Forward pass. + + This method takes head, relation and tail indices and calculates the corresponding score. + + All indices which are not None, have to be either 1-element or have the same shape, which is the batch size. + + :param h_indices: + The head indices. None indicates to use all. + :param r_indices: + The relation indices. None indicates to use all. + :param t_indices: + The tail indices. None indicates to use all. + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The score for each triple. + """ + h, r, t = [ + [ + representation.get_in_canonical_shape(indices=indices) + for representation in representations + ] + for indices, representations in ( + (h_indices, self.entity_representations), + (r_indices, self.relation_representations), + (t_indices, self.entity_representations), + ) + ] + # normalization + h, r, t = [x[0] if len(x) == 1 else x for x in (h, r, t)] + + scores = self.interaction(h=h, r=r, t=t) + if len(self.relation_representations) == 0: + # same score for all relations + repeats = [1, 1, 1, 1] + if r_indices is None: + repeats[2] = self.num_relations + else: + relation_batch_size = len(r_indices) + if scores.shape[0] < relation_batch_size: + repeats[0] = relation_batch_size + scores = scores.repeat(*repeats) + + return scores + + +class EntityEmbeddingModel(ERModel, autoreset=False): """A base module for most KGE models that have one embedding for entities.""" # TODO: deprecated @@ -1198,7 +1286,7 @@ def __init__( ) -class EntityRelationEmbeddingModel(Model, ABC): +class EntityRelationEmbeddingModel(ERModel, autoreset=False): """A base module for KGE models that have different embeddings for entities and relations.""" # TODO: Deprecated. @@ -1254,16 +1342,12 @@ def __init__( ) -def _can_slice(fn) -> bool: - return 'slice_size' in inspect.getfullargspec(fn).args - - -class MultimodalModel(EntityRelationEmbeddingModel): +class MultimodalModel(EntityRelationEmbeddingModel, autoreset=False): """A multimodal KGE model.""" -class SingleVectorEmbeddingModel(Model, ABC): - """A base class for embedding models which store a single vector for each entity and relation.""" +class SingleVectorEmbeddingModel(ERModel, autoreset=False): + """A KGEM that stores one :class:`pykeen.nn.Embedding` for each entities and relations.""" def __init__( self, @@ -1339,8 +1423,14 @@ def relation_dim(self) -> int: # noqa:D401 return embedding.embedding_dim -class DoubleRelationEmbeddingModel(Model, ABC): - """A model with one vector for each entity and two vectors for each relation.""" +class DoubleRelationEmbeddingModel(ERModel, autoreset=False): + """A KGEM that stores one :class:`pykeen.nn.Embedding` for entities and two for relations. + + .. seealso:: + + - :class:`pykeen.models.StructuredEmbedding` + - :class:`pykeen.models.TransH` + """ def __init__( self, @@ -1395,8 +1485,14 @@ def __init__( ) -class TwoVectorEmbeddingModel(Model, ABC): - """A model with two vectors for each entity and relation.""" +class TwoVectorEmbeddingModel(ERModel, autoreset=False): + """A KGEM that stores two :class:`pykeen.nn.Embedding` for each entities and relations. + + .. seealso:: + + - :class:`pykeen.models.KG2E` + - :class:`pykeen.models.TransD` + """ def __init__( self, @@ -1457,8 +1553,13 @@ def __init__( ) -class TwoSideEmbeddingModel(Model): - """A model which averages scores for forward and backward model.""" +class TwoSideEmbeddingModel(ERModel, autoreset=False): + """A KGEM with two sub-KGEMs that serve as a "forwards" and "backwards" model. + + Stores two :class:`pykeen.nn.Embedding` for each entities and relations. + + .. seealso:: :class:`pykeen.models.SimplE` + """ def __init__( self, diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index ad44399946..d517345fde 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -8,7 +8,7 @@ import torch from torch import nn -from .. import Model +from .. import ERModel from ...losses import BCEAfterSigmoidLoss, Loss from ...nn import Embedding from ...nn.emb import EmbeddingSpecification @@ -25,7 +25,7 @@ logger = logging.getLogger(__name__) -class ConvE(Model): +class ConvE(ERModel): r"""An implementation of ConvE from [dettmers2018]_. ConvE is a CNN-based approach. For each triple $(h,r,t)$, the input to ConvE is a matrix diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index a096108ccc..91d6f75e1e 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -6,7 +6,7 @@ from torch import nn -from .. import Model +from .. import ERModel from ...losses import Loss from ...nn import Embedding from ...nn.modules import NTNInteraction @@ -19,7 +19,7 @@ ] -class NTN(Model): +class NTN(ERModel): r"""An implementation of NTN from [socher2013]_. NTN uses a bilinear tensor layer instead of a standard linear neural network layer: diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index df6619630f..ed6c4f4043 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -12,7 +12,7 @@ from . import ComplEx, DistMult, ERMLP from .. import EntityEmbeddingModel -from ..base import Model +from ..base import ERModel from ...losses import Loss from ...nn import Embedding, RepresentationModule from ...nn.modules import DistMultInteraction @@ -419,7 +419,7 @@ def reset_parameters(self): act.reset_parameters() -class RGCN(Model): +class RGCN(ERModel): """An implementation of R-GCN from [schlichtkrull2018]_. This model uses graph convolutions with relation-specific weights. diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 8a2fa628b5..af75005e15 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -6,7 +6,7 @@ from torch.nn import functional -from .. import Model +from .. import ERModel from ...losses import Loss from ...nn import Embedding from ...nn.emb import EmbeddingSpecification @@ -22,7 +22,7 @@ ] -class TransR(Model): +class TransR(ERModel): r"""An implementation of TransR from [lin2015]_. TransR is an extension of :class:`pykeen.models.TransH` that explicitly considers entities and relations as diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index d91a410811..4262ab607d 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -4,7 +4,7 @@ from typing import Optional -from .. import Model +from .. import ERModel from ...losses import Loss from ...nn import Embedding from ...nn.init import xavier_normal_ @@ -18,7 +18,7 @@ ] -class UnstructuredModel(Model): +class UnstructuredModel(ERModel): r"""An implementation of the Unstructured Model (UM) published by [bordes2014]_. UM computes the distance between head and tail entities then applies the $l_p$ norm. diff --git a/src/pykeen/nn/__init__.py b/src/pykeen/nn/__init__.py index 2b24107f24..7de87500b3 100644 --- a/src/pykeen/nn/__init__.py +++ b/src/pykeen/nn/__init__.py @@ -4,10 +4,13 @@ from . import functional, init from .emb import Embedding, RepresentationModule +from .modules import Interaction, StatelessInteraction __all__ = [ 'Embedding', 'RepresentationModule', + 'Interaction', + 'StatelessInteraction', 'init', 'functional', ] diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index b707321629..66f1e2fbc7 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -18,24 +18,29 @@ logger = logging.getLogger(__name__) __all__ = [ - "ComplExInteraction", - "ConvEInteraction", - "ConvKBInteraction", - "DistMultInteraction", - "ERMLPInteraction", - "ERMLPEInteraction", - "HolEInteraction", - "Interaction", - "KG2EInteraction", - "NTNInteraction", - "ProjEInteraction", - "RESCALInteraction", - "RotatEInteraction", - "StructuredEmbeddingInteraction", - "TransDInteraction", - "TransEInteraction", - "TransHInteraction", - "TransRInteraction", + # Base Classes + 'Interaction', + 'StatelessInteraction', + 'TranslationalInteraction', + # Concrete Classes + 'ComplExInteraction', + 'ConvEInteraction', + 'ConvKBInteraction', + 'DistMultInteraction', + 'ERMLPInteraction', + 'ERMLPEInteraction', + 'HolEInteraction', + 'KG2EInteraction', + 'NTNInteraction', + 'ProjEInteraction', + 'RESCALInteraction', + 'RotatEInteraction', + 'StructuredEmbeddingInteraction', + 'TransDInteraction', + 'TransEInteraction', + 'TransHInteraction', + 'TransRInteraction', + 'UnstructuredModelInteraction', ] diff --git a/tests/test_models.py b/tests/test_models.py index 4f29550c30..d6385e12af 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -22,9 +22,9 @@ import pykeen.models from pykeen.datasets.kinships import KINSHIPS_TRAIN_PATH from pykeen.datasets.nations import NATIONS_TEST_PATH, NATIONS_TRAIN_PATH, Nations -from pykeen.models import _MODELS +from pykeen.models import _BASE_MODELS, _MODELS from pykeen.models.base import ( - DoubleRelationEmbeddingModel, EntityEmbeddingModel, + DoubleRelationEmbeddingModel, ERModel, EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, MultimodalModel, @@ -491,32 +491,33 @@ def test_reset_parameters_constructor_call(self): def test_custom_representations(self): """Tests whether we can provide custom representations.""" - if isinstance(self.model, EntityEmbeddingModel): - old_embeddings = self.model.entity_embeddings - self.model.entity_embeddings = _CustomRepresentations( + if not isinstance(self.model, ERModel): + self.skipTest(f'Not testing custom representations for model: {self.model.__class__.__name__}') + + old_entity_reps = self.model.entity_representations + self.model.entity_representations = nn.ModuleList([ + _CustomRepresentations( num_entities=self.factory.num_entities, - embedding_dim=old_embeddings.embedding_dim, + embedding_dim=er.embedding_dim, ) - # call some functions - self.model.reset_parameters_() - self.test_score_hrt() - self.test_score_t() - # reset to old state - self.model.entity_embeddings = old_embeddings - elif isinstance(self.model, EntityRelationEmbeddingModel): - old_embeddings = self.model.relation_embeddings - self.model.relation_embeddings = _CustomRepresentations( - num_entities=self.factory.num_relations, - embedding_dim=old_embeddings.embedding_dim, + for er in old_entity_reps + ]) + old_relation_reps = self.model.relation_representations + self.model.relation_representations = nn.ModuleList([ + _CustomRepresentations( + num_entities=self.factory.num_entities, + embedding_dim=er.embedding_dim, ) - # call some functions - self.model.reset_parameters_() - self.test_score_hrt() - self.test_score_t() - # reset to old state - self.model.relation_embeddings = old_embeddings - else: - self.skipTest(f'Not testing custom representations for model: {self.model.__class__.__name__}') + for er in old_relation_reps + ]) + + # call some functions + self.model.reset_parameters_() + self.test_score_hrt() + self.test_score_t() + # reset to old state + self.model.entity_representations = old_entity_reps + self.model.relation_representations = old_relation_reps class _DistanceModelTestCase(_ModelTestCase): @@ -706,9 +707,9 @@ class _BaseNTNTest(_ModelTestCase, unittest.TestCase): def test_can_slice(self): """Test that the slicing properties are calculated correctly.""" - self.assertTrue(self.model.can_slice_h) - self.assertTrue(self.model.can_slice_r) - self.assertTrue(self.model.can_slice_t) + self.assertTrue(self.model.can_slice_h, msg='Unable to slice on heads') + self.assertTrue(self.model.can_slice_r, msg='Unable to slice on relations') + self.assertTrue(self.model.can_slice_t, msg='Unable to slice on tails') class TestNTNLowMemory(_BaseNTNTest): @@ -1109,10 +1110,7 @@ class TestRandom(unittest.TestCase): def test_abstract(self): """Test that classes are checked as abstract properly.""" - self.assertTrue(Model._is_abstract()) - self.assertTrue(EntityEmbeddingModel._is_abstract()) - self.assertTrue(EntityRelationEmbeddingModel._is_abstract()) + for model_cls in _BASE_MODELS: + self.assertTrue(model_cls._is_abstract) for model_cls in _MODELS: - if issubclass(model_cls, MultimodalModel): - continue - self.assertFalse(model_cls._is_abstract(), msg=f'{model_cls.__name__} should not be abstract') + self.assertFalse(model_cls._is_abstract, msg=f'{model_cls.__name__} should not be abstract') From d561ee49e63c311a96e3a61faf0ce7eb08688d9d Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 03:36:49 +0100 Subject: [PATCH 291/690] Update .travis.yml --- .travis.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.travis.yml b/.travis.yml index 7a0b6dcc20..8ca6084948 100644 --- a/.travis.yml +++ b/.travis.yml @@ -2,7 +2,7 @@ os: linux cache: pip dist: xenial language: python -python: "3.7" +python: "3.8" jobs: include: # lint stage From 105eea062c123e4fa0fbc4496ef14fb938d9e067 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 03:39:06 +0100 Subject: [PATCH 292/690] Update docs --- docs/source/index.rst | 1 + docs/source/reference/interactions.rst | 13 ++++++++++ docs/source/reference/models.rst | 13 +--------- src/pykeen/nn/modules.py | 35 +++++++++++++------------- 4 files changed, 33 insertions(+), 29 deletions(-) create mode 100644 docs/source/reference/interactions.rst diff --git a/docs/source/index.rst b/docs/source/index.rst index 7cef1efd38..ee433108b9 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -24,6 +24,7 @@ PyKEEN :maxdepth: 2 reference/pipeline + reference/interactions reference/models reference/datasets reference/triples diff --git a/docs/source/reference/interactions.rst b/docs/source/reference/interactions.rst new file mode 100644 index 0000000000..d051f7e473 --- /dev/null +++ b/docs/source/reference/interactions.rst @@ -0,0 +1,13 @@ +Interactions +============ +Functional Interface +~~~~~~~~~~~~~~~~~~~~ +.. automodapi:: pykeen.nn.functional + :no-heading: + :headings: -- + +Module Interface +~~~~~~~~~~~~~~~~ +.. automodapi:: pykeen.nn.modules + :no-heading: + :headings: -- diff --git a/docs/source/reference/models.rst b/docs/source/reference/models.rst index 8e3baac80c..73799cb6b4 100644 --- a/docs/source/reference/models.rst +++ b/docs/source/reference/models.rst @@ -20,15 +20,4 @@ Extra Modules ------------- .. automodule:: pykeen.nn :members: - -Interactors ------------ -Functional Interface -~~~~~~~~~~~~~~~~~~~~ -.. automodule:: pykeen.nn.functional - :members: - -Module Interface -~~~~~~~~~~~~~~~~ -.. automodule:: pykeen.nn.modules - :members: + :exclude-members: Interaction, StatelessInteraction diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 66f1e2fbc7..c28fd4ce99 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -5,7 +5,7 @@ import itertools import logging import math -from abc import ABC +from abc import ABC, abstractmethod from typing import Callable, Generic, List, Optional, Sequence, Tuple, Union import torch @@ -52,7 +52,7 @@ def _unpack_singletons(*xs: Tuple) -> Sequence[Tuple]: return [x[0] if len(x) == 1 else x for x in xs] -class Interaction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation]): +class Interaction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation], ABC): """Base class for interaction functions.""" # Dimensions @@ -70,6 +70,7 @@ class Interaction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, #: The symbolic shapes for relation representations relation_shape: Sequence[str] = ("d",) + @abstractmethod def forward( self, h: HeadRepresentation, @@ -336,6 +337,10 @@ class StatelessInteraction(Interaction[HeadRepresentation, RelationRepresentatio """Interaction function without state.""" def __init__(self, f: Callable[..., torch.FloatTensor]): + """Instantiate the stateless interaction module. + + :param f: The interaction function, like ones from :mod:`pykeen.nn.functional`. + """ super().__init__() self.f = f @@ -346,14 +351,11 @@ def forward( t: TailRepresentation, ) -> torch.FloatTensor: # noqa: D102 # normalization - h, r, t = _ensure_tuple(h, r, t) + h, r, t = _ensure_tuple(h, r, t) # TODO provide example of non-simple case return self.f(*h, *r, *t) -class TranslationalInteraction( - Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], - ABC, -): +class TranslationalInteraction(Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], ABC): """The translational interaction function shared by the TransE, TransR, TransH, and other Trans models.""" def __init__(self, p: int, power_norm: bool = False): @@ -438,9 +440,7 @@ def _calculate_missing_shape_information( return input_channels, width, height -class ConvEInteraction( - Interaction[torch.FloatTensor, torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor]], -): +class ConvEInteraction(Interaction[torch.FloatTensor, torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor]]): """ConvE interaction function.""" tail_entity_shape = ("d", "k") # with k=1 @@ -546,7 +546,10 @@ def forward( class ConvKBInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): - """Interaction function of ConvKB.""" + """Interaction function of ConvKB. + + .. seealso:: :func:`pykeen.nn.functional.convkb_interaction`` + """ def __init__( self, @@ -593,15 +596,14 @@ def forward( class DistMultInteraction(StatelessInteraction[FloatTensor, FloatTensor, FloatTensor]): - """Interaction function of DistMult.""" + """A module wrapping the DistMult interaction function at :func:`pykeen.nn.functional.distmult_interaction`.""" def __init__(self): super().__init__(f=pkf.distmult_interaction) class ERMLPInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): - """ - Interaction function of ER-MLP. + """A module wrapping the ER-MLP interaction function from :func:`pykeen.nn.functional.ermlp_interaction`. .. math :: f(h, r, t) = W_2 ReLU(W_1 cat(h, r, t) + b_1) + b_2 @@ -612,8 +614,7 @@ def __init__( embedding_dim: int, hidden_dim: int, ): - """ - Initialize the interaction function. + """Initialize the interaction function. :param embedding_dim: The embedding vector dimension. @@ -657,7 +658,7 @@ def reset_parameters(self): # noqa: D102 class ERMLPEInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): - """Interaction function of ER-MLP.""" + """Interaction function of ER-MLP (E).""" def __init__( self, From 86e1fd3f5bf50505505249e940b7fab9f16bc405 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 12:02:53 +0100 Subject: [PATCH 293/690] Remove obsolete type annotations for RGCN --- src/pykeen/models/unimodal/rgcn.py | 33 +++++------------------------- 1 file changed, 5 insertions(+), 28 deletions(-) diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index ed6c4f4043..ff034ebd18 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -11,10 +11,9 @@ from torch.nn import functional from . import ComplEx, DistMult, ERMLP -from .. import EntityEmbeddingModel from ..base import ERModel from ...losses import Loss -from ...nn import Embedding, RepresentationModule +from ...nn import Embedding, Interaction, RepresentationModule from ...nn.modules import DistMultInteraction from ...triples import TriplesFactory from ...typing import DeviceHint @@ -432,31 +431,6 @@ class RGCN(ERModel): `_ """ - #: Interaction model used as decoder - base_model: EntityEmbeddingModel - - #: The blocks of the relation-specific weight matrices - #: shape: (num_relations, num_blocks, embedding_dim//num_blocks, embedding_dim//num_blocks) - blocks: Optional[nn.ParameterList] - - #: The base weight matrices to generate relation-specific weights - #: shape: (num_bases, embedding_dim, embedding_dim) - bases: Optional[nn.ParameterList] - - #: The relation-specific weights for each base - #: shape: (num_relations, num_bases) - att: Optional[nn.ParameterList] - - #: The biases for each layer (if used) - #: shape of each element: (embedding_dim,) - biases: Optional[nn.ParameterList] - - #: Batch normalization for each layer (if used) - batch_norms: Optional[nn.ModuleList] - - #: Activations for each layer (if used) - activations: Optional[nn.ModuleList] - #: The default strategy for optimizing the model's hyper-parameters hpo_default = dict( embedding_dim=dict(type=int, low=50, high=1000, q=50), @@ -480,6 +454,7 @@ class RGCN(ERModel): def __init__( self, triples_factory: TriplesFactory, + interaction: Optional[Interaction] = None, embedding_dim: int = 500, automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, @@ -504,6 +479,8 @@ def __init__( ): if triples_factory.create_inverse_triples: raise ValueError('R-GCN handles edges in an undirected manner.') + if interaction is None: + interaction = DistMultInteraction() entity_representations = RGCNRepresentations( triples_factory=triples_factory, @@ -533,7 +510,7 @@ def __init__( predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, random_seed=random_seed, - interaction=DistMultInteraction(), + interaction=interaction, entity_representations=entity_representations, relation_representations=relation_representations, ) From a10d21063aefca1bfdaf49bdc1d051f78ca1c318 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 12:06:42 +0100 Subject: [PATCH 294/690] Move get_in_canonical_shape to RepresentationModule base class --- src/pykeen/nn/emb.py | 37 +++++++++++++++++++++++++++++-------- 1 file changed, 29 insertions(+), 8 deletions(-) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index fc92fce828..dc0425da51 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -37,6 +37,31 @@ def forward(self, indices: Optional[torch.LongTensor] = None) -> torch.FloatTens """ raise NotImplementedError + def get_in_canonical_shape( + self, + indices: Optional[torch.LongTensor] = None, + reshape_dim: Optional[Sequence[int]] = None, + ) -> torch.FloatTensor: + """Get representations in canonical shape. + + :param indices: + The indices. If None, return all embeddings. + :param reshape_dim: + Optionally reshape the last dimension. + + :return: shape: (batch_size, num_embeddings, d) + """ + x = self(indices=indices) + if indices is None: + x = x.unsqueeze(dim=0) + else: + x = x.unsqueeze(dim=1) + if len(self.shape) > 1 and reshape_dim is None: + reshape_dim = self.shape + if reshape_dim is not None: + x = x.view(*x.shape[:-1], *reshape_dim) + return x + def reset_parameters(self) -> None: """Reset the module's parameters.""" @@ -236,13 +261,9 @@ def get_in_canonical_shape( :return: shape: (batch_size, num_embeddings, d) """ - x = self(indices=indices) - if indices is None: - x = x.unsqueeze(dim=0) - else: - x = x.unsqueeze(dim=1) if len(self.shape) > 1 and reshape_dim is None: reshape_dim = self.shape - if reshape_dim is not None: - x = x.view(*x.shape[:-1], *reshape_dim) - return x + return super().get_in_canonical_shape( + indices=indices, + reshape_dim=reshape_dim, + ) From 982cda401746b5ffc1f136783d615b4952800a81 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 12:07:10 +0100 Subject: [PATCH 295/690] Simplify RGCN --- src/pykeen/models/unimodal/rgcn.py | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index ff034ebd18..c58767258e 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -514,14 +514,3 @@ def __init__( entity_representations=entity_representations, relation_representations=relation_representations, ) - - def forward( - self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], - ) -> torch.FloatTensor: # noqa: D102 - h = self.entity_representations[0](indices=h_indices) - r = self.relation_representations[0](indices=r_indices) - t = self.entity_representations[0](indices=t_indices) - return self.decoder(h, r, t).unsqueeze(dim=-1) From ba337c44f9486ac9c99ca75f0dd5376ba99c30b5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 12:08:38 +0100 Subject: [PATCH 296/690] Remove Embedding specific code from RepresentationModule --- src/pykeen/nn/emb.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index dc0425da51..fbb0c412fa 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -56,8 +56,6 @@ def get_in_canonical_shape( x = x.unsqueeze(dim=0) else: x = x.unsqueeze(dim=1) - if len(self.shape) > 1 and reshape_dim is None: - reshape_dim = self.shape if reshape_dim is not None: x = x.view(*x.shape[:-1], *reshape_dim) return x From 3d9ccbfb3fd8df71f556e52bbceb1f4163fa896e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 12:09:14 +0100 Subject: [PATCH 297/690] Fix _check_constraint for RGCN --- tests/test_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index d6385e12af..374126e766 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -753,7 +753,7 @@ def _check_constraints(self): Enriched embeddings have to be reset. """ - assert self.model.entity_representations.enriched_embeddings is None + assert self.model.entity_representations[0].enriched_embeddings is None class TestRGCNBasis(_TestRGCN, unittest.TestCase): From 236ca616f7a761b66bbd097855b646f1caef698c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 12:11:41 +0100 Subject: [PATCH 298/690] Fix docstring --- src/pykeen/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 85cd2227f3..e54987290f 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -235,7 +235,7 @@ class Model(nn.Module, ABC): Subclasses of Model can decide however they want on how to store entities' and relations' representations, how they want to be looked up, and how they should - be scored. The :class:`GeneralModel` provides a commonly useful implementation + be scored. The :class:`ERModel` provides a commonly useful implementation which allows for the specification of one or more entity representations and one or more relation representations in the form of :class:`pykeen.nn.Embedding` as well as a matching instance of a :class:`pykeen.nn.Interaction`. From 8a2b7927e17bcb2926badc5b8290b17c93d0cea2 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 12:20:56 +0100 Subject: [PATCH 299/690] Make test_save_load_model_state more generic --- tests/test_models.py | 26 ++++++++++++++------------ 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 374126e766..be56538c04 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -296,23 +296,25 @@ def test_save_load_model_state(self): **(self.model_kwargs or {}), ).to_device_() - def _equal_embeddings(a: RepresentationModule, b: RepresentationModule) -> bool: - """Test whether two embeddings are equal.""" - return (a(indices=None) == b(indices=None)).all() - - if isinstance(original_model, EntityEmbeddingModel): - assert not _equal_embeddings(original_model.entity_embeddings, loaded_model.entity_embeddings) - if isinstance(original_model, EntityRelationEmbeddingModel): - assert not _equal_embeddings(original_model.relation_embeddings, loaded_model.relation_embeddings) + def _equal_weights(a: nn.Module, b: nn.Module) -> bool: + """Test whether two modules are equal.""" + a_state = a.state_dict() + b_state = b.state_dict() + if a_state.keys() != b_state.keys(): + return False + for key, original_value in a_state.items(): + if not torch.allclose(original_value, b_state[key]): + return False + return True + + assert not _equal_weights(original_model, loaded_model) with tempfile.TemporaryDirectory() as tmpdirname: file_path = os.path.join(tmpdirname, 'test.pt') original_model.save_state(path=file_path) loaded_model.load_state(path=file_path) - if isinstance(original_model, EntityEmbeddingModel): - assert _equal_embeddings(original_model.entity_embeddings, loaded_model.entity_embeddings) - if isinstance(original_model, EntityRelationEmbeddingModel): - assert _equal_embeddings(original_model.relation_embeddings, loaded_model.relation_embeddings) + + assert _equal_weights(original_model, loaded_model) @property def cli_extras(self): From 92ceb88e516441e06326ea4208b6e89993a5752d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 12:22:39 +0100 Subject: [PATCH 300/690] Remove EntityEmbeddingModel --- src/pykeen/cli.py | 7 ++--- src/pykeen/models/__init__.py | 4 +-- src/pykeen/models/base.py | 58 +---------------------------------- tests/test_models.py | 6 +--- 4 files changed, 6 insertions(+), 69 deletions(-) diff --git a/src/pykeen/cli.py b/src/pykeen/cli.py index 188ecb8eff..ea292bc52a 100644 --- a/src/pykeen/cli.py +++ b/src/pykeen/cli.py @@ -28,8 +28,8 @@ from .hpo.cli import optimize from .hpo.samplers import samplers as hpo_samplers_dict from .losses import losses as losses_dict -from .models import models as models_dict -from .models.base import EntityEmbeddingModel, EntityRelationEmbeddingModel, Model +from .models import ERModel, models as models_dict +from .models.base import Model from .models.cli import build_cli_from_cls from .optimizers import optimizers as optimizers_dict from .regularizers import regularizers as regularizers_dict @@ -94,8 +94,7 @@ def parameters(): base_parameters = set(chain( Model.__init__.__annotations__, - EntityEmbeddingModel.__init__.__annotations__, - EntityRelationEmbeddingModel.__init__.__annotations__, + ERModel.__init__.__annotations__, )) _hyperparameter_usage = sorted( (k, v) diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index 25d58a71fb..0250174d94 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -9,7 +9,7 @@ from typing import Mapping, Set, Type, Union from .base import ( # noqa:F401 - DoubleRelationEmbeddingModel, ERModel, EntityEmbeddingModel, EntityRelationEmbeddingModel, Model, MultimodalModel, + DoubleRelationEmbeddingModel, ERModel, Model, MultimodalModel, SingleVectorEmbeddingModel, TwoSideEmbeddingModel, TwoVectorEmbeddingModel, ) from .multimodal import ComplExLiteral, DistMultLiteral @@ -71,8 +71,6 @@ SingleVectorEmbeddingModel, DoubleRelationEmbeddingModel, TwoSideEmbeddingModel, - EntityEmbeddingModel, - EntityRelationEmbeddingModel, TwoVectorEmbeddingModel, MultimodalModel, } diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index e54987290f..fffce6475f 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -26,16 +26,13 @@ from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory from ..typing import ( - Constrainer, DeviceHint, HeadRepresentation, Initializer, MappedTriples, Normalizer, - RelationRepresentation, Representation, TailRepresentation, + DeviceHint, HeadRepresentation, MappedTriples, RelationRepresentation, Representation, TailRepresentation, ) from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed __all__ = [ 'Model', 'ERModel', - 'EntityEmbeddingModel', - 'EntityRelationEmbeddingModel', 'SingleVectorEmbeddingModel', 'DoubleRelationEmbeddingModel', 'TwoVectorEmbeddingModel', @@ -1233,59 +1230,6 @@ def forward( return scores -class EntityEmbeddingModel(ERModel, autoreset=False): - """A base module for most KGE models that have one embedding for entities.""" - - # TODO: deprecated - - def __init__( - self, - triples_factory: TriplesFactory, - interaction: Interaction[Representation, RelationRepresentation, Representation], - embedding_dim: int = 50, - loss: Optional[Loss] = None, - predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, - preferred_device: DeviceHint = None, - random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, - entity_initializer: Optional[Initializer] = None, - entity_initializer_kwargs: Optional[Mapping[str, Any]] = None, - entity_normalizer: Optional[Normalizer] = None, - entity_normalizer_kwargs: Optional[Mapping[str, Any]] = None, - entity_constrainer: Optional[Constrainer] = None, - entity_constrainer_kwargs: Optional[Mapping[str, Any]] = None, - ) -> None: - """Initialize the entity embedding model. - - :param embedding_dim: - The embedding dimensionality. Exact usages depends on the specific model subclass. - - .. seealso:: Constructor of the base class :class:`pykeen.models.Model` - """ - self.embedding_dim = embedding_dim - super().__init__( - triples_factory=triples_factory, - interaction=interaction, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, - regularizer=regularizer, - predict_with_sigmoid=predict_with_sigmoid, - entity_representations=Embedding( - num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, - initializer=entity_initializer, - initializer_kwargs=entity_initializer_kwargs, - normalizer=entity_normalizer, - normalizer_kwargs=entity_normalizer_kwargs, - constrainer=entity_constrainer, - constrainer_kwargs=entity_constrainer_kwargs, - ), - ) - - class EntityRelationEmbeddingModel(ERModel, autoreset=False): """A base module for KGE models that have different embeddings for entities and relations.""" diff --git a/tests/test_models.py b/tests/test_models.py index be56538c04..12483d75e0 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -24,9 +24,7 @@ from pykeen.datasets.nations import NATIONS_TEST_PATH, NATIONS_TRAIN_PATH, Nations from pykeen.models import _BASE_MODELS, _MODELS from pykeen.models.base import ( - DoubleRelationEmbeddingModel, ERModel, EntityEmbeddingModel, - EntityRelationEmbeddingModel, - Model, + DoubleRelationEmbeddingModel, ERModel, Model, MultimodalModel, SingleVectorEmbeddingModel, TwoSideEmbeddingModel, @@ -51,8 +49,6 @@ 'DummyModel', MultimodalModel.__name__, DoubleRelationEmbeddingModel.__name__, - EntityEmbeddingModel.__name__, - EntityRelationEmbeddingModel.__name__, SingleVectorEmbeddingModel.__name__, TwoVectorEmbeddingModel.__name__, TwoSideEmbeddingModel.__name__, From bca891c1a2885cf12f70238468cd10c0edca3c85 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 12:28:46 +0100 Subject: [PATCH 301/690] Update MockModel --- tests/test_early_stopping.py | 39 ++++++++++++++++++------------------ 1 file changed, 19 insertions(+), 20 deletions(-) diff --git a/tests/test_early_stopping.py b/tests/test_early_stopping.py index 587e250530..6310cc0ddf 100644 --- a/tests/test_early_stopping.py +++ b/tests/test_early_stopping.py @@ -13,7 +13,7 @@ from pykeen.evaluation import Evaluator, MetricResults, RankBasedEvaluator, RankBasedMetricResults from pykeen.evaluation.rank_based_evaluator import RANK_TYPES, SIDES from pykeen.models import TransE -from pykeen.models.base import EntityRelationEmbeddingModel, Model +from pykeen.models.base import Model from pykeen.stoppers.early_stopping import EarlyStopper, is_improvement from pykeen.trackers import MLFlowResultTracker from pykeen.training import SLCWATrainingLoop @@ -102,33 +102,32 @@ def __repr__(self): # noqa: D105 return f'{self.__class__.__name__}(losses={self.losses})' -class MockModel(EntityRelationEmbeddingModel): +class MockModel(Model): """A mock model returning fake scores.""" def __init__(self, triples_factory: TriplesFactory, automatic_memory_optimization: bool): super().__init__( triples_factory=triples_factory, automatic_memory_optimization=automatic_memory_optimization, - interaction=None, ) - num_entities = self.num_entities - self.scores = torch.arange(num_entities, dtype=torch.float) - def _generate_fake_scores(self, batch: torch.LongTensor) -> torch.FloatTensor: - """Generate fake scores s[b, i] = i of size (batch_size, num_entities).""" - batch_size = batch.shape[0] - batch_scores = self.scores.view(1, -1).repeat(batch_size, 1) - assert batch_scores.shape == (batch_size, self.num_entities) - return batch_scores - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._generate_fake_scores(batch=hrt_batch) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._generate_fake_scores(batch=hr_batch) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._generate_fake_scores(batch=rt_batch) + def forward( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: # noqa: D102 + # (batch_size, num_heads, num_relations, num_tails) + bss, (nh, nr, nt) = zip(*( + (1, num) if ind is None else (ind.shape[0], 1) + for ind, num in ( + (h_indices, self.num_entities), + (r_indices, self.num_relations), + (t_indices, self.num_entities), + ) + )) + bs = max(bss) + return torch.rand(bs, nh, nr, nt) class LogCallWrapper: From 06f5bfab20210126e9a3feae93cda0e5e8b056b7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 12:30:35 +0100 Subject: [PATCH 302/690] Move MockModel and re-use it --- src/pykeen/models/base.py | 29 +++++++++++++++++++++++++++++ tests/test_early_stopping.py | 31 +------------------------------ tests/test_evaluators.py | 33 ++------------------------------- 3 files changed, 32 insertions(+), 61 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index fffce6475f..80dfccdc7a 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1576,3 +1576,32 @@ def forward( (self.entity_representations[1], self.relation_representations[1], self.entity_representations[0]), ) ) + + +class MockModel(Model): + """A mock model returning fake scores.""" + # TODO: Where to put this? + + def __init__(self, triples_factory: TriplesFactory, automatic_memory_optimization: bool): + super().__init__( + triples_factory=triples_factory, + automatic_memory_optimization=automatic_memory_optimization, + ) + + def forward( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: # noqa: D102 + # (batch_size, num_heads, num_relations, num_tails) + bss, (nh, nr, nt) = zip(*( + (1, num) if ind is None else (ind.shape[0], 1) + for ind, num in ( + (h_indices, self.num_entities), + (r_indices, self.num_relations), + (t_indices, self.num_entities), + ) + )) + bs = max(bss) + return torch.rand(bs, nh, nr, nt, requires_grad=True) \ No newline at end of file diff --git a/tests/test_early_stopping.py b/tests/test_early_stopping.py index 6310cc0ddf..e62dfa4979 100644 --- a/tests/test_early_stopping.py +++ b/tests/test_early_stopping.py @@ -13,11 +13,10 @@ from pykeen.evaluation import Evaluator, MetricResults, RankBasedEvaluator, RankBasedMetricResults from pykeen.evaluation.rank_based_evaluator import RANK_TYPES, SIDES from pykeen.models import TransE -from pykeen.models.base import Model +from pykeen.models.base import MockModel, Model from pykeen.stoppers.early_stopping import EarlyStopper, is_improvement from pykeen.trackers import MLFlowResultTracker from pykeen.training import SLCWATrainingLoop -from pykeen.triples import TriplesFactory from pykeen.typing import MappedTriples @@ -102,34 +101,6 @@ def __repr__(self): # noqa: D105 return f'{self.__class__.__name__}(losses={self.losses})' -class MockModel(Model): - """A mock model returning fake scores.""" - - def __init__(self, triples_factory: TriplesFactory, automatic_memory_optimization: bool): - super().__init__( - triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, - ) - - def forward( - self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], - ) -> torch.FloatTensor: # noqa: D102 - # (batch_size, num_heads, num_relations, num_tails) - bss, (nh, nr, nt) = zip(*( - (1, num) if ind is None else (ind.shape[0], 1) - for ind, num in ( - (h_indices, self.num_entities), - (r_indices, self.num_relations), - (t_indices, self.num_entities), - ) - )) - bs = max(bss) - return torch.rand(bs, nh, nr, nt) - - class LogCallWrapper: """An object which wraps functions and checks whether they have been called.""" diff --git a/tests/test_evaluators.py b/tests/test_evaluators.py index 54107db0c2..da9418b0bb 100644 --- a/tests/test_evaluators.py +++ b/tests/test_evaluators.py @@ -15,7 +15,7 @@ from pykeen.evaluation.rank_based_evaluator import RANK_TYPES, SIDES, compute_rank_from_scores from pykeen.evaluation.sklearn import SklearnEvaluator, SklearnMetricResults from pykeen.models import TransE -from pykeen.models.base import EntityRelationEmbeddingModel, Model +from pykeen.models.base import MockModel, Model from pykeen.triples import TriplesFactory from pykeen.typing import MappedTriples @@ -428,35 +428,6 @@ def __repr__(self): # noqa: D105 return f'{self.__class__.__name__}(losses={self.losses})' -class DummyModel(EntityRelationEmbeddingModel): - """A dummy model returning fake scores.""" - - def __init__(self, triples_factory: TriplesFactory, automatic_memory_optimization: bool): - super().__init__( - triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, - interaction=..., - ) - num_entities = self.num_entities - self.scores = torch.arange(num_entities, dtype=torch.float) - - def _generate_fake_scores(self, batch: torch.LongTensor) -> torch.FloatTensor: - """Generate fake scores s[b, i] = i of size (batch_size, num_entities).""" - batch_size = batch.shape[0] - batch_scores = self.scores.view(1, -1).repeat(batch_size, 1) - assert batch_scores.shape == (batch_size, self.num_entities) - return batch_scores - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._generate_fake_scores(batch=hrt_batch) - - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._generate_fake_scores(batch=hr_batch) - - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa: D102 - return self._generate_fake_scores(batch=rt_batch) - - class TestEvaluationStructure(unittest.TestCase): """Tests for testing the correct structure of the evaluation procedure.""" @@ -465,7 +436,7 @@ def setUp(self): self.counter = 1337 self.evaluator = DummyEvaluator(counter=self.counter, filtered=True) self.triples_factory = Nations().training - self.model = DummyModel(triples_factory=self.triples_factory, automatic_memory_optimization=False) + self.model = MockModel(triples_factory=self.triples_factory, automatic_memory_optimization=False) def test_evaluation_structure(self): """Test if the evaluator has a balanced call of head and tail processors.""" From 8140cf6036272319f2733167d5c34282d80876f7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 12:34:00 +0100 Subject: [PATCH 303/690] Update Literal model --- src/pykeen/models/base.py | 2 +- src/pykeen/models/multimodal/complex_literal.py | 10 +++++++++- src/pykeen/models/multimodal/distmult_literal.py | 10 +++++++++- 3 files changed, 19 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 80dfccdc7a..31d6536012 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1286,7 +1286,7 @@ def __init__( ) -class MultimodalModel(EntityRelationEmbeddingModel, autoreset=False): +class MultimodalModel(ERModel, autoreset=False): """A multimodal KGE model.""" diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index bb2588ad3f..53ae838546 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -10,6 +10,7 @@ from ..base import MultimodalModel from ...losses import BCEWithLogitsLoss, Loss +from ...nn import Embedding from ...nn.modules import ComplExInteraction from ...triples import TriplesNumericLiteralsFactory from ...typing import DeviceHint @@ -52,11 +53,18 @@ def __init__( super().__init__( triples_factory=triples_factory, interaction=ComplExInteraction(), - embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, + entity_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + ), + relation_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=embedding_dim, + ), ) # Literal diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 8586af010d..7996331b6e 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -10,6 +10,7 @@ from ..base import MultimodalModel from ...losses import Loss +from ...nn import Embedding from ...nn.modules import DistMultInteraction from ...triples import TriplesNumericLiteralsFactory from ...typing import DeviceHint @@ -45,11 +46,18 @@ def __init__( super().__init__( interaction=DistMultInteraction(), triples_factory=triples_factory, - embedding_dim=embedding_dim, automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, + entity_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + ), + relation_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=embedding_dim, + ), ) numeric_literals = triples_factory.numeric_literals From 50ed1be109451bbf85ea0158610639608dec5bfc Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 12:48:56 +0100 Subject: [PATCH 304/690] Update test_model_mode.py --- src/pykeen/models/base.py | 23 ++++++++++++++++------- tests/test_model_mode.py | 38 ++++++++++++-------------------------- 2 files changed, 28 insertions(+), 33 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 31d6536012..33c7e3ca07 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1580,9 +1580,10 @@ def forward( class MockModel(Model): """A mock model returning fake scores.""" + # TODO: Where to put this? - def __init__(self, triples_factory: TriplesFactory, automatic_memory_optimization: bool): + def __init__(self, triples_factory: TriplesFactory, automatic_memory_optimization: bool = True): super().__init__( triples_factory=triples_factory, automatic_memory_optimization=automatic_memory_optimization, @@ -1595,13 +1596,21 @@ def forward( t_indices: Optional[torch.LongTensor], ) -> torch.FloatTensor: # noqa: D102 # (batch_size, num_heads, num_relations, num_tails) - bss, (nh, nr, nt) = zip(*( - (1, num) if ind is None else (ind.shape[0], 1) - for ind, num in ( + scores = torch.zeros(1, 1, 1, 1, requires_grad=True) # for requires_grad + # reproducible scores + for i, (ind, num) in enumerate( + ( (h_indices, self.num_entities), (r_indices, self.num_relations), (t_indices, self.num_entities), ) - )) - bs = max(bss) - return torch.rand(bs, nh, nr, nt, requires_grad=True) \ No newline at end of file + ): + shape = [1, 1, 1, 1] + if ind is None: + shape[i + 1] = num + delta = torch.arange(num) + else: + shape[0] = len(ind) + delta = ind + scores = scores + delta.float().view(*shape) + return scores diff --git a/tests/test_model_mode.py b/tests/test_model_mode.py index 34ff0df60b..0ea1ccb3de 100644 --- a/tests/test_model_mode.py +++ b/tests/test_model_mode.py @@ -4,13 +4,13 @@ import unittest from dataclasses import dataclass +from unittest.mock import MagicMock import torch from pykeen.datasets import Nations -from pykeen.models import TransE -from pykeen.models.base import EntityRelationEmbeddingModel -from pykeen.nn.modules import Interaction +from pykeen.models import Model, TransE +from pykeen.models.base import MockModel from pykeen.triples import TriplesFactory from pykeen.utils import resolve_device @@ -21,7 +21,7 @@ class TestBaseModel(unittest.TestCase): batch_size: int embedding_dim: int factory: TriplesFactory - model: EntityRelationEmbeddingModel + model: Model def setUp(self) -> None: """Set up the test case with a triples factory and TransE as an example model.""" @@ -81,9 +81,9 @@ class TestBaseModelScoringFunctions(unittest.TestCase): def setUp(self): """Prepare for testing the scoring functions.""" self.generator = torch.random.manual_seed(seed=42) - self.triples_factory = MinimalTriplesFactory + self.triples_factory = MagicMock(num_relations=2, num_entities=2) self.device = resolve_device() - self.model = SimpleInteractionModel(triples_factory=self.triples_factory).to(self.device) + self.model = MockModel(triples_factory=self.triples_factory).to(self.device) def test_alignment_of_score_t_fall_back(self) -> None: """Test if ``BaseModule.score_t`` aligns with ``BaseModule.score_hrt``.""" @@ -106,8 +106,8 @@ def test_alignment_of_score_t_fall_back(self) -> None: device=self.device, ) scores_t_function = self.model.score_t(hr_batch=hr_batch).flatten() - scores_hrt_function = self.model.score_hrt(hrt_batch=hrt_batch) - assert all(scores_t_function == scores_hrt_function) + scores_hrt_function = self.model.score_hrt(hrt_batch=hrt_batch).flatten() + assert (scores_t_function == scores_hrt_function).all() def test_alignment_of_score_h_fall_back(self) -> None: """Test if ``BaseModule.score_h`` aligns with ``BaseModule.score_hrt``.""" @@ -130,8 +130,8 @@ def test_alignment_of_score_h_fall_back(self) -> None: device=self.device, ) scores_h_function = self.model.score_h(rt_batch=rt_batch).flatten() - scores_hrt_function = self.model.score_hrt(hrt_batch=hrt_batch) - assert all(scores_h_function == scores_hrt_function) + scores_hrt_function = self.model.score_hrt(hrt_batch=hrt_batch).flatten() + assert (scores_h_function == scores_hrt_function).all() def test_alignment_of_score_r_fall_back(self) -> None: """Test if ``BaseModule.score_r`` aligns with ``BaseModule.score_hrt``.""" @@ -154,22 +154,8 @@ def test_alignment_of_score_r_fall_back(self) -> None: device=self.device, ) scores_r_function = self.model.score_r(ht_batch=ht_batch).flatten() - scores_hrt_function = self.model.score_hrt(hrt_batch=hrt_batch) - assert all(scores_r_function == scores_hrt_function) - - -class SimpleInteraction(Interaction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor]): - """A simple interaction as the sum of the vectors.""" - - def forward(self, h, r, t) -> torch.FloatTensor: # noqa:D102 - return torch.sum(h + r + t, dim=1) - - -class SimpleInteractionModel(EntityRelationEmbeddingModel): - """A model with a simple interaction function for testing the base model.""" - - def __init__(self, *args, **kwargs): - super().__init__(*args, interaction=SimpleInteraction(), **kwargs) + scores_hrt_function = self.model.score_hrt(hrt_batch=hrt_batch).flatten() + assert (scores_r_function == scores_hrt_function).all() @dataclass From 351e072af543d31cea1183b77989c523292a6d64 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 12:49:18 +0100 Subject: [PATCH 305/690] Remove deprecated base class --- src/pykeen/models/base.py | 58 +-------------------------------------- 1 file changed, 1 insertion(+), 57 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 33c7e3ca07..e4ba16f045 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -26,7 +26,7 @@ from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory from ..typing import ( - DeviceHint, HeadRepresentation, MappedTriples, RelationRepresentation, Representation, TailRepresentation, + DeviceHint, HeadRepresentation, MappedTriples, RelationRepresentation, TailRepresentation, ) from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed @@ -1230,62 +1230,6 @@ def forward( return scores -class EntityRelationEmbeddingModel(ERModel, autoreset=False): - """A base module for KGE models that have different embeddings for entities and relations.""" - - # TODO: Deprecated. - - def __init__( - self, - triples_factory: TriplesFactory, - interaction: Interaction[Representation, Representation, Representation], - embedding_dim: int = 50, - relation_dim: Optional[int] = None, - loss: Optional[Loss] = None, - predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, - preferred_device: DeviceHint = None, - random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, - embedding_specification: Optional[EmbeddingSpecification] = None, - relation_embedding_specification: Optional[EmbeddingSpecification] = None, - ) -> None: - """Initialize the entity embedding model. - - :param relation_dim: - The relation embedding dimensionality. If not given, defaults to same size as entity embedding - dimension. - - .. seealso:: Constructor of the base class :class:`pykeen.models.Model` - .. seealso:: Constructor of the base class :class:`pykeen.models.EntityEmbeddingModel` - """ - self.embedding_dim = embedding_dim - # Default for relation dimensionality - if relation_dim is None: - relation_dim = embedding_dim - self.relation_dim = relation_dim - super().__init__( - triples_factory=triples_factory, - interaction=interaction, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, - regularizer=regularizer, - predict_with_sigmoid=predict_with_sigmoid, - entity_representations=Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, - specification=embedding_specification, - ), - relation_representations=Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - embedding_dim=relation_dim, - specification=relation_embedding_specification, - ), - ) - - class MultimodalModel(ERModel, autoreset=False): """A multimodal KGE model.""" From d93e783a74683416b0d046acb095302ac98e8c4f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 13:08:49 +0100 Subject: [PATCH 306/690] Extract helpers methods --- src/pykeen/models/base.py | 60 ++++++++++++++++++++++++++++++--------- 1 file changed, 46 insertions(+), 14 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index e4ba16f045..31639df223 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1201,6 +1201,51 @@ def forward( :return: shape: (batch_size, num_heads, num_relations, num_tails) The score for each triple. """ + h, r, t = self._get_representations(h_indices, r_indices, t_indices) + scores = self.interaction(h=h, r=r, t=t) + scores = self._repeat_scores_if_necessary(scores, h_indices, r_indices, t_indices) + return scores + + def _repeat_scores_if_necessary( + self, + scores: torch.FloatTensor, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: + repeat_relations = len(self.relation_representations) == 0 + repeat_entities = len(self.entity_representations) == 0 + + if not (repeat_entities or repeat_relations): + return scores + + repeats = [1, 1, 1, 1] + + for i, (flag, ind, num) in enumerate(( + (repeat_entities, h_indices, self.num_entities), + (repeat_relations, r_indices, self.num_relations), + (repeat_entities, t_indices, self.num_entities), + ), start=1): + if flag: + if ind is None: + repeats[i] = num + else: + batch_size = len(ind) + if scores.shape[0] < batch_size: + repeats[0] = batch_size + + return scores.repeat(*repeats) + + def _get_representations( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> Tuple[ + Union[torch.FloatTensor, Sequence[torch.FloatTensor]], + Union[torch.FloatTensor, Sequence[torch.FloatTensor]], + Union[torch.FloatTensor, Sequence[torch.FloatTensor]] + ]: h, r, t = [ [ representation.get_in_canonical_shape(indices=indices) @@ -1214,20 +1259,7 @@ def forward( ] # normalization h, r, t = [x[0] if len(x) == 1 else x for x in (h, r, t)] - - scores = self.interaction(h=h, r=r, t=t) - if len(self.relation_representations) == 0: - # same score for all relations - repeats = [1, 1, 1, 1] - if r_indices is None: - repeats[2] = self.num_relations - else: - relation_batch_size = len(r_indices) - if scores.shape[0] < relation_batch_size: - repeats[0] = relation_batch_size - scores = scores.repeat(*repeats) - - return scores + return h, r, t class MultimodalModel(ERModel, autoreset=False): From a2a6ecfec82c96f9b3d77b5004343da60b956e1e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 13:13:34 +0100 Subject: [PATCH 307/690] Refactor literal model --- src/pykeen/models/base.py | 3 +- .../models/multimodal/distmult_literal.py | 219 ++++++++---------- 2 files changed, 102 insertions(+), 120 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 31639df223..c903b56387 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1203,8 +1203,7 @@ def forward( """ h, r, t = self._get_representations(h_indices, r_indices, t_indices) scores = self.interaction(h=h, r=r, t=t) - scores = self._repeat_scores_if_necessary(scores, h_indices, r_indices, t_indices) - return scores + return self._repeat_scores_if_necessary(scores, h_indices, r_indices, t_indices) def _repeat_scores_if_necessary( self, diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 7996331b6e..cf02591260 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -10,19 +10,102 @@ from ..base import MultimodalModel from ...losses import Loss -from ...nn import Embedding +from ...nn import Embedding, Interaction +from ...nn.emb import EmbeddingSpecification from ...nn.modules import DistMultInteraction +from ...regularizers import Regularizer from ...triples import TriplesNumericLiteralsFactory -from ...typing import DeviceHint -from ...utils import slice_triples +from ...typing import DeviceHint, HeadRepresentation, RelationRepresentation, TailRepresentation __all__ = [ 'DistMultLiteral', ] +class LiteralRepresentations(Embedding): + """Literal representations.""" + + def __init__( + self, + numeric_literals: torch.FloatTensor, + ): + num_embeddings, embedding_dim = numeric_literals.shape + super().__init__( + num_embeddings=num_embeddings, + embedding_dim=embedding_dim, + initializer=lambda x: numeric_literals, # initialize with the literals + ) + # freeze + self._embeddings.requires_grad_(False) + + +class LiteralModel(MultimodalModel): + def __init__( + self, + triples_factory: TriplesNumericLiteralsFactory, + embedding_dim: int, + interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], + dropout: float = 0.0, + entity_specification: Optional[EmbeddingSpecification] = None, + relation_specification: Optional[EmbeddingSpecification] = None, + loss: Optional[Loss] = None, + predict_with_sigmoid: bool = False, + automatic_memory_optimization: Optional[bool] = None, + preferred_device: DeviceHint = None, + random_seed: Optional[int] = None, + regularizer: Optional[Regularizer] = None, + ): + super().__init__( + triples_factory=triples_factory, + interaction=interaction, + loss=loss, + predict_with_sigmoid=predict_with_sigmoid, + automatic_memory_optimization=automatic_memory_optimization, + preferred_device=preferred_device, + random_seed=random_seed, + regularizer=regularizer, + entity_representations=[ + # entity embeddings + Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + specification=entity_specification, + ), + # Entity literals + LiteralRepresentations( + numeric_literals=torch.as_tensor(triples_factory.numeric_literals, dtype=torch.float32), + ), + ], + relation_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=embedding_dim, + specification=relation_specification, + ), + ) + self.trans = nn.Sequential( + nn.Linear(embedding_dim + triples_factory.numeric_literals.shape[1], embedding_dim), + nn.Dropout(dropout), + ) + + def forward( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + ) -> torch.FloatTensor: # noqa: D102 + h, r, t = self._get_representations(h_indices, r_indices, t_indices) + # combine entity embeddings + literals + h, t = [ + self.trans(torch.cat(x, dim=-1)) + for x in (h, t) + ] + scores = self.interaction(h=h, r=r, t=t) + return self._repeat_scores_if_necessary(scores, h_indices, r_indices, t_indices) + + # TODO: Check entire build of the model -class DistMultLiteral(MultimodalModel): +# TODO: There are no tests +class DistMultLiteral(LiteralModel): """An implementation of DistMultLiteral from [agustinus2018]_.""" #: The default strategy for optimizing the model's hyper-parameters @@ -42,124 +125,24 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, + predict_with_sigmoid: bool = False, + regularizer: Optional[Regularizer] = None, ) -> None: super().__init__( - interaction=DistMultInteraction(), triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, + embedding_dim=embedding_dim, + interaction=DistMultInteraction(), + dropout=input_dropout, + entity_specification=EmbeddingSpecification( + initializer=xavier_normal_, + ), + relation_specification=EmbeddingSpecification( + initializer=xavier_normal_, + ), loss=loss, + predict_with_sigmoid=predict_with_sigmoid, + automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, - entity_representations=Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, - ), - relation_representations=Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim, - ), - ) - - numeric_literals = triples_factory.numeric_literals - - # Embeddings - self.relation_embeddings = None - self.numeric_literals = nn.Embedding.from_pretrained( - torch.tensor(numeric_literals, dtype=torch.float, device=self.device), freeze=True, - ) - # Number of columns corresponds to number of literals - self.num_of_literals = self.numeric_literals.weight.data.shape[1] - self.linear_transformation = nn.Linear(self.embedding_dim + self.num_of_literals, self.embedding_dim) - self.input_dropout = torch.nn.Dropout(input_dropout) - self._init_embeddings() - - def _init_embeddings(self): - """Initialize the entities and relation embeddings based on the XAVIER initialization.""" - super()._init_embeddings() - self.relation_embeddings = nn.Embedding(self.num_relations, self.embedding_dim) - xavier_normal_(self.entity_embeddings.weight.data) - xavier_normal_(self.relation_embeddings.weight.data) - - @staticmethod - def _get_embeddings(elements, embedding_module, embedding_dim): - return embedding_module(elements).view(-1, embedding_dim) - - def _get_literals(self, heads, tails): - return ( - self._get_embeddings( - elements=heads, - embedding_module=self.numeric_literals, - embedding_dim=self.num_of_literals, - ), - self._get_embeddings( - elements=tails, - embedding_module=self.numeric_literals, - embedding_dim=self.num_of_literals, - ), - ) - - def _get_triple_embeddings(self, heads, relations, tails): - return ( - self._get_embeddings( - elements=heads, - embedding_module=self.entity_embeddings, - embedding_dim=self.embedding_dim, - ), - self._get_embeddings( - elements=relations, - embedding_module=self.relation_embeddings, - embedding_dim=self.embedding_dim, - ), - self._get_embeddings( - elements=tails, - embedding_module=self.entity_embeddings, - embedding_dim=self.embedding_dim, - ), - ) - - def _apply_g_function(self, entity_embeddings, literals): - """Concatenate the entities with its literals and apply the g function which is a linear transformation in this model. - - :param entity_embeddings: batch_size x self.embedding_dim - :param literals: batch_size x self.num_literals - :return: - """ - return self.linear_transformation(torch.cat([entity_embeddings, literals], dim=1)) - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa:D102 - raise NotImplementedError - - def score_t(self, hr_batch: torch.Tensor) -> torch.Tensor: - """Forward pass using right side (tail) prediction for training with the LCWA.""" - heads, relations, tails = slice_triples(hr_batch) - head_embs, relation_embs, tail_embs = self._get_triple_embeddings( - heads=heads, - relations=relations, - tails=tails, + regularizer=regularizer, ) - head_literals, tail_literals = self._get_literals(heads=heads, tails=tails) - - g_heads = self._apply_g_function(entity_embeddings=head_embs, literals=head_literals) - g_tails = self._apply_g_function(entity_embeddings=tail_embs, literals=tail_literals) - - # apply dropout - g_heads = self.input_dropout(g_heads) - g_tails = self.input_dropout(g_tails) - - # -, because lower score shall correspond to a more plausible triple. - scores = - torch.sum(g_heads * relation_embs * g_tails, dim=1) - return scores - - # TODO check if this is the same as the BaseModule - def compute_mr_loss(self, positive_scores: torch.Tensor, negative_scores: torch.Tensor) -> torch.Tensor: - """Compute the mean ranking loss for the positive and negative scores.""" - # Choose y = -1 since a smaller score is better. - # In TransE for example, the scores represent distances - if not self.compute_mr_loss: - raise RuntimeError( - 'The chosen loss does not allow the calculation of Margin Ranking losses. ' - 'Please use the compute_label_loss method instead', - ) - y = torch.ones_like(negative_scores, device=self.device) * -1 - loss = self.loss(positive_scores, negative_scores, y) - return loss From b1e6aa0f6f76c6747de6eefe8995484289e4feb7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 13:14:11 +0100 Subject: [PATCH 308/690] Add docstring --- src/pykeen/models/multimodal/distmult_literal.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index cf02591260..6cecf0c4ac 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -8,7 +8,7 @@ import torch.nn as nn from torch.nn.init import xavier_normal_ -from ..base import MultimodalModel +from .. import ERModel from ...losses import Loss from ...nn import Embedding, Interaction from ...nn.emb import EmbeddingSpecification @@ -39,7 +39,9 @@ def __init__( self._embeddings.requires_grad_(False) -class LiteralModel(MultimodalModel): +class LiteralModel(ERModel): + """Base class for models with entity literals.""" + def __init__( self, triples_factory: TriplesNumericLiteralsFactory, From 6d8da6de63c4d23dac3c7df1f0233e03099fe69a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 13:24:16 +0100 Subject: [PATCH 309/690] Refactor literal models --- .../models/multimodal/complex_literal.py | 153 +++++++----------- .../models/multimodal/distmult_literal.py | 16 +- src/pykeen/utils.py | 8 + 3 files changed, 73 insertions(+), 104 deletions(-) diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index 53ae838546..e5353e0c37 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -8,21 +8,56 @@ import torch.nn as nn from torch.nn.init import xavier_normal_ -from ..base import MultimodalModel +from .distmult_literal import LiteralModel from ...losses import BCEWithLogitsLoss, Loss -from ...nn import Embedding +from ...nn.emb import EmbeddingSpecification from ...nn.modules import ComplExInteraction +from ...regularizers import Regularizer from ...triples import TriplesNumericLiteralsFactory from ...typing import DeviceHint -from ...utils import slice_doubles +from ...utils import combine_complex, split_complex __all__ = [ 'ComplExLiteral', ] +class ComplexLiteralCombination(nn.Module): + """Separately transform real and imaginary part.""" + + def __init__( + self, + embedding_dim: int, + num_of_literals: int, + dropout: float = 0.0, + ): + super().__init__() + self.real = nn.Sequential( + nn.Dropout(dropout), + nn.Linear(embedding_dim + num_of_literals, embedding_dim), + torch.nn.Tanh(), + ) + self.imag = nn.Sequential( + nn.Dropout(dropout), + nn.Linear(embedding_dim + num_of_literals, embedding_dim), + torch.nn.Tanh(), + ) + self.embedding_dim = embedding_dim + + def forward( + self, + x: torch.FloatTensor, + ) -> torch.FloatTensor: + x, literal = x[..., :self.embedding_dim], x[..., self.embedding_dim:] + x_re, x_im = split_complex(x) + x_re = self.real(torch.cat([x_re, literal], dim=-1)) + x_im = self.real(torch.cat([x_im, literal], dim=-1)) + return combine_complex(x_re=x_re, x_im=x_im) + + # TODO: Check entire build of the model -class ComplExLiteral(MultimodalModel): +# TODO: There are no tests. +class ComplExLiteral(LiteralModel): """An implementation of ComplexLiteral from [agustinus2018]_ based on the LCWA training approach.""" #: The default strategy for optimizing the model's hyper-parameters @@ -46,107 +81,31 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, input_dropout: float = 0.2, loss: Optional[Loss] = None, + predict_with_sigmoid: bool = False, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, + regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model.""" super().__init__( triples_factory=triples_factory, + embedding_dim=2 * embedding_dim, # complex interaction=ComplExInteraction(), - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, - entity_representations=Embedding.from_specification( - num_embeddings=triples_factory.num_entities, + combination=ComplexLiteralCombination( embedding_dim=embedding_dim, + num_of_literals=triples_factory.numeric_literals.shape[-1], + dropout=input_dropout, ), - relation_representations=Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim, + entity_specification=EmbeddingSpecification( + initializer=xavier_normal_, ), + relation_specification=EmbeddingSpecification( + initializer=xavier_normal_, + ), + loss=loss, + predict_with_sigmoid=predict_with_sigmoid, + automatic_memory_optimization=automatic_memory_optimization, + preferred_device=preferred_device, + random_seed=random_seed, + regularizer=regularizer, ) - - # Literal - # num_ent x num_lit - numeric_literals = triples_factory.numeric_literals - self.numeric_literals = nn.Embedding.from_pretrained( - torch.tensor(numeric_literals, dtype=torch.float, device=self.device), freeze=True, - ) - # Number of columns corresponds to number of literals - self.num_of_literals = self.numeric_literals.weight.data.shape[1] - - self.real_non_lin_transf = torch.nn.Sequential( - nn.Linear(self.embedding_dim + self.num_of_literals, self.embedding_dim), - torch.nn.Tanh(), - ) - - self.img_non_lin_transf = torch.nn.Sequential( - nn.Linear(self.embedding_dim + self.num_of_literals, self.embedding_dim), - torch.nn.Tanh(), - ) - - self.inp_drop = torch.nn.Dropout(input_dropout) - - self.entity_embs_real = nn.Embedding(self.num_entities, self.embedding_dim, padding_idx=0) - self.entity_embs_img = nn.Embedding(self.num_entities, self.embedding_dim, padding_idx=0) - self.relation_embs_real = nn.Embedding(self.num_relations, self.embedding_dim, padding_idx=0) - self.relation_embs_img = nn.Embedding(self.num_relations, self.embedding_dim, padding_idx=0) - xavier_normal_(self.entity_embs_real.weight.data) - xavier_normal_(self.entity_embs_img.weight.data) - xavier_normal_(self.relation_embs_real.weight.data) - xavier_normal_(self.relation_embs_img.weight.data) - - def _apply_g_function(self, real_embs, img_embs, literals): - real = self.real_non_lin_transf(torch.cat([real_embs, literals], 1)) - img = self.img_non_lin_transf(torch.cat([img_embs, literals], 1)) - return real, img - - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: # noqa:D102 - raise NotImplementedError - - def score_t(self, doubles: torch.Tensor) -> torch.Tensor: - """Forward pass using right side (tail) prediction for training with the LCWA.""" - batch_heads, batch_relations = slice_doubles(doubles) - - heads_embedded_real = self.inp_drop(self.entity_embs_real(batch_heads)).view(-1, self.embedding_dim) - rels_embedded_real = self.inp_drop(self.relation_embs_real(batch_relations)).view( - -1, - self.embedding_dim, - ) - heads_embedded_img = self.inp_drop(self.entity_embs_img(batch_heads)).view(-1, self.embedding_dim) - relations_embedded_img = self.inp_drop(self.relation_embs_img(batch_relations)).view( - -1, - self.embedding_dim, - ) - # Literals - head_literals = self.numeric_literals(batch_heads).view(-1, self.num_of_literals) - heads_embedded_real, heads_embedded_img = self._apply_g_function( - real_embs=heads_embedded_real, - img_embs=heads_embedded_img, - literals=head_literals, - ) - - e2_multi_emb_real = self.real_non_lin_transf( - torch.cat([self.entity_embs_real.weight, self.numeric_literals.weight], 1), - ) - e2_multi_emb_img = self.img_non_lin_transf( - torch.cat([self.entity_embs_img.weight, self.numeric_literals.weight], 1), - ) - - # End literals - - heads_embedded_real = self.inp_drop(heads_embedded_real) - rels_embedded_real = self.inp_drop(rels_embedded_real) - heads_embedded_img = self.inp_drop(heads_embedded_img) - relations_embedded_img = self.inp_drop(relations_embedded_img) - - real_real_real = torch.mm(heads_embedded_real * rels_embedded_real, e2_multi_emb_real.t()) - real_img_img = torch.mm(heads_embedded_real * relations_embedded_img, e2_multi_emb_img.t()) - img_real_img = torch.mm(heads_embedded_img * heads_embedded_real, e2_multi_emb_img.t()) - img_img_real = torch.mm(heads_embedded_img * relations_embedded_img, e2_multi_emb_real.t()) - - predictions = real_real_real + real_img_img + img_real_img - img_img_real - predictions = torch.sigmoid(predictions) - - return predictions diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 6cecf0c4ac..180ccaf7b3 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -42,12 +42,14 @@ def __init__( class LiteralModel(ERModel): """Base class for models with entity literals.""" + # TODO: Move to other file? + def __init__( self, triples_factory: TriplesNumericLiteralsFactory, embedding_dim: int, interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], - dropout: float = 0.0, + combination: nn.Module, entity_specification: Optional[EmbeddingSpecification] = None, relation_specification: Optional[EmbeddingSpecification] = None, loss: Optional[Loss] = None, @@ -84,10 +86,7 @@ def __init__( specification=relation_specification, ), ) - self.trans = nn.Sequential( - nn.Linear(embedding_dim + triples_factory.numeric_literals.shape[1], embedding_dim), - nn.Dropout(dropout), - ) + self.combination = combination def forward( self, @@ -98,7 +97,7 @@ def forward( h, r, t = self._get_representations(h_indices, r_indices, t_indices) # combine entity embeddings + literals h, t = [ - self.trans(torch.cat(x, dim=-1)) + self.combination(torch.cat(x, dim=-1)) for x in (h, t) ] scores = self.interaction(h=h, r=r, t=t) @@ -134,7 +133,10 @@ def __init__( triples_factory=triples_factory, embedding_dim=embedding_dim, interaction=DistMultInteraction(), - dropout=input_dropout, + combination=nn.Sequential( + nn.Linear(embedding_dim + triples_factory.numeric_literals.shape[1], embedding_dim), + nn.Dropout(input_dropout), + ), entity_specification=EmbeddingSpecification( initializer=xavier_normal_, ), diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index f8a1537f28..05b2f4d29b 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -353,6 +353,14 @@ def split_complex( return x[..., :dim], x[..., dim:] +def combine_complex( + x_re: torch.FloatTensor, + x_im: torch.FloatTensor, +) -> Tuple[torch.FloatTensor, torch.FloatTensor]: + """Combine a complex tensor from real and imaginary part.""" + return torch.cat([x_re, x_im], dim=-1) + + def real_part( x: torch.FloatTensor, ) -> torch.FloatTensor: From 61063c5d7cf2829abed3d0f823d5e9a2bc727618 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 13:47:15 +0100 Subject: [PATCH 310/690] Add slice_size to model's score_* methods --- src/pykeen/models/base.py | 30 +++++++++++++++++++++++++----- src/pykeen/nn/modules.py | 34 ++++++++++++++++++++++++++++++++++ 2 files changed, 59 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index c903b56387..395ba207c1 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1001,6 +1001,8 @@ def forward( h_indices: Optional[torch.LongTensor], r_indices: Optional[torch.LongTensor], t_indices: Optional[torch.LongTensor], + slice_size: Optional[int] = None, + slice_dim: Optional[str] = None, ) -> torch.FloatTensor: """Forward pass. @@ -1014,13 +1016,17 @@ def forward( The relation indices. None indicates to use all. :param t_indices: The tail indices. None indicates to use all. + :param slice_size: + The slice size. + :param slice_dim: + The dimension along which to slice. From {"h", "r", "t"} :return: shape: (batch_size, num_heads, num_relations, num_tails) The score for each triple. """ raise NotImplementedError - def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: + def score_hrt(self, hrt_batch: torch.LongTensor, slice_size: Optional[int] = None) -> torch.FloatTensor: """Forward pass. This method takes head, relation and tail of each triple and calculates the corresponding score. @@ -1035,9 +1041,11 @@ def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2], + slice_size=slice_size, + slice_dim="h", ).view(hrt_batch.shape[0], 1) - def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: + def score_t(self, hr_batch: torch.LongTensor, slice_size: Optional[int] = None) -> torch.FloatTensor: """Forward pass using right side (tail) prediction. This method calculates the score for all possible tails for each (head, relation) pair. @@ -1052,9 +1060,11 @@ def score_t(self, hr_batch: torch.LongTensor) -> torch.FloatTensor: h_indices=hr_batch[:, 0], r_indices=hr_batch[:, 1], t_indices=None, + slice_size=slice_size, + slice_dim="h", ).view(hr_batch.shape[0], self.num_entities) - def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: + def score_h(self, rt_batch: torch.LongTensor, slice_size: Optional[int] = None) -> torch.FloatTensor: """Forward pass using left side (head) prediction. This method calculates the score for all possible heads for each (relation, tail) pair. @@ -1069,9 +1079,11 @@ def score_h(self, rt_batch: torch.LongTensor) -> torch.FloatTensor: h_indices=None, r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1], + slice_size=slice_size, + slice_dim="h", ).view(rt_batch.shape[0], self.num_entities) - def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: + def score_r(self, ht_batch: torch.LongTensor, slice_size: Optional[int] = None) -> torch.FloatTensor: """Forward pass using middle (relation) prediction. This method calculates the score for all possible relations for each (head, tail) pair. @@ -1086,6 +1098,8 @@ def score_r(self, ht_batch: torch.LongTensor) -> torch.FloatTensor: h_indices=ht_batch[:, 0], r_indices=None, t_indices=ht_batch[:, 1], + slice_size=slice_size, + slice_dim="h", ).view(ht_batch.shape[0], self.num_relations) def get_grad_params(self) -> Iterable[nn.Parameter]: @@ -1184,6 +1198,8 @@ def forward( h_indices: Optional[torch.LongTensor], r_indices: Optional[torch.LongTensor], t_indices: Optional[torch.LongTensor], + slice_size: Optional[int] = None, + slice_dim: Optional[str] = None, ) -> torch.FloatTensor: """Forward pass. @@ -1197,12 +1213,16 @@ def forward( The relation indices. None indicates to use all. :param t_indices: The tail indices. None indicates to use all. + :param slice_size: + The slice size. + :param slice_dim: + The dimension along which to slice. From {"h", "r", "t"} :return: shape: (batch_size, num_heads, num_relations, num_tails) The score for each triple. """ h, r, t = self._get_representations(h_indices, r_indices, t_indices) - scores = self.interaction(h=h, r=r, t=t) + scores = self.interaction.score(h=h, r=r, t=t, slice_size=slice_size, slice_dim=slice_dim) return self._repeat_scores_if_necessary(scores, h_indices, r_indices, t_indices) def _repeat_scores_if_necessary( diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index c28fd4ce99..ffac281f1c 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -173,6 +173,40 @@ def _check_shapes( raise_on_errors=raise_on_errors, ) + def score( + self, + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + slice_size: Optional[int] = None, + slice_dim: Optional[str] = None, + ) -> torch.FloatTensor: + """ + Compute broadcasted triple scores with optional slicing. + + .. note :: + At most one of the slice sizes may be not None. + + :param h: shape: (batch_size, num_heads, ``*``) + The head representations. + :param r: shape: (batch_size, num_relations, ``*``) + The relation representations. + :param t: shape: (batch_size, num_tails, ``*``) + The tail representations. + :param slice_size: + The slice size. + :param slice_dim: + The dimension along which to slice. From {"h", "r", "t"} + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ + kwargs = { + f"{d}_prefix": "b" if (slice_size is None or slice_dim != d) else "n" + for d in "hrt" + } + return self._score(h=h, r=r, t=t, **kwargs, slice_size=slice_size) + def _score( self, h: HeadRepresentation, From e5ad437622e006bd55fda69c589f9baac7030be9 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 13:49:30 +0100 Subject: [PATCH 311/690] Update subclasses' forward signature --- src/pykeen/models/base.py | 16 +++++++++++++++- src/pykeen/models/multimodal/distmult_literal.py | 4 +++- src/pykeen/models/unimodal/conv_e.py | 4 +++- src/pykeen/models/unimodal/simple.py | 4 ++++ 4 files changed, 25 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 395ba207c1..5d46b0a0b9 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1033,6 +1033,8 @@ def score_hrt(self, hrt_batch: torch.LongTensor, slice_size: Optional[int] = Non :param hrt_batch: shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples. + :param slice_size: + The slice size. :return: shape: (batch_size, 1), dtype: float The score for each triple. @@ -1052,6 +1054,8 @@ def score_t(self, hr_batch: torch.LongTensor, slice_size: Optional[int] = None) :param hr_batch: shape: (batch_size, 2), dtype: long The indices of (head, relation) pairs. + :param slice_size: + The slice size. :return: shape: (batch_size, num_entities), dtype: float For each h-r pair, the scores for all possible tails. @@ -1071,6 +1075,8 @@ def score_h(self, rt_batch: torch.LongTensor, slice_size: Optional[int] = None) :param rt_batch: shape: (batch_size, 2), dtype: long The indices of (relation, tail) pairs. + :param slice_size: + The slice size. :return: shape: (batch_size, num_entities), dtype: float For each r-t pair, the scores for all possible heads. @@ -1090,6 +1096,8 @@ def score_r(self, ht_batch: torch.LongTensor, slice_size: Optional[int] = None) :param ht_batch: shape: (batch_size, 2), dtype: long The indices of (head, tail) pairs. + :param slice_size: + The slice size. :return: shape: (batch_size, num_relations), dtype: float For each h-t pair, the scores for all possible relations. @@ -1559,12 +1567,16 @@ def forward( h_indices: Optional[torch.LongTensor], r_indices: Optional[torch.LongTensor], t_indices: Optional[torch.LongTensor], + slice_size: Optional[int] = None, + slice_dim: Optional[str] = None, ) -> torch.FloatTensor: # noqa: D102 return 0.5 * sum( - self.interaction( + self.interaction.score( h_source.get_in_canonical_shape(indices=h_indices), r_source.get_in_canonical_shape(indices=r_indices), t_source.get_in_canonical_shape(indices=t_indices), + slice_size=slice_size, + slice_dim=slice_dim, ) for h_source, r_source, t_source in ( (self.entity_representations[0], self.relation_representations[0], self.entity_representations[1]), @@ -1589,6 +1601,8 @@ def forward( h_indices: Optional[torch.LongTensor], r_indices: Optional[torch.LongTensor], t_indices: Optional[torch.LongTensor], + slice_size: Optional[int] = None, + slice_dim: Optional[str] = None, ) -> torch.FloatTensor: # noqa: D102 # (batch_size, num_heads, num_relations, num_tails) scores = torch.zeros(1, 1, 1, 1, requires_grad=True) # for requires_grad diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 180ccaf7b3..0337e53612 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -93,6 +93,8 @@ def forward( h_indices: Optional[torch.LongTensor], r_indices: Optional[torch.LongTensor], t_indices: Optional[torch.LongTensor], + slice_size: Optional[int] = None, + slice_dim: Optional[str] = None, ) -> torch.FloatTensor: # noqa: D102 h, r, t = self._get_representations(h_indices, r_indices, t_indices) # combine entity embeddings + literals @@ -100,7 +102,7 @@ def forward( self.combination(torch.cat(x, dim=-1)) for x in (h, t) ] - scores = self.interaction(h=h, r=r, t=t) + scores = self.interaction.score(h=h, r=r, t=t, slice_size=slice_size, slice_dim=slice_dim) return self._repeat_scores_if_necessary(scores, h_indices, r_indices, t_indices) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index d517345fde..de66334368 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -186,10 +186,12 @@ def forward( h_indices: Optional[torch.LongTensor], r_indices: Optional[torch.LongTensor], t_indices: Optional[torch.LongTensor], + slice_size: Optional[int] = None, + slice_dim: Optional[str] = None, ) -> torch.FloatTensor: # noqa: D102 h = self.entity_representations[0].get_in_canonical_shape(indices=h_indices) r = self.relation_representations[0].get_in_canonical_shape(indices=r_indices) t = self.entity_representations[0].get_in_canonical_shape(indices=t_indices) t_bias = self.entity_representations[1].get_in_canonical_shape(indices=t_indices) self.regularize_if_necessary(h, r, t) - return self.interaction(h=h, r=r, t=(t, t_bias)) + return self.interaction.score(h=h, r=r, t=(t, t_bias), slice_size=slice_size, slice_dim=slice_dim) diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index bbb7bcc8b2..1b05f2366f 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -92,11 +92,15 @@ def forward( h_indices: Optional[torch.LongTensor], r_indices: Optional[torch.LongTensor], t_indices: Optional[torch.LongTensor], + slice_size: Optional[int] = None, + slice_dim: Optional[str] = None, ) -> torch.FloatTensor: # noqa: D102 scores = super().forward( h_indices=h_indices, r_indices=r_indices, t_indices=t_indices, + slice_size=slice_size, + slice_dim=slice_dim, ) # Note: In the code in their repository, the score is clamped to [-20, 20]. From 134157c14c915778aefbf19243a6944ffe23d9e5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 13:50:56 +0100 Subject: [PATCH 312/690] Fix slice size usage --- src/pykeen/models/base.py | 10 +++------- src/pykeen/nn/modules.py | 6 ++++++ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 5d46b0a0b9..a38dd225b0 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1026,15 +1026,13 @@ def forward( """ raise NotImplementedError - def score_hrt(self, hrt_batch: torch.LongTensor, slice_size: Optional[int] = None) -> torch.FloatTensor: + def score_hrt(self, hrt_batch: torch.LongTensor) -> torch.FloatTensor: """Forward pass. This method takes head, relation and tail of each triple and calculates the corresponding score. :param hrt_batch: shape: (batch_size, 3), dtype: long The indices of (head, relation, tail) triples. - :param slice_size: - The slice size. :return: shape: (batch_size, 1), dtype: float The score for each triple. @@ -1043,8 +1041,6 @@ def score_hrt(self, hrt_batch: torch.LongTensor, slice_size: Optional[int] = Non h_indices=hrt_batch[:, 0], r_indices=hrt_batch[:, 1], t_indices=hrt_batch[:, 2], - slice_size=slice_size, - slice_dim="h", ).view(hrt_batch.shape[0], 1) def score_t(self, hr_batch: torch.LongTensor, slice_size: Optional[int] = None) -> torch.FloatTensor: @@ -1086,7 +1082,7 @@ def score_h(self, rt_batch: torch.LongTensor, slice_size: Optional[int] = None) r_indices=rt_batch[:, 0], t_indices=rt_batch[:, 1], slice_size=slice_size, - slice_dim="h", + slice_dim="r", ).view(rt_batch.shape[0], self.num_entities) def score_r(self, ht_batch: torch.LongTensor, slice_size: Optional[int] = None) -> torch.FloatTensor: @@ -1107,7 +1103,7 @@ def score_r(self, ht_batch: torch.LongTensor, slice_size: Optional[int] = None) r_indices=None, t_indices=ht_batch[:, 1], slice_size=slice_size, - slice_dim="h", + slice_dim="t", ).view(ht_batch.shape[0], self.num_relations) def get_grad_params(self) -> Iterable[nn.Parameter]: diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index ffac281f1c..3694352e39 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -308,6 +308,8 @@ def score_h( The relation representations. :param t: shape: (batch_size, d_e) The tail representations. + :param slice_size: + The slice size. :return: shape: (batch_size, num_entities) The scores. @@ -330,6 +332,8 @@ def score_r( The relation representations. :param t: shape: (batch_size, d_e) The tail representations. + :param slice_size: + The slice size. :return: shape: (batch_size, num_entities) The scores. @@ -352,6 +356,8 @@ def score_t( The relation representations. :param all_entities: shape: (num_entities, d_e) The tail representations. + :param slice_size: + The slice size. :return: shape: (batch_size, num_entities) The scores. From 46a7c9fe0e90b21cf6f514d34c873fd345a03f0a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 13:51:58 +0100 Subject: [PATCH 313/690] Add two notes for Model.forward --- src/pykeen/models/base.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index a38dd225b0..fca3b35235 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1008,7 +1008,11 @@ def forward( This method takes head, relation and tail indices and calculates the corresponding score. - All indices which are not None, have to be either 1-element or have the same shape, which is the batch size. + .. note :: + All indices which are not None, have to be either 1-element or have the same shape, which is the batch size. + + .. note :: + If slicing is requested, the corresponding indices have to be None. :param h_indices: The head indices. None indicates to use all. @@ -1019,7 +1023,7 @@ def forward( :param slice_size: The slice size. :param slice_dim: - The dimension along which to slice. From {"h", "r", "t"} + The dimension along which to slice. From {"h", "r", "t"}. :return: shape: (batch_size, num_heads, num_relations, num_tails) The score for each triple. From ae571b383e7d431646a820976352c160bb87abd4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 13:54:24 +0100 Subject: [PATCH 314/690] Fix type annotation --- src/pykeen/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 05b2f4d29b..cee50bafcb 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -356,7 +356,7 @@ def split_complex( def combine_complex( x_re: torch.FloatTensor, x_im: torch.FloatTensor, -) -> Tuple[torch.FloatTensor, torch.FloatTensor]: +) -> torch.FloatTensor: """Combine a complex tensor from real and imaginary part.""" return torch.cat([x_re, x_im], dim=-1) From 6b729ca4d8c55ea3723b46865a16dd3d37b696f4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 13:55:29 +0100 Subject: [PATCH 315/690] Add test for complex tensor utilities --- tests/test_utils.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index 795408edeb..d2a4a9721f 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -11,10 +11,10 @@ from pykeen.nn import Embedding from pykeen.utils import ( clamp_norm, - compact_mapping, + combine_complex, compact_mapping, flatten_dictionary, get_until_first_blank, - l2_regularization, + imag_part, l2_regularization, real_part, split_complex, ) @@ -217,3 +217,15 @@ def test_clamp_norm(): norm = x.norm(p=p, dim=dim) mask = torch.stack([(norm < max_norm)] * x.shape[dim], dim=dim) assert (x_c[mask] == x[mask]).all() + + +def test_complex_utils(): + """Test complex tensor utilities.""" + re = torch.rand(20, 10) + im = torch.rand(20, 10) + x = combine_complex(x_re=re, x_im=im) + assert (real_part(x) == re).all() + assert (imag_part(x) == im).all() + re2, im2 = split_complex(x) + assert (re2 == re).all() + assert (im2 == im).all() From bb288ba5f93d06b90baa1b5a28ce470ee7c4a6c2 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 14:49:32 +0100 Subject: [PATCH 316/690] Pass flake8/mypy --- src/pykeen/models/base.py | 23 +++++++------ .../models/multimodal/distmult_literal.py | 5 ++- src/pykeen/nn/emb.py | 4 ++- src/pykeen/nn/modules.py | 32 ++++++++++++------- 4 files changed, 38 insertions(+), 26 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index fca3b35235..bad21c548b 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -10,7 +10,7 @@ from collections import defaultdict from operator import itemgetter from typing import ( - Any, ClassVar, Collection, Dict, Generic, Iterable, List, Mapping, Optional, Sequence, Set, Tuple, + Any, ClassVar, Collection, Dict, Generic, Iterable, List, Mapping, Optional, Sequence, Set, TYPE_CHECKING, Tuple, Type, Union, ) @@ -25,11 +25,12 @@ from ..nn.modules import Interaction from ..regularizers import NoRegularizer, Regularizer from ..triples import TriplesFactory -from ..typing import ( - DeviceHint, HeadRepresentation, MappedTriples, RelationRepresentation, TailRepresentation, -) +from ..typing import DeviceHint, HeadRepresentation, MappedTriples, RelationRepresentation, TailRepresentation from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed +if TYPE_CHECKING: + from ..typing import Representation # noqa + __all__ = [ 'Model', 'ERModel', @@ -1271,7 +1272,7 @@ def _get_representations( ) -> Tuple[ Union[torch.FloatTensor, Sequence[torch.FloatTensor]], Union[torch.FloatTensor, Sequence[torch.FloatTensor]], - Union[torch.FloatTensor, Sequence[torch.FloatTensor]] + Union[torch.FloatTensor, Sequence[torch.FloatTensor]], ]: h, r, t = [ [ @@ -1607,13 +1608,11 @@ def forward( # (batch_size, num_heads, num_relations, num_tails) scores = torch.zeros(1, 1, 1, 1, requires_grad=True) # for requires_grad # reproducible scores - for i, (ind, num) in enumerate( - ( - (h_indices, self.num_entities), - (r_indices, self.num_relations), - (t_indices, self.num_entities), - ) - ): + for i, (ind, num) in enumerate(( + (h_indices, self.num_entities), + (r_indices, self.num_relations), + (t_indices, self.num_entities), + )): shape = [1, 1, 1, 1] if ind is None: shape[i + 1] = num diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 0337e53612..368268d432 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -2,7 +2,7 @@ """Implementation of the DistMultLiteral model.""" -from typing import Optional +from typing import Optional, TYPE_CHECKING import torch import torch.nn as nn @@ -17,6 +17,9 @@ from ...triples import TriplesNumericLiteralsFactory from ...typing import DeviceHint, HeadRepresentation, RelationRepresentation, TailRepresentation +if TYPE_CHECKING: + from ...typing import Representation # noqa + __all__ = [ 'DistMultLiteral', ] diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index fbb0c412fa..1706eedf9d 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -5,6 +5,7 @@ import dataclasses import functools import logging +from abc import ABC, abstractmethod from typing import Any, Mapping, Optional, Sequence, Union import numpy @@ -23,9 +24,10 @@ logger = logging.getLogger(__name__) -class RepresentationModule(nn.Module): +class RepresentationModule(nn.Module, ABC): """A base class for obtaining representations for entities/relations.""" + @abstractmethod def forward(self, indices: Optional[torch.LongTensor] = None) -> torch.FloatTensor: """Get representations for indices. diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 3694352e39..73550329e2 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -52,6 +52,10 @@ def _unpack_singletons(*xs: Tuple) -> Sequence[Tuple]: return [x[0] if len(x) == 1 else x for x in xs] +def _get_prefix(slice_size, slice_dim, d) -> str: + return "b" if (slice_size is None or slice_dim != d) else "n" + + class Interaction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation], ABC): """Base class for interaction functions.""" @@ -131,9 +135,9 @@ def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: def _check_shapes( self, - h: HeadRepresentation = tuple(), - r: RelationRepresentation = tuple(), - t: TailRepresentation = tuple(), + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, h_prefix: str = "b", r_prefix: str = "b", t_prefix: str = "b", @@ -201,11 +205,15 @@ def score( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - kwargs = { - f"{d}_prefix": "b" if (slice_size is None or slice_dim != d) else "n" - for d in "hrt" - } - return self._score(h=h, r=r, t=t, **kwargs, slice_size=slice_size) + return self._score( + h=h, + r=r, + t=t, + h_prefix=_get_prefix(slice_size=slice_size, slice_dim=slice_dim, d='h'), + r_prefix=_get_prefix(slice_size=slice_size, slice_dim=slice_dim, d='r'), + t_prefix=_get_prefix(slice_size=slice_size, slice_dim=slice_dim, d='t'), + slice_size=slice_size, + ) def _score( self, @@ -378,7 +386,7 @@ class StatelessInteraction(Interaction[HeadRepresentation, RelationRepresentatio def __init__(self, f: Callable[..., torch.FloatTensor]): """Instantiate the stateless interaction module. - + :param f: The interaction function, like ones from :mod:`pykeen.nn.functional`. """ super().__init__() @@ -1012,9 +1020,9 @@ def __init__( def forward( self, - h: HeadRepresentation, - r: RelationRepresentation, - t: TailRepresentation, + h: Tuple[torch.FloatTensor, torch.FloatTensor], + r: Tuple[torch.FloatTensor, torch.FloatTensor], + t: Tuple[torch.FloatTensor, torch.FloatTensor], ) -> torch.FloatTensor: # noqa:D102 h_mean, h_var = h r_mean, r_var = r From 4088fcd764430672ed9a83bd8fb397879481bcf4 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 15:24:04 +0100 Subject: [PATCH 317/690] Remove unnecessary default arguments and function cleanup --- src/pykeen/nn/modules.py | 47 +++++++++++++++++++++++++--------------- 1 file changed, 29 insertions(+), 18 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 73550329e2..fbe8273a5c 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -15,8 +15,6 @@ from ..typing import HeadRepresentation, RelationRepresentation, Representation, TailRepresentation from ..utils import check_shapes -logger = logging.getLogger(__name__) - __all__ = [ # Base Classes 'Interaction', @@ -43,17 +41,29 @@ 'UnstructuredModelInteraction', ] +logger = logging.getLogger(__name__) + -def _ensure_tuple(*x: Union[Representation, Sequence[Representation]]) -> Tuple[Sequence[Representation], ...]: - return tuple(xx if isinstance(xx, Sequence) else (xx,) for xx in x) +def _upgrade_to_sequence(x: Union[FloatTensor, Sequence[FloatTensor]]) -> Sequence[FloatTensor]: + return x if isinstance(x, Sequence) else (x,) + + +def _ensure_tuple(*x: Union[Representation, Sequence[Representation]]) -> Sequence[Sequence[Representation]]: + return tuple(_upgrade_to_sequence(xx) for xx in x) def _unpack_singletons(*xs: Tuple) -> Sequence[Tuple]: - return [x[0] if len(x) == 1 else x for x in xs] + return [ + x[0] if len(x) == 1 else x + for x in xs + ] def _get_prefix(slice_size, slice_dim, d) -> str: - return "b" if (slice_size is None or slice_dim != d) else "n" + if slice_size is None or slice_dim != d: + return 'b' + else: + return 'n' class Interaction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation], ABC): @@ -246,6 +256,7 @@ def _score( scores = self(h=h, r=r, t=t) else: assert slice_dim is not None + # TODO externalize logic into function scores = [] if slice_dim == "h": cat_dim = self.HEAD_DIM @@ -281,9 +292,9 @@ def _score( def score_hrt( self, - h: HeadRepresentation = tuple(), - r: RelationRepresentation = tuple(), - t: TailRepresentation = tuple(), + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, ) -> torch.FloatTensor: """ Score a batch of triples.. @@ -302,9 +313,9 @@ def score_hrt( def score_h( self, - all_entities: HeadRepresentation = tuple(), - r: RelationRepresentation = tuple(), - t: TailRepresentation = tuple(), + all_entities: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, slice_size: Optional[int] = None, ) -> torch.FloatTensor: """ @@ -326,9 +337,9 @@ def score_h( def score_r( self, - h: HeadRepresentation = tuple(), - all_relations: RelationRepresentation = tuple(), - t: TailRepresentation = tuple(), + h: HeadRepresentation, + all_relations: RelationRepresentation, + t: TailRepresentation, slice_size: Optional[int] = None, ) -> torch.FloatTensor: """ @@ -350,9 +361,9 @@ def score_r( def score_t( self, - h: HeadRepresentation = tuple(), - r: RelationRepresentation = tuple(), - all_entities: TailRepresentation = tuple(), + h: HeadRepresentation, + r: RelationRepresentation, + all_entities: TailRepresentation, slice_size: Optional[int] = None, ) -> torch.FloatTensor: """ From 745f61fe05c65244ebabf4807de16b23c87df4e1 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 15:39:33 +0100 Subject: [PATCH 318/690] Prepare for more typing --- src/pykeen/nn/modules.py | 11 +++++++++-- src/pykeen/typing.py | 2 ++ 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index fbe8273a5c..cf9f31bacf 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -243,7 +243,11 @@ def _score( if prefix == "n" ] slice_dim: Optional[str] = slice_dims[0] if len(slice_dims) == 1 else None - h, r, t = _ensure_tuple(h, r, t) + + # FIXME typing does not work well for this + h = _upgrade_to_sequence(h) + r = _upgrade_to_sequence(r) + t = _upgrade_to_sequence(t) assert self._check_shapes(h=h, r=r, t=t, h_prefix=h_prefix, r_prefix=r_prefix, t_prefix=t_prefix) # prepare input to generic score function: bh*, br*, bt* @@ -410,7 +414,10 @@ def forward( t: TailRepresentation, ) -> torch.FloatTensor: # noqa: D102 # normalization - h, r, t = _ensure_tuple(h, r, t) # TODO provide example of non-simple case + h = _upgrade_to_sequence(h) + r = _upgrade_to_sequence(r) + t = _upgrade_to_sequence(t) + # TODO provide example of non-simple case return self.f(*h, *r, *t) diff --git a/src/pykeen/typing.py b/src/pykeen/typing.py index eef6133202..80485d1e7e 100644 --- a/src/pykeen/typing.py +++ b/src/pykeen/typing.py @@ -36,6 +36,8 @@ DeviceHint = Union[None, str, torch.device] Representation = torch.FloatTensor +# TODO upgrade to use bound=... +# HeadRepresentation = TypeVar("HeadRepresentation", bound=Union[Representation, Sequence[Representation]]) HeadRepresentation = TypeVar("HeadRepresentation", Representation, Sequence[Representation]) # type: ignore RelationRepresentation = TypeVar("RelationRepresentation", Representation, Sequence[Representation]) # type: ignore TailRepresentation = TypeVar("TailRepresentation", Representation, Sequence[Representation]) # type: ignore From cab70b2fd4c3c92baaf759ca0978146dab6130ef Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 15:39:52 +0100 Subject: [PATCH 319/690] Clean up batch scoring --- src/pykeen/nn/modules.py | 42 +++++++++++++++++----------------------- 1 file changed, 18 insertions(+), 24 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index cf9f31bacf..126dfd23e1 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -66,6 +66,13 @@ def _get_prefix(slice_size, slice_dim, d) -> str: return 'n' +def _get_batches(z, slice_size): + for batch in zip(*(hh.split(slice_size, dim=1) for hh in _ensure_tuple(z)[0])): + if len(batch) == 1: + batch = batch[0] + yield batch + + class Interaction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation], ABC): """Base class for interaction functions.""" @@ -258,31 +265,18 @@ def _score( # get scores if slice_size is None: scores = self(h=h, r=r, t=t) + elif slice_dim == "h": + batch_scores = [self(h=h_batch, r=r, t=t) for h_batch in _get_batches(h, slice_size)] + scores = torch.cat(batch_scores, dim=self.HEAD_DIM) + elif slice_dim == "r": + batch_scores = [self(h=h, r=r_batch, t=t) for r_batch in _get_batches(r, slice_size)] + scores = torch.cat(batch_scores, dim=self.RELATION_DIM) + elif slice_dim == "t": + batch_scores = [self(h=h, r=r, t=t_batch) for t_batch in _get_batches(t, slice_size)] + scores = torch.cat(batch_scores, dim=self.TAIL_DIM) else: - assert slice_dim is not None - # TODO externalize logic into function - scores = [] - if slice_dim == "h": - cat_dim = self.HEAD_DIM - for h_batch in zip(*(hh.split(slice_size, dim=1) for hh in _ensure_tuple(h)[0])): - if len(h_batch) == 1: - h_batch = h_batch[0] - scores.append(self(h=h_batch, r=r, t=t)) - elif slice_dim == "r": - cat_dim = self.RELATION_DIM - for r_batch in zip(*(rr.split(slice_size, dim=1) for rr in _ensure_tuple(r)[0])): - if len(r_batch) == 1: - r_batch = r_batch[0] - scores.append(self(h=h, r=r_batch, t=t)) - elif slice_dim == "t": - cat_dim = self.TAIL_DIM - for t_batch in zip(*(tt.split(slice_size, dim=1) for tt in _ensure_tuple(t)[0])): - if len(t_batch) == 1: - t_batch = t_batch[0] - scores.append(self(h=h, r=r, t=t_batch)) - else: - raise ValueError(slice_dim) - scores = torch.cat(scores, dim=cat_dim) + raise ValueError(f'Invalid slice_dim: {slice_dim}') + remove_dims = [ dim for dim, prefix in zip( From 422bad7eb4413959debe549c39bb61013e95f76b Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 15:56:14 +0100 Subject: [PATCH 320/690] Reorganize new models --- src/pykeen/models/__init__.py | 6 +- src/pykeen/models/base.py | 81 +++++++++++------ .../models/multimodal/complex_literal.py | 2 +- .../models/multimodal/distmult_literal.py | 90 +------------------ src/pykeen/nn/emb.py | 18 ++++ src/pykeen/testing/__init__.py | 3 + src/pykeen/testing/mocks.py | 50 +++++++++++ tests/test_early_stopping.py | 3 +- tests/test_evaluators.py | 3 +- tests/test_model_mode.py | 2 +- tests/test_models.py | 12 +-- 11 files changed, 142 insertions(+), 128 deletions(-) create mode 100644 src/pykeen/testing/__init__.py create mode 100644 src/pykeen/testing/mocks.py diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index 0250174d94..d433331506 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -9,8 +9,8 @@ from typing import Mapping, Set, Type, Union from .base import ( # noqa:F401 - DoubleRelationEmbeddingModel, ERModel, Model, MultimodalModel, - SingleVectorEmbeddingModel, TwoSideEmbeddingModel, TwoVectorEmbeddingModel, + DoubleRelationEmbeddingModel, ERModel, LiteralModel, Model, SingleVectorEmbeddingModel, + TwoSideEmbeddingModel, TwoVectorEmbeddingModel, ) from .multimodal import ComplExLiteral, DistMultLiteral from .unimodal import ( @@ -72,7 +72,7 @@ DoubleRelationEmbeddingModel, TwoSideEmbeddingModel, TwoVectorEmbeddingModel, - MultimodalModel, + LiteralModel, } diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index bad21c548b..5dd928112c 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -21,10 +21,10 @@ from ..losses import Loss, MarginRankingLoss, NSSALoss from ..nn import Embedding, RepresentationModule -from ..nn.emb import EmbeddingSpecification +from ..nn.emb import EmbeddingSpecification, LiteralRepresentations from ..nn.modules import Interaction from ..regularizers import NoRegularizer, Regularizer -from ..triples import TriplesFactory +from ..triples import TriplesFactory, TriplesNumericLiteralsFactory from ..typing import DeviceHint, HeadRepresentation, MappedTriples, RelationRepresentation, TailRepresentation from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed @@ -38,7 +38,7 @@ 'DoubleRelationEmbeddingModel', 'TwoVectorEmbeddingModel', 'TwoSideEmbeddingModel', - 'MultimodalModel', + 'LiteralModel', ] logger = logging.getLogger(__name__) @@ -1290,10 +1290,6 @@ def _get_representations( return h, r, t -class MultimodalModel(ERModel, autoreset=False): - """A multimodal KGE model.""" - - class SingleVectorEmbeddingModel(ERModel, autoreset=False): """A KGEM that stores one :class:`pykeen.nn.Embedding` for each entities and relations.""" @@ -1586,16 +1582,54 @@ def forward( ) -class MockModel(Model): - """A mock model returning fake scores.""" +class LiteralModel(ERModel, autoreset=False): + """Base class for models with entity literals.""" - # TODO: Where to put this? + # TODO: Move to other file? - def __init__(self, triples_factory: TriplesFactory, automatic_memory_optimization: bool = True): + def __init__( + self, + triples_factory: TriplesNumericLiteralsFactory, + embedding_dim: int, + interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], + combination: nn.Module, + entity_specification: Optional[EmbeddingSpecification] = None, + relation_specification: Optional[EmbeddingSpecification] = None, + loss: Optional[Loss] = None, + predict_with_sigmoid: bool = False, + automatic_memory_optimization: Optional[bool] = None, + preferred_device: DeviceHint = None, + random_seed: Optional[int] = None, + regularizer: Optional[Regularizer] = None, + ): super().__init__( triples_factory=triples_factory, + interaction=interaction, + loss=loss, + predict_with_sigmoid=predict_with_sigmoid, automatic_memory_optimization=automatic_memory_optimization, + preferred_device=preferred_device, + random_seed=random_seed, + regularizer=regularizer, + entity_representations=[ + # entity embeddings + Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + specification=entity_specification, + ), + # Entity literals + LiteralRepresentations( + numeric_literals=torch.as_tensor(triples_factory.numeric_literals, dtype=torch.float32), + ), + ], + relation_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=embedding_dim, + specification=relation_specification, + ), ) + self.combination = combination def forward( self, @@ -1605,20 +1639,11 @@ def forward( slice_size: Optional[int] = None, slice_dim: Optional[str] = None, ) -> torch.FloatTensor: # noqa: D102 - # (batch_size, num_heads, num_relations, num_tails) - scores = torch.zeros(1, 1, 1, 1, requires_grad=True) # for requires_grad - # reproducible scores - for i, (ind, num) in enumerate(( - (h_indices, self.num_entities), - (r_indices, self.num_relations), - (t_indices, self.num_entities), - )): - shape = [1, 1, 1, 1] - if ind is None: - shape[i + 1] = num - delta = torch.arange(num) - else: - shape[0] = len(ind) - delta = ind - scores = scores + delta.float().view(*shape) - return scores + h, r, t = self._get_representations(h_indices, r_indices, t_indices) + # combine entity embeddings + literals + h, t = [ + self.combination(torch.cat(x, dim=-1)) + for x in (h, t) + ] + scores = self.interaction.score(h=h, r=r, t=t, slice_size=slice_size, slice_dim=slice_dim) + return self._repeat_scores_if_necessary(scores, h_indices, r_indices, t_indices) diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index e5353e0c37..31c5dcdc93 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -8,7 +8,7 @@ import torch.nn as nn from torch.nn.init import xavier_normal_ -from .distmult_literal import LiteralModel +from ..base import LiteralModel from ...losses import BCEWithLogitsLoss, Loss from ...nn.emb import EmbeddingSpecification from ...nn.modules import ComplExInteraction diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 368268d432..85a67977b1 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -4,18 +4,16 @@ from typing import Optional, TYPE_CHECKING -import torch import torch.nn as nn from torch.nn.init import xavier_normal_ -from .. import ERModel +from ..base import LiteralModel from ...losses import Loss -from ...nn import Embedding, Interaction from ...nn.emb import EmbeddingSpecification from ...nn.modules import DistMultInteraction from ...regularizers import Regularizer from ...triples import TriplesNumericLiteralsFactory -from ...typing import DeviceHint, HeadRepresentation, RelationRepresentation, TailRepresentation +from ...typing import DeviceHint if TYPE_CHECKING: from ...typing import Representation # noqa @@ -25,90 +23,6 @@ ] -class LiteralRepresentations(Embedding): - """Literal representations.""" - - def __init__( - self, - numeric_literals: torch.FloatTensor, - ): - num_embeddings, embedding_dim = numeric_literals.shape - super().__init__( - num_embeddings=num_embeddings, - embedding_dim=embedding_dim, - initializer=lambda x: numeric_literals, # initialize with the literals - ) - # freeze - self._embeddings.requires_grad_(False) - - -class LiteralModel(ERModel): - """Base class for models with entity literals.""" - - # TODO: Move to other file? - - def __init__( - self, - triples_factory: TriplesNumericLiteralsFactory, - embedding_dim: int, - interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], - combination: nn.Module, - entity_specification: Optional[EmbeddingSpecification] = None, - relation_specification: Optional[EmbeddingSpecification] = None, - loss: Optional[Loss] = None, - predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, - preferred_device: DeviceHint = None, - random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, - ): - super().__init__( - triples_factory=triples_factory, - interaction=interaction, - loss=loss, - predict_with_sigmoid=predict_with_sigmoid, - automatic_memory_optimization=automatic_memory_optimization, - preferred_device=preferred_device, - random_seed=random_seed, - regularizer=regularizer, - entity_representations=[ - # entity embeddings - Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, - specification=entity_specification, - ), - # Entity literals - LiteralRepresentations( - numeric_literals=torch.as_tensor(triples_factory.numeric_literals, dtype=torch.float32), - ), - ], - relation_representations=Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim, - specification=relation_specification, - ), - ) - self.combination = combination - - def forward( - self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], - slice_size: Optional[int] = None, - slice_dim: Optional[str] = None, - ) -> torch.FloatTensor: # noqa: D102 - h, r, t = self._get_representations(h_indices, r_indices, t_indices) - # combine entity embeddings + literals - h, t = [ - self.combination(torch.cat(x, dim=-1)) - for x in (h, t) - ] - scores = self.interaction.score(h=h, r=r, t=t, slice_size=slice_size, slice_dim=slice_dim) - return self._repeat_scores_if_necessary(scores, h_indices, r_indices, t_indices) - - # TODO: Check entire build of the model # TODO: There are no tests class DistMultLiteral(LiteralModel): diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index 1706eedf9d..8e3ecffc63 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -19,6 +19,7 @@ 'RepresentationModule', 'Embedding', 'EmbeddingSpecification', + 'LiteralRepresentations', ] logger = logging.getLogger(__name__) @@ -267,3 +268,20 @@ def get_in_canonical_shape( indices=indices, reshape_dim=reshape_dim, ) + + +class LiteralRepresentations(Embedding): + """Literal representations.""" + + def __init__( + self, + numeric_literals: torch.FloatTensor, + ): + num_embeddings, embedding_dim = numeric_literals.shape + super().__init__( + num_embeddings=num_embeddings, + embedding_dim=embedding_dim, + initializer=lambda x: numeric_literals, # initialize with the literals + ) + # freeze + self._embeddings.requires_grad_(False) diff --git a/src/pykeen/testing/__init__.py b/src/pykeen/testing/__init__.py new file mode 100644 index 0000000000..0ac80fed9a --- /dev/null +++ b/src/pykeen/testing/__init__.py @@ -0,0 +1,3 @@ +# -*- coding: utf-8 -*- + +"""Code for testing of KGEMs and PyKEEN.""" diff --git a/src/pykeen/testing/mocks.py b/src/pykeen/testing/mocks.py new file mode 100644 index 0000000000..f476f1a2c4 --- /dev/null +++ b/src/pykeen/testing/mocks.py @@ -0,0 +1,50 @@ +# -*- coding: utf-8 -*- + +"""Mocks for testing PyKEEN.""" + +from typing import Optional + +import torch + +from pykeen.models.base import Model +from pykeen.triples import TriplesFactory + +__all__ = [ + 'MockModel', +] + + +class MockModel(Model): + """A mock model returning fake scores.""" + + def __init__(self, triples_factory: TriplesFactory, automatic_memory_optimization: bool = True): + super().__init__( + triples_factory=triples_factory, + automatic_memory_optimization=automatic_memory_optimization, + ) + + def forward( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + slice_size: Optional[int] = None, + slice_dim: Optional[str] = None, + ) -> torch.FloatTensor: # noqa: D102 + # (batch_size, num_heads, num_relations, num_tails) + scores = torch.zeros(1, 1, 1, 1, requires_grad=True) # for requires_grad + # reproducible scores + for i, (ind, num) in enumerate(( + (h_indices, self.num_entities), + (r_indices, self.num_relations), + (t_indices, self.num_entities), + )): + shape = [1, 1, 1, 1] + if ind is None: + shape[i + 1] = num + delta = torch.arange(num) + else: + shape[0] = len(ind) + delta = ind + scores = scores + delta.float().view(*shape) + return scores diff --git a/tests/test_early_stopping.py b/tests/test_early_stopping.py index e62dfa4979..6c8a4cd651 100644 --- a/tests/test_early_stopping.py +++ b/tests/test_early_stopping.py @@ -13,8 +13,9 @@ from pykeen.evaluation import Evaluator, MetricResults, RankBasedEvaluator, RankBasedMetricResults from pykeen.evaluation.rank_based_evaluator import RANK_TYPES, SIDES from pykeen.models import TransE -from pykeen.models.base import MockModel, Model +from pykeen.models.base import Model from pykeen.stoppers.early_stopping import EarlyStopper, is_improvement +from pykeen.testing.mocks import MockModel from pykeen.trackers import MLFlowResultTracker from pykeen.training import SLCWATrainingLoop from pykeen.typing import MappedTriples diff --git a/tests/test_evaluators.py b/tests/test_evaluators.py index da9418b0bb..7d5d4e3a09 100644 --- a/tests/test_evaluators.py +++ b/tests/test_evaluators.py @@ -15,7 +15,8 @@ from pykeen.evaluation.rank_based_evaluator import RANK_TYPES, SIDES, compute_rank_from_scores from pykeen.evaluation.sklearn import SklearnEvaluator, SklearnMetricResults from pykeen.models import TransE -from pykeen.models.base import MockModel, Model +from pykeen.models.base import Model +from pykeen.testing.mocks import MockModel from pykeen.triples import TriplesFactory from pykeen.typing import MappedTriples diff --git a/tests/test_model_mode.py b/tests/test_model_mode.py index 0ea1ccb3de..b40b0e9a50 100644 --- a/tests/test_model_mode.py +++ b/tests/test_model_mode.py @@ -10,7 +10,7 @@ from pykeen.datasets import Nations from pykeen.models import Model, TransE -from pykeen.models.base import MockModel +from pykeen.testing.mocks import MockModel from pykeen.triples import TriplesFactory from pykeen.utils import resolve_device diff --git a/tests/test_models.py b/tests/test_models.py index 12483d75e0..a55b958c25 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -24,8 +24,10 @@ from pykeen.datasets.nations import NATIONS_TEST_PATH, NATIONS_TRAIN_PATH, Nations from pykeen.models import _BASE_MODELS, _MODELS from pykeen.models.base import ( - DoubleRelationEmbeddingModel, ERModel, Model, - MultimodalModel, + DoubleRelationEmbeddingModel, + ERModel, + LiteralModel, + Model, SingleVectorEmbeddingModel, TwoSideEmbeddingModel, TwoVectorEmbeddingModel, @@ -47,7 +49,7 @@ SKIP_MODULES = { Model.__name__, 'DummyModel', - MultimodalModel.__name__, + LiteralModel.__name__, DoubleRelationEmbeddingModel.__name__, SingleVectorEmbeddingModel.__name__, TwoVectorEmbeddingModel.__name__, @@ -56,7 +58,7 @@ 'models', 'get_model_cls', } -for cls in MultimodalModel.__subclasses__(): +for cls in LiteralModel.__subclasses__(): SKIP_MODULES.add(cls.__name__) _EPSILON = 1.0e-07 @@ -953,7 +955,7 @@ def test_testing(self): isinstance(value, type) and issubclass(value, _ModelTestCase) and not name.startswith('_') - and not issubclass(value.model_cls, MultimodalModel) + and not issubclass(value.model_cls, LiteralModel) ) } tested_model_names -= SKIP_MODULES From 7cf80e8c407583bc1294f85891a3de1f6615e334 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 16:34:49 +0100 Subject: [PATCH 321/690] Fix dimension calculation for ConvE --- src/pykeen/nn/modules.py | 32 +++++++++++++++++--------------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 126dfd23e1..6695b9f8dd 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -457,8 +457,7 @@ def _calculate_missing_shape_information( width: Optional[int] = None, height: Optional[int] = None, ) -> Tuple[int, int, int]: - """ - Automatically calculates missing dimensions for ConvE. + """Automatically calculates missing dimensions for ConvE. :param embedding_dim: :param input_channels: @@ -473,26 +472,29 @@ def _calculate_missing_shape_information( # Store initial input for error message original = (input_channels, width, height) - # All are None - if all(factor is None for factor in [input_channels, width, height]): + # All are None -> try and make closest to square + if input_channels is None and width is None and height is None: input_channels = 1 result_sqrt = math.floor(math.sqrt(embedding_dim)) height = max(factor for factor in range(1, result_sqrt + 1) if embedding_dim % factor == 0) width = embedding_dim // height - - # input_channels is None, and any of height or width is None -> set input_channels=1 - if input_channels is None and any(remaining is None for remaining in [width, height]): - input_channels = 1 - - # input channels is not None, and one of height or width is None - if width is None and height is not None and input_channels is not None: + # Only input channels is None + elif input_channels is None and width is not None and height is not None: + input_channels = embedding_dim // (width * height) + # Only width is None + elif input_channels is not None and width is None and height is not None: width = embedding_dim // (height * input_channels) + # Only height is none elif height is None and width is not None and input_channels is not None: height = embedding_dim // (width * input_channels) - elif input_channels is None and height is not None and width is not None: - input_channels = embedding_dim // (width * height) - else: - raise ValueError('More than one of width, height, and input_channels was None') + # Width and input_channels are None -> set input_channels to 1 and calculage height + elif input_channels is None and height is None and width is not None: + input_channels = 1 + height = embedding_dim // width + # Width and input channels are None -> set input channels to 1 and calculate width + elif input_channels is None and height is not None and width is None: + input_channels = 1 + width = embedding_dim // height if input_channels * width * height != embedding_dim: raise ValueError(f'Could not resolve {original} to a valid factorization of {embedding_dim}.') From 13d0f71bbb92b2cb3c490b1c712289e534ee8e9c Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 16:42:59 +0100 Subject: [PATCH 322/690] Can't be them so join them --- src/pykeen/nn/modules.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 6695b9f8dd..41d2937d3a 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -496,10 +496,10 @@ def _calculate_missing_shape_information( input_channels = 1 width = embedding_dim // height - if input_channels * width * height != embedding_dim: + if input_channels * width * height != embedding_dim: # type: ignore raise ValueError(f'Could not resolve {original} to a valid factorization of {embedding_dim}.') - return input_channels, width, height + return input_channels, width, height # type: ignore class ConvEInteraction(Interaction[torch.FloatTensor, torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor]]): From 5f08ded9a5e79fb50068b0fe74be1218cc7564a5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 17:47:44 +0100 Subject: [PATCH 323/690] Extract one more method --- src/pykeen/nn/modules.py | 32 +++++++++++++++++--------------- 1 file changed, 17 insertions(+), 15 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 41d2937d3a..a6f0286ee5 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -222,14 +222,12 @@ def score( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return self._score( + return self._forward( h=h, r=r, t=t, - h_prefix=_get_prefix(slice_size=slice_size, slice_dim=slice_dim, d='h'), - r_prefix=_get_prefix(slice_size=slice_size, slice_dim=slice_dim, d='r'), - t_prefix=_get_prefix(slice_size=slice_size, slice_dim=slice_dim, d='t'), slice_size=slice_size, + slice_dim=slice_dim, ) def _score( @@ -262,6 +260,20 @@ def _score( r = self._add_dim(*r, dim=self.BATCH_DIM if r_prefix == "n" else self.NUM_DIM) t = self._add_dim(*t, dim=self.BATCH_DIM if t_prefix == "n" else self.NUM_DIM) + scores = self._forward(h, r, t, slice_dim, slice_size) + + remove_dims = [ + dim + for dim, prefix in zip( + (self.HEAD_DIM, self.RELATION_DIM, self.TAIL_DIM), + (h_prefix, r_prefix, t_prefix), + ) + if prefix == "b" + ] + # prepare output shape + return self._remove_dim(scores, *remove_dims) + + def _forward(self, h, r, t, slice_size, slice_dim): # get scores if slice_size is None: scores = self(h=h, r=r, t=t) @@ -276,17 +288,7 @@ def _score( scores = torch.cat(batch_scores, dim=self.TAIL_DIM) else: raise ValueError(f'Invalid slice_dim: {slice_dim}') - - remove_dims = [ - dim - for dim, prefix in zip( - (self.HEAD_DIM, self.RELATION_DIM, self.TAIL_DIM), - (h_prefix, r_prefix, t_prefix), - ) - if prefix == "b" - ] - # prepare output shape - return self._remove_dim(scores, *remove_dims) + return scores def score_hrt( self, From e3a0f16204eee905ae9eed06168959ea24710cc6 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 17:52:41 +0100 Subject: [PATCH 324/690] Add docstring to _forward_slicing_wrapper --- src/pykeen/nn/modules.py | 59 +++++++++++++++++++++++++++++----------- 1 file changed, 43 insertions(+), 16 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index a6f0286ee5..831f5724e5 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -222,13 +222,7 @@ def score( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return self._forward( - h=h, - r=r, - t=t, - slice_size=slice_size, - slice_dim=slice_dim, - ) + return self._forward_slicing_wrapper(h=h, r=r, t=t, slice_size=slice_size, slice_dim=slice_dim) def _score( self, @@ -260,7 +254,7 @@ def _score( r = self._add_dim(*r, dim=self.BATCH_DIM if r_prefix == "n" else self.NUM_DIM) t = self._add_dim(*t, dim=self.BATCH_DIM if t_prefix == "n" else self.NUM_DIM) - scores = self._forward(h, r, t, slice_dim, slice_size) + scores = self._forward_slicing_wrapper(h=h, r=r, t=t, slice_dim=slice_dim, slice_size=slice_size) remove_dims = [ dim @@ -273,19 +267,52 @@ def _score( # prepare output shape return self._remove_dim(scores, *remove_dims) - def _forward(self, h, r, t, slice_size, slice_dim): - # get scores + def _forward_slicing_wrapper( + self, + h: Union[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor, ...]], + r: Union[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor, ...]], + t: Union[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor, ...]], + slice_size: Optional[int], + slice_dim: str, + ) -> torch.FloatTensor: + """ + Compute broadcasted triple scores with optional slicing for representations in canonical shape. + + .. note :: + Depending on the interaction function, there may be more than one representation for h/r/t. In that case, + a tuple of at least two tensors is passed. + + :param h: shape: (batch_size, num_heads, ``*``) + The head representations. + :param r: shape: (batch_size, num_relations, ``*``) + The relation representations. + :param t: shape: (batch_size, num_tails, ``*``) + The tail representations. + :param slice_size: + The slice size. + :param slice_dim: + The dimension along which to slice. From {"h", "r", "t"} + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ if slice_size is None: scores = self(h=h, r=r, t=t) elif slice_dim == "h": - batch_scores = [self(h=h_batch, r=r, t=t) for h_batch in _get_batches(h, slice_size)] - scores = torch.cat(batch_scores, dim=self.HEAD_DIM) + scores = torch.cat([ + self(h=h_batch, r=r, t=t) + for h_batch in _get_batches(h, slice_size) + ], dim=self.HEAD_DIM) elif slice_dim == "r": - batch_scores = [self(h=h, r=r_batch, t=t) for r_batch in _get_batches(r, slice_size)] - scores = torch.cat(batch_scores, dim=self.RELATION_DIM) + scores = torch.cat([ + self(h=h, r=r_batch, t=t) + for r_batch in _get_batches(r, slice_size) + ], dim=self.RELATION_DIM) elif slice_dim == "t": - batch_scores = [self(h=h, r=r, t=t_batch) for t_batch in _get_batches(t, slice_size)] - scores = torch.cat(batch_scores, dim=self.TAIL_DIM) + scores = torch.cat([ + self(h=h, r=r, t=t_batch) + for t_batch in _get_batches(t, slice_size) + ], dim=self.TAIL_DIM) else: raise ValueError(f'Invalid slice_dim: {slice_dim}') return scores From c5cc4933829823bd60a193f7c3e33303e692de8e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 17:53:52 +0100 Subject: [PATCH 325/690] Update type annotation to show that the tuple contains more than one tensor --- src/pykeen/nn/modules.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 831f5724e5..06c4772b54 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -113,7 +113,7 @@ def forward( raise NotImplementedError @staticmethod - def _add_dim(*x: torch.FloatTensor, dim: int) -> Sequence[torch.FloatTensor]: + def _add_dim(*x: torch.FloatTensor, dim: int) -> Union[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor, ...]]: """ Add a dimension to tensors. From c9ae83006c4765ff1406a9da97cd6ba6e220a37e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 17:58:26 +0100 Subject: [PATCH 326/690] Add docstring --- src/pykeen/nn/modules.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 06c4772b54..5e5a6b5839 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -85,7 +85,8 @@ class Interaction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, #: The symbolic shapes for entity representations entity_shape: Sequence[str] = ("d",) - # TODO add docstring + + #: The symbolic shapes for entity representations for tail entities, if different. This is ony relevant for ConvE. tail_entity_shape: Optional[Sequence[str]] = None #: The symbolic shapes for relation representations @@ -126,7 +127,7 @@ def _add_dim(*x: torch.FloatTensor, dim: int) -> Union[torch.FloatTensor, Tuple[ out = [xx.unsqueeze(dim=dim) for xx in x] if len(x) == 1: return out[0] - return out + return tuple(out) @staticmethod def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: From 05a36720bce444358818489a780e5b1b2949bc2a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 18:05:04 +0100 Subject: [PATCH 327/690] Add regularizer to Embedding --- src/pykeen/nn/emb.py | 23 ++++++++++++++++------- 1 file changed, 16 insertions(+), 7 deletions(-) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index 8e3ecffc63..a1a9313383 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -13,6 +13,7 @@ import torch.nn from torch import nn +from ..regularizers import Regularizer from ..typing import Constrainer, Initializer, Normalizer __all__ = [ @@ -86,7 +87,7 @@ class EmbeddingSpecification: constrainer: Optional[Constrainer] = None constrainer_kwargs: Optional[Mapping[str, Any]] = None - # regularizer: Optional[Regularizer] = None + regularizer: Optional[Regularizer] = None def make( self, @@ -105,6 +106,7 @@ def make( normalizer_kwargs=self.normalizer_kwargs, constrainer=self.constrainer, constrainer_kwargs=self.constrainer_kwargs, + regularizer=self.regularizer, ) @@ -118,6 +120,7 @@ class Embedding(RepresentationModule): shape: Sequence[int] normalizer: Optional[Normalizer] constrainer: Optional[Constrainer] + regularizer: Optional[Regularizer] def __init__( self, @@ -130,6 +133,7 @@ def __init__( normalizer_kwargs: Optional[Mapping[str, Any]] = None, constrainer: Optional[Constrainer] = None, constrainer_kwargs: Optional[Mapping[str, Any]] = None, + regularizer: Optional[Regularizer] = None, ): """Instantiate an embedding with extended functionality. @@ -169,19 +173,20 @@ def __init__( if initializer is None: initializer = nn.init.normal_ + if initializer_kwargs: initializer = functools.partial(initializer, **initializer_kwargs) self.initializer = initializer if constrainer is not None and constrainer_kwargs: - self.constrainer: Constrainer = functools.partial(constrainer, **constrainer_kwargs) - else: - self.constrainer = constrainer + constrainer: Constrainer = functools.partial(constrainer, **constrainer_kwargs) + self.constrainer = constrainer if normalizer is not None and normalizer_kwargs: - self.normalizer: Normalizer = functools.partial(normalizer, **normalizer_kwargs) - else: - self.normalizer = normalizer + normalizer: Normalizer = functools.partial(normalizer, **normalizer_kwargs) + self.normalizer = normalizer + + self.regularizer = regularizer self._embeddings = torch.nn.Embedding( num_embeddings=num_embeddings, @@ -235,6 +240,8 @@ def post_parameter_update(self): # noqa: D102 # apply constraints in-place if self.constrainer is not None: self._embeddings.weight.data = self.constrainer(self._embeddings.weight.data) + if self.regularizer is not None: + self.regularizer.reset() def forward( self, @@ -246,6 +253,8 @@ def forward( x = self._embeddings(indices) if self.normalizer is not None: x = self.normalizer(x) + if self.regularizer is not None: + self.regularizer.update(x) return x def get_in_canonical_shape( From e348d44680e28c3eee3af1ad508606a6768681b4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 18:05:41 +0100 Subject: [PATCH 328/690] Add TODO --- src/pykeen/nn/emb.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index a1a9313383..187fd48bf0 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -182,6 +182,7 @@ def __init__( constrainer: Constrainer = functools.partial(constrainer, **constrainer_kwargs) self.constrainer = constrainer + # TODO: Move regularizer and normalizer to RepresentationModule? if normalizer is not None and normalizer_kwargs: normalizer: Normalizer = functools.partial(normalizer, **normalizer_kwargs) self.normalizer = normalizer From d2afe36063c92229f50ff7a1b97653a539cc3eb8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 18:07:17 +0100 Subject: [PATCH 329/690] Do not update regularizer outside of training --- src/pykeen/regularizers.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pykeen/regularizers.py b/src/pykeen/regularizers.py index c746c7142d..3a4fad17c2 100644 --- a/src/pykeen/regularizers.py +++ b/src/pykeen/regularizers.py @@ -74,6 +74,8 @@ def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: def update(self, *tensors: torch.FloatTensor) -> None: """Update the regularization term based on passed tensors.""" + if not self.training: + return if self.apply_only_once and self.updated: return self.regularization_term = self.regularization_term + sum(self(x) for x in tensors) From f2072ba64efb2590083eec70840bde326c0decfb Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 18:07:46 +0100 Subject: [PATCH 330/690] Do not update regularizer if no gradients are tracked --- src/pykeen/regularizers.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/pykeen/regularizers.py b/src/pykeen/regularizers.py index 3a4fad17c2..976dcbe566 100644 --- a/src/pykeen/regularizers.py +++ b/src/pykeen/regularizers.py @@ -74,9 +74,7 @@ def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: def update(self, *tensors: torch.FloatTensor) -> None: """Update the regularization term based on passed tensors.""" - if not self.training: - return - if self.apply_only_once and self.updated: + if not self.training or not torch.is_grad_enabled() or (self.apply_only_once and self.updated): return self.regularization_term = self.regularization_term + sum(self(x) for x in tensors) self.updated = True From b7002fe26d4d77b91fa75ea4a9a95e790d2ccd51 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 18:26:26 +0100 Subject: [PATCH 331/690] Refactor regularizer --- src/pykeen/models/base.py | 52 ++----------------- .../models/multimodal/complex_literal.py | 1 - .../models/multimodal/distmult_literal.py | 1 - src/pykeen/models/unimodal/complex.py | 1 - src/pykeen/models/unimodal/conv_e.py | 2 - src/pykeen/models/unimodal/conv_kb.py | 1 - src/pykeen/models/unimodal/distmult.py | 1 - src/pykeen/models/unimodal/ermlp.py | 1 - src/pykeen/models/unimodal/ermlpe.py | 1 - src/pykeen/models/unimodal/hole.py | 1 - src/pykeen/models/unimodal/kg2e.py | 1 - src/pykeen/models/unimodal/ntn.py | 1 - src/pykeen/models/unimodal/proj_e.py | 1 - src/pykeen/models/unimodal/rescal.py | 1 - src/pykeen/models/unimodal/rotate.py | 1 - src/pykeen/models/unimodal/simple.py | 1 - .../models/unimodal/structured_embedding.py | 1 - src/pykeen/models/unimodal/trans_d.py | 1 - src/pykeen/models/unimodal/trans_e.py | 1 - src/pykeen/models/unimodal/trans_h.py | 1 - src/pykeen/models/unimodal/trans_r.py | 1 - src/pykeen/models/unimodal/tucker.py | 1 - .../models/unimodal/unstructured_model.py | 1 - src/pykeen/nn/emb.py | 2 - src/pykeen/regularizers.py | 50 +++++++++--------- src/pykeen/training/training_loop.py | 5 -- tests/test_regularizers.py | 16 ++---- 27 files changed, 32 insertions(+), 116 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 5dd928112c..137ce63464 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -23,7 +23,7 @@ from ..nn import Embedding, RepresentationModule from ..nn.emb import EmbeddingSpecification, LiteralRepresentations from ..nn.modules import Interaction -from ..regularizers import NoRegularizer, Regularizer +from ..regularizers import collect_regularization_terms from ..triples import TriplesFactory, TriplesNumericLiteralsFactory from ..typing import DeviceHint, HeadRepresentation, MappedTriples, RelationRepresentation, TailRepresentation from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed @@ -255,14 +255,6 @@ class Model(nn.Module, ABC): #: The instance of the loss loss: Loss - # FIXME problem with typing (remove type: ignore) - #: The default regularizer class - regularizer_default: ClassVar[Type[Regularizer]] = NoRegularizer # type: ignore - #: The default parameters for the default regularizer class - regularizer_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = None - #: The instance of the regularizer - regularizer: Regularizer - def __init__( self, triples_factory: TriplesFactory, @@ -271,7 +263,6 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the module. @@ -291,8 +282,6 @@ def __init__( The preferred device for model training and inference. :param random_seed: A random seed to use for initialising the model's weights. **Should** be set when aiming at reproducibility. - :param regularizer: - A regularizer to use for training. """ super().__init__() @@ -316,15 +305,6 @@ def __init__( # TODO: Check loss functions that require 1 and -1 as label but only self.is_mr_loss: bool = isinstance(self.loss, MarginRankingLoss) - - # Regularizer - if regularizer is None: - regularizer = self.regularizer_default( - device=self.device, - **(self.regularizer_default_kwargs or {}), - ) - self.regularizer = regularizer - self.is_nssa_loss: bool = isinstance(self.loss, NSSALoss) # The triples factory facilitates access to the dataset. @@ -439,7 +419,6 @@ def _set_device(self, device: DeviceHint = None) -> None: def to_device_(self) -> 'Model': """Transfer model to device.""" self.to(self.device) - self.regularizer.to(self.device) torch.cuda.empty_cache() return self @@ -894,21 +873,12 @@ def make_labeled_df( def post_parameter_update(self) -> None: """Has to be called after each parameter update.""" - self.regularizer.reset() for module in self.modules(): if module is self: continue if hasattr(module, "post_parameter_update"): module.post_parameter_update() - def regularize_if_necessary(self, *tensors: torch.FloatTensor) -> None: - """Update the regularizer's term given some tensors, if regularization is requested. - - :param tensors: The tensors that should be passed to the regularizer to update its term. - """ - if self.training: - self.regularizer.update(*tensors) - def compute_mr_loss( self, positive_scores: torch.FloatTensor, @@ -931,7 +901,7 @@ def compute_mr_loss( ' losses. Please use the compute_loss method instead.', ) y = torch.ones_like(negative_scores, device=self.device) - return self.loss(positive_scores, negative_scores, y) + self.regularizer.term + return self.loss(positive_scores, negative_scores, y) + collect_regularization_terms(self) def compute_label_loss( self, @@ -994,7 +964,7 @@ def _compute_loss( 'The chosen loss does not allow the calculation of margin label' ' losses. Please use the compute_mr_loss method instead.', ) - return self.loss(tensor_1, tensor_2) + self.regularizer.term + return self.loss(tensor_1, tensor_2) + collect_regularization_terms(self) @abstractmethod def forward( @@ -1156,7 +1126,6 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, entity_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, relation_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, ) -> None: @@ -1178,8 +1147,6 @@ def __init__( The preferred device for model training and inference. :param random_seed: A random seed to use for initialising the model's weights. **Should** be set when aiming at reproducibility. - :param regularizer: - A regularizer to use for training. """ super().__init__( triples_factory=triples_factory, @@ -1187,7 +1154,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, predict_with_sigmoid=predict_with_sigmoid, ) @@ -1304,7 +1270,6 @@ def __init__( predict_with_sigmoid: bool = False, preferred_device: Optional[str] = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, embedding_specification: Optional[EmbeddingSpecification] = None, relation_embedding_specification: Optional[EmbeddingSpecification] = None, ) -> None: @@ -1325,8 +1290,6 @@ def __init__( The default device where to model is located. :param random_seed: An optional random seed to set before the initialization of weights. - :param regularizer: - The regularizer to use. """ # Default for relation dimensionality if relation_dim is None: @@ -1338,7 +1301,6 @@ def __init__( predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, interaction=interaction, entity_representations=Embedding.from_specification( num_embeddings=triples_factory.num_entities, @@ -1391,7 +1353,6 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, embedding_specification: Optional[EmbeddingSpecification] = None, relation_embedding_specification: Optional[EmbeddingSpecification] = None, second_relation_embedding_specification: Optional[EmbeddingSpecification] = None, @@ -1406,7 +1367,6 @@ def __init__( predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, entity_representations=[ Embedding.from_specification( num_embeddings=triples_factory.num_entities, @@ -1453,7 +1413,6 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, embedding_specification: Optional[EmbeddingSpecification] = None, relation_embedding_specification: Optional[EmbeddingSpecification] = None, second_embedding_specification: Optional[EmbeddingSpecification] = None, @@ -1468,7 +1427,6 @@ def __init__( predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, entity_representations=[ Embedding.from_specification( num_embeddings=triples_factory.num_entities, @@ -1516,7 +1474,6 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, embedding_specification: Optional[EmbeddingSpecification] = None, relation_embedding_specification: Optional[EmbeddingSpecification] = None, second_embedding_specification: Optional[EmbeddingSpecification] = None, @@ -1531,7 +1488,6 @@ def __init__( predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, entity_representations=[ Embedding.from_specification( num_embeddings=triples_factory.num_entities, @@ -1600,7 +1556,6 @@ def __init__( automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ): super().__init__( triples_factory=triples_factory, @@ -1610,7 +1565,6 @@ def __init__( automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, entity_representations=[ # entity embeddings Embedding.from_specification( diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index 31c5dcdc93..37d670c52a 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -107,5 +107,4 @@ def __init__( automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, ) diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 85a67977b1..327d30993a 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -67,5 +67,4 @@ def __init__( automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, ) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 98c6f70a18..4f3ba49bc7 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -102,7 +102,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, # initialize with entity and relation embeddings with standard normal distribution, cf. # https://github.com/ttrouill/complex/blob/dc4eb93408d9a5288c986695b58488ac80b1cc17/efe/models.py#L481-L487 embedding_specification=EmbeddingSpecification( diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index de66334368..2e7ec0f74f 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -141,7 +141,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, entity_representations=[ Embedding.from_specification( num_embeddings=triples_factory.num_entities, @@ -193,5 +192,4 @@ def forward( r = self.relation_representations[0].get_in_canonical_shape(indices=r_indices) t = self.entity_representations[0].get_in_canonical_shape(indices=t_indices) t_bias = self.entity_representations[1].get_in_canonical_shape(indices=t_indices) - self.regularize_if_necessary(h, r, t) return self.interaction.score(h=h, r=r, t=(t, t_bias), slice_size=slice_size, slice_dim=slice_dim) diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 0fb2441d50..55929611f8 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -96,6 +96,5 @@ def __init__( automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, ) logger.warning('To be consistent with the paper, initialize entity and relation embeddings from TransE.') diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index c5a7389845..d54d772431 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -89,7 +89,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, embedding_specification=EmbeddingSpecification( # xavier uniform, cf. # https://github.com/thunlp/OpenKE/blob/adeed2c0d2bef939807ed4f69c1ea4db35fd149b/models/DistMult.py#L16-L17 diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index b85f34aa51..6d0bb8e2f0 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -64,5 +64,4 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, ) diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index 1386ff4eb2..d26c0f5c62 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -77,5 +77,4 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, ) diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 509f240dbe..71ea5867a2 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -71,7 +71,6 @@ def __init__( automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, # Initialisation, cf. https://github.com/mnick/scikit-kge/blob/master/skge/param.py#L18-L27 embedding_specification=EmbeddingSpecification( initializer=xavier_uniform_, diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 4e73f17f92..a3c39c5f1b 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -88,7 +88,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, embedding_specification=EmbeddingSpecification( constrainer=clamp_norm, # type: ignore constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 91d6f75e1e..96f5e98992 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -97,7 +97,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, interaction=NTNInteraction( non_linearity=non_linearity, ), diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index 2445fcd6a3..5bae08baf6 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -76,7 +76,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, embedding_specification=EmbeddingSpecification( initializer=xavier_uniform_, ), diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index 9b9fb28627..816903c832 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -77,5 +77,4 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, ) diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index 7404b40c8d..896dd29d4b 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -102,7 +102,6 @@ def __init__( automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, embedding_specification=EmbeddingSpecification( initializer=xavier_uniform_, ), diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index 1b05f2366f..ddf47894ed 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -81,7 +81,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, ) if isinstance(clamp_score, float): clamp_score = (-clamp_score, clamp_score) diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index 0334f68dc0..1070636cec 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -81,7 +81,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, embedding_specification=EmbeddingSpecification( initializer=xavier_uniform_, constrainer=functional.normalize, diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index a1c90834e5..6b2c825b5b 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -77,7 +77,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, embedding_specification=EmbeddingSpecification( initializer=xavier_normal_, constrainer=clamp_norm, # type: ignore diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 3d5333ac5d..473911a007 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -75,7 +75,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, embedding_specification=EmbeddingSpecification( initializer=xavier_uniform_, constrainer=functional.normalize, diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index afa289e771..572c000f15 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -90,7 +90,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, predict_with_sigmoid=predict_with_sigmoid, embedding_specification=None, relation_embedding_specification=None, diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index af75005e15..0300a80d88 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -79,7 +79,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, entity_representations=Embedding.from_specification( num_embeddings=triples_factory.num_entities, shape=embedding_dim, diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index 0bd9fa158f..57c29a0829 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -103,7 +103,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, embedding_specification=EmbeddingSpecification( initializer=xavier_normal_, ), diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index 4262ab607d..deebb82777 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -67,7 +67,6 @@ def __init__( predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, interaction=UnstructuredModelInteraction(p=scoring_fct_norm), entity_representations=Embedding( num_embeddings=triples_factory.num_entities, diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index 187fd48bf0..65b138c3d3 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -241,8 +241,6 @@ def post_parameter_update(self): # noqa: D102 # apply constraints in-place if self.constrainer is not None: self._embeddings.weight.data = self.constrainer(self._embeddings.weight.data) - if self.regularizer is not None: - self.regularizer.reset() def forward( self, diff --git a/src/pykeen/regularizers.py b/src/pykeen/regularizers.py index 976dcbe566..125ef43f65 100644 --- a/src/pykeen/regularizers.py +++ b/src/pykeen/regularizers.py @@ -18,6 +18,7 @@ 'CombinedRegularizer', 'PowerSumRegularizer', 'TransHRegularizer', + 'collect_regularization_terms', 'get_regularizer_cls', ] @@ -31,7 +32,7 @@ class Regularizer(nn.Module, ABC): weight: torch.FloatTensor #: The current regularization term (a scalar) - regularization_term: torch.FloatTensor + regularization_term: Union[torch.FloatTensor, float] #: Should the regularization only be applied once? This was used for ConvKB and defaults to False. apply_only_once: bool @@ -41,32 +42,19 @@ class Regularizer(nn.Module, ABC): def __init__( self, - device: torch.device, weight: float = 1.0, apply_only_once: bool = False, ): super().__init__() - self.device = device self.register_buffer(name='weight', tensor=torch.as_tensor(weight, device=self.device)) self.apply_only_once = apply_only_once - self.reset() - - def to(self, *args, **kwargs) -> 'Regularizer': # noqa: D102 - super().to(*args, **kwargs) - self.device = torch._C._nn._parse_to(*args, **kwargs)[0] - self.reset() - return self + self.pop_regularization_term() @classmethod def get_normalized_name(cls) -> str: """Get the normalized name of the regularizer class.""" return normalize_string(cls.__name__, suffix=_REGULARIZER_SUFFIX) - def reset(self) -> None: - """Reset the regularization term to zero.""" - self.regularization_term = torch.zeros(1, dtype=torch.float, device=self.device) - self.updated = False - @abstractmethod def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: """Compute the regularization term for one tensor.""" @@ -79,10 +67,18 @@ def update(self, *tensors: torch.FloatTensor) -> None: self.regularization_term = self.regularization_term + sum(self(x) for x in tensors) self.updated = True + def pop_regularization_term(self) -> torch.FloatTensor: + """Return the weighted regularization term, and clear it afterwards.""" + term = self.regularization_term + self.regularization_term = 0. + self.updated = False + return self.weight * term + @property def term(self) -> torch.FloatTensor: """Return the weighted regularization term.""" - return self.regularization_term * self.weight + # TODO: Deprecated + raise RuntimeError class NoRegularizer(Regularizer): @@ -100,7 +96,7 @@ def update(self, *tensors: torch.FloatTensor) -> None: # noqa: D102 def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102 # always return zero - return torch.zeros(1, dtype=x.dtype, device=x.device) + return x.new_zeros(1) class LpRegularizer(Regularizer): @@ -120,14 +116,13 @@ class LpRegularizer(Regularizer): def __init__( self, - device: torch.device, weight: float = 1.0, dim: Optional[int] = -1, normalize: bool = False, p: float = 2., apply_only_once: bool = False, ): - super().__init__(device=device, weight=weight, apply_only_once=apply_only_once) + super().__init__(weight=weight, apply_only_once=apply_only_once) self.dim = dim self.normalize = normalize self.p = p @@ -160,14 +155,13 @@ class PowerSumRegularizer(Regularizer): def __init__( self, - device: torch.device, weight: float = 1.0, dim: Optional[int] = -1, normalize: bool = False, p: float = 2., apply_only_once: bool = False, ): - super().__init__(device=device, weight=weight, apply_only_once=apply_only_once) + super().__init__(weight=weight, apply_only_once=apply_only_once) self.dim = dim self.normalize = normalize self.p = p @@ -190,13 +184,12 @@ class TransHRegularizer(Regularizer): def __init__( self, - device: torch.device, weight: float = 0.05, epsilon: float = 1e-5, ): # The regularization in TransH enforces the defined soft constraints that should computed only for every batch. # Therefore, apply_only_once is always set to True. - super().__init__(device=device, weight=weight, apply_only_once=True) + super().__init__(weight=weight, apply_only_once=True) self.epsilon = epsilon def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102 @@ -233,7 +226,7 @@ def __init__( total_weight: float = 1.0, apply_only_once: bool = False, ): - super().__init__(weight=total_weight, device=device, apply_only_once=apply_only_once) + super().__init__(weight=total_weight, apply_only_once=apply_only_once) self.regularizers = nn.ModuleList(regularizers) for r in self.regularizers: if isinstance(r, NoRegularizer): @@ -274,3 +267,12 @@ def get_regularizer_cls(query: Union[None, str, Type[Regularizer]]) -> Type[Regu default=NoRegularizer, suffix=_REGULARIZER_SUFFIX, ) + + +def collect_regularization_terms(main_module: nn.Module) -> Union[float, torch.FloatTensor]: + """Recursively collect regularization terms from attached regularizers, and clear their accumulator.""" + return sum( + module.pop_regularization_term() + for module in main_module.modules() + if isinstance(module, Regularizer) + ) diff --git a/src/pykeen/training/training_loop.py b/src/pykeen/training/training_loop.py index 8c26f3acb3..5bd89dd2a1 100644 --- a/src/pykeen/training/training_loop.py +++ b/src/pykeen/training/training_loop.py @@ -442,9 +442,6 @@ def _forward_pass(self, batch, start, stop, current_batch_size, label_smoothing, loss.backward() current_epoch_loss = loss.item() - # reset the regularizer to free the computational graph - self.model.regularizer.reset() - return current_epoch_loss @staticmethod @@ -671,7 +668,5 @@ def to_embeddingdb(self, session=None, use_tqdm: bool = False): return self.model.to_embeddingdb(session=session, use_tqdm=use_tqdm) def _free_graph_and_cache(self): - # The regularizer has to be reset to free the computational graph - self.model.regularizer.reset() # The cache of the previous run has to be freed to allow accurate memory availability estimates torch.cuda.empty_cache() diff --git a/tests/test_regularizers.py b/tests/test_regularizers.py index cd562fedf2..b2be89b92a 100644 --- a/tests/test_regularizers.py +++ b/tests/test_regularizers.py @@ -70,13 +70,6 @@ def test_model(self) -> None: # Check if regularization term is reset self.assertEqual(0., model.regularizer.term) - def test_reset(self) -> None: - """Test method `reset`.""" - # Call method - self.regularizer.reset() - - self.assertEqual(0., self.regularizer.regularization_term) - def test_update(self) -> None: """Test method `update`.""" # Generate random tensors @@ -168,8 +161,8 @@ class CombinedRegularizerTest(_RegularizerTestCase, unittest.TestCase): regularizer_cls = CombinedRegularizer regularizer_kwargs = { 'regularizers': [ - LpRegularizer(weight=0.1, p=1, device=resolve_device()), - LpRegularizer(weight=0.7, p=2, device=resolve_device()), + LpRegularizer(weight=0.1, p=1), + LpRegularizer(weight=0.7, p=2), ], } @@ -212,7 +205,6 @@ def setUp(self) -> None: self.device = resolve_device() self.regularizer_kwargs = {'weight': .5, 'epsilon': 1e-5} self.regularizer = TransHRegularizer( - device=self.device, **(self.regularizer_kwargs or {}), ) self.num_entities = 10 @@ -275,7 +267,6 @@ def test_lp(self): self.assertIn('apply_only_once', ConvKB.regularizer_default_kwargs) self.assertTrue(ConvKB.regularizer_default_kwargs['apply_only_once']) regularizer = LpRegularizer( - device=self.device, **ConvKB.regularizer_default_kwargs, ) self._help_test_regularizer(regularizer) @@ -284,7 +275,6 @@ def test_transh_regularizer(self): """Test the TransH regularizer only updates once.""" self.assertNotIn('apply_only_once', TransH.regularizer_default_kwargs) regularizer = TransHRegularizer( - device=self.device, **TransH.regularizer_default_kwargs, ) self._help_test_regularizer(regularizer) @@ -312,6 +302,6 @@ def _help_test_regularizer(self, regularizer: Regularizer, n_tensors: int = 3): self.assertTrue(regularizer.updated) self.assertEqual(term, regularizer.regularization_term) - regularizer.reset() + # regularizer.reset() self.assertFalse(regularizer.updated) self.assertEqual(0.0, regularizer.regularization_term.item()) From 327190a428c927784898f9ef1d4a3e91423395dd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 18:32:55 +0100 Subject: [PATCH 332/690] Adjust regularizers from ComplEx - ER-MLPE --- src/pykeen/models/unimodal/complex.py | 15 +++++++++++---- src/pykeen/models/unimodal/conv_e.py | 1 - src/pykeen/models/unimodal/conv_kb.py | 4 ++-- src/pykeen/models/unimodal/distmult.py | 9 +++++++-- src/pykeen/models/unimodal/ermlp.py | 2 -- src/pykeen/models/unimodal/ermlpe.py | 2 -- 6 files changed, 20 insertions(+), 13 deletions(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 4f3ba49bc7..46b473c22a 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -10,7 +10,7 @@ from ...losses import Loss, SoftplusLoss from ...nn.emb import EmbeddingSpecification from ...nn.modules import ComplExInteraction -from ...regularizers import LpRegularizer, Regularizer +from ...regularizers import LpRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -74,7 +74,6 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: """Initialize ComplEx. @@ -91,8 +90,6 @@ def __init__( The default device where to model is located. :param random_seed: int (optional) An optional random seed to set before the initialization of weights. - :param regularizer: BaseRegularizer - The regularizer to use. """ super().__init__( triples_factory=triples_factory, @@ -106,8 +103,18 @@ def __init__( # https://github.com/ttrouill/complex/blob/dc4eb93408d9a5288c986695b58488ac80b1cc17/efe/models.py#L481-L487 embedding_specification=EmbeddingSpecification( initializer=nn.init.normal_, + regularizer=LpRegularizer( + weight=0.01, + p=2.0, + normalize=True, + ), ), relation_embedding_specification=EmbeddingSpecification( initializer=nn.init.normal_, + regularizer=LpRegularizer( + weight=0.01, + p=2.0, + normalize=True, + ), ), ) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 2e7ec0f74f..aac26a4787 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -124,7 +124,6 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, apply_batch_normalization: bool = True, ) -> None: """Initialize the model.""" diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 55929611f8..da12b3aeb3 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -8,7 +8,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss from ...nn.modules import ConvKBInteraction -from ...regularizers import LpRegularizer, Regularizer +from ...regularizers import LpRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -78,12 +78,12 @@ def __init__( preferred_device: DeviceHint = None, num_filters: int = 400, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model. To be consistent with the paper, pass entity and relation embeddings pre-trained from TransE. """ + # TODO: regularize weight and bias super().__init__( triples_factory=triples_factory, interaction=ConvKBInteraction( diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index d54d772431..6b1c188089 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -11,7 +11,7 @@ from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.modules import DistMultInteraction -from ...regularizers import LpRegularizer, Regularizer +from ...regularizers import LpRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint from ...utils import compose @@ -75,7 +75,6 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: r"""Initialize DistMult. @@ -102,5 +101,11 @@ def __init__( nn.init.xavier_uniform_, functional.normalize, ), + # Only relation embeddings are regularized + regularizer=LpRegularizer( + weight=0.1, + p=2.0, + normalize=True, + ), ), ) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 6d0bb8e2f0..bd575f4310 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -7,7 +7,6 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss from ...nn.modules import ERMLPInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -47,7 +46,6 @@ def __init__( preferred_device: DeviceHint = None, random_seed: Optional[int] = None, hidden_dim: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: """Initialize ERMLP.""" if hidden_dim is None: diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index d26c0f5c62..340c43a3b8 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -7,7 +7,6 @@ from ..base import SingleVectorEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss from ...nn.modules import ERMLPEInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -62,7 +61,6 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: super().__init__( triples_factory=triples_factory, From e0d24658acb5a0a7eb2e82d9ece9ee8b50b58563 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 18:43:47 +0100 Subject: [PATCH 333/690] Add regularizer to HolE - RESCAL --- src/pykeen/models/unimodal/hole.py | 2 -- src/pykeen/models/unimodal/kg2e.py | 1 - src/pykeen/models/unimodal/ntn.py | 2 -- src/pykeen/models/unimodal/proj_e.py | 2 -- src/pykeen/models/unimodal/rescal.py | 19 ++++++++++++++++--- tests/test_regularizers.py | 1 - 6 files changed, 16 insertions(+), 11 deletions(-) diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 71ea5867a2..b041d030a2 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -9,7 +9,6 @@ from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import HolEInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint from ...utils import clamp_norm @@ -60,7 +59,6 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model.""" super().__init__( diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index a3c39c5f1b..35312bd72f 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -69,7 +69,6 @@ def __init__( dist_similarity: Optional[str] = None, c_min: float = 0.05, c_max: float = 5., - regularizer: Optional[Regularizer] = None, ) -> None: r"""Initialize KG2E. diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 96f5e98992..a0b278ad9e 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -10,7 +10,6 @@ from ...losses import Loss from ...nn import Embedding from ...nn.modules import NTNInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -62,7 +61,6 @@ def __init__( preferred_device: DeviceHint = None, random_seed: Optional[int] = None, non_linearity: Optional[nn.Module] = None, - regularizer: Optional[Regularizer] = None, ) -> None: r"""Initialize NTN. diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index 5bae08baf6..d75c526ffa 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -11,7 +11,6 @@ from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import ProjEInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -63,7 +62,6 @@ def __init__( preferred_device: DeviceHint = None, random_seed: Optional[int] = None, inner_non_linearity: Optional[nn.Module] = None, - regularizer: Optional[Regularizer] = None, ) -> None: super().__init__( triples_factory=triples_factory, diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index 816903c832..e5493a4724 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -6,8 +6,9 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss +from ...nn.emb import EmbeddingSpecification from ...nn.modules import RESCALInteraction -from ...regularizers import LpRegularizer, Regularizer +from ...regularizers import LpRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -57,7 +58,6 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: r"""Initialize RESCAL. @@ -67,7 +67,6 @@ def __init__( - OpenKE `implementation of RESCAL `_ """ - # TODO: regularization super().__init__( triples_factory=triples_factory, interaction=RESCALInteraction(), @@ -77,4 +76,18 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, + embedding_specification=EmbeddingSpecification( + regularizer=LpRegularizer( + weight=10, + p=2., + normalize=True, + ), + ), + relation_embedding_specification=EmbeddingSpecification( + regularizer=LpRegularizer( + weight=10, + p=2., + normalize=True, + ), + ), ) diff --git a/tests/test_regularizers.py b/tests/test_regularizers.py index b2be89b92a..2075aed461 100644 --- a/tests/test_regularizers.py +++ b/tests/test_regularizers.py @@ -55,7 +55,6 @@ def test_model(self) -> None: # Use RESCAL as it regularizes multiple tensors of different shape. model = RESCAL( triples_factory=self.triples_factory, - regularizer=self.regularizer, ).to(self.device) # Check if regularizer is stored correctly. From 6b3a0bd0bf142e5afcba6e6f307ed385e4cfffce Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 18:51:21 +0100 Subject: [PATCH 334/690] Add regularizers --- src/pykeen/models/unimodal/complex.py | 13 +++---------- src/pykeen/models/unimodal/rescal.py | 13 +++---------- src/pykeen/models/unimodal/rotate.py | 3 --- src/pykeen/models/unimodal/simple.py | 17 +++++++++++++++-- .../models/unimodal/structured_embedding.py | 2 -- src/pykeen/models/unimodal/trans_d.py | 2 -- src/pykeen/models/unimodal/trans_e.py | 2 -- src/pykeen/models/unimodal/trans_h.py | 4 ++-- src/pykeen/models/unimodal/trans_r.py | 2 -- src/pykeen/models/unimodal/tucker.py | 2 -- .../models/unimodal/unstructured_model.py | 2 -- 11 files changed, 23 insertions(+), 39 deletions(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 46b473c22a..1d853ce6a7 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -91,6 +91,7 @@ def __init__( :param random_seed: int (optional) An optional random seed to set before the initialization of weights. """ + regularizer = LpRegularizer(weight=0.01, p=2.0, normalize=True) super().__init__( triples_factory=triples_factory, interaction=ComplExInteraction(), @@ -103,18 +104,10 @@ def __init__( # https://github.com/ttrouill/complex/blob/dc4eb93408d9a5288c986695b58488ac80b1cc17/efe/models.py#L481-L487 embedding_specification=EmbeddingSpecification( initializer=nn.init.normal_, - regularizer=LpRegularizer( - weight=0.01, - p=2.0, - normalize=True, - ), + regularizer=regularizer, ), relation_embedding_specification=EmbeddingSpecification( initializer=nn.init.normal_, - regularizer=LpRegularizer( - weight=0.01, - p=2.0, - normalize=True, - ), + regularizer=regularizer, ), ) diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index e5493a4724..f9070dbaff 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -67,6 +67,7 @@ def __init__( - OpenKE `implementation of RESCAL `_ """ + regularizer = LpRegularizer(weight=10, p=2., normalize=True) super().__init__( triples_factory=triples_factory, interaction=RESCALInteraction(), @@ -77,17 +78,9 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, embedding_specification=EmbeddingSpecification( - regularizer=LpRegularizer( - weight=10, - p=2., - normalize=True, - ), + regularizer=regularizer, ), relation_embedding_specification=EmbeddingSpecification( - regularizer=LpRegularizer( - weight=10, - p=2., - normalize=True, - ), + regularizer=regularizer, ), ) diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index 896dd29d4b..4759f53f83 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -14,7 +14,6 @@ from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import RotatEInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -91,9 +90,7 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: - # TODO: regularization super().__init__( triples_factory=triples_factory, interaction=RotatEInteraction(), diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index ddf47894ed..bed7026498 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -8,8 +8,9 @@ from ..base import TwoSideEmbeddingModel from ...losses import Loss, SoftplusLoss +from ...nn.emb import EmbeddingSpecification from ...nn.modules import DistMultInteraction -from ...regularizers import PowerSumRegularizer, Regularizer +from ...regularizers import PowerSumRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -70,9 +71,9 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, clamp_score: Optional[Union[float, Tuple[float, float]]] = None, ) -> None: + regularizer = PowerSumRegularizer(weight=20, p=2.0, normalize=True) super().__init__( triples_factory=triples_factory, interaction=DistMultInteraction(), @@ -81,6 +82,18 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, + embedding_specification=EmbeddingSpecification( + regularizer=regularizer, + ), + relation_embedding_specification=EmbeddingSpecification( + regularizer=regularizer, + ), + second_embedding_specification=EmbeddingSpecification( + regularizer=regularizer, + ), + second_relation_embedding_specification=EmbeddingSpecification( + regularizer=regularizer, + ), ) if isinstance(clamp_score, float): clamp_score = (-clamp_score, clamp_score) diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index 1070636cec..fae414096e 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -14,7 +14,6 @@ from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import StructuredEmbeddingInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint from ...utils import compose @@ -55,7 +54,6 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: r"""Initialize SE. diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 6b2c825b5b..05c8bf973b 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -9,7 +9,6 @@ from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_normal_ from ...nn.modules import TransDInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint from ...utils import clamp_norm @@ -66,7 +65,6 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: super().__init__( triples_factory=triples_factory, diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 473911a007..5df196eb33 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -11,7 +11,6 @@ from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import TransEInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint from ...utils import compose @@ -56,7 +55,6 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: r"""Initialize TransE. diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 572c000f15..4520e93467 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -10,7 +10,7 @@ from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.modules import TransHInteraction -from ...regularizers import Regularizer, TransHRegularizer +from ...regularizers import TransHRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -72,13 +72,13 @@ def __init__( predict_with_sigmoid: bool = False, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: r"""Initialize TransH. :param embedding_dim: The entity embedding dimension $d$. Is usually $d \in [50, 300]$. :param scoring_fct_norm: The :math:`l_p` norm applied in the interaction function. Is usually ``1`` or ``2.``. """ + # TODO: Regularizer super().__init__( triples_factory=triples_factory, interaction=TransHInteraction( diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 0300a80d88..0f2c93aac1 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -12,7 +12,6 @@ from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import TransRInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint from ...utils import clamp_norm, compose @@ -70,7 +69,6 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model.""" super().__init__( diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index 57c29a0829..45986a4784 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -9,7 +9,6 @@ from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_normal_ from ...nn.modules import TuckerInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -75,7 +74,6 @@ def __init__( dropout_0: float = 0.3, dropout_1: float = 0.4, dropout_2: float = 0.5, - regularizer: Optional[Regularizer] = None, apply_batch_normalization: bool = True, ) -> None: """Initialize the model. diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index deebb82777..0e2c69603d 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -9,7 +9,6 @@ from ...nn import Embedding from ...nn.init import xavier_normal_ from ...nn.modules import UnstructuredModelInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -52,7 +51,6 @@ def __init__( predict_with_sigmoid: bool = False, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: r"""Initialize UM. From 25be75889a8b4d4f9569c5424e9e39698595d3ea Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 18:58:25 +0100 Subject: [PATCH 335/690] Add weight regularizer to ConvKB --- src/pykeen/models/unimodal/conv_kb.py | 2 +- src/pykeen/regularizers.py | 21 ++++++++++++++++----- 2 files changed, 17 insertions(+), 6 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index da12b3aeb3..529654b91e 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -83,7 +83,6 @@ def __init__( To be consistent with the paper, pass entity and relation embeddings pre-trained from TransE. """ - # TODO: regularize weight and bias super().__init__( triples_factory=triples_factory, interaction=ConvKBInteraction( @@ -97,4 +96,5 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, ) + self.regularizer = LpRegularizer(parameters=list(self.interaction.parameters())) logger.warning('To be consistent with the paper, initialize entity and relation embeddings from TransE.') diff --git a/src/pykeen/regularizers.py b/src/pykeen/regularizers.py index 125ef43f65..801ad33aee 100644 --- a/src/pykeen/regularizers.py +++ b/src/pykeen/regularizers.py @@ -3,7 +3,7 @@ """Regularization in PyKEEN.""" from abc import ABC, abstractmethod -from typing import Any, ClassVar, Collection, Iterable, Mapping, Optional, Type, Union +from typing import Any, ClassVar, Collection, Iterable, Mapping, Optional, Sequence, Type, Union import torch from torch import nn @@ -40,15 +40,20 @@ class Regularizer(nn.Module, ABC): #: The default strategy for optimizing the regularizer's hyper-parameters hpo_default: ClassVar[Mapping[str, Any]] + #: weights which should be regularized + parameters: Optional[Collection[nn.Parameter]] = None + def __init__( self, weight: float = 1.0, apply_only_once: bool = False, + parameters: Optional[Sequence[nn.Parameter]] = None, ): super().__init__() self.register_buffer(name='weight', tensor=torch.as_tensor(weight, device=self.device)) self.apply_only_once = apply_only_once self.pop_regularization_term() + self.parameters = parameters @classmethod def get_normalized_name(cls) -> str: @@ -69,6 +74,8 @@ def update(self, *tensors: torch.FloatTensor) -> None: def pop_regularization_term(self) -> torch.FloatTensor: """Return the weighted regularization term, and clear it afterwards.""" + if self.parameters is not None: + self.update(*self.parameters) term = self.regularization_term self.regularization_term = 0. self.updated = False @@ -121,8 +128,9 @@ def __init__( normalize: bool = False, p: float = 2., apply_only_once: bool = False, + parameters: Optional[Sequence[nn.Parameter]] = None, ): - super().__init__(weight=weight, apply_only_once=apply_only_once) + super().__init__(weight=weight, apply_only_once=apply_only_once, parameters=parameters) self.dim = dim self.normalize = normalize self.p = p @@ -160,8 +168,9 @@ def __init__( normalize: bool = False, p: float = 2., apply_only_once: bool = False, + parameters: Optional[Sequence[nn.Parameter]] = None, ): - super().__init__(weight=weight, apply_only_once=apply_only_once) + super().__init__(weight=weight, apply_only_once=apply_only_once, parameters=parameters) self.dim = dim self.normalize = normalize self.p = p @@ -186,10 +195,11 @@ def __init__( self, weight: float = 0.05, epsilon: float = 1e-5, + parameters: Optional[Sequence[nn.Parameter]] = None, ): # The regularization in TransH enforces the defined soft constraints that should computed only for every batch. # Therefore, apply_only_once is always set to True. - super().__init__(weight=weight, apply_only_once=True) + super().__init__(weight=weight, apply_only_once=True, parameters=parameters) self.epsilon = epsilon def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102 @@ -225,8 +235,9 @@ def __init__( device: torch.device, total_weight: float = 1.0, apply_only_once: bool = False, + parameters: Optional[Sequence[nn.Parameter]] = None, ): - super().__init__(weight=total_weight, apply_only_once=apply_only_once) + super().__init__(weight=total_weight, apply_only_once=apply_only_once, parameters=parameters) self.regularizers = nn.ModuleList(regularizers) for r in self.regularizers: if isinstance(r, NoRegularizer): From 2e5042c20f1aac451731b135a36e855292a18b4b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 19:03:02 +0100 Subject: [PATCH 336/690] Add TransH regularizer --- src/pykeen/models/unimodal/trans_h.py | 8 +++++++- src/pykeen/regularizers.py | 25 ++++++++++++------------- 2 files changed, 19 insertions(+), 14 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 4520e93467..74864b81d7 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -78,7 +78,6 @@ def __init__( :param embedding_dim: The entity embedding dimension $d$. Is usually $d \in [50, 300]$. :param scoring_fct_norm: The :math:`l_p` norm applied in the interaction function. Is usually ``1`` or ``2.``. """ - # TODO: Regularizer super().__init__( triples_factory=triples_factory, interaction=TransHInteraction( @@ -98,3 +97,10 @@ def __init__( constrainer=functional.normalize, ), ) + self.regularizer = TransHRegularizer( + weight=0.05, + epsilon=1e-5, + entity_embeddings=list(self.entity_representations[0].parameters()).pop(), + relation_embeddings=list(self.relation_representations[0].parameters()).pop(), + normal_vector_embeddings=list(self.relation_representations[1].parameters()).pop(), + ) diff --git a/src/pykeen/regularizers.py b/src/pykeen/regularizers.py index 801ad33aee..779fb61a01 100644 --- a/src/pykeen/regularizers.py +++ b/src/pykeen/regularizers.py @@ -193,34 +193,33 @@ class TransHRegularizer(Regularizer): def __init__( self, + entity_embeddings: nn.Parameter, + normal_vector_embeddings: nn.Parameter, + relation_embeddings: nn.Parameter, weight: float = 0.05, epsilon: float = 1e-5, - parameters: Optional[Sequence[nn.Parameter]] = None, ): # The regularization in TransH enforces the defined soft constraints that should computed only for every batch. # Therefore, apply_only_once is always set to True. - super().__init__(weight=weight, apply_only_once=True, parameters=parameters) + super().__init__(weight=weight, apply_only_once=False, parameters=[]) + self.normal_vector_embeddings = normal_vector_embeddings + self.relation_embeddings = relation_embeddings + self.entity_embeddings = entity_embeddings self.epsilon = epsilon def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102 raise NotImplementedError('TransH regularizer is order-sensitive!') - def update(self, *tensors: torch.FloatTensor) -> None: # noqa: D102 - if len(tensors) != 3: - raise KeyError('Expects exactly three tensors') - if self.apply_only_once and self.updated: - return - entity_embeddings, normal_vector_embeddings, relation_embeddings = tensors + def pop_regularization_term(self, *tensors: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102 # Entity soft constraint - self.regularization_term += torch.sum(functional.relu(torch.norm(entity_embeddings, dim=-1) ** 2 - 1.0)) + self.regularization_term += torch.sum(functional.relu(torch.norm(self.entity_embeddings, dim=-1) ** 2 - 1.0)) # Orthogonality soft constraint - d_r_n = functional.normalize(relation_embeddings, dim=-1) + d_r_n = functional.normalize(self.relation_embeddings, dim=-1) self.regularization_term += torch.sum( - functional.relu(torch.sum((normal_vector_embeddings * d_r_n) ** 2, dim=-1) - self.epsilon), + functional.relu(torch.sum((self.normal_vector_embeddings * d_r_n) ** 2, dim=-1) - self.epsilon), ) - - self.updated = True + return super().pop_regularization_term() class CombinedRegularizer(Regularizer): From a0141143ed7cf917bd90c64643620954555ab3d0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 16 Nov 2020 19:28:40 +0100 Subject: [PATCH 337/690] Fix tests --- src/pykeen/nn/modules.py | 8 ++-- src/pykeen/regularizers.py | 41 +++++++++----------- tests/test_regularizers.py | 77 +++++++++++++++----------------------- 3 files changed, 51 insertions(+), 75 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 5e5a6b5839..be7b9b2852 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -114,7 +114,7 @@ def forward( raise NotImplementedError @staticmethod - def _add_dim(*x: torch.FloatTensor, dim: int) -> Union[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor, ...]]: + def _add_dim(*x: torch.FloatTensor, dim: int) -> Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]: """ Add a dimension to tensors. @@ -270,9 +270,9 @@ def _score( def _forward_slicing_wrapper( self, - h: Union[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor, ...]], - r: Union[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor, ...]], - t: Union[torch.FloatTensor, Tuple[torch.FloatTensor, torch.FloatTensor, ...]], + h: Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]], + r: Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]], + t: Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]], slice_size: Optional[int], slice_dim: str, ) -> torch.FloatTensor: diff --git a/src/pykeen/regularizers.py b/src/pykeen/regularizers.py index 779fb61a01..bdf7798064 100644 --- a/src/pykeen/regularizers.py +++ b/src/pykeen/regularizers.py @@ -41,7 +41,7 @@ class Regularizer(nn.Module, ABC): hpo_default: ClassVar[Mapping[str, Any]] #: weights which should be regularized - parameters: Optional[Collection[nn.Parameter]] = None + tracked_parameters: Optional[Collection[nn.Parameter]] = None def __init__( self, @@ -50,10 +50,14 @@ def __init__( parameters: Optional[Sequence[nn.Parameter]] = None, ): super().__init__() - self.register_buffer(name='weight', tensor=torch.as_tensor(weight, device=self.device)) + self.register_buffer(name='weight', tensor=torch.as_tensor(weight)) self.apply_only_once = apply_only_once - self.pop_regularization_term() - self.parameters = parameters + self.tracked_parameters = parameters + self._clear() + + def _clear(self): + self.regularization_term = 0. + self.updated = False @classmethod def get_normalized_name(cls) -> str: @@ -65,42 +69,33 @@ def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: """Compute the regularization term for one tensor.""" raise NotImplementedError - def update(self, *tensors: torch.FloatTensor) -> None: + def update(self, *tensors: torch.FloatTensor) -> bool: """Update the regularization term based on passed tensors.""" if not self.training or not torch.is_grad_enabled() or (self.apply_only_once and self.updated): - return + return False self.regularization_term = self.regularization_term + sum(self(x) for x in tensors) self.updated = True + return True def pop_regularization_term(self) -> torch.FloatTensor: """Return the weighted regularization term, and clear it afterwards.""" - if self.parameters is not None: - self.update(*self.parameters) + if self.tracked_parameters is not None: + self.update(*self.tracked_parameters) term = self.regularization_term - self.regularization_term = 0. - self.updated = False + self._clear() return self.weight * term - @property - def term(self) -> torch.FloatTensor: - """Return the weighted regularization term.""" - # TODO: Deprecated - raise RuntimeError - class NoRegularizer(Regularizer): """A regularizer which does not perform any regularization. Used to simplify code. """ + # TODO: Deprecated #: The default strategy for optimizing the regularizer's hyper-parameters hpo_default: ClassVar[Mapping[str, Any]] = {} - def update(self, *tensors: torch.FloatTensor) -> None: # noqa: D102 - # no need to compute anything - pass - def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102 # always return zero return x.new_zeros(1) @@ -231,7 +226,6 @@ class CombinedRegularizer(Regularizer): def __init__( self, regularizers: Iterable[Regularizer], - device: torch.device, total_weight: float = 1.0, apply_only_once: bool = False, parameters: Optional[Sequence[nn.Parameter]] = None, @@ -241,9 +235,8 @@ def __init__( for r in self.regularizers: if isinstance(r, NoRegularizer): raise TypeError('Can not combine a no-op regularizer') - self.register_buffer(name='normalization_factor', tensor=torch.as_tensor( - sum(r.weight for r in self.regularizers), device=device, - ).reciprocal()) + normalization_factor = torch.as_tensor(sum(r.weight for r in self.regularizers)).reciprocal() + self.register_buffer(name='normalization_factor', tensor=normalization_factor) @property def normalize(self): # noqa: D102 diff --git a/tests/test_regularizers.py b/tests/test_regularizers.py index 2075aed461..2f861cb99b 100644 --- a/tests/test_regularizers.py +++ b/tests/test_regularizers.py @@ -7,13 +7,14 @@ from typing import Any, ClassVar, Dict, Optional, Type import torch +from torch import nn from torch.nn import functional from pykeen.datasets import Nations -from pykeen.models import ConvKB, RESCAL, TransH +from pykeen.models import ConvKB, RESCAL from pykeen.regularizers import ( CombinedRegularizer, LpRegularizer, NoRegularizer, PowerSumRegularizer, Regularizer, - TransHRegularizer, + TransHRegularizer, collect_regularization_terms, ) from pykeen.triples import TriplesFactory from pykeen.typing import MappedTriples @@ -45,9 +46,8 @@ def setUp(self) -> None: self.triples_factory = Nations().training self.device = resolve_device() self.regularizer = self.regularizer_cls( - device=self.device, **(self.regularizer_kwargs or {}), - ) + ).to(self.device) self.positive_batch = self.triples_factory.mapped_triples[:self.batch_size, :].to(device=self.device) def test_model(self) -> None: @@ -57,17 +57,20 @@ def test_model(self) -> None: triples_factory=self.triples_factory, ).to(self.device) - # Check if regularizer is stored correctly. - self.assertEqual(model.regularizer, self.regularizer) + # check for regularizer + assert sum(1 for m in model.modules() if isinstance(m, Regularizer)) > 0 # Forward pass (should update regularizer) model.score_hrt(hrt_batch=self.positive_batch) - # Call post_parameter_update (should reset regularizer) - model.post_parameter_update() + # check that regularization term is accessible + term = collect_regularization_terms(model) + assert torch.is_tensor(term) + assert term.requires_grad - # Check if regularization term is reset - self.assertEqual(0., model.regularizer.term) + # second time should be 0. + term = collect_regularization_terms(model) + assert term == 0. def test_update(self) -> None: """Test method `update`.""" @@ -76,17 +79,18 @@ def test_update(self) -> None: b = torch.rand(self.batch_size, 20, device=self.device, generator=self.generator) # Call update - self.regularizer.update(a, b) + assert self.regularizer.update(a, b) # check shape - self.assertEqual((1,), self.regularizer.term.shape) + assert 1 == self.regularizer.regularization_term.numel() # compute expected term exp_penalties = torch.stack([self._expected_penalty(x) for x in (a, b)]) expected_term = torch.sum(exp_penalties).view(1) * self.regularizer.weight assert expected_term.shape == (1,) - self.assertAlmostEqual(self.regularizer.term.item(), expected_term.item()) + observed = self.regularizer.pop_regularization_term() + self.assertAlmostEqual(observed.item(), expected_term.item()) def test_forward(self) -> None: """Test the regularizer's `forward` method.""" @@ -202,39 +206,26 @@ def setUp(self) -> None: """Set up the test case.""" self.generator = torch.random.manual_seed(seed=42) self.device = resolve_device() + self.num_entities = 10 + self.num_relations = 5 self.regularizer_kwargs = {'weight': .5, 'epsilon': 1e-5} + self.entities_weight = nn.Parameter(torch.rand(self.num_entities, 10, device=self.device, generator=self.generator)) + self.relations_weight = nn.Parameter(torch.rand(self.num_relations, 20, device=self.device, generator=self.generator)) + self.normal_vector_weight = nn.Parameter(torch.rand(self.num_relations, 20, device=self.device, generator=self.generator)) + self.regularizer_kwargs["entity_embeddings"] = self.entities_weight + self.regularizer_kwargs["normal_vector_embeddings"] = self.normal_vector_weight + self.regularizer_kwargs["relation_embeddings"] = self.relations_weight self.regularizer = TransHRegularizer( **(self.regularizer_kwargs or {}), ) - self.num_entities = 10 - self.num_relations = 5 - self.entities_weight = torch.rand(self.num_entities, 10, device=self.device, generator=self.generator) - self.relations_weight = torch.rand(self.num_relations, 20, device=self.device, generator=self.generator) - self.normal_vector_weight = torch.rand(self.num_relations, 20, device=self.device, generator=self.generator) def test_update(self): """Test update function of TransHRegularizer.""" - # Tests that exception will be thrown when more than or less than three tensors are passed - with self.assertRaises(KeyError) as context: - self.regularizer.update( - self.entities_weight, - self.normal_vector_weight, - self.relations_weight, - torch.rand(self.num_entities, 10, device=self.device, generator=self.generator), - ) - self.assertTrue('Expects exactly three tensors' in context.exception) - - self.regularizer.update( - self.entities_weight, - self.normal_vector_weight, - ) - self.assertTrue('Expects exactly three tensors' in context.exception) - # Test that regularization term is computed correctly - self.regularizer.update(self.entities_weight, self.normal_vector_weight, self.relations_weight) expected_term = self._expected_penalty() weight = self.regularizer_kwargs.get('weight') - self.assertAlmostEqual(self.regularizer.term.item(), weight * expected_term.item()) + observed_term = self.regularizer.pop_regularization_term() + assert torch.allclose(observed_term, weight * expected_term) def _expected_penalty(self) -> torch.FloatTensor: # noqa: D102 # Entity soft constraint @@ -270,17 +261,9 @@ def test_lp(self): ) self._help_test_regularizer(regularizer) - def test_transh_regularizer(self): - """Test the TransH regularizer only updates once.""" - self.assertNotIn('apply_only_once', TransH.regularizer_default_kwargs) - regularizer = TransHRegularizer( - **TransH.regularizer_default_kwargs, - ) - self._help_test_regularizer(regularizer) - def _help_test_regularizer(self, regularizer: Regularizer, n_tensors: int = 3): self.assertFalse(regularizer.updated) - self.assertEqual(0.0, regularizer.regularization_term.item()) + assert 0.0 == regularizer.regularization_term # After first update, should change the term first_tensors = [ @@ -301,6 +284,6 @@ def _help_test_regularizer(self, regularizer: Regularizer, n_tensors: int = 3): self.assertTrue(regularizer.updated) self.assertEqual(term, regularizer.regularization_term) - # regularizer.reset() + regularizer.pop_regularization_term() self.assertFalse(regularizer.updated) - self.assertEqual(0.0, regularizer.regularization_term.item()) + assert 0.0 == regularizer.regularization_term From 24c9b588e7fb485c3b1c97787e042416d0959e50 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 19:42:09 +0100 Subject: [PATCH 338/690] Pass flake8 --- src/pykeen/models/unimodal/conv_e.py | 1 - src/pykeen/models/unimodal/kg2e.py | 1 - src/pykeen/regularizers.py | 1 + tests/test_regularizers.py | 9 ++++++--- 4 files changed, 7 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index aac26a4787..feae6e6bee 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -14,7 +14,6 @@ from ...nn.emb import EmbeddingSpecification from ...nn.init import xavier_normal_ from ...nn.modules import ConvEInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 35312bd72f..153c4e82fd 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -12,7 +12,6 @@ from ...losses import Loss from ...nn.emb import EmbeddingSpecification from ...nn.modules import KG2EInteraction -from ...regularizers import Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint from ...utils import clamp_norm diff --git a/src/pykeen/regularizers.py b/src/pykeen/regularizers.py index bdf7798064..c8012acde8 100644 --- a/src/pykeen/regularizers.py +++ b/src/pykeen/regularizers.py @@ -91,6 +91,7 @@ class NoRegularizer(Regularizer): Used to simplify code. """ + # TODO: Deprecated #: The default strategy for optimizing the regularizer's hyper-parameters diff --git a/tests/test_regularizers.py b/tests/test_regularizers.py index 2f861cb99b..3a0c10eb19 100644 --- a/tests/test_regularizers.py +++ b/tests/test_regularizers.py @@ -209,9 +209,9 @@ def setUp(self) -> None: self.num_entities = 10 self.num_relations = 5 self.regularizer_kwargs = {'weight': .5, 'epsilon': 1e-5} - self.entities_weight = nn.Parameter(torch.rand(self.num_entities, 10, device=self.device, generator=self.generator)) - self.relations_weight = nn.Parameter(torch.rand(self.num_relations, 20, device=self.device, generator=self.generator)) - self.normal_vector_weight = nn.Parameter(torch.rand(self.num_relations, 20, device=self.device, generator=self.generator)) + self.entities_weight = self._rand_param(10) + self.relations_weight = self._rand_param(20) + self.normal_vector_weight = self._rand_param(20) self.regularizer_kwargs["entity_embeddings"] = self.entities_weight self.regularizer_kwargs["normal_vector_embeddings"] = self.normal_vector_weight self.regularizer_kwargs["relation_embeddings"] = self.relations_weight @@ -219,6 +219,9 @@ def setUp(self) -> None: **(self.regularizer_kwargs or {}), ) + def _rand_param(self, n): + return nn.Parameter(torch.rand(self.num_entities, 10, device=self.device, generator=self.generator)) + def test_update(self): """Test update function of TransHRegularizer.""" # Test that regularization term is computed correctly From 7834dce28d7193efc18d7f45a690971f4835d35a Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 19:44:58 +0100 Subject: [PATCH 339/690] Pass mypy --- src/pykeen/nn/emb.py | 4 ++-- src/pykeen/nn/modules.py | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index 65b138c3d3..101412918c 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -179,12 +179,12 @@ def __init__( self.initializer = initializer if constrainer is not None and constrainer_kwargs: - constrainer: Constrainer = functools.partial(constrainer, **constrainer_kwargs) + constrainer = functools.partial(constrainer, **constrainer_kwargs) self.constrainer = constrainer # TODO: Move regularizer and normalizer to RepresentationModule? if normalizer is not None and normalizer_kwargs: - normalizer: Normalizer = functools.partial(normalizer, **normalizer_kwargs) + normalizer = functools.partial(normalizer, **normalizer_kwargs) self.normalizer = normalizer self.regularizer = regularizer diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index be7b9b2852..9f2f2e1888 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -274,7 +274,7 @@ def _forward_slicing_wrapper( r: Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]], t: Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]], slice_size: Optional[int], - slice_dim: str, + slice_dim: Optional[str], ) -> torch.FloatTensor: """ Compute broadcasted triple scores with optional slicing for representations in canonical shape. From 9bf1d047f03f808234e2457b2bea98493d338209 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 16 Nov 2020 20:14:39 +0100 Subject: [PATCH 340/690] Clean up utilities, imports, and add notes for dead code --- src/pykeen/models/base.py | 3 +- .../models/multimodal/complex_literal.py | 4 +- .../models/multimodal/distmult_literal.py | 4 +- src/pykeen/models/unimodal/complex.py | 2 +- src/pykeen/models/unimodal/conv_e.py | 3 +- src/pykeen/models/unimodal/distmult.py | 2 +- src/pykeen/models/unimodal/hole.py | 2 +- src/pykeen/models/unimodal/kg2e.py | 4 +- src/pykeen/models/unimodal/proj_e.py | 2 +- src/pykeen/models/unimodal/rescal.py | 2 +- src/pykeen/models/unimodal/rotate.py | 40 +-- src/pykeen/models/unimodal/simple.py | 2 +- .../models/unimodal/structured_embedding.py | 2 +- src/pykeen/models/unimodal/trans_d.py | 2 +- src/pykeen/models/unimodal/trans_e.py | 2 +- src/pykeen/models/unimodal/trans_h.py | 2 +- src/pykeen/models/unimodal/trans_r.py | 4 +- src/pykeen/models/unimodal/tucker.py | 2 +- src/pykeen/nn/__init__.py | 4 +- src/pykeen/nn/functional.py | 298 ++---------------- src/pykeen/nn/init.py | 9 + src/pykeen/nn/modules.py | 1 + src/pykeen/nn/sim.py | 134 ++++++++ src/pykeen/typing.py | 10 +- src/pykeen/utils.py | 169 ++++++++-- tests/test_models.py | 5 +- 26 files changed, 355 insertions(+), 359 deletions(-) create mode 100644 src/pykeen/nn/sim.py diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 137ce63464..2170908298 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -20,8 +20,7 @@ from torch import nn from ..losses import Loss, MarginRankingLoss, NSSALoss -from ..nn import Embedding, RepresentationModule -from ..nn.emb import EmbeddingSpecification, LiteralRepresentations +from ..nn import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule from ..nn.modules import Interaction from ..regularizers import collect_regularization_terms from ..triples import TriplesFactory, TriplesNumericLiteralsFactory diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index 37d670c52a..dfe74a90f7 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -10,9 +10,8 @@ from ..base import LiteralModel from ...losses import BCEWithLogitsLoss, Loss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.modules import ComplExInteraction -from ...regularizers import Regularizer from ...triples import TriplesNumericLiteralsFactory from ...typing import DeviceHint from ...utils import combine_complex, split_complex @@ -84,7 +83,6 @@ def __init__( predict_with_sigmoid: bool = False, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - regularizer: Optional[Regularizer] = None, ) -> None: """Initialize the model.""" super().__init__( diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 327d30993a..27fe67df09 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -9,9 +9,8 @@ from ..base import LiteralModel from ...losses import Loss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.modules import DistMultInteraction -from ...regularizers import Regularizer from ...triples import TriplesNumericLiteralsFactory from ...typing import DeviceHint @@ -46,7 +45,6 @@ def __init__( preferred_device: DeviceHint = None, random_seed: Optional[int] = None, predict_with_sigmoid: bool = False, - regularizer: Optional[Regularizer] = None, ) -> None: super().__init__( triples_factory=triples_factory, diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 1d853ce6a7..a4ee8c43f0 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -8,7 +8,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss, SoftplusLoss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.modules import ComplExInteraction from ...regularizers import LpRegularizer from ...triples import TriplesFactory diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index feae6e6bee..05b0f66523 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -10,8 +10,7 @@ from .. import ERModel from ...losses import BCEAfterSigmoidLoss, Loss -from ...nn import Embedding -from ...nn.emb import EmbeddingSpecification +from ...nn import Embedding, EmbeddingSpecification from ...nn.init import xavier_normal_ from ...nn.modules import ConvEInteraction from ...triples import TriplesFactory diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 6b1c188089..bcc4c491b6 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -9,7 +9,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.modules import DistMultInteraction from ...regularizers import LpRegularizer from ...triples import TriplesFactory diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index b041d030a2..40233db82e 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -6,7 +6,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import HolEInteraction from ...triples import TriplesFactory diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 153c4e82fd..b73113b34e 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -10,7 +10,7 @@ from ..base import TwoVectorEmbeddingModel from ...losses import Loss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.modules import KG2EInteraction from ...triples import TriplesFactory from ...typing import DeviceHint @@ -20,7 +20,7 @@ 'KG2E', ] -_LOG_2_PI = math.log(2. * math.pi) +_LOG_2_PI = math.log(math.tau) class KG2E(TwoVectorEmbeddingModel): diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index d75c526ffa..3c7ef1a486 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -8,7 +8,7 @@ from .. import SingleVectorEmbeddingModel from ...losses import Loss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import ProjEInteraction from ...triples import TriplesFactory diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index f9070dbaff..836050e46a 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -6,7 +6,7 @@ from ..base import SingleVectorEmbeddingModel from ...losses import Loss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.modules import RESCALInteraction from ...regularizers import LpRegularizer from ...triples import TriplesFactory diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index 4759f53f83..56908e00bd 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -4,54 +4,20 @@ from typing import Optional -import numpy as np -import torch -import torch.autograd -from torch.nn import functional - from .. import SingleVectorEmbeddingModel from ...losses import Loss -from ...nn.emb import EmbeddingSpecification -from ...nn.init import xavier_uniform_ +from ...nn import EmbeddingSpecification +from ...nn.init import init_phases, xavier_uniform_ from ...nn.modules import RotatEInteraction from ...triples import TriplesFactory from ...typing import DeviceHint +from ...utils import complex_normalize __all__ = [ 'RotatE', ] -def init_phases(x: torch.Tensor) -> torch.Tensor: - r"""Generate random phases between 0 and :math:`2\pi`.""" - phases = 2 * np.pi * torch.rand_like(x[..., :x.shape[-1] // 2]) - return torch.cat([torch.cos(phases), torch.sin(phases)], dim=-1).detach() - - -def complex_normalize(x: torch.Tensor) -> torch.Tensor: - r"""Normalize the length of relation vectors, if the forward constraint has not been applied yet. - - The `modulus of complex number `_ is given as: - - .. math:: - - |a + ib| = \sqrt{a^2 + b^2} - - $l_2$ norm of complex vector $x \in \mathbb{C}^d$: - - .. math:: - \|x\|^2 = \sum_{i=1}^d |x_i|^2 - = \sum_{i=1}^d \left(\operatorname{Re}(x_i)^2 + \operatorname{Im}(x_i)^2\right) - = \left(\sum_{i=1}^d \operatorname{Re}(x_i)^2) + (\sum_{i=1}^d \operatorname{Im}(x_i)^2\right) - = \|\operatorname{Re}(x)\|^2 + \|\operatorname{Im}(x)\|^2 - = \| [\operatorname{Re}(x); \operatorname{Im}(x)] \|^2 - """ - y = x.data.view(x.shape[0], -1, 2) - y = functional.normalize(y, p=2, dim=-1) - x.data = y.view(*x.shape) - return x - - class RotatE(SingleVectorEmbeddingModel): r"""An implementation of RotatE from [sun2019]_. diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index bed7026498..8c07da910c 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -8,7 +8,7 @@ from ..base import TwoSideEmbeddingModel from ...losses import Loss, SoftplusLoss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.modules import DistMultInteraction from ...regularizers import PowerSumRegularizer from ...triples import TriplesFactory diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index fae414096e..d3abc3b2a9 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -11,7 +11,7 @@ from ..base import DoubleRelationEmbeddingModel from ...losses import Loss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import StructuredEmbeddingInteraction from ...triples import TriplesFactory diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 05c8bf973b..bd8a4d2a1c 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -6,7 +6,7 @@ from ..base import TwoVectorEmbeddingModel from ...losses import Loss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.init import xavier_normal_ from ...nn.modules import TransDInteraction from ...triples import TriplesFactory diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 5df196eb33..742e7f51d3 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -8,7 +8,7 @@ from .. import SingleVectorEmbeddingModel from ...losses import Loss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import TransEInteraction from ...triples import TriplesFactory diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 74864b81d7..6c0e8d3ba7 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -8,7 +8,7 @@ from ..base import DoubleRelationEmbeddingModel from ...losses import Loss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.modules import TransHInteraction from ...regularizers import TransHRegularizer from ...triples import TriplesFactory diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 0f2c93aac1..f5a484bb25 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- """Implementation of TransR.""" + import logging from typing import Optional @@ -8,8 +9,7 @@ from .. import ERModel from ...losses import Loss -from ...nn import Embedding -from ...nn.emb import EmbeddingSpecification +from ...nn import Embedding, EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import TransRInteraction from ...triples import TriplesFactory diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index 45986a4784..739d7ff529 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -6,7 +6,7 @@ from .. import SingleVectorEmbeddingModel from ...losses import BCEAfterSigmoidLoss, Loss -from ...nn.emb import EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.init import xavier_normal_ from ...nn.modules import TuckerInteraction from ...triples import TriplesFactory diff --git a/src/pykeen/nn/__init__.py b/src/pykeen/nn/__init__.py index 7de87500b3..eeef54bdb2 100644 --- a/src/pykeen/nn/__init__.py +++ b/src/pykeen/nn/__init__.py @@ -3,11 +3,13 @@ """PyKEEN internal "nn" module.""" from . import functional, init -from .emb import Embedding, RepresentationModule +from .emb import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule from .modules import Interaction, StatelessInteraction __all__ = [ 'Embedding', + 'EmbeddingSpecification', + 'LiteralRepresentations', 'RepresentationModule', 'Interaction', 'StatelessInteraction', diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index bf8df457fe..b7ba59081e 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -2,14 +2,18 @@ """Functional forms of interaction methods.""" -import math -from typing import NamedTuple, Optional, SupportsFloat, Tuple, Union +from typing import Optional, Tuple, Union import torch import torch.fft from torch import nn -from ..utils import broadcast_cat, clamp_norm, is_cudnn_error, split_complex +from .sim import KG2E_SIMILARITIES +from ..typing import GaussianDistribution +from ..utils import ( + broadcast_cat, clamp_norm, extended_einsum, is_cudnn_error, negative_norm_of_sum, project_entity, + split_complex, view_complex, +) __all__ = [ "complex_interaction", @@ -35,6 +39,7 @@ ] +# TODO @mberr documentation def _extract_sizes(h, r, t) -> Tuple[int, int, int, int, int]: num_heads, num_relations, num_tails = [xx.shape[1] for xx in (h, r, t)] d_e = h.shape[-1] @@ -42,44 +47,7 @@ def _extract_sizes(h, r, t) -> Tuple[int, int, int, int, int]: return num_heads, num_relations, num_tails, d_e, d_r -def _negative_norm_of_sum( - *x: torch.FloatTensor, - p: Union[str, int] = 2, - power_norm: bool = False, -) -> torch.FloatTensor: - """ - Evaluate negative norm of a sum of vectors on already broadcasted representations. - - :param x: shape: (batch_size, num_heads, num_relations, num_tails, dim) - The representations. - :param p: - The p for the norm. cf. torch.norm. - :param power_norm: - Whether to return $|x-y|_p^p$, cf. https://github.com/pytorch/pytorch/issues/28119 - - :return: shape: (batch_size, num_heads, num_relations, num_tails) - The scores. - """ - d: torch.FloatTensor = sum(x) - if power_norm: - assert isinstance(p, SupportsFloat) - return -(d.abs() ** p).sum(dim=-1) - else: - if torch.is_complex(d): - assert isinstance(p, SupportsFloat) - # workaround for complex numbers: manually compute norm - return -(d.abs() ** p).sum(dim=-1) ** (1 / p) - else: - return -d.norm(p=p, dim=-1) - - -class GaussianDistribution(NamedTuple): - """A gaussian distribution with diagonal covariance matrix.""" - - mean: torch.FloatTensor - diagonal_covariance: torch.FloatTensor - - +# TODO @mberr documentation def _apply_optional_bn_to_tensor( batch_norm: Optional[nn.BatchNorm1d], output_dropout: nn.Dropout, @@ -94,206 +62,6 @@ def _apply_optional_bn_to_tensor( return tensor -def _project_entity( - e: torch.FloatTensor, - e_p: torch.FloatTensor, - r_p: torch.FloatTensor, -) -> torch.FloatTensor: - r"""Project entity relation-specific. - - .. math:: - - e_{\bot} = M_{re} e - = (r_p e_p^T + I^{d_r \times d_e}) e - = r_p e_p^T e + I^{d_r \times d_e} e - = r_p (e_p^T e) + e' - - and additionally enforces - - .. math:: - - \|e_{\bot}\|_2 \leq 1 - - :param e: shape: (..., d_e) - The entity embedding. - :param e_p: shape: (..., d_e) - The entity projection. - :param r_p: shape: (..., d_r) - The relation projection. - - :return: shape: (..., d_r) - - """ - # The dimensions affected by e' - change_dim = min(e.shape[-1], r_p.shape[-1]) - - # Project entities - # r_p (e_p.T e) + e' - e_bot = r_p * torch.sum(e_p * e, dim=-1, keepdim=True) - e_bot[..., :change_dim] += e[..., :change_dim] - - # Enforce constraints - e_bot = clamp_norm(e_bot, p=2, dim=-1, maxnorm=1) - - return e_bot - - -def _expected_likelihood( - e: GaussianDistribution, - r: GaussianDistribution, - epsilon: float = 1.0e-10, - exact: bool = True, -) -> torch.FloatTensor: - r"""Compute the similarity based on expected likelihood. - - .. math:: - - D((\mu_e, \Sigma_e), (\mu_r, \Sigma_r))) - = \frac{1}{2} \left( - (\mu_e - \mu_r)^T(\Sigma_e + \Sigma_r)^{-1}(\mu_e - \mu_r) - + \log \det (\Sigma_e + \Sigma_r) + d \log (2 \pi) - \right) - = \frac{1}{2} \left( - \mu^T\Sigma^{-1}\mu - + \log \det \Sigma + d \log (2 \pi) - \right) - - :param e: shape: (batch_size, num_heads, num_tails, d) - The entity Gaussian distribution. - :param r: shape: (batch_size, num_relations, d) - The relation Gaussian distribution. - :param epsilon: float (default=1.0) - Small constant used to avoid numerical issues when dividing. - :param exact: - Whether to return the exact similarity, or leave out constant offsets. - - :return: torch.Tensor, shape: (s_1, ..., s_k) - The similarity. - """ - # subtract, shape: (batch_size, num_heads, num_relations, num_tails, dim) - r_shape = r.mean.shape - r_shape = (r_shape[0], 1, r_shape[1], 1, r_shape[2]) - var = r.diagonal_covariance.view(*r_shape) + e.diagonal_covariance.unsqueeze(dim=2) - mean = e.mean.unsqueeze(dim=2) - r.mean.view(*r_shape) - - #: a = \mu^T\Sigma^{-1}\mu - safe_sigma = torch.clamp_min(var, min=epsilon) - sigma_inv = torch.reciprocal(safe_sigma) - sim = torch.sum(sigma_inv * mean ** 2, dim=-1) - - #: b = \log \det \Sigma - sim = sim + safe_sigma.log().sum(dim=-1) - if exact: - sim = sim + sim.shape[-1] * math.log(2. * math.pi) - return sim - - -def _kullback_leibler_similarity( - e: GaussianDistribution, - r: GaussianDistribution, - epsilon: float = 1.0e-10, - exact: bool = True, -) -> torch.FloatTensor: - r"""Compute the similarity based on KL divergence. - - This is done between two Gaussian distributions given by mean `mu_*` and diagonal covariance matrix `sigma_*`. - - .. math:: - - D((\mu_e, \Sigma_e), (\mu_r, \Sigma_r))) - = \frac{1}{2} \left( - tr(\Sigma_r^{-1}\Sigma_e) - + (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) - - \log \frac{det(\Sigma_e)}{det(\Sigma_r)} - k_e - \right) - - Note: The sign of the function has been flipped as opposed to the description in the paper, as the - Kullback Leibler divergence is large if the distributions are dissimilar. - - :param e: shape: (batch_size, num_heads, num_tails, d) - The entity Gaussian distributions, as mean/diagonal covariance pairs. - :param r: shape: (batch_size, num_relations, d) - The relation Gaussian distributions, as mean/diagonal covariance pairs. - :param epsilon: float (default=1.0) - Small constant used to avoid numerical issues when dividing. - :param exact: - Whether to return the exact similarity, or leave out constant offsets. - - :return: torch.Tensor, shape: (s_1, ..., s_k) - The similarity. - """ - # invert covariance, shape: (batch_size, num_relations, d) - safe_sigma_r = torch.clamp_min(r.diagonal_covariance, min=epsilon) - sigma_r_inv = torch.reciprocal(safe_sigma_r) - - #: a = tr(\Sigma_r^{-1}\Sigma_e), (batch_size, num_heads, num_relations, num_tails) - # [(b, h, t, d), (b, r, d) -> (b, 1, r, d) -> (b, 1, d, r)] -> (b, h, t, r) -> (b, h, r, t) - sim = (e.diagonal_covariance @ sigma_r_inv.unsqueeze(dim=1).transpose(-2, -1)).transpose(-2, -1) - - #: b = (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) - r_shape = r.mean.shape - # mu.shape: (b, h, r, t, d) - mu = r.mean.view(r_shape[0], 1, r_shape[1], 1, r_shape[2]) - e.mean.unsqueeze(dim=2) - sim = sim + (mu ** 2 @ sigma_r_inv.view(r_shape[0], 1, r_shape[1], r_shape[2], 1)).squeeze(dim=-1) - - #: c = \log \frac{det(\Sigma_e)}{det(\Sigma_r)} - # = sum log (sigma_e)_i - sum log (sigma_r)_i - # ce.shape: (b, h, t) - ce = e.diagonal_covariance.clamp_min(min=epsilon).log().sum(dim=-1) - # cr.shape: (b, r) - cr = safe_sigma_r.log().sum(dim=-1) - sim = sim + ce.unsqueeze(dim=2) - cr.view(r_shape[0], 1, r_shape[1], 1) - - if exact: - sim = sim - e.mean.shape[-1] - sim = 0.5 * sim - - return sim - - -KG2E_SIMILARITIES = { - 'KL': _kullback_leibler_similarity, - 'EL': _expected_likelihood, -} - - -def _extended_einsum( - eq: str, - *tensors, -) -> torch.FloatTensor: - """Drop dimensions of size 1 to allow broadcasting.""" - # TODO: check if einsum is still very slow. - lhs, rhs = eq.split("->") - mod_ops, mod_t = [], [] - for op, t in zip(lhs.split(","), tensors): - mod_op = "" - assert len(op) == len(t.shape) - for i, c in reversed(list(enumerate(op))): - if t.shape[i] == 1: - t = t.squeeze(dim=i) - else: - mod_op = c + mod_op - mod_ops.append(mod_op) - mod_t.append(t) - m_lhs = ",".join(mod_ops) - r_keep_dims = set("".join(mod_ops)) - m_rhs = "".join(c for c in rhs if c in r_keep_dims) - m_eq = f"{m_lhs}->{m_rhs}" - mod_r = torch.einsum(m_eq, *mod_t) - # unsqueeze - for i, c in enumerate(rhs): - if c not in r_keep_dims: - mod_r = mod_r.unsqueeze(dim=i) - return mod_r - - -def _view_complex( - x: torch.FloatTensor, -) -> torch.Tensor: - real, imag = split_complex(x=x) - return torch.complex(real=real, imag=imag) - - def _translational_interaction( h: torch.FloatTensor, r: torch.FloatTensor, @@ -318,7 +86,7 @@ def _translational_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return _negative_norm_of_sum(h, r, -t, p=p, power_norm=power_norm) + return negative_norm_of_sum(h, r, -t, p=p, power_norm=power_norm) def _add_cuda_warning(func): @@ -358,7 +126,7 @@ def complex_interaction( """ (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] return sum( - _extended_einsum("bhd,brd,btd->bhrt", hh, rr, tt) + extended_einsum("bhd,brd,btd->bhrt", hh, rr, tt) for hh, rr, tt in [ (h_re, r_re, t_re), (h_re, r_im, t_im), @@ -559,7 +327,7 @@ def distmult_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return _extended_einsum("bhd,brd,btd->bhrt", h, r, t) + return extended_einsum("bhd,brd,btd->bhrt", h, r, t) def ermlp_interaction( @@ -670,7 +438,7 @@ def hole_interaction( composite = torch.fft.irfft(p_fft, n=h.shape[-1], dim=-1) # inner product with relation embedding - return _extended_einsum("bhtd,brd->bhrt", composite, r) + return extended_einsum("bhtd,brd->bhrt", composite, r) def kg2e_interaction( @@ -754,11 +522,11 @@ def ntn_interaction( """ # save sizes num_heads, num_relations, num_tails, _, k = _extract_sizes(h, b, t) - x = _extended_einsum("bhd,brkde,bte->bhrtk", h, w, t) - x = x + _extended_einsum("brkd,bhd->bhk", vh, h).view(-1, num_heads, 1, 1, k) - x = x + _extended_einsum("brkd,btd->btk", vt, t).view(-1, 1, 1, num_tails, k) + x = extended_einsum("bhd,brkde,bte->bhrtk", h, w, t) + x = x + extended_einsum("brkd,bhd->bhk", vh, h).view(-1, num_heads, 1, 1, k) + x = x + extended_einsum("brkd,btd->btk", vt, t).view(-1, 1, 1, num_tails, k) x = activation(x) - x = _extended_einsum("bhrtk,brk->bhrt", x, u) + x = extended_einsum("bhrtk,brk->bhrt", x, u) return x @@ -828,7 +596,7 @@ def rescal_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return _extended_einsum("bhd,brde,bte->bhrt", h, r, t) + return extended_einsum("bhd,brde,bte->bhrt", h, r, t) def rotate_interaction( @@ -854,13 +622,13 @@ def rotate_interaction( # # Equivalently, we can rotate the tail by the inverse relation, and measure the distance to the head, i.e. # # |h * r - t| = |h - conj(r) * t| # r_inv = torch.stack([r[:, :, :, 0], -r[:, :, :, 1]], dim=-1) - h, r, t = [_view_complex(x) for x in (h, r, t)] + h, r, t = [view_complex(x) for x in (h, r, t)] # Rotate (=Hadamard product in complex space). - hr = _extended_einsum("bhd,brd->bhrd", h, r) + hr = extended_einsum("bhd,brd->bhrd", h, r) # Workaround until https://github.com/pytorch/pytorch/issues/30704 is fixed - return _negative_norm_of_sum( + return negative_norm_of_sum( hr.unsqueeze(dim=3), t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]), p=2, @@ -935,9 +703,9 @@ def structured_embedding_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return _negative_norm_of_sum( - _extended_einsum("brde,bhd->bhre", r_h, h).unsqueeze(dim=3), - -_extended_einsum("brde,btd->brte", r_t, t).unsqueeze(dim=1), + return negative_norm_of_sum( + extended_einsum("brde,bhd->bhre", r_h, h).unsqueeze(dim=3), + -extended_einsum("brde,btd->brte", r_t, t).unsqueeze(dim=1), p=p, power_norm=power_norm, ) @@ -978,13 +746,13 @@ def transd_interaction( """ # Project entities # shape: (b, h, r, 1, d_r) - h_bot = _project_entity( + h_bot = project_entity( e=h.unsqueeze(dim=2), e_p=h_p.unsqueeze(dim=2), r_p=r_p.unsqueeze(dim=1), ).unsqueeze(dim=-2) # shape: (b, 1, r, t, d_r) - t_bot = _project_entity( + t_bot = project_entity( e=t.unsqueeze(dim=1), e_p=t_p.unsqueeze(dim=1), r_p=r_p.unsqueeze(dim=2), @@ -1017,7 +785,7 @@ def transe_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - num_heads, num_relations, num_tails, embedding_dim = _extract_sizes(h, r, t)[:4] + num_heads, num_relations, num_tails, embedding_dim, _ = _extract_sizes(h, r, t) return _translational_interaction( h=h.view(-1, num_heads, 1, 1, embedding_dim), r=r.view(-1, 1, num_relations, 1, embedding_dim), @@ -1056,9 +824,9 @@ def transh_interaction( """ # Project to hyperplane return _translational_interaction( - h=(h.unsqueeze(dim=2) - _extended_einsum("bhd,brd,bre->bhre", h, w_r, w_r)).unsqueeze(dim=3), + h=(h.unsqueeze(dim=2) - extended_einsum("bhd,brd,bre->bhre", h, w_r, w_r)).unsqueeze(dim=3), r=d_r.view(d_r.shape[0], 1, d_r.shape[1], 1, d_r.shape[2]), - t=(t.unsqueeze(dim=1) - _extended_einsum("btd,brd,bre->brte", t, w_r, w_r)).unsqueeze(dim=1), + t=(t.unsqueeze(dim=1) - extended_einsum("btd,brd,bre->brte", t, w_r, w_r)).unsqueeze(dim=1), p=p, power_norm=power_norm, ) @@ -1148,17 +916,17 @@ def tucker_interaction( The scores. """ # Compute wr = DO(W x_2 r) - x = do0(_extended_einsum("idj,brd->brij", core_tensor, r)) + x = do0(extended_einsum("idj,brd->brij", core_tensor, r)) # Compute h_n = DO(BN(h)) h = _apply_optional_bn_to_tensor(batch_norm=bn1, output_dropout=do1, tensor=h) # compute whr = DO(BN(h_n x_1 wr)) - x = _extended_einsum("brid,bhd->bhri", x, h) + x = extended_einsum("brid,bhd->bhri", x, h) x = _apply_optional_bn_to_tensor(batch_norm=bn2, tensor=x, output_dropout=do2) # Compute whr x_3 t - return _extended_einsum("bhrd,btd->bhrt", x, t) + return extended_einsum("bhrd,btd->bhrt", x, t) def unstructured_model_interaction( @@ -1184,4 +952,4 @@ def unstructured_model_interaction( """ h = h.unsqueeze(dim=2).unsqueeze(dim=3) t = t.unsqueeze(dim=1).unsqueeze(dim=2) - return _negative_norm_of_sum(h, -t, p=p, power_norm=power_norm) + return negative_norm_of_sum(h, -t, p=p, power_norm=power_norm) diff --git a/src/pykeen/nn/init.py b/src/pykeen/nn/init.py index b1d5cc4980..b01af783d6 100644 --- a/src/pykeen/nn/init.py +++ b/src/pykeen/nn/init.py @@ -4,12 +4,15 @@ import math +import numpy as np +import torch import torch.nn import torch.nn.init __all__ = [ 'xavier_uniform_', 'xavier_normal_', + 'init_phases', ] @@ -51,3 +54,9 @@ def xavier_normal_(tensor: torch.Tensor, gain: float = 1.0) -> torch.Tensor: std = gain * 2 / math.sqrt(tensor.shape[-1]) torch.nn.init.normal_(tensor, mean=0., std=std) return tensor + + +def init_phases(x: torch.Tensor) -> torch.Tensor: + r"""Generate random phases between 0 and :math:`2\pi`.""" + phases = 2 * np.pi * torch.rand_like(x[..., :x.shape[-1] // 2]) + return torch.cat([torch.cos(phases), torch.sin(phases)], dim=-1).detach() diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 9f2f2e1888..8aad7a7954 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -38,6 +38,7 @@ 'TransEInteraction', 'TransHInteraction', 'TransRInteraction', + 'TuckerInteraction', 'UnstructuredModelInteraction', ] diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py new file mode 100644 index 0000000000..2cfedb2314 --- /dev/null +++ b/src/pykeen/nn/sim.py @@ -0,0 +1,134 @@ +# -*- coding: utf-8 -*- + +"""Similarity functions.""" + +import math + +import torch + +from ..typing import GaussianDistribution + +__all__ = [ + 'expected_likelihood', + 'kullback_leibler_similarity', + 'KG2E_SIMILARITIES', +] + + +def expected_likelihood( + e: GaussianDistribution, + r: GaussianDistribution, + epsilon: float = 1.0e-10, + exact: bool = True, +) -> torch.FloatTensor: + r"""Compute the similarity based on expected likelihood. + + .. math:: + + D((\mu_e, \Sigma_e), (\mu_r, \Sigma_r))) + = \frac{1}{2} \left( + (\mu_e - \mu_r)^T(\Sigma_e + \Sigma_r)^{-1}(\mu_e - \mu_r) + + \log \det (\Sigma_e + \Sigma_r) + d \log (2 \pi) + \right) + = \frac{1}{2} \left( + \mu^T\Sigma^{-1}\mu + + \log \det \Sigma + d \log (2 \pi) + \right) + + :param e: shape: (batch_size, num_heads, num_tails, d) + The entity Gaussian distribution. + :param r: shape: (batch_size, num_relations, d) + The relation Gaussian distribution. + :param epsilon: float (default=1.0) + Small constant used to avoid numerical issues when dividing. + :param exact: + Whether to return the exact similarity, or leave out constant offsets. + + :return: torch.Tensor, shape: (s_1, ..., s_k) + The similarity. + """ + # subtract, shape: (batch_size, num_heads, num_relations, num_tails, dim) + r_shape = r.mean.shape + r_shape = (r_shape[0], 1, r_shape[1], 1, r_shape[2]) + var = r.diagonal_covariance.view(*r_shape) + e.diagonal_covariance.unsqueeze(dim=2) + mean = e.mean.unsqueeze(dim=2) - r.mean.view(*r_shape) + + #: a = \mu^T\Sigma^{-1}\mu + safe_sigma = torch.clamp_min(var, min=epsilon) + sigma_inv = torch.reciprocal(safe_sigma) + sim = torch.sum(sigma_inv * mean ** 2, dim=-1) + + #: b = \log \det \Sigma + sim = sim + safe_sigma.log().sum(dim=-1) + if exact: + sim = sim + sim.shape[-1] * math.log(2. * math.pi) + return sim + + +def kullback_leibler_similarity( + e: GaussianDistribution, + r: GaussianDistribution, + epsilon: float = 1.0e-10, + exact: bool = True, +) -> torch.FloatTensor: + r"""Compute the similarity based on KL divergence. + + This is done between two Gaussian distributions given by mean `mu_*` and diagonal covariance matrix `sigma_*`. + + .. math:: + + D((\mu_e, \Sigma_e), (\mu_r, \Sigma_r))) + = \frac{1}{2} \left( + tr(\Sigma_r^{-1}\Sigma_e) + + (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) + - \log \frac{det(\Sigma_e)}{det(\Sigma_r)} - k_e + \right) + + Note: The sign of the function has been flipped as opposed to the description in the paper, as the + Kullback Leibler divergence is large if the distributions are dissimilar. + + :param e: shape: (batch_size, num_heads, num_tails, d) + The entity Gaussian distributions, as mean/diagonal covariance pairs. + :param r: shape: (batch_size, num_relations, d) + The relation Gaussian distributions, as mean/diagonal covariance pairs. + :param epsilon: float (default=1.0) + Small constant used to avoid numerical issues when dividing. + :param exact: + Whether to return the exact similarity, or leave out constant offsets. + + :return: torch.Tensor, shape: (s_1, ..., s_k) + The similarity. + """ + # invert covariance, shape: (batch_size, num_relations, d) + safe_sigma_r = torch.clamp_min(r.diagonal_covariance, min=epsilon) + sigma_r_inv = torch.reciprocal(safe_sigma_r) + + #: a = tr(\Sigma_r^{-1}\Sigma_e), (batch_size, num_heads, num_relations, num_tails) + # [(b, h, t, d), (b, r, d) -> (b, 1, r, d) -> (b, 1, d, r)] -> (b, h, t, r) -> (b, h, r, t) + sim = (e.diagonal_covariance @ sigma_r_inv.unsqueeze(dim=1).transpose(-2, -1)).transpose(-2, -1) + + #: b = (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) + r_shape = r.mean.shape + # mu.shape: (b, h, r, t, d) + mu = r.mean.view(r_shape[0], 1, r_shape[1], 1, r_shape[2]) - e.mean.unsqueeze(dim=2) + sim = sim + (mu ** 2 @ sigma_r_inv.view(r_shape[0], 1, r_shape[1], r_shape[2], 1)).squeeze(dim=-1) + + #: c = \log \frac{det(\Sigma_e)}{det(\Sigma_r)} + # = sum log (sigma_e)_i - sum log (sigma_r)_i + # ce.shape: (b, h, t) + ce = e.diagonal_covariance.clamp_min(min=epsilon).log().sum(dim=-1) + # cr.shape: (b, r) + cr = safe_sigma_r.log().sum(dim=-1) + sim = sim + ce.unsqueeze(dim=2) - cr.view(r_shape[0], 1, r_shape[1], 1) + + if exact: + sim = sim - e.mean.shape[-1] + sim = 0.5 * sim + + return sim + + +KG2E_SIMILARITIES = { + 'KL': kullback_leibler_similarity, + 'EL': expected_likelihood, +} diff --git a/src/pykeen/typing.py b/src/pykeen/typing.py index 80485d1e7e..706e564a8b 100644 --- a/src/pykeen/typing.py +++ b/src/pykeen/typing.py @@ -2,7 +2,7 @@ """Type hints for PyKEEN.""" -from typing import Callable, Mapping, Sequence, TypeVar, Union +from typing import Callable, Mapping, NamedTuple, Sequence, TypeVar, Union import numpy as np import torch @@ -20,6 +20,7 @@ 'RelationRepresentation', 'Representation', 'TailRepresentation', + 'GaussianDistribution', ] LabeledTriples = np.ndarray @@ -41,3 +42,10 @@ HeadRepresentation = TypeVar("HeadRepresentation", Representation, Sequence[Representation]) # type: ignore RelationRepresentation = TypeVar("RelationRepresentation", Representation, Sequence[Representation]) # type: ignore TailRepresentation = TypeVar("TailRepresentation", Representation, Sequence[Representation]) # type: ignore + + +class GaussianDistribution(NamedTuple): + """A gaussian distribution with diagonal covariance matrix.""" + + mean: torch.FloatTensor + diagonal_covariance: torch.FloatTensor diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index cee50bafcb..8a5be593b7 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -7,13 +7,14 @@ import logging import random from io import BytesIO -from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Type, TypeVar, Union +from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, SupportsFloat, Tuple, Type, TypeVar, Union import numpy import numpy as np import pandas as pd import torch import torch.nn +from torch.nn import functional from .typing import DeviceHint @@ -22,12 +23,15 @@ 'compose', 'check_shapes', 'clamp_norm', + 'combine_complex', 'compact_mapping', + 'complex_normalize', + 'fix_dataclass_init_docs', 'imag_part', 'invert_mapping', 'is_cudnn_error', - 'l2_regularization', 'is_cuda_oom_error', + 'l2_regularization', 'random_non_negative_int', 'real_part', 'resolve_device', @@ -38,17 +42,13 @@ 'split_list_in_batches', 'normalize_string', 'normalized_lookup', + 'negative_norm_of_sum', 'get_cls', 'get_until_first_blank', 'flatten_dictionary', 'set_random_seed', 'NoRandomSeedNecessary', 'Result', - 'fix_dataclass_init_docs', - 'get_hrt_indices', - 'get_hr_indices', - 'get_ht_indices', - 'get_rt_indices', ] logger = logging.getLogger(__name__) @@ -59,6 +59,7 @@ _CUDA_OOM_ERROR = 'CUDA out of memory.' +# TODO remove (unused) def l2_regularization( *xs: torch.Tensor, normalize: bool = False, @@ -103,6 +104,7 @@ def slice_triples(triples): ) +# TODO remove (unused/untested) def slice_doubles(doubles): """Get the heads and relations from a matrix of doubles.""" return ( @@ -114,6 +116,7 @@ def slice_doubles(doubles): X = TypeVar('X') +# TODO remove (unused) def split_list_in_batches(input_list: List[X], batch_size: int) -> List[List[X]]: """Split a list of instances in batches of size batch_size.""" return list(split_list_in_batches_iter(input_list=input_list, batch_size=batch_size)) @@ -353,6 +356,12 @@ def split_complex( return x[..., :dim], x[..., dim:] +def view_complex(x: torch.FloatTensor) -> torch.Tensor: + """Convert a PyKEEN complex tensor representation into a torch one.""" + real, imag = split_complex(x=x) + return torch.complex(real=real, imag=imag) + + def combine_complex( x_re: torch.FloatTensor, x_im: torch.FloatTensor, @@ -361,6 +370,7 @@ def combine_complex( return torch.cat([x_re, x_im], dim=-1) +# TODO remove (unused) def real_part( x: torch.FloatTensor, ) -> torch.FloatTensor: @@ -369,6 +379,7 @@ def real_part( return x[..., :dim] +# TODO remove (unused) def imag_part( x: torch.FloatTensor, ) -> torch.FloatTensor: @@ -518,37 +529,141 @@ def broadcast_cat( return torch.cat([x.repeat(*x_rep), y.repeat(*y_rep)], dim=dim) -# These three following methods could be used throughout PyKEEN to improve -# the implicit documentation of functions +def get_subclasses(cls: Type[X]) -> Iterable[Type[X]]: + """ + Get all subclasses. + Credit to: https://stackoverflow.com/a/33607093. -def get_hrt_indices(hrt_batch: torch.LongTensor) -> Tuple[torch.LongTensor, torch.LongTensor, torch.LongTensor]: - """Get indices from a head/relation/tail batch.""" - return hrt_batch[:, 0], hrt_batch[:, 1], hrt_batch[:, 2] + """ + for subclass in cls.__subclasses__(): + yield from get_subclasses(subclass) + yield subclass -def get_hr_indices(hr_batch: torch.LongTensor) -> Tuple[torch.LongTensor, torch.LongTensor, None]: - """Get indices from a head/relation batch.""" - return hr_batch[:, 0], hr_batch[:, 1], None +def complex_normalize(x: torch.Tensor) -> torch.Tensor: + r"""Normalize the length of relation vectors, if the forward constraint has not been applied yet. + The `modulus of complex number `_ is given as: -def get_ht_indices(ht_batch: torch.LongTensor) -> Tuple[torch.LongTensor, None, torch.LongTensor]: - """Get indices from a head/tail batch.""" - return ht_batch[:, 0], None, ht_batch[:, 1] + .. math:: + |a + ib| = \sqrt{a^2 + b^2} -def get_rt_indices(rt_batch: torch.LongTensor) -> Tuple[None, torch.LongTensor, torch.LongTensor]: - """Get indices from a relation/tail batch.""" - return None, rt_batch[:, 0], rt_batch[:, 1] + $l_2$ norm of complex vector $x \in \mathbb{C}^d$: + .. math:: + \|x\|^2 = \sum_{i=1}^d |x_i|^2 + = \sum_{i=1}^d \left(\operatorname{Re}(x_i)^2 + \operatorname{Im}(x_i)^2\right) + = \left(\sum_{i=1}^d \operatorname{Re}(x_i)^2) + (\sum_{i=1}^d \operatorname{Im}(x_i)^2\right) + = \|\operatorname{Re}(x)\|^2 + \|\operatorname{Im}(x)\|^2 + = \| [\operatorname{Re}(x); \operatorname{Im}(x)] \|^2 + """ + y = x.data.view(x.shape[0], -1, 2) + y = functional.normalize(y, p=2, dim=-1) + x.data = y.view(*x.shape) + return x -def get_subclasses(cls: Type[X]) -> Iterable[Type[X]]: + +def negative_norm_of_sum( + *x: torch.FloatTensor, + p: Union[str, int] = 2, + power_norm: bool = False, +) -> torch.FloatTensor: + """Evaluate negative norm of a sum of vectors on already broadcasted representations. + + :param x: shape: (batch_size, num_heads, num_relations, num_tails, dim) + The representations. + :param p: + The p for the norm. cf. torch.norm. + :param power_norm: + Whether to return $|x-y|_p^p$, cf. https://github.com/pytorch/pytorch/issues/28119 + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. """ - Get all subclasses. + d: torch.FloatTensor = sum(x) + if power_norm: + assert isinstance(p, SupportsFloat) + return -(d.abs() ** p).sum(dim=-1) - Credit to: https://stackoverflow.com/a/33607093. + if torch.is_complex(d): + assert isinstance(p, SupportsFloat) + # workaround for complex numbers: manually compute norm + return -(d.abs() ** p).sum(dim=-1) ** (1 / p) + + return -d.norm(p=p, dim=-1) + + +def extended_einsum( + eq: str, + *tensors, +) -> torch.FloatTensor: + """Drop dimensions of size 1 to allow broadcasting.""" + # TODO: check if einsum is still very slow. + lhs, rhs = eq.split("->") + mod_ops, mod_t = [], [] + for op, t in zip(lhs.split(","), tensors): + mod_op = "" + assert len(op) == len(t.shape) + for i, c in reversed(list(enumerate(op))): + if t.shape[i] == 1: + t = t.squeeze(dim=i) + else: + mod_op = c + mod_op + mod_ops.append(mod_op) + mod_t.append(t) + m_lhs = ",".join(mod_ops) + r_keep_dims = set("".join(mod_ops)) + m_rhs = "".join(c for c in rhs if c in r_keep_dims) + m_eq = f"{m_lhs}->{m_rhs}" + mod_r = torch.einsum(m_eq, *mod_t) + # unsqueeze + for i, c in enumerate(rhs): + if c not in r_keep_dims: + mod_r = mod_r.unsqueeze(dim=i) + return mod_r + + +def project_entity( + e: torch.FloatTensor, + e_p: torch.FloatTensor, + r_p: torch.FloatTensor, +) -> torch.FloatTensor: + r"""Project entity relation-specific. + + .. math:: + + e_{\bot} = M_{re} e + = (r_p e_p^T + I^{d_r \times d_e}) e + = r_p e_p^T e + I^{d_r \times d_e} e + = r_p (e_p^T e) + e' + + and additionally enforces + + .. math:: + + \|e_{\bot}\|_2 \leq 1 + + :param e: shape: (..., d_e) + The entity embedding. + :param e_p: shape: (..., d_e) + The entity projection. + :param r_p: shape: (..., d_r) + The relation projection. + + :return: shape: (..., d_r) """ - for subclass in cls.__subclasses__(): - yield from get_subclasses(subclass) - yield subclass + # The dimensions affected by e' + change_dim = min(e.shape[-1], r_p.shape[-1]) + + # Project entities + # r_p (e_p.T e) + e' + e_bot = r_p * torch.sum(e_p * e, dim=-1, keepdim=True) + e_bot[..., :change_dim] += e[..., :change_dim] + + # Enforce constraints + e_bot = clamp_norm(e_bot, p=2, dim=-1, maxnorm=1) + + return e_bot diff --git a/tests/test_models.py b/tests/test_models.py index a55b958c25..3c6e689cec 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -41,10 +41,9 @@ symmetric_edge_weights, ) from pykeen.nn import RepresentationModule -from pykeen.nn.functional import _project_entity from pykeen.training import LCWATrainingLoop, SLCWATrainingLoop, TrainingLoop from pykeen.triples import TriplesFactory -from pykeen.utils import all_in_bounds, clamp_norm, set_random_seed +from pykeen.utils import all_in_bounds, clamp_norm, project_entity, set_random_seed SKIP_MODULES = { Model.__name__, @@ -862,7 +861,7 @@ def test_project_entity(self): r_p = torch.rand(self.batch_size, 1, self.model.relation_dim, generator=self.generator) # project - e_bot = _project_entity(e=e, e_p=e_p, r_p=r_p) + e_bot = project_entity(e=e, e_p=e_p, r_p=r_p) # check shape: assert e_bot.shape == (self.batch_size, self.model.num_entities, self.model.relation_dim) From 4a2720c720f4561f2eb8d2e97fc7645aaf7ce48a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 10:10:17 +0100 Subject: [PATCH 341/690] Remove slight code duplication --- src/pykeen/models/base.py | 22 ++++++++++++---------- 1 file changed, 12 insertions(+), 10 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 2170908298..d2d657973c 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1107,6 +1107,16 @@ def load_state(self, path: str) -> None: self.load_state_dict(torch.load(path, map_location=self.device)) +def _prepare_representation_module_list( + representations: Union[None, RepresentationModule, Sequence[RepresentationModule]], +) -> Sequence[RepresentationModule]: + """Normalize list of representations and wrap into nn.ModuleList.""" + # Important: use ModuleList to ensure that Pytorch correctly handles their devices and parameters + if representations is not None and not isinstance(representations, Sequence): + representations = [representations] + return nn.ModuleList(representations) + + class ERModel(Model, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation], autoreset=False): """A commonly useful base for KGEMs using embeddings and interaction modules.""" @@ -1155,16 +1165,8 @@ def __init__( random_seed=random_seed, predict_with_sigmoid=predict_with_sigmoid, ) - - # normalization - if entity_representations is not None and not isinstance(entity_representations, Sequence): - entity_representations = [entity_representations] - if relation_representations is not None and not isinstance(relation_representations, Sequence): - relation_representations = [relation_representations] - # Important: use ModuleList to ensure that Pytorch correctly handles their devices and parameters - self.entity_representations = nn.ModuleList(entity_representations) - self.relation_representations = nn.ModuleList(relation_representations) - + self.entity_representations = _prepare_representation_module_list(representations=entity_representations) + self.relation_representations = _prepare_representation_module_list(representations=relation_representations) self.interaction = interaction def forward( From 5d24094bfaee721fda7629afced1453decab69eb Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 10:13:56 +0100 Subject: [PATCH 342/690] Subclass TransR from DoubleRelationEmbeddingModel --- src/pykeen/models/base.py | 3 ++ src/pykeen/models/unimodal/trans_r.py | 55 +++++++++++---------------- 2 files changed, 26 insertions(+), 32 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index d2d657973c..119f1a9d31 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1349,6 +1349,7 @@ def __init__( ], embedding_dim: int = 50, relation_dim: Union[None, int, Sequence[int]] = None, + second_relation_dim: Union[None, int, Sequence[int]] = None, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, automatic_memory_optimization: Optional[bool] = None, @@ -1360,6 +1361,8 @@ def __init__( ) -> None: if relation_dim is None: relation_dim = embedding_dim + if second_relation_dim is None: + second_relation_dim = relation_dim super().__init__( triples_factory=triples_factory, interaction=interaction, diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index f5a484bb25..8f49fd88fd 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -7,9 +7,9 @@ from torch.nn import functional -from .. import ERModel +from .. import DoubleRelationEmbeddingModel from ...losses import Loss -from ...nn import Embedding, EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import TransRInteraction from ...triples import TriplesFactory @@ -21,7 +21,7 @@ ] -class TransR(ERModel): +class TransR(DoubleRelationEmbeddingModel): r"""An implementation of TransR from [lin2015]_. TransR is an extension of :class:`pykeen.models.TransH` that explicitly considers entities and relations as @@ -77,37 +77,28 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - entity_representations=Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - shape=embedding_dim, - specification=EmbeddingSpecification( - initializer=xavier_uniform_, - constrainer=clamp_norm, # type: ignore - constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), - ), + # Entity embeddings + embedding_dim=embedding_dim, + embedding_specification=EmbeddingSpecification( + initializer=xavier_uniform_, + constrainer=clamp_norm, # type: ignore + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), - relation_representations=[ - Embedding.from_specification( - triples_factory.num_relations, - shape=relation_dim, - specification=EmbeddingSpecification( - initializer=compose( - xavier_uniform_, - functional.normalize, - ), - constrainer=clamp_norm, # type: ignore - constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), - ), - ), - # Relation projections - Embedding.from_specification( - triples_factory.num_relations, - shape=(relation_dim, embedding_dim), - specification=EmbeddingSpecification( - initializer=xavier_uniform_, - ), + # Relation embeddings + relation_dim=relation_dim, + relation_embedding_specification=EmbeddingSpecification( + initializer=compose( + xavier_uniform_, + functional.normalize, ), - ], + constrainer=clamp_norm, # type: ignore + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + ), + # Relation projections + second_relation_dim=(relation_dim, embedding_dim), + second_relation_embedding_specification=EmbeddingSpecification( + initializer=xavier_uniform_, + ), interaction=TransRInteraction( p=scoring_fct_norm, ), From ca514a93db2aaa5d3cc73ac47f0e6c1f02e0bba8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 10:14:12 +0100 Subject: [PATCH 343/690] Fix import --- src/pykeen/models/unimodal/trans_r.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 8f49fd88fd..d632646b1f 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -7,7 +7,7 @@ from torch.nn import functional -from .. import DoubleRelationEmbeddingModel +from ..base import DoubleRelationEmbeddingModel from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.init import xavier_uniform_ From a138455f6aa30aca3076e24d4633219f3b6851f5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 10:14:29 +0100 Subject: [PATCH 344/690] Fix usage of second_relation_dim --- src/pykeen/models/base.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 119f1a9d31..4493fd8f8f 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1318,6 +1318,7 @@ def __init__( @property def embedding_dim(self) -> int: # noqa:D401 """The entity embedding dim.""" + # TODO: Deprecated; directly use self.entity_representations[0].embedding_dim instead? embedding = self.entity_representations[0] assert isinstance(embedding, Embedding) return embedding.embedding_dim @@ -1325,6 +1326,7 @@ def embedding_dim(self) -> int: # noqa:D401 @property def relation_dim(self) -> int: # noqa:D401 """The relation embedding dim.""" + # TODO: Deprecated; directly use self.relation_representations[0].embedding_dim instead? embedding = self.relation_representations[0] assert isinstance(embedding, Embedding) return embedding.embedding_dim @@ -1386,7 +1388,7 @@ def __init__( ), Embedding.from_specification( num_embeddings=triples_factory.num_relations, - shape=relation_dim, + shape=second_relation_dim, specification=second_relation_embedding_specification, ), ], From 94846c13939ab1d335adbc3baad03b645d0adff5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 10:16:50 +0100 Subject: [PATCH 345/690] Expose embedding specification for ComplEx --- src/pykeen/models/unimodal/complex.py | 26 ++++++++++++++++---------- 1 file changed, 16 insertions(+), 10 deletions(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index a4ee8c43f0..a728a21b7f 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -74,6 +74,8 @@ def __init__( loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, + embedding_specification: Optional[EmbeddingSpecification] = None, + relation_embedding_specification: Optional[EmbeddingSpecification] = None, ) -> None: """Initialize ComplEx. @@ -92,6 +94,18 @@ def __init__( An optional random seed to set before the initialization of weights. """ regularizer = LpRegularizer(weight=0.01, p=2.0, normalize=True) + # initialize with entity and relation embeddings with standard normal distribution, cf. + # https://github.com/ttrouill/complex/blob/dc4eb93408d9a5288c986695b58488ac80b1cc17/efe/models.py#L481-L487 + if embedding_specification is None: + embedding_specification = EmbeddingSpecification( + initializer=nn.init.normal_, + regularizer=regularizer, + ) + if relation_embedding_specification is None: + relation_embedding_specification = EmbeddingSpecification( + initializer=nn.init.normal_, + regularizer=regularizer, + ), super().__init__( triples_factory=triples_factory, interaction=ComplExInteraction(), @@ -100,14 +114,6 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - # initialize with entity and relation embeddings with standard normal distribution, cf. - # https://github.com/ttrouill/complex/blob/dc4eb93408d9a5288c986695b58488ac80b1cc17/efe/models.py#L481-L487 - embedding_specification=EmbeddingSpecification( - initializer=nn.init.normal_, - regularizer=regularizer, - ), - relation_embedding_specification=EmbeddingSpecification( - initializer=nn.init.normal_, - regularizer=regularizer, - ), + embedding_specification=embedding_specification, + relation_embedding_specification=relation_embedding_specification, ) From 2e961f9fb66046f8ae0b6af2134ef343f3569cec Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 10:24:38 +0100 Subject: [PATCH 346/690] Add option to add weigh regularizers by name --- src/pykeen/models/base.py | 14 +++++++++++++- src/pykeen/regularizers.py | 14 +++++++++++--- 2 files changed, 24 insertions(+), 4 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 4493fd8f8f..6f5a883fec 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -22,7 +22,7 @@ from ..losses import Loss, MarginRankingLoss, NSSALoss from ..nn import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule from ..nn.modules import Interaction -from ..regularizers import collect_regularization_terms +from ..regularizers import Regularizer, collect_regularization_terms from ..triples import TriplesFactory, TriplesNumericLiteralsFactory from ..typing import DeviceHint, HeadRepresentation, MappedTriples, RelationRepresentation, TailRepresentation from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed @@ -1126,6 +1126,9 @@ class ERModel(Model, Generic[HeadRepresentation, RelationRepresentation, TailRep #: The relation representations relation_representations: Sequence[RepresentationModule] + #: The weight regularizers + weight_regularizers: List[Regularizer] + def __init__( self, triples_factory: TriplesFactory, @@ -1168,6 +1171,15 @@ def __init__( self.entity_representations = _prepare_representation_module_list(representations=entity_representations) self.relation_representations = _prepare_representation_module_list(representations=relation_representations) self.interaction = interaction + self.weight_regularizers = nn.ModuleList() + + def add_weight_regularizer(self, parameter_name: str, regularizer: Regularizer) -> None: + weights = dict(self.named_parameters()) + if parameter_name not in weights.keys(): + raise ValueError(f"Invalid parameter_name={parameter_name}. Available are: {sorted(weights.keys())}.") + parameter: nn.Parameter = weights[parameter_name] + regularizer.add_parameter(parameter=parameter) + self.weight_regularizers.append(regularizer) def forward( self, diff --git a/src/pykeen/regularizers.py b/src/pykeen/regularizers.py index c8012acde8..b805a258ff 100644 --- a/src/pykeen/regularizers.py +++ b/src/pykeen/regularizers.py @@ -3,7 +3,7 @@ """Regularization in PyKEEN.""" from abc import ABC, abstractmethod -from typing import Any, ClassVar, Collection, Iterable, Mapping, Optional, Sequence, Type, Union +from typing import Any, ClassVar, Collection, Iterable, List, Mapping, Optional, Sequence, Type, Union import torch from torch import nn @@ -41,7 +41,7 @@ class Regularizer(nn.Module, ABC): hpo_default: ClassVar[Mapping[str, Any]] #: weights which should be regularized - tracked_parameters: Optional[Collection[nn.Parameter]] = None + tracked_parameters: List[nn.Parameter] def __init__( self, @@ -52,6 +52,10 @@ def __init__( super().__init__() self.register_buffer(name='weight', tensor=torch.as_tensor(weight)) self.apply_only_once = apply_only_once + if parameters is None: + parameters = [] + else: + parameters = list(parameters) self.tracked_parameters = parameters self._clear() @@ -59,6 +63,10 @@ def _clear(self): self.regularization_term = 0. self.updated = False + def add_parameter(self, parameter: nn.Parameter) -> None: + """Add a parameter for regularization.""" + self.tracked_parameters.append(parameter) + @classmethod def get_normalized_name(cls) -> str: """Get the normalized name of the regularizer class.""" @@ -79,7 +87,7 @@ def update(self, *tensors: torch.FloatTensor) -> bool: def pop_regularization_term(self) -> torch.FloatTensor: """Return the weighted regularization term, and clear it afterwards.""" - if self.tracked_parameters is not None: + if len(self.tracked_parameters) > 0: self.update(*self.tracked_parameters) term = self.regularization_term self._clear() From ba1719b6f5f07e37e72ebd1893e80f0b2bc2dbb9 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 10:27:32 +0100 Subject: [PATCH 347/690] Add comment --- src/pykeen/models/base.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 6f5a883fec..59d845a544 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1171,6 +1171,8 @@ def __init__( self.entity_representations = _prepare_representation_module_list(representations=entity_representations) self.relation_representations = _prepare_representation_module_list(representations=relation_representations) self.interaction = interaction + # Comment: it is important that the regularizers are stored in a module list, in order to appear in + # model.modules(). Thereby, we can collect them automatically. self.weight_regularizers = nn.ModuleList() def add_weight_regularizer(self, parameter_name: str, regularizer: Regularizer) -> None: From f957f6a637d7cf29c1bf7d5168fa9b3719cc0129 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 10:30:48 +0100 Subject: [PATCH 348/690] Add docstring --- src/pykeen/models/base.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 59d845a544..b622e531bf 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1175,7 +1175,18 @@ def __init__( # model.modules(). Thereby, we can collect them automatically. self.weight_regularizers = nn.ModuleList() - def add_weight_regularizer(self, parameter_name: str, regularizer: Regularizer) -> None: + def add_weight_regularizer( + self, + parameter_name: str, + regularizer: Regularizer, + ) -> None: + """Add a model weight to a regularizer's weight list, and register the regularizer with the model. + + :param parameter_name: + The parameter name. Available parameter names are shown by `dict(self.named_parameters()).keys()`. + :param regularizer: + The regularizer instance which will regularize the weights. + """ weights = dict(self.named_parameters()) if parameter_name not in weights.keys(): raise ValueError(f"Invalid parameter_name={parameter_name}. Available are: {sorted(weights.keys())}.") From 281189d733aeb3214e75a795f1816f3d70f4785c Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 13:08:52 +0100 Subject: [PATCH 349/690] Cleanup test --- tests/test_regularizers.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/tests/test_regularizers.py b/tests/test_regularizers.py index 3a0c10eb19..a2766aa8f7 100644 --- a/tests/test_regularizers.py +++ b/tests/test_regularizers.py @@ -198,9 +198,9 @@ class TransHRegularizerTest(unittest.TestCase): regularizer_kwargs: Dict num_entities: int num_relations: int - entities_weight: torch.Tensor - relations_weight: torch.Tensor - normal_vector_weight: torch.Tensor + entities_weight: nn.Parameter + relations_weight: nn.Parameter + normal_vector_weight: nn.Parameter def setUp(self) -> None: """Set up the test case.""" @@ -208,37 +208,37 @@ def setUp(self) -> None: self.device = resolve_device() self.num_entities = 10 self.num_relations = 5 - self.regularizer_kwargs = {'weight': .5, 'epsilon': 1e-5} self.entities_weight = self._rand_param(10) self.relations_weight = self._rand_param(20) self.normal_vector_weight = self._rand_param(20) - self.regularizer_kwargs["entity_embeddings"] = self.entities_weight - self.regularizer_kwargs["normal_vector_embeddings"] = self.normal_vector_weight - self.regularizer_kwargs["relation_embeddings"] = self.relations_weight + self.weight = .5 + self.epsilon = 1e-5 + self.regularizer_kwargs = dict() self.regularizer = TransHRegularizer( - **(self.regularizer_kwargs or {}), + weight=self.weight, epsilon=self.epsilon, + entity_embeddings=self.entities_weight, + normal_vector_embeddings=self.normal_vector_weight, + relation_embeddings=self.relations_weight, ) - def _rand_param(self, n): - return nn.Parameter(torch.rand(self.num_entities, 10, device=self.device, generator=self.generator)) + def _rand_param(self, n) -> nn.Parameter: + return nn.Parameter(torch.rand(self.num_entities, n, device=self.device, generator=self.generator)) def test_update(self): """Test update function of TransHRegularizer.""" # Test that regularization term is computed correctly expected_term = self._expected_penalty() - weight = self.regularizer_kwargs.get('weight') observed_term = self.regularizer.pop_regularization_term() - assert torch.allclose(observed_term, weight * expected_term) + assert torch.allclose(observed_term, self.weight * expected_term) def _expected_penalty(self) -> torch.FloatTensor: # noqa: D102 # Entity soft constraint regularization_term = torch.sum(functional.relu(torch.norm(self.entities_weight, dim=-1) ** 2 - 1.0)) - epsilon = self.regularizer_kwargs.get('epsilon') # # Orthogonality soft constraint d_r_n = functional.normalize(self.relations_weight, dim=-1) regularization_term += torch.sum( - functional.relu(torch.sum((self.normal_vector_weight * d_r_n) ** 2, dim=-1) - epsilon), + functional.relu(torch.sum((self.normal_vector_weight * d_r_n) ** 2, dim=-1) - self.epsilon), ) return regularization_term From e885fadb105c0f17381fbb3a1d54f7ff6786c929 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 13:09:46 +0100 Subject: [PATCH 350/690] Add default regularizer back --- src/pykeen/models/base.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index b622e531bf..bbac25245c 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -254,6 +254,11 @@ class Model(nn.Module, ABC): #: The instance of the loss loss: Loss + #: The default regularizer class + regularizer_default: ClassVar[Optional[Type[Regularizer]]] = None + #: The default parameters for the default regularizer class + regularizer_default_kwargs: ClassVar[Optional[Mapping[str, Any]]] = None + def __init__( self, triples_factory: TriplesFactory, From ebab5827b14c9fae1646e710dc7033d66e0fc77b Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 13:10:53 +0100 Subject: [PATCH 351/690] Clean up regularizers --- src/pykeen/regularizers.py | 27 +++++++++++++-------------- 1 file changed, 13 insertions(+), 14 deletions(-) diff --git a/src/pykeen/regularizers.py b/src/pykeen/regularizers.py index b805a258ff..3aa51bd955 100644 --- a/src/pykeen/regularizers.py +++ b/src/pykeen/regularizers.py @@ -34,6 +34,9 @@ class Regularizer(nn.Module, ABC): #: The current regularization term (a scalar) regularization_term: Union[torch.FloatTensor, float] + #: Has this regularizer been updated? + updated: bool + #: Should the regularization only be applied once? This was used for ConvKB and defaults to False. apply_only_once: bool @@ -47,16 +50,12 @@ def __init__( self, weight: float = 1.0, apply_only_once: bool = False, - parameters: Optional[Sequence[nn.Parameter]] = None, + parameters: Optional[Iterable[nn.Parameter]] = None, ): super().__init__() self.register_buffer(name='weight', tensor=torch.as_tensor(weight)) self.apply_only_once = apply_only_once - if parameters is None: - parameters = [] - else: - parameters = list(parameters) - self.tracked_parameters = parameters + self.tracked_parameters = list(parameters) if parameters else [] self._clear() def _clear(self): @@ -102,7 +101,7 @@ class NoRegularizer(Regularizer): # TODO: Deprecated - #: The default strategy for optimizing the regularizer's hyper-parameters + #: The default strategy for optimizing the no-op regularizer's hyper-parameters hpo_default: ClassVar[Mapping[str, Any]] = {} def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102 @@ -120,8 +119,8 @@ class LpRegularizer(Regularizer): #: This allows dimensionality-independent weight tuning. normalize: bool - #: The default strategy for optimizing the regularizer's hyper-parameters - hpo_default = dict( + #: The default strategy for optimizing the LP regularizer's hyper-parameters + hpo_default: ClassVar[Mapping[str, Any]] = dict( weight=dict(type=float, low=0.01, high=1.0, scale='log'), ) @@ -160,8 +159,8 @@ class PowerSumRegularizer(Regularizer): Has some nice properties, cf. e.g. https://github.com/pytorch/pytorch/issues/28119. """ - #: The default strategy for optimizing the regularizer's hyper-parameters - hpo_default = dict( + #: The default strategy for optimizing the power sum regularizer's hyper-parameters + hpo_default: ClassVar[Mapping[str, Any]] = dict( weight=dict(type=float, low=0.01, high=1.0, scale='log'), ) @@ -190,8 +189,8 @@ def forward(self, x: torch.FloatTensor) -> torch.FloatTensor: # noqa: D102 class TransHRegularizer(Regularizer): """A regularizer for the soft constraints in TransH.""" - #: The default strategy for optimizing the regularizer's hyper-parameters - hpo_default = dict( + #: The default strategy for optimizing the TransH regularizer's hyper-parameters + hpo_default: ClassVar[Mapping[str, Any]] = dict( weight=dict(type=float, low=0.01, high=1.0, scale='log'), ) @@ -242,7 +241,7 @@ def __init__( super().__init__(weight=total_weight, apply_only_once=apply_only_once, parameters=parameters) self.regularizers = nn.ModuleList(regularizers) for r in self.regularizers: - if isinstance(r, NoRegularizer): + if r is None or isinstance(r, NoRegularizer): raise TypeError('Can not combine a no-op regularizer') normalization_factor = torch.as_tensor(sum(r.weight for r in self.regularizers)).reciprocal() self.register_buffer(name='normalization_factor', tensor=normalization_factor) From 4f9dcfe60ac58f70cb00ea257f227d0319364d65 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 13:21:08 +0100 Subject: [PATCH 352/690] Fix trailing comma --- src/pykeen/models/unimodal/complex.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index a728a21b7f..899d55618b 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -105,7 +105,7 @@ def __init__( relation_embedding_specification = EmbeddingSpecification( initializer=nn.init.normal_, regularizer=regularizer, - ), + ) super().__init__( triples_factory=triples_factory, interaction=ComplExInteraction(), From 5c1ba555d0d40999f5fbd65d6bec8134075cc7e5 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 13:22:08 +0100 Subject: [PATCH 353/690] Update instantiation of regularizer With this commit, the default regularizer and kwargs stored in the class are used instead of hard coding the values again. This makes it a bit more dynamic and will allow for user input in a later commit. --- src/pykeen/models/base.py | 10 ++++++++++ src/pykeen/models/unimodal/complex.py | 2 +- src/pykeen/models/unimodal/conv_kb.py | 2 +- src/pykeen/models/unimodal/distmult.py | 6 +----- src/pykeen/models/unimodal/rescal.py | 2 +- src/pykeen/models/unimodal/simple.py | 2 +- src/pykeen/models/unimodal/trans_h.py | 4 +--- 7 files changed, 16 insertions(+), 12 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index bbac25245c..911be3942e 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -420,6 +420,16 @@ def _set_device(self, device: DeviceHint = None) -> None: """Set the Torch device to use.""" self.device = resolve_device(device=device) + def _instantiate_default_regularizer(self, **kwargs) -> Optional[Regularizer]: + """Instantiate the regularizer from this class's default settings. + + If the default regularizer is None, None is returned. + Handles the corner case when the default regularizer's keyword arguments are None + Additional keyword arguments can be passed through to the `__init__()` function + """ + if self.regularizer_default is not None: + return self.regularizer_default(**(self.regularizer_default_kwargs or {}), **kwargs) + def to_device_(self) -> 'Model': """Transfer model to device.""" self.to(self.device) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 899d55618b..708c5bed60 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -93,7 +93,7 @@ def __init__( :param random_seed: int (optional) An optional random seed to set before the initialization of weights. """ - regularizer = LpRegularizer(weight=0.01, p=2.0, normalize=True) + regularizer = self._instantiate_default_regularizer() # initialize with entity and relation embeddings with standard normal distribution, cf. # https://github.com/ttrouill/complex/blob/dc4eb93408d9a5288c986695b58488ac80b1cc17/efe/models.py#L481-L487 if embedding_specification is None: diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 529654b91e..5fd14928e8 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -96,5 +96,5 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, ) - self.regularizer = LpRegularizer(parameters=list(self.interaction.parameters())) + self.regularizer = self._instantiate_default_regularizer(parameters=list(self.interaction.parameters())) logger.warning('To be consistent with the paper, initialize entity and relation embeddings from TransE.') diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index bcc4c491b6..6dfb4330be 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -102,10 +102,6 @@ def __init__( functional.normalize, ), # Only relation embeddings are regularized - regularizer=LpRegularizer( - weight=0.1, - p=2.0, - normalize=True, - ), + regularizer=self._instantiate_default_regularizer(), ), ) diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index 836050e46a..e91b6ed0f0 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -67,7 +67,7 @@ def __init__( - OpenKE `implementation of RESCAL `_ """ - regularizer = LpRegularizer(weight=10, p=2., normalize=True) + regularizer = self._instantiate_default_regularizer() super().__init__( triples_factory=triples_factory, interaction=RESCALInteraction(), diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index 8c07da910c..6e3c6e3a92 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -73,7 +73,7 @@ def __init__( random_seed: Optional[int] = None, clamp_score: Optional[Union[float, Tuple[float, float]]] = None, ) -> None: - regularizer = PowerSumRegularizer(weight=20, p=2.0, normalize=True) + regularizer = self._instantiate_default_regularizer() super().__init__( triples_factory=triples_factory, interaction=DistMultInteraction(), diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 6c0e8d3ba7..56d0664c04 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -97,9 +97,7 @@ def __init__( constrainer=functional.normalize, ), ) - self.regularizer = TransHRegularizer( - weight=0.05, - epsilon=1e-5, + self.regularizer = self._instantiate_default_regularizer( entity_embeddings=list(self.entity_representations[0].parameters()).pop(), relation_embeddings=list(self.relation_representations[0].parameters()).pop(), normal_vector_embeddings=list(self.relation_representations[1].parameters()).pop(), From b86fce85b430302992a80cbc6c0806d491cfa979 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 13:27:33 +0100 Subject: [PATCH 354/690] Fix typing --- src/pykeen/models/base.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 911be3942e..52de2777c9 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -427,8 +427,12 @@ def _instantiate_default_regularizer(self, **kwargs) -> Optional[Regularizer]: Handles the corner case when the default regularizer's keyword arguments are None Additional keyword arguments can be passed through to the `__init__()` function """ - if self.regularizer_default is not None: - return self.regularizer_default(**(self.regularizer_default_kwargs or {}), **kwargs) + if self.regularizer_default is None: + return None + + _kwargs = dict(self.regularizer_default_kwargs or {}) + _kwargs.update(kwargs) + return self.regularizer_default(**_kwargs) def to_device_(self) -> 'Model': """Transfer model to device.""" From 4f42439d94c9cbb7bcd53656ccbe4d5ec2346c9a Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 13:48:43 +0100 Subject: [PATCH 355/690] Update CLI --- src/pykeen/models/cli/builders.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/cli/builders.py b/src/pykeen/models/cli/builders.py index fd266b449a..e4f5d41d09 100644 --- a/src/pykeen/models/cli/builders.py +++ b/src/pykeen/models/cli/builders.py @@ -14,8 +14,7 @@ from . import options from .options import CLI_OPTIONS from ..base import Model -from ...regularizers import Regularizer, _REGULARIZER_SUFFIX, regularizers -from ...utils import normalize_string +from ...regularizers import Regularizer, regularizers __all__ = [ 'build_cli_from_cls', @@ -54,7 +53,7 @@ def _decorate_model_kwargs(command: click.Command) -> click.Command: parameter = signature.parameters[name] option = click.option( '--regularizer', - type=str, + type=click.Choice(regularizers), default=parameter.default, show_default=True, help=f'The name of the regularizer preset for {model.__name__}', @@ -88,8 +87,8 @@ def _decorate_model_kwargs(command: click.Command) -> click.Command: regularizer_option = click.option( '--regularizer', type=click.Choice(regularizers), - help=f'The name of the regularizer. Defaults to' - f' {normalize_string(model.regularizer_default.__name__, suffix=_REGULARIZER_SUFFIX)}', + default=model.regularizer_default.get_normalized_name() if model.regularizer_default else 'no', + show_default=True, ) @click.command(help=f'CLI for {model.__name__}', name=model.__name__.lower()) From 20a2451da77293bd8bf3d38eb2472f26a52e091d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 14:19:49 +0100 Subject: [PATCH 356/690] Add test for pop_regularization_term --- tests/test_regularizers.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test_regularizers.py b/tests/test_regularizers.py index a2766aa8f7..44ab0f9c75 100644 --- a/tests/test_regularizers.py +++ b/tests/test_regularizers.py @@ -110,6 +110,22 @@ def test_forward(self) -> None: else: assert (expected_penalty == penalty).all() + def test_pop_regularization_term(self): + """Test pop_regularization_term.""" + regularization_term = self.regularizer.pop_regularization_term() + + # check type + assert isinstance(regularization_term, float) or torch.is_tensor(regularization_term) + + # float only if there is not real regularization term + if isinstance(regularization_term, float): + assert regularization_term == 0.0 + + # check that the regularizer has been clear + assert isinstance(self.regularizer.regularization_term, float) + assert self.regularizer.regularization_term == 0.0 + assert self.regularizer.updated is False + def _expected_penalty(self, x: torch.FloatTensor) -> torch.FloatTensor: """Compute expected penalty for given tensor.""" return None From 18556904d7107673b496b6c5e82392b31a8442e7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 14:35:53 +0100 Subject: [PATCH 357/690] Add test for add_weight_regularizer --- tests/test_models.py | 63 +++++++++++++++++++++++++++++++++++++------- 1 file changed, 53 insertions(+), 10 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 3c6e689cec..79ec22f6c5 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -8,7 +8,7 @@ import traceback import unittest from typing import Any, ClassVar, Mapping, Optional, Type -from unittest.mock import patch +from unittest.mock import MagicMock, patch import numpy import pytest @@ -41,6 +41,7 @@ symmetric_edge_weights, ) from pykeen.nn import RepresentationModule +from pykeen.regularizers import LpRegularizer, Regularizer, collect_regularization_terms from pykeen.training import LCWATrainingLoop, SLCWATrainingLoop, TrainingLoop from pykeen.triples import TriplesFactory from pykeen.utils import all_in_bounds, clamp_norm, project_entity, set_random_seed @@ -86,6 +87,48 @@ def get_in_canonical_shape( return x.unsqueeze(dim=1) +class ERModelTests(unittest.TestCase): + """Test basic functionality of ERModel.""" + + def setUp(self) -> None: + """Setup the test instance.""" + self.model = ERModel( + triples_factory=MagicMock(), + interaction=MagicMock(), + ) + + def test_add_weight_regularizer_non_existing(self): + """Test add_weight_regularizer.""" + # try to add regularizer to non-existing weight + with self.assertRaises(ValueError): + self.model.add_weight_regularizer( + parameter_name="this.weight.does.not.exist", + regularizer=..., + ) + + def test_add_weight_regularizer(self): + """Test add_weight_regularizer.""" + # add weighted submodules + self.model.linear = nn.Linear(3, 2) + self.model.sub_model = nn.Sequential( + nn.Linear(2, 3), + nn.LeakyReLU(), + self.model.linear, + ) + + regularizer = LpRegularizer() + + # try to add regularizer to existing weight + self.model.add_weight_regularizer( + parameter_name="linear.weight", + regularizer=regularizer, + ) + + # check it gets found by collect + term = collect_regularization_terms(self.model) + assert torch.is_tensor(term) + + class _ModelTestCase: """A test case for quickly defining common tests for KGE models.""" @@ -385,16 +428,16 @@ def test_has_hpo_defaults(self): else: self.assertIsInstance(d, dict) - def test_post_parameter_update_regularizer(self): - """Test whether post_parameter_update resets the regularization term.""" - # set regularizer term - self.model.regularizer.regularization_term = None - - # call post_parameter_update - self.model.post_parameter_update() + def test_collect_regularization_terms(self): + """Test whether collect_regularization_terms resets all regularization terms.""" + # TODO: Does this need to be tested for all models, since the code paths are the same? + # retrieve all regularization terms + collect_regularization_terms(self.model) - # assert that the regularization term has been reset - assert self.model.regularizer.regularization_term == torch.zeros(1, dtype=torch.float, device=self.model.device) + # check that all terms are reset + for module in self.model.modules(): + if isinstance(module, Regularizer): + assert module.regularization_term == 0.0 def test_post_parameter_update(self): """Test whether post_parameter_update correctly enforces model constraints.""" From ca68f030547c6d40f5b426f222351b8c873ccf09 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 14:37:58 +0100 Subject: [PATCH 358/690] Fix docstring --- tests/test_models.py | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 79ec22f6c5..62fa8d88f6 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -98,8 +98,7 @@ def setUp(self) -> None: ) def test_add_weight_regularizer_non_existing(self): - """Test add_weight_regularizer.""" - # try to add regularizer to non-existing weight + """Test that an assertion is raised for add_weight_regularizer for a non-existing weight.""" with self.assertRaises(ValueError): self.model.add_weight_regularizer( parameter_name="this.weight.does.not.exist", @@ -116,12 +115,10 @@ def test_add_weight_regularizer(self): self.model.linear, ) - regularizer = LpRegularizer() - # try to add regularizer to existing weight self.model.add_weight_regularizer( parameter_name="linear.weight", - regularizer=regularizer, + regularizer=LpRegularizer(), ) # check it gets found by collect From 8c7727868c6151fdf43fd9e31b6c3ddc868486f4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 14:47:48 +0100 Subject: [PATCH 359/690] Add test for collect_regularization_terms --- tests/test_models.py | 13 +-------- tests/test_regularizers.py | 60 +++++++++++++++++++++++++++++++++++++- 2 files changed, 60 insertions(+), 13 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 62fa8d88f6..8aeb7f156a 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -41,7 +41,7 @@ symmetric_edge_weights, ) from pykeen.nn import RepresentationModule -from pykeen.regularizers import LpRegularizer, Regularizer, collect_regularization_terms +from pykeen.regularizers import LpRegularizer, collect_regularization_terms from pykeen.training import LCWATrainingLoop, SLCWATrainingLoop, TrainingLoop from pykeen.triples import TriplesFactory from pykeen.utils import all_in_bounds, clamp_norm, project_entity, set_random_seed @@ -425,17 +425,6 @@ def test_has_hpo_defaults(self): else: self.assertIsInstance(d, dict) - def test_collect_regularization_terms(self): - """Test whether collect_regularization_terms resets all regularization terms.""" - # TODO: Does this need to be tested for all models, since the code paths are the same? - # retrieve all regularization terms - collect_regularization_terms(self.model) - - # check that all terms are reset - for module in self.model.modules(): - if isinstance(module, Regularizer): - assert module.regularization_term == 0.0 - def test_post_parameter_update(self): """Test whether post_parameter_update correctly enforces model constraints.""" # do one optimization step diff --git a/tests/test_regularizers.py b/tests/test_regularizers.py index 44ab0f9c75..7f8b3be854 100644 --- a/tests/test_regularizers.py +++ b/tests/test_regularizers.py @@ -5,13 +5,15 @@ import logging import unittest from typing import Any, ClassVar, Dict, Optional, Type +from unittest.mock import MagicMock import torch from torch import nn from torch.nn import functional from pykeen.datasets import Nations -from pykeen.models import ConvKB, RESCAL +from pykeen.models import ConvKB, ERModel, RESCAL +from pykeen.nn import EmbeddingSpecification from pykeen.regularizers import ( CombinedRegularizer, LpRegularizer, NoRegularizer, PowerSumRegularizer, Regularizer, TransHRegularizer, collect_regularization_terms, @@ -306,3 +308,59 @@ def _help_test_regularizer(self, regularizer: Regularizer, n_tensors: int = 3): regularizer.pop_regularization_term() self.assertFalse(regularizer.updated) assert 0.0 == regularizer.regularization_term + + +def test_collect_regularization_terms(): + """Test whether collect_regularization_terms finds and resets all regularization terms.""" + regularizers = [ + LpRegularizer(), + PowerSumRegularizer(), + LpRegularizer(p=1, normalize=True, apply_only_once=True), + PowerSumRegularizer(normalize=True), + ] + model = ERModel( + triples_factory=MagicMock(), + interaction=MagicMock(), + entity_representations=EmbeddingSpecification( + regularizer=regularizers[0], + ).make(num_embeddings=3, embedding_dim=2, shape=None), + relation_representations=EmbeddingSpecification( + regularizer=regularizers[1], + ).make(num_embeddings=3, embedding_dim=2, shape=None), + ) + + # add weighted modules + linear = nn.Linear(3, 2) + model.sub_module = nn.ModuleList([ + nn.Sequential( + linear, + nn.Linear(2, 3), + ), + nn.BatchNorm1d(2), + linear, # one module occuring twice + ]) + + # add weight regularizer + model.add_weight_regularizer( + parameter_name="sub_module.0.0.bias", + regularizer=regularizers[2], + ) + model.add_weight_regularizer( + parameter_name="entity_representations.0._embeddings.weight", + regularizer=regularizers[3], + ) + + # retrieve all regularization terms + collect_regularization_terms(model) + + # check that all terms are reset + found_regularizers = set() + for module in model.modules(): + if isinstance(module, Regularizer): + term = module.regularization_term + assert isinstance(term, float) + assert term == 0.0 + found_regularizers.add(id(module)) + + # check that all regularizers were found + assert found_regularizers == set(map(id, regularizers)) From 5db61c46fc673501a88c70c4c7ce5a6b8128dbda Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 14:49:22 +0100 Subject: [PATCH 360/690] Fix apply_only_once for TransH --- src/pykeen/regularizers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/regularizers.py b/src/pykeen/regularizers.py index 3aa51bd955..1784805141 100644 --- a/src/pykeen/regularizers.py +++ b/src/pykeen/regularizers.py @@ -204,7 +204,7 @@ def __init__( ): # The regularization in TransH enforces the defined soft constraints that should computed only for every batch. # Therefore, apply_only_once is always set to True. - super().__init__(weight=weight, apply_only_once=False, parameters=[]) + super().__init__(weight=weight, apply_only_once=True, parameters=[]) self.normal_vector_embeddings = normal_vector_embeddings self.relation_embeddings = relation_embeddings self.entity_embeddings = entity_embeddings From 5c2ecd9df93d6b3d6fd1ea479701b742331d5139 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 14:50:09 +0100 Subject: [PATCH 361/690] Fix typo in ComplexLiteralCombination --- src/pykeen/models/multimodal/complex_literal.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index dfe74a90f7..f0147cd0f1 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -50,7 +50,7 @@ def forward( x, literal = x[..., :self.embedding_dim], x[..., self.embedding_dim:] x_re, x_im = split_complex(x) x_re = self.real(torch.cat([x_re, literal], dim=-1)) - x_im = self.real(torch.cat([x_im, literal], dim=-1)) + x_im = self.imag(torch.cat([x_im, literal], dim=-1)) return combine_complex(x_re=x_re, x_im=x_im) From 31136b4a5ab99d494b13316868bbaa2444910e85 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 14:52:28 +0100 Subject: [PATCH 362/690] Add explanation for canonical shape --- src/pykeen/nn/emb.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index 101412918c..dba2ac4fe8 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -48,12 +48,16 @@ def get_in_canonical_shape( ) -> torch.FloatTensor: """Get representations in canonical shape. + The canonical shape is given by (batch_size, 1, *) if indices is not None, where batch_size=len(indices), + or (1, num, *) if indices is None with num equal to the total number of embeddings. + + :param indices: The indices. If None, return all embeddings. :param reshape_dim: Optionally reshape the last dimension. - :return: shape: (batch_size, num_embeddings, d) + :return: shape: (batch_size, num_embeddings, *) """ x = self(indices=indices) if indices is None: From a2364f1c9b1692f4c408a68f8eaf01b4e4049e15 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 15:16:59 +0100 Subject: [PATCH 363/690] Pass CI --- src/pykeen/nn/emb.py | 6 +++--- tests/test_models.py | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index dba2ac4fe8..41e056e74e 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -48,8 +48,8 @@ def get_in_canonical_shape( ) -> torch.FloatTensor: """Get representations in canonical shape. - The canonical shape is given by (batch_size, 1, *) if indices is not None, where batch_size=len(indices), - or (1, num, *) if indices is None with num equal to the total number of embeddings. + The canonical shape is given by (batch_size, 1, ``*``) if indices is not None, where batch_size=len(indices), + or (1, num, ``*``) if indices is None with num equal to the total number of embeddings. :param indices: @@ -57,7 +57,7 @@ def get_in_canonical_shape( :param reshape_dim: Optionally reshape the last dimension. - :return: shape: (batch_size, num_embeddings, *) + :return: shape: (batch_size, num_embeddings, ``*``) """ x = self(indices=indices) if indices is None: diff --git a/tests/test_models.py b/tests/test_models.py index 8aeb7f156a..3098601e7e 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -91,7 +91,7 @@ class ERModelTests(unittest.TestCase): """Test basic functionality of ERModel.""" def setUp(self) -> None: - """Setup the test instance.""" + """Set up the test instance.""" self.model = ERModel( triples_factory=MagicMock(), interaction=MagicMock(), From 619e25e0c4abc261639cf40186f09b8fa9e0bd0f Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 15:39:52 +0100 Subject: [PATCH 364/690] Remove regularizer from CLI --- src/pykeen/models/cli/builders.py | 37 ++++++++----------------------- 1 file changed, 9 insertions(+), 28 deletions(-) diff --git a/src/pykeen/models/cli/builders.py b/src/pykeen/models/cli/builders.py index e4f5d41d09..83f6b50431 100644 --- a/src/pykeen/models/cli/builders.py +++ b/src/pykeen/models/cli/builders.py @@ -46,34 +46,25 @@ def build_cli_from_cls(model: Type[Model]) -> click.Command: # noqa: D202 def _decorate_model_kwargs(command: click.Command) -> click.Command: for name, annotation in model.__init__.__annotations__.items(): + parameter = signature.parameters[name] + if name in _SKIP_ARGS or annotation in _SKIP_ANNOTATIONS: continue - if annotation == Union[None, str, Regularizer]: # a model that has preset regularization - parameter = signature.parameters[name] - option = click.option( - '--regularizer', - type=click.Choice(regularizers), - default=parameter.default, - show_default=True, - help=f'The name of the regularizer preset for {model.__name__}', - ) - elif name in CLI_OPTIONS: option = CLI_OPTIONS[name] elif annotation in {Optional[int], Optional[str]}: option = click.option(f'--{name.replace("_", "-")}', type=_OPTIONAL_MAP[annotation]) - else: - parameter = signature.parameters[name] - if parameter.default is None: - logger.warning( - f'Missing handler in {model.__name__} for {name}: ' - f'type={annotation} default={parameter.default}', - ) - continue + elif parameter.default is None: + logger.warning( + f'Missing handler in {model.__name__} for {name}: ' + f'type={annotation} default={parameter.default}', + ) + continue + else: option = click.option(f'--{name.replace("_", "-")}', type=annotation, default=parameter.default) try: @@ -84,13 +75,6 @@ def _decorate_model_kwargs(command: click.Command) -> click.Command: return command - regularizer_option = click.option( - '--regularizer', - type=click.Choice(regularizers), - default=model.regularizer_default.get_normalized_name() if model.regularizer_default else 'no', - show_default=True, - ) - @click.command(help=f'CLI for {model.__name__}', name=model.__name__.lower()) @options.device_option @options.dataset_option @@ -98,7 +82,6 @@ def _decorate_model_kwargs(command: click.Command) -> click.Command: @options.testing_option @options.valiadation_option @options.optimizer_option - @regularizer_option @options.training_loop_option @options.number_epochs_option @options.batch_size_option @@ -117,7 +100,6 @@ def main( device, training_loop, optimizer, - regularizer, number_epochs, batch_size, learning_rate, @@ -156,7 +138,6 @@ def main( device=device, model=model, model_kwargs=model_kwargs, - regularizer=regularizer, dataset=dataset, training=training_triples_factory, testing=testing_triples_factory or training_triples_factory, From da04fa2668a124985c51bfee8d0fa31626c5df4a Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 15:40:07 +0100 Subject: [PATCH 365/690] Clean up tests --- src/pykeen/utils.py | 3 ++- tests/test_early_stopping.py | 1 + tests/test_interactions.py | 6 +++--- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 8a5be593b7..040a49b3a6 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -605,7 +605,8 @@ def extended_einsum( mod_ops, mod_t = [], [] for op, t in zip(lhs.split(","), tensors): mod_op = "" - assert len(op) == len(t.shape) + if len(op) != len(t.shape): + raise ValueError(f'Shapes not equal: op={op} and t.shape={t.shape}') for i, c in reversed(list(enumerate(op))): if t.shape[i] == 1: t = t.squeeze(dim=i) diff --git a/tests/test_early_stopping.py b/tests/test_early_stopping.py index 6c8a4cd651..36fd95650a 100644 --- a/tests/test_early_stopping.py +++ b/tests/test_early_stopping.py @@ -243,6 +243,7 @@ def test_early_stopping(self): num_epochs=self.max_num_epochs, batch_size=self.batch_size, stopper=stopper, + use_tqdm=False, ) self.assertEqual(stopper.number_results, len(losses) // stopper.frequency) self.assertEqual(self.stop_epoch, len(losses), msg='Did not stop early like it should have') diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 35e5735b88..5ae7444609 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -123,7 +123,7 @@ def test_score_h_slicing(self): ) scores = self.instance.score_h(all_entities=h, r=r, t=t, slice_size=self.num_entities // 2 + 1) scores_no_slice = self.instance.score_h(all_entities=h, r=r, t=t, slice_size=None) - assert torch.allclose(scores, scores_no_slice) + assert torch.allclose(scores, scores_no_slice), f'Differences: {scores - scores_no_slice}' def test_score_r(self): """Test score_r.""" @@ -152,7 +152,7 @@ def test_score_r_slicing(self): ) scores = self.instance.score_r(h=h, all_relations=r, t=t, slice_size=self.num_relations // 2 + 1) scores_no_slice = self.instance.score_r(h=h, all_relations=r, t=t, slice_size=None) - assert torch.allclose(scores, scores_no_slice) + assert torch.allclose(scores, scores_no_slice), f'Differences: {scores - scores_no_slice}' def test_score_t(self): """Test score_t.""" @@ -175,7 +175,7 @@ def test_score_t_slicing(self): ) scores = self.instance.score_t(h=h, r=r, all_entities=t, slice_size=self.num_entities // 2 + 1) scores_no_slice = self.instance.score_t(h=h, r=r, all_entities=t, slice_size=None) - assert torch.allclose(scores, scores_no_slice) + assert torch.allclose(scores, scores_no_slice), f'Differences: {scores - scores_no_slice}' def test_forward(self): """Test forward.""" From dc53ea3a12ed1d0ab4f5892b1d0ace01118c926a Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 15:40:28 +0100 Subject: [PATCH 366/690] Fix TransD test and meta model test --- tests/test_models.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 3098601e7e..d49ceae736 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -40,7 +40,7 @@ inverse_outdegree_edge_weights, symmetric_edge_weights, ) -from pykeen.nn import RepresentationModule +from pykeen.nn import Embedding, RepresentationModule from pykeen.regularizers import LpRegularizer, collect_regularization_terms from pykeen.training import LCWATrainingLoop, SLCWATrainingLoop, TrainingLoop from pykeen.triples import TriplesFactory @@ -48,12 +48,13 @@ SKIP_MODULES = { Model.__name__, - 'DummyModel', + ERModel.__name__, LiteralModel.__name__, DoubleRelationEmbeddingModel.__name__, SingleVectorEmbeddingModel.__name__, TwoVectorEmbeddingModel.__name__, TwoSideEmbeddingModel.__name__, + 'DummyModel', 'MockModel', 'models', 'get_model_cls', @@ -881,19 +882,25 @@ def _check_constraints(self): def test_project_entity(self): """Test _project_entity.""" + self.assertIsInstance(self.model, pykeen.models.TransD) + # random entity embeddings & projections e = torch.rand(1, self.model.num_entities, self.embedding_dim, generator=self.generator) e = clamp_norm(e, maxnorm=1, p=2, dim=-1) e_p = torch.rand(1, self.model.num_entities, self.embedding_dim, generator=self.generator) # random relation embeddings & projections - r_p = torch.rand(self.batch_size, 1, self.model.relation_dim, generator=self.generator) + relation_rep: Embedding = self.model.relation_representations[0] + self.assertIsInstance(relation_rep, Embedding) + relation_dim = relation_rep.embedding_dim + + r_p = torch.rand(self.batch_size, 1, relation_dim, generator=self.generator) # project e_bot = project_entity(e=e, e_p=e_p, r_p=r_p) # check shape: - assert e_bot.shape == (self.batch_size, self.model.num_entities, self.model.relation_dim) + assert e_bot.shape == (self.batch_size, self.model.num_entities, relation_dim) # check normalization assert (torch.norm(e_bot, dim=-1, p=2) <= 1.0 + 1.0e-06).all() From 3698e3e53c44feab2aaed9bcaf096ea7f38e60d9 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 15:41:15 +0100 Subject: [PATCH 367/690] Update builders.py --- src/pykeen/models/cli/builders.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/pykeen/models/cli/builders.py b/src/pykeen/models/cli/builders.py index 83f6b50431..14aa43a119 100644 --- a/src/pykeen/models/cli/builders.py +++ b/src/pykeen/models/cli/builders.py @@ -6,7 +6,7 @@ import json import logging import sys -from typing import Optional, Type, Union +from typing import Optional, Type import click from torch import nn @@ -14,7 +14,6 @@ from . import options from .options import CLI_OPTIONS from ..base import Model -from ...regularizers import Regularizer, regularizers __all__ = [ 'build_cli_from_cls', From 5073afc83edb9c0c7472ebd8a4ecb8b7953edafd Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 15:47:38 +0100 Subject: [PATCH 368/690] Remembered why it's like this --- src/pykeen/models/cli/builders.py | 17 ++++++++--------- 1 file changed, 8 insertions(+), 9 deletions(-) diff --git a/src/pykeen/models/cli/builders.py b/src/pykeen/models/cli/builders.py index 14aa43a119..30a1dea349 100644 --- a/src/pykeen/models/cli/builders.py +++ b/src/pykeen/models/cli/builders.py @@ -45,8 +45,6 @@ def build_cli_from_cls(model: Type[Model]) -> click.Command: # noqa: D202 def _decorate_model_kwargs(command: click.Command) -> click.Command: for name, annotation in model.__init__.__annotations__.items(): - parameter = signature.parameters[name] - if name in _SKIP_ARGS or annotation in _SKIP_ANNOTATIONS: continue @@ -56,14 +54,15 @@ def _decorate_model_kwargs(command: click.Command) -> click.Command: elif annotation in {Optional[int], Optional[str]}: option = click.option(f'--{name.replace("_", "-")}', type=_OPTIONAL_MAP[annotation]) - elif parameter.default is None: - logger.warning( - f'Missing handler in {model.__name__} for {name}: ' - f'type={annotation} default={parameter.default}', - ) - continue - else: + parameter = signature.parameters[name] + if parameter.default is None: + logger.warning( + f'Missing handler in {model.__name__} for {name}: ' + f'type={annotation} default={parameter.default}', + ) + continue + option = click.option(f'--{name.replace("_", "-")}', type=annotation, default=parameter.default) try: From e02e8be74d5e4081ad5ff6f749231c2f8c27adfe Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 15:49:21 +0100 Subject: [PATCH 369/690] Cut regularizer out of pipeline and HPO --- src/pykeen/hpo/hpo.py | 34 ++++++++++++++++++---------------- src/pykeen/pipeline.py | 18 ++++++++---------- 2 files changed, 26 insertions(+), 26 deletions(-) diff --git a/src/pykeen/hpo/hpo.py b/src/pykeen/hpo/hpo.py index 0ec20b3cc5..ebc0983221 100644 --- a/src/pykeen/hpo/hpo.py +++ b/src/pykeen/hpo/hpo.py @@ -149,13 +149,13 @@ def __call__(self, trial: Trial) -> Optional[float]: kwargs_ranges=self.loss_kwargs_ranges, ) # 4. Regularizer - _regularizer_kwargs = _get_kwargs( - trial=trial, - prefix='regularizer', - default_kwargs_ranges=self.regularizer.hpo_default, - kwargs=self.regularizer_kwargs, - kwargs_ranges=self.regularizer_kwargs_ranges, - ) + # _regularizer_kwargs = _get_kwargs( + # trial=trial, + # prefix='regularizer', + # default_kwargs_ranges=self.regularizer.hpo_default, + # kwargs=self.regularizer_kwargs, + # kwargs_ranges=self.regularizer_kwargs_ranges, + # ) # 5. Optimizer _optimizer_kwargs = _get_kwargs( trial=trial, @@ -205,8 +205,8 @@ def __call__(self, trial: Trial) -> Optional[float]: loss=self.loss, loss_kwargs=_loss_kwargs, # 4. Regularizer - regularizer=self.regularizer, - regularizer_kwargs=_regularizer_kwargs, + # regularizer=self.regularizer, + # regularizer_kwargs=_regularizer_kwargs, clear_optimizer=True, # 5. Optimizer optimizer=self.optimizer, @@ -612,13 +612,15 @@ def hpo_pipeline( study.set_user_attr('loss', normalize_string(loss.__name__, suffix=_LOSS_SUFFIX)) logger.info(f'Using loss: {loss}') # 4. Regularizer - regularizer: Type[Regularizer] = ( - model.regularizer_default - if regularizer is None else - get_regularizer_cls(regularizer) - ) - study.set_user_attr('regularizer', regularizer.get_normalized_name()) - logger.info(f'Using regularizer: {regularizer}') + if regularizer is not None: + logger.warning('Usage of the regularizer with the HPO is currently under maitenance.') + # regularizer: Type[Regularizer] = ( + # model.regularizer_default + # if regularizer is None else + # get_regularizer_cls(regularizer) + # ) + # study.set_user_attr('regularizer', regularizer.get_normalized_name()) + # logger.info(f'Using regularizer: {regularizer}') # 5. Optimizer optimizer: Type[Optimizer] = get_optimizer_cls(optimizer) study.set_user_attr('optimizer', normalize_string(optimizer.__name__)) diff --git a/src/pykeen/pipeline.py b/src/pykeen/pipeline.py index 472c9c76c8..ad45c789ca 100644 --- a/src/pykeen/pipeline.py +++ b/src/pykeen/pipeline.py @@ -871,15 +871,13 @@ def pipeline( # noqa: C901 model_kwargs.setdefault('random_seed', random_seed) if regularizer is not None: + logger.warning('Specification of the regularizer from the pipeline() is currently under maitenance') # FIXME this should never happen. - if 'regularizer' in model_kwargs: - logger.warning('Can not specify regularizer in kwargs and model_kwargs. removing from model_kwargs') - del model_kwargs['regularizer'] - regularizer_cls: Type[Regularizer] = get_regularizer_cls(regularizer) - model_kwargs['regularizer'] = regularizer_cls( - device=device, - **(regularizer_kwargs or {}), - ) + # if 'regularizer' in model_kwargs: + # logger.warning('Can not specify regularizer in kwargs and model_kwargs. removing from model_kwargs') + # del model_kwargs['regularizer'] + # regularizer_cls: Type[Regularizer] = get_regularizer_cls(regularizer) + # model_kwargs['regularizer'] = regularizer_cls(**(regularizer_kwargs or {})) if loss is not None: if 'loss' in model_kwargs: # FIXME @@ -983,8 +981,8 @@ def pipeline( # noqa: C901 logging.debug(f"model_kwargs: {model_kwargs}") logging.debug(f"loss: {loss}") logging.debug(f"loss_kwargs: {loss_kwargs}") - logging.debug(f"regularizer: {regularizer}") - logging.debug(f"regularizer_kwargs: {regularizer_kwargs}") + # logging.debug(f"regularizer: {regularizer}") + # logging.debug(f"regularizer_kwargs: {regularizer_kwargs}") logging.debug(f"optimizer: {optimizer}") logging.debug(f"optimizer_kwargs: {optimizer_kwargs}") logging.debug(f"training_loop: {training_loop}") From 8b24039667f6d03031f6c484f830e7cdd842c107 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 15:50:56 +0100 Subject: [PATCH 370/690] More regularizer cleaning --- src/pykeen/hpo/hpo.py | 8 ++++---- src/pykeen/pipeline.py | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/src/pykeen/hpo/hpo.py b/src/pykeen/hpo/hpo.py index ebc0983221..a2c2de0921 100644 --- a/src/pykeen/hpo/hpo.py +++ b/src/pykeen/hpo/hpo.py @@ -26,7 +26,7 @@ from ..models.base import Model from ..optimizers import Optimizer, get_optimizer_cls, optimizers_hpo_defaults from ..pipeline import pipeline, replicate_pipeline_from_config -from ..regularizers import Regularizer, get_regularizer_cls +from ..regularizers import Regularizer from ..sampling import NegativeSampler, get_negative_sampler_cls from ..stoppers import EarlyStopper, Stopper, get_stopper_cls from ..trackers import ResultTracker, get_result_tracker_cls @@ -57,7 +57,7 @@ class Objective: dataset: Union[None, str, Type[DataSet]] # 1. model: Type[Model] # 2. loss: Type[Loss] # 3. - regularizer: Type[Regularizer] # 4. + # regularizer: Type[Regularizer] # 4. optimizer: Type[Optimizer] # 5. training_loop: Type[TrainingLoop] # 6. evaluator: Type[Evaluator] # 8. @@ -77,8 +77,8 @@ class Objective: loss_kwargs: Optional[Mapping[str, Any]] = None loss_kwargs_ranges: Optional[Mapping[str, Any]] = None # 4. Regularizer - regularizer_kwargs: Optional[Mapping[str, Any]] = None - regularizer_kwargs_ranges: Optional[Mapping[str, Any]] = None + # regularizer_kwargs: Optional[Mapping[str, Any]] = None + # regularizer_kwargs_ranges: Optional[Mapping[str, Any]] = None # 5. Optimizer optimizer_kwargs: Optional[Mapping[str, Any]] = None optimizer_kwargs_ranges: Optional[Mapping[str, Any]] = None diff --git a/src/pykeen/pipeline.py b/src/pykeen/pipeline.py index ad45c789ca..bddc96a13a 100644 --- a/src/pykeen/pipeline.py +++ b/src/pykeen/pipeline.py @@ -183,7 +183,7 @@ from .models import get_model_cls from .models.base import Model from .optimizers import get_optimizer_cls -from .regularizers import Regularizer, get_regularizer_cls +from .regularizers import Regularizer from .sampling import NegativeSampler, get_negative_sampler_cls from .stoppers import EarlyStopper, Stopper, get_stopper_cls from .trackers import ResultTracker, get_result_tracker_cls From 7403b1ba0288fe5f966ac3bb5b7e1ba73d4d5118 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 15:56:56 +0100 Subject: [PATCH 371/690] Update hpo.py --- src/pykeen/hpo/hpo.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/hpo/hpo.py b/src/pykeen/hpo/hpo.py index a2c2de0921..0be641f9b3 100644 --- a/src/pykeen/hpo/hpo.py +++ b/src/pykeen/hpo/hpo.py @@ -57,7 +57,7 @@ class Objective: dataset: Union[None, str, Type[DataSet]] # 1. model: Type[Model] # 2. loss: Type[Loss] # 3. - # regularizer: Type[Regularizer] # 4. + regularizer: Type[Regularizer] # 4. optimizer: Type[Optimizer] # 5. training_loop: Type[TrainingLoop] # 6. evaluator: Type[Evaluator] # 8. @@ -77,8 +77,8 @@ class Objective: loss_kwargs: Optional[Mapping[str, Any]] = None loss_kwargs_ranges: Optional[Mapping[str, Any]] = None # 4. Regularizer - # regularizer_kwargs: Optional[Mapping[str, Any]] = None - # regularizer_kwargs_ranges: Optional[Mapping[str, Any]] = None + regularizer_kwargs: Optional[Mapping[str, Any]] = None + regularizer_kwargs_ranges: Optional[Mapping[str, Any]] = None # 5. Optimizer optimizer_kwargs: Optional[Mapping[str, Any]] = None optimizer_kwargs_ranges: Optional[Mapping[str, Any]] = None From 09a86f30e1afc912051ecff8fe5fea5927461bab Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 15:59:15 +0100 Subject: [PATCH 372/690] Remove deprecated pipeline test --- tests/test_pipeline.py | 24 ------------------------ 1 file changed, 24 deletions(-) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 70c06b3dc0..39d67a69e1 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -105,27 +105,3 @@ def test_predict_all_with_novelties(self): self.assertEqual(possible, len(all_df.index)) self.assertEqual(self.model.triples_factory.num_triples, all_df['in_training'].sum()) self.assertEqual(self.testing_mapped_triples.shape[0], all_df['in_testing'].sum()) - - -class TestAttributes(unittest.TestCase): - """Test that the keywords given to the pipeline make it through.""" - - def test_specify_regularizer(self): - """Test a pipeline that uses a regularizer.""" - for regularizer, cls in [ - (None, pykeen.regularizers.NoRegularizer), - ('no', pykeen.regularizers.NoRegularizer), - (NoRegularizer, pykeen.regularizers.NoRegularizer), - ('powersum', pykeen.regularizers.PowerSumRegularizer), - ('lp', pykeen.regularizers.LpRegularizer), - ]: - with self.subTest(regularizer=regularizer): - pipeline_result = pipeline( - model='TransE', - dataset='Nations', - regularizer=regularizer, - training_kwargs=dict(num_epochs=1), - ) - self.assertIsInstance(pipeline_result, PipelineResult) - self.assertIsInstance(pipeline_result.model, Model) - self.assertIsInstance(pipeline_result.model.regularizer, cls) From d4755388ae1ad00c9b27f6c65a9187acbd1aaa08 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 15:59:31 +0100 Subject: [PATCH 373/690] Update test_pipeline.py --- tests/test_pipeline.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 39d67a69e1..43a382b592 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -6,11 +6,8 @@ import pandas as pd -import pykeen.regularizers from pykeen.datasets import Nations -from pykeen.models.base import Model -from pykeen.pipeline import PipelineResult, pipeline -from pykeen.regularizers import NoRegularizer +from pykeen.pipeline import pipeline class TestPipeline(unittest.TestCase): From 8eb6bd0c15e37ad173dc34bcf9e9215ffc3e3e94 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 16:21:44 +0100 Subject: [PATCH 374/690] Add predict with sigmoid default value for TransH --- src/pykeen/experiments/transh/wang2014_transh_fb15k.json | 1 + src/pykeen/experiments/transh/wang2014_transh_wn18.json | 1 + 2 files changed, 2 insertions(+) diff --git a/src/pykeen/experiments/transh/wang2014_transh_fb15k.json b/src/pykeen/experiments/transh/wang2014_transh_fb15k.json index b33724de1c..c685ca8c99 100644 --- a/src/pykeen/experiments/transh/wang2014_transh_fb15k.json +++ b/src/pykeen/experiments/transh/wang2014_transh_fb15k.json @@ -6,6 +6,7 @@ "dataset": "fb15k", "model": "TransH", "model_kwargs": { + "predict_with_sigmoid": false, "embedding_dim": 100, "scoring_fct_norm": 2 }, diff --git a/src/pykeen/experiments/transh/wang2014_transh_wn18.json b/src/pykeen/experiments/transh/wang2014_transh_wn18.json index 87e710b874..4902d4115b 100644 --- a/src/pykeen/experiments/transh/wang2014_transh_wn18.json +++ b/src/pykeen/experiments/transh/wang2014_transh_wn18.json @@ -6,6 +6,7 @@ "dataset": "wn18", "model": "TransH", "model_kwargs": { + "predict_with_sigmoid": false, "embedding_dim": 50, "scoring_fct_norm": 2 }, From 40fd54e328c64b8ee77a439f1d8f30580d1c7423 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 16:36:22 +0100 Subject: [PATCH 375/690] Remove obsolete can_slice test --- tests/test_models.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index d49ceae736..57e5d1f33d 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -734,12 +734,6 @@ class _BaseNTNTest(_ModelTestCase, unittest.TestCase): model_cls = pykeen.models.NTN - def test_can_slice(self): - """Test that the slicing properties are calculated correctly.""" - self.assertTrue(self.model.can_slice_h, msg='Unable to slice on heads') - self.assertTrue(self.model.can_slice_r, msg='Unable to slice on relations') - self.assertTrue(self.model.can_slice_t, msg='Unable to slice on tails') - class TestNTNLowMemory(_BaseNTNTest): """Test the NTN model with automatic memory optimization.""" From 359269d2084c3b7995e43bf11d0777cc37946e10 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 16:42:15 +0100 Subject: [PATCH 376/690] Fix _CustomRepresentations --- tests/test_models.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 57e5d1f33d..287c21b070 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -7,7 +7,7 @@ import tempfile import traceback import unittest -from typing import Any, ClassVar, Mapping, Optional, Type +from typing import Any, ClassVar, Mapping, Optional, Sequence, Type from unittest.mock import MagicMock, patch import numpy @@ -68,24 +68,27 @@ class _CustomRepresentations(RepresentationModule): """A custom representation module with minimal implementation.""" - def __init__(self, num_entities: int, embedding_dim: int = 2): + def __init__(self, num_entities: int, shape: Sequence[int]): super().__init__() self.num_embeddings = num_entities - self.embedding_dim = embedding_dim - self.x = nn.Parameter(torch.rand(embedding_dim)) + self.shape = shape + self.x = nn.Parameter(torch.rand(numpy.prod(shape))) def forward(self, indices: Optional[torch.LongTensor] = None) -> torch.FloatTensor: n = self.num_embeddings if indices is None else indices.shape[0] - return self.x.unsqueeze(dim=0).repeat(n, 1) + return self.x.unsqueeze(dim=0).repeat(n, 1).view(-1, *self.shape) def get_in_canonical_shape( self, indices: Optional[torch.LongTensor] = None, + reshape_dim: Optional[Sequence[int]] = None, ) -> torch.FloatTensor: x = self(indices=indices) if indices is None: - return x.unsqueeze(dim=0) - return x.unsqueeze(dim=1) + x = x.unsqueeze(dim=0) + else: + x = x.unsqueeze(dim=1) + return x class ERModelTests(unittest.TestCase): @@ -527,7 +530,7 @@ def test_custom_representations(self): self.model.entity_representations = nn.ModuleList([ _CustomRepresentations( num_entities=self.factory.num_entities, - embedding_dim=er.embedding_dim, + shape=er.shape, ) for er in old_entity_reps ]) @@ -535,7 +538,7 @@ def test_custom_representations(self): self.model.relation_representations = nn.ModuleList([ _CustomRepresentations( num_entities=self.factory.num_entities, - embedding_dim=er.embedding_dim, + shape=er.shape, ) for er in old_relation_reps ]) From 80020d7a531ef1482847deee07241f6d2f7de73c Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 16:53:36 +0100 Subject: [PATCH 377/690] Remove slicing checks This is possible everywhere now --- src/pykeen/evaluation/evaluator.py | 14 ----------- src/pykeen/models/base.py | 40 +++--------------------------- src/pykeen/training/lcwa.py | 17 ------------- 3 files changed, 4 insertions(+), 67 deletions(-) diff --git a/src/pykeen/evaluation/evaluator.py b/src/pykeen/evaluation/evaluator.py index 6af4860259..ab562ec3cf 100644 --- a/src/pykeen/evaluation/evaluator.py +++ b/src/pykeen/evaluation/evaluator.py @@ -258,7 +258,6 @@ def _param_size_search( values_dict[key] = start_value values_dict['slice_size'] = None elif key == 'slice_size': - self._check_slicing_availability(model, batch_size=1) values_dict[key] = start_value values_dict['batch_size'] = 1 else: @@ -320,19 +319,6 @@ def _param_size_search( return values_dict[key], evaluated_once - @staticmethod - def _check_slicing_availability(model: Model, batch_size: int) -> None: - # Test if slicing is implemented for the required functions of this model - if model.triples_factory.create_inverse_triples: - if not model.can_slice_t: - raise MemoryError(f"The current model can't be evaluated on this hardware with these parameters, as " - f"evaluation batch_size={batch_size} is too big and slicing is not implemented for " - f"this model yet.") - elif not model.can_slice_t or not model.can_slice_h: - raise MemoryError(f"The current model can't be evaluated on this hardware with these parameters, as " - f"evaluation batch_size={batch_size} is too big and slicing is not implemented for this " - f"model yet.") - def create_sparse_positive_filter_( hrt_batch: MappedTriples, diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 52de2777c9..8f956ceed4 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -3,7 +3,6 @@ """Base module for all KGE models.""" import functools -import inspect import itertools as itt import logging from abc import ABC, abstractmethod @@ -194,10 +193,6 @@ def _process_remove_known(df: pd.DataFrame, remove_known: bool, testing: Optiona return df -def _can_slice(fn) -> bool: - return 'slice_size' in inspect.getfullargspec(fn).args - - def _track_hyperparameters(cls: Type['Model']) -> None: """Initialize the subclass while keeping track of hyper-parameters.""" # Keep track of the hyper-parameters that are used across all @@ -329,21 +324,6 @@ def __init_subclass__(cls, autoreset: bool = True, **kwargs): # noqa:D105 _track_hyperparameters(cls) _add_post_reset_parameters(cls) - @property - def can_slice_h(self) -> bool: - """Whether score_h supports slicing.""" - return _can_slice(self.score_h) - - @property - def can_slice_r(self) -> bool: - """Whether score_r supports slicing.""" - return _can_slice(self.score_r) - - @property - def can_slice_t(self) -> bool: - """Whether score_t supports slicing.""" - return _can_slice(self.score_t) - @property def modules_not_supporting_sub_batching(self) -> Collection[nn.Module]: """Return all modules not supporting sub-batching.""" @@ -491,10 +471,7 @@ def predict_scores_all_tails( """ # Enforce evaluation mode self.eval() - if slice_size is not None and self.can_slice_t: - scores = self.score_t(hr_batch, slice_size=slice_size) # type: ignore - else: - scores = self.score_t(hr_batch) + scores = self.score_t(hr_batch, slice_size=slice_size) # type: ignore if self.predict_with_sigmoid: scores = torch.sigmoid(scores) return scores @@ -630,10 +607,7 @@ def predict_scores_all_relations( """ # Enforce evaluation mode self.eval() - if slice_size is not None and self.can_slice_r: - scores = self.score_r(ht_batch, slice_size=slice_size) # type: ignore - else: - scores = self.score_r(ht_batch) + scores = self.score_r(ht_batch, slice_size=slice_size) # type: ignore if self.predict_with_sigmoid: scores = torch.sigmoid(scores) return scores @@ -666,10 +640,7 @@ def predict_scores_all_heads( for a (tail, inverse_relation) pair. ''' if not self.triples_factory.create_inverse_triples: - if slice_size is not None and self.can_slice_h: - scores = self.score_h(rt_batch, slice_size=slice_size) # type: ignore - else: - scores = self.score_h(rt_batch) + scores = self.score_h(rt_batch, slice_size=slice_size) if self.predict_with_sigmoid: scores = torch.sigmoid(scores) return scores @@ -691,10 +662,7 @@ def predict_scores_all_heads( # The score_t function requires (entity, relation) pairs instead of (relation, entity) pairs rt_batch_cloned = rt_batch_cloned.flip(1) - if slice_size is not None and self.can_slice_t: - scores = self.score_t(rt_batch_cloned, slice_size=slice_size) # type: ignore - else: - scores = self.score_t(rt_batch_cloned) + scores = self.score_t(rt_batch_cloned, slice_size=slice_size) # type: ignore if self.predict_with_sigmoid: scores = torch.sigmoid(scores) diff --git a/src/pykeen/training/lcwa.py b/src/pykeen/training/lcwa.py index 14cdda9b66..2ae50efb3b 100644 --- a/src/pykeen/training/lcwa.py +++ b/src/pykeen/training/lcwa.py @@ -121,7 +121,6 @@ def _slice_size_search( sub_batch_size: int, supports_sub_batching: bool, ) -> int: # noqa: D102 - self._check_slicing_availability(supports_sub_batching) reached_max = False evaluated_once = False logger.info("Trying slicing now.") @@ -166,19 +165,3 @@ def _slice_size_search( evaluated_once = True return slice_size - - def _check_slicing_availability(self, supports_sub_batching: bool): - if self.model.can_slice_t: - return - elif supports_sub_batching: - report = ( - "This model supports sub-batching, but it also requires slicing," - " which is not implemented for this model yet." - ) - else: - report = ( - "This model doesn't support sub-batching and slicing is not" - " implemented for this model yet." - ) - logger.warning(report) - raise MemoryError("The current model can't be trained on this hardware with these parameters.") From 6ccdc9609f9ef289adbc2f084662ea903c8162c2 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 16:53:53 +0100 Subject: [PATCH 378/690] Cleanup tests --- src/pykeen/utils.py | 29 +---- tests/test_handling_of_cuda_exceptions.py | 30 ----- tests/test_hpo.py | 10 +- tests/test_nn.py | 67 ++++++++++ tests/test_utils.py | 148 ++++++++-------------- 5 files changed, 131 insertions(+), 153 deletions(-) delete mode 100644 tests/test_handling_of_cuda_exceptions.py create mode 100644 tests/test_nn.py diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 040a49b3a6..8fea998997 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -6,6 +6,7 @@ import json import logging import random +from abc import ABC, abstractmethod from io import BytesIO from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, SupportsFloat, Tuple, Type, TypeVar, Union @@ -27,13 +28,11 @@ 'compact_mapping', 'complex_normalize', 'fix_dataclass_init_docs', - 'imag_part', 'invert_mapping', 'is_cudnn_error', 'is_cuda_oom_error', 'l2_regularization', 'random_non_negative_int', - 'real_part', 'resolve_device', 'slice_triples', 'slice_doubles', @@ -327,17 +326,18 @@ def compact_mapping( return translated, translation -class Result: +class Result(ABC): """A superclass of results that can be saved to a directory.""" + @abstractmethod def save_to_directory(self, directory: str, **kwargs) -> None: """Save the results to the directory.""" - raise NotImplementedError + @abstractmethod def save_to_ftp(self, directory: str, ftp: ftplib.FTP) -> None: """Save the results to the directory in an FTP server.""" - raise NotImplementedError + @abstractmethod def save_to_s3(self, directory: str, bucket: str, s3=None) -> None: """Save all artifacts to the given directory in an S3 Bucket. @@ -345,7 +345,6 @@ def save_to_s3(self, directory: str, bucket: str, s3=None) -> None: :param bucket: The name of the S3 bucket :param s3: A client from :func:`boto3.client`, if already instantiated """ - raise NotImplementedError def split_complex( @@ -370,24 +369,6 @@ def combine_complex( return torch.cat([x_re, x_im], dim=-1) -# TODO remove (unused) -def real_part( - x: torch.FloatTensor, -) -> torch.FloatTensor: - """Get the real part from a complex tensor.""" - dim = x.shape[-1] // 2 - return x[..., :dim] - - -# TODO remove (unused) -def imag_part( - x: torch.FloatTensor, -) -> torch.FloatTensor: - """Get the imaginary part from a complex tensor.""" - dim = x.shape[-1] // 2 - return x[..., dim:] - - def fix_dataclass_init_docs(cls: Type) -> Type: """Fix the ``__init__`` documentation for a :class:`dataclasses.dataclass`. diff --git a/tests/test_handling_of_cuda_exceptions.py b/tests/test_handling_of_cuda_exceptions.py deleted file mode 100644 index b75021eef0..0000000000 --- a/tests/test_handling_of_cuda_exceptions.py +++ /dev/null @@ -1,30 +0,0 @@ -# -*- coding: utf-8 -*- - -"""Test that CUDA exceptions are processed properly.""" - -import unittest - -from pykeen.utils import _CUDA_OOM_ERROR, _CUDNN_ERROR, is_cuda_oom_error, is_cudnn_error - - -class TestCudaExceptionsHandling(unittest.TestCase): - """Test handling of CUDA exceptions.""" - - not_cuda_error = Exception("Something else.") - - def test_is_cuda_oom_error(self): - """Test handling of a CUDA out of memory exception.""" - error = RuntimeError(_CUDA_OOM_ERROR) - self.assertTrue(is_cuda_oom_error(runtime_error=error)) - self.assertFalse(is_cudnn_error(runtime_error=error)) - - self.assertFalse(is_cuda_oom_error(runtime_error=self.not_cuda_error)) - - def test_is_cudnn_error(self): - """Test handling of a cuDNN error.""" - error = RuntimeError(_CUDNN_ERROR) - self.assertTrue(is_cudnn_error(runtime_error=error)) - self.assertFalse(is_cuda_oom_error(runtime_error=error)) - - error = Exception("Something else.") - self.assertFalse(is_cudnn_error(runtime_error=self.not_cuda_error)) diff --git a/tests/test_hpo.py b/tests/test_hpo.py index f9c9027b8a..0e1819f9cb 100644 --- a/tests/test_hpo.py +++ b/tests/test_hpo.py @@ -34,7 +34,7 @@ def test_run(self): hpo_pipeline_result = hpo_pipeline( dataset='nations', model='TransE', - training_kwargs=dict(num_epochs=5), + training_kwargs=dict(num_epochs=5, use_tqdm=False), n_trials=2, ) df = hpo_pipeline_result.study.trials_dataframe(multi_index=True) @@ -51,7 +51,7 @@ def test_specified_model_hyperparameter(self): dataset='nations', model='TransE', model_kwargs=dict(embedding_dim=target_embedding_dim), - training_kwargs=dict(num_epochs=5), + training_kwargs=dict(num_epochs=5, use_tqdm=False), n_trials=2, ) df = hpo_pipeline_result.study.trials_dataframe(multi_index=True) @@ -66,7 +66,7 @@ def test_specified_loss_hyperparameter(self): dataset='nations', model='TransE', loss_kwargs=dict(margin=1.0), - training_kwargs=dict(num_epochs=5), + training_kwargs=dict(num_epochs=5, use_tqdm=False), n_trials=2, ) df = hpo_pipeline_result.study.trials_dataframe(multi_index=True) @@ -84,7 +84,7 @@ def test_specified_loss_and_model_hyperparameter(self): model_kwargs=dict(embedding_dim=target_embedding_dim), loss='MarginRankingLoss', loss_kwargs=dict(margin=1.0), - training_kwargs=dict(num_epochs=5), + training_kwargs=dict(num_epochs=5, use_tqdm=False), n_trials=2, ) df = hpo_pipeline_result.study.trials_dataframe(multi_index=True) @@ -104,7 +104,7 @@ def test_specified_range(self): loss_kwargs_ranges=dict( margin=dict(type=int, low=1, high=2), ), - training_kwargs=dict(num_epochs=5), + training_kwargs=dict(num_epochs=5, use_tqdm=False), n_trials=2, ) df = hpo_pipeline_result.study.trials_dataframe(multi_index=True) diff --git a/tests/test_nn.py b/tests/test_nn.py new file mode 100644 index 0000000000..5cba448604 --- /dev/null +++ b/tests/test_nn.py @@ -0,0 +1,67 @@ +# -*- coding: utf-8 -*- + +"""Unittest for the :mod:`pykeen.nn` module.""" + +import unittest + +import torch + +from pykeen.nn import Embedding + + +class EmbeddingsInCanonicalShapeTests(unittest.TestCase): + """Test get_embedding_in_canonical_shape().""" + + #: The number of embeddings + num_embeddings: int = 3 + + #: The embedding dimension + embedding_dim: int = 2 + + def setUp(self) -> None: + """Initialize embedding.""" + self.embedding = Embedding(num_embeddings=self.num_embeddings, embedding_dim=self.embedding_dim) + self.generator = torch.manual_seed(42) + self.embedding._embeddings.weight.data = torch.rand( + self.num_embeddings, + self.embedding_dim, + generator=self.generator, + ) + + def test_no_indices(self): + """Test getting all embeddings.""" + emb = self.embedding.get_in_canonical_shape(indices=None) + + # check shape + assert emb.shape == (1, self.num_embeddings, self.embedding_dim) + + # check values + exp = self.embedding(indices=None).view(1, self.num_embeddings, self.embedding_dim) + assert torch.allclose(emb, exp) + + def _test_with_indices(self, indices: torch.Tensor) -> None: + """Help tests with index.""" + emb = self.embedding.get_in_canonical_shape(indices=indices) + + # check shape + num_ind = indices.shape[0] + assert emb.shape == (num_ind, 1, self.embedding_dim) + + # check values + exp = torch.stack([self.embedding(i) for i in indices], dim=0).view(num_ind, 1, self.embedding_dim) + assert torch.allclose(emb, exp) + + def test_with_consecutive_indices(self): + """Test to retrieve all embeddings with consecutive indices.""" + indices = torch.arange(self.num_embeddings, dtype=torch.long) + self._test_with_indices(indices=indices) + + def test_with_indices_with_duplicates(self): + """Test to retrieve embeddings at random positions with duplicate indices.""" + indices = torch.randint( + self.num_embeddings, + size=(2 * self.num_embeddings,), + dtype=torch.long, + generator=self.generator, + ) + self._test_with_indices(indices=indices) diff --git a/tests/test_utils.py b/tests/test_utils.py index d2a4a9721f..0affa4fb36 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -"""Unittest for for global utilities.""" +"""Tests for the :mod:`pykeen.utils` module.""" import string import unittest @@ -8,13 +8,9 @@ import numpy import torch -from pykeen.nn import Embedding from pykeen.utils import ( - clamp_norm, - combine_complex, compact_mapping, - flatten_dictionary, - get_until_first_blank, - imag_part, l2_regularization, real_part, split_complex, + _CUDA_OOM_ERROR, _CUDNN_ERROR, clamp_norm, combine_complex, compact_mapping, flatten_dictionary, + get_until_first_blank, is_cuda_oom_error, is_cudnn_error, l2_regularization, split_complex, ) @@ -128,62 +124,8 @@ def test_regular(self): self.assertEqual("Broken line.", r) -class EmbeddingsInCanonicalShapeTests(unittest.TestCase): - """Test get_embedding_in_canonical_shape().""" - - #: The number of embeddings - num_embeddings: int = 3 - - #: The embedding dimension - embedding_dim: int = 2 - - def setUp(self) -> None: - """Initialize embedding.""" - self.embedding = Embedding(num_embeddings=self.num_embeddings, embedding_dim=self.embedding_dim) - self.generator = torch.manual_seed(42) - self.embedding._embeddings.weight.data = torch.rand( - self.num_embeddings, - self.embedding_dim, - generator=self.generator, - ) - - def test_no_indices(self): - """Test getting all embeddings.""" - emb = self.embedding.get_in_canonical_shape(indices=None) - - # check shape - assert emb.shape == (1, self.num_embeddings, self.embedding_dim) - - # check values - exp = self.embedding(indices=None).view(1, self.num_embeddings, self.embedding_dim) - assert torch.allclose(emb, exp) - - def _test_with_indices(self, indices: torch.Tensor) -> None: - """Help tests with index.""" - emb = self.embedding.get_in_canonical_shape(indices=indices) - - # check shape - num_ind = indices.shape[0] - assert emb.shape == (num_ind, 1, self.embedding_dim) - - # check values - exp = torch.stack([self.embedding(i) for i in indices], dim=0).view(num_ind, 1, self.embedding_dim) - assert torch.allclose(emb, exp) - - def test_with_consecutive_indices(self): - """Test to retrieve all embeddings with consecutive indices.""" - indices = torch.arange(self.num_embeddings, dtype=torch.long) - self._test_with_indices(indices=indices) - - def test_with_indices_with_duplicates(self): - """Test to retrieve embeddings at random positions with duplicate indices.""" - indices = torch.randint( - self.num_embeddings, - size=(2 * self.num_embeddings,), - dtype=torch.long, - generator=self.generator, - ) - self._test_with_indices(indices=indices) +class TestUtils(unittest.TestCase): + """Tests for :mod:`pykeen.utils`.""" def test_compact_mapping(self): """Test ``compact_mapping()``.""" @@ -198,34 +140,52 @@ def test_compact_mapping(self): self.assertEqual(set(id_remapping.keys()), set(mapping.values())) self.assertEqual(set(id_remapping.values()), set(compacted_mapping.values())) - -def test_clamp_norm(): - """Test clamp_norm() .""" - max_norm = 1.0 - gen = torch.manual_seed(42) - eps = 1.0e-06 - for p in [1, 2, float('inf')]: - for _ in range(10): - x = torch.rand(10, 20, 30, generator=gen) - for dim in range(x.ndimension()): - x_c = clamp_norm(x, maxnorm=max_norm, p=p, dim=dim) - - # check maximum norm constraint - assert (x_c.norm(p=p, dim=dim) <= max_norm + eps).all() - - # unchanged values for small norms - norm = x.norm(p=p, dim=dim) - mask = torch.stack([(norm < max_norm)] * x.shape[dim], dim=dim) - assert (x_c[mask] == x[mask]).all() - - -def test_complex_utils(): - """Test complex tensor utilities.""" - re = torch.rand(20, 10) - im = torch.rand(20, 10) - x = combine_complex(x_re=re, x_im=im) - assert (real_part(x) == re).all() - assert (imag_part(x) == im).all() - re2, im2 = split_complex(x) - assert (re2 == re).all() - assert (im2 == im).all() + def test_clamp_norm(self): + """Test :func:`pykeen.utils.clamp_norm`.""" + max_norm = 1.0 + gen = torch.manual_seed(42) + eps = 1.0e-06 + for p in [1, 2, float('inf')]: + for _ in range(10): + x = torch.rand(10, 20, 30, generator=gen) + for dim in range(x.ndimension()): + x_c = clamp_norm(x, maxnorm=max_norm, p=p, dim=dim) + + # check maximum norm constraint + assert (x_c.norm(p=p, dim=dim) <= max_norm + eps).all() + + # unchanged values for small norms + norm = x.norm(p=p, dim=dim) + mask = torch.stack([(norm < max_norm)] * x.shape[dim], dim=dim) + assert (x_c[mask] == x[mask]).all() + + def test_complex_utils(self): + """Test complex tensor utilities.""" + re = torch.rand(20, 10) + im = torch.rand(20, 10) + x = combine_complex(x_re=re, x_im=im) + re2, im2 = split_complex(x) + assert (re2 == re).all() + assert (im2 == im).all() + + +class TestCudaExceptionsHandling(unittest.TestCase): + """Test handling of CUDA exceptions.""" + + not_cuda_error = RuntimeError("Something else.") + + def test_is_cuda_oom_error(self): + """Test handling of a CUDA out of memory exception.""" + error = RuntimeError(_CUDA_OOM_ERROR) + self.assertTrue(is_cuda_oom_error(runtime_error=error)) + self.assertFalse(is_cudnn_error(runtime_error=error)) + + self.assertFalse(is_cuda_oom_error(runtime_error=self.not_cuda_error)) + + def test_is_cudnn_error(self): + """Test handling of a cuDNN error.""" + error = RuntimeError(_CUDNN_ERROR) + self.assertTrue(is_cudnn_error(runtime_error=error)) + self.assertFalse(is_cuda_oom_error(runtime_error=error)) + + self.assertFalse(is_cudnn_error(runtime_error=self.not_cuda_error)) From ca3e74d406104bd5f1704e7db2cb94bcc53555f1 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 17:05:23 +0100 Subject: [PATCH 379/690] Update TransH regularizer @mberr please check this looks right. Before it was with a dictionary for some reason --- src/pykeen/models/unimodal/trans_h.py | 11 ++++++----- src/pykeen/utils.py | 10 ++++++++++ 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 56d0664c04..3dc7355794 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -13,6 +13,7 @@ from ...regularizers import TransHRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint +from ...utils import pop_only __all__ = [ 'TransH', @@ -97,8 +98,8 @@ def __init__( constrainer=functional.normalize, ), ) - self.regularizer = self._instantiate_default_regularizer( - entity_embeddings=list(self.entity_representations[0].parameters()).pop(), - relation_embeddings=list(self.relation_representations[0].parameters()).pop(), - normal_vector_embeddings=list(self.relation_representations[1].parameters()).pop(), - ) + self.regularizer = self._instantiate_default_regularizer(parameters=[ + pop_only(self.entity_representations[0].parameters()), + pop_only(self.relation_representations[0].parameters()), + pop_only(self.relation_representations[1].parameters()), + ]) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 8fea998997..2fb12f09e6 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -649,3 +649,13 @@ def project_entity( e_bot = clamp_norm(e_bot, p=2, dim=-1, maxnorm=1) return e_bot + + +def pop_only(elements: Iterable[X]) -> X: + """Unpack a one element list, or raise an error.""" + elements = tuple(elements) + if len(elements) == 0: + raise ValueError('Empty sequence given') + if len(elements) > 1: + raise ValueError(f'More than one element: {elements}') + return elements[0] From 7d235ab0dfa4ced08c3352f6ce49e69a74af5478 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 17:22:13 +0100 Subject: [PATCH 380/690] Fix TransH... again --- src/pykeen/models/unimodal/trans_h.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 3dc7355794..488f1e1c65 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -98,8 +98,9 @@ def __init__( constrainer=functional.normalize, ), ) - self.regularizer = self._instantiate_default_regularizer(parameters=[ - pop_only(self.entity_representations[0].parameters()), - pop_only(self.relation_representations[0].parameters()), - pop_only(self.relation_representations[1].parameters()), - ]) + # Note that the TransH regularizer has a different interface + self.regularizer = self._instantiate_default_regularizer( + entity_embeddings=pop_only(self.entity_representations[0].parameters()), + relation_embeddings=pop_only(self.relation_representations[0].parameters()), + normal_vector_embeddings=pop_only(self.relation_representations[1].parameters()), + ) From 4809c8af8ce3f6ea794c457c5b8cf18a1dbc9f6c Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 17 Nov 2020 17:45:41 +0100 Subject: [PATCH 381/690] Hack in fix to RGCN test --- tests/test_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index 287c21b070..21420ebc87 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -530,7 +530,7 @@ def test_custom_representations(self): self.model.entity_representations = nn.ModuleList([ _CustomRepresentations( num_entities=self.factory.num_entities, - shape=er.shape, + shape=er.base_embeddings.shape if isinstance(er, RGCNRepresentations) else er.shape, ) for er in old_entity_reps ]) From 34dcf51014c9e0c207437ec331bac8eea44edd6c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 17 Nov 2020 17:58:05 +0100 Subject: [PATCH 382/690] Fix missing import --- tests/test_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index 21420ebc87..c8c63d2c10 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -36,7 +36,7 @@ ) from pykeen.models.cli import build_cli_from_cls from pykeen.models.unimodal.rgcn import ( - inverse_indegree_edge_weights, + RGCNRepresentations, inverse_indegree_edge_weights, inverse_outdegree_edge_weights, symmetric_edge_weights, ) From 7e3b9e3b908a17802a6540893d6d32b900d39633 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 10:17:41 +0100 Subject: [PATCH 383/690] change num_slices to be divisor-free to all other sizes --- tests/test_interactions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 5ae7444609..f33cc82400 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -298,7 +298,7 @@ class NTNTests(InteractionTests, unittest.TestCase): cls = pykeen.nn.modules.NTNInteraction - num_slices: int = 2 + num_slices: int = 11 shape_kwargs = dict( k=2, ) From f26d7810366be03406e20b7701477b6493994757 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 10:17:59 +0100 Subject: [PATCH 384/690] Also do this for the kwargs --- tests/test_interactions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index f33cc82400..8781cff5c0 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -300,7 +300,7 @@ class NTNTests(InteractionTests, unittest.TestCase): num_slices: int = 11 shape_kwargs = dict( - k=2, + k=11, ) From 5966d95c902f112ee6980542aa972c7050e2cb00 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 10:53:09 +0100 Subject: [PATCH 385/690] Add tests for functional form consistency --- tests/test_interactions.py | 82 ++++++++++++++++++++++++++++++-------- 1 file changed, 66 insertions(+), 16 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 8781cff5c0..785034e8e8 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -3,12 +3,13 @@ """Tests for interaction functions.""" import unittest -from typing import Any, Collection, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union +from typing import Any, Callable, Collection, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union from unittest.case import SkipTest import torch import pykeen.nn.modules +from pykeen.nn import functional as pkf from pykeen.nn.modules import Interaction, StatelessInteraction, TranslationalInteraction from pykeen.typing import Representation from pykeen.utils import get_subclasses @@ -63,6 +64,8 @@ class InteractionTests(GenericTests[pykeen.nn.modules.Interaction]): shape_kwargs = dict() + functional_form: Callable[..., torch.FloatTensor] + def _get_hrt( self, *shapes: Tuple[int, ...], @@ -177,35 +180,63 @@ def test_score_t_slicing(self): scores_no_slice = self.instance.score_t(h=h, r=r, all_entities=t, slice_size=None) assert torch.allclose(scores, scores_no_slice), f'Differences: {scores - scores_no_slice}' - def test_forward(self): - """Test forward.""" - for hs, rs, ts in [ - [ + def _get_test_shapes(self) -> Collection[Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]]]: + """Return a set of test shapes for (h, r, t).""" + return ( + ( # single score + (1, 1), + (1, 1), + (1, 1), + ), + ( # score_r with multi-t (self.batch_size, 1), (1, self.num_relations), - (self.batch_size, self.num_entities), - ], - [ + (self.batch_size, self.num_entities // 2 + 1), + ), + ( # score_r with multi-t and broadcasted head (1, 1), (1, self.num_relations), (self.batch_size, self.num_entities), - ], - [ + ), + ( # full cwa (1, self.num_entities), (1, self.num_relations), (1, self.num_entities), - ], - ]: + ), + ) + + def _get_output_shape( + self, + hs: Tuple[int, int], + rs: Tuple[int, int], + ts: Tuple[int, int], + ) -> Tuple[int, int, int, int]: + batch_size = max(hs[0], rs[0], ts[0]) + nh, nr, nt = hs[1], rs[1], ts[1] + if len(self.cls.relation_shape) == 0: + nr = 1 + if len(self.cls.entity_shape) == 0: + nh = nt = 1 + return batch_size, nh, nr, nt + + def _prepare_functional_input(self, h, r, t) -> Mapping[str, Any]: + return dict(h=h, r=r, t=t) + + def test_forward(self): + """Test forward.""" + for hs, rs, ts in self._get_test_shapes(): if any(isinstance(m, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d)) for m in self.instance.modules()): + # TODO: do we need to skip this for every combination? or only if batch_size = 1? continue with self.subTest(f"forward({hs}, {rs}, {ts})"): - expected_shape = [max(hs[0], rs[0], ts[0]), hs[1], rs[1], ts[1]] - if len(self.cls.relation_shape) == 0: - expected_shape[2] = 1 - expected_shape = tuple(expected_shape) h, r, t = self._get_hrt(hs, rs, ts) scores = self.instance(h=h, r=r, t=t) + expected_shape = self._get_output_shape(hs, rs, ts) self._check_scores(scores=scores, exp_shape=expected_shape) + with self.subTest(f"forward({hs}, {rs}, {ts}) - consistency with functional"): + kwargs = self._prepare_functional_input(h, r, t) + scores_f = self.__class__.functional_form(**kwargs) + assert torch.allclose(scores, scores_f) def test_scores(self): """Test individual scores.""" @@ -224,12 +255,14 @@ class ComplExTests(InteractionTests, unittest.TestCase): """Tests for ComplEx interaction function.""" cls = pykeen.nn.modules.ComplExInteraction + functional_form = pkf.complex_interaction class ConvETests(InteractionTests, unittest.TestCase): """Tests for ConvE interaction function.""" cls = pykeen.nn.modules.ConvEInteraction + functional_form = pkf.conve_interaction kwargs = dict( embedding_height=1, embedding_width=2, @@ -252,6 +285,7 @@ class ConvKBTests(InteractionTests, unittest.TestCase): """Tests for ConvKB interaction function.""" cls = pykeen.nn.modules.ConvKBInteraction + functional_form = pkf.convkb_interaction kwargs = dict( embedding_dim=InteractionTests.dim, num_filters=2 * InteractionTests.dim - 1, @@ -262,6 +296,7 @@ class DistMultTests(InteractionTests, unittest.TestCase): """Tests for DistMult interaction function.""" cls = pykeen.nn.modules.DistMultInteraction + functional_form = pkf.distmult_interaction def _exp_score(self, h, r, t) -> torch.FloatTensor: return (h * r * t).sum(dim=-1) @@ -271,6 +306,7 @@ class ERMLPTests(InteractionTests, unittest.TestCase): """Tests for ERMLP interaction function.""" cls = pykeen.nn.modules.ERMLPInteraction + functional_form = pkf.ermlp_interaction kwargs = dict( embedding_dim=InteractionTests.dim, hidden_dim=2 * InteractionTests.dim - 1, @@ -281,6 +317,7 @@ class ERMLPETests(InteractionTests, unittest.TestCase): """Tests for ERMLP-E interaction function.""" cls = pykeen.nn.modules.ERMLPEInteraction + functional_form = pkf.ermlpe_interaction kwargs = dict( embedding_dim=InteractionTests.dim, hidden_dim=2 * InteractionTests.dim - 1, @@ -291,12 +328,14 @@ class HolETests(InteractionTests, unittest.TestCase): """Tests for HolE interaction function.""" cls = pykeen.nn.modules.HolEInteraction + functional_form = pkf.hole_interaction class NTNTests(InteractionTests, unittest.TestCase): """Tests for NTN interaction function.""" cls = pykeen.nn.modules.NTNInteraction + functional_form = pkf.ntn_interaction num_slices: int = 11 shape_kwargs = dict( @@ -308,6 +347,7 @@ class ProjETests(InteractionTests, unittest.TestCase): """Tests for ProjE interaction function.""" cls = pykeen.nn.modules.ProjEInteraction + functional_form = pkf.proje_interaction kwargs = dict( embedding_dim=InteractionTests.dim, ) @@ -317,18 +357,21 @@ class RESCALTests(InteractionTests, unittest.TestCase): """Tests for RESCAL interaction function.""" cls = pykeen.nn.modules.RESCALInteraction + functional_form = pkf.rescal_interaction class KG2ETests(InteractionTests, unittest.TestCase): """Tests for KG2E interaction function.""" cls = pykeen.nn.modules.KG2EInteraction + functional_form = pkf.kg2e_interaction class TuckerTests(InteractionTests, unittest.TestCase): """Tests for Tucker interaction function.""" cls = pykeen.nn.modules.TuckerInteraction + functional_form = pkf.tucker_interaction kwargs = dict( embedding_dim=InteractionTests.dim, ) @@ -338,6 +381,7 @@ class RotatETests(InteractionTests, unittest.TestCase): """Tests for RotatE interaction function.""" cls = pykeen.nn.modules.RotatEInteraction + functional_form = pkf.rotate_interaction class TranslationalInteractionTests(InteractionTests): @@ -355,6 +399,7 @@ class TransDTests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransD interaction function.""" cls = pykeen.nn.modules.TransDInteraction + functional_form = pkf.transd_interaction shape_kwargs = dict( e=3, ) @@ -393,6 +438,7 @@ class TransETests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransE interaction function.""" cls = pykeen.nn.modules.TransEInteraction + functional_form = pkf.transe_interaction def _exp_score(self, h, r, t) -> torch.FloatTensor: return -(h + r - t).norm(p=2, dim=-1) @@ -402,12 +448,14 @@ class TransHTests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransH interaction function.""" cls = pykeen.nn.modules.TransHInteraction + functional_form = pkf.transh_interaction class TransRTests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransR interaction function.""" cls = pykeen.nn.modules.TransRInteraction + functional_form = pkf.transr_interaction shape_kwargs = dict( e=3, ) @@ -428,12 +476,14 @@ class SETests(TranslationalInteractionTests, unittest.TestCase): """Tests for SE interaction function.""" cls = pykeen.nn.modules.StructuredEmbeddingInteraction + functional_form = pkf.structured_embedding_interaction class UMTests(TranslationalInteractionTests, unittest.TestCase): """Tests for UM interaction function.""" cls = pykeen.nn.modules.UnstructuredModelInteraction + functional_form = pkf.unstructured_model_interaction class InteractionTestsTest(TestsTest[Interaction], unittest.TestCase): From 0365433fd47815e7493475e63021246ee29e3a4d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 10:54:27 +0100 Subject: [PATCH 386/690] Add type annotation --- tests/test_interactions.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 785034e8e8..18153d91d3 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -219,7 +219,12 @@ def _get_output_shape( nh = nt = 1 return batch_size, nh, nr, nt - def _prepare_functional_input(self, h, r, t) -> Mapping[str, Any]: + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: return dict(h=h, r=r, t=t) def test_forward(self): From 2d7cc5269db33514638e02913a483fe5088660ae Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 11:12:29 +0100 Subject: [PATCH 387/690] Add tests for functional for for all models --- tests/test_interactions.py | 133 ++++++++++++++++++++++++++++++++++++- 1 file changed, 132 insertions(+), 1 deletion(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 18153d91d3..bd2f36e1c0 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -280,7 +280,7 @@ def _get_hrt( self, *shapes: Tuple[int, ...], **kwargs, - ) -> Tuple[Union[Representation, Sequence[Representation]], ...]: + ) -> Tuple[Union[Representation, Sequence[Representation]], ...]: # noqa: D102 h, r, t = super()._get_hrt(*shapes, **kwargs) t_bias = torch.rand_like(t[..., 0, None]) return h, r, (t, t_bias) @@ -296,6 +296,22 @@ class ConvKBTests(InteractionTests, unittest.TestCase): num_filters=2 * InteractionTests.dim - 1, ) + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: # noqa: D102 + return dict( + h=h, + r=r, + t=t, + conv=self.instance.conv, + activation=self.instance.activation, + hidden_dropout=self.instance.hidden_dropout, + linear=self.instance.linear, + ) + class DistMultTests(InteractionTests, unittest.TestCase): """Tests for DistMult interaction function.""" @@ -317,6 +333,21 @@ class ERMLPTests(InteractionTests, unittest.TestCase): hidden_dim=2 * InteractionTests.dim - 1, ) + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: # noqa: D102 + return dict( + h=h, + r=r, + t=t, + hidden=self.instance.hidden, + activation=self.instance.activation, + final=self.instance.hidden_to_score, + ) + class ERMLPETests(InteractionTests, unittest.TestCase): """Tests for ERMLP-E interaction function.""" @@ -347,6 +378,14 @@ class NTNTests(InteractionTests, unittest.TestCase): k=11, ) + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: # noqa: D102 + return dict(h=h, t=t, w=r[0], b=r[1], u=r[2], vh=r[3], vt=r[4], activation=self.instance.non_linearity) + class ProjETests(InteractionTests, unittest.TestCase): """Tests for ProjE interaction function.""" @@ -357,6 +396,23 @@ class ProjETests(InteractionTests, unittest.TestCase): embedding_dim=InteractionTests.dim, ) + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: # noqa: D102 + return dict( + h=h, + r=r, + t=t, + d_e=self.instance.d_e, + d_r=self.instance.d_r, + b_c=self.instance.b_c, + b_p=self.instance.b_p, + activation=self.instance.inner_non_linearity, + ) + class RESCALTests(InteractionTests, unittest.TestCase): """Tests for RESCAL interaction function.""" @@ -371,6 +427,14 @@ class KG2ETests(InteractionTests, unittest.TestCase): cls = pykeen.nn.modules.KG2EInteraction functional_form = pkf.kg2e_interaction + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: # noqa: D102 + return dict(h_mean=h[0], h_var=h[1], r_mean=r[0], r_var=r[1], t_mean=t[0], t_var=t[1]) + class TuckerTests(InteractionTests, unittest.TestCase): """Tests for Tucker interaction function.""" @@ -399,6 +463,20 @@ class TranslationalInteractionTests(InteractionTests): def _additional_score_checks(self, scores): assert (scores <= 0).all() + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: # noqa: D102 + return dict( + h=h, + r=r, + t=t, + p=self.instance.p, + power_norm=self.instance.power_norm, + ) + class TransDTests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransD interaction function.""" @@ -438,6 +516,16 @@ def test_manual_big_relation_dim(self): scores = self.instance.score_hrt(h=(h, h_p), r=(r, r_p), t=(t, t_p)) self.assertAlmostEqual(scores.item(), -27, delta=0.01) + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: # noqa: D102 + kwargs = dict(super()._prepare_functional_input(h=h, r=r, t=t)) + kwargs.update(h=h[0], r=r[0], t=t[0], h_p=h[1], r_p=r[1], t_p=t[1]) + return kwargs + class TransETests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransE interaction function.""" @@ -455,6 +543,17 @@ class TransHTests(TranslationalInteractionTests, unittest.TestCase): cls = pykeen.nn.modules.TransHInteraction functional_form = pkf.transh_interaction + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: # noqa: D102 + kwargs = dict(super()._prepare_functional_input(h=h, r=r, t=t)) + w_r, d_r = kwargs.pop("r") + kwargs.update(w_r=w_r, d_r=d_r) + return kwargs + class TransRTests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransR interaction function.""" @@ -465,6 +564,17 @@ class TransRTests(TranslationalInteractionTests, unittest.TestCase): e=3, ) + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: # noqa: D102 + kwargs = dict(super()._prepare_functional_input(h=h, r=r, t=t)) + r, m_r = kwargs.pop("r") + kwargs.update(r=r, m_r=m_r) + return kwargs + def test_manual(self): """Manually test the value of the interaction function.""" # Compute Scores @@ -483,6 +593,17 @@ class SETests(TranslationalInteractionTests, unittest.TestCase): cls = pykeen.nn.modules.StructuredEmbeddingInteraction functional_form = pkf.structured_embedding_interaction + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: # noqa: D102 + kwargs = dict(super()._prepare_functional_input(h=h, r=r, t=t)) + r_h, r_t = kwargs.pop("r") + kwargs.update(r_h=r_h, r_t=r_t) + return kwargs + class UMTests(TranslationalInteractionTests, unittest.TestCase): """Tests for UM interaction function.""" @@ -490,6 +611,16 @@ class UMTests(TranslationalInteractionTests, unittest.TestCase): cls = pykeen.nn.modules.UnstructuredModelInteraction functional_form = pkf.unstructured_model_interaction + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: # noqa: D102 + kwargs = dict(super()._prepare_functional_input(h=h, r=r, t=t)) + kwargs.pop("r") + return kwargs + class InteractionTestsTest(TestsTest[Interaction], unittest.TestCase): """Test for tests for all interaction functions.""" From 78d201d066607477c9cbcb1c3adb50db5ec7e62c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 11:46:13 +0100 Subject: [PATCH 388/690] Simplify ConvE model --- src/pykeen/nn/functional.py | 63 +++++--------------------- src/pykeen/nn/modules.py | 88 +++++++++++++++++++------------------ tests/test_interactions.py | 53 ++++++++++++++++++++-- 3 files changed, 105 insertions(+), 99 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index b7ba59081e..10965b4fa7 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -4,6 +4,7 @@ from typing import Optional, Tuple, Union +import numpy import torch import torch.fft from torch import nn @@ -145,16 +146,8 @@ def conve_interaction( input_channels: int, embedding_height: int, embedding_width: int, - num_in_features: int, - bn0: Optional[nn.BatchNorm1d], - bn1: Optional[nn.BatchNorm1d], - bn2: Optional[nn.BatchNorm1d], - inp_drop: nn.Dropout, - feature_map_drop: nn.Dropout2d, - hidden_drop: nn.Dropout, - conv1: nn.Conv2d, - activation: nn.Module, - fc: nn.Linear, + hr2d: nn.Module, + hr1d: nn.Module, ) -> torch.FloatTensor: """ Evaluate the ConvE interaction function. @@ -173,26 +166,10 @@ def conve_interaction( The height of the reshaped embedding. :param embedding_width: The width of the reshaped embedding. - :param num_in_features: - The number of output features of the final layer (calculated with kernel and embedding dimensions). - :param bn0: - The first batch normalization layer. - :param bn1: - The second batch normalization layer. - :param bn2: - The third batch normalization layer. - :param inp_drop: - The input dropout layer. - :param feature_map_drop: - The feature map dropout layer. - :param hidden_drop: - The hidden dropout layer. - :param conv1: - The convolution layer. - :param activation: - The activation function. - :param fc: - The final fully connected layer. + :param hr2d: + The first module, transforming the 2D stacked head-relation "image". + :param hr1d: + The second module, transforming the 1D flattened output of the 2D module. :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. @@ -212,30 +189,12 @@ def conve_interaction( r = r.view(*r.shape[:-1], input_channels, embedding_height, embedding_width) x = broadcast_cat(h, r, dim=-2).view(-1, input_channels, 2 * embedding_height, embedding_width) - # batch_size, num_input_channels, 2*height, width - if bn0 is not None: - x = bn0(x) - - # batch_size, num_input_channels, 2*height, width - x = inp_drop(x) - - # (N,C_out,H_out,W_out) - x = conv1(x) - - if bn1 is not None: - x = bn1(x) - - x = activation(x) - x = feature_map_drop(x) + # batch_size', num_input_channels, 2*height, width + x = hr2d(x) # batch_size', num_output_channels * (2 * height - kernel_height + 1) * (width - kernel_width + 1) - x = x.view(-1, num_in_features) - x = fc(x) - x = hidden_drop(x) - - if bn2 is not None: - x = bn2(x) - x = activation(x) + x = x.view(-1, numpy.prod(x.shape[-3:])) + x = hr1d(x) # reshape: (batch_size', embedding_dim) -> (b, h, r, 1, d) x = x.view(-1, num_heads, num_relations, 1, embedding_dim) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 8aad7a7954..feeab3013a 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -538,6 +538,12 @@ class ConvEInteraction(Interaction[torch.FloatTensor, torch.FloatTensor, Tuple[t tail_entity_shape = ("d", "k") # with k=1 + #: The head-relation encoder operating on 2D "images" + hr2d: nn.Module + + #: The head-relation encoder operating on the 1D flattened version + hr1d: nn.Module + def __init__( self, input_channels: Optional[int] = None, @@ -568,46 +574,50 @@ def __init__( height=embedding_height, ) logger.info(f'Resolved to {input_channels} * {embedding_width} * {embedding_height} = {embedding_dim}.') - self.embedding_dim = embedding_dim - self.embedding_height = embedding_height - self.embedding_width = embedding_width - self.input_channels = input_channels - if self.input_channels * self.embedding_height * self.embedding_width != self.embedding_dim: + if input_channels * embedding_height * embedding_width != embedding_dim: raise ValueError( - f'Product of input channels ({self.input_channels}), height ({self.embedding_height}), and width ' - f'({self.embedding_width}) does not equal target embedding dimension ({self.embedding_dim})', + f'Product of input channels ({input_channels}), height ({embedding_height}), and width ' + f'({embedding_width}) does not equal target embedding dimension ({embedding_dim})', ) - self.inp_drop = nn.Dropout(input_dropout) - self.hidden_drop = nn.Dropout(output_dropout) - self.feature_map_drop = nn.Dropout2d(feature_map_dropout) - - self.conv1 = torch.nn.Conv2d( - in_channels=self.input_channels, - out_channels=output_channels, - kernel_size=(kernel_height, kernel_width), - stride=1, - padding=0, - bias=True, - ) + # encoders + # 1: 2D encoder: BN?, DO, Conv, BN?, Act, DO + hr2d_layers = [ + nn.BatchNorm2d(input_channels) if apply_batch_normalization else None, + nn.Dropout(input_dropout), + nn.Conv2d( + in_channels=input_channels, + out_channels=output_channels, + kernel_size=(kernel_height, kernel_width), + stride=1, + padding=0, + bias=True, + ), + nn.BatchNorm2d(output_channels) if apply_batch_normalization else None, + nn.ReLU(), + nn.Dropout2d(feature_map_dropout), + ] + self.hr2d = nn.Sequential(*(layer for layer in hr2d_layers if layer is not None)) - self.apply_batch_normalization = apply_batch_normalization - if self.apply_batch_normalization: - self.bn0 = nn.BatchNorm2d(self.input_channels) - self.bn1 = nn.BatchNorm2d(output_channels) - self.bn2 = nn.BatchNorm1d(self.embedding_dim) - else: - self.bn0 = None - self.bn1 = None - self.bn2 = None - self.num_in_features = ( + # 2: 1D encoder: FC, DO, BN?, Act + num_in_features = ( output_channels - * (2 * self.embedding_height - kernel_height + 1) - * (self.embedding_width - kernel_width + 1) + * (2 * embedding_height - kernel_height + 1) + * (embedding_width - kernel_width + 1) ) - self.fc = nn.Linear(self.num_in_features, self.embedding_dim) - self.activation = nn.ReLU() + hr1d_layers = [ + nn.Linear(num_in_features, embedding_dim), + nn.Dropout(output_dropout), + nn.BatchNorm1d(embedding_dim) if apply_batch_normalization else None, + nn.ReLU(), + ] + self.hr1d = nn.Sequential(*(layer for layer in hr1d_layers if layer is not None)) + + # store reshaping dimensions + self.embedding_height = embedding_height + self.embedding_width = embedding_width + self.input_channels = input_channels def forward( self, @@ -625,16 +635,8 @@ def forward( input_channels=self.input_channels, embedding_height=self.embedding_height, embedding_width=self.embedding_width, - num_in_features=self.num_in_features, - bn0=self.bn0, - bn1=self.bn1, - bn2=self.bn2, - inp_drop=self.inp_drop, - feature_map_drop=self.feature_map_drop, - hidden_drop=self.hidden_drop, - conv1=self.conv1, - activation=self.activation, - fc=self.fc, + hr2d=self.hr2d, + hr1d=self.hr1d, ) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index bd2f36e1c0..b2c8e94d52 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -3,6 +3,7 @@ """Tests for interaction functions.""" import unittest +from operator import itemgetter from typing import Any, Callable, Collection, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union from unittest.case import SkipTest @@ -248,11 +249,16 @@ def test_scores(self): self.instance.eval() for _ in range(10): h, r, t = self._get_hrt((1, 1), (1, 1), (1, 1)) - scores = self.instance(h=h, r=r, t=t) - exp_score = self._exp_score(h, r, t).item() - assert scores.item() == exp_score + kwargs = self._prepare_functional_input(h, r, t) - def _exp_score(self, h, r, t) -> torch.FloatTensor: + # calculate by functional + scores_f = self.__class__.functional_form(**kwargs).item() + + # calculate manually + scores_f_manual = self._exp_score(**kwargs).item() + assert scores_f_manual == scores_f + + def _exp_score(self, **kwargs) -> torch.FloatTensor: raise SkipTest("No score check implemented.") @@ -285,6 +291,45 @@ def _get_hrt( t_bias = torch.rand_like(t[..., 0, None]) return h, r, (t, t_bias) + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: # noqa: D102 + return dict( + h=h, + r=r, + t=t[0], + t_bias=t[1], + input_channels=self.instance.input_channels, + embedding_height=self.instance.embedding_height, + embedding_width=self.instance.embedding_width, + num_in_features=self.instance.num_in_features, + bn0=self.instance.bn0, + bn1=self.instance.bn1, + bn2=self.instance.bn2, + inp_drop=self.instance.inp_drop, + feature_map_drop=self.instance.feature_map_drop, + hidden_drop=self.instance.hidden_drop, + conv1=self.instance.conv1, + activation=self.instance.activation, + fc=self.instance.fc, + ) + + def _exp_score(self, **kwargs) -> torch.FloatTensor: + # unpack + activation, bn0, bn1, bn2, conv1, embedding_height, embedding_width, fc, feature_map_drop, h, hidden_drop, \ + inp_drop, input_channels, num_in_features, r, t, t_bias = \ + list(map(itemgetter(0), sorted(kwargs.items(), key=itemgetter(1)))) + bn0, bn1, bn2 = [lambda x:x if bn is None else bn for bn in (bn0, bn1, bn2)] + hr = torch.cat([ + h.view(1, input_channels, embedding_height, embedding_width), + r.view(1, input_channels, embedding_height, embedding_width) + ], dim=2) + x = feature_map_drop(activation(bn1(conv1(inp_drop(bn0(hr)))))) + + class ConvKBTests(InteractionTests, unittest.TestCase): """Tests for ConvKB interaction function.""" From cfb6419c21ae059b5315222c8c7a1e2e2c0984c1 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 11:52:20 +0100 Subject: [PATCH 389/690] Add ConvE manual score check --- tests/test_interactions.py | 27 ++++++++++----------------- 1 file changed, 10 insertions(+), 17 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index b2c8e94d52..91eb4aed2a 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -7,6 +7,7 @@ from typing import Any, Callable, Collection, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union from unittest.case import SkipTest +import numpy import torch import pykeen.nn.modules @@ -305,30 +306,22 @@ def _prepare_functional_input( input_channels=self.instance.input_channels, embedding_height=self.instance.embedding_height, embedding_width=self.instance.embedding_width, - num_in_features=self.instance.num_in_features, - bn0=self.instance.bn0, - bn1=self.instance.bn1, - bn2=self.instance.bn2, - inp_drop=self.instance.inp_drop, - feature_map_drop=self.instance.feature_map_drop, - hidden_drop=self.instance.hidden_drop, - conv1=self.instance.conv1, - activation=self.instance.activation, - fc=self.instance.fc, + hr2d=self.instance.hr2d, + hr1d=self.instance.hr1d, ) def _exp_score(self, **kwargs) -> torch.FloatTensor: # unpack - activation, bn0, bn1, bn2, conv1, embedding_height, embedding_width, fc, feature_map_drop, h, hidden_drop, \ - inp_drop, input_channels, num_in_features, r, t, t_bias = \ - list(map(itemgetter(0), sorted(kwargs.items(), key=itemgetter(1)))) - bn0, bn1, bn2 = [lambda x:x if bn is None else bn for bn in (bn0, bn1, bn2)] - hr = torch.cat([ + sorted_kwargs = list(map(itemgetter(1), sorted(kwargs.items(), key=itemgetter(0)))) + embedding_height, embedding_width, h, hr1d, hr2d, input_channels, r, t, t_bias = sorted_kwargs + x = torch.cat([ h.view(1, input_channels, embedding_height, embedding_width), r.view(1, input_channels, embedding_height, embedding_width) ], dim=2) - x = feature_map_drop(activation(bn1(conv1(inp_drop(bn0(hr)))))) - + x = hr2d(x) + x = x.view(-1, numpy.prod(x.shape[-3:])) + x = hr1d(x) + return (x.view(1, -1) * t.view(1, -1)).sum() + t_bias class ConvKBTests(InteractionTests, unittest.TestCase): From 622490f0cd06c84dd99859201036608212de9cc5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 11:52:39 +0100 Subject: [PATCH 390/690] Add docstring --- tests/test_interactions.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 91eb4aed2a..5a7c736590 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -260,6 +260,7 @@ def test_scores(self): assert scores_f_manual == scores_f def _exp_score(self, **kwargs) -> torch.FloatTensor: + """Compute the expected score for a single-score batch.""" raise SkipTest("No score check implemented.") From 7b057cb9911f25037064cc0114f96e37b36f1b84 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 11:53:32 +0100 Subject: [PATCH 391/690] Add _prepare_functional_input to ER-MLP E --- tests/test_interactions.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 5a7c736590..a7d724ba04 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -398,6 +398,14 @@ class ERMLPETests(InteractionTests, unittest.TestCase): hidden_dim=2 * InteractionTests.dim - 1, ) + def _prepare_functional_input( + self, + h: Union[Representation, Sequence[Representation]], + r: Union[Representation, Sequence[Representation]], + t: Union[Representation, Sequence[Representation]], + ) -> Mapping[str, Any]: # noqa: D102 + return dict(h=h, r=r, t=t, mlp=self.instance.mlp) + class HolETests(InteractionTests, unittest.TestCase): """Tests for HolE interaction function.""" From 63a106218bf3de8fa7a9909493e3b9b9977cffac Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 12:03:17 +0100 Subject: [PATCH 392/690] Add exp_score to NTN --- tests/test_interactions.py | 27 ++++++++++++++++++++++----- 1 file changed, 22 insertions(+), 5 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index a7d724ba04..f8040d168a 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -271,6 +271,10 @@ class ComplExTests(InteractionTests, unittest.TestCase): functional_form = pkf.complex_interaction +def _get_key_sorted_kwargs_values(kwargs: Mapping[str, Any]) -> Sequence[Any]: + return list(map(itemgetter(1), sorted(kwargs.items(), key=itemgetter(0)))) + + class ConvETests(InteractionTests, unittest.TestCase): """Tests for ConvE interaction function.""" @@ -312,12 +316,10 @@ def _prepare_functional_input( ) def _exp_score(self, **kwargs) -> torch.FloatTensor: - # unpack - sorted_kwargs = list(map(itemgetter(1), sorted(kwargs.items(), key=itemgetter(0)))) - embedding_height, embedding_width, h, hr1d, hr2d, input_channels, r, t, t_bias = sorted_kwargs + height, width, h, hr1d, hr2d, input_channels, r, t, t_bias = _get_key_sorted_kwargs_values(kwargs) x = torch.cat([ - h.view(1, input_channels, embedding_height, embedding_width), - r.view(1, input_channels, embedding_height, embedding_width) + h.view(1, input_channels, height, width), + r.view(1, input_channels, height, width) ], dim=2) x = hr2d(x) x = x.view(-1, numpy.prod(x.shape[-3:])) @@ -433,6 +435,21 @@ def _prepare_functional_input( ) -> Mapping[str, Any]: # noqa: D102 return dict(h=h, t=t, w=r[0], b=r[1], u=r[2], vh=r[3], vt=r[4], activation=self.instance.non_linearity) + def _exp_score(self, **kwargs) -> torch.FloatTensor: + # f(h,r,t) = u_r^T act(h W_r t + V_r h + V_r' t + b_r) + # shapes: + # w: (k, dim, dim) + # vh/vt: (k, dim) + # b/u: (k,) + activation, b, h, t, u, vh, vt, w = _get_key_sorted_kwargs_values(kwargs) + h, t = h.view(1, self.dim, 1), t.view(1, self.dim, 1) + w = w.view(self.num_slices, self.dim, self.dim) + vh, vt = [v.view(self.num_slices, 1, self.dim) for v in (vh, vt)] + b = b.view(self.num_slices, 1, 1) + u = u.view(self.num_slices, ) + x = activation(h.transpose(-2, -1) @ w @ t + vh @ h + vt @ t + b).view(self.num_slices) + return (x * u).sum() + class ProjETests(InteractionTests, unittest.TestCase): """Tests for ProjE interaction function.""" From 48e43044a22e3ccd199d7f71c326f19e70010386 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 12:05:14 +0100 Subject: [PATCH 393/690] Re-order arguments --- src/pykeen/models/unimodal/ntn.py | 18 +++++++++--------- src/pykeen/nn/functional.py | 8 ++++---- src/pykeen/nn/modules.py | 2 +- 3 files changed, 14 insertions(+), 14 deletions(-) diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index a0b278ad9e..273541ca26 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -73,14 +73,6 @@ def __init__( num_embeddings=triples_factory.num_relations, shape=(num_slices, embedding_dim, embedding_dim), ) - b = Embedding( - num_embeddings=triples_factory.num_relations, - embedding_dim=num_slices, - ) - u = Embedding( - num_embeddings=triples_factory.num_relations, - embedding_dim=num_slices, - ) vh = Embedding( num_embeddings=triples_factory.num_relations, shape=(num_slices, embedding_dim), @@ -89,6 +81,14 @@ def __init__( num_embeddings=triples_factory.num_relations, shape=(num_slices, embedding_dim), ) + b = Embedding( + num_embeddings=triples_factory.num_relations, + embedding_dim=num_slices, + ) + u = Embedding( + num_embeddings=triples_factory.num_relations, + embedding_dim=num_slices, + ) super().__init__( triples_factory=triples_factory, automatic_memory_optimization=automatic_memory_optimization, @@ -102,5 +102,5 @@ def __init__( num_embeddings=triples_factory.num_entities, embedding_dim=embedding_dim, ), - relation_representations=(w, b, u, vh, vt), + relation_representations=(w, vh, vt, b, u), ) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 10965b4fa7..11100b919a 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -446,10 +446,10 @@ def ntn_interaction( h: torch.FloatTensor, t: torch.FloatTensor, w: torch.FloatTensor, - b: torch.FloatTensor, - u: torch.FloatTensor, vh: torch.FloatTensor, vt: torch.FloatTensor, + b: torch.FloatTensor, + u: torch.FloatTensor, activation: nn.Module, ) -> torch.FloatTensor: r""" @@ -461,12 +461,12 @@ def ntn_interaction( :param h: shape: (batch_size, num_heads, dim) The head representations. + :param w: shape: (batch_size, num_relations, k, dim, dim) + The relation specific transformation matrix W_r. :param vh: shape: (batch_size, num_relations, k, dim) The head transformation matrix V_h. :param vt: shape: (batch_size, num_relations, k, dim) The tail transformation matrix V_h. - :param w: shape: (batch_size, num_relations, k, dim, dim) - The relation specific transformation matrix W_r. :param b: shape: (batch_size, num_relations, k) The relation specific offset b_r. :param u: shape: (batch_size, num_relations, k) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index feeab3013a..487f8215ae 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -1036,7 +1036,7 @@ def forward( r: Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], t: torch.FloatTensor, ) -> torch.FloatTensor: # noqa:D102 - w, b, u, vh, vt = r + w, vh, vt, b, u = r return pkf.ntn_interaction(h=h, t=t, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity) From cab86929dd9e8b4540643ca2e67d4f35a76ffa6e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 12:09:07 +0100 Subject: [PATCH 394/690] Fix NTN --- src/pykeen/nn/functional.py | 6 +++--- src/pykeen/nn/modules.py | 2 +- tests/test_interactions.py | 3 ++- 3 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 11100b919a..6fe98ef031 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -480,10 +480,10 @@ def ntn_interaction( The scores. """ # save sizes - num_heads, num_relations, num_tails, _, k = _extract_sizes(h, b, t) + num_heads, num_relations, num_tails, _, num_slices = _extract_sizes(h, b, t) x = extended_einsum("bhd,brkde,bte->bhrtk", h, w, t) - x = x + extended_einsum("brkd,bhd->bhk", vh, h).view(-1, num_heads, 1, 1, k) - x = x + extended_einsum("brkd,btd->btk", vt, t).view(-1, 1, 1, num_tails, k) + x = x + extended_einsum("brkd,bhd->bhrk", vh, h).view(-1, num_heads, 1, 1, num_slices) + x = x + extended_einsum("brkd,btd->brtk", vt, t).view(-1, 1, 1, num_tails, num_slices) x = activation(x) x = extended_einsum("bhrtk,brk->bhrt", x, u) return x diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 487f8215ae..6c0863eeb2 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -1019,7 +1019,7 @@ class NTNInteraction( ): """The interaction function of NTN.""" - relation_shape = ("kdd", "k", "k", "kd", "kd") + relation_shape = ("kdd", "kd", "kd", "k", "k") def __init__( self, diff --git a/tests/test_interactions.py b/tests/test_interactions.py index f8040d168a..2b0ec2f0a2 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -433,7 +433,8 @@ def _prepare_functional_input( r: Union[Representation, Sequence[Representation]], t: Union[Representation, Sequence[Representation]], ) -> Mapping[str, Any]: # noqa: D102 - return dict(h=h, t=t, w=r[0], b=r[1], u=r[2], vh=r[3], vt=r[4], activation=self.instance.non_linearity) + w, vh, vt, b, u = r + return dict(h=h, t=t, w=w, vh=vh, vt=vt, b=b, u=u, activation=self.instance.non_linearity) def _exp_score(self, **kwargs) -> torch.FloatTensor: # f(h,r,t) = u_r^T act(h W_r t + V_r h + V_r' t + b_r) From 39a09fd32575dd3a915abbe9d1a955a503cc8a4c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 12:11:41 +0100 Subject: [PATCH 395/690] Add static unpack method for NTN --- src/pykeen/nn/modules.py | 11 ++++++++--- tests/test_interactions.py | 3 +-- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 6c0863eeb2..8cde1a5e8f 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -6,7 +6,7 @@ import logging import math from abc import ABC, abstractmethod -from typing import Callable, Generic, List, Optional, Sequence, Tuple, Union +from typing import Callable, Generic, List, Mapping, Optional, Sequence, Tuple, Union import torch from torch import FloatTensor, nn @@ -1030,14 +1030,19 @@ def __init__( non_linearity = nn.Tanh() self.non_linearity = non_linearity + @staticmethod + def unpack(h, r, t) -> Mapping[str, torch.FloatTensor]: + """Unpack (h, r, t) to arguments for functional ntn_interaction.""" + w, vh, vt, b, u = r + return dict(h=h, t=t, w=w, b=b, u=u, vh=vh, vt=vt) + def forward( self, h: torch.FloatTensor, r: Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], t: torch.FloatTensor, ) -> torch.FloatTensor: # noqa:D102 - w, vh, vt, b, u = r - return pkf.ntn_interaction(h=h, t=t, w=w, b=b, u=u, vh=vh, vt=vt, activation=self.non_linearity) + return pkf.ntn_interaction(**self.unpack(h=h, r=r, t=t), activation=self.non_linearity) class KG2EInteraction( diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 2b0ec2f0a2..3b24fe5e69 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -433,8 +433,7 @@ def _prepare_functional_input( r: Union[Representation, Sequence[Representation]], t: Union[Representation, Sequence[Representation]], ) -> Mapping[str, Any]: # noqa: D102 - w, vh, vt, b, u = r - return dict(h=h, t=t, w=w, vh=vh, vt=vt, b=b, u=u, activation=self.instance.non_linearity) + return dict(**self.cls.unpack(h=h, r=r, t=t), activation=self.instance.non_linearity) def _exp_score(self, **kwargs) -> torch.FloatTensor: # f(h,r,t) = u_r^T act(h W_r t + V_r h + V_r' t + b_r) From 7fd7f4a8f091246a71609bdb5172437867298830 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 12:22:49 +0100 Subject: [PATCH 396/690] Add conversion utilities --- src/pykeen/nn/functional.py | 4 ++-- src/pykeen/nn/modules.py | 27 ++++++++++++++++++++++++++- tests/test_interactions.py | 15 +++++++-------- 3 files changed, 35 insertions(+), 11 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 6fe98ef031..6056d4330c 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -482,8 +482,8 @@ def ntn_interaction( # save sizes num_heads, num_relations, num_tails, _, num_slices = _extract_sizes(h, b, t) x = extended_einsum("bhd,brkde,bte->bhrtk", h, w, t) - x = x + extended_einsum("brkd,bhd->bhrk", vh, h).view(-1, num_heads, 1, 1, num_slices) - x = x + extended_einsum("brkd,btd->brtk", vt, t).view(-1, 1, 1, num_tails, num_slices) + x = x + extended_einsum("brkd,bhd->bhrk", vh, h).unsqueeze(dim=3) + x = x + extended_einsum("brkd,btd->brtk", vt, t).unsqueeze(dim=1) x = activation(x) x = extended_einsum("bhrtk,brk->bhrt", x, u) return x diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 8cde1a5e8f..323c47514e 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -6,7 +6,7 @@ import logging import math from abc import ABC, abstractmethod -from typing import Callable, Generic, List, Mapping, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Generic, List, Mapping, MutableMapping, Optional, Sequence, Tuple, Union import torch from torch import FloatTensor, nn @@ -93,6 +93,31 @@ class Interaction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, #: The symbolic shapes for relation representations relation_shape: Sequence[str] = ("d",) + def _prepare_hrt_for_functional( + self, + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + ) -> MutableMapping[str, torch.FloatTensor]: + """Conversion utility to prepare the h/r/t representations for the functional form.""" + assert all(torch.is_tensor(x) for x in (h, r, t)) + return dict(h=h, r=r, t=t) + + def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: + """Conversion utility to prepare the state to be passed to the functional form.""" + return dict() + + def _prepare_for_functional( + self, + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + ) -> Mapping[str, torch.FloatTensor]: + """Conversion utility to prepare the arguments for the functional form.""" + kwargs = self._prepare_hrt_for_functional(h=h, r=r, t=t) + kwargs.update(self._prepare_state_for_functional()) + return kwargs + @abstractmethod def forward( self, diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 3b24fe5e69..109865cded 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -227,6 +227,7 @@ def _prepare_functional_input( r: Union[Representation, Sequence[Representation]], t: Union[Representation, Sequence[Representation]], ) -> Mapping[str, Any]: + # TODO: Move this as class method to Interaction return dict(h=h, r=r, t=t) def test_forward(self): @@ -441,14 +442,12 @@ def _exp_score(self, **kwargs) -> torch.FloatTensor: # w: (k, dim, dim) # vh/vt: (k, dim) # b/u: (k,) - activation, b, h, t, u, vh, vt, w = _get_key_sorted_kwargs_values(kwargs) - h, t = h.view(1, self.dim, 1), t.view(1, self.dim, 1) - w = w.view(self.num_slices, self.dim, self.dim) - vh, vt = [v.view(self.num_slices, 1, self.dim) for v in (vh, vt)] - b = b.view(self.num_slices, 1, 1) - u = u.view(self.num_slices, ) - x = activation(h.transpose(-2, -1) @ w @ t + vh @ h + vt @ t + b).view(self.num_slices) - return (x * u).sum() + h, t = [kwargs[name].view(1, self.dim, 1) for name in "ht"] + w = kwargs["w"].view(self.num_slices, self.dim, self.dim) + vh, vt = [kwargs[name].view(self.num_slices, 1, self.dim) for name in ("vh", "vt")] + b = kwargs["b"].view(self.num_slices, 1, 1) + u = kwargs["u"].view(1, self.num_slices) + return u @ kwargs["activation"](h.transpose(-2, -1) @ w @ t + vh @ h + vt @ t + b).view(self.num_slices, 1) class ProjETests(InteractionTests, unittest.TestCase): From c81192639a0223116835e154e24e0e58dbd6b375 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 12:25:09 +0100 Subject: [PATCH 397/690] Add simplification of Interaction modules --- src/pykeen/nn/modules.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 323c47514e..401169c8b3 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -5,7 +5,7 @@ import itertools import logging import math -from abc import ABC, abstractmethod +from abc import ABC from typing import Any, Callable, Generic, List, Mapping, MutableMapping, Optional, Sequence, Tuple, Union import torch @@ -93,6 +93,9 @@ class Interaction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, #: The symbolic shapes for relation representations relation_shape: Sequence[str] = ("d",) + #: The functional interaction form + func: Callable[..., torch.FloatTensor] + def _prepare_hrt_for_functional( self, h: HeadRepresentation, @@ -118,7 +121,6 @@ def _prepare_for_functional( kwargs.update(self._prepare_state_for_functional()) return kwargs - @abstractmethod def forward( self, h: HeadRepresentation, @@ -137,7 +139,7 @@ def forward( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - raise NotImplementedError + return self.__class__.func(**self._prepare_for_functional(h=h, r=r, t=t)) @staticmethod def _add_dim(*x: torch.FloatTensor, dim: int) -> Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]: @@ -500,11 +502,10 @@ def forward( return pkf.transe_interaction(h=h, r=r, t=t, p=self.p, power_norm=self.power_norm) -class ComplExInteraction(StatelessInteraction[FloatTensor, FloatTensor, FloatTensor]): +class ComplExInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function of ComplEx.""" - def __init__(self): - super().__init__(f=pkf.complex_interaction) + func = pkf.complex_interaction def _calculate_missing_shape_information( From f97bd5f177d9ff36e8c759e7419d2107d5b10b72 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 12:26:44 +0100 Subject: [PATCH 398/690] Simplify ConvE --- src/pykeen/nn/modules.py | 24 ++++++++++++------------ 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 401169c8b3..d52360fb81 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -570,6 +570,9 @@ class ConvEInteraction(Interaction[torch.FloatTensor, torch.FloatTensor, Tuple[t #: The head-relation encoder operating on the 1D flattened version hr1d: nn.Module + #: The interaction function + func = pkf.conve_interaction + def __init__( self, input_channels: Optional[int] = None, @@ -645,19 +648,16 @@ def __init__( self.embedding_width = embedding_width self.input_channels = input_channels - def forward( + def _prepare_hrt_for_functional( self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: Tuple[torch.FloatTensor, torch.FloatTensor], - ) -> torch.FloatTensor: # noqa: D102 - # get tail bias term - t, t_bias = t - return pkf.conve_interaction( - h=h, - r=r, - t=t, - t_bias=t_bias, + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + ) -> MutableMapping[str, torch.FloatTensor]: + return dict(h=h, r=r, t=t[0], t_bias=t[1]) + + def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: + return dict( input_channels=self.input_channels, embedding_height=self.embedding_height, embedding_width=self.embedding_width, From 90cf3220872be29396a5fb7399599e8a4ba549fd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 12:31:07 +0100 Subject: [PATCH 399/690] Add noqas --- src/pykeen/nn/modules.py | 18 ++++++------------ 1 file changed, 6 insertions(+), 12 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index d52360fb81..def34a22bb 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -653,10 +653,10 @@ def _prepare_hrt_for_functional( h: HeadRepresentation, r: RelationRepresentation, t: TailRepresentation, - ) -> MutableMapping[str, torch.FloatTensor]: + ) -> MutableMapping[str, torch.FloatTensor]: # noqa: D102 return dict(h=h, r=r, t=t[0], t_bias=t[1]) - def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: + def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: # noqa: D102 return dict( input_channels=self.input_channels, embedding_height=self.embedding_height, @@ -672,6 +672,8 @@ class ConvKBInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): .. seealso:: :func:`pykeen.nn.functional.convkb_interaction`` """ + func = pkf.convkb_interaction + def __init__( self, hidden_dropout_rate: float = 0., @@ -699,16 +701,8 @@ def reset_parameters(self): # noqa: D102 nn.init.constant_(self.conv.weight[..., 2], -0.1) nn.init.zeros_(self.conv.bias) - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa: D102 - return pkf.convkb_interaction( - h=h, - r=r, - t=t, + def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: # noqa: D102 + return dict( conv=self.conv, activation=self.activation, hidden_dropout=self.hidden_dropout, From 4e6ed025a1719de160609396cb302e2548626272 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 12:32:52 +0100 Subject: [PATCH 400/690] Remove stateless interaction --- src/pykeen/nn/__init__.py | 3 +-- src/pykeen/nn/modules.py | 47 +++++++------------------------------- tests/test_interactions.py | 3 +-- 3 files changed, 10 insertions(+), 43 deletions(-) diff --git a/src/pykeen/nn/__init__.py b/src/pykeen/nn/__init__.py index eeef54bdb2..2a591716a9 100644 --- a/src/pykeen/nn/__init__.py +++ b/src/pykeen/nn/__init__.py @@ -4,7 +4,7 @@ from . import functional, init from .emb import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule -from .modules import Interaction, StatelessInteraction +from .modules import Interaction __all__ = [ 'Embedding', @@ -12,7 +12,6 @@ 'LiteralRepresentations', 'RepresentationModule', 'Interaction', - 'StatelessInteraction', 'init', 'functional', ] diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index def34a22bb..527361e2a6 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -18,7 +18,6 @@ __all__ = [ # Base Classes 'Interaction', - 'StatelessInteraction', 'TranslationalInteraction', # Concrete Classes 'ComplExInteraction', @@ -448,31 +447,6 @@ def reset_parameters(self): mod.reset_parameters() -class StatelessInteraction(Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation]): - """Interaction function without state.""" - - def __init__(self, f: Callable[..., torch.FloatTensor]): - """Instantiate the stateless interaction module. - - :param f: The interaction function, like ones from :mod:`pykeen.nn.functional`. - """ - super().__init__() - self.f = f - - def forward( - self, - h: HeadRepresentation, - r: RelationRepresentation, - t: TailRepresentation, - ) -> torch.FloatTensor: # noqa: D102 - # normalization - h = _upgrade_to_sequence(h) - r = _upgrade_to_sequence(r) - t = _upgrade_to_sequence(t) - # TODO provide example of non-simple case - return self.f(*h, *r, *t) - - class TranslationalInteraction(Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], ABC): """The translational interaction function shared by the TransE, TransR, TransH, and other Trans models.""" @@ -710,11 +684,10 @@ def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: # noqa: D1 ) -class DistMultInteraction(StatelessInteraction[FloatTensor, FloatTensor, FloatTensor]): +class DistMultInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): """A module wrapping the DistMult interaction function at :func:`pykeen.nn.functional.distmult_interaction`.""" - def __init__(self): - super().__init__(f=pkf.distmult_interaction) + func = pkf.distmult_interaction class ERMLPInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): @@ -828,18 +801,16 @@ def forward( return pkf.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p, power_norm=self.power_norm) -class RotatEInteraction(StatelessInteraction[FloatTensor, FloatTensor, FloatTensor]): +class RotatEInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function of RotatE.""" - def __init__(self): - super().__init__(f=pkf.rotate_interaction) + func = pkf.rotate_interaction -class HolEInteraction(StatelessInteraction[FloatTensor, FloatTensor, FloatTensor]): +class HolEInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function for HolE.""" - def __init__(self): - super().__init__(f=pkf.hole_interaction) + func = pkf.hole_interaction class ProjEInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): @@ -886,13 +857,11 @@ def forward( ) -class RESCALInteraction(StatelessInteraction[FloatTensor, FloatTensor, FloatTensor]): +class RESCALInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function of RESCAL.""" relation_shape = ("dd",) - - def __init__(self): - super().__init__(f=pkf.rescal_interaction) + func = pkf.rescal_interaction class StructuredEmbeddingInteraction( diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 109865cded..56ef99da22 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -12,7 +12,7 @@ import pykeen.nn.modules from pykeen.nn import functional as pkf -from pykeen.nn.modules import Interaction, StatelessInteraction, TranslationalInteraction +from pykeen.nn.modules import Interaction, TranslationalInteraction from pykeen.typing import Representation from pykeen.utils import get_subclasses @@ -692,5 +692,4 @@ class InteractionTestsTest(TestsTest[Interaction], unittest.TestCase): base_test = InteractionTests skip_cls = { TranslationalInteraction, - StatelessInteraction, } From 5fb19316435cfab77fb4b22a22ec58d3aa4a34e7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 12:48:10 +0100 Subject: [PATCH 401/690] Simplify tests by binding functional form to all interaction modules --- src/pykeen/nn/modules.py | 183 ++++++++++++++++------------------ tests/test_interactions.py | 198 +------------------------------------ 2 files changed, 88 insertions(+), 293 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 527361e2a6..db7e82a2d7 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -95,8 +95,8 @@ class Interaction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, #: The functional interaction form func: Callable[..., torch.FloatTensor] + @staticmethod def _prepare_hrt_for_functional( - self, h: HeadRepresentation, r: RelationRepresentation, t: TailRepresentation, @@ -463,17 +463,14 @@ def __init__(self, p: int, power_norm: bool = False): self.p = p self.power_norm = power_norm + def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: # noqa: D102 + return dict(p=self.p, power_norm=self.power_norm) + class TransEInteraction(TranslationalInteraction[FloatTensor, FloatTensor, FloatTensor]): """The TransE interaction function.""" - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa:D102 - return pkf.transe_interaction(h=h, r=r, t=t, p=self.p, power_norm=self.power_norm) + func = pkf.transe_interaction class ComplExInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): @@ -622,8 +619,8 @@ def __init__( self.embedding_width = embedding_width self.input_channels = input_channels + @staticmethod def _prepare_hrt_for_functional( - self, h: HeadRepresentation, r: RelationRepresentation, t: TailRepresentation, @@ -697,6 +694,8 @@ class ERMLPInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): f(h, r, t) = W_2 ReLU(W_1 cat(h, r, t) + b_1) + b_2 """ + func = pkf.ermlp_interaction + def __init__( self, embedding_dim: int, @@ -718,16 +717,8 @@ def __init__( self.activation = nn.ReLU() self.hidden_to_score = nn.Linear(in_features=hidden_dim, out_features=1, bias=True) - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa: D102 - return pkf.ermlp_interaction( - h=h, - r=r, - t=t, + def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: # noqa: D102 + return dict( hidden=self.hidden, activation=self.activation, final=self.hidden_to_score, @@ -748,6 +739,8 @@ def reset_parameters(self): # noqa: D102 class ERMLPEInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function of ER-MLP (E).""" + func = pkf.ermlpe_interaction + def __init__( self, hidden_dim: int = 300, @@ -768,13 +761,8 @@ def __init__( nn.ReLU(), ) - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa: D102 - return pkf.ermlpe_interaction(h=h, r=r, t=t, mlp=self.mlp) + def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: # noqa: D102 + return dict(mlp=self.mlp) class TransRInteraction( @@ -787,18 +775,18 @@ class TransRInteraction( """The TransR interaction function.""" relation_shape = ("e", "de") + func = pkf.transr_interaction def __init__(self, p: int, power_norm: bool = True): super().__init__(p=p, power_norm=power_norm) - def forward( - self, - h: torch.FloatTensor, - r: Tuple[torch.FloatTensor, torch.FloatTensor], - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa:D102 - r, m_r = r - return pkf.transr_interaction(h=h, r=r, t=t, m_r=m_r, p=self.p, power_norm=self.power_norm) + @staticmethod + def _prepare_hrt_for_functional( + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + ) -> MutableMapping[str, torch.FloatTensor]: # noqa: D102 + return dict(h=h, r=r, t=t, m_r=m_r) class RotatEInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): @@ -816,6 +804,8 @@ class HolEInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): class ProjEInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function for ProjE.""" + func = pkf.proje_interaction + def __init__( self, embedding_dim: int = 50, @@ -845,16 +835,8 @@ def reset_parameters(self): # noqa: D102 for p in self.parameters(): nn.init.uniform_(p, a=-bound, b=bound) - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa:D102 - return pkf.proje_interaction( - h=h, r=r, t=t, - d_e=self.d_e, d_r=self.d_r, b_c=self.b_c, b_p=self.b_p, activation=self.inner_non_linearity, - ) + def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: + return dict(d_e=self.d_e, d_r=self.d_r, b_c=self.b_c, b_p=self.b_p, activation=self.inner_non_linearity) class RESCALInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): @@ -874,20 +856,22 @@ class StructuredEmbeddingInteraction( """Interaction function of Structured Embedding.""" relation_shape = ("dd", "dd") + func = pkf.structured_embedding_interaction - def forward( - self, - h: torch.FloatTensor, - r: Tuple[torch.FloatTensor, torch.FloatTensor], - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa:D102 - rh, rt = r - return pkf.structured_embedding_interaction(h=h, r_h=rh, r_t=rt, t=t, p=self.p, power_norm=self.power_norm) + @staticmethod + def _prepare_hrt_for_functional( + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + ) -> MutableMapping[str, torch.FloatTensor]: # noqa: D102 + return dict(h=h, r=r, r_h=r[0], r_t=r[1]) class TuckerInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): """Interaction function of Tucker.""" + func = pkf.tucker_interaction + def __init__( self, embedding_dim: int = 200, @@ -933,16 +917,8 @@ def reset_parameters(self): # noqa:D102 # Initialize core tensor, cf. https://github.com/ibalazevic/TuckER/blob/master/model.py#L12 nn.init.uniform_(self.core_tensor, -1., 1.) - def forward( - self, - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa:D102 - return pkf.tucker_interaction( - h=h, - r=r, - t=t, + def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: + return dict( core_tensor=self.core_tensor, do0=self.input_dropout, do1=self.hidden_dropout_1, @@ -960,16 +936,18 @@ class UnstructuredModelInteraction( # shapes relation_shape: Sequence[str] = tuple() + func = pkf.unstructured_model_interaction + def __init__(self, p: int, power_norm: bool = True): super().__init__(p=p, power_norm=power_norm) - def forward( - self, - h: torch.FloatTensor, - r: None, - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa:D102 - return pkf.unstructured_model_interaction(h, t, p=self.p, power_norm=self.power_norm) + @staticmethod + def _prepare_hrt_for_functional( + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + ) -> MutableMapping[str, torch.FloatTensor]: # noqa: D102 + return dict(h=h, t=t) class TransDInteraction( @@ -983,20 +961,21 @@ class TransDInteraction( entity_shape = ("d", "d") relation_shape = ("e", "e") + func = pkf.transd_interaction def __init__(self, p: int = 2, power_norm: bool = True): super().__init__(p=p, power_norm=power_norm) - def forward( - self, - h: Tuple[torch.FloatTensor, torch.FloatTensor], - r: Tuple[torch.FloatTensor, torch.FloatTensor], - t: Tuple[torch.FloatTensor, torch.FloatTensor], - ) -> torch.FloatTensor: # noqa:D102 + @staticmethod + def _prepare_hrt_for_functional( + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + ) -> MutableMapping[str, torch.FloatTensor]: # noqa: D102 h, h_p = h r, r_p = r t, t_p = t - return pkf.transd_interaction(h=h, r=r, t=t, h_p=h_p, r_p=r_p, t_p=t_p, p=self.p, power_norm=self.power_norm) + return dict(h=h, r=r, t=t, h_p=h_p, r_p=r_p, t_p=t_p) class NTNInteraction( @@ -1009,6 +988,7 @@ class NTNInteraction( """The interaction function of NTN.""" relation_shape = ("kdd", "kd", "kd", "k", "k") + func = pkf.ntn_interaction def __init__( self, @@ -1020,18 +1000,16 @@ def __init__( self.non_linearity = non_linearity @staticmethod - def unpack(h, r, t) -> Mapping[str, torch.FloatTensor]: - """Unpack (h, r, t) to arguments for functional ntn_interaction.""" + def _prepare_hrt_for_functional( + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + ) -> MutableMapping[str, torch.FloatTensor]: # noqa: D102 w, vh, vt, b, u = r return dict(h=h, t=t, w=w, b=b, u=u, vh=vh, vt=vt) - def forward( - self, - h: torch.FloatTensor, - r: Tuple[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa:D102 - return pkf.ntn_interaction(**self.unpack(h=h, r=r, t=t), activation=self.non_linearity) + def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: # noqa: D102 + return dict(activation=self.non_linearity) class KG2EInteraction( @@ -1047,6 +1025,7 @@ class KG2EInteraction( relation_shape = ("d", "d") similarity: str exact: bool + func = pkf.kg2e_interaction def __init__( self, @@ -1059,22 +1038,26 @@ def __init__( self.similarity = similarity self.exact = exact - def forward( - self, - h: Tuple[torch.FloatTensor, torch.FloatTensor], - r: Tuple[torch.FloatTensor, torch.FloatTensor], - t: Tuple[torch.FloatTensor, torch.FloatTensor], - ) -> torch.FloatTensor: # noqa:D102 + @staticmethod + def _prepare_hrt_for_functional( + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + ) -> MutableMapping[str, torch.FloatTensor]: h_mean, h_var = h r_mean, r_var = r t_mean, t_var = t - return pkf.kg2e_interaction( + return dict( h_mean=h_mean, h_var=h_var, r_mean=r_mean, r_var=r_var, t_mean=t_mean, t_var=t_var, + ) + + def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: + return dict( similarity=self.similarity, exact=self.exact, ) @@ -1084,12 +1067,12 @@ class TransHInteraction(TranslationalInteraction[FloatTensor, Tuple[FloatTensor, """Interaction function of TransH.""" relation_shape = ("d", "d") + func = pkf.transh_interaction - def forward( - self, - h: torch.FloatTensor, - r: Tuple[torch.FloatTensor, torch.FloatTensor], - t: torch.FloatTensor, - ) -> torch.FloatTensor: # noqa: D102 - w_r, d_r = r - return pkf.transh_interaction(h, w_r, d_r, t, p=self.p, power_norm=self.power_norm) + @staticmethod + def _prepare_hrt_for_functional( + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + ) -> MutableMapping[str, torch.FloatTensor]: # noqa: D102 + return dict(h=h, w_r=r[0], d_r=r[1], t=t) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 56ef99da22..38f6b59c8d 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -4,14 +4,13 @@ import unittest from operator import itemgetter -from typing import Any, Callable, Collection, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union +from typing import Any, Collection, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union from unittest.case import SkipTest import numpy import torch import pykeen.nn.modules -from pykeen.nn import functional as pkf from pykeen.nn.modules import Interaction, TranslationalInteraction from pykeen.typing import Representation from pykeen.utils import get_subclasses @@ -66,8 +65,6 @@ class InteractionTests(GenericTests[pykeen.nn.modules.Interaction]): shape_kwargs = dict() - functional_form: Callable[..., torch.FloatTensor] - def _get_hrt( self, *shapes: Tuple[int, ...], @@ -221,15 +218,6 @@ def _get_output_shape( nh = nt = 1 return batch_size, nh, nr, nt - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: - # TODO: Move this as class method to Interaction - return dict(h=h, r=r, t=t) - def test_forward(self): """Test forward.""" for hs, rs, ts in self._get_test_shapes(): @@ -242,8 +230,8 @@ def test_forward(self): expected_shape = self._get_output_shape(hs, rs, ts) self._check_scores(scores=scores, exp_shape=expected_shape) with self.subTest(f"forward({hs}, {rs}, {ts}) - consistency with functional"): - kwargs = self._prepare_functional_input(h, r, t) - scores_f = self.__class__.functional_form(**kwargs) + kwargs = self.instance._prepare_for_functional(h=h, r=r, t=t) + scores_f = self.cls.func(**kwargs) assert torch.allclose(scores, scores_f) def test_scores(self): @@ -251,10 +239,10 @@ def test_scores(self): self.instance.eval() for _ in range(10): h, r, t = self._get_hrt((1, 1), (1, 1), (1, 1)) - kwargs = self._prepare_functional_input(h, r, t) + kwargs = self.instance._prepare_for_functional(h=h, r=r, t=t) # calculate by functional - scores_f = self.__class__.functional_form(**kwargs).item() + scores_f = self.cls.func(**kwargs).item() # calculate manually scores_f_manual = self._exp_score(**kwargs).item() @@ -269,7 +257,6 @@ class ComplExTests(InteractionTests, unittest.TestCase): """Tests for ComplEx interaction function.""" cls = pykeen.nn.modules.ComplExInteraction - functional_form = pkf.complex_interaction def _get_key_sorted_kwargs_values(kwargs: Mapping[str, Any]) -> Sequence[Any]: @@ -280,7 +267,6 @@ class ConvETests(InteractionTests, unittest.TestCase): """Tests for ConvE interaction function.""" cls = pykeen.nn.modules.ConvEInteraction - functional_form = pkf.conve_interaction kwargs = dict( embedding_height=1, embedding_width=2, @@ -298,24 +284,6 @@ def _get_hrt( t_bias = torch.rand_like(t[..., 0, None]) return h, r, (t, t_bias) - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: # noqa: D102 - return dict( - h=h, - r=r, - t=t[0], - t_bias=t[1], - input_channels=self.instance.input_channels, - embedding_height=self.instance.embedding_height, - embedding_width=self.instance.embedding_width, - hr2d=self.instance.hr2d, - hr1d=self.instance.hr1d, - ) - def _exp_score(self, **kwargs) -> torch.FloatTensor: height, width, h, hr1d, hr2d, input_channels, r, t, t_bias = _get_key_sorted_kwargs_values(kwargs) x = torch.cat([ @@ -332,34 +300,16 @@ class ConvKBTests(InteractionTests, unittest.TestCase): """Tests for ConvKB interaction function.""" cls = pykeen.nn.modules.ConvKBInteraction - functional_form = pkf.convkb_interaction kwargs = dict( embedding_dim=InteractionTests.dim, num_filters=2 * InteractionTests.dim - 1, ) - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: # noqa: D102 - return dict( - h=h, - r=r, - t=t, - conv=self.instance.conv, - activation=self.instance.activation, - hidden_dropout=self.instance.hidden_dropout, - linear=self.instance.linear, - ) - class DistMultTests(InteractionTests, unittest.TestCase): """Tests for DistMult interaction function.""" cls = pykeen.nn.modules.DistMultInteraction - functional_form = pkf.distmult_interaction def _exp_score(self, h, r, t) -> torch.FloatTensor: return (h * r * t).sum(dim=-1) @@ -369,73 +319,38 @@ class ERMLPTests(InteractionTests, unittest.TestCase): """Tests for ERMLP interaction function.""" cls = pykeen.nn.modules.ERMLPInteraction - functional_form = pkf.ermlp_interaction kwargs = dict( embedding_dim=InteractionTests.dim, hidden_dim=2 * InteractionTests.dim - 1, ) - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: # noqa: D102 - return dict( - h=h, - r=r, - t=t, - hidden=self.instance.hidden, - activation=self.instance.activation, - final=self.instance.hidden_to_score, - ) - class ERMLPETests(InteractionTests, unittest.TestCase): """Tests for ERMLP-E interaction function.""" cls = pykeen.nn.modules.ERMLPEInteraction - functional_form = pkf.ermlpe_interaction kwargs = dict( embedding_dim=InteractionTests.dim, hidden_dim=2 * InteractionTests.dim - 1, ) - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: # noqa: D102 - return dict(h=h, r=r, t=t, mlp=self.instance.mlp) - class HolETests(InteractionTests, unittest.TestCase): """Tests for HolE interaction function.""" cls = pykeen.nn.modules.HolEInteraction - functional_form = pkf.hole_interaction class NTNTests(InteractionTests, unittest.TestCase): """Tests for NTN interaction function.""" cls = pykeen.nn.modules.NTNInteraction - functional_form = pkf.ntn_interaction num_slices: int = 11 shape_kwargs = dict( k=11, ) - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: # noqa: D102 - return dict(**self.cls.unpack(h=h, r=r, t=t), activation=self.instance.non_linearity) - def _exp_score(self, **kwargs) -> torch.FloatTensor: # f(h,r,t) = u_r^T act(h W_r t + V_r h + V_r' t + b_r) # shapes: @@ -454,56 +369,27 @@ class ProjETests(InteractionTests, unittest.TestCase): """Tests for ProjE interaction function.""" cls = pykeen.nn.modules.ProjEInteraction - functional_form = pkf.proje_interaction kwargs = dict( embedding_dim=InteractionTests.dim, ) - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: # noqa: D102 - return dict( - h=h, - r=r, - t=t, - d_e=self.instance.d_e, - d_r=self.instance.d_r, - b_c=self.instance.b_c, - b_p=self.instance.b_p, - activation=self.instance.inner_non_linearity, - ) - class RESCALTests(InteractionTests, unittest.TestCase): """Tests for RESCAL interaction function.""" cls = pykeen.nn.modules.RESCALInteraction - functional_form = pkf.rescal_interaction class KG2ETests(InteractionTests, unittest.TestCase): """Tests for KG2E interaction function.""" cls = pykeen.nn.modules.KG2EInteraction - functional_form = pkf.kg2e_interaction - - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: # noqa: D102 - return dict(h_mean=h[0], h_var=h[1], r_mean=r[0], r_var=r[1], t_mean=t[0], t_var=t[1]) class TuckerTests(InteractionTests, unittest.TestCase): """Tests for Tucker interaction function.""" cls = pykeen.nn.modules.TuckerInteraction - functional_form = pkf.tucker_interaction kwargs = dict( embedding_dim=InteractionTests.dim, ) @@ -513,7 +399,6 @@ class RotatETests(InteractionTests, unittest.TestCase): """Tests for RotatE interaction function.""" cls = pykeen.nn.modules.RotatEInteraction - functional_form = pkf.rotate_interaction class TranslationalInteractionTests(InteractionTests): @@ -526,26 +411,11 @@ class TranslationalInteractionTests(InteractionTests): def _additional_score_checks(self, scores): assert (scores <= 0).all() - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: # noqa: D102 - return dict( - h=h, - r=r, - t=t, - p=self.instance.p, - power_norm=self.instance.power_norm, - ) - class TransDTests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransD interaction function.""" cls = pykeen.nn.modules.TransDInteraction - functional_form = pkf.transd_interaction shape_kwargs = dict( e=3, ) @@ -579,22 +449,11 @@ def test_manual_big_relation_dim(self): scores = self.instance.score_hrt(h=(h, h_p), r=(r, r_p), t=(t, t_p)) self.assertAlmostEqual(scores.item(), -27, delta=0.01) - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: # noqa: D102 - kwargs = dict(super()._prepare_functional_input(h=h, r=r, t=t)) - kwargs.update(h=h[0], r=r[0], t=t[0], h_p=h[1], r_p=r[1], t_p=t[1]) - return kwargs - class TransETests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransE interaction function.""" cls = pykeen.nn.modules.TransEInteraction - functional_form = pkf.transe_interaction def _exp_score(self, h, r, t) -> torch.FloatTensor: return -(h + r - t).norm(p=2, dim=-1) @@ -604,40 +463,16 @@ class TransHTests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransH interaction function.""" cls = pykeen.nn.modules.TransHInteraction - functional_form = pkf.transh_interaction - - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: # noqa: D102 - kwargs = dict(super()._prepare_functional_input(h=h, r=r, t=t)) - w_r, d_r = kwargs.pop("r") - kwargs.update(w_r=w_r, d_r=d_r) - return kwargs class TransRTests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransR interaction function.""" cls = pykeen.nn.modules.TransRInteraction - functional_form = pkf.transr_interaction shape_kwargs = dict( e=3, ) - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: # noqa: D102 - kwargs = dict(super()._prepare_functional_input(h=h, r=r, t=t)) - r, m_r = kwargs.pop("r") - kwargs.update(r=r, m_r=m_r) - return kwargs - def test_manual(self): """Manually test the value of the interaction function.""" # Compute Scores @@ -654,35 +489,12 @@ class SETests(TranslationalInteractionTests, unittest.TestCase): """Tests for SE interaction function.""" cls = pykeen.nn.modules.StructuredEmbeddingInteraction - functional_form = pkf.structured_embedding_interaction - - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: # noqa: D102 - kwargs = dict(super()._prepare_functional_input(h=h, r=r, t=t)) - r_h, r_t = kwargs.pop("r") - kwargs.update(r_h=r_h, r_t=r_t) - return kwargs class UMTests(TranslationalInteractionTests, unittest.TestCase): """Tests for UM interaction function.""" cls = pykeen.nn.modules.UnstructuredModelInteraction - functional_form = pkf.unstructured_model_interaction - - def _prepare_functional_input( - self, - h: Union[Representation, Sequence[Representation]], - r: Union[Representation, Sequence[Representation]], - t: Union[Representation, Sequence[Representation]], - ) -> Mapping[str, Any]: # noqa: D102 - kwargs = dict(super()._prepare_functional_input(h=h, r=r, t=t)) - kwargs.pop("r") - return kwargs class InteractionTestsTest(TestsTest[Interaction], unittest.TestCase): From f10a06ef3363fb71f567ab5db38531112e50376f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:02:49 +0100 Subject: [PATCH 402/690] Fix _prepare_hrt_for_functional for SE --- src/pykeen/nn/modules.py | 4 ++-- tests/test_interactions.py | 5 +++-- 2 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index db7e82a2d7..911ec68f73 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -786,7 +786,7 @@ def _prepare_hrt_for_functional( r: RelationRepresentation, t: TailRepresentation, ) -> MutableMapping[str, torch.FloatTensor]: # noqa: D102 - return dict(h=h, r=r, t=t, m_r=m_r) + return dict(h=h, r=r[0], t=t, m_r=r[1]) class RotatEInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): @@ -864,7 +864,7 @@ def _prepare_hrt_for_functional( r: RelationRepresentation, t: TailRepresentation, ) -> MutableMapping[str, torch.FloatTensor]: # noqa: D102 - return dict(h=h, r=r, r_h=r[0], r_t=r[1]) + return dict(h=h, t=t, r_h=r[0], r_t=r[1]) class TuckerInteraction(Interaction[FloatTensor, FloatTensor, FloatTensor]): diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 38f6b59c8d..fe457830c7 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -455,8 +455,9 @@ class TransETests(TranslationalInteractionTests, unittest.TestCase): cls = pykeen.nn.modules.TransEInteraction - def _exp_score(self, h, r, t) -> torch.FloatTensor: - return -(h + r - t).norm(p=2, dim=-1) + def _exp_score(self, h, r, t, p, power_norm) -> torch.FloatTensor: + assert not power_norm + return -(h + r - t).norm(p=p, dim=-1) class TransHTests(TranslationalInteractionTests, unittest.TestCase): From 9aa003aaa561bd0de790f07819122f15649ae198 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:06:17 +0100 Subject: [PATCH 403/690] Add _exp_score to complex --- tests/test_interactions.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index fe457830c7..08b26f93af 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -3,7 +3,6 @@ """Tests for interaction functions.""" import unittest -from operator import itemgetter from typing import Any, Collection, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union from unittest.case import SkipTest @@ -13,7 +12,7 @@ import pykeen.nn.modules from pykeen.nn.modules import Interaction, TranslationalInteraction from pykeen.typing import Representation -from pykeen.utils import get_subclasses +from pykeen.utils import get_subclasses, split_complex T = TypeVar("T") @@ -258,9 +257,9 @@ class ComplExTests(InteractionTests, unittest.TestCase): cls = pykeen.nn.modules.ComplExInteraction - -def _get_key_sorted_kwargs_values(kwargs: Mapping[str, Any]) -> Sequence[Any]: - return list(map(itemgetter(1), sorted(kwargs.items(), key=itemgetter(0)))) + def _exp_score(self, h, r, t) -> torch.FloatTensor: # noqa: D102 + h, r, t = [torch.complex(*split_complex(x)) for x in (h, r, t)] + return (h * r * t).sum().real class ConvETests(InteractionTests, unittest.TestCase): @@ -284,8 +283,7 @@ def _get_hrt( t_bias = torch.rand_like(t[..., 0, None]) return h, r, (t, t_bias) - def _exp_score(self, **kwargs) -> torch.FloatTensor: - height, width, h, hr1d, hr2d, input_channels, r, t, t_bias = _get_key_sorted_kwargs_values(kwargs) + def _exp_score(self, height, width, h, hr1d, hr2d, input_channels, r, t, t_bias) -> torch.FloatTensor: x = torch.cat([ h.view(1, input_channels, height, width), r.view(1, input_channels, height, width) From b9219739721301e6e6f69d703553b066d610140d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:10:09 +0100 Subject: [PATCH 404/690] Add _exp_score to ConvKB --- tests/test_interactions.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 08b26f93af..e1584368b0 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -303,6 +303,11 @@ class ConvKBTests(InteractionTests, unittest.TestCase): num_filters=2 * InteractionTests.dim - 1, ) + def _exp_score(self, h, r, t, conv, activation, hidden_dropout, linear) -> torch.FloatTensor: # noqa: D102 + # W_L drop(act(W_C \ast ([h; r; t]) + b_C)) + b_L + x = torch.stack([x.view(-1) for x in (h, r, t)], dim=1).view(1, 1, -1, 3) + return linear(hidden_dropout(activation(conv(x).view(1, -1)))) + class DistMultTests(InteractionTests, unittest.TestCase): """Tests for DistMult interaction function.""" From 8b72a499472eb0af20fdecd6418430d65a17f62a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:12:10 +0100 Subject: [PATCH 405/690] Add _exp_score to ERMLP --- tests/test_interactions.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index e1584368b0..f261d93771 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -327,6 +327,10 @@ class ERMLPTests(InteractionTests, unittest.TestCase): hidden_dim=2 * InteractionTests.dim - 1, ) + def _exp_score(self, h, r, t, hidden, activation, final) -> torch.FloatTensor: + x = torch.cat([x.view(-1) for x in (h, r, t)]) + return final(activation(hidden(x))) + class ERMLPETests(InteractionTests, unittest.TestCase): """Tests for ERMLP-E interaction function.""" From f538ee273fd9e63345ded20217d4db9ddc8c7620 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:12:26 +0100 Subject: [PATCH 406/690] Only check for approximate equality (due to fp operations) --- tests/test_interactions.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index f261d93771..1c8a920119 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -241,11 +241,11 @@ def test_scores(self): kwargs = self.instance._prepare_for_functional(h=h, r=r, t=t) # calculate by functional - scores_f = self.cls.func(**kwargs).item() + scores_f = self.cls.func(**kwargs) # calculate manually - scores_f_manual = self._exp_score(**kwargs).item() - assert scores_f_manual == scores_f + scores_f_manual = self._exp_score(**kwargs) + assert torch.allclose(scores_f_manual, scores_f) def _exp_score(self, **kwargs) -> torch.FloatTensor: """Compute the expected score for a single-score batch.""" From c8b8efb45c8d2ea66e26b220031e333ad11085d5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:13:59 +0100 Subject: [PATCH 407/690] Fix ConvE _exp_score --- tests/test_interactions.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 1c8a920119..a759deb90a 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -283,10 +283,12 @@ def _get_hrt( t_bias = torch.rand_like(t[..., 0, None]) return h, r, (t, t_bias) - def _exp_score(self, height, width, h, hr1d, hr2d, input_channels, r, t, t_bias) -> torch.FloatTensor: + def _exp_score( + self, embedding_height, embedding_width, h, hr1d, hr2d, input_channels, r, t, t_bias + ) -> torch.FloatTensor: x = torch.cat([ - h.view(1, input_channels, height, width), - r.view(1, input_channels, height, width) + h.view(1, input_channels, embedding_height, embedding_width), + r.view(1, input_channels, embedding_height, embedding_width) ], dim=2) x = hr2d(x) x = x.view(-1, numpy.prod(x.shape[-3:])) From 01a5284129e4a09c6d8b48f864fb92eaf591483b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:16:44 +0100 Subject: [PATCH 408/690] Add _exp_score to ERMLPE --- tests/test_interactions.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index a759deb90a..0a74450951 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -343,6 +343,10 @@ class ERMLPETests(InteractionTests, unittest.TestCase): hidden_dim=2 * InteractionTests.dim - 1, ) + def _exp_score(self, h, r, t, mlp) -> torch.FloatTensor: # noqa: D102 + x = torch.cat([x.view(1, -1) for x in (h, r)], dim=-1) + return mlp(x).view(1, -1) @ t.view(-1, 1) + class HolETests(InteractionTests, unittest.TestCase): """Tests for HolE interaction function.""" From d792d81bc9eddae77638ae5a5f9858c0454c5d0f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:19:31 +0100 Subject: [PATCH 409/690] Add _exp_score to HolE --- tests/test_interactions.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 0a74450951..1a8e5d89f1 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -353,6 +353,12 @@ class HolETests(InteractionTests, unittest.TestCase): cls = pykeen.nn.modules.HolEInteraction + def _exp_score(self, h, r, t) -> torch.FloatTensor: # noqa: D102 + h, t = [torch.fft.rfft(x.view(1, -1), dim=-1) for x in (h, t)] + h = torch.conj(h) + c = torch.fft.irfft(h * t, n=h.shape[-1], dim=-1) + return (c * r).sum() + class NTNTests(InteractionTests, unittest.TestCase): """Tests for NTN interaction function.""" From e48ed12b41d55e2fe8eb92970669a723dde71550 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:24:08 +0100 Subject: [PATCH 410/690] Simplify NTN _exp_score --- tests/test_interactions.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 1a8e5d89f1..eca53eab8b 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -370,18 +370,20 @@ class NTNTests(InteractionTests, unittest.TestCase): k=11, ) - def _exp_score(self, **kwargs) -> torch.FloatTensor: + def _exp_score(self, h, t, w, vt, vh, b, u, activation) -> torch.FloatTensor: # f(h,r,t) = u_r^T act(h W_r t + V_r h + V_r' t + b_r) - # shapes: - # w: (k, dim, dim) - # vh/vt: (k, dim) - # b/u: (k,) - h, t = [kwargs[name].view(1, self.dim, 1) for name in "ht"] - w = kwargs["w"].view(self.num_slices, self.dim, self.dim) - vh, vt = [kwargs[name].view(self.num_slices, 1, self.dim) for name in ("vh", "vt")] - b = kwargs["b"].view(self.num_slices, 1, 1) - u = kwargs["u"].view(1, self.num_slices) - return u @ kwargs["activation"](h.transpose(-2, -1) @ w @ t + vh @ h + vt @ t + b).view(self.num_slices, 1) + # shapes: w: (k, dim, dim), vh/vt: (k, dim), b/u: (k,), h/t: (dim,) + # remove batch/num dimension + h, t, w, vt, vh, b, u = [x.view(*x.shape[2:]) for x in (h, t, w, vt, vh, b, u)] + score = 0. + for i in range(u.shape[-1]): + score = score + u[i] * activation( + h.view(1, self.dim) @ w[i] @ t.view(self.dim, 1) + + (vh[i] * h.view(-1)).sum() + + (vt[i] * t.view(-1)).sum() + + b[i] + ) + return score class ProjETests(InteractionTests, unittest.TestCase): From 08028b1ff8cfe00eb3375cb28e463227eec3b5b4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:29:08 +0100 Subject: [PATCH 411/690] Add _exp_score to ProjE --- tests/test_interactions.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index eca53eab8b..bddd78e9e2 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -360,6 +360,10 @@ def _exp_score(self, h, r, t) -> torch.FloatTensor: # noqa: D102 return (c * r).sum() +def _strip_dim(*x): + return [xx.view(xx.shape[2:]) for xx in x] + + class NTNTests(InteractionTests, unittest.TestCase): """Tests for NTN interaction function.""" @@ -374,7 +378,7 @@ def _exp_score(self, h, t, w, vt, vh, b, u, activation) -> torch.FloatTensor: # f(h,r,t) = u_r^T act(h W_r t + V_r h + V_r' t + b_r) # shapes: w: (k, dim, dim), vh/vt: (k, dim), b/u: (k,), h/t: (dim,) # remove batch/num dimension - h, t, w, vt, vh, b, u = [x.view(*x.shape[2:]) for x in (h, t, w, vt, vh, b, u)] + h, t, w, vt, vh, b, u = _strip_dim(h, t, w, vt, vh, b, u) score = 0. for i in range(u.shape[-1]): score = score + u[i] * activation( @@ -394,6 +398,11 @@ class ProjETests(InteractionTests, unittest.TestCase): embedding_dim=InteractionTests.dim, ) + def _exp_score(self, h, r, t, d_e, d_r, b_c, b_p, activation) -> torch.FloatTensor: + # f(h, r, t) = g(t z(D_e h + D_r r + b_c) + b_p) + h, r, t = _strip_dim(h, r, t) + return (t * activation((d_e * h) + (d_r * r) + b_c)).sum() + b_p + class RESCALTests(InteractionTests, unittest.TestCase): """Tests for RESCAL interaction function.""" From 39b7fa4faceb16330263e269d019702fe2b23c19 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:30:01 +0100 Subject: [PATCH 412/690] Add _exp_score to RESCAL --- tests/test_interactions.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index bddd78e9e2..848e3f70ab 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -409,6 +409,11 @@ class RESCALTests(InteractionTests, unittest.TestCase): cls = pykeen.nn.modules.RESCALInteraction + def _exp_score(self, h, r, t) -> torch.FloatTensor: + # f(h, r, t) = h @ r @ t + h, r, t = _strip_dim(h, r, t) + return h.view(1, -1) @ r @ t.view(-1, 1) + class KG2ETests(InteractionTests, unittest.TestCase): """Tests for KG2E interaction function.""" From bbb5391ef6fe971da39b10294e4922e8f66ad974 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:32:59 +0100 Subject: [PATCH 413/690] Re-use view_complex --- tests/test_interactions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 848e3f70ab..c407607f3a 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -12,7 +12,7 @@ import pykeen.nn.modules from pykeen.nn.modules import Interaction, TranslationalInteraction from pykeen.typing import Representation -from pykeen.utils import get_subclasses, split_complex +from pykeen.utils import get_subclasses, view_complex T = TypeVar("T") @@ -258,7 +258,7 @@ class ComplExTests(InteractionTests, unittest.TestCase): cls = pykeen.nn.modules.ComplExInteraction def _exp_score(self, h, r, t) -> torch.FloatTensor: # noqa: D102 - h, r, t = [torch.complex(*split_complex(x)) for x in (h, r, t)] + h, r, t = [view_complex(x) for x in (h, r, t)] return (h * r * t).sum().real From 17c8b01489a8a820090ca57a69ecd65e41c183b0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:36:25 +0100 Subject: [PATCH 414/690] Fix rotate_interaction --- src/pykeen/nn/functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 6056d4330c..1953819543 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -589,7 +589,7 @@ def rotate_interaction( # Workaround until https://github.com/pytorch/pytorch/issues/30704 is fixed return negative_norm_of_sum( hr.unsqueeze(dim=3), - t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]), + -t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]), p=2, power_norm=False, ) From 2729925b81a6767eeb41abb37b35786a1399189e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:36:34 +0100 Subject: [PATCH 415/690] Add _exp_score for RotatE --- tests/test_interactions.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index c407607f3a..2bfd7f8b10 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -435,6 +435,12 @@ class RotatETests(InteractionTests, unittest.TestCase): cls = pykeen.nn.modules.RotatEInteraction + def _exp_score(self, h, r, t) -> torch.FloatTensor: # noqa: D102 + h, r, t = _strip_dim(*(view_complex(x) for x in (h, r, t))) + hr = h * r + d = hr - t + return -(d.abs() ** 2).sum().sqrt() + class TranslationalInteractionTests(InteractionTests): """Common tests for translational interaction.""" From 3e6857647e5550e97d83bcdbc66189753ece93b5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:37:24 +0100 Subject: [PATCH 416/690] Add comment to RotatE optimization --- src/pykeen/nn/functional.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 1953819543..59f8b0e39d 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -575,6 +575,7 @@ def rotate_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ + # TODO: rotate by inverse relation, if the cost(dist(h, t * r_inv)) < cost(dist(h * r, t)) # # r expresses a rotation in complex plane. # # The inverse rotation is expressed by the complex conjugate of r. # # The score is computed as the distance of the relation-rotated head to the tail. From 453aa1fcafcd5521ef88e8999ed9de874b965ef3 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:45:16 +0100 Subject: [PATCH 417/690] Add equivalence check to test_project_entity --- tests/test_models.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/tests/test_models.py b/tests/test_models.py index c8c63d2c10..9adf78bbda 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -902,6 +902,13 @@ def test_project_entity(self): # check normalization assert (torch.norm(e_bot, dim=-1, p=2) <= 1.0 + 1.0e-06).all() + # check equivalence of re-formulation + # e_{\bot} = M_{re} e = (r_p e_p^T + I^{d_r \times d_e}) e + # = r_p (e_p^T e) + e' + M_re = (r_p.unsqueeze(dim=-1) @ e_p.unsqueeze(dim=-2) + torch.eye(relation_dim, self.embedding_dim).unsqueeze(dim=0)) + e_vanilla = (M_re @ e.unsqueeze(dim=-1)).squeeze(dim=-1) + assert torch.allclose(e_vanilla, e_bot) + class TestTransE(_DistanceModelTestCase, unittest.TestCase): """Test the TransE model.""" From b17d266c91a7caeadab3567ff0f9236c82dc1b09 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:47:05 +0100 Subject: [PATCH 418/690] Add _exp_score to TransD --- tests/test_interactions.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 2bfd7f8b10..35c4b92773 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -12,7 +12,7 @@ import pykeen.nn.modules from pykeen.nn.modules import Interaction, TranslationalInteraction from pykeen.typing import Representation -from pykeen.utils import get_subclasses, view_complex +from pykeen.utils import get_subclasses, project_entity, view_complex T = TypeVar("T") @@ -490,6 +490,12 @@ def test_manual_big_relation_dim(self): scores = self.instance.score_hrt(h=(h, h_p), r=(r, r_p), t=(t, t_p)) self.assertAlmostEqual(scores.item(), -27, delta=0.01) + def _exp_score(self, h, r, t, h_p, r_p, t_p, p, power_norm) -> torch.FloatTensor: # noqa: D102 + assert power_norm + h_bot = project_entity(e=h, e_p=h_p, r_p=r_p) + t_bot = project_entity(e=t, e_p=t_p, r_p=r_p) + return -((h_bot + r - t_bot) ** p).sum() + class TransETests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransE interaction function.""" From 3b2bcc66cf3d00c92245f45fc0c1b32cb761560c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:50:00 +0100 Subject: [PATCH 419/690] Move utils test to test_utils --- tests/test_models.py | 36 ++---------------------------------- tests/test_utils.py | 34 +++++++++++++++++++++++++++++++++- 2 files changed, 35 insertions(+), 35 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 9adf78bbda..2e32c0111f 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -40,11 +40,11 @@ inverse_outdegree_edge_weights, symmetric_edge_weights, ) -from pykeen.nn import Embedding, RepresentationModule +from pykeen.nn import RepresentationModule from pykeen.regularizers import LpRegularizer, collect_regularization_terms from pykeen.training import LCWATrainingLoop, SLCWATrainingLoop, TrainingLoop from pykeen.triples import TriplesFactory -from pykeen.utils import all_in_bounds, clamp_norm, project_entity, set_random_seed +from pykeen.utils import all_in_bounds, set_random_seed SKIP_MODULES = { Model.__name__, @@ -877,38 +877,6 @@ def _check_constraints(self): for emb in (self.model.entity_representations[0], self.model.relation_representations[0]): assert all_in_bounds(emb(indices=None).norm(p=2, dim=-1), high=1., a_tol=_EPSILON) - def test_project_entity(self): - """Test _project_entity.""" - self.assertIsInstance(self.model, pykeen.models.TransD) - - # random entity embeddings & projections - e = torch.rand(1, self.model.num_entities, self.embedding_dim, generator=self.generator) - e = clamp_norm(e, maxnorm=1, p=2, dim=-1) - e_p = torch.rand(1, self.model.num_entities, self.embedding_dim, generator=self.generator) - - # random relation embeddings & projections - relation_rep: Embedding = self.model.relation_representations[0] - self.assertIsInstance(relation_rep, Embedding) - relation_dim = relation_rep.embedding_dim - - r_p = torch.rand(self.batch_size, 1, relation_dim, generator=self.generator) - - # project - e_bot = project_entity(e=e, e_p=e_p, r_p=r_p) - - # check shape: - assert e_bot.shape == (self.batch_size, self.model.num_entities, relation_dim) - - # check normalization - assert (torch.norm(e_bot, dim=-1, p=2) <= 1.0 + 1.0e-06).all() - - # check equivalence of re-formulation - # e_{\bot} = M_{re} e = (r_p e_p^T + I^{d_r \times d_e}) e - # = r_p (e_p^T e) + e' - M_re = (r_p.unsqueeze(dim=-1) @ e_p.unsqueeze(dim=-2) + torch.eye(relation_dim, self.embedding_dim).unsqueeze(dim=0)) - e_vanilla = (M_re @ e.unsqueeze(dim=-1)).squeeze(dim=-1) - assert torch.allclose(e_vanilla, e_bot) - class TestTransE(_DistanceModelTestCase, unittest.TestCase): """Test the TransE model.""" diff --git a/tests/test_utils.py b/tests/test_utils.py index 0affa4fb36..71f76c5565 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -10,7 +10,7 @@ from pykeen.utils import ( _CUDA_OOM_ERROR, _CUDNN_ERROR, clamp_norm, combine_complex, compact_mapping, flatten_dictionary, - get_until_first_blank, is_cuda_oom_error, is_cudnn_error, l2_regularization, split_complex, + get_until_first_blank, is_cuda_oom_error, is_cudnn_error, l2_regularization, project_entity, split_complex, ) @@ -189,3 +189,35 @@ def test_is_cudnn_error(self): self.assertFalse(is_cuda_oom_error(runtime_error=error)) self.assertFalse(is_cudnn_error(runtime_error=self.not_cuda_error)) + + +def test_project_entity(): + """Test _project_entity.""" + batch_size = 2 + embedding_dim = 3 + relation_dim = 5 + num_entities = 7 + + # random entity embeddings & projections + e = torch.rand(1, num_entities, embedding_dim) + e = clamp_norm(e, maxnorm=1, p=2, dim=-1) + e_p = torch.rand(1, num_entities, embedding_dim) + + # random relation embeddings & projections + r_p = torch.rand(batch_size, 1, relation_dim) + + # project + e_bot = project_entity(e=e, e_p=e_p, r_p=r_p) + + # check shape: + assert e_bot.shape == (batch_size, num_entities, relation_dim) + + # check normalization + assert (torch.norm(e_bot, dim=-1, p=2) <= 1.0 + 1.0e-06).all() + + # check equivalence of re-formulation + # e_{\bot} = M_{re} e = (r_p e_p^T + I^{d_r \times d_e}) e + # = r_p (e_p^T e) + e' + M_re = (r_p.unsqueeze(dim=-1) @ e_p.unsqueeze(dim=-2) + torch.eye(relation_dim, embedding_dim).unsqueeze(dim=0)) + e_vanilla = (M_re @ e.unsqueeze(dim=-1)).squeeze(dim=-1) + assert torch.allclose(e_vanilla, e_bot) From dae5654e5369c9c19955ed33d62187839560870a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:52:36 +0100 Subject: [PATCH 420/690] Fix unittest --- tests/test_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index 71f76c5565..2d3ef12a28 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -218,6 +218,9 @@ def test_project_entity(): # check equivalence of re-formulation # e_{\bot} = M_{re} e = (r_p e_p^T + I^{d_r \times d_e}) e # = r_p (e_p^T e) + e' - M_re = (r_p.unsqueeze(dim=-1) @ e_p.unsqueeze(dim=-2) + torch.eye(relation_dim, embedding_dim).unsqueeze(dim=0)) + M_re = r_p.unsqueeze(dim=-1) @ e_p.unsqueeze(dim=-2) + M_re = M_re + torch.eye(relation_dim, embedding_dim).view(1, 1, relation_dim, embedding_dim) + assert M_re.shape == (batch_size, num_entities, relation_dim, embedding_dim) e_vanilla = (M_re @ e.unsqueeze(dim=-1)).squeeze(dim=-1) + e_vanilla = clamp_norm(e_vanilla, p=2, dim=-1, maxnorm=1) assert torch.allclose(e_vanilla, e_bot) From 26ee2acc2ba32c206ecaf5425335d3fb7168d4d2 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:56:16 +0100 Subject: [PATCH 421/690] Add _exp_score to TransH --- tests/test_interactions.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 35c4b92773..0f76e2d0c0 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -512,6 +512,12 @@ class TransHTests(TranslationalInteractionTests, unittest.TestCase): cls = pykeen.nn.modules.TransHInteraction + def _exp_score(self, h, w_r, d_r, t, p, power_norm) -> torch.FloatTensor: # noqa: D102 + assert not power_norm + h, w_r, d_r, t = _strip_dim(h, w_r, d_r, t) + h, t = [x - (x * w_r).sum() * w_r for x in (h, t)] + return -(h + d_r - t).norm(p=p) + class TransRTests(TranslationalInteractionTests, unittest.TestCase): """Tests for TransR interaction function.""" From c03c78cd9ca97f0154c64d6f9627a75b49f7f9dc Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 13:58:50 +0100 Subject: [PATCH 422/690] Add _exp_score to TransR --- tests/test_interactions.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 0f76e2d0c0..1c4ce06c63 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -12,7 +12,7 @@ import pykeen.nn.modules from pykeen.nn.modules import Interaction, TranslationalInteraction from pykeen.typing import Representation -from pykeen.utils import get_subclasses, project_entity, view_complex +from pykeen.utils import clamp_norm, get_subclasses, project_entity, view_complex T = TypeVar("T") @@ -538,6 +538,12 @@ def test_manual(self): first_score = scores[0].item() self.assertAlmostEqual(first_score, -32, delta=1.0e-04) + def _exp_score(self, h, r, m_r, t, p, power_norm) -> torch.FloatTensor: + assert power_norm + h, r, m_r, t = _strip_dim(h, r, m_r, t) + h_bot, t_bot = [clamp_norm(x.unsqueeze(dim=0) @ m_r, p=2, dim=-1, maxnorm=1.) for x in (h, t)] + return -((h_bot + r - t_bot) ** p).sum() + class SETests(TranslationalInteractionTests, unittest.TestCase): """Tests for SE interaction function.""" From 4b9cdd6ffc4575b48d011900839465a22378b981 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 14:19:15 +0100 Subject: [PATCH 423/690] Fix docstring --- src/pykeen/nn/functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 59f8b0e39d..9e507c850e 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -645,7 +645,7 @@ def structured_embedding_interaction( Evaluate the Structured Embedding interaction function. .. math :: - f(h, r, t) = \|R_h r - R_t t\| + f(h, r, t) = -\|R_h h - R_t t\| :param h: shape: (batch_size, num_heads, dim) The head representations. From 714de0e0dc2d781457bee7ffd8ecd08f5f2cade6 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 14:20:13 +0100 Subject: [PATCH 424/690] Add _exp_score for SE --- tests/test_interactions.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 1c4ce06c63..3295eb687b 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -550,6 +550,14 @@ class SETests(TranslationalInteractionTests, unittest.TestCase): cls = pykeen.nn.modules.StructuredEmbeddingInteraction + def _exp_score(self, h, t, r_h, r_t, p, power_norm) -> torch.FloatTensor: + assert not power_norm + # -\|R_h h - R_t t\| + h, t, r_h, r_t = _strip_dim(h, t, r_h, r_t) + h = r_h @ h.unsqueeze(dim=-1) + t = r_t @ t.unsqueeze(dim=-1) + return -(h - t).norm(p) + class UMTests(TranslationalInteractionTests, unittest.TestCase): """Tests for UM interaction function.""" From 44732379305b32a2b9634f1a78f75f6969b0cd57 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 14:21:16 +0100 Subject: [PATCH 425/690] Add _exp_score for UM --- tests/test_interactions.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 3295eb687b..afbfdb7e8d 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -564,6 +564,12 @@ class UMTests(TranslationalInteractionTests, unittest.TestCase): cls = pykeen.nn.modules.UnstructuredModelInteraction + def _exp_score(self, h, t, p, power_norm) -> torch.FloatTensor: + assert power_norm + # -\|h - t\| + h, t = _strip_dim(h, t) + return -(h - t).pow(p).sum() + class InteractionTestsTest(TestsTest[Interaction], unittest.TestCase): """Test for tests for all interaction functions.""" From e769b36d68f67d30ac14f3f9caaa7c28c7a83e26 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 14:21:53 +0100 Subject: [PATCH 426/690] Make _exp_score necessary --- tests/test_interactions.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index afbfdb7e8d..3e0c20b92e 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -3,6 +3,7 @@ """Tests for interaction functions.""" import unittest +from abc import abstractmethod from typing import Any, Collection, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union from unittest.case import SkipTest @@ -247,9 +248,10 @@ def test_scores(self): scores_f_manual = self._exp_score(**kwargs) assert torch.allclose(scores_f_manual, scores_f) + @abstractmethod def _exp_score(self, **kwargs) -> torch.FloatTensor: """Compute the expected score for a single-score batch.""" - raise SkipTest("No score check implemented.") + raise NotImplementedError class ComplExTests(InteractionTests, unittest.TestCase): From 7be40cb589a49c6344e00a733e6efc8c19859340 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 14:32:08 +0100 Subject: [PATCH 427/690] add _exp_score for TuckER --- tests/test_interactions.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 3e0c20b92e..5c3d07ceb1 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -251,7 +251,7 @@ def test_scores(self): @abstractmethod def _exp_score(self, **kwargs) -> torch.FloatTensor: """Compute the expected score for a single-score batch.""" - raise NotImplementedError + raise NotImplementedError(f"{self.cls.__name__}({sorted(kwargs.keys())})") class ComplExTests(InteractionTests, unittest.TestCase): @@ -431,6 +431,16 @@ class TuckerTests(InteractionTests, unittest.TestCase): embedding_dim=InteractionTests.dim, ) + def _exp_score(self, bn1, bn2, core_tensor, do0, do1, do2, h, r, t) -> torch.FloatTensor: + # DO(BN(DO(BN(h)) x_1 DO(W x_2 r))) x_3 t + # = A x_3 t + # with A = DO(BN(B)), B = C x_1 D = , C=DO(BN(h)), D = DO(W x_2 r) + h, r, t = _strip_dim(h, r, t) + c = do0((core_tensor * r[None, :, None]).sum(dim=1)) # shape: (embedding_dim, embedding_dim) + b = do1(bn1(h.view(1, -1))).view(-1) # shape: (embedding_dim) + aa = (b[:, None] * c).sum(dim=1) # shape: (embedding_dim) + return do2(bn2((aa.view(1, -1)))) + class RotatETests(InteractionTests, unittest.TestCase): """Tests for RotatE interaction function.""" From c35f8b00264005ac08ef6451b2846ce1648c4f3d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 14:46:05 +0100 Subject: [PATCH 428/690] Add _exp_score for KG2E --- tests/test_interactions.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 5c3d07ceb1..31ccd9474b 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -377,7 +377,7 @@ class NTNTests(InteractionTests, unittest.TestCase): ) def _exp_score(self, h, t, w, vt, vh, b, u, activation) -> torch.FloatTensor: - # f(h,r,t) = u_r^T act(h W_r t + V_r h + V_r' t + b_r) + # f(h,r,t) = u_r^T act(h W_r t + V_r h + V_r t + b_r) # shapes: w: (k, dim, dim), vh/vt: (k, dim), b/u: (k,), h/t: (dim,) # remove batch/num dimension h, t, w, vt, vh, b, u = _strip_dim(h, t, w, vt, vh, b, u) @@ -422,6 +422,14 @@ class KG2ETests(InteractionTests, unittest.TestCase): cls = pykeen.nn.modules.KG2EInteraction + def _exp_score(self, exact, h_mean, h_var, r_mean, r_var, similarity, t_mean, t_var): + assert similarity == "KL" + h_mean, h_var, r_mean, r_var, t_mean, t_var = _strip_dim(h_mean, h_var, r_mean, r_var, t_mean, t_var) + e_mean, e_var = h_mean - t_mean, h_var + t_var + p = torch.distributions.MultivariateNormal(loc=e_mean, covariance_matrix=torch.diag(e_var)) + q = torch.distributions.MultivariateNormal(loc=r_mean, covariance_matrix=torch.diag(r_var)) + return torch.distributions.kl.kl_divergence(p, q) + class TuckerTests(InteractionTests, unittest.TestCase): """Tests for Tucker interaction function.""" From b2c7ad0f539e92a03a70eb8cb2211bb6325f3bb3 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 14:57:14 +0100 Subject: [PATCH 429/690] Do not skip forward tests for all shapes for batch norm interaction modules --- tests/test_interactions.py | 32 ++++++++++++++++++++------------ 1 file changed, 20 insertions(+), 12 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 31ccd9474b..dc40fe21b9 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -2,6 +2,7 @@ """Tests for interaction functions.""" +import logging import unittest from abc import abstractmethod from typing import Any, Collection, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union @@ -16,6 +17,7 @@ from pykeen.utils import clamp_norm, get_subclasses, project_entity, view_complex T = TypeVar("T") +logger = logging.getLogger(__name__) class GenericTests(Generic[T]): @@ -221,18 +223,24 @@ def _get_output_shape( def test_forward(self): """Test forward.""" for hs, rs, ts in self._get_test_shapes(): - if any(isinstance(m, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d)) for m in self.instance.modules()): - # TODO: do we need to skip this for every combination? or only if batch_size = 1? - continue - with self.subTest(f"forward({hs}, {rs}, {ts})"): - h, r, t = self._get_hrt(hs, rs, ts) - scores = self.instance(h=h, r=r, t=t) - expected_shape = self._get_output_shape(hs, rs, ts) - self._check_scores(scores=scores, exp_shape=expected_shape) - with self.subTest(f"forward({hs}, {rs}, {ts}) - consistency with functional"): - kwargs = self.instance._prepare_for_functional(h=h, r=r, t=t) - scores_f = self.cls.func(**kwargs) - assert torch.allclose(scores, scores_f) + try: + with self.subTest(f"forward({hs}, {rs}, {ts})"): + h, r, t = self._get_hrt(hs, rs, ts) + scores = self.instance(h=h, r=r, t=t) + expected_shape = self._get_output_shape(hs, rs, ts) + self._check_scores(scores=scores, exp_shape=expected_shape) + with self.subTest(f"forward({hs}, {rs}, {ts}) - consistency with functional"): + kwargs = self.instance._prepare_for_functional(h=h, r=r, t=t) + scores_f = self.cls.func(**kwargs) + assert torch.allclose(scores, scores_f) + except ValueError as error: + # check whether the error originates from batch norm for single element batches + small_batch_size = any(s[0] == 1 for s in (hs, rs, ts)) + has_batch_norm = any(isinstance(m, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d)) for m in self.instance.modules()) + if small_batch_size and has_batch_norm: + logger.warning(f"Skipping test for shapes {hs}, {rs}, {ts}") + continue + raise error def test_scores(self): """Test individual scores.""" From 173e70b6ec29746207222ab8c419e9d247ce4346 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 14:58:58 +0100 Subject: [PATCH 430/690] Remove unused deprecated base Model property relation_dim --- src/pykeen/models/base.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 8f956ceed4..c61883d082 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1335,14 +1335,6 @@ def embedding_dim(self) -> int: # noqa:D401 assert isinstance(embedding, Embedding) return embedding.embedding_dim - @property - def relation_dim(self) -> int: # noqa:D401 - """The relation embedding dim.""" - # TODO: Deprecated; directly use self.relation_representations[0].embedding_dim instead? - embedding = self.relation_representations[0] - assert isinstance(embedding, Embedding) - return embedding.embedding_dim - class DoubleRelationEmbeddingModel(ERModel, autoreset=False): """A KGEM that stores one :class:`pykeen.nn.Embedding` for entities and two for relations. From d21138cdbc58200854cbcf7e379147c15493cd15 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:01:36 +0100 Subject: [PATCH 431/690] Move LiteralModel base class to multimodal module --- src/pykeen/models/__init__.py | 3 +- src/pykeen/models/base.py | 70 +----------------- src/pykeen/models/multimodal/base.py | 74 +++++++++++++++++++ .../models/multimodal/complex_literal.py | 7 +- .../models/multimodal/distmult_literal.py | 7 +- tests/test_models.py | 3 +- 6 files changed, 85 insertions(+), 79 deletions(-) create mode 100644 src/pykeen/models/multimodal/base.py diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index d433331506..e2b6f0bffd 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -9,9 +9,10 @@ from typing import Mapping, Set, Type, Union from .base import ( # noqa:F401 - DoubleRelationEmbeddingModel, ERModel, LiteralModel, Model, SingleVectorEmbeddingModel, + DoubleRelationEmbeddingModel, ERModel, Model, SingleVectorEmbeddingModel, TwoSideEmbeddingModel, TwoVectorEmbeddingModel, ) +from .multimodal.base import LiteralModel from .multimodal import ComplExLiteral, DistMultLiteral from .unimodal import ( ComplEx, diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index c61883d082..0d7577f35a 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -19,10 +19,10 @@ from torch import nn from ..losses import Loss, MarginRankingLoss, NSSALoss -from ..nn import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule +from ..nn import Embedding, EmbeddingSpecification, RepresentationModule from ..nn.modules import Interaction from ..regularizers import Regularizer, collect_regularization_terms -from ..triples import TriplesFactory, TriplesNumericLiteralsFactory +from ..triples import TriplesFactory from ..typing import DeviceHint, HeadRepresentation, MappedTriples, RelationRepresentation, TailRepresentation from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed @@ -36,7 +36,6 @@ 'DoubleRelationEmbeddingModel', 'TwoVectorEmbeddingModel', 'TwoSideEmbeddingModel', - 'LiteralModel', ] logger = logging.getLogger(__name__) @@ -1546,68 +1545,3 @@ def forward( (self.entity_representations[1], self.relation_representations[1], self.entity_representations[0]), ) ) - - -class LiteralModel(ERModel, autoreset=False): - """Base class for models with entity literals.""" - - # TODO: Move to other file? - - def __init__( - self, - triples_factory: TriplesNumericLiteralsFactory, - embedding_dim: int, - interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], - combination: nn.Module, - entity_specification: Optional[EmbeddingSpecification] = None, - relation_specification: Optional[EmbeddingSpecification] = None, - loss: Optional[Loss] = None, - predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, - preferred_device: DeviceHint = None, - random_seed: Optional[int] = None, - ): - super().__init__( - triples_factory=triples_factory, - interaction=interaction, - loss=loss, - predict_with_sigmoid=predict_with_sigmoid, - automatic_memory_optimization=automatic_memory_optimization, - preferred_device=preferred_device, - random_seed=random_seed, - entity_representations=[ - # entity embeddings - Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, - specification=entity_specification, - ), - # Entity literals - LiteralRepresentations( - numeric_literals=torch.as_tensor(triples_factory.numeric_literals, dtype=torch.float32), - ), - ], - relation_representations=Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim, - specification=relation_specification, - ), - ) - self.combination = combination - - def forward( - self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], - slice_size: Optional[int] = None, - slice_dim: Optional[str] = None, - ) -> torch.FloatTensor: # noqa: D102 - h, r, t = self._get_representations(h_indices, r_indices, t_indices) - # combine entity embeddings + literals - h, t = [ - self.combination(torch.cat(x, dim=-1)) - for x in (h, t) - ] - scores = self.interaction.score(h=h, r=r, t=t, slice_size=slice_size, slice_dim=slice_dim) - return self._repeat_scores_if_necessary(scores, h_indices, r_indices, t_indices) diff --git a/src/pykeen/models/multimodal/base.py b/src/pykeen/models/multimodal/base.py new file mode 100644 index 0000000000..1be83b9acf --- /dev/null +++ b/src/pykeen/models/multimodal/base.py @@ -0,0 +1,74 @@ +"""Base classes for multi-modal models.""" +from typing import Optional + +import torch +from torch import nn + +from pykeen.losses import Loss +from pykeen.models import ERModel +from pykeen.nn import Embedding, EmbeddingSpecification, Interaction, LiteralRepresentations +from pykeen.triples import TriplesNumericLiteralsFactory +from pykeen.typing import DeviceHint, HeadRepresentation, RelationRepresentation, TailRepresentation + + +class LiteralModel(ERModel, autoreset=False): + """Base class for models with entity literals.""" + + def __init__( + self, + triples_factory: TriplesNumericLiteralsFactory, + embedding_dim: int, + interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], + combination: nn.Module, + entity_specification: Optional[EmbeddingSpecification] = None, + relation_specification: Optional[EmbeddingSpecification] = None, + loss: Optional[Loss] = None, + predict_with_sigmoid: bool = False, + automatic_memory_optimization: Optional[bool] = None, + preferred_device: DeviceHint = None, + random_seed: Optional[int] = None, + ): + super().__init__( + triples_factory=triples_factory, + interaction=interaction, + loss=loss, + predict_with_sigmoid=predict_with_sigmoid, + automatic_memory_optimization=automatic_memory_optimization, + preferred_device=preferred_device, + random_seed=random_seed, + entity_representations=[ + # entity embeddings + Embedding.from_specification( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + specification=entity_specification, + ), + # Entity literals + LiteralRepresentations( + numeric_literals=torch.as_tensor(triples_factory.numeric_literals, dtype=torch.float32), + ), + ], + relation_representations=Embedding.from_specification( + num_embeddings=triples_factory.num_relations, + embedding_dim=embedding_dim, + specification=relation_specification, + ), + ) + self.combination = combination + + def forward( + self, + h_indices: Optional[torch.LongTensor], + r_indices: Optional[torch.LongTensor], + t_indices: Optional[torch.LongTensor], + slice_size: Optional[int] = None, + slice_dim: Optional[str] = None, + ) -> torch.FloatTensor: # noqa: D102 + h, r, t = self._get_representations(h_indices, r_indices, t_indices) + # combine entity embeddings + literals + h, t = [ + self.combination(torch.cat(x, dim=-1)) + for x in (h, t) + ] + scores = self.interaction.score(h=h, r=r, t=t, slice_size=slice_size, slice_dim=slice_dim) + return self._repeat_scores_if_necessary(scores, h_indices, r_indices, t_indices) diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index f0147cd0f1..c9600a87d6 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -6,9 +6,8 @@ import torch import torch.nn as nn -from torch.nn.init import xavier_normal_ -from ..base import LiteralModel +from .base import LiteralModel from ...losses import BCEWithLogitsLoss, Loss from ...nn import EmbeddingSpecification from ...nn.modules import ComplExInteraction @@ -95,10 +94,10 @@ def __init__( dropout=input_dropout, ), entity_specification=EmbeddingSpecification( - initializer=xavier_normal_, + initializer=nn.init.xavier_normal_, ), relation_specification=EmbeddingSpecification( - initializer=xavier_normal_, + initializer=nn.init.xavier_normal_, ), loss=loss, predict_with_sigmoid=predict_with_sigmoid, diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 27fe67df09..7236aecf0e 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -5,9 +5,8 @@ from typing import Optional, TYPE_CHECKING import torch.nn as nn -from torch.nn.init import xavier_normal_ -from ..base import LiteralModel +from .base import LiteralModel from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.modules import DistMultInteraction @@ -55,10 +54,10 @@ def __init__( nn.Dropout(input_dropout), ), entity_specification=EmbeddingSpecification( - initializer=xavier_normal_, + initializer=nn.init.xavier_normal_, ), relation_specification=EmbeddingSpecification( - initializer=xavier_normal_, + initializer=nn.init.xavier_normal_, ), loss=loss, predict_with_sigmoid=predict_with_sigmoid, diff --git a/tests/test_models.py b/tests/test_models.py index 2e32c0111f..024ca35f67 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -22,11 +22,10 @@ import pykeen.models from pykeen.datasets.kinships import KINSHIPS_TRAIN_PATH from pykeen.datasets.nations import NATIONS_TEST_PATH, NATIONS_TRAIN_PATH, Nations -from pykeen.models import _BASE_MODELS, _MODELS +from pykeen.models import LiteralModel, _BASE_MODELS, _MODELS from pykeen.models.base import ( DoubleRelationEmbeddingModel, ERModel, - LiteralModel, Model, SingleVectorEmbeddingModel, TwoSideEmbeddingModel, From 1b7f611f85ff4becd14161c0946634ef1c64cce3 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:04:06 +0100 Subject: [PATCH 432/690] Add documentation to _extract_sizes --- src/pykeen/nn/functional.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 9e507c850e..f8279a0da1 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -40,8 +40,12 @@ ] -# TODO @mberr documentation -def _extract_sizes(h, r, t) -> Tuple[int, int, int, int, int]: +def _extract_sizes( + h: torch.Tensor, + r: torch.Tensor, + t: torch.Tensor, +) -> Tuple[int, int, int, int, int]: + """Utility to extract size dimensions from head/relation/tail representations.""" num_heads, num_relations, num_tails = [xx.shape[1] for xx in (h, r, t)] d_e = h.shape[-1] d_r = r.shape[-1] From 2acfbacd07933348e2fb9de7bb4b02dfc3c2b2f5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:05:15 +0100 Subject: [PATCH 433/690] Add documentation to _apply_optional_bn_to_tensor --- src/pykeen/nn/functional.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index f8279a0da1..eedacb9fa2 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -52,19 +52,18 @@ def _extract_sizes( return num_heads, num_relations, num_tails, d_e, d_r -# TODO @mberr documentation def _apply_optional_bn_to_tensor( batch_norm: Optional[nn.BatchNorm1d], output_dropout: nn.Dropout, tensor: torch.FloatTensor, ) -> torch.FloatTensor: + """Apply optional batch normalization and dropout layer. Supports multiple batch dimensions.""" if batch_norm is not None: shape = tensor.shape tensor = tensor.reshape(-1, shape[-1]) tensor = batch_norm(tensor) tensor = tensor.view(*shape) - tensor = output_dropout(tensor) - return tensor + return output_dropout(tensor) def _translational_interaction( From 59b97a9c5cfe9ca53348977cdb95fafd8292621b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:06:26 +0100 Subject: [PATCH 434/690] Reword todo notice --- src/pykeen/regularizers.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/regularizers.py b/src/pykeen/regularizers.py index 1784805141..af41e7b51f 100644 --- a/src/pykeen/regularizers.py +++ b/src/pykeen/regularizers.py @@ -99,7 +99,7 @@ class NoRegularizer(Regularizer): Used to simplify code. """ - # TODO: Deprecated + # TODO(cthoyt): Deprecated, but used in notebooks / README #: The default strategy for optimizing the no-op regularizer's hyper-parameters hpo_default: ClassVar[Mapping[str, Any]] = {} From 912bc8d7ca0618a214855c4e3771e00c75540de9 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:07:08 +0100 Subject: [PATCH 435/690] Remove unused utility --- src/pykeen/utils.py | 26 -------------------------- tests/test_utils.py | 28 +--------------------------- 2 files changed, 1 insertion(+), 53 deletions(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 2fb12f09e6..1b23a59ee6 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -10,7 +10,6 @@ from io import BytesIO from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, SupportsFloat, Tuple, Type, TypeVar, Union -import numpy import numpy as np import pandas as pd import torch @@ -31,7 +30,6 @@ 'invert_mapping', 'is_cudnn_error', 'is_cuda_oom_error', - 'l2_regularization', 'random_non_negative_int', 'resolve_device', 'slice_triples', @@ -58,30 +56,6 @@ _CUDA_OOM_ERROR = 'CUDA out of memory.' -# TODO remove (unused) -def l2_regularization( - *xs: torch.Tensor, - normalize: bool = False, -) -> torch.Tensor: - """ - Compute squared L2-regularization term. - - :param xs: a list of torch.Tensor - The tensors for which to compute the regularization. - :param normalize: - Whether to divide the term by the total number of elements in the tensors. - - :return: The sum of squared value across all tensors. - """ - regularization_term = sum(x.pow(2).sum() for x in xs) - - # Normalize by the number of elements in the tensors for dimensionality-independent weight tuning. - if normalize: - regularization_term /= sum(numpy.prod(x.shape) for x in xs) - - return regularization_term - - def resolve_device(device: DeviceHint = None) -> torch.device: """Resolve a torch.device given a desired device (string).""" if device is None or device == 'gpu': diff --git a/tests/test_utils.py b/tests/test_utils.py index 2d3ef12a28..0e5cfcd3f2 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -5,40 +5,14 @@ import string import unittest -import numpy import torch from pykeen.utils import ( _CUDA_OOM_ERROR, _CUDNN_ERROR, clamp_norm, combine_complex, compact_mapping, flatten_dictionary, - get_until_first_blank, is_cuda_oom_error, is_cudnn_error, l2_regularization, project_entity, split_complex, + get_until_first_blank, is_cuda_oom_error, is_cudnn_error, project_entity, split_complex, ) -class L2RegularizationTest(unittest.TestCase): - """Test L2 regularization.""" - - def test_one_tensor(self): - """Test if output is correct for a single tensor.""" - t = torch.ones(1, 2, 3, 4) - reg = l2_regularization(t) - self.assertAlmostEqual(float(reg), float(numpy.prod(t.shape))) - - def test_many_tensors(self): - """Test if output is correct for var-args.""" - ts = [] - exp_reg = 0. - for i, shape in enumerate([ - (1, 2, 3), - (2, 3, 4), - (3, 4, 5), - ]): - t = torch.ones(*shape) * (i + 1) - ts.append(t) - exp_reg += numpy.prod(t.shape) * (i + 1) ** 2 - reg = l2_regularization(*ts) - self.assertAlmostEqual(float(reg), exp_reg) - - class FlattenDictionaryTest(unittest.TestCase): """Test flatten_dictionary.""" From aa00e52e10861f82c1bfb824c3c85f6c657ef1ec Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:08:20 +0100 Subject: [PATCH 436/690] Remove unused utilities --- src/pykeen/triples/triples_factory.py | 10 ++++------ src/pykeen/utils.py | 20 -------------------- 2 files changed, 4 insertions(+), 26 deletions(-) diff --git a/src/pykeen/triples/triples_factory.py b/src/pykeen/triples/triples_factory.py index c8ac719cf1..a920d97bd7 100644 --- a/src/pykeen/triples/triples_factory.py +++ b/src/pykeen/triples/triples_factory.py @@ -17,7 +17,7 @@ from .instances import LCWAInstances, SLCWAInstances from .utils import load_triples from ..typing import EntityMapping, LabeledTriples, MappedTriples, RelationMapping -from ..utils import compact_mapping, invert_mapping, random_non_negative_int, slice_triples +from ..utils import compact_mapping, invert_mapping, random_non_negative_int __all__ = [ 'TriplesFactory', @@ -135,14 +135,12 @@ def _map_triples_elements_to_ids( logger.warning('Provided empty triples to map.') return torch.empty(0, 3, dtype=torch.long) - heads, relations, tails = slice_triples(triples) - # When triples that don't exist are trying to be mapped, they get the id "-1" entity_getter = np.vectorize(entity_to_id.get) - head_column = entity_getter(heads, [-1]) - tail_column = entity_getter(tails, [-1]) + head_column = entity_getter(triples[:, 0], [-1]) + tail_column = entity_getter(triples[:, 2], [-1]) relation_getter = np.vectorize(relation_to_id.get) - relation_column = relation_getter(relations, [-1]) + relation_column = relation_getter(triples[:, 1], [-1]) # Filter all non-existent triples head_filter = head_column < 0 diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 1b23a59ee6..a18eda9bea 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -32,8 +32,6 @@ 'is_cuda_oom_error', 'random_non_negative_int', 'resolve_device', - 'slice_triples', - 'slice_doubles', 'split_complex', 'split_list_in_batches_iter', 'split_list_in_batches', @@ -68,24 +66,6 @@ def resolve_device(device: DeviceHint = None) -> torch.device: return device -def slice_triples(triples): - """Get the heads, relations, and tails from a matrix of triples.""" - return ( - triples[:, 0:1], # heads - triples[:, 1:2], # relations - triples[:, 2:3], # tails - ) - - -# TODO remove (unused/untested) -def slice_doubles(doubles): - """Get the heads and relations from a matrix of doubles.""" - return ( - doubles[:, 0:1], # heads - doubles[:, 1:2], # relations - ) - - X = TypeVar('X') From 9789e8e585206ecacdb48f3bc2f0136377cb3587 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:08:59 +0100 Subject: [PATCH 437/690] Remove unused method --- src/pykeen/utils.py | 7 ------- 1 file changed, 7 deletions(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index a18eda9bea..5a86154344 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -34,7 +34,6 @@ 'resolve_device', 'split_complex', 'split_list_in_batches_iter', - 'split_list_in_batches', 'normalize_string', 'normalized_lookup', 'negative_norm_of_sum', @@ -69,12 +68,6 @@ def resolve_device(device: DeviceHint = None) -> torch.device: X = TypeVar('X') -# TODO remove (unused) -def split_list_in_batches(input_list: List[X], batch_size: int) -> List[List[X]]: - """Split a list of instances in batches of size batch_size.""" - return list(split_list_in_batches_iter(input_list=input_list, batch_size=batch_size)) - - def split_list_in_batches_iter(input_list: List[X], batch_size: int) -> Iterable[List[X]]: """Split a list of instances in batches of size batch_size.""" return ( From f73a0bee714d764bfa8ae35db8efd1c8e060a733 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:13:00 +0100 Subject: [PATCH 438/690] Fix slice_triples replacement --- src/pykeen/triples/triples_factory.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/triples/triples_factory.py b/src/pykeen/triples/triples_factory.py index a920d97bd7..5575d86ed7 100644 --- a/src/pykeen/triples/triples_factory.py +++ b/src/pykeen/triples/triples_factory.py @@ -137,10 +137,10 @@ def _map_triples_elements_to_ids( # When triples that don't exist are trying to be mapped, they get the id "-1" entity_getter = np.vectorize(entity_to_id.get) - head_column = entity_getter(triples[:, 0], [-1]) - tail_column = entity_getter(triples[:, 2], [-1]) + head_column = entity_getter(triples[:, 0:1], [-1]) + tail_column = entity_getter(triples[:, 2:3], [-1]) relation_getter = np.vectorize(relation_to_id.get) - relation_column = relation_getter(triples[:, 1], [-1]) + relation_column = relation_getter(triples[:, 1:2], [-1]) # Filter all non-existent triples head_filter = head_column < 0 From ff5394fb3965407dc6d857cdbd8b136856b34cd7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:14:55 +0100 Subject: [PATCH 439/690] Add TODO --- tests/test_models.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_models.py b/tests/test_models.py index 024ca35f67..ac470422d8 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -222,6 +222,7 @@ def test_reset_parameters_(self): def _check_scores(self, batch, scores) -> None: """Check the scores produced by a forward function.""" + # TODO: Move score checks to Interaction tests? # check for finite values by default self.assertTrue(torch.all(torch.isfinite(scores)).item(), f'Some scores were not finite:\n{scores}') From 6824ae12b5b6d44b852dc2eaef3a9b9024253c84 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:18:59 +0100 Subject: [PATCH 440/690] Split functional consistency test --- tests/test_interactions.py | 24 +++++++++++++++--------- 1 file changed, 15 insertions(+), 9 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index dc40fe21b9..15905c971b 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -224,15 +224,10 @@ def test_forward(self): """Test forward.""" for hs, rs, ts in self._get_test_shapes(): try: - with self.subTest(f"forward({hs}, {rs}, {ts})"): - h, r, t = self._get_hrt(hs, rs, ts) - scores = self.instance(h=h, r=r, t=t) - expected_shape = self._get_output_shape(hs, rs, ts) - self._check_scores(scores=scores, exp_shape=expected_shape) - with self.subTest(f"forward({hs}, {rs}, {ts}) - consistency with functional"): - kwargs = self.instance._prepare_for_functional(h=h, r=r, t=t) - scores_f = self.cls.func(**kwargs) - assert torch.allclose(scores, scores_f) + h, r, t = self._get_hrt(hs, rs, ts) + scores = self.instance(h=h, r=r, t=t) + expected_shape = self._get_output_shape(hs, rs, ts) + self._check_scores(scores=scores, exp_shape=expected_shape) except ValueError as error: # check whether the error originates from batch norm for single element batches small_batch_size = any(s[0] == 1 for s in (hs, rs, ts)) @@ -242,6 +237,17 @@ def test_forward(self): continue raise error + def test_forward_consistency_with_functional(self): + """Test forward's consistency with functional.""" + # set in eval mode (otherwise there are non-deterministic factors like Dropout + self.instance.eval() + for hs, rs, ts in self._get_test_shapes(): + h, r, t = self._get_hrt(hs, rs, ts) + scores = self.instance(h=h, r=r, t=t) + kwargs = self.instance._prepare_for_functional(h=h, r=r, t=t) + scores_f = self.cls.func(**kwargs) + assert torch.allclose(scores, scores_f) + def test_scores(self): """Test individual scores.""" self.instance.eval() From 96c49e6b64d4fae6b600ec4a6464c45bbe62a946 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:25:08 +0100 Subject: [PATCH 441/690] Fix ComplEx interaction function --- src/pykeen/nn/functional.py | 20 +++++++++++++------- tests/test_interactions.py | 2 +- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index eedacb9fa2..7b6c18cc83 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -115,9 +115,12 @@ def complex_interaction( r: torch.FloatTensor, t: torch.FloatTensor, ) -> torch.FloatTensor: - """ + r""" Evaluate the ComplEx interaction function. + .. math :: + Re(\langle h, r, conj(t) \rangle) + :param h: shape: (batch_size, num_heads, `2*dim`) The complex head representations. :param r: shape: (batch_size, num_relations, 2*dim) @@ -128,14 +131,17 @@ def complex_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ + # Re() + # = - - - + # = + + - (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] return sum( - extended_einsum("bhd,brd,btd->bhrt", hh, rr, tt) - for hh, rr, tt in [ - (h_re, r_re, t_re), - (h_re, r_im, t_im), - (h_im, r_re, t_im), - (h_im, r_im, t_re), + factor * extended_einsum("bhd,brd,btd->bhrt", hh, rr, tt) + for factor, hh, rr, tt in [ + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), ] ) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 15905c971b..66f247d43f 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -275,7 +275,7 @@ class ComplExTests(InteractionTests, unittest.TestCase): def _exp_score(self, h, r, t) -> torch.FloatTensor: # noqa: D102 h, r, t = [view_complex(x) for x in (h, r, t)] - return (h * r * t).sum().real + return (h * r * torch.conj(t)).sum().real class ConvETests(InteractionTests, unittest.TestCase): From e05cc7351104e2559226bb0a6043bdf031f9c623 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:36:50 +0100 Subject: [PATCH 442/690] Extract utility for tensor sum (for future optimization) --- src/pykeen/nn/functional.py | 55 +++++++++++++++++++++---------------- src/pykeen/utils.py | 8 +++++- 2 files changed, 38 insertions(+), 25 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 7b6c18cc83..d5a75b6fd0 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -13,7 +13,7 @@ from ..typing import GaussianDistribution from ..utils import ( broadcast_cat, clamp_norm, extended_einsum, is_cudnn_error, negative_norm_of_sum, project_entity, - split_complex, view_complex, + split_complex, tensor_sum, view_complex, ) __all__ = [ @@ -135,15 +135,15 @@ def complex_interaction( # = - - - # = + + - (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] - return sum( + return tensor_sum(*( factor * extended_einsum("bhd,brd,btd->bhrt", hh, rr, tt) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] - ) + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] + )) @_add_cuda_warning @@ -214,9 +214,7 @@ def conve_interaction( x = (x @ t).squeeze(dim=-2) # add bias term - x = x + t_bias.view(t.shape[0], 1, 1, num_tails) - - return x + return tensor_sum(x, t_bias.view(t.shape[0], 1, 1, num_tails)) def convkb_interaction( @@ -266,7 +264,7 @@ def convkb_interaction( h = (h.view(h.shape[0], h.shape[1], 1, 1, embedding_dim, 1) * conv_head.view(1, 1, 1, 1, 1, num_filters)) r = (r.view(r.shape[0], 1, r.shape[1], 1, embedding_dim, 1) * conv_rel.view(1, 1, 1, 1, 1, num_filters)) t = (t.view(t.shape[0], 1, 1, t.shape[1], embedding_dim, 1) * conv_tail.view(1, 1, 1, 1, 1, num_filters)) - x = activation(conv_bias + h + r + t) + x = activation(tensor_sum(conv_bias, h, r, t)) # Apply dropout, cf. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L54-L56 x = hidden_dropout(x) @@ -333,9 +331,7 @@ def ermlp_interaction( h = h.view(-1, num_heads, 1, 1, embedding_dim) @ head_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) r = r.view(-1, 1, num_relations, 1, embedding_dim) @ rel_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) t = t.view(-1, 1, 1, num_tails, embedding_dim) @ tail_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) - # TODO: Choosing which to combine first, h/r, h/t or r/t, depending on the shape might further improve - # performance in a 1:n scenario. - return final(activation(bias + h + r + t)).squeeze(dim=-1) + return final(activation(tensor_sum(bias, h, r, t))).squeeze(dim=-1) def ermlpe_interaction( @@ -490,10 +486,11 @@ def ntn_interaction( """ # save sizes num_heads, num_relations, num_tails, _, num_slices = _extract_sizes(h, b, t) - x = extended_einsum("bhd,brkde,bte->bhrtk", h, w, t) - x = x + extended_einsum("brkd,bhd->bhrk", vh, h).unsqueeze(dim=3) - x = x + extended_einsum("brkd,btd->brtk", vt, t).unsqueeze(dim=1) - x = activation(x) + x = activation(tensor_sum( + extended_einsum("bhd,brkde,bte->bhrtk", h, w, t), + extended_einsum("brkd,bhd->bhrk", vh, h).unsqueeze(dim=3), + extended_einsum("brkd,btd->brtk", vt, t).unsqueeze(dim=1), + )) x = extended_einsum("bhrtk,brk->bhrt", x, u) return x @@ -540,7 +537,11 @@ def proje_interaction( h = h * d_e.view(1, 1, dim) r = r * d_r.view(1, 1, dim) # combination, shape: (b, h, r, d) - x = h.unsqueeze(dim=2) + r.unsqueeze(dim=1) + b_c.view(1, 1, 1, dim) + x = tensor_sum( + h.unsqueeze(dim=2), + r.unsqueeze(dim=1), + b_c.view(1, 1, 1, dim), + ) x = activation(x) # dot product with t, shape: (b, h, r, t) return (x @ t.unsqueeze(dim=1).transpose(-2, -1)) + b_p @@ -633,6 +634,7 @@ def simple_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ + # TODO: unused scores = 0.5 * (distmult_interaction(h=h, r=r, t=t) + distmult_interaction(h=h_inv, r=r_inv, t=t_inv)) # Note: In the code in their repository, the score is clamped to [-20, 20]. # That is not mentioned in the paper, so it is made optional here. @@ -792,10 +794,15 @@ def transh_interaction( The scores. """ # Project to hyperplane - return _translational_interaction( - h=(h.unsqueeze(dim=2) - extended_einsum("bhd,brd,bre->bhre", h, w_r, w_r)).unsqueeze(dim=3), - r=d_r.view(d_r.shape[0], 1, d_r.shape[1], 1, d_r.shape[2]), - t=(t.unsqueeze(dim=1) - extended_einsum("btd,brd,bre->brte", t, w_r, w_r)).unsqueeze(dim=1), + return negative_norm_of_sum( + # h + h.unsqueeze(dim=2), + -extended_einsum("bhd,brd,bre->bhre", h, w_r, w_r).unsqueeze(dim=3), + # r + d_r.view(d_r.shape[0], 1, d_r.shape[1], 1, d_r.shape[2]), + # -t + -t.unsqueeze(dim=1), + +extended_einsum("btd,brd,bre->brte", t, w_r, w_r).unsqueeze(dim=1), p=p, power_norm=power_norm, ) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 5a86154344..0ebef8a53e 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -493,6 +493,12 @@ def complex_normalize(x: torch.Tensor) -> torch.Tensor: return x +def tensor_sum(*x: torch.FloatTensor) -> torch.FloatTensor: + """Compute sum of tensors in brodcastable shape.""" + # TODO: Optimize order + return sum(x) + + def negative_norm_of_sum( *x: torch.FloatTensor, p: Union[str, int] = 2, @@ -510,7 +516,7 @@ def negative_norm_of_sum( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - d: torch.FloatTensor = sum(x) + d: torch.FloatTensor = tensor_sum(x) if power_norm: assert isinstance(p, SupportsFloat) return -(d.abs() ** p).sum(dim=-1) From c7728045a0d604debca37dd8b045f155dba73c78 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:40:46 +0100 Subject: [PATCH 443/690] Post refactoring fixes --- src/pykeen/nn/functional.py | 9 +++++---- src/pykeen/utils.py | 2 +- 2 files changed, 6 insertions(+), 5 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index d5a75b6fd0..c2d49e79e0 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -793,16 +793,17 @@ def transh_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ + num_heads, num_relations, num_tails, dim = _extract_sizes(h, w_r, t)[:4] # Project to hyperplane return negative_norm_of_sum( # h - h.unsqueeze(dim=2), + h.view(h.shape[0], num_heads, 1, 1, dim), -extended_einsum("bhd,brd,bre->bhre", h, w_r, w_r).unsqueeze(dim=3), # r - d_r.view(d_r.shape[0], 1, d_r.shape[1], 1, d_r.shape[2]), + d_r.view(d_r.shape[0], 1, num_relations, 1, dim), # -t - -t.unsqueeze(dim=1), - +extended_einsum("btd,brd,bre->brte", t, w_r, w_r).unsqueeze(dim=1), + -t.view(t.shape[0], 1, 1, num_tails, dim), + extended_einsum("btd,brd,bre->brte", t, w_r, w_r).unsqueeze(dim=1), p=p, power_norm=power_norm, ) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 0ebef8a53e..6d19f30fce 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -516,7 +516,7 @@ def negative_norm_of_sum( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - d: torch.FloatTensor = tensor_sum(x) + d: torch.FloatTensor = tensor_sum(*x) if power_norm: assert isinstance(p, SupportsFloat) return -(d.abs() ** p).sum(dim=-1) From f9899c988cd65ca0511ea37625bf2dad7ec00527 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 15:56:49 +0100 Subject: [PATCH 444/690] Fix TuckER tests, and improve variable names --- src/pykeen/models/unimodal/tucker.py | 6 +-- src/pykeen/nn/functional.py | 78 +++++++++++++++++----------- src/pykeen/nn/modules.py | 41 +++++++++------ tests/test_interactions.py | 15 +++--- 4 files changed, 81 insertions(+), 59 deletions(-) diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index 739d7ff529..3df8914d92 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -90,9 +90,9 @@ def __init__( interaction=TuckerInteraction( embedding_dim=embedding_dim, relation_dim=relation_dim, - dropout_0=dropout_0, - dropout_1=dropout_1, - dropout_2=dropout_2, + head_dropout=dropout_0, + relation_dropout=dropout_1, + head_relation_dropout=dropout_2, apply_batch_normalization=apply_batch_normalization, ), embedding_dim=embedding_dim, diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index c2d49e79e0..419cfa7740 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -53,17 +53,17 @@ def _extract_sizes( def _apply_optional_bn_to_tensor( - batch_norm: Optional[nn.BatchNorm1d], + x: torch.FloatTensor, output_dropout: nn.Dropout, - tensor: torch.FloatTensor, + batch_norm: Optional[nn.BatchNorm1d] = None, ) -> torch.FloatTensor: """Apply optional batch normalization and dropout layer. Supports multiple batch dimensions.""" if batch_norm is not None: - shape = tensor.shape - tensor = tensor.reshape(-1, shape[-1]) - tensor = batch_norm(tensor) - tensor = tensor.view(*shape) - return output_dropout(tensor) + shape = x.shape + x = x.reshape(-1, shape[-1]) + x = batch_norm(x) + x = x.view(*shape) + return output_dropout(x) def _translational_interaction( @@ -855,18 +855,20 @@ def tucker_interaction( r: torch.FloatTensor, t: torch.FloatTensor, core_tensor: torch.FloatTensor, - do0: nn.Dropout, - do1: nn.Dropout, - do2: nn.Dropout, - bn1: Optional[nn.BatchNorm1d], - bn2: Optional[nn.BatchNorm1d], + do_h: nn.Dropout, + do_r: nn.Dropout, + do_hr: nn.Dropout, + bn_h: Optional[nn.BatchNorm1d], + bn_hr: Optional[nn.BatchNorm1d], ) -> torch.FloatTensor: - """ + r""" Evaluate the TuckEr interaction function. Compute scoring function W x_1 h x_2 r x_3 t as in the official implementation, i.e. as - DO(BN(DO(BN(h)) x_1 DO(W x_2 r))) x_3 t + .. math :: + + DO_{hr}(BN_{hr}(DO_h(BN_h(h)) x_1 DO_r(W x_2 r))) x_3 t where BN denotes BatchNorm and DO denotes Dropout @@ -878,32 +880,46 @@ def tucker_interaction( The tail representations. :param core_tensor: shape: (d_e, d_r, d_e) The core tensor. - :param do1: + :param do_h: The dropout layer for the head representations. - :param do0: + :param do_r: The first hidden dropout. - :param do2: + :param do_hr: The second hidden dropout. - :param bn1: + :param bn_h: The first batch normalization layer. - :param bn2: + :param bn_hr: The second batch normalization layer. :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # Compute wr = DO(W x_2 r) - x = do0(extended_einsum("idj,brd->brij", core_tensor, r)) - - # Compute h_n = DO(BN(h)) - h = _apply_optional_bn_to_tensor(batch_norm=bn1, output_dropout=do1, tensor=h) - - # compute whr = DO(BN(h_n x_1 wr)) - x = extended_einsum("brid,bhd->bhri", x, h) - x = _apply_optional_bn_to_tensor(batch_norm=bn2, tensor=x, output_dropout=do2) - - # Compute whr x_3 t - return extended_einsum("bhrd,btd->bhrt", x, t) + return extended_einsum( + # x_3 contraction + "bhrd,btd->bhrt", + _apply_optional_bn_to_tensor( + x=extended_einsum( + # x_1 contraction + "brid,bhd->bhri", + _apply_optional_bn_to_tensor( + x=extended_einsum( + # x_2 contraction + "idj,brd->brij", + core_tensor, + r, + ), + output_dropout=do_r, + ), + _apply_optional_bn_to_tensor( + x=h, + batch_norm=bn_h, + output_dropout=do_h, + )), + batch_norm=bn_hr, + output_dropout=do_hr, + ), + t, + ) def unstructured_model_interaction( diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 911ec68f73..2c523562d8 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -876,19 +876,25 @@ def __init__( self, embedding_dim: int = 200, relation_dim: Optional[int] = None, - dropout_0: float = 0.3, - dropout_1: float = 0.4, - dropout_2: float = 0.5, + head_dropout: float = 0.3, + relation_dropout: float = 0.4, + head_relation_dropout: float = 0.5, apply_batch_normalization: bool = True, ): """Initialize the Tucker interaction function. :param embedding_dim: + The entity embedding dimension. :param relation_dim: - :param dropout_0: - :param dropout_1: - :param dropout_2: + The relation embedding dimension. + :param head_dropout: + The dropout rate applied to the head representations. + :param relation_dropout: + The dropout rate applied to the relation representations. + :param head_relation_dropout: + The dropout rate applied to the combined head and relation representations. :param apply_batch_normalization: + Whether to use batch normalization on head representations and the combination of head and relation. """ super().__init__() @@ -903,28 +909,29 @@ def __init__( ) # Dropout - self.input_dropout = nn.Dropout(dropout_0) - self.hidden_dropout_1 = nn.Dropout(dropout_1) - self.hidden_dropout_2 = nn.Dropout(dropout_2) + self.head_dropout = nn.Dropout(head_dropout) + self.relation_dropout = nn.Dropout(relation_dropout) + self.head_relation_dropout = nn.Dropout(head_relation_dropout) if apply_batch_normalization: - self.bn1 = nn.BatchNorm1d(embedding_dim) - self.bn2 = nn.BatchNorm1d(embedding_dim) + self.head_batch_norm = nn.BatchNorm1d(embedding_dim) + self.head_relation_batchnorm = nn.BatchNorm1d(embedding_dim) else: - self.bn1 = self.bn2 = None + self.head_batch_norm = self.head_relation_batchnorm = None def reset_parameters(self): # noqa:D102 # Initialize core tensor, cf. https://github.com/ibalazevic/TuckER/blob/master/model.py#L12 nn.init.uniform_(self.core_tensor, -1., 1.) + # batch norm gets reset automatically, since it defines reset_parameters def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: return dict( core_tensor=self.core_tensor, - do0=self.input_dropout, - do1=self.hidden_dropout_1, - do2=self.hidden_dropout_2, - bn1=self.bn1, - bn2=self.bn2, + do_h=self.head_dropout, + do_r=self.relation_dropout, + do_hr=self.head_relation_dropout, + bn_h=self.head_batch_norm, + bn_hr=self.head_relation_batchnorm, ) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 66f247d43f..737854f91d 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -453,15 +453,14 @@ class TuckerTests(InteractionTests, unittest.TestCase): embedding_dim=InteractionTests.dim, ) - def _exp_score(self, bn1, bn2, core_tensor, do0, do1, do2, h, r, t) -> torch.FloatTensor: - # DO(BN(DO(BN(h)) x_1 DO(W x_2 r))) x_3 t - # = A x_3 t - # with A = DO(BN(B)), B = C x_1 D = , C=DO(BN(h)), D = DO(W x_2 r) + def _exp_score(self, bn_h, bn_hr, core_tensor, do_h, do_r, do_hr, h, r, t) -> torch.FloatTensor: + # DO_{hr}(BN_{hr}(DO_h(BN_h(h)) x_1 DO_r(W x_2 r))) x_3 t h, r, t = _strip_dim(h, r, t) - c = do0((core_tensor * r[None, :, None]).sum(dim=1)) # shape: (embedding_dim, embedding_dim) - b = do1(bn1(h.view(1, -1))).view(-1) # shape: (embedding_dim) - aa = (b[:, None] * c).sum(dim=1) # shape: (embedding_dim) - return do2(bn2((aa.view(1, -1)))) + a = do_r((core_tensor * r[None, :, None]).sum(dim=1)) # shape: (embedding_dim, embedding_dim) + b = do_h(bn_h(h.view(1, -1))).view(-1) # shape: (embedding_dim) + c = (b[:, None] * a).sum(dim=1) # shape: (embedding_dim) + d = do_hr(bn_hr((c.view(1, -1)))).view(-1) # shape: (embedding_dim) + return (d * t).sum() class RotatETests(InteractionTests, unittest.TestCase): From 800fe9626a33bf3f2c08f62e18d9de3455861dc5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 16:55:35 +0100 Subject: [PATCH 445/690] Try to simplify ConvKB test code --- src/pykeen/nn/functional.py | 12 ++++++------ tests/test_interactions.py | 9 ++++++--- 2 files changed, 12 insertions(+), 9 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 419cfa7740..ca186a6c46 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -261,18 +261,18 @@ def convkb_interaction( conv_head, conv_rel, conv_tail = conv.weight[:, 0, 0, :].t() conv_bias = conv.bias.view(1, 1, 1, 1, 1, num_filters) # h.shape: (b, nh, d), conv_head.shape: (o), out.shape: (b, nh, d, o) - h = (h.view(h.shape[0], h.shape[1], 1, 1, embedding_dim, 1) * conv_head.view(1, 1, 1, 1, 1, num_filters)) - r = (r.view(r.shape[0], 1, r.shape[1], 1, embedding_dim, 1) * conv_rel.view(1, 1, 1, 1, 1, num_filters)) - t = (t.view(t.shape[0], 1, 1, t.shape[1], embedding_dim, 1) * conv_tail.view(1, 1, 1, 1, 1, num_filters)) + h = (h.view(h.shape[0], num_heads, 1, 1, embedding_dim, 1) * conv_head.view(1, 1, 1, 1, 1, num_filters)) + r = (r.view(r.shape[0], 1, num_relations, 1, embedding_dim, 1) * conv_rel.view(1, 1, 1, 1, 1, num_filters)) + t = (t.view(t.shape[0], 1, 1, num_tails, embedding_dim, 1) * conv_tail.view(1, 1, 1, 1, 1, num_filters)) x = activation(tensor_sum(conv_bias, h, r, t)) # Apply dropout, cf. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L54-L56 x = hidden_dropout(x) # Linear layer for final scores - return linear( - x.view(-1, embedding_dim * num_filters), - ).view(-1, num_heads, num_relations, num_tails) + x = x.view(*x.shape[:-2], embedding_dim * num_filters) + x = linear(x) + return x.view(-1, num_heads, num_relations, num_tails) def distmult_interaction( diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 737854f91d..d68a6938c1 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -256,10 +256,10 @@ def test_scores(self): kwargs = self.instance._prepare_for_functional(h=h, r=r, t=t) # calculate by functional - scores_f = self.cls.func(**kwargs) + scores_f = self.cls.func(**kwargs).view(-1) # calculate manually - scores_f_manual = self._exp_score(**kwargs) + scores_f_manual = self._exp_score(**kwargs).view(-1) assert torch.allclose(scores_f_manual, scores_f) @abstractmethod @@ -323,8 +323,11 @@ class ConvKBTests(InteractionTests, unittest.TestCase): def _exp_score(self, h, r, t, conv, activation, hidden_dropout, linear) -> torch.FloatTensor: # noqa: D102 # W_L drop(act(W_C \ast ([h; r; t]) + b_C)) + b_L + # prepare conv input (N, C, H, W) x = torch.stack([x.view(-1) for x in (h, r, t)], dim=1).view(1, 1, -1, 3) - return linear(hidden_dropout(activation(conv(x).view(1, -1)))) + x = conv(x) + x = hidden_dropout(activation(x)) + return linear(x.view(1, -1)) class DistMultTests(InteractionTests, unittest.TestCase): From e1761cba8436c8c8cf958bfb9d47e3e5c4cfaebb Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 16:59:21 +0100 Subject: [PATCH 446/690] Fix SE --- src/pykeen/nn/functional.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index ca186a6c46..6b43ca4a9c 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -675,8 +675,8 @@ def structured_embedding_interaction( The scores. """ return negative_norm_of_sum( - extended_einsum("brde,bhd->bhre", r_h, h).unsqueeze(dim=3), - -extended_einsum("brde,btd->brte", r_t, t).unsqueeze(dim=1), + extended_einsum("bred,bhd->bhre", r_h, h).unsqueeze(dim=3), + -extended_einsum("bred,btd->brte", r_t, t).unsqueeze(dim=1), p=p, power_norm=power_norm, ) From 801b214823a6ce42d798ce15932cab0ec25bb77b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 17:03:11 +0100 Subject: [PATCH 447/690] Fix NTN --- src/pykeen/nn/functional.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 6b43ca4a9c..2275f5c350 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -490,6 +490,7 @@ def ntn_interaction( extended_einsum("bhd,brkde,bte->bhrtk", h, w, t), extended_einsum("brkd,bhd->bhrk", vh, h).unsqueeze(dim=3), extended_einsum("brkd,btd->brtk", vt, t).unsqueeze(dim=1), + b.view(b.shape[0], 1, num_relations, 1, num_slices) )) x = extended_einsum("bhrtk,brk->bhrt", x, u) return x From fcf7d22e8ca8e66be60eb090362e12c5d6d37e3d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 17:21:34 +0100 Subject: [PATCH 448/690] Add tests for vectorized KL divergence --- tests/test_nn.py | 54 ++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/tests/test_nn.py b/tests/test_nn.py index 5cba448604..9d3faca586 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -7,6 +7,8 @@ import torch from pykeen.nn import Embedding +from pykeen.nn.sim import kullback_leibler_similarity +from pykeen.typing import GaussianDistribution class EmbeddingsInCanonicalShapeTests(unittest.TestCase): @@ -65,3 +67,55 @@ def test_with_indices_with_duplicates(self): generator=self.generator, ) self._test_with_indices(indices=indices) + + +class KullbackLeiblerTests(unittest.TestCase): + """Tests for the vectorized computation of KL divergences.""" + + d: int = 3 + + def setUp(self) -> None: + self.e_mean = torch.rand(self.d) + self.e_var = torch.rand(self.d) + self.r_mean = torch.rand(self.d) + self.r_var = torch.rand(self.d) + + def get_e(self, pre_shape=(1, 1, 1)): + return GaussianDistribution( + mean=self.e_mean.view(*pre_shape, self.d), + diagonal_covariance=self.e_var.view(*pre_shape, self.d), + ) + + def get_r(self, pre_shape=(1, 1)): + return GaussianDistribution( + mean=self.r_mean.view(*pre_shape, self.d), + diagonal_covariance=self.r_var.view(*pre_shape, self.d), + ) + + def test_against_torch_builtin(self): + """Compare value against torch.distributions.""" + # r: (batch_size, num_heads, num_tails, d) + e = self.get_e() + # r: (batch_size, num_relations, d) + r = self.get_r() + sim = kullback_leibler_similarity(e=e, r=r, exact=True).view(-1) + + p = torch.distributions.MultivariateNormal(loc=self.e_mean, covariance_matrix=torch.diag(self.e_var)) + q = torch.distributions.MultivariateNormal(loc=self.r_mean, covariance_matrix=torch.diag(self.r_var)) + sim2 = torch.distributions.kl_divergence(p=p, q=q).view(-1) + assert torch.allclose(sim, sim2) + + def test_self_similarity(self): + """Check value of similarity to self.""" + # e: (batch_size, num_heads, num_tails, d) + e = self.get_e() + r = self.get_e(pre_shape=(1, 1)) + sim = kullback_leibler_similarity(e=e, r=r, exact=True) + assert torch.allclose(sim, torch.zeros_like(sim)) + + def test_value_range(self): + """Check the value range.""" + e = self.get_e() + r = self.get_r() + sim = kullback_leibler_similarity(e=e, r=r, exact=True) + assert (sim <= 0).all() From dd14f233751c6ff4f179fd0ad34fea2ca11f0fbd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 18:23:57 +0100 Subject: [PATCH 449/690] Fix KL divergence --- src/pykeen/nn/sim.py | 85 +++++++++++++++++++++++++++----------------- tests/test_nn.py | 10 ++++-- 2 files changed, 60 insertions(+), 35 deletions(-) diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index 2cfedb2314..fb5e946f74 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -71,21 +71,24 @@ def kullback_leibler_similarity( epsilon: float = 1.0e-10, exact: bool = True, ) -> torch.FloatTensor: - r"""Compute the similarity based on KL divergence. + r"""Compute the negative KL divergence. This is done between two Gaussian distributions given by mean `mu_*` and diagonal covariance matrix `sigma_*`. - + .. math:: + + D((\mu_0, \Sigma_0), (\mu_1, \Sigma_1)) = 0.5 * ( + tr(\Sigma_1^-1 \Sigma_0) + + (\mu_1 - \mu_0) * \Sigma_1^-1 (\mu_1 - \mu_0) + - k + + ln (det(\Sigma_1) / det(\Sigma_0)) + ) - D((\mu_e, \Sigma_e), (\mu_r, \Sigma_r))) - = \frac{1}{2} \left( - tr(\Sigma_r^{-1}\Sigma_e) - + (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) - - \log \frac{det(\Sigma_e)}{det(\Sigma_r)} - k_e - \right) + .. note :: + This methods assumes diagonal covariance matrices :math:`\Sigma`. - Note: The sign of the function has been flipped as opposed to the description in the paper, as the - Kullback Leibler divergence is large if the distributions are dissimilar. + .. seealso :: + https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Kullback%E2%80%93Leibler_divergence :param e: shape: (batch_size, num_heads, num_tails, d) The entity Gaussian distributions, as mean/diagonal covariance pairs. @@ -99,33 +102,51 @@ def kullback_leibler_similarity( :return: torch.Tensor, shape: (s_1, ..., s_k) The similarity. """ - # invert covariance, shape: (batch_size, num_relations, d) - safe_sigma_r = torch.clamp_min(r.diagonal_covariance, min=epsilon) - sigma_r_inv = torch.reciprocal(safe_sigma_r) - #: a = tr(\Sigma_r^{-1}\Sigma_e), (batch_size, num_heads, num_relations, num_tails) - # [(b, h, t, d), (b, r, d) -> (b, 1, r, d) -> (b, 1, d, r)] -> (b, h, t, r) -> (b, h, r, t) - sim = (e.diagonal_covariance @ sigma_r_inv.unsqueeze(dim=1).transpose(-2, -1)).transpose(-2, -1) + assert (e.diagonal_covariance > 0).all() and (r.diagonal_covariance > 0).all() + + # broadcast shapes to (batch_size, num_heads, num_relations, num_tails, dim) + e_shape = e.mean.shape # (batch_size, num_heads, num_tails, dim) + e_mean = e.mean.view(e_shape[0], e_shape[1], 1, e_shape[2], e_shape[3]) + e_var = e.diagonal_covariance.view(e_shape[0], e_shape[1], 1, e_shape[2], e_shape[3]) + + r_shape = r.mean.shape # (batch_size, num_relations, dim) + r_mean = r.mean.view(r_shape[0], 1, r_shape[1], 1, r_shape[2]) + r_var: torch.FloatTensor = r.diagonal_covariance.view(r_shape[0], 1, r_shape[1], 1, r_shape[2]) + + # 1. Component + # tr(sigma_1^-1 sigma_0) = sum (sigma_1^-1 sigma_0)[i, i] + # since sigma_0, sigma_1 are diagonal matrices: + # = sum (sigma_1^-1[i] sigma_0[i]) = sum (sigma_0[i] / sigma_1[i]) + r_var_safe = r_var.clamp_min(min=epsilon) + result = (e_var / r_var_safe).sum(dim=-1) + + # 2. Component + # (mu_1 - mu_0) * Sigma_1^-1 (mu_1 - mu_0) + # with mu = (mu_1 - mu_0) + # = mu * Sigma_1^-1 mu + # since Sigma_1 is diagonal + # = mu**2 / sigma_1 + mu = r_mean - e_mean + result = result + (mu.pow(2) / r_var_safe).sum(dim=-1) + + # 3. Component + if exact: + result = result - e_shape[-1] - #: b = (\mu_r - \mu_e)^T\Sigma_r^{-1}(\mu_r - \mu_e) - r_shape = r.mean.shape - # mu.shape: (b, h, r, t, d) - mu = r.mean.view(r_shape[0], 1, r_shape[1], 1, r_shape[2]) - e.mean.unsqueeze(dim=2) - sim = sim + (mu ** 2 @ sigma_r_inv.view(r_shape[0], 1, r_shape[1], r_shape[2], 1)).squeeze(dim=-1) - - #: c = \log \frac{det(\Sigma_e)}{det(\Sigma_r)} - # = sum log (sigma_e)_i - sum log (sigma_r)_i - # ce.shape: (b, h, t) - ce = e.diagonal_covariance.clamp_min(min=epsilon).log().sum(dim=-1) - # cr.shape: (b, r) - cr = safe_sigma_r.log().sum(dim=-1) - sim = sim + ce.unsqueeze(dim=2) - cr.view(r_shape[0], 1, r_shape[1], 1) + # 4. Component + # ln (det(\Sigma_1) / det(\Sigma_0)) + # = ln det Sigma_1 - ln det Sigma_0 + # since Sigma is diagonal, we have det Sigma = prod Sigma[ii] + # = ln prod Sigma_1[ii] - ln prod Sigma_0[ii] + # = sum ln Sigma_1[ii] - sum ln Sigma_0[ii] + e_var_safe = e_var.clamp_min(min=epsilon) + result = result + r_var_safe.log().sum(dim=-1) - e_var_safe.log().sum(dim=-1) if exact: - sim = sim - e.mean.shape[-1] - sim = 0.5 * sim + result = 0.5 * result - return sim + return -result KG2E_SIMILARITIES = { diff --git a/tests/test_nn.py b/tests/test_nn.py index 9d3faca586..310f3f95f9 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -76,9 +76,9 @@ class KullbackLeiblerTests(unittest.TestCase): def setUp(self) -> None: self.e_mean = torch.rand(self.d) - self.e_var = torch.rand(self.d) + self.e_var = torch.rand(self.d).exp() self.r_mean = torch.rand(self.d) - self.r_var = torch.rand(self.d) + self.r_var = torch.rand(self.d).exp() def get_e(self, pre_shape=(1, 1, 1)): return GaussianDistribution( @@ -102,12 +102,14 @@ def test_against_torch_builtin(self): p = torch.distributions.MultivariateNormal(loc=self.e_mean, covariance_matrix=torch.diag(self.e_var)) q = torch.distributions.MultivariateNormal(loc=self.r_mean, covariance_matrix=torch.diag(self.r_var)) - sim2 = torch.distributions.kl_divergence(p=p, q=q).view(-1) + sim2 = -torch.distributions.kl_divergence(p=p, q=q).view(-1) assert torch.allclose(sim, sim2) def test_self_similarity(self): """Check value of similarity to self.""" # e: (batch_size, num_heads, num_tails, d) + # https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence#Properties + # divergence = 0 => similarity = -divergence = 0 e = self.get_e() r = self.get_e(pre_shape=(1, 1)) sim = kullback_leibler_similarity(e=e, r=r, exact=True) @@ -115,6 +117,8 @@ def test_self_similarity(self): def test_value_range(self): """Check the value range.""" + # https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence#Properties + # divergence >= 0 => similarity = -divergence <= 0 e = self.get_e() r = self.get_r() sim = kullback_leibler_similarity(e=e, r=r, exact=True) From 728fd8bfdd8d9f2ff6576f6fddca4b9d032600fd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 18:27:18 +0100 Subject: [PATCH 450/690] Use tensor_sum --- src/pykeen/nn/sim.py | 22 ++++++++++++++++++---- 1 file changed, 18 insertions(+), 4 deletions(-) diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index fb5e946f74..ce4ab00ea8 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -14,6 +14,8 @@ 'KG2E_SIMILARITIES', ] +from ..utils import tensor_sum + def expected_likelihood( e: GaussianDistribution, @@ -114,12 +116,16 @@ def kullback_leibler_similarity( r_mean = r.mean.view(r_shape[0], 1, r_shape[1], 1, r_shape[2]) r_var: torch.FloatTensor = r.diagonal_covariance.view(r_shape[0], 1, r_shape[1], 1, r_shape[2]) + terms = [] + # 1. Component # tr(sigma_1^-1 sigma_0) = sum (sigma_1^-1 sigma_0)[i, i] # since sigma_0, sigma_1 are diagonal matrices: # = sum (sigma_1^-1[i] sigma_0[i]) = sum (sigma_0[i] / sigma_1[i]) r_var_safe = r_var.clamp_min(min=epsilon) - result = (e_var / r_var_safe).sum(dim=-1) + terms.append( + (e_var / r_var_safe).sum(dim=-1) + ) # 2. Component # (mu_1 - mu_0) * Sigma_1^-1 (mu_1 - mu_0) @@ -128,11 +134,15 @@ def kullback_leibler_similarity( # since Sigma_1 is diagonal # = mu**2 / sigma_1 mu = r_mean - e_mean - result = result + (mu.pow(2) / r_var_safe).sum(dim=-1) + terms.append( + (mu.pow(2) / r_var_safe).sum(dim=-1) + ) # 3. Component if exact: - result = result - e_shape[-1] + terms.append( + -e_shape[-1] + ) # 4. Component # ln (det(\Sigma_1) / det(\Sigma_0)) @@ -141,8 +151,12 @@ def kullback_leibler_similarity( # = ln prod Sigma_1[ii] - ln prod Sigma_0[ii] # = sum ln Sigma_1[ii] - sum ln Sigma_0[ii] e_var_safe = e_var.clamp_min(min=epsilon) - result = result + r_var_safe.log().sum(dim=-1) - e_var_safe.log().sum(dim=-1) + terms.extend(( + r_var_safe.log().sum(dim=-1), + -e_var_safe.log().sum(dim=-1) + )) + result = tensor_sum(*terms) if exact: result = 0.5 * result From f178bb10305e3a1ef6fab502d7a4327a6af83129 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 18:28:27 +0100 Subject: [PATCH 451/690] Fix KG2E test --- tests/test_interactions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index d68a6938c1..0342d8ecc5 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -445,7 +445,7 @@ def _exp_score(self, exact, h_mean, h_var, r_mean, r_var, similarity, t_mean, t_ e_mean, e_var = h_mean - t_mean, h_var + t_var p = torch.distributions.MultivariateNormal(loc=e_mean, covariance_matrix=torch.diag(e_var)) q = torch.distributions.MultivariateNormal(loc=r_mean, covariance_matrix=torch.diag(r_var)) - return torch.distributions.kl.kl_divergence(p, q) + return -torch.distributions.kl.kl_divergence(p, q) class TuckerTests(InteractionTests, unittest.TestCase): From d7780448dffb7871111749e2e53ba8cf28172ef8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 18:47:27 +0100 Subject: [PATCH 452/690] Fix ConvKB functional form dimension ordering --- src/pykeen/nn/functional.py | 35 +++++++++++++++++++++++++---------- 1 file changed, 25 insertions(+), 10 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 2275f5c350..f343ad0195 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -258,21 +258,36 @@ def convkb_interaction( assert conv.weight.shape == (num_filters, 1, 1, 3) # compute conv(stack(h, r, t)) - conv_head, conv_rel, conv_tail = conv.weight[:, 0, 0, :].t() - conv_bias = conv.bias.view(1, 1, 1, 1, 1, num_filters) - # h.shape: (b, nh, d), conv_head.shape: (o), out.shape: (b, nh, d, o) - h = (h.view(h.shape[0], num_heads, 1, 1, embedding_dim, 1) * conv_head.view(1, 1, 1, 1, 1, num_filters)) - r = (r.view(r.shape[0], 1, num_relations, 1, embedding_dim, 1) * conv_rel.view(1, 1, 1, 1, 1, num_filters)) - t = (t.view(t.shape[0], 1, 1, num_tails, embedding_dim, 1) * conv_tail.view(1, 1, 1, 1, 1, num_filters)) - x = activation(tensor_sum(conv_bias, h, r, t)) + # prepare input shapes for broadcasting + h = h.view(h.shape[0], num_heads, 1, 1, 1, embedding_dim) + r = r.view(r.shape[0], 1, num_relations, 1, 1, embedding_dim) + t = t.view(t.shape[0], 1, 1, num_tails, 1, embedding_dim) + + # conv.weight.shape = (C_out, C_in, kernel_size[0], kernel_size[1]) + # here, kernel_size = (1, 3), C_in = 1, C_out = num_filters + # -> conv_head, conv_rel, conv_tail shapes: (num_filters,) + # conv_head, conv_rel, conv_tail = conv.weight[:, 0, 0, :].t() + # conv_bias = conv.bias.view(1, 1, 1, 1, num_filters, 1) + conv_head, conv_rel, conv_tail, conv_bias = [ + c.view(1, 1, 1, 1, num_filters, 1) + for c in list(conv.weight[:, 0, 0, :].t()) + [conv.bias] + ] + + # convolve -> output.shape: (*, embedding_dim, num_filters) + h = conv_head @ h + r = conv_rel @ r + t = conv_tail @ t + + x = tensor_sum(conv_bias, h, r, t) + x = activation(x) # Apply dropout, cf. https://github.com/daiquocnguyen/ConvKB/blob/master/model.py#L54-L56 x = hidden_dropout(x) - # Linear layer for final scores - x = x.view(*x.shape[:-2], embedding_dim * num_filters) + # Linear layer for final scores; use flattened representations, shape: (b, h, r, t, d * f) + x = x.view(x.shape[0], num_heads, num_relations, num_tails, embedding_dim * num_filters) x = linear(x) - return x.view(-1, num_heads, num_relations, num_tails) + return x.squeeze(dim=-1) def distmult_interaction( From 158224d0661c2f81740c693b3e3f9b6f196daf30 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 19:00:11 +0100 Subject: [PATCH 453/690] code style - part 1 --- src/pykeen/nn/sim.py | 21 +++++++-------------- 1 file changed, 7 insertions(+), 14 deletions(-) diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index ce4ab00ea8..14f1dbe05b 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -7,6 +7,7 @@ import torch from ..typing import GaussianDistribution +from ..utils import tensor_sum __all__ = [ 'expected_likelihood', @@ -14,8 +15,6 @@ 'KG2E_SIMILARITIES', ] -from ..utils import tensor_sum - def expected_likelihood( e: GaussianDistribution, @@ -76,9 +75,9 @@ def kullback_leibler_similarity( r"""Compute the negative KL divergence. This is done between two Gaussian distributions given by mean `mu_*` and diagonal covariance matrix `sigma_*`. - + .. math:: - + D((\mu_0, \Sigma_0), (\mu_1, \Sigma_1)) = 0.5 * ( tr(\Sigma_1^-1 \Sigma_0) + (\mu_1 - \mu_0) * \Sigma_1^-1 (\mu_1 - \mu_0) @@ -123,9 +122,7 @@ def kullback_leibler_similarity( # since sigma_0, sigma_1 are diagonal matrices: # = sum (sigma_1^-1[i] sigma_0[i]) = sum (sigma_0[i] / sigma_1[i]) r_var_safe = r_var.clamp_min(min=epsilon) - terms.append( - (e_var / r_var_safe).sum(dim=-1) - ) + terms.append((e_var / r_var_safe).sum(dim=-1)) # 2. Component # (mu_1 - mu_0) * Sigma_1^-1 (mu_1 - mu_0) @@ -134,15 +131,11 @@ def kullback_leibler_similarity( # since Sigma_1 is diagonal # = mu**2 / sigma_1 mu = r_mean - e_mean - terms.append( - (mu.pow(2) / r_var_safe).sum(dim=-1) - ) + terms.append((mu.pow(2) / r_var_safe).sum(dim=-1)) # 3. Component if exact: - terms.append( - -e_shape[-1] - ) + terms.append(-e_shape[-1]) # 4. Component # ln (det(\Sigma_1) / det(\Sigma_0)) @@ -153,7 +146,7 @@ def kullback_leibler_similarity( e_var_safe = e_var.clamp_min(min=epsilon) terms.extend(( r_var_safe.log().sum(dim=-1), - -e_var_safe.log().sum(dim=-1) + -e_var_safe.log().sum(dim=-1), )) result = tensor_sum(*terms) From c460b880bf8951459b1f09921829cf9155ad8a68 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 19:00:40 +0100 Subject: [PATCH 454/690] code style - part 2 --- src/pykeen/nn/functional.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index f343ad0195..840a32a9e5 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -45,7 +45,7 @@ def _extract_sizes( r: torch.Tensor, t: torch.Tensor, ) -> Tuple[int, int, int, int, int]: - """Utility to extract size dimensions from head/relation/tail representations.""" + """Extract size dimensions from head/relation/tail representations.""" num_heads, num_relations, num_tails = [xx.shape[1] for xx in (h, r, t)] d_e = h.shape[-1] d_r = r.shape[-1] @@ -138,11 +138,11 @@ def complex_interaction( return tensor_sum(*( factor * extended_einsum("bhd,brd,btd->bhrt", hh, rr, tt) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) @@ -505,7 +505,7 @@ def ntn_interaction( extended_einsum("bhd,brkde,bte->bhrtk", h, w, t), extended_einsum("brkd,bhd->bhrk", vh, h).unsqueeze(dim=3), extended_einsum("brkd,btd->brtk", vt, t).unsqueeze(dim=1), - b.view(b.shape[0], 1, num_relations, 1, num_slices) + b.view(b.shape[0], 1, num_relations, 1, num_slices), )) x = extended_einsum("bhrtk,brk->bhrt", x, u) return x From b79c6d01b03678937b59bf6891ed8540c15cf97f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 19:01:17 +0100 Subject: [PATCH 455/690] Fix import order --- src/pykeen/models/__init__.py | 3 +-- src/pykeen/models/multimodal/__init__.py | 2 ++ 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index e2b6f0bffd..5d2e2c23b7 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -12,8 +12,7 @@ DoubleRelationEmbeddingModel, ERModel, Model, SingleVectorEmbeddingModel, TwoSideEmbeddingModel, TwoVectorEmbeddingModel, ) -from .multimodal.base import LiteralModel -from .multimodal import ComplExLiteral, DistMultLiteral +from .multimodal import ComplExLiteral, DistMultLiteral, LiteralModel from .unimodal import ( ComplEx, ConvE, diff --git a/src/pykeen/models/multimodal/__init__.py b/src/pykeen/models/multimodal/__init__.py index c81a03b74b..e0122ad108 100644 --- a/src/pykeen/models/multimodal/__init__.py +++ b/src/pykeen/models/multimodal/__init__.py @@ -6,10 +6,12 @@ `_ arXiv preprint arXiv:1802.00934. """ +from .base import LiteralModel from .complex_literal import ComplExLiteral from .distmult_literal import DistMultLiteral __all__ = [ 'ComplExLiteral', 'DistMultLiteral', + 'LiteralModel', ] From 2282e33cc3d34044897b1ad3302299c4be342651 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 19:03:02 +0100 Subject: [PATCH 456/690] Code style --- tests/test_interactions.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 0342d8ecc5..8118e6273e 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -231,7 +231,10 @@ def test_forward(self): except ValueError as error: # check whether the error originates from batch norm for single element batches small_batch_size = any(s[0] == 1 for s in (hs, rs, ts)) - has_batch_norm = any(isinstance(m, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d)) for m in self.instance.modules()) + has_batch_norm = any( + isinstance(m, (torch.nn.BatchNorm1d, torch.nn.BatchNorm2d)) + for m in self.instance.modules() + ) if small_batch_size and has_batch_norm: logger.warning(f"Skipping test for shapes {hs}, {rs}, {ts}") continue @@ -401,10 +404,10 @@ def _exp_score(self, h, t, w, vt, vh, b, u, activation) -> torch.FloatTensor: score = 0. for i in range(u.shape[-1]): score = score + u[i] * activation( - h.view(1, self.dim) @ w[i] @ t.view(self.dim, 1) + - (vh[i] * h.view(-1)).sum() + - (vt[i] * t.view(-1)).sum() + - b[i] + h.view(1, self.dim) @ w[i] @ t.view(self.dim, 1) + + (vh[i] * h.view(-1)).sum() + + (vt[i] * t.view(-1)).sum() + + b[i] ) return score From 791e996c6303c9f33835073a11bbbcff4130df03 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 19:06:04 +0100 Subject: [PATCH 457/690] more code style --- src/pykeen/nn/sim.py | 1 - tests/test_interactions.py | 14 ++++++-------- tests/test_nn.py | 18 +++++++++--------- tests/test_utils.py | 8 ++++---- 4 files changed, 19 insertions(+), 22 deletions(-) diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index 14f1dbe05b..07f0978d82 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -103,7 +103,6 @@ def kullback_leibler_similarity( :return: torch.Tensor, shape: (s_1, ..., s_k) The similarity. """ - assert (e.diagonal_covariance > 0).all() and (r.diagonal_covariance > 0).all() # broadcast shapes to (batch_size, num_heads, num_relations, num_tails, dim) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 8118e6273e..3ef19aa02d 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -303,11 +303,11 @@ def _get_hrt( return h, r, (t, t_bias) def _exp_score( - self, embedding_height, embedding_width, h, hr1d, hr2d, input_channels, r, t, t_bias + self, embedding_height, embedding_width, h, hr1d, hr2d, input_channels, r, t, t_bias, ) -> torch.FloatTensor: x = torch.cat([ h.view(1, input_channels, embedding_height, embedding_width), - r.view(1, input_channels, embedding_height, embedding_width) + r.view(1, input_channels, embedding_height, embedding_width), ], dim=2) x = hr2d(x) x = x.view(-1, numpy.prod(x.shape[-3:])) @@ -403,12 +403,10 @@ def _exp_score(self, h, t, w, vt, vh, b, u, activation) -> torch.FloatTensor: h, t, w, vt, vh, b, u = _strip_dim(h, t, w, vt, vh, b, u) score = 0. for i in range(u.shape[-1]): - score = score + u[i] * activation( - h.view(1, self.dim) @ w[i] @ t.view(self.dim, 1) - + (vh[i] * h.view(-1)).sum() - + (vt[i] * t.view(-1)).sum() - + b[i] - ) + first_part = h.view(1, self.dim) @ w[i] @ t.view(self.dim, 1) + second_part = (vh[i] * h.view(-1)).sum() + third_part = (vt[i] * t.view(-1)).sum() + score = score + u[i] * activation(first_part + second_part + third_part + b[i]) return score diff --git a/tests/test_nn.py b/tests/test_nn.py index 310f3f95f9..3a7b8ba8fd 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -74,19 +74,19 @@ class KullbackLeiblerTests(unittest.TestCase): d: int = 3 - def setUp(self) -> None: + def setUp(self) -> None: # noqa: D102 self.e_mean = torch.rand(self.d) self.e_var = torch.rand(self.d).exp() self.r_mean = torch.rand(self.d) self.r_var = torch.rand(self.d).exp() - def get_e(self, pre_shape=(1, 1, 1)): + def _get_e(self, pre_shape=(1, 1, 1)): return GaussianDistribution( mean=self.e_mean.view(*pre_shape, self.d), diagonal_covariance=self.e_var.view(*pre_shape, self.d), ) - def get_r(self, pre_shape=(1, 1)): + def _get_r(self, pre_shape=(1, 1)): return GaussianDistribution( mean=self.r_mean.view(*pre_shape, self.d), diagonal_covariance=self.r_var.view(*pre_shape, self.d), @@ -95,9 +95,9 @@ def get_r(self, pre_shape=(1, 1)): def test_against_torch_builtin(self): """Compare value against torch.distributions.""" # r: (batch_size, num_heads, num_tails, d) - e = self.get_e() + e = self._get_e() # r: (batch_size, num_relations, d) - r = self.get_r() + r = self._get_r() sim = kullback_leibler_similarity(e=e, r=r, exact=True).view(-1) p = torch.distributions.MultivariateNormal(loc=self.e_mean, covariance_matrix=torch.diag(self.e_var)) @@ -110,8 +110,8 @@ def test_self_similarity(self): # e: (batch_size, num_heads, num_tails, d) # https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence#Properties # divergence = 0 => similarity = -divergence = 0 - e = self.get_e() - r = self.get_e(pre_shape=(1, 1)) + e = self._get_e() + r = self._get_e(pre_shape=(1, 1)) sim = kullback_leibler_similarity(e=e, r=r, exact=True) assert torch.allclose(sim, torch.zeros_like(sim)) @@ -119,7 +119,7 @@ def test_value_range(self): """Check the value range.""" # https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence#Properties # divergence >= 0 => similarity = -divergence <= 0 - e = self.get_e() - r = self.get_r() + e = self._get_e() + r = self._get_r() sim = kullback_leibler_similarity(e=e, r=r, exact=True) assert (sim <= 0).all() diff --git a/tests/test_utils.py b/tests/test_utils.py index 0e5cfcd3f2..8d02a0472b 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -192,9 +192,9 @@ def test_project_entity(): # check equivalence of re-formulation # e_{\bot} = M_{re} e = (r_p e_p^T + I^{d_r \times d_e}) e # = r_p (e_p^T e) + e' - M_re = r_p.unsqueeze(dim=-1) @ e_p.unsqueeze(dim=-2) - M_re = M_re + torch.eye(relation_dim, embedding_dim).view(1, 1, relation_dim, embedding_dim) - assert M_re.shape == (batch_size, num_entities, relation_dim, embedding_dim) - e_vanilla = (M_re @ e.unsqueeze(dim=-1)).squeeze(dim=-1) + m_re = r_p.unsqueeze(dim=-1) @ e_p.unsqueeze(dim=-2) + m_re = m_re + torch.eye(relation_dim, embedding_dim).view(1, 1, relation_dim, embedding_dim) + assert m_re.shape == (batch_size, num_entities, relation_dim, embedding_dim) + e_vanilla = (m_re @ e.unsqueeze(dim=-1)).squeeze(dim=-1) e_vanilla = clamp_norm(e_vanilla, p=2, dim=-1, maxnorm=1) assert torch.allclose(e_vanilla, e_bot) From d5737061c8f650e009c02958c7229fffaf77c3dd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 18 Nov 2020 19:09:10 +0100 Subject: [PATCH 458/690] Fix mypy issue --- src/pykeen/models/multimodal/base.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/multimodal/base.py b/src/pykeen/models/multimodal/base.py index 1be83b9acf..8becb00986 100644 --- a/src/pykeen/models/multimodal/base.py +++ b/src/pykeen/models/multimodal/base.py @@ -1,15 +1,22 @@ """Base classes for multi-modal models.""" -from typing import Optional +from typing import Optional, TYPE_CHECKING import torch from torch import nn from pykeen.losses import Loss -from pykeen.models import ERModel +from pykeen.models.base import ERModel from pykeen.nn import Embedding, EmbeddingSpecification, Interaction, LiteralRepresentations from pykeen.triples import TriplesNumericLiteralsFactory from pykeen.typing import DeviceHint, HeadRepresentation, RelationRepresentation, TailRepresentation +__all__ = [ + "LiteralModel", +] + +if TYPE_CHECKING: + from ...typing import Representation # noqa + class LiteralModel(ERModel, autoreset=False): """Base class for models with entity literals.""" From 6866c143eb4327a86a42f4e27a9e3cf044185769 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 19 Nov 2020 13:23:24 +0100 Subject: [PATCH 459/690] Reorganize testing code --- src/pykeen/testing/mocks.py | 32 ++- src/pykeen/utils.py | 5 + tests/resources/hpo_complex_nations.json | 182 ------------------ tests/test_early_stopping.py | 46 +++-- tests/test_interactions.py | 28 ++- tests/test_model_mode.py | 17 -- tests/test_models.py | 133 +++++-------- tests/test_regularizers.py | 111 +++++------ .../test_utils.py => test_training_utils.py} | 0 tests/test_utils.py | 69 ++++--- 10 files changed, 216 insertions(+), 407 deletions(-) delete mode 100644 tests/resources/hpo_complex_nations.json rename tests/{training/test_utils.py => test_training_utils.py} (100%) diff --git a/src/pykeen/testing/mocks.py b/src/pykeen/testing/mocks.py index f476f1a2c4..ec10b697ea 100644 --- a/src/pykeen/testing/mocks.py +++ b/src/pykeen/testing/mocks.py @@ -2,15 +2,19 @@ """Mocks for testing PyKEEN.""" -from typing import Optional +from typing import Optional, Sequence +import numpy import torch +from torch import nn from pykeen.models.base import Model +from pykeen.nn import RepresentationModule from pykeen.triples import TriplesFactory __all__ = [ 'MockModel', + 'MockRepresentations', ] @@ -48,3 +52,29 @@ def forward( delta = ind scores = scores + delta.float().view(*shape) return scores + + +class MockRepresentations(RepresentationModule): + """A custom representation module with minimal implementation.""" + + def __init__(self, num_entities: int, shape: Sequence[int]): + super().__init__() + self.num_embeddings = num_entities + self.shape = shape + self.x = nn.Parameter(torch.rand(numpy.prod(shape))) + + def forward(self, indices: Optional[torch.LongTensor] = None) -> torch.FloatTensor: # noqa: D102 + n = self.num_embeddings if indices is None else indices.shape[0] + return self.x.unsqueeze(dim=0).repeat(n, 1).view(-1, *self.shape) + + def get_in_canonical_shape( + self, + indices: Optional[torch.LongTensor] = None, + reshape_dim: Optional[Sequence[int]] = None, + ) -> torch.FloatTensor: # noqa: D102 + x = self(indices=indices) + if indices is None: + x = x.unsqueeze(dim=0) + else: + x = x.unsqueeze(dim=1) + return x diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 6d19f30fce..5ee7000ce9 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -612,3 +612,8 @@ def pop_only(elements: Iterable[X]) -> X: if len(elements) > 1: raise ValueError(f'More than one element: {elements}') return elements[0] + + +def strip_dim(*x): + """Strip the last dimension.""" + return [xx.view(xx.shape[2:]) for xx in x] diff --git a/tests/resources/hpo_complex_nations.json b/tests/resources/hpo_complex_nations.json deleted file mode 100644 index 022bb22a65..0000000000 --- a/tests/resources/hpo_complex_nations.json +++ /dev/null @@ -1,182 +0,0 @@ -{ - "metadata": { - "title": "HPO Over Nations for ComplEx" - }, - "ablation": { - "datasets": [ - "nations" - ], - "models": [ - "ComplEx" - ], - "model_kwargs": { - "ComplEx": {} - }, - "model_kwargs_ranges": { - "ComplEx": { - "embedding_dim": { - "type": "int", - "low": 10, - "high": 30, - "q": 10 - } - } - }, - "training_loops": [ - "slcwa", - "lcwa" - ], - "optimizers": [ - "adam", - "adadelta" - ], - "optimizer_kwargs": { - "ComplEx": { - "adam": {}, - "adadelta": {} - } - }, - "optimizer_kwargs_ranges": { - "ComplEx": { - "adam": { - "lr": { - "type": "float", - "low": 0.001, - "high": 0.1, - "scale": "log" - } - }, - "adadelta": { - "lr": { - "type": "float", - "low": 0.001, - "high": 0.1, - "scale": "log" - } - } - } - }, - "loss_functions": [ - "MarginRankingLoss", - "BCEAfterSigmoidLoss" - ], - "loss_kwargs": { - "ComplEx": { - "MarginRankingLoss": {}, - "BCEAfterSigmoidLoss": {} - } - }, - "loss_kwargs_ranges": { - "ComplEx": { - "MarginRankingLoss": { - "margin": { - "type": "float", - "low": 0.5, - "high": 1.5, - "q": 0.1 - } - }, - "BCEAfterSigmoidLoss": {} - } - }, - "regularizers": [ - "NoRegularizer", - "PowerSumRegularizer" - ], - "regularizer_kwargs": { - "ComplEx": { - "NoRegularizer": {}, - "PowerSumRegularizer": { - "p": 2.0 - } - } - }, - "regularizer_kwargs_ranges": { - "ComplEx": { - "NoRegularizer": {}, - "PowerSumRegularizer": { - "weight": { - "type": "float", - "low": 0.01, - "high": 1.0, - "scale": "log" - } - } - } - }, - "negative_sampler": "BasicNegativeSampler", - "negative_sampler_kwargs": { - "ComplEx": { - "BasicNegativeSampler": {} - } - }, - "negative_sampler_kwargs_ranges": { - "ComplEx": { - "BasicNegativeSampler": { - "num_negs_per_pos": { - "type": "int", - "low": 1, - "high": 5, - "q": 1 - } - } - } - }, - "create_inverse_triples": [ - true, - false - ], - "evaluator": "RankBasedEvaluator", - "evaluator_kwargs": { - "filtered": true - }, - "evaluation_kwargs": { - "batch_size": 16 - }, - "training_kwargs": { - "ComplEx": { - "slcwa": { - "num_epochs": 10 - }, - "lcwa": { - "num_epochs": 10, - "label_smoothing": 0.0 - } - } - }, - "training_kwargs_ranges": { - "ComplEx": { - "slcwa": { - "batch_size": { - "type": "int", - "low": 128, - "high": 512, - "q": 100 - } - }, - "lcwa": { - "batch_size": { - "type": "int", - "low": 128, - "high": 200, - "q": 10 - } - } - } - }, - "stopper": "early", - "stopper_kwargs": { - "frequency": 2, - "patience": 2, - "delta": 0.002 - } - }, - "optuna": { - "n_trials": 2, - "timeout": 10, - "metric": "hits@10", - "direction": "maximize", - "sampler": "random", - "pruner": "nop" - } -} \ No newline at end of file diff --git a/tests/test_early_stopping.py b/tests/test_early_stopping.py index 36fd95650a..4cbfb1b67e 100644 --- a/tests/test_early_stopping.py +++ b/tests/test_early_stopping.py @@ -21,24 +21,34 @@ from pykeen.typing import MappedTriples -def test_is_improvement(): - """Test is_improvement().""" - for best_value, current_value, larger_is_better, relative_delta, is_better in [ - # equal value; larger is better - (1.0, 1.0, True, 0.0, False), - # equal value; smaller is better - (1.0, 1.0, False, 0.0, False), - # larger is better; improvement - (1.0, 1.1, True, 0.0, True), - # larger is better; improvement; but not significant - (1.0, 1.1, True, 0.1, False), - ]: - assert is_better == is_improvement( - best_value=best_value, - current_value=current_value, - larger_is_better=larger_is_better, - relative_delta=relative_delta, - ) +class TestRandom(unittest.TestCase): + """Random tests for early stopper.""" + + def test_is_improvement(self): + """Test is_improvement().""" + for best_value, current_value, larger_is_better, relative_delta, is_better in [ + # equal value; larger is better + (1.0, 1.0, True, 0.0, False), + # equal value; smaller is better + (1.0, 1.0, False, 0.0, False), + # larger is better; improvement + (1.0, 1.1, True, 0.0, True), + # larger is better; improvement; but not significant + (1.0, 1.1, True, 0.1, False), + ]: + with self.subTest( + best_value=best_value, + current_value=current_value, + larger_is_better=larger_is_better, + relative_delta=relative_delta, + is_better=is_better, + ): + self.assertEqual(is_better, is_improvement( + best_value=best_value, + current_value=current_value, + larger_is_better=larger_is_better, + relative_delta=relative_delta, + )) class MockEvaluator(Evaluator): diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 3ef19aa02d..2995276907 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -14,7 +14,7 @@ import pykeen.nn.modules from pykeen.nn.modules import Interaction, TranslationalInteraction from pykeen.typing import Representation -from pykeen.utils import clamp_norm, get_subclasses, project_entity, view_complex +from pykeen.utils import clamp_norm, get_subclasses, project_entity, strip_dim, view_complex T = TypeVar("T") logger = logging.getLogger(__name__) @@ -263,7 +263,7 @@ def test_scores(self): # calculate manually scores_f_manual = self._exp_score(**kwargs).view(-1) - assert torch.allclose(scores_f_manual, scores_f) + assert torch.allclose(scores_f_manual, scores_f), f'Diff: {scores_f_manual - scores_f}' @abstractmethod def _exp_score(self, **kwargs) -> torch.FloatTensor: @@ -382,10 +382,6 @@ def _exp_score(self, h, r, t) -> torch.FloatTensor: # noqa: D102 return (c * r).sum() -def _strip_dim(*x): - return [xx.view(xx.shape[2:]) for xx in x] - - class NTNTests(InteractionTests, unittest.TestCase): """Tests for NTN interaction function.""" @@ -400,7 +396,7 @@ def _exp_score(self, h, t, w, vt, vh, b, u, activation) -> torch.FloatTensor: # f(h,r,t) = u_r^T act(h W_r t + V_r h + V_r t + b_r) # shapes: w: (k, dim, dim), vh/vt: (k, dim), b/u: (k,), h/t: (dim,) # remove batch/num dimension - h, t, w, vt, vh, b, u = _strip_dim(h, t, w, vt, vh, b, u) + h, t, w, vt, vh, b, u = strip_dim(h, t, w, vt, vh, b, u) score = 0. for i in range(u.shape[-1]): first_part = h.view(1, self.dim) @ w[i] @ t.view(self.dim, 1) @@ -420,7 +416,7 @@ class ProjETests(InteractionTests, unittest.TestCase): def _exp_score(self, h, r, t, d_e, d_r, b_c, b_p, activation) -> torch.FloatTensor: # f(h, r, t) = g(t z(D_e h + D_r r + b_c) + b_p) - h, r, t = _strip_dim(h, r, t) + h, r, t = strip_dim(h, r, t) return (t * activation((d_e * h) + (d_r * r) + b_c)).sum() + b_p @@ -431,7 +427,7 @@ class RESCALTests(InteractionTests, unittest.TestCase): def _exp_score(self, h, r, t) -> torch.FloatTensor: # f(h, r, t) = h @ r @ t - h, r, t = _strip_dim(h, r, t) + h, r, t = strip_dim(h, r, t) return h.view(1, -1) @ r @ t.view(-1, 1) @@ -442,7 +438,7 @@ class KG2ETests(InteractionTests, unittest.TestCase): def _exp_score(self, exact, h_mean, h_var, r_mean, r_var, similarity, t_mean, t_var): assert similarity == "KL" - h_mean, h_var, r_mean, r_var, t_mean, t_var = _strip_dim(h_mean, h_var, r_mean, r_var, t_mean, t_var) + h_mean, h_var, r_mean, r_var, t_mean, t_var = strip_dim(h_mean, h_var, r_mean, r_var, t_mean, t_var) e_mean, e_var = h_mean - t_mean, h_var + t_var p = torch.distributions.MultivariateNormal(loc=e_mean, covariance_matrix=torch.diag(e_var)) q = torch.distributions.MultivariateNormal(loc=r_mean, covariance_matrix=torch.diag(r_var)) @@ -459,7 +455,7 @@ class TuckerTests(InteractionTests, unittest.TestCase): def _exp_score(self, bn_h, bn_hr, core_tensor, do_h, do_r, do_hr, h, r, t) -> torch.FloatTensor: # DO_{hr}(BN_{hr}(DO_h(BN_h(h)) x_1 DO_r(W x_2 r))) x_3 t - h, r, t = _strip_dim(h, r, t) + h, r, t = strip_dim(h, r, t) a = do_r((core_tensor * r[None, :, None]).sum(dim=1)) # shape: (embedding_dim, embedding_dim) b = do_h(bn_h(h.view(1, -1))).view(-1) # shape: (embedding_dim) c = (b[:, None] * a).sum(dim=1) # shape: (embedding_dim) @@ -473,7 +469,7 @@ class RotatETests(InteractionTests, unittest.TestCase): cls = pykeen.nn.modules.RotatEInteraction def _exp_score(self, h, r, t) -> torch.FloatTensor: # noqa: D102 - h, r, t = _strip_dim(*(view_complex(x) for x in (h, r, t))) + h, r, t = strip_dim(*(view_complex(x) for x in (h, r, t))) hr = h * r d = hr - t return -(d.abs() ** 2).sum().sqrt() @@ -551,7 +547,7 @@ class TransHTests(TranslationalInteractionTests, unittest.TestCase): def _exp_score(self, h, w_r, d_r, t, p, power_norm) -> torch.FloatTensor: # noqa: D102 assert not power_norm - h, w_r, d_r, t = _strip_dim(h, w_r, d_r, t) + h, w_r, d_r, t = strip_dim(h, w_r, d_r, t) h, t = [x - (x * w_r).sum() * w_r for x in (h, t)] return -(h + d_r - t).norm(p=p) @@ -577,7 +573,7 @@ def test_manual(self): def _exp_score(self, h, r, m_r, t, p, power_norm) -> torch.FloatTensor: assert power_norm - h, r, m_r, t = _strip_dim(h, r, m_r, t) + h, r, m_r, t = strip_dim(h, r, m_r, t) h_bot, t_bot = [clamp_norm(x.unsqueeze(dim=0) @ m_r, p=2, dim=-1, maxnorm=1.) for x in (h, t)] return -((h_bot + r - t_bot) ** p).sum() @@ -590,7 +586,7 @@ class SETests(TranslationalInteractionTests, unittest.TestCase): def _exp_score(self, h, t, r_h, r_t, p, power_norm) -> torch.FloatTensor: assert not power_norm # -\|R_h h - R_t t\| - h, t, r_h, r_t = _strip_dim(h, t, r_h, r_t) + h, t, r_h, r_t = strip_dim(h, t, r_h, r_t) h = r_h @ h.unsqueeze(dim=-1) t = r_t @ t.unsqueeze(dim=-1) return -(h - t).norm(p) @@ -604,7 +600,7 @@ class UMTests(TranslationalInteractionTests, unittest.TestCase): def _exp_score(self, h, t, p, power_norm) -> torch.FloatTensor: assert power_norm # -\|h - t\| - h, t = _strip_dim(h, t) + h, t = strip_dim(h, t) return -(h - t).pow(p).sum() diff --git a/tests/test_model_mode.py b/tests/test_model_mode.py index b40b0e9a50..51478feaee 100644 --- a/tests/test_model_mode.py +++ b/tests/test_model_mode.py @@ -3,7 +3,6 @@ """Test that models are set in the right mode when they're training.""" import unittest -from dataclasses import dataclass from unittest.mock import MagicMock import torch @@ -156,19 +155,3 @@ def test_alignment_of_score_r_fall_back(self) -> None: scores_r_function = self.model.score_r(ht_batch=ht_batch).flatten() scores_hrt_function = self.model.score_hrt(hrt_batch=hrt_batch).flatten() assert (scores_r_function == scores_hrt_function).all() - - -@dataclass -class MinimalTriplesFactory: - """A triples factory with minial attributes to allow the model to initiate.""" - - relation_to_id = { - "0": 0, - "1": 1, - } - entity_to_id = { - "0": 0, - "1": 1, - } - num_entities = 2 - num_relations = 2 diff --git a/tests/test_models.py b/tests/test_models.py index ac470422d8..9fa72e3a69 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -3,11 +3,12 @@ """Test that models can be executed.""" import importlib +import itertools as itt import os import tempfile import traceback import unittest -from typing import Any, ClassVar, Mapping, Optional, Sequence, Type +from typing import Any, ClassVar, Mapping, Optional, Type from unittest.mock import MagicMock, patch import numpy @@ -27,8 +28,6 @@ DoubleRelationEmbeddingModel, ERModel, Model, - SingleVectorEmbeddingModel, - TwoSideEmbeddingModel, TwoVectorEmbeddingModel, _extend_batch, get_novelty_mask, @@ -39,57 +38,25 @@ inverse_outdegree_edge_weights, symmetric_edge_weights, ) -from pykeen.nn import RepresentationModule from pykeen.regularizers import LpRegularizer, collect_regularization_terms +from pykeen.testing.mocks import MockRepresentations from pykeen.training import LCWATrainingLoop, SLCWATrainingLoop, TrainingLoop from pykeen.triples import TriplesFactory from pykeen.utils import all_in_bounds, set_random_seed SKIP_MODULES = { Model.__name__, - ERModel.__name__, - LiteralModel.__name__, - DoubleRelationEmbeddingModel.__name__, - SingleVectorEmbeddingModel.__name__, - TwoVectorEmbeddingModel.__name__, - TwoSideEmbeddingModel.__name__, 'DummyModel', 'MockModel', 'models', 'get_model_cls', } -for cls in LiteralModel.__subclasses__(): +for cls in itt.chain(_BASE_MODELS, LiteralModel.__subclasses__()): SKIP_MODULES.add(cls.__name__) _EPSILON = 1.0e-07 -class _CustomRepresentations(RepresentationModule): - """A custom representation module with minimal implementation.""" - - def __init__(self, num_entities: int, shape: Sequence[int]): - super().__init__() - self.num_embeddings = num_entities - self.shape = shape - self.x = nn.Parameter(torch.rand(numpy.prod(shape))) - - def forward(self, indices: Optional[torch.LongTensor] = None) -> torch.FloatTensor: - n = self.num_embeddings if indices is None else indices.shape[0] - return self.x.unsqueeze(dim=0).repeat(n, 1).view(-1, *self.shape) - - def get_in_canonical_shape( - self, - indices: Optional[torch.LongTensor] = None, - reshape_dim: Optional[Sequence[int]] = None, - ) -> torch.FloatTensor: - x = self(indices=indices) - if indices is None: - x = x.unsqueeze(dim=0) - else: - x = x.unsqueeze(dim=1) - return x - - class ERModelTests(unittest.TestCase): """Test basic functionality of ERModel.""" @@ -528,7 +495,7 @@ def test_custom_representations(self): old_entity_reps = self.model.entity_representations self.model.entity_representations = nn.ModuleList([ - _CustomRepresentations( + MockRepresentations( num_entities=self.factory.num_entities, shape=er.base_embeddings.shape if isinstance(er, RGCNRepresentations) else er.shape, ) @@ -536,7 +503,7 @@ def test_custom_representations(self): ]) old_relation_reps = self.model.relation_representations self.model.relation_representations = nn.ModuleList([ - _CustomRepresentations( + MockRepresentations( num_entities=self.factory.num_entities, shape=er.shape, ) @@ -1025,32 +992,6 @@ def test_models_have_experiments(self): self.fail(f'Missing experimental configuration directories for the following models:\n{_s}') -def test_extend_batch(): - """Test `_extend_batch()`.""" - batch = torch.tensor([[a, b] for a in range(3) for b in range(4)]).view(-1, 2) - all_ids = [2 * i for i in range(5)] - - batch_size = batch.shape[0] - num_choices = len(all_ids) - - for dim in range(3): - h_ext_batch = _extend_batch(batch=batch, all_ids=all_ids, dim=dim) - - # check shape - assert h_ext_batch.shape == (batch_size * num_choices, 3) - - # check content - actual_content = set(tuple(map(int, hrt)) for hrt in h_ext_batch) - exp_content = set() - for i in all_ids: - for b in batch: - c = list(map(int, b)) - c.insert(dim, i) - exp_content.add(tuple(c)) - - assert actual_content == exp_content - - class MessageWeightingTests(unittest.TestCase): """unittests for message weighting.""" @@ -1093,25 +1034,6 @@ def test_symmetric_edge_weights(self): self._test_message_weighting(weight_func=symmetric_edge_weights) -def test_get_novelty_mask(): - """Test `get_novelty_mask()`.""" - num_triples = 7 - base = torch.arange(num_triples) - mapped_triples = torch.stack([base, base, 3 * base], dim=-1) - query_ids = torch.randperm(num_triples).numpy()[:num_triples // 2] - exp_novel = query_ids != 0 - col = 2 - other_col_ids = numpy.asarray([0, 0]) - mask = get_novelty_mask( - mapped_triples=mapped_triples, - query_ids=query_ids, - col=col, - other_col_ids=other_col_ids, - ) - assert mask.shape == query_ids.shape - assert (mask == exp_novel).all() - - class TestRandom(unittest.TestCase): """Extra tests.""" @@ -1121,3 +1043,46 @@ def test_abstract(self): self.assertTrue(model_cls._is_abstract) for model_cls in _MODELS: self.assertFalse(model_cls._is_abstract, msg=f'{model_cls.__name__} should not be abstract') + + def test_get_novelty_mask(self): + """Test `get_novelty_mask()`.""" + num_triples = 7 + base = torch.arange(num_triples) + mapped_triples = torch.stack([base, base, 3 * base], dim=-1) + query_ids = torch.randperm(num_triples).numpy()[:num_triples // 2] + exp_novel = query_ids != 0 + col = 2 + other_col_ids = numpy.asarray([0, 0]) + mask = get_novelty_mask( + mapped_triples=mapped_triples, + query_ids=query_ids, + col=col, + other_col_ids=other_col_ids, + ) + assert mask.shape == query_ids.shape + assert (mask == exp_novel).all() + + def test_extend_batch(self): + """Test `_extend_batch()`.""" + batch = torch.tensor([[a, b] for a in range(3) for b in range(4)]).view(-1, 2) + all_ids = [2 * i for i in range(5)] + + batch_size = batch.shape[0] + num_choices = len(all_ids) + + for dim in range(3): + h_ext_batch = _extend_batch(batch=batch, all_ids=all_ids, dim=dim) + + # check shape + assert h_ext_batch.shape == (batch_size * num_choices, 3) + + # check content + actual_content = set(tuple(map(int, hrt)) for hrt in h_ext_batch) + exp_content = set() + for i in all_ids: + for b in batch: + c = list(map(int, b)) + c.insert(dim, i) + exp_content.add(tuple(c)) + + assert actual_content == exp_content diff --git a/tests/test_regularizers.py b/tests/test_regularizers.py index 7f8b3be854..5c65b9d705 100644 --- a/tests/test_regularizers.py +++ b/tests/test_regularizers.py @@ -310,57 +310,60 @@ def _help_test_regularizer(self, regularizer: Regularizer, n_tensors: int = 3): assert 0.0 == regularizer.regularization_term -def test_collect_regularization_terms(): - """Test whether collect_regularization_terms finds and resets all regularization terms.""" - regularizers = [ - LpRegularizer(), - PowerSumRegularizer(), - LpRegularizer(p=1, normalize=True, apply_only_once=True), - PowerSumRegularizer(normalize=True), - ] - model = ERModel( - triples_factory=MagicMock(), - interaction=MagicMock(), - entity_representations=EmbeddingSpecification( - regularizer=regularizers[0], - ).make(num_embeddings=3, embedding_dim=2, shape=None), - relation_representations=EmbeddingSpecification( - regularizer=regularizers[1], - ).make(num_embeddings=3, embedding_dim=2, shape=None), - ) - - # add weighted modules - linear = nn.Linear(3, 2) - model.sub_module = nn.ModuleList([ - nn.Sequential( - linear, - nn.Linear(2, 3), - ), - nn.BatchNorm1d(2), - linear, # one module occuring twice - ]) - - # add weight regularizer - model.add_weight_regularizer( - parameter_name="sub_module.0.0.bias", - regularizer=regularizers[2], - ) - model.add_weight_regularizer( - parameter_name="entity_representations.0._embeddings.weight", - regularizer=regularizers[3], - ) - - # retrieve all regularization terms - collect_regularization_terms(model) - - # check that all terms are reset - found_regularizers = set() - for module in model.modules(): - if isinstance(module, Regularizer): - term = module.regularization_term - assert isinstance(term, float) - assert term == 0.0 - found_regularizers.add(id(module)) - - # check that all regularizers were found - assert found_regularizers == set(map(id, regularizers)) +class TestRandom(unittest.TestCase): + """Test random regularization utilities.""" + + def test_collect_regularization_terms(self): + """Test whether collect_regularization_terms finds and resets all regularization terms.""" + regularizers = [ + LpRegularizer(), + PowerSumRegularizer(), + LpRegularizer(p=1, normalize=True, apply_only_once=True), + PowerSumRegularizer(normalize=True), + ] + model = ERModel( + triples_factory=MagicMock(), + interaction=MagicMock(), + entity_representations=EmbeddingSpecification( + regularizer=regularizers[0], + ).make(num_embeddings=3, embedding_dim=2, shape=None), + relation_representations=EmbeddingSpecification( + regularizer=regularizers[1], + ).make(num_embeddings=3, embedding_dim=2, shape=None), + ) + + # add weighted modules + linear = nn.Linear(3, 2) + model.sub_module = nn.ModuleList([ + nn.Sequential( + linear, + nn.Linear(2, 3), + ), + nn.BatchNorm1d(2), + linear, # one module occuring twice + ]) + + # add weight regularizer + model.add_weight_regularizer( + parameter_name="sub_module.0.0.bias", + regularizer=regularizers[2], + ) + model.add_weight_regularizer( + parameter_name="entity_representations.0._embeddings.weight", + regularizer=regularizers[3], + ) + + # retrieve all regularization terms + collect_regularization_terms(model) + + # check that all terms are reset + found_regularizers = set() + for module in model.modules(): + if isinstance(module, Regularizer): + term = module.regularization_term + assert isinstance(term, float) + assert term == 0.0 + found_regularizers.add(id(module)) + + # check that all regularizers were found + self.assertEqual(found_regularizers, set(map(id, regularizers))) diff --git a/tests/training/test_utils.py b/tests/test_training_utils.py similarity index 100% rename from tests/training/test_utils.py rename to tests/test_training_utils.py diff --git a/tests/test_utils.py b/tests/test_utils.py index 8d02a0472b..876ce365e9 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -142,6 +142,40 @@ def test_complex_utils(self): assert (re2 == re).all() assert (im2 == im).all() + def test_project_entity(self): + """Test _project_entity.""" + batch_size = 2 + embedding_dim = 3 + relation_dim = 5 + num_entities = 7 + + # random entity embeddings & projections + e = torch.rand(1, num_entities, embedding_dim) + e = clamp_norm(e, maxnorm=1, p=2, dim=-1) + e_p = torch.rand(1, num_entities, embedding_dim) + + # random relation embeddings & projections + r_p = torch.rand(batch_size, 1, relation_dim) + + # project + e_bot = project_entity(e=e, e_p=e_p, r_p=r_p) + + # check shape: + assert e_bot.shape == (batch_size, num_entities, relation_dim) + + # check normalization + assert (torch.norm(e_bot, dim=-1, p=2) <= 1.0 + 1.0e-06).all() + + # check equivalence of re-formulation + # e_{\bot} = M_{re} e = (r_p e_p^T + I^{d_r \times d_e}) e + # = r_p (e_p^T e) + e' + m_re = r_p.unsqueeze(dim=-1) @ e_p.unsqueeze(dim=-2) + m_re = m_re + torch.eye(relation_dim, embedding_dim).view(1, 1, relation_dim, embedding_dim) + assert m_re.shape == (batch_size, num_entities, relation_dim, embedding_dim) + e_vanilla = (m_re @ e.unsqueeze(dim=-1)).squeeze(dim=-1) + e_vanilla = clamp_norm(e_vanilla, p=2, dim=-1, maxnorm=1) + assert torch.allclose(e_vanilla, e_bot) + class TestCudaExceptionsHandling(unittest.TestCase): """Test handling of CUDA exceptions.""" @@ -163,38 +197,3 @@ def test_is_cudnn_error(self): self.assertFalse(is_cuda_oom_error(runtime_error=error)) self.assertFalse(is_cudnn_error(runtime_error=self.not_cuda_error)) - - -def test_project_entity(): - """Test _project_entity.""" - batch_size = 2 - embedding_dim = 3 - relation_dim = 5 - num_entities = 7 - - # random entity embeddings & projections - e = torch.rand(1, num_entities, embedding_dim) - e = clamp_norm(e, maxnorm=1, p=2, dim=-1) - e_p = torch.rand(1, num_entities, embedding_dim) - - # random relation embeddings & projections - r_p = torch.rand(batch_size, 1, relation_dim) - - # project - e_bot = project_entity(e=e, e_p=e_p, r_p=r_p) - - # check shape: - assert e_bot.shape == (batch_size, num_entities, relation_dim) - - # check normalization - assert (torch.norm(e_bot, dim=-1, p=2) <= 1.0 + 1.0e-06).all() - - # check equivalence of re-formulation - # e_{\bot} = M_{re} e = (r_p e_p^T + I^{d_r \times d_e}) e - # = r_p (e_p^T e) + e' - m_re = r_p.unsqueeze(dim=-1) @ e_p.unsqueeze(dim=-2) - m_re = m_re + torch.eye(relation_dim, embedding_dim).view(1, 1, relation_dim, embedding_dim) - assert m_re.shape == (batch_size, num_entities, relation_dim, embedding_dim) - e_vanilla = (m_re @ e.unsqueeze(dim=-1)).squeeze(dim=-1) - e_vanilla = clamp_norm(e_vanilla, p=2, dim=-1, maxnorm=1) - assert torch.allclose(e_vanilla, e_bot) From cc0397cfb1a771095569f4ed91d30d80305685fd Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 19 Nov 2020 13:32:13 +0100 Subject: [PATCH 460/690] Update meta make sure still works on 37 sorry @mberr --- setup.cfg | 1 + src/pykeen/utils.py | 6 +++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/setup.cfg b/setup.cfg index e9a1a59fdf..0f06fde2e8 100644 --- a/setup.cfg +++ b/setup.cfg @@ -36,6 +36,7 @@ classifiers = Programming Language :: Python :: 3.6 Programming Language :: Python :: 3.7 Programming Language :: Python :: 3.8 + Programming Language :: Python :: 3.9 Programming Language :: Python :: 3 :: Only Topic :: Scientific/Engineering :: Artificial Intelligence Topic :: Scientific/Engineering :: Chemistry diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 5ee7000ce9..03eaae5ea1 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -8,7 +8,7 @@ import random from abc import ABC, abstractmethod from io import BytesIO -from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, SupportsFloat, Tuple, Type, TypeVar, Union +from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Type, TypeVar, Union import numpy as np import pandas as pd @@ -518,11 +518,11 @@ def negative_norm_of_sum( """ d: torch.FloatTensor = tensor_sum(*x) if power_norm: - assert isinstance(p, SupportsFloat) + assert not isinstance(p, str) return -(d.abs() ** p).sum(dim=-1) if torch.is_complex(d): - assert isinstance(p, SupportsFloat) + assert not isinstance(p, str) # workaround for complex numbers: manually compute norm return -(d.abs() ** p).sum(dim=-1) ** (1 / p) From 3bb9f1619f6dd675902d4c42a6139ae736f14a32 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 19 Nov 2020 13:42:02 +0100 Subject: [PATCH 461/690] Update adding regularizers @mberr the weight_regularizers instance was only used in tests, so I updated the TransH class to use it explicitly. I know that the regularizers are looked up automatically, but this makes it more obvious why the weight_regularizers variable is even there --- src/pykeen/models/base.py | 4 ++-- src/pykeen/models/unimodal/trans_h.py | 4 +++- tests/test_models.py | 4 ++-- tests/test_regularizers.py | 4 ++-- 4 files changed, 9 insertions(+), 7 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 0d7577f35a..0560b69fb3 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1161,7 +1161,7 @@ def __init__( # model.modules(). Thereby, we can collect them automatically. self.weight_regularizers = nn.ModuleList() - def add_weight_regularizer( + def append_weight_regularizer( self, parameter_name: str, regularizer: Regularizer, @@ -1173,7 +1173,7 @@ def add_weight_regularizer( :param regularizer: The regularizer instance which will regularize the weights. """ - weights = dict(self.named_parameters()) + weights: Mapping[str, nn.Parameter] = dict(self.named_parameters()) if parameter_name not in weights.keys(): raise ValueError(f"Invalid parameter_name={parameter_name}. Available are: {sorted(weights.keys())}.") parameter: nn.Parameter = weights[parameter_name] diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 488f1e1c65..d8a0976365 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -98,9 +98,11 @@ def __init__( constrainer=functional.normalize, ), ) + # Note that the TransH regularizer has a different interface - self.regularizer = self._instantiate_default_regularizer( + regularizer = self._instantiate_default_regularizer( entity_embeddings=pop_only(self.entity_representations[0].parameters()), relation_embeddings=pop_only(self.relation_representations[0].parameters()), normal_vector_embeddings=pop_only(self.relation_representations[1].parameters()), ) + self.weight_regularizers.append(regularizer) diff --git a/tests/test_models.py b/tests/test_models.py index 9fa72e3a69..d2f4c38d91 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -70,7 +70,7 @@ def setUp(self) -> None: def test_add_weight_regularizer_non_existing(self): """Test that an assertion is raised for add_weight_regularizer for a non-existing weight.""" with self.assertRaises(ValueError): - self.model.add_weight_regularizer( + self.model.append_weight_regularizer( parameter_name="this.weight.does.not.exist", regularizer=..., ) @@ -86,7 +86,7 @@ def test_add_weight_regularizer(self): ) # try to add regularizer to existing weight - self.model.add_weight_regularizer( + self.model.append_weight_regularizer( parameter_name="linear.weight", regularizer=LpRegularizer(), ) diff --git a/tests/test_regularizers.py b/tests/test_regularizers.py index 5c65b9d705..7a2bb457eb 100644 --- a/tests/test_regularizers.py +++ b/tests/test_regularizers.py @@ -344,11 +344,11 @@ def test_collect_regularization_terms(self): ]) # add weight regularizer - model.add_weight_regularizer( + model.append_weight_regularizer( parameter_name="sub_module.0.0.bias", regularizer=regularizers[2], ) - model.add_weight_regularizer( + model.append_weight_regularizer( parameter_name="entity_representations.0._embeddings.weight", regularizer=regularizers[3], ) From 362e3fbbf98324ff9b5512e5248fa78a8b395765 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 19 Nov 2020 13:49:44 +0100 Subject: [PATCH 462/690] Update trans_h.py --- src/pykeen/models/unimodal/trans_h.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index d8a0976365..488f1e1c65 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -98,11 +98,9 @@ def __init__( constrainer=functional.normalize, ), ) - # Note that the TransH regularizer has a different interface - regularizer = self._instantiate_default_regularizer( + self.regularizer = self._instantiate_default_regularizer( entity_embeddings=pop_only(self.entity_representations[0].parameters()), relation_embeddings=pop_only(self.relation_representations[0].parameters()), normal_vector_embeddings=pop_only(self.relation_representations[1].parameters()), ) - self.weight_regularizers.append(regularizer) From be130bfeb7062c4ee98912be91300ee539fad82e Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 19 Nov 2020 13:57:57 +0100 Subject: [PATCH 463/690] Update name for clarity --- src/pykeen/models/__init__.py | 2 +- src/pykeen/models/base.py | 6 +++--- tests/test_models.py | 7 +++++-- 3 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index 5d2e2c23b7..e09320c0bf 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -78,7 +78,7 @@ def _concrete_subclasses(cls: Type[Model]): for subcls in cls.__subclasses__(): - if not subcls._is_abstract and subcls not in _BASE_MODELS: + if not subcls._is_base_model and subcls not in _BASE_MODELS: yield subcls yield from _concrete_subclasses(subcls) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 0560b69fb3..7560322090 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -236,7 +236,7 @@ class Model(nn.Module, ABC): _hyperparameter_usage: ClassVar[Dict[str, Set[str]]] = defaultdict(set) #: Keep track of if this is a base model - _is_abstract: ClassVar[bool] + _is_base_model: ClassVar[bool] #: The default strategy for optimizing the model's hyper-parameters hpo_default: ClassVar[Mapping[str, Any]] @@ -318,8 +318,8 @@ def __init__( self.automatic_memory_optimization = automatic_memory_optimization def __init_subclass__(cls, autoreset: bool = True, **kwargs): # noqa:D105 - cls._is_abstract = not autoreset - if not cls._is_abstract: + cls._is_base_model = not autoreset + if not cls._is_base_model: _track_hyperparameters(cls) _add_post_reset_parameters(cls) diff --git a/tests/test_models.py b/tests/test_models.py index d2f4c38d91..afeaf1288a 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1040,9 +1040,12 @@ class TestRandom(unittest.TestCase): def test_abstract(self): """Test that classes are checked as abstract properly.""" for model_cls in _BASE_MODELS: - self.assertTrue(model_cls._is_abstract) + self.assertTrue(model_cls._is_base_model) for model_cls in _MODELS: - self.assertFalse(model_cls._is_abstract, msg=f'{model_cls.__name__} should not be abstract') + self.assertFalse( + model_cls._is_base_model, + msg=f'{model_cls.__name__} should not be marked as a a base model', + ) def test_get_novelty_mask(self): """Test `get_novelty_mask()`.""" From 7dda1037857dd456c383e6e691146633c62a7b9a Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 19 Nov 2020 14:39:42 +0100 Subject: [PATCH 464/690] Fix british --- src/pykeen/models/unimodal/hole.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 40233db82e..4e07aef62c 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -69,7 +69,7 @@ def __init__( automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, - # Initialisation, cf. https://github.com/mnick/scikit-kge/blob/master/skge/param.py#L18-L27 + # Initialization, cf. https://github.com/mnick/scikit-kge/blob/master/skge/param.py#L18-L27 embedding_specification=EmbeddingSpecification( initializer=xavier_uniform_, constrainer=clamp_norm, # type: ignore From e6d945bfba67e01e5730f01d3c022d9415ce306d Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 19 Nov 2020 14:39:49 +0100 Subject: [PATCH 465/690] Mark slow test --- tests/test_early_stopping.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_early_stopping.py b/tests/test_early_stopping.py index 4cbfb1b67e..2a8395e731 100644 --- a/tests/test_early_stopping.py +++ b/tests/test_early_stopping.py @@ -6,6 +6,7 @@ from typing import Iterable, List, Optional import numpy +import pytest import torch from torch.optim import Adam @@ -231,6 +232,7 @@ def setUp(self) -> None: torch.manual_seed(seed=self.seed) numpy.random.seed(seed=self.seed) + @pytest.mark.slow def test_early_stopping(self): """Tests early stopping.""" # Set automatic_memory_optimization to false during testing From 3bb92c81e06cdffcd5c9f866e0ab621160d2c0e4 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 19 Nov 2020 14:43:33 +0100 Subject: [PATCH 466/690] Update .travis.yml --- .travis.yml | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/.travis.yml b/.travis.yml index 8ca6084948..d1c317c748 100644 --- a/.travis.yml +++ b/.travis.yml @@ -29,6 +29,19 @@ jobs: - name: "Test fast unit tests" stage: fast_tests env: TOXENV=py + # borrowed windows testing config from https://github.com/shaypal5/cachier/blob/master/.travis.yml + - name: "Test fast unit tests (Windows)" + stage: fast_tests + os: windows # Windows 10.0.17134 N/A Build 17134 + language: shell # 'language: python' is an error on Travis CI Windows + before_install: + - choco install python --version 3.7.4 + - python --version + install: + - pip install tox codecov coverage + env: PATH=/c/Python37:/c/Python37/Scripts:$PATH TOXENV=py + after_success: + - tox -e coverage-report - name: "Test slow unit tests" stage: slow_tests env: TOXENV=integration From e0e123d74f8f14b7a1b788c4d52544adc7126a71 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 19 Nov 2020 15:01:30 +0100 Subject: [PATCH 467/690] Update .travis.yml --- .travis.yml | 13 ------------- 1 file changed, 13 deletions(-) diff --git a/.travis.yml b/.travis.yml index d1c317c748..8ca6084948 100644 --- a/.travis.yml +++ b/.travis.yml @@ -29,19 +29,6 @@ jobs: - name: "Test fast unit tests" stage: fast_tests env: TOXENV=py - # borrowed windows testing config from https://github.com/shaypal5/cachier/blob/master/.travis.yml - - name: "Test fast unit tests (Windows)" - stage: fast_tests - os: windows # Windows 10.0.17134 N/A Build 17134 - language: shell # 'language: python' is an error on Travis CI Windows - before_install: - - choco install python --version 3.7.4 - - python --version - install: - - pip install tox codecov coverage - env: PATH=/c/Python37:/c/Python37/Scripts:$PATH TOXENV=py - after_success: - - tox -e coverage-report - name: "Test slow unit tests" stage: slow_tests env: TOXENV=integration From e5d33d8db85c7e0c907259e006efaf61b9e1d19c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 19 Nov 2020 16:29:22 +0100 Subject: [PATCH 468/690] Extend append_weight_regularizer input types --- src/pykeen/models/base.py | 20 +++++++++++++------- tests/test_models.py | 4 ++-- tests/test_regularizers.py | 4 ++-- 3 files changed, 17 insertions(+), 11 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 7560322090..8c87fb5e11 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1163,21 +1163,27 @@ def __init__( def append_weight_regularizer( self, - parameter_name: str, + parameter: Union[str, nn.Parameter, Iterable[Union[str, nn.Parameter]]], regularizer: Regularizer, ) -> None: """Add a model weight to a regularizer's weight list, and register the regularizer with the model. - :param parameter_name: - The parameter name. Available parameter names are shown by `dict(self.named_parameters()).keys()`. + :param parameter: + The parameter, either as name, or as nn.Parameter object. A list of available parameter names is shown by + `sorted(dict(self.named_parameters()).keys())`. :param regularizer: The regularizer instance which will regularize the weights. """ + # normalize input + if isinstance(parameter, (str, nn.Parameter)): + parameter = list(parameter) weights: Mapping[str, nn.Parameter] = dict(self.named_parameters()) - if parameter_name not in weights.keys(): - raise ValueError(f"Invalid parameter_name={parameter_name}. Available are: {sorted(weights.keys())}.") - parameter: nn.Parameter = weights[parameter_name] - regularizer.add_parameter(parameter=parameter) + for param in parameter: + if isinstance(param, str): + if parameter not in weights.keys(): + raise ValueError(f"Invalid parameter_name={parameter}. Available are: {sorted(weights.keys())}.") + param: nn.Parameter = weights[param] + regularizer.add_parameter(parameter=param) self.weight_regularizers.append(regularizer) def forward( diff --git a/tests/test_models.py b/tests/test_models.py index afeaf1288a..09cceb7380 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -71,7 +71,7 @@ def test_add_weight_regularizer_non_existing(self): """Test that an assertion is raised for add_weight_regularizer for a non-existing weight.""" with self.assertRaises(ValueError): self.model.append_weight_regularizer( - parameter_name="this.weight.does.not.exist", + parameter="this.weight.does.not.exist", regularizer=..., ) @@ -87,7 +87,7 @@ def test_add_weight_regularizer(self): # try to add regularizer to existing weight self.model.append_weight_regularizer( - parameter_name="linear.weight", + parameter="linear.weight", regularizer=LpRegularizer(), ) diff --git a/tests/test_regularizers.py b/tests/test_regularizers.py index 7a2bb457eb..a5f1732898 100644 --- a/tests/test_regularizers.py +++ b/tests/test_regularizers.py @@ -345,11 +345,11 @@ def test_collect_regularization_terms(self): # add weight regularizer model.append_weight_regularizer( - parameter_name="sub_module.0.0.bias", + parameter="sub_module.0.0.bias", regularizer=regularizers[2], ) model.append_weight_regularizer( - parameter_name="entity_representations.0._embeddings.weight", + parameter="entity_representations.0._embeddings.weight", regularizer=regularizers[3], ) From ccfeb08c4a76ad50fb6dd5c9a317b16644c6c873 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 19 Nov 2020 16:29:41 +0100 Subject: [PATCH 469/690] Re-introduce ConvKB regularizer comment, and use append_weight_regularizer --- src/pykeen/models/unimodal/conv_kb.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 5fd14928e8..4961cbf7ee 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -96,5 +96,10 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, ) - self.regularizer = self._instantiate_default_regularizer(parameters=list(self.interaction.parameters())) + # In the code base only the weights of the output layer are used for regularization + # c.f. https://github.com/daiquocnguyen/ConvKB/blob/73a22bfa672f690e217b5c18536647c7cf5667f1/model.py#L60-L66 + self.append_weight_regularizer( + parameter=self.interaction.parameters(), + regularizer=self._instantiate_default_regularizer(), + ) logger.warning('To be consistent with the paper, initialize entity and relation embeddings from TransE.') From 7d8eef3703ae6101ec9a684d50166f3bc0107695 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 19 Nov 2020 16:40:10 +0100 Subject: [PATCH 470/690] Cleanup utils.__all__ --- src/pykeen/utils.py | 18 +++++++++++------- tests/test_interactions.py | 4 +++- 2 files changed, 14 insertions(+), 8 deletions(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 03eaae5ea1..cddcc95578 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -18,6 +18,7 @@ from .typing import DeviceHint +# get_subclasses, project_entity, set_random_seed, strip_dim, view_complex __all__ = [ 'broadcast_cat', 'compose', @@ -27,20 +28,23 @@ 'compact_mapping', 'complex_normalize', 'fix_dataclass_init_docs', + 'flatten_dictionary', + 'get_cls', + 'get_subclasses', + 'get_until_first_blank', 'invert_mapping', 'is_cudnn_error', 'is_cuda_oom_error', + 'project_entity', + 'negative_norm_of_sum', + 'normalize_string', + 'normalized_lookup', 'random_non_negative_int', 'resolve_device', + 'set_random_seed', 'split_complex', 'split_list_in_batches_iter', - 'normalize_string', - 'normalized_lookup', - 'negative_norm_of_sum', - 'get_cls', - 'get_until_first_blank', - 'flatten_dictionary', - 'set_random_seed', + 'view_complex', 'NoRandomSeedNecessary', 'Result', ] diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 2995276907..78b13c9214 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -14,7 +14,7 @@ import pykeen.nn.modules from pykeen.nn.modules import Interaction, TranslationalInteraction from pykeen.typing import Representation -from pykeen.utils import clamp_norm, get_subclasses, project_entity, strip_dim, view_complex +from pykeen.utils import clamp_norm, get_subclasses, project_entity, set_random_seed, strip_dim, view_complex T = TypeVar("T") logger = logging.getLogger(__name__) @@ -29,6 +29,8 @@ class GenericTests(Generic[T]): def setUp(self) -> None: """Set up the generic testing method.""" + # fix seeds for reproducibility + set_random_seed(seed=42) kwargs = self.kwargs or {} kwargs = self._pre_instantiation_hook(kwargs=dict(kwargs)) self.instance = self.cls(**kwargs) From 4ec4e2bd5ec23dd1017c7f66d358654cb040ec4f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 19 Nov 2020 16:44:10 +0100 Subject: [PATCH 471/690] Add comment to test_scores --- tests/test_interactions.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 78b13c9214..22323e30bd 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -255,6 +255,7 @@ def test_forward_consistency_with_functional(self): def test_scores(self): """Test individual scores.""" + # set in eval mode (otherwise there are non-deterministic factors like Dropout self.instance.eval() for _ in range(10): h, r, t = self._get_hrt((1, 1), (1, 1), (1, 1)) From 71944da4a271a15b9e1d1e55fdfa04cbbbd12ca8 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 19 Nov 2020 18:59:16 +0100 Subject: [PATCH 472/690] Fix bug in upgrading --- src/pykeen/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 8c87fb5e11..d72760e7cc 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1176,7 +1176,7 @@ def append_weight_regularizer( """ # normalize input if isinstance(parameter, (str, nn.Parameter)): - parameter = list(parameter) + parameter = [parameter] weights: Mapping[str, nn.Parameter] = dict(self.named_parameters()) for param in parameter: if isinstance(param, str): From 859c0850626623095e085ef85bc23fd8f7878d7d Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 19 Nov 2020 18:59:23 +0100 Subject: [PATCH 473/690] Bandage on tricky typing --- src/pykeen/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index d72760e7cc..53b3665789 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1182,7 +1182,7 @@ def append_weight_regularizer( if isinstance(param, str): if parameter not in weights.keys(): raise ValueError(f"Invalid parameter_name={parameter}. Available are: {sorted(weights.keys())}.") - param: nn.Parameter = weights[param] + param: nn.Parameter = weights[param] # type: ignore regularizer.add_parameter(parameter=param) self.weight_regularizers.append(regularizer) From 47da63ec91fb70d097f2ed4d1d17e7bbc05a940c Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 19 Nov 2020 19:00:41 +0100 Subject: [PATCH 474/690] Make typing correct It's possible that regularization could be turned off on ConvKB (though not correct) so this makes mypy happy --- src/pykeen/models/unimodal/conv_kb.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 4961cbf7ee..600d211b61 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -96,10 +96,12 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, ) + regularizer = self._instantiate_default_regularizer() # In the code base only the weights of the output layer are used for regularization # c.f. https://github.com/daiquocnguyen/ConvKB/blob/73a22bfa672f690e217b5c18536647c7cf5667f1/model.py#L60-L66 - self.append_weight_regularizer( - parameter=self.interaction.parameters(), - regularizer=self._instantiate_default_regularizer(), - ) + if regularizer is not None: + self.append_weight_regularizer( + parameter=self.interaction.parameters(), + regularizer=regularizer, + ) logger.warning('To be consistent with the paper, initialize entity and relation embeddings from TransE.') From f1fbdeb56641cf2c3144661f1029f68c2bfce590 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 19 Nov 2020 19:57:50 +0100 Subject: [PATCH 475/690] Fix typo --- src/pykeen/models/base.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 53b3665789..9a1154eae2 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1180,8 +1180,8 @@ def append_weight_regularizer( weights: Mapping[str, nn.Parameter] = dict(self.named_parameters()) for param in parameter: if isinstance(param, str): - if parameter not in weights.keys(): - raise ValueError(f"Invalid parameter_name={parameter}. Available are: {sorted(weights.keys())}.") + if param not in weights: + raise KeyError(f"Invalid parameter_name={parameter}. Available are: {sorted(weights.keys())}.") param: nn.Parameter = weights[param] # type: ignore regularizer.add_parameter(parameter=param) self.weight_regularizers.append(regularizer) From af97e03527882b68982f250acadd69fa9c3ee621 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:17:47 +0100 Subject: [PATCH 476/690] Fix unittest catching error --- tests/test_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index 09cceb7380..000bf03349 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -69,7 +69,7 @@ def setUp(self) -> None: def test_add_weight_regularizer_non_existing(self): """Test that an assertion is raised for add_weight_regularizer for a non-existing weight.""" - with self.assertRaises(ValueError): + with self.assertRaises(KeyError): self.model.append_weight_regularizer( parameter="this.weight.does.not.exist", regularizer=..., From 32df6b6cecc7f0422f113f3e74d8a8f455be9d7a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:20:59 +0100 Subject: [PATCH 477/690] Rename test case --- tests/test_models.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 000bf03349..c2e1219ac8 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1034,8 +1034,8 @@ def test_symmetric_edge_weights(self): self._test_message_weighting(weight_func=symmetric_edge_weights) -class TestRandom(unittest.TestCase): - """Extra tests.""" +class TestModelUtilities(unittest.TestCase): + """Extra tests for utilit functions.""" def test_abstract(self): """Test that classes are checked as abstract properly.""" From 4c6fed4562ab0b97ad61dc0f143868fe4866c4ae Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:24:23 +0100 Subject: [PATCH 478/690] Test multiple different initializations for test_scores --- tests/test_interactions.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 22323e30bd..ad9e082e71 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -258,6 +258,8 @@ def test_scores(self): # set in eval mode (otherwise there are non-deterministic factors like Dropout self.instance.eval() for _ in range(10): + # test multiple different initializations + self.instance.reset_parameters() h, r, t = self._get_hrt((1, 1), (1, 1), (1, 1)) kwargs = self.instance._prepare_for_functional(h=h, r=r, t=t) From 9a8bc32236419f3ed4c75b01861c79a2bae4c7bd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:28:04 +0100 Subject: [PATCH 479/690] Fix typo --- src/pykeen/nn/modules.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 2c523562d8..dd43501c90 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -915,9 +915,9 @@ def __init__( if apply_batch_normalization: self.head_batch_norm = nn.BatchNorm1d(embedding_dim) - self.head_relation_batchnorm = nn.BatchNorm1d(embedding_dim) + self.head_relation_batch_norm = nn.BatchNorm1d(embedding_dim) else: - self.head_batch_norm = self.head_relation_batchnorm = None + self.head_batch_norm = self.head_relation_batch_norm = None def reset_parameters(self): # noqa:D102 # Initialize core tensor, cf. https://github.com/ibalazevic/TuckER/blob/master/model.py#L12 @@ -931,7 +931,7 @@ def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: do_r=self.relation_dropout, do_hr=self.head_relation_dropout, bn_h=self.head_batch_norm, - bn_hr=self.head_relation_batchnorm, + bn_hr=self.head_relation_batch_norm, ) From 80e059c6c8f73d49b3d8699b1f114374f5189c59 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:28:32 +0100 Subject: [PATCH 480/690] Add TODO --- src/pykeen/nn/modules.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index dd43501c90..665061b9e2 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -923,6 +923,7 @@ def reset_parameters(self): # noqa:D102 # Initialize core tensor, cf. https://github.com/ibalazevic/TuckER/blob/master/model.py#L12 nn.init.uniform_(self.core_tensor, -1., 1.) # batch norm gets reset automatically, since it defines reset_parameters + # TODO: This is only the case if bound to a Model, but not if used without it def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: return dict( From 302c8db733f0659084b4783198b710a414540cb5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:29:28 +0100 Subject: [PATCH 481/690] Remove wrong TODO --- src/pykeen/nn/modules.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 665061b9e2..dd43501c90 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -923,7 +923,6 @@ def reset_parameters(self): # noqa:D102 # Initialize core tensor, cf. https://github.com/ibalazevic/TuckER/blob/master/model.py#L12 nn.init.uniform_(self.core_tensor, -1., 1.) # batch norm gets reset automatically, since it defines reset_parameters - # TODO: This is only the case if bound to a Model, but not if used without it def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: return dict( From 7725749ed413fd6eb06ca9f01bb1ce7b544de10b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:32:14 +0100 Subject: [PATCH 482/690] Equivalentl rename dimensions for Tucker functional --- src/pykeen/nn/functional.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 840a32a9e5..eeb813442f 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -912,15 +912,15 @@ def tucker_interaction( """ return extended_einsum( # x_3 contraction - "bhrd,btd->bhrt", + "bhrk,btk->bhrt", _apply_optional_bn_to_tensor( x=extended_einsum( # x_1 contraction - "brid,bhd->bhri", + "brik,bhi->bhrk", _apply_optional_bn_to_tensor( x=extended_einsum( # x_2 contraction - "idj,brd->brij", + "ijk,brj->brik", core_tensor, r, ), From 3c5be20d527553a012de3743b7e0d35763a2df6a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:34:30 +0100 Subject: [PATCH 483/690] Fix exp_score for TuckerTests --- tests/test_interactions.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index ad9e082e71..7018e9b83a 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -461,11 +461,11 @@ class TuckerTests(InteractionTests, unittest.TestCase): def _exp_score(self, bn_h, bn_hr, core_tensor, do_h, do_r, do_hr, h, r, t) -> torch.FloatTensor: # DO_{hr}(BN_{hr}(DO_h(BN_h(h)) x_1 DO_r(W x_2 r))) x_3 t h, r, t = strip_dim(h, r, t) - a = do_r((core_tensor * r[None, :, None]).sum(dim=1)) # shape: (embedding_dim, embedding_dim) + a = do_r((core_tensor * r[None, :, None]).sum(dim=1, keepdims=True)) # shape: (embedding_dim, 1, embedding_dim) b = do_h(bn_h(h.view(1, -1))).view(-1) # shape: (embedding_dim) - c = (b[:, None] * a).sum(dim=1) # shape: (embedding_dim) - d = do_hr(bn_hr((c.view(1, -1)))).view(-1) # shape: (embedding_dim) - return (d * t).sum() + c = (b[:, None, None] * a).sum(dim=0, keepdims=True) # shape: (1, 1, embedding_dim) + d = do_hr(bn_hr((c.view(1, -1)))).view(1, 1, -1) # shape: (1, 1, 1, embedding_dim) + return (d * t[None, None, :]).sum() class RotatETests(InteractionTests, unittest.TestCase): From e43cefb1415bddbdb78014d77ce1a3659830e08c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:52:39 +0100 Subject: [PATCH 484/690] Fix some darglint issues for nn.emb docstrings --- src/pykeen/nn/emb.py | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index 41e056e74e..a719471f2b 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -145,6 +145,8 @@ def __init__( The number of embeddings. :param embedding_dim: >0 The embedding dimensionality. + :param shape: + The embedding shape. If given, shape supersedes embedding_dim, with setting embedding_dim = prod(shape). :param initializer: An optional initializer, which takes an uninitialized (num_embeddings, embedding_dim) tensor as input, and returns an initialized tensor of same shape and dtype (which may be the same, i.e. the @@ -161,6 +163,9 @@ def __init__( to be in-place, but the weight tensor is modified in-place. :param constrainer_kwargs: Additional keyword arguments passed to the constrainer + + :raises ValueError: + If neither shape nor embedding_dim are given. """ super().__init__() if shape is None and embedding_dim is None: @@ -208,10 +213,12 @@ def from_specification( ) -> 'Embedding': """Create an embedding based on a specification. - :param num_embeddings: + :param num_embeddings: >0 The number of embeddings. - :param embedding_dim: + :param embedding_dim: >0 The embedding dimension. + :param shape: + The embedding shape. If given, shape supersedes embedding_dim, with setting embedding_dim = prod(shape). :param specification: The specification. :return: From 9ff8104e00b4a1a5436b3719f5fb7bd271df97ad Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:54:00 +0100 Subject: [PATCH 485/690] Add missing parameter to docstring --- src/pykeen/nn/functional.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index eeb813442f..0c877568fe 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -138,11 +138,11 @@ def complex_interaction( return tensor_sum(*( factor * extended_einsum("bhd,brd,btd->bhrt", hh, rr, tt) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) @@ -646,6 +646,8 @@ def simple_interaction( The relation representations. :param t_inv: shape: (batch_size, num_tails, dim) The tail representations. + :param clamp: + Clamp the scores to the given range. :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. From e2d26666a6cb33f59fd3c41e7f12dc93a27fdd84 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:55:25 +0100 Subject: [PATCH 486/690] Improve docstring of _calculate_missing_shape_information --- src/pykeen/nn/modules.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index dd43501c90..5c9a7fb02d 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -488,13 +488,18 @@ def _calculate_missing_shape_information( """Automatically calculates missing dimensions for ConvE. :param embedding_dim: + The embedding dimension. :param input_channels: + The number of input channels for the convolution. :param width: + The width of the embedding "image". :param height: + The height of the embedding "image". :return: (input_channels, width, height), such that `embedding_dim = input_channels * width * height` - :raises: + + :raises ValueError: If no factorization could be found. """ # Store initial input for error message From 0fdbd83e7d087e39c82a0b0a8e0e51184fdbda68 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:56:23 +0100 Subject: [PATCH 487/690] Add raises to docstring --- src/pykeen/nn/modules.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 5c9a7fb02d..5d96fad809 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -323,6 +323,9 @@ def _forward_slicing_wrapper( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. + + :raises ValueError: + If slice_dim is invalid. """ if slice_size is None: scores = self(h=h, r=r, t=t) From baca36ed780dad5ae4e2df72d0e6d61209389900 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:57:09 +0100 Subject: [PATCH 488/690] Add raises to docstring --- src/pykeen/nn/modules.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 5d96fad809..8f5840d8da 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -168,6 +168,10 @@ def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: :return: The squeezed tensor. + + :raises ValueError: + If there are duplicates in dims (after normalizing the dimensions, i.e. resolving negative dimension + indices). """ # normalize dimensions dims = tuple(d if d >= 0 else len(x.shape) + d for d in dims) From 26c270bd96feb348859cb52a33bc1cae4ff1502a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 10:58:17 +0100 Subject: [PATCH 489/690] Fix docstring for _add_dim --- src/pykeen/nn/modules.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 8f5840d8da..b8b9994788 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -145,11 +145,13 @@ def _add_dim(*x: torch.FloatTensor, dim: int) -> Union[torch.FloatTensor, Tuple[ """ Add a dimension to tensors. - :param x: shape: (d1, ..., dk) - The tensor. + :param x: + The tensors. + :param dim: + The dimension to unsqueeze. - :return: shape: (1, d1, ..., dk) - The tensor with batch dimension. + :return: + The tensors with an additional 1-element dimension. """ out = [xx.unsqueeze(dim=dim) for xx in x] if len(x) == 1: From 99090eff400185ae3314edc0179b5e4ca234d3b6 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 20 Nov 2020 12:01:53 +0100 Subject: [PATCH 490/690] asaf --- src/pykeen/nn/functional.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 0c877568fe..43a3f9754b 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -138,11 +138,11 @@ def complex_interaction( return tensor_sum(*( factor * extended_einsum("bhd,brd,btd->bhrt", hh, rr, tt) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) From 35b690df2e6b26ce269532b4896c94c1739fab69 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Fri, 20 Nov 2020 12:23:50 +0100 Subject: [PATCH 491/690] Improve logging on tests that sometimes fail --- tests/test_interactions.py | 15 +++++++++++---- 1 file changed, 11 insertions(+), 4 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 7018e9b83a..b734854b20 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -129,7 +129,7 @@ def test_score_h_slicing(self): ) scores = self.instance.score_h(all_entities=h, r=r, t=t, slice_size=self.num_entities // 2 + 1) scores_no_slice = self.instance.score_h(all_entities=h, r=r, t=t, slice_size=None) - assert torch.allclose(scores, scores_no_slice), f'Differences: {scores - scores_no_slice}' + self._check_close_scores(scores=scores, scores_no_slice=scores_no_slice) def test_score_r(self): """Test score_r.""" @@ -158,7 +158,7 @@ def test_score_r_slicing(self): ) scores = self.instance.score_r(h=h, all_relations=r, t=t, slice_size=self.num_relations // 2 + 1) scores_no_slice = self.instance.score_r(h=h, all_relations=r, t=t, slice_size=None) - assert torch.allclose(scores, scores_no_slice), f'Differences: {scores - scores_no_slice}' + self._check_close_scores(scores=scores, scores_no_slice=scores_no_slice) def test_score_t(self): """Test score_t.""" @@ -181,7 +181,12 @@ def test_score_t_slicing(self): ) scores = self.instance.score_t(h=h, r=r, all_entities=t, slice_size=self.num_entities // 2 + 1) scores_no_slice = self.instance.score_t(h=h, r=r, all_entities=t, slice_size=None) - assert torch.allclose(scores, scores_no_slice), f'Differences: {scores - scores_no_slice}' + self._check_close_scores(scores=scores, scores_no_slice=scores_no_slice) + + def _check_close_scores(self, scores, scores_no_slice): + self.assertFalse(torch.isnan(scores).any(), msg=f'Normal scores had nan:\n\t{scores}') + self.assertFalse(torch.isnan(scores_no_slice).any(), msg=f'Slice scores had nan\n\t{scores}') + self.assertTrue(torch.allclose(scores, scores_no_slice), msg=f'Differences: {scores - scores_no_slice}') def _get_test_shapes(self) -> Collection[Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]]]: """Return a set of test shapes for (h, r, t).""" @@ -238,7 +243,9 @@ def test_forward(self): for m in self.instance.modules() ) if small_batch_size and has_batch_norm: - logger.warning(f"Skipping test for shapes {hs}, {rs}, {ts}") + logger.warning( + f"Skipping test for shapes {hs}, {rs}, {ts} becuase too small batch size for batch norm", + ) continue raise error From d6ba8411851c1fe64a022d51c621cd49eaa1812a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 13:11:54 +0100 Subject: [PATCH 492/690] Add reset_parameters to InteractionTests --- tests/test_interactions.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index b734854b20..336e4a4896 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -69,6 +69,10 @@ class InteractionTests(GenericTests[pykeen.nn.modules.Interaction]): shape_kwargs = dict() + def post_instantiation_hook(self) -> None: + """Initialize parameters.""" + self.instance.reset_parameters() + def _get_hrt( self, *shapes: Tuple[int, ...], From e0c6f33cfc3dd5caf962e0e2865e4823f09074b9 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 18:57:03 +0100 Subject: [PATCH 493/690] Use isfinite to also check for +/-inf --- tests/test_interactions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 336e4a4896..1ca8d24f44 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -188,8 +188,8 @@ def test_score_t_slicing(self): self._check_close_scores(scores=scores, scores_no_slice=scores_no_slice) def _check_close_scores(self, scores, scores_no_slice): - self.assertFalse(torch.isnan(scores).any(), msg=f'Normal scores had nan:\n\t{scores}') - self.assertFalse(torch.isnan(scores_no_slice).any(), msg=f'Slice scores had nan\n\t{scores}') + self.assertTrue(torch.isfinite(scores).all(), msg=f'Normal scores had nan:\n\t{scores}') + self.assertTrue(torch.isfinite(scores_no_slice).all(), msg=f'Slice scores had nan\n\t{scores}') self.assertTrue(torch.allclose(scores, scores_no_slice), msg=f'Differences: {scores - scores_no_slice}') def _get_test_shapes(self) -> Collection[Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]]]: From c1b619437ed0e52c4f57f308d9f27e3961d7f6a9 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Fri, 20 Nov 2020 18:57:24 +0100 Subject: [PATCH 494/690] Fix typo --- tests/test_interactions.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 1ca8d24f44..ed89727a60 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -248,7 +248,7 @@ def test_forward(self): ) if small_batch_size and has_batch_norm: logger.warning( - f"Skipping test for shapes {hs}, {rs}, {ts} becuase too small batch size for batch norm", + f"Skipping test for shapes {hs}, {rs}, {ts} because too small batch size for batch norm", ) continue raise error From 2e5549bf181115b3cb4cc3600aa8b8612aa039f5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sun, 22 Nov 2020 10:44:46 +0100 Subject: [PATCH 495/690] Add shape to RepresentationModule --- src/pykeen/models/unimodal/rgcn.py | 2 +- src/pykeen/nn/emb.py | 16 +++++++++++----- src/pykeen/testing/mocks.py | 5 ++--- 3 files changed, 14 insertions(+), 9 deletions(-) diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index c58767258e..e89adaf1d6 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -131,7 +131,7 @@ def __init__( buffer_messages: bool = True, base_representations: Optional[RepresentationModule] = None, ): - super().__init__() + super().__init__(shape=(self.embedding_dim,)) self.triples_factory = triples_factory diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index a719471f2b..2b01dde7ac 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -29,6 +29,13 @@ class RepresentationModule(nn.Module, ABC): """A base class for obtaining representations for entities/relations.""" + #: The shape of a single representation + shape: Sequence[int] + + def __init__(self, shape: Sequence[int]): + super().__init__() + self.shape = shape + @abstractmethod def forward(self, indices: Optional[torch.LongTensor] = None) -> torch.FloatTensor: """Get representations for indices. @@ -121,7 +128,6 @@ class Embedding(RepresentationModule): can be used throughout PyKEEN as a more fully featured drop-in replacement. """ - shape: Sequence[int] normalizer: Optional[Normalizer] constrainer: Optional[Constrainer] regularizer: Optional[Regularizer] @@ -167,18 +173,18 @@ def __init__( :raises ValueError: If neither shape nor embedding_dim are given. """ - super().__init__() if shape is None and embedding_dim is None: raise ValueError('Missing both shape and embedding_dim') elif shape is not None: if isinstance(shape, int): - self.shape = (shape,) + shape = (shape,) else: - self.shape = shape + shape = shape embedding_dim = numpy.prod(shape) else: assert embedding_dim is not None - self.shape = (embedding_dim,) + shape = (embedding_dim,) + super().__init__(shape=shape) if initializer is None: initializer = nn.init.normal_ diff --git a/src/pykeen/testing/mocks.py b/src/pykeen/testing/mocks.py index ec10b697ea..d190c7f0ef 100644 --- a/src/pykeen/testing/mocks.py +++ b/src/pykeen/testing/mocks.py @@ -58,10 +58,9 @@ class MockRepresentations(RepresentationModule): """A custom representation module with minimal implementation.""" def __init__(self, num_entities: int, shape: Sequence[int]): - super().__init__() + super().__init__(shape=shape) self.num_embeddings = num_entities - self.shape = shape - self.x = nn.Parameter(torch.rand(numpy.prod(shape))) + self.x = nn.Parameter(torch.rand(numpy.prod(self.shape))) def forward(self, indices: Optional[torch.LongTensor] = None) -> torch.FloatTensor: # noqa: D102 n = self.num_embeddings if indices is None else indices.shape[0] From c33e1e59cc86f8f58f7a9d94a3327f70501e9119 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Sun, 22 Nov 2020 10:52:46 +0100 Subject: [PATCH 496/690] Add max_id to RepresentationModule --- src/pykeen/models/unimodal/rgcn.py | 2 +- src/pykeen/nn/emb.py | 10 +++++++--- src/pykeen/testing/mocks.py | 7 +++---- 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index e89adaf1d6..bc1aadd93d 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -131,7 +131,7 @@ def __init__( buffer_messages: bool = True, base_representations: Optional[RepresentationModule] = None, ): - super().__init__(shape=(self.embedding_dim,)) + super().__init__(shape=(self.embedding_dim,), max_id=triples_factory.num_entities) self.triples_factory = triples_factory diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index 2b01dde7ac..6f4d21bae4 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -32,9 +32,13 @@ class RepresentationModule(nn.Module, ABC): #: The shape of a single representation shape: Sequence[int] - def __init__(self, shape: Sequence[int]): + #: The maximum admissible ID (excl.) + max_id: int + + def __init__(self, shape: Sequence[int], max_id: int): super().__init__() self.shape = shape + self.max_id = max_id @abstractmethod def forward(self, indices: Optional[torch.LongTensor] = None) -> torch.FloatTensor: @@ -184,7 +188,7 @@ def __init__( else: assert embedding_dim is not None shape = (embedding_dim,) - super().__init__(shape=shape) + super().__init__(shape=shape, max_id=num_embeddings) if initializer is None: initializer = nn.init.normal_ @@ -241,7 +245,7 @@ def from_specification( @property def num_embeddings(self) -> int: # noqa: D401 """The total number of representations (i.e. the maximum ID).""" - return self._embeddings.num_embeddings + return self.max_id @property def embedding_dim(self) -> int: # noqa: D401 diff --git a/src/pykeen/testing/mocks.py b/src/pykeen/testing/mocks.py index d190c7f0ef..07554538a0 100644 --- a/src/pykeen/testing/mocks.py +++ b/src/pykeen/testing/mocks.py @@ -58,12 +58,11 @@ class MockRepresentations(RepresentationModule): """A custom representation module with minimal implementation.""" def __init__(self, num_entities: int, shape: Sequence[int]): - super().__init__(shape=shape) - self.num_embeddings = num_entities - self.x = nn.Parameter(torch.rand(numpy.prod(self.shape))) + super().__init__(shape=shape, max_id=num_entities) + self.x = nn.Parameter(torch.rand(int(numpy.prod(self.shape)))) def forward(self, indices: Optional[torch.LongTensor] = None) -> torch.FloatTensor: # noqa: D102 - n = self.num_embeddings if indices is None else indices.shape[0] + n = self.max_id if indices is None else indices.shape[0] return self.x.unsqueeze(dim=0).repeat(n, 1).view(-1, *self.shape) def get_in_canonical_shape( From 3a4af878d51cf46cc94dcb381614c8841cdd263a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 12:09:54 +0100 Subject: [PATCH 497/690] Simplify LCWA training loop --- src/pykeen/training/lcwa.py | 10 ++-------- 1 file changed, 2 insertions(+), 8 deletions(-) diff --git a/src/pykeen/training/lcwa.py b/src/pykeen/training/lcwa.py index 2ae50efb3b..a8b7604c83 100644 --- a/src/pykeen/training/lcwa.py +++ b/src/pykeen/training/lcwa.py @@ -43,19 +43,13 @@ def _process_batch( # Send batch to device batch_pairs = batch_pairs[start:stop].to(device=self.device) + predictions = self.model.score_t(hr_batch=batch_pairs) batch_labels_full = batch_labels_full[start:stop].to(device=self.device) - - if slice_size is None: - predictions = self.model.score_t(hr_batch=batch_pairs) - else: - predictions = self.model.score_t(hr_batch=batch_pairs, slice_size=slice_size) - - loss = self._loss_helper( + return self._loss_helper( predictions, batch_labels_full, label_smoothing, ) - return loss def _label_loss_helper( self, From 28888b302532c26c87cdfe0f1ffd7b4d657c81d1 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 12:10:04 +0100 Subject: [PATCH 498/690] Fix simplification --- src/pykeen/training/lcwa.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/training/lcwa.py b/src/pykeen/training/lcwa.py index a8b7604c83..addd1faafa 100644 --- a/src/pykeen/training/lcwa.py +++ b/src/pykeen/training/lcwa.py @@ -43,7 +43,7 @@ def _process_batch( # Send batch to device batch_pairs = batch_pairs[start:stop].to(device=self.device) - predictions = self.model.score_t(hr_batch=batch_pairs) + predictions = self.model.score_t(hr_batch=batch_pairs, slice_size=slice_size) batch_labels_full = batch_labels_full[start:stop].to(device=self.device) return self._loss_helper( predictions, From 88d7c308473f47256cf62dff5a13edbb105a7a90 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 13:04:32 +0100 Subject: [PATCH 499/690] Let training instance creation respect use_tqdm --- src/pykeen/training/training_loop.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/training/training_loop.py b/src/pykeen/training/training_loop.py index 5bd89dd2a1..b882260ee8 100644 --- a/src/pykeen/training/training_loop.py +++ b/src/pykeen/training/training_loop.py @@ -174,7 +174,7 @@ def train( """ # Create training instances # During size probing the training instances should not show the tqdm progress bar - self.training_instances = self._create_instances(use_tqdm=not only_size_probing) + self.training_instances = self._create_instances(use_tqdm=use_tqdm and not only_size_probing) # In some cases, e.g. using Optuna for HPO, the cuda cache from a previous run is not cleared torch.cuda.empty_cache() From 02f8412ad0b6a6fc93824f860aad59089240fcbc Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 14:52:09 +0100 Subject: [PATCH 500/690] More explicit viewing --- src/pykeen/nn/functional.py | 13 ++++++------- 1 file changed, 6 insertions(+), 7 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 43a3f9754b..dc11df0980 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -184,7 +184,6 @@ def conve_interaction( The scores. """ # bind sizes - # batch_size = max(x.shape[0] for x in (h, r, t)) num_heads = h.shape[1] num_relations = r.shape[1] num_tails = t.shape[1] @@ -192,11 +191,11 @@ def conve_interaction( # repeat if necessary, and concat head and relation, batch_size', num_input_channels, 2*height, width # with batch_size' = batch_size * num_heads * num_relations - h = h.unsqueeze(dim=2) - h = h.view(*h.shape[:-1], input_channels, embedding_height, embedding_width) - r = r.unsqueeze(dim=1) - r = r.view(*r.shape[:-1], input_channels, embedding_height, embedding_width) - x = broadcast_cat(h, r, dim=-2).view(-1, input_channels, 2 * embedding_height, embedding_width) + x = broadcast_cat( + h.view(h.shape[0], num_heads, 1, input_channels, embedding_height, embedding_width), + r.view(r.shape[0], 1, num_relations, input_channels, embedding_height, embedding_width), + dim=-2, + ).view(-1, input_channels, 2 * embedding_height, embedding_width) # batch_size', num_input_channels, 2*height, width x = hr2d(x) @@ -206,7 +205,7 @@ def conve_interaction( x = hr1d(x) # reshape: (batch_size', embedding_dim) -> (b, h, r, 1, d) - x = x.view(-1, num_heads, num_relations, 1, embedding_dim) + x = x.view(x.shape[0], num_heads, num_relations, 1, embedding_dim) # For efficient calculation, each of the convolved [h, r] rows has only to be multiplied with one t row # output_shape: (batch_size, num_heads, num_relations, num_tails) From a4d3a9f0c8ef1f762f4946faee3cfdbe4b52e5ac Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 23 Nov 2020 15:51:02 +0100 Subject: [PATCH 501/690] Add mypy testing --- .github/workflows/tests.yml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index dc54c3e911..049f22e7c4 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -25,6 +25,9 @@ jobs: - name: Check code quality with flake8 run: tox -e flake8 + - name: Check typing with MyPy + env: tox -e mypy + - name: Check package metadata with Pyroma run: tox -e pyroma # - name: Check static typing with MyPy From 5b042b4d35cfbef84539c7fd8cc60696823d26d1 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 19:49:56 +0100 Subject: [PATCH 502/690] Modify all interaction functions to use the broadcasted format --- src/pykeen/nn/functional.py | 328 +++++++++++++++++------------------- src/pykeen/nn/sim.py | 29 ++-- src/pykeen/utils.py | 11 +- 3 files changed, 176 insertions(+), 192 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index dc11df0980..4b1989196a 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -13,7 +13,7 @@ from ..typing import GaussianDistribution from ..utils import ( broadcast_cat, clamp_norm, extended_einsum, is_cudnn_error, negative_norm_of_sum, project_entity, - split_complex, tensor_sum, view_complex, + split_complex, tensor_product, tensor_sum, view_complex, ) __all__ = [ @@ -46,7 +46,7 @@ def _extract_sizes( t: torch.Tensor, ) -> Tuple[int, int, int, int, int]: """Extract size dimensions from head/relation/tail representations.""" - num_heads, num_relations, num_tails = [xx.shape[1] for xx in (h, r, t)] + num_heads, num_relations, num_tails = [xx.shape[i] for i, xx in enumerate((h, r, t), start=1)] d_e = h.shape[-1] d_r = r.shape[-1] return num_heads, num_relations, num_tails, d_e, d_r @@ -121,11 +121,11 @@ def complex_interaction( .. math :: Re(\langle h, r, conj(t) \rangle) - :param h: shape: (batch_size, num_heads, `2*dim`) + :param h: shape: (batch_size, num_heads, 1, 1, `2*dim`) The complex head representations. - :param r: shape: (batch_size, num_relations, 2*dim) + :param r: shape: (batch_size, 1, num_relations, 1, 2*dim) The complex relation representations. - :param t: shape: (batch_size, num_tails, 2*dim) + :param t: shape: (batch_size, 1, 1, num_tails, 2*dim) The complex tail representations. :return: shape: (batch_size, num_heads, num_relations, num_tails) @@ -136,13 +136,13 @@ def complex_interaction( # = + + - (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] return tensor_sum(*( - factor * extended_einsum("bhd,brd,btd->bhrt", hh, rr, tt) + factor * tensor_product(hh, rr, tt).sum(dim=-1) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) @@ -161,13 +161,13 @@ def conve_interaction( """ Evaluate the ConvE interaction function. - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param r: shape: (batch_size, num_relations, dim) + :param r: shape: (batch_size, 1, num_relations, 1, dim) The relation representations. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. - :param t_bias: shape: (batch_size, num_tails, dim) + :param t_bias: shape: (batch_size, 1, 1, num_tails, 1) The tail entity bias. :param input_channels: The number of input channels. @@ -185,15 +185,15 @@ def conve_interaction( """ # bind sizes num_heads = h.shape[1] - num_relations = r.shape[1] - num_tails = t.shape[1] + num_relations = r.shape[2] + num_tails = t.shape[3] embedding_dim = h.shape[-1] # repeat if necessary, and concat head and relation, batch_size', num_input_channels, 2*height, width # with batch_size' = batch_size * num_heads * num_relations x = broadcast_cat( - h.view(h.shape[0], num_heads, 1, input_channels, embedding_height, embedding_width), - r.view(r.shape[0], 1, num_relations, input_channels, embedding_height, embedding_width), + h.view(*h.shape[:-1], input_channels, embedding_height, embedding_width), + r.view(*r.shape[:-1], input_channels, embedding_height, embedding_width), dim=-2, ).view(-1, input_channels, 2 * embedding_height, embedding_width) @@ -209,11 +209,11 @@ def conve_interaction( # For efficient calculation, each of the convolved [h, r] rows has only to be multiplied with one t row # output_shape: (batch_size, num_heads, num_relations, num_tails) - t = t.view(t.shape[0], 1, 1, num_tails, embedding_dim).transpose(-1, -2) + t = t.transpose(-1, -2) x = (x @ t).squeeze(dim=-2) # add bias term - return tensor_sum(x, t_bias.view(t.shape[0], 1, 1, num_tails)) + return tensor_sum(x, t_bias.squeeze(dim=-1)) def convkb_interaction( @@ -231,11 +231,11 @@ def convkb_interaction( .. math:: W_L drop(act(W_C \ast ([h; r; t]) + b_C)) + b_L - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param r: shape: (batch_size, num_relations, dim) + :param r: shape: (batch_size, 1, num_relations, 1, dim) The relation representations. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. :param conv: The 3x1 convolution. @@ -258,9 +258,9 @@ def convkb_interaction( # compute conv(stack(h, r, t)) # prepare input shapes for broadcasting - h = h.view(h.shape[0], num_heads, 1, 1, 1, embedding_dim) - r = r.view(r.shape[0], 1, num_relations, 1, 1, embedding_dim) - t = t.view(t.shape[0], 1, 1, num_tails, 1, embedding_dim) + h = h.unsqueeze(dim=-2) + r = r.unsqueeze(dim=-2) + t = t.unsqueeze(dim=-2) # conv.weight.shape = (C_out, C_in, kernel_size[0], kernel_size[1]) # here, kernel_size = (1, 3), C_in = 1, C_out = num_filters @@ -284,7 +284,7 @@ def convkb_interaction( x = hidden_dropout(x) # Linear layer for final scores; use flattened representations, shape: (b, h, r, t, d * f) - x = x.view(x.shape[0], num_heads, num_relations, num_tails, embedding_dim * num_filters) + x = x.view(x.shape[:-2], embedding_dim * num_filters) x = linear(x) return x.squeeze(dim=-1) @@ -297,17 +297,17 @@ def distmult_interaction( """ Evaluate the DistMult interaction function. - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param r: shape: (batch_size, num_relations, dim) + :param r: shape: (batch_size, 1, num_relations, 1, dim) The relation representations. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return extended_einsum("bhd,brd,btd->bhrt", h, r, t) + return tensor_product(h, r, t) def ermlp_interaction( @@ -321,11 +321,11 @@ def ermlp_interaction( r""" Evaluate the ER-MLP interaction function. - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param r: shape: (batch_size, num_relations, dim) + :param r: shape: (batch_size, 1, num_relations, 1, dim) The relation representations. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. :param hidden: The first linear layer. @@ -342,9 +342,9 @@ def ermlp_interaction( # split, shape: (embedding_dim, hidden_dim) head_to_hidden, rel_to_hidden, tail_to_hidden = hidden.weight.t().split(embedding_dim) bias = hidden.bias.view(1, 1, 1, 1, -1) - h = h.view(-1, num_heads, 1, 1, embedding_dim) @ head_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) - r = r.view(-1, 1, num_relations, 1, embedding_dim) @ rel_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) - t = t.view(-1, 1, 1, num_tails, embedding_dim) @ tail_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) + h = h @ head_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) + r = r @ rel_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) + t = t @ tail_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) return final(activation(tensor_sum(bias, h, r, t))).squeeze(dim=-1) @@ -357,11 +357,11 @@ def ermlpe_interaction( r""" Evaluate the ER-MLPE interaction function. - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param r: shape: (batch_size, num_relations, dim) + :param r: shape: (batch_size, 1, num_relations, 1, dim) The relation representations. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. :param mlp: The MLP. @@ -369,19 +369,18 @@ def ermlpe_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # repeat if necessary, and concat head and relation, (batch_size, num_heads, num_relations, 2 * embedding_dim) - x = broadcast_cat(h.unsqueeze(dim=2), r.unsqueeze(dim=1), dim=-1) + # repeat if necessary, and concat head and relation, (batch_size, num_heads, num_relations, 1, 2 * embedding_dim) + x = broadcast_cat(h, r, dim=-1) # Predict t embedding, shape: (b, h, r, 1, d) shape = x.shape - x = mlp(x.view(-1, shape[-1])).view(*shape[:-1], -1).unsqueeze(dim=-2) + x = mlp(x.view(-1, shape[-1])).view(*shape[:-1], -1) # transpose t, (b, 1, 1, d, t) - t = t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]).transpose(-2, -1) + t = t.transpose(-2, -1) # dot product, (b, h, r, 1, t) - x = (x @ t).squeeze(dim=-2) - return x + return (x @ t).squeeze(dim=-2) def hole_interaction( @@ -392,11 +391,11 @@ def hole_interaction( """ Evaluate the HolE interaction function. - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param r: shape: (batch_size, num_relations, dim) + :param r: shape: (batch_size, 1, num_relations, 1, dim) The relation representations. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. :return: shape: (batch_size, num_heads, num_relations, num_tails) @@ -406,17 +405,20 @@ def hole_interaction( a_fft = torch.fft.rfft(h, dim=-1) b_fft = torch.fft.rfft(t, dim=-1) - # complex conjugate, shape = (b, h, d) + # complex conjugate a_fft = torch.conj(a_fft) - # Hadamard product in frequency domain, shape: (b, h, t, d) - p_fft = a_fft.unsqueeze(dim=2) * b_fft.unsqueeze(dim=1) + # Hadamard product in frequency domain + p_fft = a_fft * b_fft - # inverse real FFT, shape: (b, h, t, d) + # inverse real FFT, shape: (b, h, 1, t, d) composite = torch.fft.irfft(p_fft, n=h.shape[-1], dim=-1) + # transpose r, (b, 1, r, 1, d) -> (b, 1, 1, d, r) + r = r.permute(0, 1, 3, 4, 2) + # inner product with relation embedding - return extended_einsum("bhtd,brd->bhrt", composite, r) + return (composite @ r).squeeze(dim=-2) def kg2e_interaction( @@ -432,17 +434,17 @@ def kg2e_interaction( """ Evaluate the KG2E interaction function. - :param h_mean: shape: (batch_size, num_heads, d) + :param h_mean: shape: (batch_size, num_heads, 1, 1, d) The head entity distribution mean. - :param h_var: shape: (batch_size, num_heads, d) + :param h_var: shape: (batch_size, num_heads, 1, 1, d) The head entity distribution variance. - :param r_mean: shape: (batch_size, num_relations, d) + :param r_mean: shape: (batch_size, 1, num_relations, 1, d) The relation distribution mean. - :param r_var: shape: (batch_size, num_relations, d) + :param r_var: shape: (batch_size, 1, num_relations, 1, d) The relation distribution variance. - :param t_mean: shape: (batch_size, num_tails, d) + :param t_mean: shape: (batch_size, 1, 1, num_tails, d) The tail entity distribution mean. - :param t_var: shape: (batch_size, num_tails, d) + :param t_var: shape: (batch_size, 1, 1, num_tails, d) The tail entity distribution variance. :param similarity: The similarity measures for gaussian distributions. From {"KL", "EL"}. @@ -454,8 +456,8 @@ def kg2e_interaction( """ similarity_fn = KG2E_SIMILARITIES[similarity] # Compute entity distribution - e_mean = h_mean.unsqueeze(dim=2) - t_mean.unsqueeze(dim=1) - e_var = h_var.unsqueeze(dim=2) + t_var.unsqueeze(dim=1) + e_mean = h_mean - t_mean + e_var = h_var + t_var e = GaussianDistribution(mean=e_mean, diagonal_covariance=e_var) r = GaussianDistribution(mean=r_mean, diagonal_covariance=r_var) return similarity_fn(e=e, r=r, exact=exact) @@ -478,19 +480,19 @@ def ntn_interaction( f(h,r,t) = u_r^T act(h W_r t + V_r h + V_r' t + b_r) - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param w: shape: (batch_size, num_relations, k, dim, dim) + :param w: shape: (batch_size, 1, num_relations, 1, k, dim, dim) The relation specific transformation matrix W_r. - :param vh: shape: (batch_size, num_relations, k, dim) + :param vh: shape: (batch_size, 1, num_relations, 1, k, dim) The head transformation matrix V_h. - :param vt: shape: (batch_size, num_relations, k, dim) + :param vt: shape: (batch_size, 1, num_relations, 1, k, dim) The tail transformation matrix V_h. - :param b: shape: (batch_size, num_relations, k) + :param b: shape: (batch_size, 1, num_relations, 1, k) The relation specific offset b_r. - :param u: shape: (batch_size, num_relations, k) + :param u: shape: (batch_size, 1, num_relations, 1, k) The relation specific final linear transformation b_r. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. :param activation: The activation function. @@ -501,12 +503,12 @@ def ntn_interaction( # save sizes num_heads, num_relations, num_tails, _, num_slices = _extract_sizes(h, b, t) x = activation(tensor_sum( - extended_einsum("bhd,brkde,bte->bhrtk", h, w, t), - extended_einsum("brkd,bhd->bhrk", vh, h).unsqueeze(dim=3), - extended_einsum("brkd,btd->brtk", vt, t).unsqueeze(dim=1), - b.view(b.shape[0], 1, num_relations, 1, num_slices), + extended_einsum("bhrtd,brkde,bhrte->bhrtk", h, w, t), + extended_einsum("bhrtkd,bhrtd->bhrtk", vh, h), + extended_einsum("bhrtkd,bhrtd->bhrtk", vt, t), + b, )) - x = extended_einsum("bhrtk,brk->bhrt", x, u) + x = extended_einsum("bhrtk,bhrtk->bhrt", x, u) return x @@ -527,11 +529,11 @@ def proje_interaction( f(h, r, t) = g(t z(D_e h + D_r r + b_c) + b_p) - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param r: shape: (batch_size, num_relations, dim) + :param r: shape: (batch_size, 1, num_relations, 1, dim) The relation representations. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. :param d_e: shape: (dim,) Global entity projection. @@ -549,17 +551,14 @@ def proje_interaction( """ num_heads, num_relations, num_tails, dim, _ = _extract_sizes(h, r, t) # global projections - h = h * d_e.view(1, 1, dim) - r = r * d_r.view(1, 1, dim) + h = h * d_e.view(1, 1, 1, 1, dim) + r = r * d_r.view(1, 1, 1, 1, dim) # combination, shape: (b, h, r, d) - x = tensor_sum( - h.unsqueeze(dim=2), - r.unsqueeze(dim=1), - b_c.view(1, 1, 1, dim), - ) + x = tensor_sum(h, r, b_c) x = activation(x) # dot product with t, shape: (b, h, r, t) - return (x @ t.unsqueeze(dim=1).transpose(-2, -1)) + b_p + t = t.transpose(-2, -1) + return (x @ t) + b_p def rescal_interaction( @@ -570,17 +569,21 @@ def rescal_interaction( """ Evaluate the RESCAL interaction function. - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param r: shape: (batch_size, num_relations, dim, dim) + :param r: shape: (batch_size, 1, num_relations, 1, dim, dim) The relation representations. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return extended_einsum("bhd,brde,bte->bhrt", h, r, t) + return torch.chain_matmul( + h.unsqueeze(dim=-2), + r, + t.unsqueeze(dim=-1), + ).view(-1, h.shape[1], r.shape[2], t.shape[3]) def rotate_interaction( @@ -590,11 +593,11 @@ def rotate_interaction( ) -> torch.FloatTensor: """Evaluate the interaction function of RotatE for given embeddings. - :param h: shape: (batch_size, num_heads, 2*dim) + :param h: shape: (batch_size, num_heads, 1, 1, 2*dim) The head representations. - :param r: shape: (batch_size, num_relations, 2*dim) + :param r: shape: (batch_size, 1, num_relations, 1, 2*dim) The relation representations. - :param t: shape: (batch_size, num_tails, 2*dim) + :param t: shape: (batch_size, 1, 1, num_tails, 2*dim) The tail representations. :return: shape: (batch_size, num_heads, num_relations, num_tails) @@ -610,12 +613,12 @@ def rotate_interaction( h, r, t = [view_complex(x) for x in (h, r, t)] # Rotate (=Hadamard product in complex space). - hr = extended_einsum("bhd,brd->bhrd", h, r) + hr = h * r # Workaround until https://github.com/pytorch/pytorch/issues/30704 is fixed return negative_norm_of_sum( - hr.unsqueeze(dim=3), - -t.view(t.shape[0], 1, 1, t.shape[1], t.shape[2]), + hr, + -t, p=2, power_norm=False, ) @@ -633,17 +636,17 @@ def simple_interaction( """ Evaluate the SimplE interaction function. - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param r: shape: (batch_size, num_relations, dim, dim) + :param r: shape: (batch_size, 1, num_relations, 1, dim, dim) The relation representations. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. - :param h_inv: shape: (batch_size, num_heads, dim) + :param h_inv: shape: (batch_size, num_heads, 1, 1, dim) The inverse head representations. - :param r_inv: shape: (batch_size, num_relations, dim, dim) + :param r_inv: shape: (batch_size, 1, num_relations, 1, dim, dim) The relation representations. - :param t_inv: shape: (batch_size, num_tails, dim) + :param t_inv: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. :param clamp: Clamp the scores to the given range. @@ -675,13 +678,13 @@ def structured_embedding_interaction( .. math :: f(h, r, t) = -\|R_h h - R_t t\| - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param r_h: shape: (batch_size, num_relations, dim, rel_dim) + :param r_h: shape: (batch_size, 1, num_relations, 1, dim, rel_dim) The relation-specific head projection. - :param r_t: shape: (batch_size, num_relations, dim, rel_dim) + :param r_t: shape: (batch_size, 1, num_relations, 1, dim, rel_dim) The relation-specific tail projection. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. :param p: The p for the norm. cf. torch.norm. @@ -692,8 +695,8 @@ def structured_embedding_interaction( The scores. """ return negative_norm_of_sum( - extended_einsum("bred,bhd->bhre", r_h, h).unsqueeze(dim=3), - -extended_einsum("bred,btd->brte", r_t, t).unsqueeze(dim=1), + (h @ r_h), + -(t @ r_t), p=p, power_norm=power_norm, ) @@ -712,17 +715,17 @@ def transd_interaction( """ Evaluate the TransD interaction function. - :param h: shape: (batch_size, num_heads, d_e) + :param h: shape: (batch_size, num_heads, 1, 1, d_e) The head representations. - :param r: shape: (batch_size, num_relations, d_r) + :param r: shape: (batch_size, 1, num_relations, 1, d_r) The relation representations. - :param t: shape: (batch_size, num_tails, d_e) + :param t: shape: (batch_size, 1, 1, num_tails, d_e) The tail representations. - :param h_p: shape: (batch_size, num_heads, d_e) + :param h_p: shape: (batch_size, num_heads, 1, 1, d_e) The head projections. - :param r_p: shape: (batch_size, num_relations, d_r) + :param r_p: shape: (batch_size, 1, num_relations, 1, d_r) The relation projections. - :param t_p: shape: (batch_size, num_tails, d_e) + :param t_p: shape: (batch_size, 1, 1, num_tails, d_e) The tail projections. :param p: The parameter p for selecting the norm. @@ -733,19 +736,16 @@ def transd_interaction( The scores. """ # Project entities - # shape: (b, h, r, 1, d_r) h_bot = project_entity( - e=h.unsqueeze(dim=2), - e_p=h_p.unsqueeze(dim=2), - r_p=r_p.unsqueeze(dim=1), - ).unsqueeze(dim=-2) - # shape: (b, 1, r, t, d_r) + e=h, + e_p=h_p, + r_p=r_p, + ) t_bot = project_entity( - e=t.unsqueeze(dim=1), - e_p=t_p.unsqueeze(dim=1), - r_p=r_p.unsqueeze(dim=2), - ).unsqueeze(dim=1) - r = r.view(r.shape[0], 1, r.shape[1], 1, r.shape[2]) + e=t, + e_p=t_p, + r_p=r_p, + ) return _translational_interaction(h=h_bot, r=r, t=t_bot, p=p, power_norm=power_norm) @@ -759,11 +759,11 @@ def transe_interaction( """ Evaluate the TransE interaction function. - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param r: shape: (batch_size, num_relations, dim) + :param r: shape: (batch_size, 1, num_relations, 1, dim) The relation representations. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. :param p: The p for the norm. @@ -773,14 +773,7 @@ def transe_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - num_heads, num_relations, num_tails, embedding_dim, _ = _extract_sizes(h, r, t) - return _translational_interaction( - h=h.view(-1, num_heads, 1, 1, embedding_dim), - r=r.view(-1, 1, num_relations, 1, embedding_dim), - t=t.view(-1, 1, 1, num_tails, embedding_dim), - p=p, - power_norm=power_norm, - ) + return _translational_interaction(h=h, r=r, t=t, p=p, power_norm=power_norm) def transh_interaction( @@ -794,13 +787,13 @@ def transh_interaction( """ Evaluate the DistMult interaction function. - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param w_r: shape: (batch_size, num_relations, dim) + :param w_r: shape: (batch_size, 1, num_relations, 1, dim) The relation normal vector representations. - :param d_r: shape: (batch_size, num_relations, dim) + :param d_r: shape: (batch_size, 1, num_relations, 1, dim) The relation difference vector representations. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. :param p: The p for the norm. cf. torch.norm. @@ -810,17 +803,15 @@ def transh_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - num_heads, num_relations, num_tails, dim = _extract_sizes(h, w_r, t)[:4] - # Project to hyperplane return negative_norm_of_sum( - # h - h.view(h.shape[0], num_heads, 1, 1, dim), - -extended_einsum("bhd,brd,bre->bhre", h, w_r, w_r).unsqueeze(dim=3), + # h projection to hyperplane + h, + -(h * w_r).sum(dim=-1, keepdims=True) * w_r, # r - d_r.view(d_r.shape[0], 1, num_relations, 1, dim), - # -t - -t.view(t.shape[0], 1, 1, num_tails, dim), - extended_einsum("btd,brd,bre->brte", t, w_r, w_r).unsqueeze(dim=1), + d_r, + # -t projection to hyperplane + -t, + (t * w_r).sum(dim=-1, keepdims=True) * w_r, p=p, power_norm=power_norm, ) @@ -836,13 +827,13 @@ def transr_interaction( ) -> torch.FloatTensor: """Evaluate the interaction function for given embeddings. - :param h: shape: (batch_size, num_heads, d_e) + :param h: shape: (batch_size, num_heads, 1, 1, d_e) Head embeddings. - :param r: shape: (batch_size, num_relations, d_r) + :param r: shape: (batch_size, 1, num_relations, 1, d_r) Relation embeddings. - :param m_r: shape: (batch_size, num_relations, d_e, d_r) + :param m_r: shape: (batch_size, 1, num_relations, 1, d_e, d_r) The relation specific linear transformations. - :param t: shape: (batch_size, num_tails, d_e) + :param t: shape: (batch_size, 1, 1, num_tails, d_e) Tail embeddings. :param p: The parameter p for selecting the norm. @@ -852,18 +843,9 @@ def transr_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - num_heads, num_relations, num_tails, d_e, d_r = _extract_sizes(h=h, r=r, t=t) # project to relation specific subspace and ensure constraints - # head, shape: (b, h, r, 1, d_r) - h_bot = h.view(-1, num_heads, 1, 1, d_e) @ m_r.view(-1, 1, num_relations, d_e, d_r) - h_bot = clamp_norm(h_bot, p=2, dim=-1, maxnorm=1.) - - # head, shape: (b, 1, r, t, d_r) - t_bot = t.view(-1, 1, 1, num_tails, d_e) @ m_r.view(-1, 1, num_relations, d_e, d_r) - t_bot = clamp_norm(t_bot, p=2, dim=-1, maxnorm=1.) - - # evaluate score function, shape: (b, h, r, t) - r = r.view(-1, 1, num_relations, 1, d_r) + h_bot = clamp_norm(h @ m_r, p=2, dim=-1, maxnorm=1.) + t_bot = clamp_norm(t @ m_r, p=2, dim=-1, maxnorm=1.) return _translational_interaction(h=h_bot, r=r, t=t_bot, p=p, power_norm=power_norm) @@ -889,11 +871,11 @@ def tucker_interaction( where BN denotes BatchNorm and DO denotes Dropout - :param h: shape: (batch_size, num_heads, d_e) + :param h: shape: (batch_size, num_heads, 1, 1, d_e) The head representations. - :param r: shape: (batch_size, num_relations, d_r) + :param r: shape: (batch_size, 1, num_relations, 1, d_r) The relation representations. - :param t: shape: (batch_size, num_tails, d_e) + :param t: shape: (batch_size, 1, 1, num_tails, d_e) The tail representations. :param core_tensor: shape: (d_e, d_r, d_e) The core tensor. @@ -913,15 +895,15 @@ def tucker_interaction( """ return extended_einsum( # x_3 contraction - "bhrk,btk->bhrt", + "bhrtk,bhrtk->bhrt", _apply_optional_bn_to_tensor( x=extended_einsum( # x_1 contraction - "brik,bhi->bhrk", + "bhrtik,bhrti->bhrtk", _apply_optional_bn_to_tensor( x=extended_einsum( # x_2 contraction - "ijk,brj->brik", + "ijk,bhrtj->bhrtik", core_tensor, r, ), @@ -948,9 +930,9 @@ def unstructured_model_interaction( """ Evaluate the SimplE interaction function. - :param h: shape: (batch_size, num_heads, dim) + :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param t: shape: (batch_size, num_tails, dim) + :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. :param p: The parameter p for selecting the norm. @@ -960,6 +942,4 @@ def unstructured_model_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - h = h.unsqueeze(dim=2).unsqueeze(dim=3) - t = t.unsqueeze(dim=1).unsqueeze(dim=2) return negative_norm_of_sum(h, -t, p=p, power_norm=power_norm) diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index 07f0978d82..01c261fca3 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -36,23 +36,21 @@ def expected_likelihood( + \log \det \Sigma + d \log (2 \pi) \right) - :param e: shape: (batch_size, num_heads, num_tails, d) + :param e: shape: (batch_size, num_heads, 1, num_tails, d) The entity Gaussian distribution. - :param r: shape: (batch_size, num_relations, d) + :param r: shape: (batch_size, 1, num_relations, 1, d) The relation Gaussian distribution. :param epsilon: float (default=1.0) Small constant used to avoid numerical issues when dividing. :param exact: Whether to return the exact similarity, or leave out constant offsets. - :return: torch.Tensor, shape: (s_1, ..., s_k) + :return: torch.Tensor, shape: (batch_size, num_heads, num_relations, num_tails) The similarity. """ # subtract, shape: (batch_size, num_heads, num_relations, num_tails, dim) - r_shape = r.mean.shape - r_shape = (r_shape[0], 1, r_shape[1], 1, r_shape[2]) - var = r.diagonal_covariance.view(*r_shape) + e.diagonal_covariance.unsqueeze(dim=2) - mean = e.mean.unsqueeze(dim=2) - r.mean.view(*r_shape) + var = r.diagonal_covariance + e.diagonal_covariance + mean = e.mean - r.mean #: a = \mu^T\Sigma^{-1}\mu safe_sigma = torch.clamp_min(var, min=epsilon) @@ -91,9 +89,9 @@ def kullback_leibler_similarity( .. seealso :: https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Kullback%E2%80%93Leibler_divergence - :param e: shape: (batch_size, num_heads, num_tails, d) + :param e: shape: (batch_size, num_heads, 1, num_tails, d) The entity Gaussian distributions, as mean/diagonal covariance pairs. - :param r: shape: (batch_size, num_relations, d) + :param r: shape: (batch_size, 1, num_relations, 1, d) The relation Gaussian distributions, as mean/diagonal covariance pairs. :param epsilon: float (default=1.0) Small constant used to avoid numerical issues when dividing. @@ -105,14 +103,11 @@ def kullback_leibler_similarity( """ assert (e.diagonal_covariance > 0).all() and (r.diagonal_covariance > 0).all() - # broadcast shapes to (batch_size, num_heads, num_relations, num_tails, dim) - e_shape = e.mean.shape # (batch_size, num_heads, num_tails, dim) - e_mean = e.mean.view(e_shape[0], e_shape[1], 1, e_shape[2], e_shape[3]) - e_var = e.diagonal_covariance.view(e_shape[0], e_shape[1], 1, e_shape[2], e_shape[3]) + e_mean = e.mean + e_var = e.diagonal_covariance - r_shape = r.mean.shape # (batch_size, num_relations, dim) - r_mean = r.mean.view(r_shape[0], 1, r_shape[1], 1, r_shape[2]) - r_var: torch.FloatTensor = r.diagonal_covariance.view(r_shape[0], 1, r_shape[1], 1, r_shape[2]) + r_mean = r.mean + r_var = r.diagonal_covariance terms = [] @@ -134,7 +129,7 @@ def kullback_leibler_similarity( # 3. Component if exact: - terms.append(-e_shape[-1]) + terms.append(-e_mean.shape[-1]) # 4. Component # ln (det(\Sigma_1) / det(\Sigma_0)) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index cddcc95578..3d1292cb01 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -3,8 +3,10 @@ """Utilities for PyKEEN.""" import ftplib +import functools import json import logging +import operator import random from abc import ABC, abstractmethod from io import BytesIO @@ -498,11 +500,17 @@ def complex_normalize(x: torch.Tensor) -> torch.Tensor: def tensor_sum(*x: torch.FloatTensor) -> torch.FloatTensor: - """Compute sum of tensors in brodcastable shape.""" + """Compute elementwise sum of tensors in brodcastable shape.""" # TODO: Optimize order return sum(x) +def tensor_product(*x: torch.FloatTensor) -> torch.FloatTensor: + """"Compute elementwise product of tensors in broadcastable shape.""" + # TODO: Optimize order + return functools.reduce(operator.mul, x[1:], x[0]) + + def negative_norm_of_sum( *x: torch.FloatTensor, p: Union[str, int] = 2, @@ -545,6 +553,7 @@ def extended_einsum( mod_op = "" if len(op) != len(t.shape): raise ValueError(f'Shapes not equal: op={op} and t.shape={t.shape}') + # TODO: t_shape = list(t.shape); del t_shape[i]; t.view(*shape) -> only one reshape operation for i, c in reversed(list(enumerate(op))): if t.shape[i] == 1: t = t.squeeze(dim=i) From d4047cf7fb4a7262277a4b0d7d8c2f385ac4c1b5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 19:50:40 +0100 Subject: [PATCH 503/690] Inline nop function --- src/pykeen/nn/functional.py | 33 +++------------------------------ 1 file changed, 3 insertions(+), 30 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 4b1989196a..0d9d2e4381 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -66,33 +66,6 @@ def _apply_optional_bn_to_tensor( return output_dropout(x) -def _translational_interaction( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, - p: Union[int, str] = 2, - power_norm: bool = False, -) -> torch.FloatTensor: - """ - Evaluate a translational distance interaction function on already broadcasted representations. - - :param h: shape: (batch_size, num_heads, num_relations, num_tails, dim) - The head representations. - :param r: shape: (batch_size, num_heads, num_relations, num_tails, dim) - The relation representations. - :param t: shape: (batch_size, num_heads, num_relations, num_tails, dim) - The tail representations. - :param p: - The p for the norm. cf. torch.norm. - :param power_norm: - Whether to return $|x-y|_p^p$, cf. https://github.com/pytorch/pytorch/issues/28119 - - :return: shape: (batch_size, num_heads, num_relations, num_tails) - The scores. - """ - return negative_norm_of_sum(h, r, -t, p=p, power_norm=power_norm) - - def _add_cuda_warning(func): def wrapped(*args, **kwargs): try: @@ -746,7 +719,7 @@ def transd_interaction( e_p=t_p, r_p=r_p, ) - return _translational_interaction(h=h_bot, r=r, t=t_bot, p=p, power_norm=power_norm) + return negative_norm_of_sum(h_bot, r, -t_bot, p=p, power_norm=power_norm) def transe_interaction( @@ -773,7 +746,7 @@ def transe_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return _translational_interaction(h=h, r=r, t=t, p=p, power_norm=power_norm) + return negative_norm_of_sum(h, r, -t, p=p, power_norm=power_norm) def transh_interaction( @@ -846,7 +819,7 @@ def transr_interaction( # project to relation specific subspace and ensure constraints h_bot = clamp_norm(h @ m_r, p=2, dim=-1, maxnorm=1.) t_bot = clamp_norm(t @ m_r, p=2, dim=-1, maxnorm=1.) - return _translational_interaction(h=h_bot, r=r, t=t_bot, p=p, power_norm=power_norm) + return negative_norm_of_sum(h_bot, r, -t_bot, p=p, power_norm=power_norm) def tucker_interaction( From 71bc48a3970f8091efec1b74f3b1eb5e477b5951 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 19:57:14 +0100 Subject: [PATCH 504/690] Use matmul instead of einsum --- src/pykeen/nn/functional.py | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 0d9d2e4381..a266a160fb 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -473,16 +473,14 @@ def ntn_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # save sizes - num_heads, num_relations, num_tails, _, num_slices = _extract_sizes(h, b, t) x = activation(tensor_sum( extended_einsum("bhrtd,brkde,bhrte->bhrtk", h, w, t), - extended_einsum("bhrtkd,bhrtd->bhrtk", vh, h), - extended_einsum("bhrtkd,bhrtd->bhrtk", vt, t), + (vh @ h.unsqueeze(dim=-1)).squeeze(dim=-1), + (vt @ t.unsqueeze(dim=-1)).squeeze(dim=-1), b, )) - x = extended_einsum("bhrtk,bhrtk->bhrt", x, u) - return x + u = u.transpose(-2, -1) + return (x @ u).squeeze(dim=-2) def proje_interaction( From e68c959bcd55648bbb2bccc203c31ec3629d72dd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 20:05:02 +0100 Subject: [PATCH 505/690] rename file --- src/pykeen/nn/__init__.py | 2 +- src/pykeen/nn/modules.py | 8 ++++---- src/pykeen/nn/{emb.py => representation.py} | 0 3 files changed, 5 insertions(+), 5 deletions(-) rename src/pykeen/nn/{emb.py => representation.py} (100%) diff --git a/src/pykeen/nn/__init__.py b/src/pykeen/nn/__init__.py index 2a591716a9..e4a41f8441 100644 --- a/src/pykeen/nn/__init__.py +++ b/src/pykeen/nn/__init__.py @@ -3,7 +3,7 @@ """PyKEEN internal "nn" module.""" from . import functional, init -from .emb import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule +from .representation import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule from .modules import Interaction __all__ = [ diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index b8b9994788..0317d2504d 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -126,13 +126,13 @@ def forward( r: RelationRepresentation, t: TailRepresentation, ) -> torch.FloatTensor: - """Compute broadcasted triple scores given representations for head, relation and tails. + """Compute broadcasted triple scores given broadcasted representations for head, relation and tails. - :param h: shape: (batch_size, num_heads, ``*``) + :param h: shape: (batch_size, num_heads, 1, 1, ``*``) The head representations. - :param r: shape: (batch_size, num_relations, ``*``) + :param r: shape: (batch_size, 1, num_relations, 1, ``*``) The relation representations. - :param t: shape: (batch_size, num_tails, ``*``) + :param t: shape: (batch_size, 1, 1, num_tails, ``*``) The tail representations. :return: shape: (batch_size, num_heads, num_relations, num_tails) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/representation.py similarity index 100% rename from src/pykeen/nn/emb.py rename to src/pykeen/nn/representation.py From c55c4990b1be8d6f8cfa730fc8768bb5b86f2120 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 20:30:53 +0100 Subject: [PATCH 506/690] Adjust get_in_canonical_shape --- src/pykeen/models/base.py | 20 ++---- src/pykeen/models/multimodal/base.py | 2 - src/pykeen/models/unimodal/conv_e.py | 11 ++- src/pykeen/nn/representation.py | 100 ++++++++++++--------------- src/pykeen/testing/mocks.py | 14 +--- 5 files changed, 58 insertions(+), 89 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 9a1154eae2..0e9675a75a 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1277,6 +1277,8 @@ def _get_representations( class SingleVectorEmbeddingModel(ERModel, autoreset=False): """A KGEM that stores one :class:`pykeen.nn.Embedding` for each entities and relations.""" + # TODO: Get rid of it + def __init__( self, triples_factory: TriplesFactory, @@ -1322,12 +1324,10 @@ def __init__( interaction=interaction, entity_representations=Embedding.from_specification( num_embeddings=triples_factory.num_entities, - shape=embedding_dim, specification=embedding_specification, ), relation_representations=Embedding.from_specification( num_embeddings=triples_factory.num_relations, - shape=relation_dim, specification=relation_embedding_specification, ), ) @@ -1350,6 +1350,8 @@ class DoubleRelationEmbeddingModel(ERModel, autoreset=False): - :class:`pykeen.models.TransH` """ + # TODO: Get rid of it + def __init__( self, triples_factory: TriplesFactory, @@ -1385,19 +1387,16 @@ def __init__( entity_representations=[ Embedding.from_specification( num_embeddings=triples_factory.num_entities, - shape=embedding_dim, specification=embedding_specification, ), ], relation_representations=[ Embedding.from_specification( num_embeddings=triples_factory.num_relations, - shape=relation_dim, specification=relation_embedding_specification, ), Embedding.from_specification( num_embeddings=triples_factory.num_relations, - shape=second_relation_dim, specification=second_relation_embedding_specification, ), ], @@ -1413,6 +1412,8 @@ class TwoVectorEmbeddingModel(ERModel, autoreset=False): - :class:`pykeen.models.TransD` """ + # TODO: Get rid of it + def __init__( self, triples_factory: TriplesFactory, @@ -1445,24 +1446,20 @@ def __init__( entity_representations=[ Embedding.from_specification( num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, specification=embedding_specification, ), Embedding.from_specification( num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, specification=second_embedding_specification, ), ], relation_representations=[ Embedding.from_specification( num_embeddings=triples_factory.num_relations, - embedding_dim=relation_dim, specification=relation_embedding_specification, ), Embedding.from_specification( num_embeddings=triples_factory.num_relations, - embedding_dim=relation_dim, specification=second_relation_embedding_specification, ), ], @@ -1494,6 +1491,7 @@ def __init__( second_embedding_specification: Optional[EmbeddingSpecification] = None, second_relation_embedding_specification: Optional[EmbeddingSpecification] = None, ): + # TODO: Get rid of it if relation_dim is None: relation_dim = embedding_dim super().__init__( @@ -1506,24 +1504,20 @@ def __init__( entity_representations=[ Embedding.from_specification( num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, specification=embedding_specification, ), Embedding.from_specification( num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, specification=second_embedding_specification, ), ], relation_representations=[ Embedding.from_specification( num_embeddings=triples_factory.num_relations, - embedding_dim=relation_dim, specification=relation_embedding_specification, ), Embedding.from_specification( num_embeddings=triples_factory.num_relations, - embedding_dim=relation_dim, specification=second_relation_embedding_specification, ), ], diff --git a/src/pykeen/models/multimodal/base.py b/src/pykeen/models/multimodal/base.py index 8becb00986..03e0ded7e6 100644 --- a/src/pykeen/models/multimodal/base.py +++ b/src/pykeen/models/multimodal/base.py @@ -47,7 +47,6 @@ def __init__( # entity embeddings Embedding.from_specification( num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, specification=entity_specification, ), # Entity literals @@ -57,7 +56,6 @@ def __init__( ], relation_representations=Embedding.from_specification( num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim, specification=relation_specification, ), ) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 05b0f66523..bd19dbde7f 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -141,7 +141,6 @@ def __init__( entity_representations=[ Embedding.from_specification( num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, specification=EmbeddingSpecification( initializer=xavier_normal_, ), @@ -149,7 +148,6 @@ def __init__( # ConvE uses one bias for each entity Embedding.from_specification( num_embeddings=triples_factory.num_entities, - embedding_dim=1, specification=EmbeddingSpecification( initializer=nn.init.zeros_, ), @@ -157,7 +155,6 @@ def __init__( ], relation_representations=Embedding.from_specification( num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim, specification=EmbeddingSpecification( initializer=xavier_normal_, ), @@ -185,8 +182,8 @@ def forward( slice_size: Optional[int] = None, slice_dim: Optional[str] = None, ) -> torch.FloatTensor: # noqa: D102 - h = self.entity_representations[0].get_in_canonical_shape(indices=h_indices) - r = self.relation_representations[0].get_in_canonical_shape(indices=r_indices) - t = self.entity_representations[0].get_in_canonical_shape(indices=t_indices) - t_bias = self.entity_representations[1].get_in_canonical_shape(indices=t_indices) + h = self.entity_representations[0].get_in_canonical_shape(dim="h", indices=h_indices) + r = self.relation_representations[0].get_in_canonical_shape(dim="r", indices=r_indices) + t = self.entity_representations[0].get_in_canonical_shape(dim="t", indices=t_indices) + t_bias = self.entity_representations[1].get_in_canonical_shape(dim="t", indices=t_indices) return self.interaction.score(h=h, r=r, t=(t, t_bias), slice_size=slice_size, slice_dim=slice_dim) diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index 6f4d21bae4..639999d890 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -25,6 +25,18 @@ logger = logging.getLogger(__name__) +HEAD_DIM = 1 +RELATION_DIM = 2 +TAIL_DIM = 3 +DIMS = dict(h=HEAD_DIM, r=RELATION_DIM, t=TAIL_DIM) + + +def _normalize_dim(dim: Union[int, str]) -> int: + """Normalize the dimension selection.""" + if isinstance(dim, int): + return dim + return DIMS[dim.lower()[0]] + class RepresentationModule(nn.Module, ABC): """A base class for obtaining representations for entities/relations.""" @@ -54,30 +66,47 @@ def forward(self, indices: Optional[torch.LongTensor] = None) -> torch.FloatTens def get_in_canonical_shape( self, + dim: Union[int, str], indices: Optional[torch.LongTensor] = None, - reshape_dim: Optional[Sequence[int]] = None, ) -> torch.FloatTensor: """Get representations in canonical shape. + The canonical shape is given as + + (batch_size, d_1, d_2, d_3, ``*``) + + fulfilling the following properties: + + Let i = dim. If indices is None, the return shape is (1, d_1, d_2, d_3) with d_i = num_representations, + d_i = 1 else. If indices is not None, then batch_size = indices.shape[0], and d_i = 1 if + indices.ndimension() = 1 else d_i = indices.shape[1] + The canonical shape is given by (batch_size, 1, ``*``) if indices is not None, where batch_size=len(indices), or (1, num, ``*``) if indices is None with num equal to the total number of embeddings. + :param dim: + The dimension along which to expand for indices = None, or indices.ndimension() == 2. :param indices: - The indices. If None, return all embeddings. - :param reshape_dim: - Optionally reshape the last dimension. + The indices. Either None, in which care all embeddings are returned, or a 1 or 2 dimensional index tensor. - :return: shape: (batch_size, num_embeddings, ``*``) + :return: shape: (batch_size, d1, d2, d3, *self.shape) """ - x = self(indices=indices) + dim = _normalize_dim(dim=dim) if indices is None: - x = x.unsqueeze(dim=0) + x = self(indices=indices) + r_shape = (1, self.max_id) else: - x = x.unsqueeze(dim=1) - if reshape_dim is not None: - x = x.view(*x.shape[:-1], *reshape_dim) - return x + flat_indices = indices.view(-1) + x = self(indices=flat_indices) + if indices.ndimension() > 1: + x = x.view(*indices.shape, -1) + r_shape = tuple(indices.shape) + if len(r_shape) < 2: + r_shape = r_shape + (1,) + shape = [r_shape[0], 1, 1, 1] + shape[dim] = r_shape[1] + return x.view(*shape, *self.shape) def reset_parameters(self) -> None: """Reset the module's parameters.""" @@ -90,8 +119,7 @@ def post_parameter_update(self): class EmbeddingSpecification: """An embedding specification.""" - # embedding_dim: int - # shape: Optional[Sequence[int]] = None + shape: Optional[Sequence[int]] = None initializer: Optional[Initializer] = None initializer_kwargs: Optional[Mapping[str, Any]] = None @@ -107,14 +135,10 @@ class EmbeddingSpecification: def make( self, num_embeddings: int, - embedding_dim: Optional[int], - shape: Optional[Union[int, Sequence[int]]], ) -> 'Embedding': """Create an embedding with this specification.""" return Embedding( num_embeddings=num_embeddings, - embedding_dim=embedding_dim, - shape=shape, initializer=self.initializer, initializer_kwargs=self.initializer_kwargs, normalizer=self.normalizer, @@ -178,16 +202,13 @@ def __init__( If neither shape nor embedding_dim are given. """ if shape is None and embedding_dim is None: - raise ValueError('Missing both shape and embedding_dim') - elif shape is not None: - if isinstance(shape, int): - shape = (shape,) - else: - shape = shape + raise ValueError('Missing both, shape and embedding_dim') + elif shape is None: + shape = (embedding_dim,) + elif embedding_dim is None: embedding_dim = numpy.prod(shape) else: - assert embedding_dim is not None - shape = (embedding_dim,) + raise ValueError('Provided both, shape and embedding_dim') super().__init__(shape=shape, max_id=num_embeddings) if initializer is None: @@ -217,18 +238,12 @@ def __init__( def from_specification( cls, num_embeddings: int, - embedding_dim: Optional[int] = None, - shape: Optional[Union[int, Sequence[int]]] = None, specification: Optional[EmbeddingSpecification] = None, ) -> 'Embedding': """Create an embedding based on a specification. :param num_embeddings: >0 The number of embeddings. - :param embedding_dim: >0 - The embedding dimension. - :param shape: - The embedding shape. If given, shape supersedes embedding_dim, with setting embedding_dim = prod(shape). :param specification: The specification. :return: @@ -238,8 +253,6 @@ def from_specification( specification = EmbeddingSpecification() return specification.make( num_embeddings=num_embeddings, - embedding_dim=embedding_dim, - shape=shape, ) @property @@ -277,27 +290,6 @@ def forward( self.regularizer.update(x) return x - def get_in_canonical_shape( - self, - indices: Optional[torch.LongTensor] = None, - reshape_dim: Optional[Sequence[int]] = None, - ) -> torch.FloatTensor: - """Get embedding in canonical shape. - - :param indices: - The indices. If None, return all embeddings. - :param reshape_dim: - Optionally reshape the last dimension. - - :return: shape: (batch_size, num_embeddings, d) - """ - if len(self.shape) > 1 and reshape_dim is None: - reshape_dim = self.shape - return super().get_in_canonical_shape( - indices=indices, - reshape_dim=reshape_dim, - ) - class LiteralRepresentations(Embedding): """Literal representations.""" diff --git a/src/pykeen/testing/mocks.py b/src/pykeen/testing/mocks.py index 07554538a0..e67f460600 100644 --- a/src/pykeen/testing/mocks.py +++ b/src/pykeen/testing/mocks.py @@ -63,16 +63,4 @@ def __init__(self, num_entities: int, shape: Sequence[int]): def forward(self, indices: Optional[torch.LongTensor] = None) -> torch.FloatTensor: # noqa: D102 n = self.max_id if indices is None else indices.shape[0] - return self.x.unsqueeze(dim=0).repeat(n, 1).view(-1, *self.shape) - - def get_in_canonical_shape( - self, - indices: Optional[torch.LongTensor] = None, - reshape_dim: Optional[Sequence[int]] = None, - ) -> torch.FloatTensor: # noqa: D102 - x = self(indices=indices) - if indices is None: - x = x.unsqueeze(dim=0) - else: - x = x.unsqueeze(dim=1) - return x + return self.x.unsqueeze(dim=0).repeat(n, 1) From da0790df4cd2a4f9cbd0c0e24506583f2b3863a4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 20:31:55 +0100 Subject: [PATCH 507/690] Fix _get_representations --- src/pykeen/models/base.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 0e9675a75a..67452c2870 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1260,13 +1260,13 @@ def _get_representations( ]: h, r, t = [ [ - representation.get_in_canonical_shape(indices=indices) + representation.get_in_canonical_shape(dim=dim, indices=indices) for representation in representations ] - for indices, representations in ( - (h_indices, self.entity_representations), - (r_indices, self.relation_representations), - (t_indices, self.entity_representations), + for dim, indices, representations in ( + ("h", h_indices, self.entity_representations), + ("r", r_indices, self.relation_representations), + ("t", t_indices, self.entity_representations), ) ] # normalization From c439dc46b163372251fb50e42d2514c938cef565 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 20:37:13 +0100 Subject: [PATCH 508/690] Allow forwarding embedding specifications to ERModel --- src/pykeen/models/base.py | 36 ++++++++++++++++++++++++++++----- src/pykeen/nn/representation.py | 3 +++ 2 files changed, 34 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 67452c2870..886c5487f8 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1094,12 +1094,22 @@ def load_state(self, path: str) -> None: def _prepare_representation_module_list( - representations: Union[None, RepresentationModule, Sequence[RepresentationModule]], + representations: Union[ + None, + EmbeddingSpecification, + RepresentationModule, + Sequence[Union[EmbeddingSpecification, RepresentationModule]], + ], + num_embeddings: int, ) -> Sequence[RepresentationModule]: """Normalize list of representations and wrap into nn.ModuleList.""" # Important: use ModuleList to ensure that Pytorch correctly handles their devices and parameters if representations is not None and not isinstance(representations, Sequence): representations = [representations] + representations = [ + r if isinstance(r, RepresentationModule) else r.make(num_embeddings=num_embeddings) + for r in representations + ] return nn.ModuleList(representations) @@ -1119,13 +1129,23 @@ def __init__( self, triples_factory: TriplesFactory, interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], + entity_representations: Union[ + None, + EmbeddingSpecification, + RepresentationModule, + Sequence[Union[EmbeddingSpecification, RepresentationModule]] + ] = None, + relation_representations: Union[ + None, + EmbeddingSpecification, + RepresentationModule, + Sequence[Union[EmbeddingSpecification, RepresentationModule]] + ] = None, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, - entity_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, - relation_representations: Union[None, RepresentationModule, Sequence[RepresentationModule]] = None, ) -> None: """Initialize the module. @@ -1154,8 +1174,14 @@ def __init__( random_seed=random_seed, predict_with_sigmoid=predict_with_sigmoid, ) - self.entity_representations = _prepare_representation_module_list(representations=entity_representations) - self.relation_representations = _prepare_representation_module_list(representations=relation_representations) + self.entity_representations = _prepare_representation_module_list( + representations=entity_representations, + num_embeddings=triples_factory.num_entities, + ) + self.relation_representations = _prepare_representation_module_list( + representations=relation_representations, + num_embeddings=triples_factory.num_relations, + ) self.interaction = interaction # Comment: it is important that the regularizers are stored in a module list, in order to appear in # model.modules(). Thereby, we can collect them automatically. diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index 639999d890..4602700bb7 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -119,6 +119,7 @@ def post_parameter_update(self): class EmbeddingSpecification: """An embedding specification.""" + embedding_dim: Optional[int] = None shape: Optional[Sequence[int]] = None initializer: Optional[Initializer] = None @@ -139,6 +140,8 @@ def make( """Create an embedding with this specification.""" return Embedding( num_embeddings=num_embeddings, + embedding_dim=self.embedding_dim, + shape=self.shape, initializer=self.initializer, initializer_kwargs=self.initializer_kwargs, normalizer=self.normalizer, From 8743038b707f64e4b77227e319eb44f6e832b719 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 20:37:19 +0100 Subject: [PATCH 509/690] Update DistMult --- src/pykeen/models/unimodal/distmult.py | 19 ++++++++++--------- 1 file changed, 10 insertions(+), 9 deletions(-) diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 6dfb4330be..e00a6871b4 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -7,7 +7,7 @@ from torch import nn from torch.nn import functional -from ..base import SingleVectorEmbeddingModel +from .. import ERModel from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.modules import DistMultInteraction @@ -21,7 +21,7 @@ ] -class DistMult(SingleVectorEmbeddingModel): +class DistMult(ERModel): r"""An implementation of DistMult from [yang2014]_. This model simplifies RESCAL by restricting matrices representing relations as diagonal matrices. @@ -83,19 +83,16 @@ def __init__( super().__init__( triples_factory=triples_factory, interaction=DistMultInteraction(), - embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, - embedding_specification=EmbeddingSpecification( + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, # xavier uniform, cf. # https://github.com/thunlp/OpenKE/blob/adeed2c0d2bef939807ed4f69c1ea4db35fd149b/models/DistMult.py#L16-L17 initializer=nn.init.xavier_uniform_, # Constrain entity embeddings to unit length constrainer=functional.normalize, ), - relation_embedding_specification=EmbeddingSpecification( + relation_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, # relations are initialized to unit length (but not constraint) initializer=compose( nn.init.xavier_uniform_, @@ -104,4 +101,8 @@ def __init__( # Only relation embeddings are regularized regularizer=self._instantiate_default_regularizer(), ), + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + random_seed=random_seed, ) From e346e522ebe148657721ff83c449f961b647d1f0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 20:43:06 +0100 Subject: [PATCH 510/690] Better integration of complex embeddings --- src/pykeen/models/unimodal/complex.py | 14 +++++++++----- src/pykeen/nn/representation.py | 13 +++++++++++++ 2 files changed, 22 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 708c5bed60..6d4814a1cd 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -4,9 +4,10 @@ from typing import Optional +import torch import torch.nn as nn -from ..base import SingleVectorEmbeddingModel +from ..base import ERModel from ...losses import Loss, SoftplusLoss from ...nn import EmbeddingSpecification from ...nn.modules import ComplExInteraction @@ -19,7 +20,7 @@ ] -class ComplEx(SingleVectorEmbeddingModel): +class ComplEx(ERModel): r"""An implementation of ComplEx [trouillon2016]_. ComplEx is an extension of :class:`pykeen.models.DistMult` that uses complex valued representations for the @@ -98,22 +99,25 @@ def __init__( # https://github.com/ttrouill/complex/blob/dc4eb93408d9a5288c986695b58488ac80b1cc17/efe/models.py#L481-L487 if embedding_specification is None: embedding_specification = EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=nn.init.normal_, regularizer=regularizer, + dtype=torch.complex64, ) if relation_embedding_specification is None: relation_embedding_specification = EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=nn.init.normal_, regularizer=regularizer, + dtype=torch.complex64, ) super().__init__( triples_factory=triples_factory, interaction=ComplExInteraction(), - embedding_dim=2 * embedding_dim, # complex embeddings + entity_representations=embedding_specification, + relation_representations=relation_embedding_specification, automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, - embedding_specification=embedding_specification, - relation_embedding_specification=relation_embedding_specification, ) diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index 4602700bb7..98a28753d0 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -122,6 +122,8 @@ class EmbeddingSpecification: embedding_dim: Optional[int] = None shape: Optional[Sequence[int]] = None + dtype: Optional[torch.dtype] = None + initializer: Optional[Initializer] = None initializer_kwargs: Optional[Mapping[str, Any]] = None @@ -142,6 +144,7 @@ def make( num_embeddings=num_embeddings, embedding_dim=self.embedding_dim, shape=self.shape, + dtype=self.dtype, initializer=self.initializer, initializer_kwargs=self.initializer_kwargs, normalizer=self.normalizer, @@ -175,6 +178,7 @@ def __init__( constrainer: Optional[Constrainer] = None, constrainer_kwargs: Optional[Mapping[str, Any]] = None, regularizer: Optional[Regularizer] = None, + dtype: Optional[torch.dtype] = None, ): """Instantiate an embedding with extended functionality. @@ -212,6 +216,15 @@ def __init__( embedding_dim = numpy.prod(shape) else: raise ValueError('Provided both, shape and embedding_dim') + if dtype is None: + dtype = torch.get_default_dtype() + + # work-around until full complex support + # TODO: verify that this is our understanding of complex! + if dtype.is_complex: + shape = tuple(shape) + shape = shape[:-1] + (2 * shape[-1],) + embedding_dim = embedding_dim * 2 super().__init__(shape=shape, max_id=num_embeddings) if initializer is None: From 664bb5b7139b8c93769b18ff9d130dc06337e5c7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 20:45:20 +0100 Subject: [PATCH 511/690] Update ConvKB --- src/pykeen/models/unimodal/conv_kb.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 600d211b61..9a77b059ee 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -5,8 +5,9 @@ import logging from typing import Optional -from ..base import SingleVectorEmbeddingModel +from ..base import ERModel from ...losses import Loss +from ...nn import EmbeddingSpecification from ...nn.modules import ConvKBInteraction from ...regularizers import LpRegularizer from ...triples import TriplesFactory @@ -19,7 +20,7 @@ logger = logging.getLogger(__name__) -class ConvKB(SingleVectorEmbeddingModel): +class ConvKB(ERModel): r"""An implementation of ConvKB from [nguyen2018]_. ConvKB uses a convolutional neural network (CNN) whose feature maps capture global interactions of the input. @@ -90,7 +91,12 @@ def __init__( embedding_dim=embedding_dim, num_filters=num_filters, ), - embedding_dim=embedding_dim, + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim + ), + relation_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, + ), loss=loss, automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, From 1507c315bd46868142d48088b29210e387bdca32 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 20:51:23 +0100 Subject: [PATCH 512/690] Update ERMLP - KG2E --- src/pykeen/models/unimodal/ermlp.py | 12 ++++++-- src/pykeen/models/unimodal/ermlpe.py | 12 ++++++-- src/pykeen/models/unimodal/hole.py | 19 +++++++------ src/pykeen/models/unimodal/kg2e.py | 42 ++++++++++++++-------------- 4 files changed, 49 insertions(+), 36 deletions(-) diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index bd575f4310..3d0c8cd2be 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -4,8 +4,9 @@ from typing import Optional -from ..base import SingleVectorEmbeddingModel +from ..base import ERModel from ...losses import Loss +from ...nn import EmbeddingSpecification from ...nn.modules import ERMLPInteraction from ...triples import TriplesFactory from ...typing import DeviceHint @@ -15,7 +16,7 @@ ] -class ERMLP(SingleVectorEmbeddingModel): +class ERMLP(ERModel): r"""An implementation of ERMLP from [dong2014]_. ERMLP is a multi-layer perceptron based approach that uses a single hidden layer and represents entities and @@ -57,7 +58,12 @@ def __init__( embedding_dim=embedding_dim, hidden_dim=hidden_dim, ), - embedding_dim=embedding_dim, + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, + ), + relation_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, + ), automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index 340c43a3b8..00da44382c 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -4,8 +4,9 @@ from typing import Any, ClassVar, Mapping, Optional, Type -from ..base import SingleVectorEmbeddingModel +from ..base import ERModel from ...losses import BCEAfterSigmoidLoss, Loss +from ...nn import EmbeddingSpecification from ...nn.modules import ERMLPEInteraction from ...triples import TriplesFactory from ...typing import DeviceHint @@ -15,7 +16,7 @@ ] -class ERMLPE(SingleVectorEmbeddingModel): +class ERMLPE(ERModel): r"""An extension of ERMLP proposed by [sharifzadeh2019]_. This model uses a neural network-based approach similar to ERMLP and with slight modifications. @@ -70,7 +71,12 @@ def __init__( hidden_dropout=hidden_dropout, embedding_dim=embedding_dim, ), - embedding_dim=embedding_dim, + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, + ), + relation_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, + ), automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 4e07aef62c..dbf70612fa 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -4,7 +4,7 @@ from typing import Optional -from ..base import SingleVectorEmbeddingModel +from ..base import ERModel from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.init import xavier_uniform_ @@ -18,7 +18,7 @@ ] -class HolE(SingleVectorEmbeddingModel): +class HolE(ERModel): r"""An implementation of HolE [nickel2016]_. Holographic embeddings (HolE) make use of the circular correlation operator to compute interactions between @@ -63,19 +63,20 @@ def __init__( """Initialize the model.""" super().__init__( triples_factory=triples_factory, - embedding_dim=embedding_dim, - loss=loss, interaction=HolEInteraction(), - automatic_memory_optimization=automatic_memory_optimization, - preferred_device=preferred_device, - random_seed=random_seed, # Initialization, cf. https://github.com/mnick/scikit-kge/blob/master/skge/param.py#L18-L27 - embedding_specification=EmbeddingSpecification( + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=xavier_uniform_, constrainer=clamp_norm, # type: ignore constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), - relation_embedding_specification=EmbeddingSpecification( + relation_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=xavier_uniform_, ), + loss=loss, + automatic_memory_optimization=automatic_memory_optimization, + preferred_device=preferred_device, + random_seed=random_seed, ) diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index b73113b34e..1356d8f232 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -8,7 +8,7 @@ import torch import torch.autograd -from ..base import TwoVectorEmbeddingModel +from ..base import ERModel from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.modules import KG2EInteraction @@ -23,7 +23,7 @@ _LOG_2_PI = math.log(math.tau) -class KG2E(TwoVectorEmbeddingModel): +class KG2E(ERModel): r"""An implementation of KG2E from [he2015]_. KG2E aims to explicitly model (un)certainties in entities and relations (e.g. influenced by the number of triples @@ -76,32 +76,32 @@ def __init__( :param c_min: :param c_max: """ + # Both, entities and relations, are represented as d-dimensional Normal distributions with diagonal covariance + # matrix + representation_spec = [ + # mean of Normal distribution + EmbeddingSpecification( + embedding_dim=embedding_dim, + constrainer=clamp_norm, # type: ignore + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + ), + # diagonal covariance of Normal distribution + # Ensure positive definite covariances matrices and appropriate size by clamping + EmbeddingSpecification( + embedding_dim=embedding_dim, + constrainer=torch.clamp, + constrainer_kwargs=dict(min=c_min, max=c_max), + ), + ] super().__init__( triples_factory=triples_factory, interaction=KG2EInteraction( similarity=dist_similarity, ), - embedding_dim=embedding_dim, + entity_representations=representation_spec, + relation_representations=representation_spec, automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, - embedding_specification=EmbeddingSpecification( - constrainer=clamp_norm, # type: ignore - constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), - ), - relation_embedding_specification=EmbeddingSpecification( - constrainer=clamp_norm, # type: ignore - constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), - ), - # Ensure positive definite covariances matrices and appropriate size by clamping - second_embedding_specification=EmbeddingSpecification( - constrainer=torch.clamp, - constrainer_kwargs=dict(min=c_min, max=c_max), - ), - second_relation_embedding_specification=EmbeddingSpecification( - # Ensure positive definite covariances matrices and appropriate size by clamping - constrainer=torch.clamp, - constrainer_kwargs=dict(min=c_min, max=c_max), - ), ) From 94b79cfc774469fc4013e535d192b9efe015db79 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:00:11 +0100 Subject: [PATCH 513/690] Update NTN - RotatE --- src/pykeen/models/unimodal/ntn.py | 48 +++++++++++----------------- src/pykeen/models/unimodal/proj_e.py | 19 +++++------ src/pykeen/models/unimodal/rescal.py | 20 ++++++------ src/pykeen/models/unimodal/rgcn.py | 10 +++--- src/pykeen/models/unimodal/rotate.py | 23 +++++++------ 5 files changed, 57 insertions(+), 63 deletions(-) diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 273541ca26..59dd97d11c 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -6,9 +6,9 @@ from torch import nn -from .. import ERModel +from ..base import ERModel from ...losses import Loss -from ...nn import Embedding +from ...nn import EmbeddingSpecification from ...nn.modules import NTNInteraction from ...triples import TriplesFactory from ...typing import DeviceHint @@ -69,38 +69,28 @@ def __init__( :param non_linearity: A non-linear activation function. Defaults to the hyperbolic tangent :class:`torch.nn.Tanh`. """ - w = Embedding( - num_embeddings=triples_factory.num_relations, - shape=(num_slices, embedding_dim, embedding_dim), - ) - vh = Embedding( - num_embeddings=triples_factory.num_relations, - shape=(num_slices, embedding_dim), - ) - vt = Embedding( - num_embeddings=triples_factory.num_relations, - shape=(num_slices, embedding_dim), - ) - b = Embedding( - num_embeddings=triples_factory.num_relations, - embedding_dim=num_slices, - ) - u = Embedding( - num_embeddings=triples_factory.num_relations, - embedding_dim=num_slices, - ) super().__init__( triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, interaction=NTNInteraction( non_linearity=non_linearity, ), - entity_representations=Embedding( - num_embeddings=triples_factory.num_entities, + entity_representations=EmbeddingSpecification( embedding_dim=embedding_dim, ), - relation_representations=(w, vh, vt, b, u), + relation_representations=[ + # w: (k, d, d) + EmbeddingSpecification(shape=(num_slices, embedding_dim, embedding_dim)), + # vh: (k, d) + EmbeddingSpecification(shape=(num_slices, embedding_dim)), + # vt: (k, d) + EmbeddingSpecification(shape=(num_slices, embedding_dim)), + # b: (k,) + EmbeddingSpecification(shape=(num_slices,)), + # u: (k,) + EmbeddingSpecification(shape=(num_slices,)), + ], + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + random_seed=random_seed, ) diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index 3c7ef1a486..397b90eaf1 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -6,7 +6,7 @@ from torch import nn -from .. import SingleVectorEmbeddingModel +from ..base import ERModel from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.init import xavier_uniform_ @@ -19,7 +19,7 @@ ] -class ProjE(SingleVectorEmbeddingModel): +class ProjE(ERModel): r"""An implementation of ProjE from [shi2017]_. ProjE is a neural network-based approach with a *combination* and a *projection* layer. The interaction model @@ -69,15 +69,16 @@ def __init__( embedding_dim=embedding_dim, inner_non_linearity=inner_non_linearity, ), - embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, - embedding_specification=EmbeddingSpecification( + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=xavier_uniform_, ), - relation_embedding_specification=EmbeddingSpecification( + relation_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=xavier_uniform_, ), + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + random_seed=random_seed, ) diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index e91b6ed0f0..3ac048c9fa 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -4,7 +4,7 @@ from typing import Optional -from ..base import SingleVectorEmbeddingModel +from ..base import ERModel from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.modules import RESCALInteraction @@ -17,7 +17,7 @@ ] -class RESCAL(SingleVectorEmbeddingModel): +class RESCAL(ERModel): r"""An implementation of RESCAL from [nickel2011]_. This model represents relations as matrices and models interactions between latent features. @@ -71,16 +71,16 @@ def __init__( super().__init__( triples_factory=triples_factory, interaction=RESCALInteraction(), - embedding_dim=embedding_dim, - relation_dim=(embedding_dim, embedding_dim), - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, - embedding_specification=EmbeddingSpecification( + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, regularizer=regularizer, ), - relation_embedding_specification=EmbeddingSpecification( + relation_representations=EmbeddingSpecification( + shape=(embedding_dim, embedding_dim), regularizer=regularizer, ), + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + random_seed=random_seed, ) diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index bc1aadd93d..c04a595033 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -13,7 +13,7 @@ from . import ComplEx, DistMult, ERMLP from ..base import ERModel from ...losses import Loss -from ...nn import Embedding, Interaction, RepresentationModule +from ...nn import Embedding, EmbeddingSpecification, Interaction, RepresentationModule from ...nn.modules import DistMultInteraction from ...triples import TriplesFactory from ...typing import DeviceHint @@ -499,10 +499,6 @@ def __init__( buffer_messages=buffer_messages, base_representations=None, ) - relation_representations = Embedding( - num_embeddings=triples_factory.num_relations, - embedding_dim=embedding_dim, - ) super().__init__( triples_factory=triples_factory, automatic_memory_optimization=automatic_memory_optimization, @@ -512,5 +508,7 @@ def __init__( random_seed=random_seed, interaction=interaction, entity_representations=entity_representations, - relation_representations=relation_representations, + relation_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, + ), ) diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index 56908e00bd..36d83afa3b 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -4,7 +4,9 @@ from typing import Optional -from .. import SingleVectorEmbeddingModel +import torch + +from ..base import ERModel from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.init import init_phases, xavier_uniform_ @@ -18,7 +20,7 @@ ] -class RotatE(SingleVectorEmbeddingModel): +class RotatE(ERModel): r"""An implementation of RotatE from [sun2019]_. RotatE models relations as rotations from head to tail entities in complex space: @@ -60,16 +62,19 @@ def __init__( super().__init__( triples_factory=triples_factory, interaction=RotatEInteraction(), - embedding_dim=2 * embedding_dim, - loss=loss, - automatic_memory_optimization=automatic_memory_optimization, - preferred_device=preferred_device, - random_seed=random_seed, - embedding_specification=EmbeddingSpecification( + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=xavier_uniform_, + dtype=torch.complex64, ), - relation_embedding_specification=EmbeddingSpecification( + relation_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=init_phases, constrainer=complex_normalize, + dtype=torch.complex64, ), + loss=loss, + automatic_memory_optimization=automatic_memory_optimization, + preferred_device=preferred_device, + random_seed=random_seed, ) From 4912501ffb04d0fdf79650c30235ec6755b13f93 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:08:32 +0100 Subject: [PATCH 514/690] Fix SimplE --- src/pykeen/models/unimodal/simple.py | 74 ++++++++++++++-------------- src/pykeen/nn/functional.py | 1 - src/pykeen/nn/modules.py | 30 +++++++++++ 3 files changed, 68 insertions(+), 37 deletions(-) diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index 6e3c6e3a92..c8c3791032 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -4,12 +4,12 @@ from typing import Any, ClassVar, Mapping, Optional, Tuple, Union -import torch.autograd +import torch -from ..base import TwoSideEmbeddingModel +from ..base import ERModel from ...losses import Loss, SoftplusLoss from ...nn import EmbeddingSpecification -from ...nn.modules import DistMultInteraction +from ...nn.modules import SimplEInteraction from ...regularizers import PowerSumRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -19,7 +19,7 @@ ] -class SimplE(TwoSideEmbeddingModel): +class SimplE(ERModel): r"""An implementation of SimplE [kazemi2018]_. SimplE is an extension of canonical polyadic (CP), an early tensor factorization approach in which each entity @@ -76,28 +76,32 @@ def __init__( regularizer = self._instantiate_default_regularizer() super().__init__( triples_factory=triples_factory, - interaction=DistMultInteraction(), - embedding_dim=embedding_dim, + interaction=SimplEInteraction(clamp_score=clamp_score), + entity_representations=[ + EmbeddingSpecification( + embedding_dim=embedding_dim, + regularizer=regularizer, + ), + EmbeddingSpecification( + embedding_dim=embedding_dim, + regularizer=regularizer, + ) + ], + relation_representations=[ + EmbeddingSpecification( + embedding_dim=embedding_dim, + regularizer=regularizer, + ), + EmbeddingSpecification( + embedding_dim=embedding_dim, + regularizer=regularizer, + ), + ], automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, - embedding_specification=EmbeddingSpecification( - regularizer=regularizer, - ), - relation_embedding_specification=EmbeddingSpecification( - regularizer=regularizer, - ), - second_embedding_specification=EmbeddingSpecification( - regularizer=regularizer, - ), - second_relation_embedding_specification=EmbeddingSpecification( - regularizer=regularizer, - ), ) - if isinstance(clamp_score, float): - clamp_score = (-clamp_score, clamp_score) - self.clamp = clamp_score def forward( self, @@ -107,18 +111,16 @@ def forward( slice_size: Optional[int] = None, slice_dim: Optional[str] = None, ) -> torch.FloatTensor: # noqa: D102 - scores = super().forward( - h_indices=h_indices, - r_indices=r_indices, - t_indices=t_indices, - slice_size=slice_size, - slice_dim=slice_dim, - ) - - # Note: In the code in their repository, the score is clamped to [-20, 20]. - # That is not mentioned in the paper, so it is omitted here. - if self.clamp is not None: - min_, max_ = self.clamp - scores = scores.clamp(min=min_, max=max_) - - return scores + h, r, t = zip(*( + ( + h_source.get_in_canonical_shape(dim="h", indices=h_indices), + r_source.get_in_canonical_shape(dim="r", indices=r_indices), + t_source.get_in_canonical_shape(dim="t", indices=t_indices), + ) + for h_source, r_source, t_source in ( + (self.entity_representations[0], self.relation_representations[0], self.entity_representations[1]), + (self.entity_representations[1], self.relation_representations[1], self.entity_representations[0]), + ) + )) + scores = self.interaction.score(h=h, r=r, t=t, slice_size=slice_size, slice_dim=slice_dim) + return self._repeat_scores_if_necessary(scores, h_indices, r_indices, t_indices) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index a266a160fb..7f945ecf8f 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -625,7 +625,6 @@ def simple_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # TODO: unused scores = 0.5 * (distmult_interaction(h=h, r=r, t=t) + distmult_interaction(h=h_inv, r=r_inv, t=t_inv)) # Note: In the code in their repository, the score is clamped to [-20, 20]. # That is not mentioned in the paper, so it is made optional here. diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 0317d2504d..848d25cbbc 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -32,6 +32,7 @@ 'ProjEInteraction', 'RESCALInteraction', 'RotatEInteraction', + 'SimplEInteraction', 'StructuredEmbeddingInteraction', 'TransDInteraction', 'TransEInteraction', @@ -1097,3 +1098,32 @@ def _prepare_hrt_for_functional( t: TailRepresentation, ) -> MutableMapping[str, torch.FloatTensor]: # noqa: D102 return dict(h=h, w_r=r[0], d_r=r[1], t=t) + + +class SimplEInteraction( + Interaction[ + Tuple[torch.FloatTensor, torch.FloatTensor], + Tuple[torch.FloatTensor, torch.FloatTensor], + Tuple[torch.FloatTensor, torch.FloatTensor] + ] +): + """Interaction function of SimplE.""" + + func = pkf.simple_interaction + + def __init__(self, clamp_score: Optional[Union[float, Tuple[float, float]]] = None, ): + super().__init__() + if isinstance(clamp_score, float): + clamp_score = (-clamp_score, clamp_score) + self.clamp = clamp_score + + def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: # noqa: D102 + return dict(clamp=self.clamp_score) + + @staticmethod + def _prepare_hrt_for_functional( + h: HeadRepresentation, + r: RelationRepresentation, + t: TailRepresentation, + ) -> MutableMapping[str, torch.FloatTensor]: # noqa: D102 + return dict(h=h[0], h_inv=h[1], r=r[0], r_inv=r[1], t=t[0], t_inv=t[1]) From 2258aef8370725a2ed46b29e966051c7b07f0318 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:14:52 +0100 Subject: [PATCH 515/690] Update SE - UM --- .../models/unimodal/structured_embedding.py | 31 ++++++----- src/pykeen/models/unimodal/trans_d.py | 20 +++---- src/pykeen/models/unimodal/trans_e.py | 19 +++---- src/pykeen/models/unimodal/trans_h.py | 24 +++++---- src/pykeen/models/unimodal/trans_r.py | 52 ++++++++++--------- src/pykeen/models/unimodal/tucker.py | 20 +++---- .../models/unimodal/unstructured_model.py | 7 ++- 7 files changed, 92 insertions(+), 81 deletions(-) diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index d3abc3b2a9..73749bf6b5 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -9,7 +9,7 @@ from torch import nn from torch.nn import functional -from ..base import DoubleRelationEmbeddingModel +from ..base import ERModel from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.init import xavier_uniform_ @@ -23,7 +23,7 @@ ] -class StructuredEmbedding(DoubleRelationEmbeddingModel): +class StructuredEmbedding(ERModel): r"""An implementation of the Structured Embedding (SE) published by [bordes2011]_. SE applies role- and relation-specific projection matrices @@ -73,20 +73,23 @@ def __init__( p=scoring_fct_norm, power_norm=False, ), - embedding_dim=embedding_dim, - relation_dim=(embedding_dim, embedding_dim), + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, + initializer=xavier_uniform_, + constrainer=functional.normalize, + ), + relation_representations=[ + EmbeddingSpecification( + shape=(embedding_dim, embedding_dim), + initializer=relation_initializer, + ), + EmbeddingSpecification( + shape=(embedding_dim, embedding_dim), + initializer=relation_initializer, + ), + ], automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, - embedding_specification=EmbeddingSpecification( - initializer=xavier_uniform_, - constrainer=functional.normalize, - ), - relation_embedding_specification=EmbeddingSpecification( - initializer=relation_initializer, - ), - second_relation_embedding_specification=EmbeddingSpecification( - initializer=relation_initializer, - ), ) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index bd8a4d2a1c..1a56ed794b 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -4,7 +4,7 @@ from typing import Optional -from ..base import TwoVectorEmbeddingModel +from ..base import ERModel from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.init import xavier_normal_ @@ -18,7 +18,7 @@ ] -class TransD(TwoVectorEmbeddingModel): +class TransD(ERModel): r"""An implementation of TransD from [ji2015]_. TransD is an extension of :class:`pykeen.models.TransR` that, like TransR, considers entities and relations @@ -69,20 +69,20 @@ def __init__( super().__init__( triples_factory=triples_factory, interaction=TransDInteraction(p=2, power_norm=True), - embedding_dim=embedding_dim, - relation_dim=relation_dim, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, - embedding_specification=EmbeddingSpecification( + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=xavier_normal_, constrainer=clamp_norm, # type: ignore constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), - relation_embedding_specification=EmbeddingSpecification( + relation_representations=EmbeddingSpecification( + embedding_dim=relation_dim, initializer=xavier_normal_, constrainer=clamp_norm, # type: ignore constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + random_seed=random_seed, ) diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 742e7f51d3..279493dc56 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -6,7 +6,7 @@ from torch.nn import functional -from .. import SingleVectorEmbeddingModel +from ..base import ERModel from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.init import xavier_uniform_ @@ -20,7 +20,7 @@ ] -class TransE(SingleVectorEmbeddingModel): +class TransE(ERModel): r"""TransE models relations as a translation from head to tail entities in :math:`\textbf{e}` [bordes2013]_. .. math:: @@ -68,19 +68,20 @@ def __init__( super().__init__( triples_factory=triples_factory, interaction=TransEInteraction(p=scoring_fct_norm, power_norm=False), - embedding_dim=embedding_dim, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, - embedding_specification=EmbeddingSpecification( + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=xavier_uniform_, constrainer=functional.normalize, ), - relation_embedding_specification=EmbeddingSpecification( + relation_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=compose( xavier_uniform_, functional.normalize, ), ), + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + random_seed=random_seed, ) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 488f1e1c65..eb957a7c5d 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -6,7 +6,7 @@ from torch.nn import functional -from ..base import DoubleRelationEmbeddingModel +from ..base import ERModel from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.modules import TransHInteraction @@ -20,7 +20,7 @@ ] -class TransH(DoubleRelationEmbeddingModel): +class TransH(ERModel): r"""An implementation of TransH [wang2014]_. This model extends :class:`pykeen.models.TransE` by applying the translation from head to tail entity in a @@ -85,18 +85,24 @@ def __init__( p=scoring_fct_norm, power_norm=False, ), - embedding_dim=embedding_dim, + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, + ), + relation_representations=[ + EmbeddingSpecification( + embedding_dim=embedding_dim, + ), + EmbeddingSpecification( + embedding_dim=embedding_dim, + # Normalise the normal vectors by their l2 norms + constrainer=functional.normalize, + ), + ], automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, predict_with_sigmoid=predict_with_sigmoid, - embedding_specification=None, - relation_embedding_specification=None, - second_relation_embedding_specification=EmbeddingSpecification( - # Normalise the normal vectors by their l2 norms - constrainer=functional.normalize, - ), ) # Note that the TransH regularizer has a different interface self.regularizer = self._instantiate_default_regularizer( diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index d632646b1f..704fbe8644 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -7,7 +7,7 @@ from torch.nn import functional -from ..base import DoubleRelationEmbeddingModel +from ..base import ERModel from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.init import xavier_uniform_ @@ -21,7 +21,7 @@ ] -class TransR(DoubleRelationEmbeddingModel): +class TransR(ERModel): r"""An implementation of TransR from [lin2015]_. TransR is an extension of :class:`pykeen.models.TransH` that explicitly considers entities and relations as @@ -73,34 +73,36 @@ def __init__( """Initialize the model.""" super().__init__( triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, + interaction=TransRInteraction( + p=scoring_fct_norm, + ), # Entity embeddings - embedding_dim=embedding_dim, - embedding_specification=EmbeddingSpecification( + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=xavier_uniform_, constrainer=clamp_norm, # type: ignore constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), - # Relation embeddings - relation_dim=relation_dim, - relation_embedding_specification=EmbeddingSpecification( - initializer=compose( - xavier_uniform_, - functional.normalize, + relation_representations=[ + # Relation embeddings + EmbeddingSpecification( + embedding_dim=relation_dim, + initializer=compose( + xavier_uniform_, + functional.normalize, + ), + constrainer=clamp_norm, # type: ignore + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), ), - constrainer=clamp_norm, # type: ignore - constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), - ), - # Relation projections - second_relation_dim=(relation_dim, embedding_dim), - second_relation_embedding_specification=EmbeddingSpecification( - initializer=xavier_uniform_, - ), - interaction=TransRInteraction( - p=scoring_fct_norm, - ), + # Relation projections + EmbeddingSpecification( + shape=(relation_dim, embedding_dim), + initializer=xavier_uniform_, + ), + ], + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + random_seed=random_seed, ) logging.warning("Initialize from TransE") diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index 3df8914d92..e7733a41b8 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -4,7 +4,7 @@ from typing import Any, ClassVar, Mapping, Optional -from .. import SingleVectorEmbeddingModel +from ..base import ERModel from ...losses import BCEAfterSigmoidLoss, Loss from ...nn import EmbeddingSpecification from ...nn.init import xavier_normal_ @@ -17,7 +17,7 @@ ] -class TuckER(SingleVectorEmbeddingModel): +class TuckER(ERModel): r"""An implementation of TuckEr from [balazevic2019]_. TuckER is a linear model that is based on the tensor factorization method Tucker in which a three-mode tensor @@ -95,16 +95,16 @@ def __init__( head_relation_dropout=dropout_2, apply_batch_normalization=apply_batch_normalization, ), - embedding_dim=embedding_dim, - relation_dim=relation_dim, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, - embedding_specification=EmbeddingSpecification( + entity_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=xavier_normal_, ), - relation_embedding_specification=EmbeddingSpecification( + relation_representations=EmbeddingSpecification( + embedding_dim=relation_dim, initializer=xavier_normal_, ), + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + random_seed=random_seed, ) diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index 0e2c69603d..afd299e036 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -4,9 +4,9 @@ from typing import Optional -from .. import ERModel +from ..base import ERModel from ...losses import Loss -from ...nn import Embedding +from ...nn import EmbeddingSpecification from ...nn.init import xavier_normal_ from ...nn.modules import UnstructuredModelInteraction from ...triples import TriplesFactory @@ -66,8 +66,7 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, interaction=UnstructuredModelInteraction(p=scoring_fct_norm), - entity_representations=Embedding( - num_embeddings=triples_factory.num_entities, + entity_representations=EmbeddingSpecification( embedding_dim=embedding_dim, initializer=xavier_normal_, ), From c5081c3c2254c068301c1b299df0c7c42d7d8a1e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:18:05 +0100 Subject: [PATCH 516/690] Remove abstract base classes --- src/pykeen/models/base.py | 279 +------------------------------------- 1 file changed, 1 insertion(+), 278 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 886c5487f8..8b5c123fc5 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -19,7 +19,7 @@ from torch import nn from ..losses import Loss, MarginRankingLoss, NSSALoss -from ..nn import Embedding, EmbeddingSpecification, RepresentationModule +from ..nn import EmbeddingSpecification, RepresentationModule from ..nn.modules import Interaction from ..regularizers import Regularizer, collect_regularization_terms from ..triples import TriplesFactory @@ -32,10 +32,6 @@ __all__ = [ 'Model', 'ERModel', - 'SingleVectorEmbeddingModel', - 'DoubleRelationEmbeddingModel', - 'TwoVectorEmbeddingModel', - 'TwoSideEmbeddingModel', ] logger = logging.getLogger(__name__) @@ -1298,276 +1294,3 @@ def _get_representations( # normalization h, r, t = [x[0] if len(x) == 1 else x for x in (h, r, t)] return h, r, t - - -class SingleVectorEmbeddingModel(ERModel, autoreset=False): - """A KGEM that stores one :class:`pykeen.nn.Embedding` for each entities and relations.""" - - # TODO: Get rid of it - - def __init__( - self, - triples_factory: TriplesFactory, - interaction: Interaction[torch.FloatTensor, torch.FloatTensor, torch.FloatTensor], - embedding_dim: int = 200, - relation_dim: Union[None, int, Sequence[int]] = None, - automatic_memory_optimization: Optional[bool] = None, - loss: Optional[Loss] = None, - predict_with_sigmoid: bool = False, - preferred_device: Optional[str] = None, - random_seed: Optional[int] = None, - embedding_specification: Optional[EmbeddingSpecification] = None, - relation_embedding_specification: Optional[EmbeddingSpecification] = None, - ) -> None: - """Initialize embedding model. - - :param triples_factory: - The triple factory connected to the model. - :param interaction: - The embedding-based interaction function used to compute scores. - :param embedding_dim: - The embedding dimensionality of the entity embeddings. - :param automatic_memory_optimization: - Whether to automatically optimize the sub-batch size during training and batch size during evaluation with - regards to the hardware at hand. - :param loss: - The loss to use. - :param preferred_device: - The default device where to model is located. - :param random_seed: - An optional random seed to set before the initialization of weights. - """ - # Default for relation dimensionality - if relation_dim is None: - relation_dim = embedding_dim - super().__init__( - triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - predict_with_sigmoid=predict_with_sigmoid, - preferred_device=preferred_device, - random_seed=random_seed, - interaction=interaction, - entity_representations=Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - specification=embedding_specification, - ), - relation_representations=Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - specification=relation_embedding_specification, - ), - ) - - @property - def embedding_dim(self) -> int: # noqa:D401 - """The entity embedding dim.""" - # TODO: Deprecated; directly use self.entity_representations[0].embedding_dim instead? - embedding = self.entity_representations[0] - assert isinstance(embedding, Embedding) - return embedding.embedding_dim - - -class DoubleRelationEmbeddingModel(ERModel, autoreset=False): - """A KGEM that stores one :class:`pykeen.nn.Embedding` for entities and two for relations. - - .. seealso:: - - - :class:`pykeen.models.StructuredEmbedding` - - :class:`pykeen.models.TransH` - """ - - # TODO: Get rid of it - - def __init__( - self, - triples_factory: TriplesFactory, - interaction: Interaction[ - torch.FloatTensor, - Tuple[torch.FloatTensor, torch.FloatTensor], - torch.FloatTensor, - ], - embedding_dim: int = 50, - relation_dim: Union[None, int, Sequence[int]] = None, - second_relation_dim: Union[None, int, Sequence[int]] = None, - loss: Optional[Loss] = None, - predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, - preferred_device: DeviceHint = None, - random_seed: Optional[int] = None, - embedding_specification: Optional[EmbeddingSpecification] = None, - relation_embedding_specification: Optional[EmbeddingSpecification] = None, - second_relation_embedding_specification: Optional[EmbeddingSpecification] = None, - ) -> None: - if relation_dim is None: - relation_dim = embedding_dim - if second_relation_dim is None: - second_relation_dim = relation_dim - super().__init__( - triples_factory=triples_factory, - interaction=interaction, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - predict_with_sigmoid=predict_with_sigmoid, - preferred_device=preferred_device, - random_seed=random_seed, - entity_representations=[ - Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - specification=embedding_specification, - ), - ], - relation_representations=[ - Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - specification=relation_embedding_specification, - ), - Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - specification=second_relation_embedding_specification, - ), - ], - ) - - -class TwoVectorEmbeddingModel(ERModel, autoreset=False): - """A KGEM that stores two :class:`pykeen.nn.Embedding` for each entities and relations. - - .. seealso:: - - - :class:`pykeen.models.KG2E` - - :class:`pykeen.models.TransD` - """ - - # TODO: Get rid of it - - def __init__( - self, - triples_factory: TriplesFactory, - interaction: Interaction[ - Tuple[torch.FloatTensor, torch.FloatTensor], - Tuple[torch.FloatTensor, torch.FloatTensor], - Tuple[torch.FloatTensor, torch.FloatTensor], - ], - embedding_dim: int = 50, - relation_dim: Optional[int] = None, - loss: Optional[Loss] = None, - predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, - preferred_device: DeviceHint = None, - random_seed: Optional[int] = None, - embedding_specification: Optional[EmbeddingSpecification] = None, - relation_embedding_specification: Optional[EmbeddingSpecification] = None, - second_embedding_specification: Optional[EmbeddingSpecification] = None, - second_relation_embedding_specification: Optional[EmbeddingSpecification] = None, - ) -> None: - if relation_dim is None: - relation_dim = embedding_dim - super().__init__( - triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - predict_with_sigmoid=predict_with_sigmoid, - preferred_device=preferred_device, - random_seed=random_seed, - entity_representations=[ - Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - specification=embedding_specification, - ), - Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - specification=second_embedding_specification, - ), - ], - relation_representations=[ - Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - specification=relation_embedding_specification, - ), - Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - specification=second_relation_embedding_specification, - ), - ], - interaction=interaction, - ) - - -class TwoSideEmbeddingModel(ERModel, autoreset=False): - """A KGEM with two sub-KGEMs that serve as a "forwards" and "backwards" model. - - Stores two :class:`pykeen.nn.Embedding` for each entities and relations. - - .. seealso:: :class:`pykeen.models.SimplE` - """ - - def __init__( - self, - triples_factory: TriplesFactory, - interaction: Interaction, - embedding_dim: int = 50, - relation_dim: Optional[int] = None, - loss: Optional[Loss] = None, - predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, - preferred_device: DeviceHint = None, - random_seed: Optional[int] = None, - embedding_specification: Optional[EmbeddingSpecification] = None, - relation_embedding_specification: Optional[EmbeddingSpecification] = None, - second_embedding_specification: Optional[EmbeddingSpecification] = None, - second_relation_embedding_specification: Optional[EmbeddingSpecification] = None, - ): - # TODO: Get rid of it - if relation_dim is None: - relation_dim = embedding_dim - super().__init__( - triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - predict_with_sigmoid=predict_with_sigmoid, - preferred_device=preferred_device, - random_seed=random_seed, - entity_representations=[ - Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - specification=embedding_specification, - ), - Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - specification=second_embedding_specification, - ), - ], - relation_representations=[ - Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - specification=relation_embedding_specification, - ), - Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - specification=second_relation_embedding_specification, - ), - ], - interaction=interaction, - ) - - def forward( - self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], - slice_size: Optional[int] = None, - slice_dim: Optional[str] = None, - ) -> torch.FloatTensor: # noqa: D102 - return 0.5 * sum( - self.interaction.score( - h_source.get_in_canonical_shape(indices=h_indices), - r_source.get_in_canonical_shape(indices=r_indices), - t_source.get_in_canonical_shape(indices=t_indices), - slice_size=slice_size, - slice_dim=slice_dim, - ) - for h_source, r_source, t_source in ( - (self.entity_representations[0], self.relation_representations[0], self.entity_representations[1]), - (self.entity_representations[1], self.relation_representations[1], self.entity_representations[0]), - ) - ) From 56001d4c6b76477d70cd5e065a025154a3568474 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:19:56 +0100 Subject: [PATCH 517/690] Fix DistMult and broken imports --- src/pykeen/models/__init__.py | 9 +-------- src/pykeen/nn/functional.py | 2 +- tests/test_models.py | 2 -- 3 files changed, 2 insertions(+), 11 deletions(-) diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index e09320c0bf..1763737405 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -8,10 +8,7 @@ from typing import Mapping, Set, Type, Union -from .base import ( # noqa:F401 - DoubleRelationEmbeddingModel, ERModel, Model, SingleVectorEmbeddingModel, - TwoSideEmbeddingModel, TwoVectorEmbeddingModel, -) +from .base import ERModel, Model from .multimodal import ComplExLiteral, DistMultLiteral, LiteralModel from .unimodal import ( ComplEx, @@ -68,10 +65,6 @@ _BASE_MODELS = { ERModel, - SingleVectorEmbeddingModel, - DoubleRelationEmbeddingModel, - TwoSideEmbeddingModel, - TwoVectorEmbeddingModel, LiteralModel, } diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 7f945ecf8f..a21307baf9 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -280,7 +280,7 @@ def distmult_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return tensor_product(h, r, t) + return tensor_product(h, r, t).sum(dim=-1) def ermlp_interaction( diff --git a/tests/test_models.py b/tests/test_models.py index c2e1219ac8..403f8f94dc 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -25,10 +25,8 @@ from pykeen.datasets.nations import NATIONS_TEST_PATH, NATIONS_TRAIN_PATH, Nations from pykeen.models import LiteralModel, _BASE_MODELS, _MODELS from pykeen.models.base import ( - DoubleRelationEmbeddingModel, ERModel, Model, - TwoVectorEmbeddingModel, _extend_batch, get_novelty_mask, ) From 19ef8290330e431533f44f1deb77f5ebdad5076d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:20:18 +0100 Subject: [PATCH 518/690] Fix broken annotations --- tests/test_models.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tests/test_models.py b/tests/test_models.py index 403f8f94dc..ff59366772 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -671,7 +671,6 @@ def _check_constraints(self): * Entity and relation embeddings have to have at most unit L2 norm. * Covariances have to have values between c_min and c_max """ - self.model: TwoVectorEmbeddingModel low = self.model.entity_representations[1].constrainer.keywords['min'] high = self.model.entity_representations[1].constrainer.keywords['max'] @@ -867,7 +866,6 @@ def _check_constraints(self): Entity embeddings have to have unit L2 norm. """ - self.model: DoubleRelationEmbeddingModel entity_norms = self.model.relation_representations[1](indices=None).norm(p=2, dim=-1) assert torch.allclose(entity_norms, torch.ones_like(entity_norms)) From 013adc0c3fcfaf21cc83a6b973510ff502ed8576 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:22:50 +0100 Subject: [PATCH 519/690] Fix _prepare_representation_module_list --- src/pykeen/models/base.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 8b5c123fc5..860d6b257e 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1105,6 +1105,7 @@ def _prepare_representation_module_list( representations = [ r if isinstance(r, RepresentationModule) else r.make(num_embeddings=num_embeddings) for r in representations + if representations is not None ] return nn.ModuleList(representations) From e3ce1cfe397dfe51664513a9e476ba1eadb78980 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:24:50 +0100 Subject: [PATCH 520/690] Fix _prepare_representation_module_list - again --- src/pykeen/models/base.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 860d6b257e..1f717047c1 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1100,13 +1100,13 @@ def _prepare_representation_module_list( ) -> Sequence[RepresentationModule]: """Normalize list of representations and wrap into nn.ModuleList.""" # Important: use ModuleList to ensure that Pytorch correctly handles their devices and parameters - if representations is not None and not isinstance(representations, Sequence): - representations = [representations] - representations = [ - r if isinstance(r, RepresentationModule) else r.make(num_embeddings=num_embeddings) - for r in representations - if representations is not None - ] + if representations is not None: + if not isinstance(representations, Sequence): + representations = [representations] + representations = [ + r if isinstance(r, RepresentationModule) else r.make(num_embeddings=num_embeddings) + for r in representations + ] return nn.ModuleList(representations) From 661e932bde68e3c788fb03fc6e8595a8afd405d4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:29:49 +0100 Subject: [PATCH 521/690] Fix SE shape --- src/pykeen/nn/functional.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index a21307baf9..f2a917f6f4 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -665,8 +665,8 @@ def structured_embedding_interaction( The scores. """ return negative_norm_of_sum( - (h @ r_h), - -(t @ r_t), + (h @ r_h.squeeze(dim=-3)), + -(t @ r_t.squeeze(dim=-3)), p=p, power_norm=power_norm, ) From 8314c757cf35afa5ab39a70731485553fb0deaa6 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 23 Nov 2020 21:31:35 +0100 Subject: [PATCH 522/690] Comment out unused code @mberr why was this not used? --- src/pykeen/nn/functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index a21307baf9..1ba1bb0338 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -159,7 +159,7 @@ def conve_interaction( # bind sizes num_heads = h.shape[1] num_relations = r.shape[2] - num_tails = t.shape[3] + # num_tails = t.shape[3] # FIXME unused! embedding_dim = h.shape[-1] # repeat if necessary, and concat head and relation, batch_size', num_input_channels, 2*height, width From bde23b0fe3becac1c0bce6ae796b5db196223373 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 23 Nov 2020 21:31:58 +0100 Subject: [PATCH 523/690] pass flake8 --- src/pykeen/models/base.py | 4 ++-- src/pykeen/models/unimodal/conv_kb.py | 2 +- src/pykeen/models/unimodal/kg2e.py | 8 ++++---- src/pykeen/models/unimodal/proj_e.py | 17 +++++++++++++++-- src/pykeen/models/unimodal/simple.py | 2 +- src/pykeen/nn/__init__.py | 2 +- src/pykeen/nn/functional.py | 16 +++++++--------- src/pykeen/nn/modules.py | 6 +++--- src/pykeen/utils.py | 2 +- 9 files changed, 35 insertions(+), 24 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 1f717047c1..f8e3921ff7 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1130,13 +1130,13 @@ def __init__( None, EmbeddingSpecification, RepresentationModule, - Sequence[Union[EmbeddingSpecification, RepresentationModule]] + Sequence[Union[EmbeddingSpecification, RepresentationModule]], ] = None, relation_representations: Union[ None, EmbeddingSpecification, RepresentationModule, - Sequence[Union[EmbeddingSpecification, RepresentationModule]] + Sequence[Union[EmbeddingSpecification, RepresentationModule]], ] = None, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 9a77b059ee..4d4a20c925 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -92,7 +92,7 @@ def __init__( num_filters=num_filters, ), entity_representations=EmbeddingSpecification( - embedding_dim=embedding_dim + embedding_dim=embedding_dim, ), relation_representations=EmbeddingSpecification( embedding_dim=embedding_dim, diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 1356d8f232..ff89882a22 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -61,13 +61,13 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, - loss: Optional[Loss] = None, - preferred_device: DeviceHint = None, - random_seed: Optional[int] = None, dist_similarity: Optional[str] = None, c_min: float = 0.05, c_max: float = 5., + loss: Optional[Loss] = None, + automatic_memory_optimization: Optional[bool] = None, + preferred_device: DeviceHint = None, + random_seed: Optional[int] = None, ) -> None: r"""Initialize KG2E. diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index 397b90eaf1..1e2cae2457 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -56,13 +56,26 @@ class ProjE(ERModel): def __init__( self, triples_factory: TriplesFactory, + # ProjE parameters embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, + inner_non_linearity: Optional[nn.Module] = None, + # Loss loss: Optional[Loss] = None, + # Model parameters preferred_device: DeviceHint = None, + automatic_memory_optimization: Optional[bool] = None, random_seed: Optional[int] = None, - inner_non_linearity: Optional[nn.Module] = None, ) -> None: + """Initialize :class:`ERModel` using :class:`ProjEInteraction`. + + :param triples_factory: + :param embedding_dim: + :param inner_non_linearity: + :param loss: + :param automatic_memory_optimization: + :param preferred_device: + :param random_seed: + """ super().__init__( triples_factory=triples_factory, interaction=ProjEInteraction( diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index c8c3791032..0adf5cb3f8 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -85,7 +85,7 @@ def __init__( EmbeddingSpecification( embedding_dim=embedding_dim, regularizer=regularizer, - ) + ), ], relation_representations=[ EmbeddingSpecification( diff --git a/src/pykeen/nn/__init__.py b/src/pykeen/nn/__init__.py index e4a41f8441..dada0a07ca 100644 --- a/src/pykeen/nn/__init__.py +++ b/src/pykeen/nn/__init__.py @@ -3,8 +3,8 @@ """PyKEEN internal "nn" module.""" from . import functional, init -from .representation import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule from .modules import Interaction +from .representation import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule __all__ = [ 'Embedding', diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 1ba1bb0338..312398648d 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -111,11 +111,11 @@ def complex_interaction( return tensor_sum(*( factor * tensor_product(hh, rr, tt).sum(dim=-1) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) @@ -131,8 +131,7 @@ def conve_interaction( hr2d: nn.Module, hr1d: nn.Module, ) -> torch.FloatTensor: - """ - Evaluate the ConvE interaction function. + """Evaluate the ConvE interaction function. :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. @@ -198,8 +197,7 @@ def convkb_interaction( hidden_dropout: nn.Dropout, linear: nn.Linear, ) -> torch.FloatTensor: - r""" - Evaluate the ConvKB interaction function. + r"""Evaluate the ConvKB interaction function. .. math:: W_L drop(act(W_C \ast ([h; r; t]) + b_C)) + b_L diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 848d25cbbc..c0cfbcb053 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -1104,14 +1104,14 @@ class SimplEInteraction( Interaction[ Tuple[torch.FloatTensor, torch.FloatTensor], Tuple[torch.FloatTensor, torch.FloatTensor], - Tuple[torch.FloatTensor, torch.FloatTensor] - ] + Tuple[torch.FloatTensor, torch.FloatTensor], + ], ): """Interaction function of SimplE.""" func = pkf.simple_interaction - def __init__(self, clamp_score: Optional[Union[float, Tuple[float, float]]] = None, ): + def __init__(self, clamp_score: Union[None, float, Tuple[float, float]] = None): super().__init__() if isinstance(clamp_score, float): clamp_score = (-clamp_score, clamp_score) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 3d1292cb01..51ead3f2db 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -506,7 +506,7 @@ def tensor_sum(*x: torch.FloatTensor) -> torch.FloatTensor: def tensor_product(*x: torch.FloatTensor) -> torch.FloatTensor: - """"Compute elementwise product of tensors in broadcastable shape.""" + """Compute elementwise product of tensors in broadcastable shape.""" # TODO: Optimize order return functools.reduce(operator.mul, x[1:], x[0]) From 9c58aeb75f9b02189aeea3d906f15df9354e6ec3 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:33:16 +0100 Subject: [PATCH 524/690] Fix ConvKB --- src/pykeen/nn/functional.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index f2a917f6f4..120d795e18 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -231,6 +231,7 @@ def convkb_interaction( # compute conv(stack(h, r, t)) # prepare input shapes for broadcasting + # (b, h, r, t, 1, d) h = h.unsqueeze(dim=-2) r = r.unsqueeze(dim=-2) t = t.unsqueeze(dim=-2) @@ -238,8 +239,7 @@ def convkb_interaction( # conv.weight.shape = (C_out, C_in, kernel_size[0], kernel_size[1]) # here, kernel_size = (1, 3), C_in = 1, C_out = num_filters # -> conv_head, conv_rel, conv_tail shapes: (num_filters,) - # conv_head, conv_rel, conv_tail = conv.weight[:, 0, 0, :].t() - # conv_bias = conv.bias.view(1, 1, 1, 1, num_filters, 1) + # reshape to (1, 1, 1, 1, f, 1) conv_head, conv_rel, conv_tail, conv_bias = [ c.view(1, 1, 1, 1, num_filters, 1) for c in list(conv.weight[:, 0, 0, :].t()) + [conv.bias] @@ -257,7 +257,7 @@ def convkb_interaction( x = hidden_dropout(x) # Linear layer for final scores; use flattened representations, shape: (b, h, r, t, d * f) - x = x.view(x.shape[:-2], embedding_dim * num_filters) + x = x.view(*x.shape[:-2], embedding_dim * num_filters) x = linear(x) return x.squeeze(dim=-1) From 40002d222fce1a9043b36a3f5e4c409b94272555 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:34:33 +0100 Subject: [PATCH 525/690] inline sizes for ConvE --- src/pykeen/nn/functional.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 4dfc9fd98f..249a5859a4 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -155,12 +155,6 @@ def conve_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # bind sizes - num_heads = h.shape[1] - num_relations = r.shape[2] - # num_tails = t.shape[3] # FIXME unused! - embedding_dim = h.shape[-1] - # repeat if necessary, and concat head and relation, batch_size', num_input_channels, 2*height, width # with batch_size' = batch_size * num_heads * num_relations x = broadcast_cat( @@ -177,7 +171,7 @@ def conve_interaction( x = hr1d(x) # reshape: (batch_size', embedding_dim) -> (b, h, r, 1, d) - x = x.view(x.shape[0], num_heads, num_relations, 1, embedding_dim) + x = x.view(x.shape[0], h.shape[1], r.shape[2], 1, h.shape[-1]) # For efficient calculation, each of the convolved [h, r] rows has only to be multiplied with one t row # output_shape: (batch_size, num_heads, num_relations, num_tails) From b7b9e71d180092f735430611effc9cc02c515936 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:35:50 +0100 Subject: [PATCH 526/690] Fix ConvE init --- src/pykeen/models/unimodal/conv_e.py | 44 ++++++++++++---------------- 1 file changed, 19 insertions(+), 25 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index bd19dbde7f..7da25f4a53 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -134,31 +134,6 @@ def __init__( ) super().__init__( triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, - loss=loss, - preferred_device=preferred_device, - random_seed=random_seed, - entity_representations=[ - Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - specification=EmbeddingSpecification( - initializer=xavier_normal_, - ), - ), - # ConvE uses one bias for each entity - Embedding.from_specification( - num_embeddings=triples_factory.num_entities, - specification=EmbeddingSpecification( - initializer=nn.init.zeros_, - ), - ), - ], - relation_representations=Embedding.from_specification( - num_embeddings=triples_factory.num_relations, - specification=EmbeddingSpecification( - initializer=xavier_normal_, - ), - ), interaction=ConvEInteraction( input_channels=input_channels, output_channels=output_channels, @@ -172,6 +147,25 @@ def __init__( embedding_dim=embedding_dim, apply_batch_normalization=apply_batch_normalization, ), + entity_representations=[ + EmbeddingSpecification( + embedding_dim=embedding_dim, + initializer=xavier_normal_, + ), + # ConvE uses one bias for each entity + EmbeddingSpecification( + embedding_dim=1, + initializer=nn.init.zeros_, + ), + ], + relation_representations=EmbeddingSpecification( + embedding_dim=embedding_dim, + initializer=xavier_normal_, + ), + automatic_memory_optimization=automatic_memory_optimization, + loss=loss, + preferred_device=preferred_device, + random_seed=random_seed, ) def forward( From 6eda30a74feed9ae9fa37ca2b7f036c59e9efe02 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:37:36 +0100 Subject: [PATCH 527/690] Fix ConvE revert reshaping --- src/pykeen/nn/functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 249a5859a4..f3f0b55a74 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -171,7 +171,7 @@ def conve_interaction( x = hr1d(x) # reshape: (batch_size', embedding_dim) -> (b, h, r, 1, d) - x = x.view(x.shape[0], h.shape[1], r.shape[2], 1, h.shape[-1]) + x = x.view(-1, h.shape[1], r.shape[2], 1, h.shape[-1]) # For efficient calculation, each of the convolved [h, r] rows has only to be multiplied with one t row # output_shape: (batch_size, num_heads, num_relations, num_tails) From 59e3bbf64c20972c53a0ca70a32526707f22fa6f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:41:31 +0100 Subject: [PATCH 528/690] Fix TransR functional form --- src/pykeen/nn/functional.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index f3f0b55a74..d19b35998a 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -793,7 +793,7 @@ def transr_interaction( Head embeddings. :param r: shape: (batch_size, 1, num_relations, 1, d_r) Relation embeddings. - :param m_r: shape: (batch_size, 1, num_relations, 1, d_e, d_r) + :param m_r: shape: (batch_size, 1, num_relations, 1, d_r, d_e) The relation specific linear transformations. :param t: shape: (batch_size, 1, 1, num_tails, d_e) Tail embeddings. @@ -806,8 +806,8 @@ def transr_interaction( The scores. """ # project to relation specific subspace and ensure constraints - h_bot = clamp_norm(h @ m_r, p=2, dim=-1, maxnorm=1.) - t_bot = clamp_norm(t @ m_r, p=2, dim=-1, maxnorm=1.) + h_bot = clamp_norm(m_r @ h.unsqueeze(dim=-1), p=2, dim=-1, maxnorm=1.).squeeze(dim=-1) + t_bot = clamp_norm(m_r @ t.unsqueeze(dim=-1), p=2, dim=-1, maxnorm=1.).squeeze(dim=-1) return negative_norm_of_sum(h_bot, r, -t_bot, p=p, power_norm=power_norm) From 7b40cc8fc0e8e087793fed1a3e59fbbb52b212a8 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 23 Nov 2020 21:43:27 +0100 Subject: [PATCH 529/690] Add generics to LiteralModel --- src/pykeen/models/multimodal/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/multimodal/base.py b/src/pykeen/models/multimodal/base.py index 03e0ded7e6..c57b8de6cc 100644 --- a/src/pykeen/models/multimodal/base.py +++ b/src/pykeen/models/multimodal/base.py @@ -18,7 +18,7 @@ from ...typing import Representation # noqa -class LiteralModel(ERModel, autoreset=False): +class LiteralModel(ERModel[HeadRepresentation, RelationRepresentation, TailRepresentation], autoreset=False): """Base class for models with entity literals.""" def __init__( From 1192bbf27bb3e46c279e4ab7cc108f1a02926e29 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 23 Nov 2020 21:43:36 +0100 Subject: [PATCH 530/690] Style cleanup --- src/pykeen/models/multimodal/base.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/multimodal/base.py b/src/pykeen/models/multimodal/base.py index c57b8de6cc..6c7efa7a4b 100644 --- a/src/pykeen/models/multimodal/base.py +++ b/src/pykeen/models/multimodal/base.py @@ -1,14 +1,17 @@ +# -*- coding: utf-8 -*- + """Base classes for multi-modal models.""" + from typing import Optional, TYPE_CHECKING import torch from torch import nn -from pykeen.losses import Loss -from pykeen.models.base import ERModel -from pykeen.nn import Embedding, EmbeddingSpecification, Interaction, LiteralRepresentations -from pykeen.triples import TriplesNumericLiteralsFactory -from pykeen.typing import DeviceHint, HeadRepresentation, RelationRepresentation, TailRepresentation +from ..base import ERModel +from ...losses import Loss +from ...nn import Embedding, EmbeddingSpecification, Interaction, LiteralRepresentations +from ...triples import TriplesNumericLiteralsFactory +from ...typing import DeviceHint, HeadRepresentation, RelationRepresentation, TailRepresentation __all__ = [ "LiteralModel", From a3c015fc31dd55d0ccb75ca93f11ab109148692a Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 23 Nov 2020 21:44:02 +0100 Subject: [PATCH 531/690] Update handling fo embedding dimension in literal models --- src/pykeen/models/multimodal/base.py | 1 - src/pykeen/models/multimodal/complex_literal.py | 5 ++++- src/pykeen/models/multimodal/distmult_literal.py | 3 ++- 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/src/pykeen/models/multimodal/base.py b/src/pykeen/models/multimodal/base.py index 6c7efa7a4b..23ac535ab6 100644 --- a/src/pykeen/models/multimodal/base.py +++ b/src/pykeen/models/multimodal/base.py @@ -27,7 +27,6 @@ class LiteralModel(ERModel[HeadRepresentation, RelationRepresentation, TailRepre def __init__( self, triples_factory: TriplesNumericLiteralsFactory, - embedding_dim: int, interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], combination: nn.Module, entity_specification: Optional[EmbeddingSpecification] = None, diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index c9600a87d6..8d57fa8760 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -86,7 +86,6 @@ def __init__( """Initialize the model.""" super().__init__( triples_factory=triples_factory, - embedding_dim=2 * embedding_dim, # complex interaction=ComplExInteraction(), combination=ComplexLiteralCombination( embedding_dim=embedding_dim, @@ -94,10 +93,14 @@ def __init__( dropout=input_dropout, ), entity_specification=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=nn.init.xavier_normal_, + dtype=torch.complex64, ), relation_specification=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=nn.init.xavier_normal_, + dtype=torch.complex64, ), loss=loss, predict_with_sigmoid=predict_with_sigmoid, diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 7236aecf0e..48f265e0a1 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -47,16 +47,17 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, - embedding_dim=embedding_dim, interaction=DistMultInteraction(), combination=nn.Sequential( nn.Linear(embedding_dim + triples_factory.numeric_literals.shape[1], embedding_dim), nn.Dropout(input_dropout), ), entity_specification=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=nn.init.xavier_normal_, ), relation_specification=EmbeddingSpecification( + embedding_dim=embedding_dim, initializer=nn.init.xavier_normal_, ), loss=loss, From ad0fe00d4ad7cdf87a73f3fefdd806109965e111 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:45:06 +0100 Subject: [PATCH 532/690] Fix usage of chain_matmul it only works for 2D matrices --- src/pykeen/nn/functional.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index d19b35998a..d2f70bee40 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -542,11 +542,7 @@ def rescal_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return torch.chain_matmul( - h.unsqueeze(dim=-2), - r, - t.unsqueeze(dim=-1), - ).view(-1, h.shape[1], r.shape[2], t.shape[3]) + return extended_einsum("bhrtd,bhrtde,bhrte->bhrt", h, r, t) def rotate_interaction( From 8caba1e7d59a36a2390a1c8d79f3234059a380b7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:47:29 +0100 Subject: [PATCH 533/690] Fix NTN einsum eq --- src/pykeen/nn/functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index d2f70bee40..a390e785bb 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -466,7 +466,7 @@ def ntn_interaction( The scores. """ x = activation(tensor_sum( - extended_einsum("bhrtd,brkde,bhrte->bhrtk", h, w, t), + extended_einsum("bhrtd,bhrtkde,bhrte->bhrtk", h, w, t), (vh @ h.unsqueeze(dim=-1)).squeeze(dim=-1), (vt @ t.unsqueeze(dim=-1)).squeeze(dim=-1), b, From 6d829f5b3a7ce00235bc1fe808daeb5c42b46721 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:54:05 +0100 Subject: [PATCH 534/690] Fix TransD --- src/pykeen/models/unimodal/trans_d.py | 30 ++++++++++++++++----------- 1 file changed, 18 insertions(+), 12 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 1a56ed794b..dffd7fbd37 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -69,18 +69,24 @@ def __init__( super().__init__( triples_factory=triples_factory, interaction=TransDInteraction(p=2, power_norm=True), - entity_representations=EmbeddingSpecification( - embedding_dim=embedding_dim, - initializer=xavier_normal_, - constrainer=clamp_norm, # type: ignore - constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), - ), - relation_representations=EmbeddingSpecification( - embedding_dim=relation_dim, - initializer=xavier_normal_, - constrainer=clamp_norm, # type: ignore - constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), - ), + entity_representations=[ + EmbeddingSpecification( + embedding_dim=embedding_dim, + initializer=xavier_normal_, + constrainer=clamp_norm, # type: ignore + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + ) + for _ in range(2) + ], + relation_representations=[ + EmbeddingSpecification( + embedding_dim=relation_dim, + initializer=xavier_normal_, + constrainer=clamp_norm, # type: ignore + constrainer_kwargs=dict(maxnorm=1., p=2, dim=-1), + ) + for _ in range(2) + ], automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, From 24a6bf7cc588798bd88364a5794b3e1fc66a39bc Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:54:16 +0100 Subject: [PATCH 535/690] Fix SimplEInteraction --- src/pykeen/nn/modules.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index c0cfbcb053..7d4d977592 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -1115,7 +1115,7 @@ def __init__(self, clamp_score: Union[None, float, Tuple[float, float]] = None): super().__init__() if isinstance(clamp_score, float): clamp_score = (-clamp_score, clamp_score) - self.clamp = clamp_score + self.clamp_score = clamp_score def _prepare_state_for_functional(self) -> MutableMapping[str, Any]: # noqa: D102 return dict(clamp=self.clamp_score) From b9f83d4468ffbb87e6bf35b1c89fd0f4d8fc0a32 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 21:55:03 +0100 Subject: [PATCH 536/690] Fix RGCNRepresentations --- src/pykeen/models/unimodal/rgcn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index c04a595033..4007462e8f 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -131,7 +131,7 @@ def __init__( buffer_messages: bool = True, base_representations: Optional[RepresentationModule] = None, ): - super().__init__(shape=(self.embedding_dim,), max_id=triples_factory.num_entities) + super().__init__(shape=(embedding_dim,), max_id=triples_factory.num_entities) self.triples_factory = triples_factory From d5a4595716f1adbb9a3d6e77a88af40b5e210218 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 22:10:20 +0100 Subject: [PATCH 537/690] Add verification to representations utility --- src/pykeen/models/base.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index f8e3921ff7..bdec43d7de 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1107,6 +1107,15 @@ def _prepare_representation_module_list( r if isinstance(r, RepresentationModule) else r.make(num_embeddings=num_embeddings) for r in representations ] + for r in representations: + if r.max_id < num_embeddings: + raise ValueError(f"{r} only provides {r.max_id} representations, but should provide {num_embeddings}.") + elif r.max_id > num_embeddings: + logger.warning( + f"{r} provides {r.max_id} representations, although only {num_embeddings} are needed. While this " + f"is not necessarily wrong, it can indicate an error where the number of representations was chosen" + f" wrong." + ) return nn.ModuleList(representations) From 53c304515328fc9a0590d041a52f51ddc1930632 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 23 Nov 2020 22:16:48 +0100 Subject: [PATCH 538/690] Pass flake8 and mypy --- .../models/multimodal/complex_literal.py | 6 ++--- src/pykeen/models/unimodal/conv_e.py | 2 +- src/pykeen/nn/representation.py | 24 +++++++++++++------ 3 files changed, 21 insertions(+), 11 deletions(-) diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index 8d57fa8760..008310fc74 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -1,6 +1,6 @@ # -*- coding: utf-8 -*- -"""Implementation of the ComplexLiteral model based on the local closed world assumption (LCWA) training approach.""" +"""Implementation of the ComplexLiteral model.""" from typing import Any, ClassVar, Mapping, Optional @@ -20,7 +20,7 @@ ] -class ComplexLiteralCombination(nn.Module): +class ComplExLiteralCombination(nn.Module): """Separately transform real and imaginary part.""" def __init__( @@ -87,7 +87,7 @@ def __init__( super().__init__( triples_factory=triples_factory, interaction=ComplExInteraction(), - combination=ComplexLiteralCombination( + combination=ComplExLiteralCombination( embedding_dim=embedding_dim, num_of_literals=triples_factory.numeric_literals.shape[-1], dropout=input_dropout, diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 7da25f4a53..5a29b01b24 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -10,7 +10,7 @@ from .. import ERModel from ...losses import BCEAfterSigmoidLoss, Loss -from ...nn import Embedding, EmbeddingSpecification +from ...nn import EmbeddingSpecification from ...nn.init import xavier_normal_ from ...nn.modules import ConvEInteraction from ...triples import TriplesFactory diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index 98a28753d0..6db04d3df6 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -6,7 +6,7 @@ import functools import logging from abc import ABC, abstractmethod -from typing import Any, Mapping, Optional, Sequence, Union +from typing import Any, Iterable, Mapping, Optional, Sequence, Tuple, Union import numpy import torch @@ -47,9 +47,9 @@ class RepresentationModule(nn.Module, ABC): #: The maximum admissible ID (excl.) max_id: int - def __init__(self, shape: Sequence[int], max_id: int): + def __init__(self, shape: Iterable[int], max_id: int): super().__init__() - self.shape = shape + self.shape = tuple(shape) self.max_id = max_id @abstractmethod @@ -93,6 +93,7 @@ def get_in_canonical_shape( :return: shape: (batch_size, d1, d2, d3, *self.shape) """ dim = _normalize_dim(dim=dim) + r_shape: Tuple[int, ...] if indices is None: x = self(indices=indices) r_shape = (1, self.max_id) @@ -210,19 +211,28 @@ def __init__( """ if shape is None and embedding_dim is None: raise ValueError('Missing both, shape and embedding_dim') - elif shape is None: + elif shape is not None and embedding_dim is not None: + raise ValueError('Provided both, shape and embedding_dim') + elif shape is None and embedding_dim is not None: shape = (embedding_dim,) - elif embedding_dim is None: + elif isinstance(shape, int) and embedding_dim is None: + embedding_dim = shape + shape = (shape,) + elif isinstance(shape, Sequence) and embedding_dim is None: + shape = tuple(shape) embedding_dim = numpy.prod(shape) else: - raise ValueError('Provided both, shape and embedding_dim') + raise TypeError(f'Invalid type for shape: ({type(shape)}) {shape}') + + assert isinstance(shape, tuple) + assert isinstance(embedding_dim, int) + if dtype is None: dtype = torch.get_default_dtype() # work-around until full complex support # TODO: verify that this is our understanding of complex! if dtype.is_complex: - shape = tuple(shape) shape = shape[:-1] + (2 * shape[-1],) embedding_dim = embedding_dim * 2 super().__init__(shape=shape, max_id=num_embeddings) From a744d14fafe6313b6c7f408068bf0abf236f2864 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 23 Nov 2020 22:27:16 +0100 Subject: [PATCH 539/690] hsh --- .github/workflows/tests.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 049f22e7c4..e0f6e89aef 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -26,7 +26,7 @@ jobs: run: tox -e flake8 - name: Check typing with MyPy - env: tox -e mypy + run: tox -e mypy - name: Check package metadata with Pyroma run: tox -e pyroma From da4a2590b52a4d14e7bdf84ed47483a28fa91eaa Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 22:27:56 +0100 Subject: [PATCH 540/690] Add shapes verification to passed representations --- src/pykeen/models/base.py | 50 +++++++++++++++++++++------------ src/pykeen/nn/modules.py | 3 ++ src/pykeen/nn/representation.py | 2 +- src/pykeen/utils.py | 12 ++++---- 4 files changed, 43 insertions(+), 24 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index bdec43d7de..7257b56a33 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -5,6 +5,7 @@ import functools import itertools as itt import logging +import string from abc import ABC, abstractmethod from collections import defaultdict from operator import itemgetter @@ -24,7 +25,7 @@ from ..regularizers import Regularizer, collect_regularization_terms from ..triples import TriplesFactory from ..typing import DeviceHint, HeadRepresentation, MappedTriples, RelationRepresentation, TailRepresentation -from ..utils import NoRandomSeedNecessary, resolve_device, set_random_seed +from ..utils import NoRandomSeedNecessary, check_shapes, resolve_device, set_random_seed if TYPE_CHECKING: from ..typing import Representation # noqa @@ -1097,26 +1098,37 @@ def _prepare_representation_module_list( Sequence[Union[EmbeddingSpecification, RepresentationModule]], ], num_embeddings: int, + shapes: Sequence[str], ) -> Sequence[RepresentationModule]: """Normalize list of representations and wrap into nn.ModuleList.""" # Important: use ModuleList to ensure that Pytorch correctly handles their devices and parameters - if representations is not None: - if not isinstance(representations, Sequence): - representations = [representations] - representations = [ - r if isinstance(r, RepresentationModule) else r.make(num_embeddings=num_embeddings) - for r in representations - ] - for r in representations: - if r.max_id < num_embeddings: - raise ValueError(f"{r} only provides {r.max_id} representations, but should provide {num_embeddings}.") - elif r.max_id > num_embeddings: - logger.warning( - f"{r} provides {r.max_id} representations, although only {num_embeddings} are needed. While this " - f"is not necessarily wrong, it can indicate an error where the number of representations was chosen" - f" wrong." - ) - return nn.ModuleList(representations) + if representations is None: + representations = [] + if not isinstance(representations, Sequence): + representations = [representations] + if len(representations) != len(shapes): + raise ValueError( + f"Interaction function requires {len(shapes)} representations, but {len(representations)} were given." + ) + modules = [] + for r in representations: + if not isinstance(r, RepresentationModule): + assert isinstance(r, EmbeddingSpecification) + r = r.make(num_embeddings=num_embeddings) + if r.max_id < num_embeddings: + raise ValueError(f"{r} only provides {r.max_id} representations, but should provide {num_embeddings}.") + elif r.max_id > num_embeddings: + logger.warning( + f"{r} provides {r.max_id} representations, although only {num_embeddings} are needed. While this " + f"is not necessarily wrong, it can indicate an error where the number of representations was chosen" + f" wrong." + ) + modules.append(r) + check_shapes(*zip( + (r.shape for r in modules), + shapes + ), raise_on_errors=True) + return nn.ModuleList(modules) class ERModel(Model, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation], autoreset=False): @@ -1183,10 +1195,12 @@ def __init__( self.entity_representations = _prepare_representation_module_list( representations=entity_representations, num_embeddings=triples_factory.num_entities, + shapes=interaction.entity_shape, ) self.relation_representations = _prepare_representation_module_list( representations=relation_representations, num_embeddings=triples_factory.num_relations, + shapes=interaction.relation_shape, ) self.interaction = interaction # Comment: it is important that the regularizers are stored in a module list, in order to appear in diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 7d4d977592..0209682cfc 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -257,6 +257,7 @@ def score( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ + # TODO: fix shapes return self._forward_slicing_wrapper(h=h, r=r, t=t, slice_size=slice_size, slice_dim=slice_dim) def _score( @@ -289,6 +290,7 @@ def _score( r = self._add_dim(*r, dim=self.BATCH_DIM if r_prefix == "n" else self.NUM_DIM) t = self._add_dim(*t, dim=self.BATCH_DIM if t_prefix == "n" else self.NUM_DIM) + # TODO: fix shapes scores = self._forward_slicing_wrapper(h=h, r=r, t=t, slice_dim=slice_dim, slice_size=slice_size) remove_dims = [ @@ -334,6 +336,7 @@ def _forward_slicing_wrapper( :raises ValueError: If slice_dim is invalid. """ + # TODO: fix shapes if slice_size is None: scores = self(h=h, r=r, t=t) elif slice_dim == "h": diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index 6db04d3df6..e75a421a4c 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -220,7 +220,7 @@ def __init__( shape = (shape,) elif isinstance(shape, Sequence) and embedding_dim is None: shape = tuple(shape) - embedding_dim = numpy.prod(shape) + embedding_dim = int(numpy.prod(shape)) else: raise TypeError(f'Invalid type for shape: ({type(shape)}) {shape}') diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 51ead3f2db..0ac11f2dcc 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -392,7 +392,7 @@ def random_non_negative_int() -> int: def check_shapes( - *x: Tuple[torch.Tensor, str], + *x: Tuple[Union[torch.Tensor, Tuple[int, ...]], str], raise_on_errors: bool = True, ) -> bool: """ @@ -413,11 +413,13 @@ def check_shapes( """ dims: Dict[str, Tuple[int, ...]] = dict() errors = [] - for tensor, shape in x: - if tensor.ndimension() != len(shape): - errors.append(f"Invalid number of dimensions: {tensor.shape} vs. {shape}") + for actual_shape, shape in x: + if torch.is_tensor(actual_shape): + actual_shape = actual_shape.shape + if len(actual_shape) != len(shape): + errors.append(f"Invalid number of dimensions: {actual_shape} vs. {shape}") continue - for dim, name in zip(tensor.shape, shape): + for dim, name in zip(actual_shape, shape): exp_dim = dims.get(name) if exp_dim is not None and exp_dim != dim: errors.append(f"{name}: {dim} vs. {exp_dim}") From 10ef4796a1fdd4a69fa3ec590ff78b626fe16ef1 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 22:29:49 +0100 Subject: [PATCH 541/690] Fix superclass order for ERModel @cthoyt please check --- src/pykeen/models/base.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 7257b56a33..e9ab786f6f 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1131,7 +1131,7 @@ def _prepare_representation_module_list( return nn.ModuleList(modules) -class ERModel(Model, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation], autoreset=False): +class ERModel(Generic[HeadRepresentation, RelationRepresentation, TailRepresentation], Model, autoreset=False): """A commonly useful base for KGEMs using embeddings and interaction modules.""" #: The entity representations From cbfab6a2187de3020d7492e96ea5b7850c140c03 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 22:31:20 +0100 Subject: [PATCH 542/690] Improve error messages --- src/pykeen/models/base.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index e9ab786f6f..f5cd80d02e 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -5,7 +5,6 @@ import functools import itertools as itt import logging -import string from abc import ABC, abstractmethod from collections import defaultdict from operator import itemgetter @@ -1099,6 +1098,7 @@ def _prepare_representation_module_list( ], num_embeddings: int, shapes: Sequence[str], + label: str, ) -> Sequence[RepresentationModule]: """Normalize list of representations and wrap into nn.ModuleList.""" # Important: use ModuleList to ensure that Pytorch correctly handles their devices and parameters @@ -1108,7 +1108,8 @@ def _prepare_representation_module_list( representations = [representations] if len(representations) != len(shapes): raise ValueError( - f"Interaction function requires {len(shapes)} representations, but {len(representations)} were given." + f"Interaction function requires {len(shapes)} {label} representations, but " + f"{len(representations)} were given." ) modules = [] for r in representations: @@ -1116,12 +1117,14 @@ def _prepare_representation_module_list( assert isinstance(r, EmbeddingSpecification) r = r.make(num_embeddings=num_embeddings) if r.max_id < num_embeddings: - raise ValueError(f"{r} only provides {r.max_id} representations, but should provide {num_embeddings}.") + raise ValueError( + f"{r} only provides {r.max_id} {label} representations, but should provide {num_embeddings}." + ) elif r.max_id > num_embeddings: logger.warning( - f"{r} provides {r.max_id} representations, although only {num_embeddings} are needed. While this " - f"is not necessarily wrong, it can indicate an error where the number of representations was chosen" - f" wrong." + f"{r} provides {r.max_id} {label} representations, although only {num_embeddings} are needed." + f"While this is not necessarily wrong, it can indicate an error where the number of {label} " + f"representations was chosen wrong." ) modules.append(r) check_shapes(*zip( @@ -1196,11 +1199,13 @@ def __init__( representations=entity_representations, num_embeddings=triples_factory.num_entities, shapes=interaction.entity_shape, + label="entity", ) self.relation_representations = _prepare_representation_module_list( representations=relation_representations, num_embeddings=triples_factory.num_relations, shapes=interaction.relation_shape, + label="relation", ) self.interaction = interaction # Comment: it is important that the regularizers are stored in a module list, in order to appear in From 9cb56b0f80b3f056ed0c34a99ea9366107c15091 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 22:39:27 +0100 Subject: [PATCH 543/690] Extract canonical shape utility --- src/pykeen/nn/modules.py | 14 +++++------- src/pykeen/nn/representation.py | 39 +++++++++++++++++++++++++++++---- src/pykeen/utils.py | 8 ++++++- 3 files changed, 47 insertions(+), 14 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 0209682cfc..3058ec409b 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -13,7 +13,7 @@ from . import functional as pkf from ..typing import HeadRepresentation, RelationRepresentation, Representation, TailRepresentation -from ..utils import check_shapes +from ..utils import check_shapes, upgrade_to_sequence __all__ = [ # Base Classes @@ -45,12 +45,8 @@ logger = logging.getLogger(__name__) -def _upgrade_to_sequence(x: Union[FloatTensor, Sequence[FloatTensor]]) -> Sequence[FloatTensor]: - return x if isinstance(x, Sequence) else (x,) - - def _ensure_tuple(*x: Union[Representation, Sequence[Representation]]) -> Sequence[Sequence[Representation]]: - return tuple(_upgrade_to_sequence(xx) for xx in x) + return tuple(upgrade_to_sequence(xx) for xx in x) def _unpack_singletons(*xs: Tuple) -> Sequence[Tuple]: @@ -280,9 +276,9 @@ def _score( slice_dim: Optional[str] = slice_dims[0] if len(slice_dims) == 1 else None # FIXME typing does not work well for this - h = _upgrade_to_sequence(h) - r = _upgrade_to_sequence(r) - t = _upgrade_to_sequence(t) + h = upgrade_to_sequence(h) + r = upgrade_to_sequence(r) + t = upgrade_to_sequence(t) assert self._check_shapes(h=h, r=r, t=t, h_prefix=h_prefix, r_prefix=r_prefix, t_prefix=t_prefix) # prepare input to generic score function: bh*, br*, bt* diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index e75a421a4c..3e6904a756 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -13,6 +13,7 @@ import torch.nn from torch import nn +from ..utils import upgrade_to_sequence from ..regularizers import Regularizer from ..typing import Constrainer, Initializer, Normalizer @@ -38,6 +39,39 @@ def _normalize_dim(dim: Union[int, str]) -> int: return DIMS[dim.lower()[0]] +def convert_to_canonical_shape( + x: torch.FloatTensor, + dim: Union[int, str], + num: Optional[int] = None, + batch_size: int = 1, + suffix_shape: Union[int, Sequence[int, ...]] = -1, +) -> torch.FloatTensor: + """ + Convert a tensor to canonical shape. + + :param x: + The tensor in compatible shape. + :param dim: + The "num" dimension. + :param batch_size: + The batch size. + :param num: + The number. + :param suffix_shape: + The suffix shape. + + :return: shape: (batch_size, num_heads, num_relations, num_tails, ``*``) + A tensor in canonical shape. + """ + if num is None: + num = x.shape[0] + suffix_shape = upgrade_to_sequence(suffix_shape) + shape = [batch_size, 1, 1, 1] + dim = _normalize_dim(dim=dim) + shape[dim] = num + return x.view(*shape, *suffix_shape) + + class RepresentationModule(nn.Module, ABC): """A base class for obtaining representations for entities/relations.""" @@ -92,7 +126,6 @@ def get_in_canonical_shape( :return: shape: (batch_size, d1, d2, d3, *self.shape) """ - dim = _normalize_dim(dim=dim) r_shape: Tuple[int, ...] if indices is None: x = self(indices=indices) @@ -105,9 +138,7 @@ def get_in_canonical_shape( r_shape = tuple(indices.shape) if len(r_shape) < 2: r_shape = r_shape + (1,) - shape = [r_shape[0], 1, 1, 1] - shape[dim] = r_shape[1] - return x.view(*shape, *self.shape) + return convert_to_canonical_shape(x=x, dim=dim, num=r_shape[1], batch_size=r_shape[0], suffix_shape=self.shape) def reset_parameters(self) -> None: """Reset the module's parameters.""" diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 0ac11f2dcc..701c428bf1 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -10,12 +10,13 @@ import random from abc import ABC, abstractmethod from io import BytesIO -from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Tuple, Type, TypeVar, Union +from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union import numpy as np import pandas as pd import torch import torch.nn +from torch._C import FloatTensor from torch.nn import functional from .typing import DeviceHint @@ -632,3 +633,8 @@ def pop_only(elements: Iterable[X]) -> X: def strip_dim(*x): """Strip the last dimension.""" return [xx.view(xx.shape[2:]) for xx in x] + + +def upgrade_to_sequence(x: Union[X, Sequence[X]]) -> Sequence[X]: + """Ensure that the input is a sequence.""" + return x if isinstance(x, Sequence) else (x,) From 93f3bfa0640aa57cc8ca10063fcacb61aa0769e7 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 23 Nov 2020 23:06:24 +0100 Subject: [PATCH 544/690] Pass flake8 / mypy --- src/pykeen/models/base.py | 8 ++++---- src/pykeen/nn/representation.py | 4 ++-- src/pykeen/utils.py | 4 ++-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index f5cd80d02e..975131a331 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1109,7 +1109,7 @@ def _prepare_representation_module_list( if len(representations) != len(shapes): raise ValueError( f"Interaction function requires {len(shapes)} {label} representations, but " - f"{len(representations)} were given." + f"{len(representations)} were given.", ) modules = [] for r in representations: @@ -1118,18 +1118,18 @@ def _prepare_representation_module_list( r = r.make(num_embeddings=num_embeddings) if r.max_id < num_embeddings: raise ValueError( - f"{r} only provides {r.max_id} {label} representations, but should provide {num_embeddings}." + f"{r} only provides {r.max_id} {label} representations, but should provide {num_embeddings}.", ) elif r.max_id > num_embeddings: logger.warning( f"{r} provides {r.max_id} {label} representations, although only {num_embeddings} are needed." f"While this is not necessarily wrong, it can indicate an error where the number of {label} " - f"representations was chosen wrong." + f"representations was chosen wrong.", ) modules.append(r) check_shapes(*zip( (r.shape for r in modules), - shapes + shapes, ), raise_on_errors=True) return nn.ModuleList(modules) diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index 3e6904a756..8080fca0e2 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -13,9 +13,9 @@ import torch.nn from torch import nn -from ..utils import upgrade_to_sequence from ..regularizers import Regularizer from ..typing import Constrainer, Initializer, Normalizer +from ..utils import upgrade_to_sequence __all__ = [ 'RepresentationModule', @@ -44,7 +44,7 @@ def convert_to_canonical_shape( dim: Union[int, str], num: Optional[int] = None, batch_size: int = 1, - suffix_shape: Union[int, Sequence[int, ...]] = -1, + suffix_shape: Union[int, Sequence[int]] = -1, ) -> torch.FloatTensor: """ Convert a tensor to canonical shape. diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 701c428bf1..506b0584fc 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -8,6 +8,7 @@ import logging import operator import random +import typing from abc import ABC, abstractmethod from io import BytesIO from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union @@ -16,7 +17,6 @@ import pandas as pd import torch import torch.nn -from torch._C import FloatTensor from torch.nn import functional from .typing import DeviceHint @@ -415,7 +415,7 @@ def check_shapes( dims: Dict[str, Tuple[int, ...]] = dict() errors = [] for actual_shape, shape in x: - if torch.is_tensor(actual_shape): + if isinstance(actual_shape, torch.Tensor): actual_shape = actual_shape.shape if len(actual_shape) != len(shape): errors.append(f"Invalid number of dimensions: {actual_shape} vs. {shape}") From 614759bb0915897b2fd3e0ddfcd2f61bd593a300 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 23:07:53 +0100 Subject: [PATCH 545/690] Simplify score_* for Interaction --- src/pykeen/nn/modules.py | 189 +++++++++------------------------------ 1 file changed, 40 insertions(+), 149 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 3058ec409b..595232d675 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -2,18 +2,18 @@ """Stateful interaction functions.""" -import itertools import logging import math from abc import ABC -from typing import Any, Callable, Generic, List, Mapping, MutableMapping, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Union import torch from torch import FloatTensor, nn from . import functional as pkf +from .representation import DIMS, convert_to_canonical_shape from ..typing import HeadRepresentation, RelationRepresentation, Representation, TailRepresentation -from ..utils import check_shapes, upgrade_to_sequence +from ..utils import upgrade_to_sequence __all__ = [ # Base Classes @@ -73,13 +73,6 @@ def _get_batches(z, slice_size): class Interaction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, TailRepresentation], ABC): """Base class for interaction functions.""" - # Dimensions - BATCH_DIM: int = 0 - NUM_DIM: int = 1 - HEAD_DIM: int = 1 - RELATION_DIM: int = 2 - TAIL_DIM: int = 3 - #: The symbolic shapes for entity representations entity_shape: Sequence[str] = ("d",) @@ -137,94 +130,6 @@ def forward( """ return self.__class__.func(**self._prepare_for_functional(h=h, r=r, t=t)) - @staticmethod - def _add_dim(*x: torch.FloatTensor, dim: int) -> Union[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]: - """ - Add a dimension to tensors. - - :param x: - The tensors. - :param dim: - The dimension to unsqueeze. - - :return: - The tensors with an additional 1-element dimension. - """ - out = [xx.unsqueeze(dim=dim) for xx in x] - if len(x) == 1: - return out[0] - return tuple(out) - - @staticmethod - def _remove_dim(x: torch.FloatTensor, *dims: int) -> torch.FloatTensor: - """ - Remove dimensions from a tensor. - - :param x: - The tensor. - :param dims: - The dimensions to remove. - - :return: - The squeezed tensor. - - :raises ValueError: - If there are duplicates in dims (after normalizing the dimensions, i.e. resolving negative dimension - indices). - """ - # normalize dimensions - dims = tuple(d if d >= 0 else len(x.shape) + d for d in dims) - if len(set(dims)) != len(dims): - raise ValueError(f"Duplicate dimensions: {dims}") - assert all(0 <= d < len(x.shape) for d in dims) - for dim in sorted(dims, reverse=True): - x = x.squeeze(dim=dim) - return x - - def _check_shapes( - self, - h: HeadRepresentation, - r: RelationRepresentation, - t: TailRepresentation, - h_prefix: str = "b", - r_prefix: str = "b", - t_prefix: str = "b", - raise_on_errors: bool = True, - ) -> bool: - entity_shape = self.entity_shape - if isinstance(entity_shape, str): - entity_shape = (entity_shape,) - relation_shape = self.relation_shape - if isinstance(relation_shape, str): - relation_shape = (relation_shape,) - tail_entity_shape = self.tail_entity_shape - if tail_entity_shape is None: - tail_entity_shape = entity_shape - if isinstance(tail_entity_shape, str): - tail_entity_shape = (tail_entity_shape,) - if len(h) != len(entity_shape): - if raise_on_errors: - raise ValueError - return False - if len(r) != len(relation_shape): - if raise_on_errors: - raise ValueError - return False - if len(t) != len(tail_entity_shape): - if raise_on_errors: - raise ValueError - return False - - # TODO make helper function + unit test - a = ((hh, h_prefix + hs) for hh, hs in zip(h, entity_shape)) # type: ignore - b = ((rr, r_prefix + rs) for rr, rs in zip(r, relation_shape)) # type: ignore - c = ((tt, t_prefix + ts) for tt, ts in zip(t, tail_entity_shape)) # type: ignore - - return check_shapes( - *itertools.chain(a, b, c), - raise_on_errors=raise_on_errors, - ) - def score( self, h: HeadRepresentation, @@ -239,11 +144,11 @@ def score( .. note :: At most one of the slice sizes may be not None. - :param h: shape: (batch_size, num_heads, ``*``) + :param h: shape: (batch_size, num_heads, `1, 1, `*``) The head representations. - :param r: shape: (batch_size, num_relations, ``*``) + :param r: shape: (batch_size, 1, num_relations, 1, ``*``) The relation representations. - :param t: shape: (batch_size, num_tails, ``*``) + :param t: shape: (batch_size, 1, 1, num_tails, ``*``) The tail representations. :param slice_size: The slice size. @@ -253,7 +158,6 @@ def score( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # TODO: fix shapes return self._forward_slicing_wrapper(h=h, r=r, t=t, slice_size=slice_size, slice_dim=slice_dim) def _score( @@ -261,44 +165,32 @@ def _score( h: HeadRepresentation, r: RelationRepresentation, t: TailRepresentation, - h_prefix: str = "b", - r_prefix: str = "b", - t_prefix: str = "b", slice_size: Optional[int] = None, + slice_dim: str = None, ) -> torch.FloatTensor: - assert {h_prefix, r_prefix, t_prefix}.issubset(list("bn")) - # at most one of h_prefix, r_prefix, t_prefix equals n - slice_dims: List[str] = [ - dim - for dim, prefix in zip("hrt", (h_prefix, r_prefix, t_prefix)) - if prefix == "n" - ] - slice_dim: Optional[str] = slice_dims[0] if len(slice_dims) == 1 else None - - # FIXME typing does not work well for this - h = upgrade_to_sequence(h) - r = upgrade_to_sequence(r) - t = upgrade_to_sequence(t) - assert self._check_shapes(h=h, r=r, t=t, h_prefix=h_prefix, r_prefix=r_prefix, t_prefix=t_prefix) - - # prepare input to generic score function: bh*, br*, bt* - h = self._add_dim(*h, dim=self.BATCH_DIM if h_prefix == "n" else self.NUM_DIM) - r = self._add_dim(*r, dim=self.BATCH_DIM if r_prefix == "n" else self.NUM_DIM) - t = self._add_dim(*t, dim=self.BATCH_DIM if t_prefix == "n" else self.NUM_DIM) - - # TODO: fix shapes - scores = self._forward_slicing_wrapper(h=h, r=r, t=t, slice_dim=slice_dim, slice_size=slice_size) - - remove_dims = [ - dim - for dim, prefix in zip( - (self.HEAD_DIM, self.RELATION_DIM, self.TAIL_DIM), - (h_prefix, r_prefix, t_prefix), - ) - if prefix == "b" - ] - # prepare output shape - return self._remove_dim(scores, *remove_dims) + """ + + :param h: shape: (b, h, *) + :param r: shape: (b, r, *) + :param t: shape: (b, t, *) + :param slice_size: ... + :param slice_dim: ... + :return: shape: (b, h, r, t) + """ + h, r, t = _unpack_singletons(*( + [ + convert_to_canonical_shape( + x=x.unsqueeze(dim=0 if dim != slice_dim else 1), + dim=dim, + num=x.shape[1], + batch_size=x.shape[0], + suffix_shape=x.shape[2:], + ) + for x in upgrade_to_sequence(xx) + ] + for xx, dim in zip((h, r, t), "hrt") + )) + return self._forward_slicing_wrapper(h=h, r=r, t=t, slice_dim=slice_dim, slice_size=slice_size) def _forward_slicing_wrapper( self, @@ -315,11 +207,11 @@ def _forward_slicing_wrapper( Depending on the interaction function, there may be more than one representation for h/r/t. In that case, a tuple of at least two tensors is passed. - :param h: shape: (batch_size, num_heads, ``*``) + :param h: shape: (batch_size, num_heads, 1, 1, ``*``) The head representations. - :param r: shape: (batch_size, num_relations, ``*``) + :param r: shape: (batch_size, 1, num_relations, 1, ``*``) The relation representations. - :param t: shape: (batch_size, num_tails, ``*``) + :param t: shape: (batch_size, 1, 1, num_tails, ``*``) The tail representations. :param slice_size: The slice size. @@ -332,24 +224,23 @@ def _forward_slicing_wrapper( :raises ValueError: If slice_dim is invalid. """ - # TODO: fix shapes if slice_size is None: scores = self(h=h, r=r, t=t) elif slice_dim == "h": scores = torch.cat([ self(h=h_batch, r=r, t=t) for h_batch in _get_batches(h, slice_size) - ], dim=self.HEAD_DIM) + ], dim=DIMS[slice_dim]) elif slice_dim == "r": scores = torch.cat([ self(h=h, r=r_batch, t=t) for r_batch in _get_batches(r, slice_size) - ], dim=self.RELATION_DIM) + ], dim=DIMS[slice_dim]) elif slice_dim == "t": scores = torch.cat([ self(h=h, r=r, t=t_batch) for t_batch in _get_batches(t, slice_size) - ], dim=self.TAIL_DIM) + ], dim=DIMS[slice_dim]) else: raise ValueError(f'Invalid slice_dim: {slice_dim}') return scores @@ -373,7 +264,7 @@ def score_hrt( :return: shape: (batch_size, 1) The scores. """ - return self._score(h=h, r=r, t=t).unsqueeze(dim=-1) + return self._score(h=h, r=r, t=t)[:, 0, 0, 0, None] def score_h( self, @@ -397,7 +288,7 @@ def score_h( :return: shape: (batch_size, num_entities) The scores. """ - return self._score(h=all_entities, r=r, t=t, h_prefix="n", slice_size=slice_size) + return self._score(h=all_entities, r=r, t=t, slice_dim="h", slice_size=slice_size)[:, :, 0, 0] def score_r( self, @@ -421,7 +312,7 @@ def score_r( :return: shape: (batch_size, num_entities) The scores. """ - return self._score(h=h, r=all_relations, t=t, r_prefix="n", slice_size=slice_size) + return self._score(h=h, r=all_relations, t=t, slice_dim="r", slice_size=slice_size)[:, 0, :, 0] def score_t( self, @@ -445,7 +336,7 @@ def score_t( :return: shape: (batch_size, num_entities) The scores. """ - return self._score(h=h, r=r, t=all_entities, t_prefix="n", slice_size=slice_size) + return self._score(h=h, r=r, t=all_entities, slice_dim="t", slice_size=slice_size)[:, 0, 0, :] def reset_parameters(self): """Reset parameters the interaction function may have.""" From 356df81a2a85ae759d20b98af6d7e73149c3b39d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 23:11:08 +0100 Subject: [PATCH 546/690] Skip shape tests with non-consistent head entity/tail entity numbers --- src/pykeen/models/base.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 975131a331..15f470daa4 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -1099,6 +1099,7 @@ def _prepare_representation_module_list( num_embeddings: int, shapes: Sequence[str], label: str, + skip_checks: bool = False, ) -> Sequence[RepresentationModule]: """Normalize list of representations and wrap into nn.ModuleList.""" # Important: use ModuleList to ensure that Pytorch correctly handles their devices and parameters @@ -1106,7 +1107,7 @@ def _prepare_representation_module_list( representations = [] if not isinstance(representations, Sequence): representations = [representations] - if len(representations) != len(shapes): + if not skip_checks and len(representations) != len(shapes): raise ValueError( f"Interaction function requires {len(shapes)} {label} representations, but " f"{len(representations)} were given.", @@ -1127,10 +1128,11 @@ def _prepare_representation_module_list( f"representations was chosen wrong.", ) modules.append(r) - check_shapes(*zip( - (r.shape for r in modules), - shapes, - ), raise_on_errors=True) + if not skip_checks: + check_shapes(*zip( + (r.shape for r in modules), + shapes, + ), raise_on_errors=True) return nn.ModuleList(modules) @@ -1200,6 +1202,7 @@ def __init__( num_embeddings=triples_factory.num_entities, shapes=interaction.entity_shape, label="entity", + skip_checks=interaction.tail_entity_shape is not None, ) self.relation_representations = _prepare_representation_module_list( representations=relation_representations, From c51662bb1f6d377f54ece756f82d30e1af6efa67 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 23:34:43 +0100 Subject: [PATCH 547/690] fix TransR functional --- src/pykeen/nn/functional.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index a390e785bb..603c02dcd4 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -789,7 +789,7 @@ def transr_interaction( Head embeddings. :param r: shape: (batch_size, 1, num_relations, 1, d_r) Relation embeddings. - :param m_r: shape: (batch_size, 1, num_relations, 1, d_r, d_e) + :param m_r: shape: (batch_size, 1, num_relations, 1, d_e, d_r) The relation specific linear transformations. :param t: shape: (batch_size, 1, 1, num_tails, d_e) Tail embeddings. @@ -802,8 +802,8 @@ def transr_interaction( The scores. """ # project to relation specific subspace and ensure constraints - h_bot = clamp_norm(m_r @ h.unsqueeze(dim=-1), p=2, dim=-1, maxnorm=1.).squeeze(dim=-1) - t_bot = clamp_norm(m_r @ t.unsqueeze(dim=-1), p=2, dim=-1, maxnorm=1.).squeeze(dim=-1) + h_bot = clamp_norm((h.unsqueeze(dim=-2) @ m_r), p=2, dim=-1, maxnorm=1.).squeeze(dim=-2) + t_bot = clamp_norm((t.unsqueeze(dim=-2) @ m_r), p=2, dim=-1, maxnorm=1.).squeeze(dim=-2) return negative_norm_of_sum(h_bot, r, -t_bot, p=p, power_norm=power_norm) From a262c30a521405f76cc3edf73bd032f8e7c20bde Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 23:34:59 +0100 Subject: [PATCH 548/690] Adjust InteractionModule and tests --- src/pykeen/nn/modules.py | 32 ++++++++++++++---------- tests/test_interactions.py | 51 ++++++++++++++++++++------------------ 2 files changed, 46 insertions(+), 37 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 595232d675..60af58020d 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -177,20 +177,26 @@ def _score( :param slice_dim: ... :return: shape: (b, h, r, t) """ - h, r, t = _unpack_singletons(*( - [ - convert_to_canonical_shape( - x=x.unsqueeze(dim=0 if dim != slice_dim else 1), - dim=dim, - num=x.shape[1], - batch_size=x.shape[0], - suffix_shape=x.shape[2:], + kwargs = dict() + for key, x in zip("hrt", (h, r, t)): + value = [] + for xx in upgrade_to_sequence(x): + # bring to (b, n, *) + xx = xx.unsqueeze(dim=1 if key != slice_dim else 0) + # bring to (b, h, r, t, *) + xx = convert_to_canonical_shape( + x=xx, + dim=key, + num=xx.shape[1], + batch_size=xx.shape[0], + suffix_shape=xx.shape[2:], ) - for x in upgrade_to_sequence(xx) - ] - for xx, dim in zip((h, r, t), "hrt") - )) - return self._forward_slicing_wrapper(h=h, r=r, t=t, slice_dim=slice_dim, slice_size=slice_size) + value.append(xx) + # unpack singleton + if len(value) == 1: + value = value[0] + kwargs[key] = value + return self._forward_slicing_wrapper(**kwargs, slice_dim=slice_dim, slice_size=slice_size) def _forward_slicing_wrapper( self, diff --git a/tests/test_interactions.py b/tests/test_interactions.py index ed89727a60..6f6ef17d90 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -192,44 +192,47 @@ def _check_close_scores(self, scores, scores_no_slice): self.assertTrue(torch.isfinite(scores_no_slice).all(), msg=f'Slice scores had nan\n\t{scores}') self.assertTrue(torch.allclose(scores, scores_no_slice), msg=f'Differences: {scores - scores_no_slice}') - def _get_test_shapes(self) -> Collection[Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]]]: + def _get_test_shapes(self) -> Collection[Tuple[ + Tuple[int, int, int, int], + Tuple[int, int, int, int], + Tuple[int, int, int, int], + ]]: """Return a set of test shapes for (h, r, t).""" return ( ( # single score - (1, 1), - (1, 1), - (1, 1), + (1, 1, 1, 1), + (1, 1, 1, 1), + (1, 1, 1, 1), ), ( # score_r with multi-t - (self.batch_size, 1), - (1, self.num_relations), - (self.batch_size, self.num_entities // 2 + 1), + (self.batch_size, 1, 1, 1), + (1, 1, self.num_relations, 1), + (self.batch_size, 1, 1, self.num_entities // 2 + 1), ), ( # score_r with multi-t and broadcasted head - (1, 1), - (1, self.num_relations), - (self.batch_size, self.num_entities), + (1, 1, 1, 1), + (1, 1, self.num_relations, 1), + (self.batch_size, 1, 1, self.num_entities), ), ( # full cwa - (1, self.num_entities), - (1, self.num_relations), - (1, self.num_entities), + (1, self.num_entities, 1, 1), + (1, 1, self.num_relations, 1), + (1, 1, 1, self.num_entities), ), ) def _get_output_shape( self, - hs: Tuple[int, int], - rs: Tuple[int, int], - ts: Tuple[int, int], + hs: Tuple[int, int, int, int], + rs: Tuple[int, int, int, int], + ts: Tuple[int, int, int, int], ) -> Tuple[int, int, int, int]: - batch_size = max(hs[0], rs[0], ts[0]) - nh, nr, nt = hs[1], rs[1], ts[1] - if len(self.cls.relation_shape) == 0: - nr = 1 - if len(self.cls.entity_shape) == 0: - nh = nt = 1 - return batch_size, nh, nr, nt + result = [max(ds) for ds in zip(hs, rs, ts)] + if len(self.instance.entity_shape) == 0: + result[1] = result[3] = 1 + if len(self.instance.relation_shape) == 0: + result[2] = 1 + return tuple(result) def test_forward(self): """Test forward.""" @@ -271,7 +274,7 @@ def test_scores(self): for _ in range(10): # test multiple different initializations self.instance.reset_parameters() - h, r, t = self._get_hrt((1, 1), (1, 1), (1, 1)) + h, r, t = self._get_hrt((1, 1, 1, 1), (1, 1, 1, 1), (1, 1, 1, 1)) kwargs = self.instance._prepare_for_functional(h=h, r=r, t=t) # calculate by functional From 36e42043577596fdb14ffb206b72eda0f1e0bf37 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 23:39:58 +0100 Subject: [PATCH 549/690] Update strip dim utility --- src/pykeen/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 506b0584fc..96e08a7b8c 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -630,9 +630,9 @@ def pop_only(elements: Iterable[X]) -> X: return elements[0] -def strip_dim(*x): - """Strip the last dimension.""" - return [xx.view(xx.shape[2:]) for xx in x] +def strip_dim(*x, num: int = 4): + """Strip the first dimensions.""" + return [xx.view(xx.shape[num:]) for xx in x] def upgrade_to_sequence(x: Union[X, Sequence[X]]) -> Sequence[X]: From 95b41a374c77d9bff1ffd1a320debb09e4cd5b19 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 23:43:31 +0100 Subject: [PATCH 550/690] Fix shapes for NTN --- src/pykeen/nn/functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 603c02dcd4..5fbed3009d 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -472,7 +472,7 @@ def ntn_interaction( b, )) u = u.transpose(-2, -1) - return (x @ u).squeeze(dim=-2) + return (x @ u).squeeze(dim=-1) def proje_interaction( From ed03d53b07280a162db185fba8a1dc39c4e263e4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 23:49:42 +0100 Subject: [PATCH 551/690] Fix ProjE --- src/pykeen/nn/functional.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 5fbed3009d..8a75c3e48d 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -516,12 +516,12 @@ def proje_interaction( # global projections h = h * d_e.view(1, 1, 1, 1, dim) r = r * d_r.view(1, 1, 1, 1, dim) - # combination, shape: (b, h, r, d) + # combination, shape: (b, h, r, 1, d) x = tensor_sum(h, r, b_c) - x = activation(x) + x = activation(x) # shape: (b, h, r, 1, d) # dot product with t, shape: (b, h, r, t) - t = t.transpose(-2, -1) - return (x @ t) + b_p + t = t.transpose(-2, -1) # shape: (b, 1, 1, d, t) + return (x @ t).squeeze(dim=-2) + b_p def rescal_interaction( From a1bd33ee2397545a77507b62a1d0a55f3252caef Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 23:51:14 +0100 Subject: [PATCH 552/690] Fix HolE --- src/pykeen/nn/functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 8a75c3e48d..cd0834b549 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -383,7 +383,7 @@ def hole_interaction( r = r.permute(0, 1, 3, 4, 2) # inner product with relation embedding - return (composite @ r).squeeze(dim=-2) + return (composite @ r).squeeze(dim=-1) def kg2e_interaction( From f287ffd5aac605b6fa51662e0e4e1dc5451fce52 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 23:55:02 +0100 Subject: [PATCH 553/690] Fix HolE --- src/pykeen/nn/functional.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index cd0834b549..ec8df83fc3 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -379,11 +379,11 @@ def hole_interaction( # inverse real FFT, shape: (b, h, 1, t, d) composite = torch.fft.irfft(p_fft, n=h.shape[-1], dim=-1) - # transpose r, (b, 1, r, 1, d) -> (b, 1, 1, d, r) - r = r.permute(0, 1, 3, 4, 2) + # transpose composite: (b, h, 1, d, t) + composite = composite.transpose(-2, -1) # inner product with relation embedding - return (composite @ r).squeeze(dim=-1) + return (r @ composite).squeeze(dim=-2) def kg2e_interaction( From d47c8f8b349df49c6fba149748f00e177e0b94a0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 23:56:50 +0100 Subject: [PATCH 554/690] Move generic test classes --- src/pykeen/testing/base.py | 45 ++++++++++++++++++++++++++++++++++++++ tests/test_interactions.py | 45 +++----------------------------------- 2 files changed, 48 insertions(+), 42 deletions(-) create mode 100644 src/pykeen/testing/base.py diff --git a/src/pykeen/testing/base.py b/src/pykeen/testing/base.py new file mode 100644 index 0000000000..e34c94493c --- /dev/null +++ b/src/pykeen/testing/base.py @@ -0,0 +1,45 @@ +"""Base classes for simplified testing.""" +from typing import Any, Collection, Generic, Mapping, MutableMapping, Optional, Type, TypeVar + +from pykeen.utils import get_subclasses, set_random_seed + +T = TypeVar("T") + + +class GenericTests(Generic[T]): + """Generic tests.""" + + cls: Type[T] + kwargs: Optional[Mapping[str, Any]] = None + instance: T + + def setUp(self) -> None: + """Set up the generic testing method.""" + # fix seeds for reproducibility + set_random_seed(seed=42) + kwargs = self.kwargs or {} + kwargs = self._pre_instantiation_hook(kwargs=dict(kwargs)) + self.instance = self.cls(**kwargs) + self.post_instantiation_hook() + + def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: + """Perform actions before instantiation, potentially modyfing kwargs.""" + return kwargs + + def post_instantiation_hook(self) -> None: + """Perform actions after instantiation.""" + + +class TestsTest(Generic[T]): + """A generic test for tests.""" + + base_cls: Type[T] + base_test: Type[GenericTests[T]] + skip_cls: Collection[T] = tuple() + + def test_testing(self): + """Check that there is a test for all subclasses.""" + to_test = set(get_subclasses(self.base_cls)).difference(self.skip_cls) + tested = (test_cls.cls for test_cls in get_subclasses(self.base_test) if hasattr(test_cls, "cls")) + not_tested = to_test.difference(tested) + assert not not_tested, not_tested diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 6f6ef17d90..9b37037474 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -5,7 +5,7 @@ import logging import unittest from abc import abstractmethod -from typing import Any, Collection, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, TypeVar, Union +from typing import Collection, Sequence, Tuple, TypeVar, Union from unittest.case import SkipTest import numpy @@ -13,52 +13,13 @@ import pykeen.nn.modules from pykeen.nn.modules import Interaction, TranslationalInteraction +from pykeen.testing.base import GenericTests, TestsTest from pykeen.typing import Representation -from pykeen.utils import clamp_norm, get_subclasses, project_entity, set_random_seed, strip_dim, view_complex +from pykeen.utils import clamp_norm, project_entity, strip_dim, view_complex -T = TypeVar("T") logger = logging.getLogger(__name__) -class GenericTests(Generic[T]): - """Generic tests.""" - - cls: Type[T] - kwargs: Optional[Mapping[str, Any]] = None - instance: T - - def setUp(self) -> None: - """Set up the generic testing method.""" - # fix seeds for reproducibility - set_random_seed(seed=42) - kwargs = self.kwargs or {} - kwargs = self._pre_instantiation_hook(kwargs=dict(kwargs)) - self.instance = self.cls(**kwargs) - self.post_instantiation_hook() - - def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: - """Perform actions before instantiation, potentially modyfing kwargs.""" - return kwargs - - def post_instantiation_hook(self) -> None: - """Perform actions after instantiation.""" - - -class TestsTest(Generic[T]): - """A generic test for tests.""" - - base_cls: Type[T] - base_test: Type[GenericTests[T]] - skip_cls: Collection[T] = tuple() - - def test_testing(self): - """Check that there is a test for all subclasses.""" - to_test = set(get_subclasses(self.base_cls)).difference(self.skip_cls) - tested = (test_cls.cls for test_cls in get_subclasses(self.base_test) if hasattr(test_cls, "cls")) - not_tested = to_test.difference(tested) - assert not not_tested, not_tested - - class InteractionTests(GenericTests[pykeen.nn.modules.Interaction]): """Generic test for interaction functions.""" From 8d7c2bb98044dc0521a42ed23afaa92a2d9f8298 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 23 Nov 2020 23:58:57 +0100 Subject: [PATCH 555/690] Add shape information to SimplE module --- src/pykeen/nn/modules.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 60af58020d..8ecec5cd27 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -1006,6 +1006,8 @@ class SimplEInteraction( """Interaction function of SimplE.""" func = pkf.simple_interaction + entity_shape = ("d", "d") + relation_shape = ("d", "d") def __init__(self, clamp_score: Union[None, float, Tuple[float, float]] = None): super().__init__() From 97866f01549d665b7d718477a9c0fd4b586dae60 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 00:00:50 +0100 Subject: [PATCH 556/690] Add SimplE test --- tests/test_interactions.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 9b37037474..6d02043a35 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -5,13 +5,14 @@ import logging import unittest from abc import abstractmethod -from typing import Collection, Sequence, Tuple, TypeVar, Union +from typing import Collection, Sequence, Tuple, Union from unittest.case import SkipTest import numpy import torch import pykeen.nn.modules +from pykeen.nn.functional import distmult_interaction from pykeen.nn.modules import Interaction, TranslationalInteraction from pykeen.testing.base import GenericTests, TestsTest from pykeen.typing import Representation @@ -584,6 +585,17 @@ def _exp_score(self, h, t, p, power_norm) -> torch.FloatTensor: return -(h - t).pow(p).sum() +class SimplEInteractionTests(InteractionTests, unittest.TestCase): + """Tests for SimplE interaction function.""" + + cls = pykeen.nn.modules.SimplEInteraction + + def _exp_score(self, h, r, t, h_inv, r_inv, t_inv, clamp) -> torch.FloatTensor: + h, r, t, h_inv, r_inv, t_inv = strip_dim(h, r, t, h_inv, r_inv, t_inv) + assert clamp is None + return 0.5 * distmult_interaction(h, r, t) + 0.5 * distmult_interaction(h_inv, r_inv, t_inv) + + class InteractionTestsTest(TestsTest[Interaction], unittest.TestCase): """Test for tests for all interaction functions.""" From e828b17b25fe3adde83999ad8ca2c5e80c903c30 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 00:08:41 +0100 Subject: [PATCH 557/690] Fix TransR relation projection matrices --- src/pykeen/models/unimodal/trans_r.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 704fbe8644..76777336bc 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -96,7 +96,7 @@ def __init__( ), # Relation projections EmbeddingSpecification( - shape=(relation_dim, embedding_dim), + shape=(embedding_dim, relation_dim), initializer=xavier_uniform_, ), ], From 7ad01b3cbe8d616f1e9318fe454a7ffcb50fe560 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 00:20:49 +0100 Subject: [PATCH 558/690] Fix SE --- src/pykeen/nn/functional.py | 20 ++++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index ec8df83fc3..9df4d30817 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -111,11 +111,11 @@ def complex_interaction( return tensor_sum(*( factor * tensor_product(hh, rr, tt).sum(dim=-1) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) @@ -518,7 +518,7 @@ def proje_interaction( r = r * d_r.view(1, 1, 1, 1, dim) # combination, shape: (b, h, r, 1, d) x = tensor_sum(h, r, b_c) - x = activation(x) # shape: (b, h, r, 1, d) + x = activation(x) # shape: (b, h, r, 1, d) # dot product with t, shape: (b, h, r, t) t = t.transpose(-2, -1) # shape: (b, 1, 1, d, t) return (x @ t).squeeze(dim=-2) + b_p @@ -638,9 +638,9 @@ def structured_embedding_interaction( :param h: shape: (batch_size, num_heads, 1, 1, dim) The head representations. - :param r_h: shape: (batch_size, 1, num_relations, 1, dim, rel_dim) + :param r_h: shape: (batch_size, 1, num_relations, 1, rel_dim, dim) The relation-specific head projection. - :param r_t: shape: (batch_size, 1, num_relations, 1, dim, rel_dim) + :param r_t: shape: (batch_size, 1, num_relations, 1, rel_dim, dim) The relation-specific tail projection. :param t: shape: (batch_size, 1, 1, num_tails, dim) The tail representations. @@ -653,8 +653,8 @@ def structured_embedding_interaction( The scores. """ return negative_norm_of_sum( - (h @ r_h.squeeze(dim=-3)), - -(t @ r_t.squeeze(dim=-3)), + (r_h @ h.unsqueeze(dim=-1)).squeeze(dim=-1), + -(r_t @ t.unsqueeze(dim=-1)).squeeze(dim=-1), p=p, power_norm=power_norm, ) From e07be941ce2782395e8c41edc220ea35ebad1db3 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 00:25:48 +0100 Subject: [PATCH 559/690] For flake --- src/pykeen/nn/functional.py | 10 +++++----- src/pykeen/nn/modules.py | 3 +++ src/pykeen/utils.py | 1 - 3 files changed, 8 insertions(+), 6 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 9df4d30817..1ce51f15cf 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -111,11 +111,11 @@ def complex_interaction( return tensor_sum(*( factor * tensor_product(hh, rr, tt).sum(dim=-1) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 8ecec5cd27..2b8e8139a1 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -169,6 +169,9 @@ def _score( slice_dim: str = None, ) -> torch.FloatTensor: """ + Compute scores for the score_* methods outside of models. + + TODO: merge this with the Model utilities? :param h: shape: (b, h, *) :param r: shape: (b, r, *) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 96e08a7b8c..4304709048 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -8,7 +8,6 @@ import logging import operator import random -import typing from abc import ABC, abstractmethod from io import BytesIO from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union From 583a620b51689c61fcf9a064e0f250cfbe0671c7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 00:29:14 +0100 Subject: [PATCH 560/690] Do not use dictionary --- src/pykeen/nn/modules.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 2b8e8139a1..1b108c786e 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -180,7 +180,7 @@ def _score( :param slice_dim: ... :return: shape: (b, h, r, t) """ - kwargs = dict() + args = [] for key, x in zip("hrt", (h, r, t)): value = [] for xx in upgrade_to_sequence(x): @@ -198,8 +198,9 @@ def _score( # unpack singleton if len(value) == 1: value = value[0] - kwargs[key] = value - return self._forward_slicing_wrapper(**kwargs, slice_dim=slice_dim, slice_size=slice_size) + args.append(value) + h, r, t = args + return self._forward_slicing_wrapper(h=h, r=r, t=t, slice_dim=slice_dim, slice_size=slice_size) def _forward_slicing_wrapper( self, From 838b76d956dd7d597877bfc50cd934478236e618 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 00:30:20 +0100 Subject: [PATCH 561/690] Add annotation for xx --- src/pykeen/nn/modules.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 1b108c786e..6939be1ef6 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -185,7 +185,7 @@ def _score( value = [] for xx in upgrade_to_sequence(x): # bring to (b, n, *) - xx = xx.unsqueeze(dim=1 if key != slice_dim else 0) + xx: torch.FloatTensor = xx.unsqueeze(dim=1 if key != slice_dim else 0) # bring to (b, h, r, t, *) xx = convert_to_canonical_shape( x=xx, From a2ccc6e1c33f751c2e9486c061fd7c403bb67e8b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 00:31:12 +0100 Subject: [PATCH 562/690] Give mypy a type for xx --- src/pykeen/nn/modules.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 6939be1ef6..0f856a44bc 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -183,9 +183,9 @@ def _score( args = [] for key, x in zip("hrt", (h, r, t)): value = [] - for xx in upgrade_to_sequence(x): + for xx in upgrade_to_sequence(x): # type: torch.FloatTensor # bring to (b, n, *) - xx: torch.FloatTensor = xx.unsqueeze(dim=1 if key != slice_dim else 0) + xx = xx.unsqueeze(dim=1 if key != slice_dim else 0) # bring to (b, h, r, t, *) xx = convert_to_canonical_shape( x=xx, From 69fe97f7f84417df6ed82e8bde181d268f46b82b Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 24 Nov 2020 01:55:25 +0100 Subject: [PATCH 563/690] Update minimum python version --- .github/workflows/tests.yml | 6 +++--- setup.cfg | 5 +---- 2 files changed, 4 insertions(+), 7 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index e0f6e89aef..c9dc52edb6 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -9,7 +9,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [ 3.6, 3.7, 3.8 ] + python-version: [ 3.7, 3.8 ] steps: - uses: actions/checkout@v2 - name: Set up Python ${{ matrix.python-version }} @@ -40,7 +40,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [ 3.6, 3.7, 3.8 ] + python-version: [ 3.7, 3.8 ] steps: - uses: actions/checkout@v2 - name: Set up Python ${{ matrix.python-version }} @@ -66,7 +66,7 @@ jobs: runs-on: ubuntu-latest strategy: matrix: - python-version: [ 3.6, 3.7, 3.8 ] + python-version: [ 3.7, 3.8 ] steps: - uses: actions/checkout@v2 - name: Set up Python ${{ matrix.python-version }} diff --git a/setup.cfg b/setup.cfg index 8e331c6b78..495df60c3e 100644 --- a/setup.cfg +++ b/setup.cfg @@ -33,10 +33,8 @@ classifiers = License :: OSI Approved :: MIT License Operating System :: OS Independent Programming Language :: Python - Programming Language :: Python :: 3.6 Programming Language :: Python :: 3.7 Programming Language :: Python :: 3.8 - Programming Language :: Python :: 3.9 Programming Language :: Python :: 3 :: Only Topic :: Scientific/Engineering :: Artificial Intelligence Topic :: Scientific/Engineering :: Chemistry @@ -51,7 +49,6 @@ keywords = [options] install_requires = - dataclasses; python_version < "3.7" dataclasses-json numpy click @@ -66,7 +63,7 @@ install_requires = zip_safe = false include_package_data = True -python_requires = >=3.6 +python_requires = >=3.7 # Where is my code packages = find: From 14b45ea5ee25bd528c3fe4796e894de3bf59b6af Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 24 Nov 2020 02:05:29 +0100 Subject: [PATCH 564/690] Update representation.py --- src/pykeen/nn/representation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index 8080fca0e2..199d06a63a 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -124,7 +124,7 @@ def get_in_canonical_shape( :param indices: The indices. Either None, in which care all embeddings are returned, or a 1 or 2 dimensional index tensor. - :return: shape: (batch_size, d1, d2, d3, *self.shape) + :return: shape: (batch_size, d1, d2, d3, ``*self.shape``) """ r_shape: Tuple[int, ...] if indices is None: From b0fe4f4989f85fbe0c7d8870d85ede0911f15134 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 10:24:12 +0100 Subject: [PATCH 565/690] Extend Embedding tests --- tests/test_nn.py | 142 +++++++++++++++++++++++++++++------------------ 1 file changed, 89 insertions(+), 53 deletions(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index 3a7b8ba8fd..9685ad0303 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -1,72 +1,108 @@ # -*- coding: utf-8 -*- """Unittest for the :mod:`pykeen.nn` module.""" - +import itertools import unittest +from typing import Iterable, Optional, Sequence import torch -from pykeen.nn import Embedding +from pykeen.nn import Embedding, RepresentationModule from pykeen.nn.sim import kullback_leibler_similarity +from pykeen.testing.base import GenericTests from pykeen.typing import GaussianDistribution -class EmbeddingsInCanonicalShapeTests(unittest.TestCase): - """Test get_embedding_in_canonical_shape().""" +class RepresentationModuleTests(GenericTests[RepresentationModule]): + """Tests for RepresentationModule.""" - #: The number of embeddings - num_embeddings: int = 3 + batch_size: int = 3 + num: int = 5 + exp_shape: Sequence[int] = (5,) - #: The embedding dimension - embedding_dim: int = 2 + def test_max_id(self): + assert self.instance.max_id == self.num - def setUp(self) -> None: - """Initialize embedding.""" - self.embedding = Embedding(num_embeddings=self.num_embeddings, embedding_dim=self.embedding_dim) - self.generator = torch.manual_seed(42) - self.embedding._embeddings.weight.data = torch.rand( - self.num_embeddings, - self.embedding_dim, - generator=self.generator, - ) + def test_shape(self): + assert self.instance.shape == self.exp_shape - def test_no_indices(self): - """Test getting all embeddings.""" - emb = self.embedding.get_in_canonical_shape(indices=None) - - # check shape - assert emb.shape == (1, self.num_embeddings, self.embedding_dim) - - # check values - exp = self.embedding(indices=None).view(1, self.num_embeddings, self.embedding_dim) - assert torch.allclose(emb, exp) - - def _test_with_indices(self, indices: torch.Tensor) -> None: - """Help tests with index.""" - emb = self.embedding.get_in_canonical_shape(indices=indices) - - # check shape - num_ind = indices.shape[0] - assert emb.shape == (num_ind, 1, self.embedding_dim) - - # check values - exp = torch.stack([self.embedding(i) for i in indices], dim=0).view(num_ind, 1, self.embedding_dim) - assert torch.allclose(emb, exp) - - def test_with_consecutive_indices(self): - """Test to retrieve all embeddings with consecutive indices.""" - indices = torch.arange(self.num_embeddings, dtype=torch.long) - self._test_with_indices(indices=indices) - - def test_with_indices_with_duplicates(self): - """Test to retrieve embeddings at random positions with duplicate indices.""" - indices = torch.randint( - self.num_embeddings, - size=(2 * self.num_embeddings,), - dtype=torch.long, - generator=self.generator, + def _test_forward(self, indices: Optional[torch.LongTensor]): + """Test the forward method.""" + assert indices is None or ( + torch.is_tensor(indices) + and indices.dtype == torch.long + and indices.ndimension() == 1 ) - self._test_with_indices(indices=indices) + x = self.instance(indices=indices) + assert torch.is_tensor(x) + assert x.dtype == torch.float32 + n = self.num if indices is None else indices.shape[0] + assert x.shape == tuple([n, *self.instance.shape]) + + def _test_indices(self) -> Iterable[torch.LongTensor]: + return [ + torch.randint(self.num, size=(self.batch_size,)), + torch.randperm(self.num), + torch.randperm(self.num).repeat(2), + ] + + def test_forward_without_indices(self): + self._test_forward(indices=None) + + def test_forward_with_indices(self): + for indices in self._test_indices(): + self._test_forward(indices=indices) + + def _test_in_canonical_shape(self, indices): + name_to_shape = dict(h=1, r=2, t=3) + for dim in itertools.chain(name_to_shape.keys(), name_to_shape.values()): + # batch_size, d1, d2, d3, * + x = self.instance.get_in_canonical_shape(dim=dim, indices=indices) + assert torch.is_tensor(x) + assert x.dtype == torch.float32 + assert x.ndimension() == 4 + len(self.exp_shape) + exp_shape = [1, 1, 1, 1] + list(self.exp_shape) + if isinstance(dim, str): + dim = name_to_shape[dim] + if indices is None: # 1-n scoring + exp_shape[dim] = self.num + if indices is not None: # batch dimension + exp_shape[0] = indices.shape[0] + if indices.ndimension() > 1: # multi-target batching + exp_shape[dim] = indices.shape[1] + assert x.shape == tuple(exp_shape) + + def test_get_in_canonical_shape_without_indices(self): + self._test_in_canonical_shape(indices=None) + + def test_get_in_canonical_shape_with_indices(self): + for indices in self._test_indices(): + self._test_in_canonical_shape(indices=indices) + + def test_get_in_canonical_shape_with_2d_indices(self): + indices = torch.randint(self.num, size=(self.batch_size, 2)) + self._test_in_canonical_shape(indices=indices) + + +class EmbeddingTests(RepresentationModuleTests, unittest.TestCase): + """Tests for Embedding.""" + + cls = Embedding + kwargs = dict( + num_embeddings=RepresentationModuleTests.num, + shape=RepresentationModuleTests.exp_shape, + ) + + +class TensorEmbeddingTests(RepresentationModuleTests, unittest.TestCase): + """Tests for Embedding with 2-dimensional shape.""" + + cls = Embedding + exp_shape = (3, 7) + kwargs = dict( + num_embeddings=RepresentationModuleTests.num, + shape=(3, 7), + ) class KullbackLeiblerTests(unittest.TestCase): From 6219d7aef2b85b24c33d595533f2df5ba9209e55 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 10:25:36 +0100 Subject: [PATCH 566/690] Fix Embedding.forward for >1d embeddings --- src/pykeen/nn/representation.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index 199d06a63a..ea228f7488 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -341,6 +341,7 @@ def forward( x = self._embeddings.weight else: x = self._embeddings(indices) + x = x.view(x.shape[0], *self.shape) if self.normalizer is not None: x = self.normalizer(x) if self.regularizer is not None: From dbe5c50b2c185006de28f6420734124ba9691079 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 10:31:48 +0100 Subject: [PATCH 567/690] Add tests for LiteralRepresentations --- tests/test_nn.py | 30 ++++++++++++++++++++++++++++-- 1 file changed, 28 insertions(+), 2 deletions(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index 9685ad0303..61a310279f 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -3,11 +3,11 @@ """Unittest for the :mod:`pykeen.nn` module.""" import itertools import unittest -from typing import Iterable, Optional, Sequence +from typing import Any, Iterable, MutableMapping, Optional, Sequence import torch -from pykeen.nn import Embedding, RepresentationModule +from pykeen.nn import Embedding, LiteralRepresentations, RepresentationModule from pykeen.nn.sim import kullback_leibler_similarity from pykeen.testing.base import GenericTests from pykeen.typing import GaussianDistribution @@ -20,6 +20,9 @@ class RepresentationModuleTests(GenericTests[RepresentationModule]): num: int = 5 exp_shape: Sequence[int] = (5,) + def post_instantiation_hook(self) -> None: # noqa: D102 + self.instance.reset_parameters() + def test_max_id(self): assert self.instance.max_id == self.num @@ -38,6 +41,11 @@ def _test_forward(self, indices: Optional[torch.LongTensor]): assert x.dtype == torch.float32 n = self.num if indices is None else indices.shape[0] assert x.shape == tuple([n, *self.instance.shape]) + self._verify_content(x=x, indices=indices) + + def _verify_content(self, x, indices): + """Additional verification.""" + assert x.requires_grad def _test_indices(self) -> Iterable[torch.LongTensor]: return [ @@ -105,6 +113,24 @@ class TensorEmbeddingTests(RepresentationModuleTests, unittest.TestCase): ) +class LiteralRepresentationsTests(RepresentationModuleTests, unittest.TestCase): + """Tests for literal embeddings.""" + + cls = LiteralRepresentations + + def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102 + kwargs = super()._pre_instantiation_hook(kwargs=kwargs) + self.numeric_literals = torch.rand(self.num, *self.exp_shape) + kwargs["numeric_literals"] = self.numeric_literals + return kwargs + + def _verify_content(self, x, indices): # noqa: D102 + exp_x = self.numeric_literals + if indices is not None: + exp_x = exp_x[indices] + assert torch.allclose(x, exp_x) + + class KullbackLeiblerTests(unittest.TestCase): """Tests for the vectorized computation of KL divergences.""" From d7dcf90d87c874b20c9583746c6386da2711bac2 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 10:39:41 +0100 Subject: [PATCH 568/690] Add utility for expected_canonical_shape --- src/pykeen/nn/representation.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index ea228f7488..3b8d21ca29 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -39,6 +39,38 @@ def _normalize_dim(dim: Union[int, str]) -> int: return DIMS[dim.lower()[0]] +def get_expected_canonical_shape( + indices_shape: Union[None, int, Tuple[int, int]], + dim: Union[str, int], + suffix_shape: Union[int, Sequence[int]], + num: Optional[int] = None, +) -> Tuple[int, ...]: + """ + Calculate the expected canonical shape for the given parameters. + + :param indices_shape: + The shape of the indices, or None, if no indices are to be provided. + :param dim: + The dimension, either symbolic, or numeric. + :param suffix_shape: + The suffix-shape. + :param num: + The number of representations, if indices_shape is None, i.e. 1-n scoring. + + :return: (batch_size, num_heads, num_relations, num_tails, ``*``). + The expected shape, a tuple of at least 5 positive integers. + """ + exp_shape = [1, 1, 1, 1] + list(suffix_shape) + dim = _normalize_dim(dim=dim) + if indices_shape is None: # 1-n scoring + exp_shape[dim] = num + else: # batch dimension + exp_shape[0] = indices_shape[0] + if len(indices_shape) > 1: # multi-target batching + exp_shape[dim] = indices_shape[1] + return tuple(exp_shape) + + def convert_to_canonical_shape( x: torch.FloatTensor, dim: Union[int, str], From 2f861ab2f8a69706a268d9e3eecc95b2a805c47a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 10:47:39 +0100 Subject: [PATCH 569/690] Improve documentation of tests --- src/pykeen/nn/representation.py | 16 ++++---- tests/test_nn.py | 67 ++++++++++++++++++++++----------- 2 files changed, 54 insertions(+), 29 deletions(-) diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index 3b8d21ca29..d1286f11b5 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -40,7 +40,7 @@ def _normalize_dim(dim: Union[int, str]) -> int: def get_expected_canonical_shape( - indices_shape: Union[None, int, Tuple[int, int]], + indices: Union[None, int, Tuple[int, int], torch.LongTensor], dim: Union[str, int], suffix_shape: Union[int, Sequence[int]], num: Optional[int] = None, @@ -48,8 +48,8 @@ def get_expected_canonical_shape( """ Calculate the expected canonical shape for the given parameters. - :param indices_shape: - The shape of the indices, or None, if no indices are to be provided. + :param indices: + The indices, their shape, or None, if no indices are to be provided. :param dim: The dimension, either symbolic, or numeric. :param suffix_shape: @@ -60,14 +60,16 @@ def get_expected_canonical_shape( :return: (batch_size, num_heads, num_relations, num_tails, ``*``). The expected shape, a tuple of at least 5 positive integers. """ + if torch.is_tensor(indices): + indices = indices.shape exp_shape = [1, 1, 1, 1] + list(suffix_shape) dim = _normalize_dim(dim=dim) - if indices_shape is None: # 1-n scoring + if indices is None: # 1-n scoring exp_shape[dim] = num else: # batch dimension - exp_shape[0] = indices_shape[0] - if len(indices_shape) > 1: # multi-target batching - exp_shape[dim] = indices_shape[1] + exp_shape[0] = indices[0] + if len(indices) > 1: # multi-target batching + exp_shape[dim] = indices[1] return tuple(exp_shape) diff --git a/tests/test_nn.py b/tests/test_nn.py index 61a310279f..0284f27ad8 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -5,9 +5,11 @@ import unittest from typing import Any, Iterable, MutableMapping, Optional, Sequence +import pytest import torch from pykeen.nn import Embedding, LiteralRepresentations, RepresentationModule +from pykeen.nn.representation import DIMS, get_expected_canonical_shape from pykeen.nn.sim import kullback_leibler_similarity from pykeen.testing.base import GenericTests from pykeen.typing import GaussianDistribution @@ -16,26 +18,28 @@ class RepresentationModuleTests(GenericTests[RepresentationModule]): """Tests for RepresentationModule.""" + #: The batch size batch_size: int = 3 + + #: The number of representations num: int = 5 + + #: The expected shape of an individual representation exp_shape: Sequence[int] = (5,) def post_instantiation_hook(self) -> None: # noqa: D102 self.instance.reset_parameters() def test_max_id(self): + """Test the maximum ID.""" assert self.instance.max_id == self.num def test_shape(self): + """Test the shape.""" assert self.instance.shape == self.exp_shape def _test_forward(self, indices: Optional[torch.LongTensor]): """Test the forward method.""" - assert indices is None or ( - torch.is_tensor(indices) - and indices.dtype == torch.long - and indices.ndimension() == 1 - ) x = self.instance(indices=indices) assert torch.is_tensor(x) assert x.dtype == torch.float32 @@ -47,47 +51,66 @@ def _verify_content(self, x, indices): """Additional verification.""" assert x.requires_grad - def _test_indices(self) -> Iterable[torch.LongTensor]: + def _valid_indices(self) -> Iterable[torch.LongTensor]: return [ torch.randint(self.num, size=(self.batch_size,)), torch.randperm(self.num), torch.randperm(self.num).repeat(2), ] + def _invalid_indices(self) -> Iterable[torch.LongTensor]: + return [ + torch.as_tensor([self.num], dtype=torch.long), # too high index + torch.randint(self.num, size=(2, 3)), # too many indices + ] + def test_forward_without_indices(self): + """Test forward without providing indices.""" self._test_forward(indices=None) def test_forward_with_indices(self): - for indices in self._test_indices(): + """Test forward with providing indices.""" + for indices in self._valid_indices(): self._test_forward(indices=indices) - def _test_in_canonical_shape(self, indices): - name_to_shape = dict(h=1, r=2, t=3) - for dim in itertools.chain(name_to_shape.keys(), name_to_shape.values()): + def test_forward_with_invalid_indices(self): + """Test whether passing invalid indices crashes.""" + for indices in self._invalid_indices(): + with pytest.raises((IndexError, RuntimeError)): + self._test_forward(indices=indices) + + def _test_in_canonical_shape(self, indices: Optional[torch.LongTensor]): + """Test get_in_canonical_shape with the given indices.""" + # test both, using the actual dimension, and its name + for dim in itertools.chain(DIMS.keys(), DIMS.values()): # batch_size, d1, d2, d3, * x = self.instance.get_in_canonical_shape(dim=dim, indices=indices) + + # data type assert torch.is_tensor(x) - assert x.dtype == torch.float32 + assert x.dtype == torch.float32 # todo: adjust? assert x.ndimension() == 4 + len(self.exp_shape) - exp_shape = [1, 1, 1, 1] + list(self.exp_shape) - if isinstance(dim, str): - dim = name_to_shape[dim] - if indices is None: # 1-n scoring - exp_shape[dim] = self.num - if indices is not None: # batch dimension - exp_shape[0] = indices.shape[0] - if indices.ndimension() > 1: # multi-target batching - exp_shape[dim] = indices.shape[1] - assert x.shape == tuple(exp_shape) + + # get expected shape + exp_shape = get_expected_canonical_shape( + indices=indices, + dim=dim, + suffix_shape=self.exp_shape, + num=self.num, + ) + assert x.shape == exp_shape def test_get_in_canonical_shape_without_indices(self): + """Test get_in_canonical_shape without indices, i.e. with 1-n scoring.""" self._test_in_canonical_shape(indices=None) def test_get_in_canonical_shape_with_indices(self): - for indices in self._test_indices(): + """Test get_in_canonical_shape with 1-dimensional indices.""" + for indices in self._valid_indices(): self._test_in_canonical_shape(indices=indices) def test_get_in_canonical_shape_with_2d_indices(self): + """Test get_in_canonical_shape with 2-dimensional indices.""" indices = torch.randint(self.num, size=(self.batch_size, 2)) self._test_in_canonical_shape(indices=indices) From 5fd64d883aca9b7df3d6f8f593033fd16d8ec442 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 10:58:18 +0100 Subject: [PATCH 570/690] Add test test for representations --- tests/test_nn.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index 0284f27ad8..e8b732c7a8 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -11,7 +11,7 @@ from pykeen.nn import Embedding, LiteralRepresentations, RepresentationModule from pykeen.nn.representation import DIMS, get_expected_canonical_shape from pykeen.nn.sim import kullback_leibler_similarity -from pykeen.testing.base import GenericTests +from pykeen.testing.base import GenericTests, TestsTest from pykeen.typing import GaussianDistribution @@ -154,6 +154,13 @@ def _verify_content(self, x, indices): # noqa: D102 assert torch.allclose(x, exp_x) +class RepresentationModuleTestsTest(TestsTest[RepresentationModule], unittest.TestCase): + """Test that there are tests for all representation modules.""" + + base_cls = RepresentationModule + base_test = RepresentationModuleTests + + class KullbackLeiblerTests(unittest.TestCase): """Tests for the vectorized computation of KL divergences.""" From 96c32c1e4efef8687d02b65bd6f1bdb9e6e58fcc Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 11:19:31 +0100 Subject: [PATCH 571/690] Add test for EmbeddingSpecification --- tests/test_nn.py | 43 ++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 42 insertions(+), 1 deletion(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index e8b732c7a8..568ba1a76b 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -4,11 +4,13 @@ import itertools import unittest from typing import Any, Iterable, MutableMapping, Optional, Sequence +from unittest.mock import Mock +import numpy import pytest import torch -from pykeen.nn import Embedding, LiteralRepresentations, RepresentationModule +from pykeen.nn import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule from pykeen.nn.representation import DIMS, get_expected_canonical_shape from pykeen.nn.sim import kullback_leibler_similarity from pykeen.testing.base import GenericTests, TestsTest @@ -161,6 +163,45 @@ class RepresentationModuleTestsTest(TestsTest[RepresentationModule], unittest.Te base_test = RepresentationModuleTests +class EmbeddingSpecificationTests(unittest.TestCase): + """Tests for EmbeddingSpecification.""" + + #: The number of embeddings + num: int = 3 + + def test_make(self): + """Test make.""" + initializer = Mock() + normalizer = Mock() + constrainer = Mock() + regularizer = Mock() + for embedding_dim, shape in [ + (None, (3,)), + (None, (3, 5)), + (3, None), + ]: + spec = EmbeddingSpecification( + embedding_dim=embedding_dim, + shape=shape, + initializer=initializer, + normalizer=normalizer, + constrainer=constrainer, + regularizer=regularizer, + ) + emb = spec.make(num_embeddings=self.num) + + # check shape + assert emb.embedding_dim == (embedding_dim or int(numpy.prod(shape))) + assert emb.shape == (shape or (embedding_dim,)) + assert emb.num_embeddings == self.num + + # check attributes + assert emb.initializer is initializer + assert emb.normalizer is normalizer + assert emb.constrainer is constrainer + assert emb.regularizer is regularizer + + class KullbackLeiblerTests(unittest.TestCase): """Tests for the vectorized computation of KL divergences.""" From 2649b95779ce16f79013572070eb0d6ca313273b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 11:21:09 +0100 Subject: [PATCH 572/690] Add test for constructor error raising --- tests/test_nn.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/tests/test_nn.py b/tests/test_nn.py index 568ba1a76b..2210d844fa 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -126,6 +126,19 @@ class EmbeddingTests(RepresentationModuleTests, unittest.TestCase): shape=RepresentationModuleTests.exp_shape, ) + def test_constructor_errors(self): + """Test error cases for constructor call.""" + for embedding_dim, shape in ( + (None, None), # neither + (3, (5, 3)), # both + ): + with pytest.raises(ValueError): + Embedding( + num_embeddings=self.num, + embedding_dim=embedding_dim, + shape=shape, + ) + class TensorEmbeddingTests(RepresentationModuleTests, unittest.TestCase): """Tests for Embedding with 2-dimensional shape.""" From c07c975e5c92e27bb22ec7dc29420703d79e1330 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 11:34:43 +0100 Subject: [PATCH 573/690] Add test for initializer --- tests/test_nn.py | 35 +++++++++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index 2210d844fa..fc670b15a5 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -3,8 +3,8 @@ """Unittest for the :mod:`pykeen.nn` module.""" import itertools import unittest -from typing import Any, Iterable, MutableMapping, Optional, Sequence -from unittest.mock import Mock +from typing import Any, Iterable, Mapping, MutableMapping, Optional, Sequence +from unittest.mock import MagicMock, Mock import numpy import pytest @@ -139,6 +139,37 @@ def test_constructor_errors(self): shape=shape, ) + def _test_initializer( + self, + initializer=torch.nn.init.normal_, + kwargs: Optional[Mapping[str, Any]] = None, + ): + wrapped_initializer = MagicMock(side_effect=initializer) + embedding_kwargs = dict() + if kwargs is not None: + embedding_kwargs["initializer_kwargs"] = kwargs + embedding = Embedding( + num_embeddings=self.num, + shape=self.exp_shape, + initializer=wrapped_initializer, + **embedding_kwargs + ) + wrapped_initializer.assert_not_called() + embedding.reset_parameters() + wrapped_initializer.assert_called_once() + # one positional + assert len(wrapped_initializer.call_args.args) == 1 + # additional key-word based + assert len(wrapped_initializer.call_args.kwargs) == len(kwargs or {}) + + def test_initializer(self): + """Test initializer.""" + self._test_initializer() + + def test_initializer_with_kwargs(self): + """Test initializer with kwargs.""" + self._test_initializer(kwargs=dict(mean=3)) + class TensorEmbeddingTests(RepresentationModuleTests, unittest.TestCase): """Tests for Embedding with 2-dimensional shape.""" From 3bf23c4a49a139d7a42a8aa4d9489d6f6f2db632 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 11:36:01 +0100 Subject: [PATCH 574/690] Add documentation --- tests/test_nn.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index fc670b15a5..953ad7f3c0 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -144,7 +144,11 @@ def _test_initializer( initializer=torch.nn.init.normal_, kwargs: Optional[Mapping[str, Any]] = None, ): + """Test initializer usage.""" + # wrap to check calls wrapped_initializer = MagicMock(side_effect=initializer) + + # instantiate embedding embedding_kwargs = dict() if kwargs is not None: embedding_kwargs["initializer_kwargs"] = kwargs @@ -154,12 +158,17 @@ def _test_initializer( initializer=wrapped_initializer, **embedding_kwargs ) + # check that initializer gets not called before reset_parameters wrapped_initializer.assert_not_called() + + # check that initializer gets called exactly once in reset_parameters embedding.reset_parameters() wrapped_initializer.assert_called_once() - # one positional + + # .. with one positional argument ... assert len(wrapped_initializer.call_args.args) == 1 - # additional key-word based + + # .. and additional key-word based arguments. assert len(wrapped_initializer.call_args.kwargs) == len(kwargs or {}) def test_initializer(self): From 4098723be6dd7f1ceaac33b5611cb6375f3a75a6 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 11:36:43 +0100 Subject: [PATCH 575/690] Add trailing comma --- tests/test_nn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index 953ad7f3c0..bee8968573 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -156,7 +156,7 @@ def _test_initializer( num_embeddings=self.num, shape=self.exp_shape, initializer=wrapped_initializer, - **embedding_kwargs + **embedding_kwargs, ) # check that initializer gets not called before reset_parameters wrapped_initializer.assert_not_called() From 8d4a682d851e471da0cc9176169a80f2a810c260 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 11:48:12 +0100 Subject: [PATCH 576/690] Add tests for normalizer --- tests/test_nn.py | 81 ++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 65 insertions(+), 16 deletions(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index bee8968573..43ce3f2c03 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -9,6 +9,7 @@ import numpy import pytest import torch +from torch.nn import functional from pykeen.nn import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule from pykeen.nn.representation import DIMS, get_expected_canonical_shape @@ -139,45 +140,93 @@ def test_constructor_errors(self): shape=shape, ) - def _test_initializer( + def _test_func_with_kwargs( self, - initializer=torch.nn.init.normal_, + name: str, + func, + reset_parameters_call: bool, + forward_call: bool, kwargs: Optional[Mapping[str, Any]] = None, ): """Test initializer usage.""" # wrap to check calls - wrapped_initializer = MagicMock(side_effect=initializer) + wrapped = MagicMock(side_effect=func) # instantiate embedding - embedding_kwargs = dict() + embedding_kwargs = {name: wrapped} if kwargs is not None: - embedding_kwargs["initializer_kwargs"] = kwargs + embedding_kwargs[f"{name}_kwargs"] = kwargs embedding = Embedding( num_embeddings=self.num, shape=self.exp_shape, - initializer=wrapped_initializer, **embedding_kwargs, ) - # check that initializer gets not called before reset_parameters - wrapped_initializer.assert_not_called() - # check that initializer gets called exactly once in reset_parameters + # check that nothing gets called in constructor + wrapped.assert_not_called() + exp_call_count = 0 + + # check call in reset_parameters embedding.reset_parameters() - wrapped_initializer.assert_called_once() + if reset_parameters_call: + exp_call_count += 1 + + # called with one positional argument ... + assert len(wrapped.call_args.args) == 1 + + # .. and additional key-word based arguments. + assert len(wrapped.call_args.kwargs) == len(kwargs or {}) + assert wrapped.call_count == exp_call_count - # .. with one positional argument ... - assert len(wrapped_initializer.call_args.args) == 1 + # check call in forward + embedding.forward(indices=None) + if forward_call: + exp_call_count += 1 - # .. and additional key-word based arguments. - assert len(wrapped_initializer.call_args.kwargs) == len(kwargs or {}) + # called with one positional argument ... + assert len(wrapped.call_args.args) == 1 + + # .. and additional key-word based arguments. + assert len(wrapped.call_args.kwargs) == len(kwargs or {}) + assert wrapped.call_count == exp_call_count def test_initializer(self): """Test initializer.""" - self._test_initializer() + self._test_func_with_kwargs( + name="initializer", + func=torch.nn.init.normal_, + reset_parameters_call=True, + forward_call=False, + ) def test_initializer_with_kwargs(self): """Test initializer with kwargs.""" - self._test_initializer(kwargs=dict(mean=3)) + self._test_func_with_kwargs( + name="initializer", + func=torch.nn.init.normal_, + reset_parameters_call=True, + forward_call=False, + kwargs=dict(mean=3) + ) + + def test_normalizer(self): + """Test normalizer.""" + self._test_func_with_kwargs( + name="normalizer", + func=functional.normalize, + reset_parameters_call=False, + forward_call=True, + ) + + def test_normalizer_kwargs(self): + """Test normalizer with kwargs.""" + self._test_func_with_kwargs( + name="normalizer", + func=functional.normalize, + reset_parameters_call=False, + forward_call=True, + kwargs=dict(p=1), + ) class TensorEmbeddingTests(RepresentationModuleTests, unittest.TestCase): From 5bac9f6844ce1a8079d7006ab45408d7be320345 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 11:55:34 +0100 Subject: [PATCH 577/690] Add tests for constrainer; simplify code --- tests/test_nn.py | 78 +++++++++++++++++++++++++++++++++++------------- 1 file changed, 58 insertions(+), 20 deletions(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index 43ce3f2c03..01b375ff01 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -144,9 +144,10 @@ def _test_func_with_kwargs( self, name: str, func, - reset_parameters_call: bool, - forward_call: bool, kwargs: Optional[Mapping[str, Any]] = None, + reset_parameters_call: bool = False, + forward_call: bool = False, + post_parameter_update_call: bool = False, ): """Test initializer usage.""" # wrap to check calls @@ -164,31 +165,55 @@ def _test_func_with_kwargs( # check that nothing gets called in constructor wrapped.assert_not_called() - exp_call_count = 0 + call_count = 0 # check call in reset_parameters embedding.reset_parameters() - if reset_parameters_call: - exp_call_count += 1 - - # called with one positional argument ... - assert len(wrapped.call_args.args) == 1 - - # .. and additional key-word based arguments. - assert len(wrapped.call_args.kwargs) == len(kwargs or {}) - assert wrapped.call_count == exp_call_count + call_count = self._check_call( + call_count=call_count, + should_be_called=reset_parameters_call, + wrapped=wrapped, + kwargs=kwargs, + ) # check call in forward embedding.forward(indices=None) - if forward_call: - exp_call_count += 1 + call_count = self._check_call( + call_count=call_count, + should_be_called=forward_call, + wrapped=wrapped, + kwargs=kwargs, + ) + + # check call in post_parameter_update + embedding.post_parameter_update() + call_count = self._check_call( + call_count=call_count, + should_be_called=post_parameter_update_call, + wrapped=wrapped, + kwargs=kwargs, + ) + + def _check_call( + self, + call_count: int, + should_be_called: bool, + wrapped: MagicMock, + kwargs: Optional[Mapping[str, Any]], + ) -> int: + if should_be_called: + call_count += 1 + + assert wrapped.call_count == call_count # called with one positional argument ... assert len(wrapped.call_args.args) == 1 # .. and additional key-word based arguments. assert len(wrapped.call_args.kwargs) == len(kwargs or {}) - assert wrapped.call_count == exp_call_count + else: + assert wrapped.call_count == call_count + return call_count def test_initializer(self): """Test initializer.""" @@ -196,7 +221,6 @@ def test_initializer(self): name="initializer", func=torch.nn.init.normal_, reset_parameters_call=True, - forward_call=False, ) def test_initializer_with_kwargs(self): @@ -204,9 +228,8 @@ def test_initializer_with_kwargs(self): self._test_func_with_kwargs( name="initializer", func=torch.nn.init.normal_, + kwargs=dict(mean=3), reset_parameters_call=True, - forward_call=False, - kwargs=dict(mean=3) ) def test_normalizer(self): @@ -214,7 +237,6 @@ def test_normalizer(self): self._test_func_with_kwargs( name="normalizer", func=functional.normalize, - reset_parameters_call=False, forward_call=True, ) @@ -223,9 +245,25 @@ def test_normalizer_kwargs(self): self._test_func_with_kwargs( name="normalizer", func=functional.normalize, - reset_parameters_call=False, + kwargs=dict(p=1), forward_call=True, + ) + + def test_constrainer(self): + """Test constrainer.""" + self._test_func_with_kwargs( + name="constrainer", + func=functional.normalize, + post_parameter_update_call=True, + ) + + def test_constrainer_kwargs(self): + """Test constrainer with kwargs.""" + self._test_func_with_kwargs( + name="constrainer", + func=functional.normalize, kwargs=dict(p=1), + post_parameter_update_call=True, ) From 5d1336950e6347ccf6fe60cec7c2a575a56a08b1 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 11:56:40 +0100 Subject: [PATCH 578/690] Extract method --- tests/test_nn.py | 63 ++++++++++++++++++++++++++++++------------------ 1 file changed, 39 insertions(+), 24 deletions(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index 01b375ff01..ae6a1aae93 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -118,6 +118,42 @@ def test_get_in_canonical_shape_with_2d_indices(self): self._test_in_canonical_shape(indices=indices) +def _check_call( + call_count: int, + should_be_called: bool, + wrapped: MagicMock, + kwargs: Optional[Mapping[str, Any]], +) -> int: + """ + Check whether a wrapped method is called. + + :param call_count: + The previous call count. + :param should_be_called: + Whether it should be called. + :param wrapped: + The wrapped method. + :param kwargs: + The expected kwargs when called. + + :return: + The updated counter. + """ + if should_be_called: + call_count += 1 + + assert wrapped.call_count == call_count + + # called with one positional argument ... + assert len(wrapped.call_args.args) == 1 + + # .. and additional key-word based arguments. + assert len(wrapped.call_args.kwargs) == len(kwargs or {}) + else: + assert wrapped.call_count == call_count + return call_count + + class EmbeddingTests(RepresentationModuleTests, unittest.TestCase): """Tests for Embedding.""" @@ -169,7 +205,7 @@ def _test_func_with_kwargs( # check call in reset_parameters embedding.reset_parameters() - call_count = self._check_call( + call_count = _check_call( call_count=call_count, should_be_called=reset_parameters_call, wrapped=wrapped, @@ -178,7 +214,7 @@ def _test_func_with_kwargs( # check call in forward embedding.forward(indices=None) - call_count = self._check_call( + call_count = _check_call( call_count=call_count, should_be_called=forward_call, wrapped=wrapped, @@ -187,34 +223,13 @@ def _test_func_with_kwargs( # check call in post_parameter_update embedding.post_parameter_update() - call_count = self._check_call( + call_count = _check_call( call_count=call_count, should_be_called=post_parameter_update_call, wrapped=wrapped, kwargs=kwargs, ) - def _check_call( - self, - call_count: int, - should_be_called: bool, - wrapped: MagicMock, - kwargs: Optional[Mapping[str, Any]], - ) -> int: - if should_be_called: - call_count += 1 - - assert wrapped.call_count == call_count - - # called with one positional argument ... - assert len(wrapped.call_args.args) == 1 - - # .. and additional key-word based arguments. - assert len(wrapped.call_args.kwargs) == len(kwargs or {}) - else: - assert wrapped.call_count == call_count - return call_count - def test_initializer(self): """Test initializer.""" self._test_func_with_kwargs( From 2607fd835d8a57468ec1f532481c031a9b5d1930 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 11:56:54 +0100 Subject: [PATCH 579/690] Delete unused variable --- tests/test_nn.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index ae6a1aae93..5f2efcf452 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -223,7 +223,7 @@ def _test_func_with_kwargs( # check call in post_parameter_update embedding.post_parameter_update() - call_count = _check_call( + _check_call( call_count=call_count, should_be_called=post_parameter_update_call, wrapped=wrapped, From f32f748ab26a0dcfaf7b4e05757f5a8b9c8a9bb0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 11:58:33 +0100 Subject: [PATCH 580/690] Subclass LiteralRepresentationsTests from correct test super class --- tests/test_nn.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index 5f2efcf452..2d60ee753f 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -293,13 +293,14 @@ class TensorEmbeddingTests(RepresentationModuleTests, unittest.TestCase): ) -class LiteralRepresentationsTests(RepresentationModuleTests, unittest.TestCase): +class LiteralRepresentationsTests(EmbeddingTests, unittest.TestCase): """Tests for literal embeddings.""" cls = LiteralRepresentations def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102 - kwargs = super()._pre_instantiation_hook(kwargs=kwargs) + # requires own kwargs + kwargs.clear() self.numeric_literals = torch.rand(self.num, *self.exp_shape) kwargs["numeric_literals"] = self.numeric_literals return kwargs From 8290fd1a30e673dbce4e04373d7648f8983be555 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:03:27 +0100 Subject: [PATCH 581/690] Move RGCN representations to representation.py --- src/pykeen/models/unimodal/rgcn.py | 401 +---------------------------- src/pykeen/nn/representation.py | 365 +++++++++++++++++++++++++- tests/test_models.py | 6 +- 3 files changed, 367 insertions(+), 405 deletions(-) diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index 4007462e8f..1e9ed5229a 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -2,19 +2,17 @@ """Implementation of the R-GCN model.""" -import logging -from os import path from typing import Any, Callable, Mapping, Optional, Type import torch from torch import nn -from torch.nn import functional from . import ComplEx, DistMult, ERMLP from ..base import ERModel from ...losses import Loss -from ...nn import Embedding, EmbeddingSpecification, Interaction, RepresentationModule +from ...nn import EmbeddingSpecification, Interaction from ...nn.modules import DistMultInteraction +from ...nn.representation import RGCNRepresentations, inverse_indegree_edge_weights, inverse_outdegree_edge_weights, symmetric_edge_weights from ...triples import TriplesFactory from ...typing import DeviceHint @@ -22,401 +20,6 @@ 'RGCN', ] -logger = logging.getLogger(name=path.basename(__file__)) - - -def _get_neighborhood( - start_nodes: torch.LongTensor, - sources: torch.LongTensor, - targets: torch.LongTensor, - k: int, - num_nodes: int, - undirected: bool = False, -) -> torch.BoolTensor: - # Construct node neighbourhood mask - node_mask = torch.zeros(num_nodes, device=start_nodes.device, dtype=torch.bool) - - # Set nodes in batch to true - node_mask[start_nodes] = True - - # Compute k-neighbourhood - for _ in range(k): - # if the target node needs an embeddings, so does the source node - node_mask[sources] |= node_mask[targets] - - if undirected: - node_mask[targets] |= node_mask[sources] - - # Create edge mask - edge_mask = node_mask[targets] - - if undirected: - edge_mask |= node_mask[sources] - - return edge_mask - - -# pylint: disable=unused-argument -def inverse_indegree_edge_weights(source: torch.LongTensor, target: torch.LongTensor) -> torch.FloatTensor: - """Normalize messages by inverse in-degree. - - :param source: shape: (num_edges,) - The source indices. - :param target: shape: (num_edges,) - The target indices. - - :return: shape: (num_edges,) - The edge weights. - """ - # Calculate in-degree, i.e. number of incoming edges - uniq, inv, cnt = torch.unique(target, return_counts=True, return_inverse=True) - return cnt[inv].float().reciprocal() - - -# pylint: disable=unused-argument -def inverse_outdegree_edge_weights(source: torch.LongTensor, target: torch.LongTensor) -> torch.FloatTensor: - """Normalize messages by inverse out-degree. - - :param source: shape: (num_edges,) - The source indices. - :param target: shape: (num_edges,) - The target indices. - - :return: shape: (num_edges,) - The edge weights. - """ - # Calculate in-degree, i.e. number of incoming edges - uniq, inv, cnt = torch.unique(source, return_counts=True, return_inverse=True) - return cnt[inv].float().reciprocal() - - -def symmetric_edge_weights(source: torch.LongTensor, target: torch.LongTensor) -> torch.FloatTensor: - """Normalize messages by product of inverse sqrt of in-degree and out-degree. - - :param source: shape: (num_edges,) - The source indices. - :param target: shape: (num_edges,) - The target indices. - - :return: shape: (num_edges,) - The edge weights. - """ - return ( - inverse_indegree_edge_weights(source=source, target=target) - * inverse_outdegree_edge_weights(source=source, target=target) - ).sqrt() - - -class RGCNRepresentations(RepresentationModule): - """Representations enriched by R-GCN.""" - - def __init__( - self, - triples_factory: TriplesFactory, - embedding_dim: int = 500, - num_bases_or_blocks: int = 5, - num_layers: int = 2, - use_bias: bool = True, - use_batch_norm: bool = False, - activation_cls: Optional[Type[nn.Module]] = None, - activation_kwargs: Optional[Mapping[str, Any]] = None, - sparse_messages_slcwa: bool = True, - edge_dropout: float = 0.4, - self_loop_dropout: float = 0.2, - edge_weighting: Callable[ - [torch.LongTensor, torch.LongTensor], - torch.FloatTensor, - ] = inverse_indegree_edge_weights, - decomposition: str = 'basis', - buffer_messages: bool = True, - base_representations: Optional[RepresentationModule] = None, - ): - super().__init__(shape=(embedding_dim,), max_id=triples_factory.num_entities) - - self.triples_factory = triples_factory - - # normalize representations - if base_representations is None: - base_representations = Embedding( - num_embeddings=triples_factory.num_entities, - embedding_dim=embedding_dim, - # https://github.com/MichSchli/RelationPrediction/blob/c77b094fe5c17685ed138dae9ae49b304e0d8d89/code/encoders/affine_transform.py#L24-L28 - initializer=nn.init.xavier_uniform_, - ) - self.base_embeddings = base_representations - self.embedding_dim = embedding_dim - - # check decomposition - self.decomposition = decomposition - if self.decomposition == 'basis': - if num_bases_or_blocks is None: - logging.info('Using a heuristic to determine the number of bases.') - num_bases_or_blocks = triples_factory.num_relations // 2 + 1 - if num_bases_or_blocks > triples_factory.num_relations: - raise ValueError('The number of bases should not exceed the number of relations.') - elif self.decomposition == 'block': - if num_bases_or_blocks is None: - logging.info('Using a heuristic to determine the number of blocks.') - num_bases_or_blocks = 2 - if embedding_dim % num_bases_or_blocks != 0: - raise ValueError( - 'With block decomposition, the embedding dimension has to be divisible by the number of' - f' blocks, but {embedding_dim} % {num_bases_or_blocks} != 0.', - ) - else: - raise ValueError(f'Unknown decomposition: "{decomposition}". Please use either "basis" or "block".') - - self.num_bases = num_bases_or_blocks - self.edge_weighting = edge_weighting - self.edge_dropout = edge_dropout - if self_loop_dropout is None: - self_loop_dropout = edge_dropout - self.self_loop_dropout = self_loop_dropout - self.use_batch_norm = use_batch_norm - if activation_cls is None: - activation_cls = nn.ReLU - self.activation_cls = activation_cls - self.activation_kwargs = activation_kwargs - if use_batch_norm: - if use_bias: - logger.warning('Disabling bias because batch normalization was used.') - use_bias = False - self.use_bias = use_bias - self.num_layers = num_layers - self.sparse_messages_slcwa = sparse_messages_slcwa - - # Save graph using buffers, such that the tensors are moved together with the model - h, r, t = self.triples_factory.mapped_triples.t() - self.register_buffer('sources', h) - self.register_buffer('targets', t) - self.register_buffer('edge_types', r) - - self.activations = nn.ModuleList([ - self.activation_cls(**(self.activation_kwargs or {})) for _ in range(self.num_layers) - ]) - - # Weights - self.bases = nn.ParameterList() - if self.decomposition == 'basis': - self.att = nn.ParameterList() - for _ in range(self.num_layers): - self.bases.append(nn.Parameter( - data=torch.empty( - self.num_bases, - self.embedding_dim, - self.embedding_dim, - ), - requires_grad=True, - )) - self.att.append(nn.Parameter( - data=torch.empty( - self.triples_factory.num_relations + 1, - self.num_bases, - ), - requires_grad=True, - )) - elif self.decomposition == 'block': - block_size = self.embedding_dim // self.num_bases - for _ in range(self.num_layers): - self.bases.append(nn.Parameter( - data=torch.empty( - self.triples_factory.num_relations + 1, - self.num_bases, - block_size, - block_size, - ), - requires_grad=True, - )) - - self.att = None - else: - raise NotImplementedError - if self.use_bias: - self.biases = nn.ParameterList([ - nn.Parameter(torch.empty(self.embedding_dim), requires_grad=True) - for _ in range(self.num_layers) - ]) - else: - self.biases = None - if self.use_batch_norm: - self.batch_norms = nn.ModuleList([ - nn.BatchNorm1d(num_features=self.embedding_dim) - for _ in range(self.num_layers) - ]) - else: - self.batch_norms = None - - # buffering of messages - self.buffer_messages = buffer_messages - self.enriched_embeddings = None - - def _get_relation_weights(self, i_layer: int, r: int) -> torch.FloatTensor: - if self.decomposition == 'block': - # allocate weight - w = torch.zeros(self.embedding_dim, self.embedding_dim, device=self.bases[i_layer].device) - - # Get blocks - this_layer_blocks = self.bases[i_layer] - - # self.bases[i_layer].shape (num_relations, num_blocks, embedding_dim/num_blocks, embedding_dim/num_blocks) - # note: embedding_dim is guaranteed to be divisible by num_bases in the constructor - block_size = self.embedding_dim // self.num_bases - for b, start in enumerate(range(0, self.embedding_dim, block_size)): - stop = start + block_size - w[start:stop, start:stop] = this_layer_blocks[r, b, :, :] - - elif self.decomposition == 'basis': - # The current basis weights, shape: (num_bases) - att = self.att[i_layer][r, :] - # the current bases, shape: (num_bases, embedding_dim, embedding_dim) - b = self.bases[i_layer] - # compute the current relation weights, shape: (embedding_dim, embedding_dim) - w = torch.sum(att[:, None, None] * b, dim=0) - - else: - raise AssertionError(f'Unknown decomposition: {self.decomposition}') - - return w - - def forward( - self, - indices: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: - # use buffered messages if applicable - if indices is None and self.enriched_embeddings is not None: - return self.enriched_embeddings - - # Bind fields - # shape: (num_entities, embedding_dim) - x = self.base_embeddings(indices=None) - sources = self.sources - targets = self.targets - edge_types = self.edge_types - - # Edge dropout: drop the same edges on all layers (only in training mode) - if self.training and self.edge_dropout is not None: - # Get random dropout mask - edge_keep_mask = torch.rand(self.sources.shape[0], device=x.device) > self.edge_dropout - - # Apply to edges - sources = sources[edge_keep_mask] - targets = targets[edge_keep_mask] - edge_types = edge_types[edge_keep_mask] - - # Different dropout for self-loops (only in training mode) - if self.training and self.self_loop_dropout is not None: - node_keep_mask = torch.rand(self.triples_factory.num_entities, device=x.device) > self.self_loop_dropout - else: - node_keep_mask = None - - for i in range(self.num_layers): - # Initialize embeddings in the next layer for all nodes - new_x = torch.zeros_like(x) - - # TODO: Can we vectorize this loop? - for r in range(self.triples_factory.num_relations): - # Choose the edges which are of the specific relation - mask = (edge_types == r) - - # No edges available? Skip rest of inner loop - if not mask.any(): - continue - - # Get source and target node indices - sources_r = sources[mask] - targets_r = targets[mask] - - # send messages in both directions - sources_r, targets_r = torch.cat([sources_r, targets_r]), torch.cat([targets_r, sources_r]) - - # Select source node embeddings - x_s = x[sources_r] - - # get relation weights - w = self._get_relation_weights(i_layer=i, r=r) - - # Compute message (b x d) * (d x d) = (b x d) - m_r = x_s @ w - - # Normalize messages by relation-specific in-degree - if self.edge_weighting is not None: - m_r *= self.edge_weighting(sources_r, targets_r).unsqueeze(dim=-1) - - # Aggregate messages in target - new_x.index_add_(dim=0, index=targets_r, source=m_r) - - # Self-loop - self_w = self._get_relation_weights(i_layer=i, r=self.triples_factory.num_relations) - if node_keep_mask is None: - new_x += x @ self_w - else: - new_x[node_keep_mask] += x[node_keep_mask] @ self_w - - # Apply bias, if requested - if self.use_bias: - bias = self.biases[i] - new_x += bias - - # Apply batch normalization, if requested - if self.use_batch_norm: - batch_norm = self.batch_norms[i] - new_x = batch_norm(new_x) - - # Apply non-linearity - if self.activations is not None: - activation = self.activations[i] - new_x = activation(new_x) - - x = new_x - - if indices is None and self.buffer_messages: - self.enriched_embeddings = x - if indices is not None: - x = x[indices] - - return x - - def post_parameter_update(self) -> None: # noqa: D102 - super().post_parameter_update() - - # invalidate enriched embeddings - self.enriched_embeddings = None - - def reset_parameters(self): - self.base_embeddings.reset_parameters() - - gain = nn.init.calculate_gain(nonlinearity=self.activation_cls.__name__.lower()) - if self.decomposition == 'basis': - for base in self.bases: - nn.init.xavier_normal_(base, gain=gain) - for att in self.att: - # Random convex-combination of bases for initialization (guarantees that initial weight matrices are - # initialized properly) - # We have one additional relation for self-loops - nn.init.uniform_(att) - functional.normalize(att.data, p=1, dim=1, out=att.data) - elif self.decomposition == 'block': - for base in self.bases: - block_size = base.shape[-1] - # Xavier Glorot initialization of each block - std = torch.sqrt(torch.as_tensor(2.)) * gain / (2 * block_size) - nn.init.normal_(base, std=std) - - # Reset biases - if self.biases is not None: - for bias in self.biases: - nn.init.zeros_(bias) - - # Reset batch norm parameters - if self.batch_norms is not None: - for bn in self.batch_norms: - bn.reset_parameters() - - # Reset activation parameters, if any - for act in self.activations: - if hasattr(act, 'reset_parameters'): - act.reset_parameters() - class RGCN(ERModel): """An implementation of R-GCN from [schlichtkrull2018]_. diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index d1286f11b5..53322cd2f0 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -6,14 +6,16 @@ import functools import logging from abc import ABC, abstractmethod -from typing import Any, Iterable, Mapping, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Iterable, Mapping, Optional, Sequence, Tuple, Type, Union import numpy import torch import torch.nn from torch import nn +from torch.nn import functional from ..regularizers import Regularizer +from ..triples import TriplesFactory from ..typing import Constrainer, Initializer, Normalizer from ..utils import upgrade_to_sequence @@ -22,6 +24,7 @@ 'Embedding', 'EmbeddingSpecification', 'LiteralRepresentations', + 'RGCNRepresentations', ] logger = logging.getLogger(__name__) @@ -398,3 +401,363 @@ def __init__( ) # freeze self._embeddings.requires_grad_(False) + + +def inverse_indegree_edge_weights(source: torch.LongTensor, target: torch.LongTensor) -> torch.FloatTensor: + """Normalize messages by inverse in-degree. + + :param source: shape: (num_edges,) + The source indices. + :param target: shape: (num_edges,) + The target indices. + + :return: shape: (num_edges,) + The edge weights. + """ + # Calculate in-degree, i.e. number of incoming edges + uniq, inv, cnt = torch.unique(target, return_counts=True, return_inverse=True) + return cnt[inv].float().reciprocal() + + +def inverse_outdegree_edge_weights(source: torch.LongTensor, target: torch.LongTensor) -> torch.FloatTensor: + """Normalize messages by inverse out-degree. + + :param source: shape: (num_edges,) + The source indices. + :param target: shape: (num_edges,) + The target indices. + + :return: shape: (num_edges,) + The edge weights. + """ + # Calculate in-degree, i.e. number of incoming edges + uniq, inv, cnt = torch.unique(source, return_counts=True, return_inverse=True) + return cnt[inv].float().reciprocal() + + +def symmetric_edge_weights(source: torch.LongTensor, target: torch.LongTensor) -> torch.FloatTensor: + """Normalize messages by product of inverse sqrt of in-degree and out-degree. + + :param source: shape: (num_edges,) + The source indices. + :param target: shape: (num_edges,) + The target indices. + + :return: shape: (num_edges,) + The edge weights. + """ + return ( + inverse_indegree_edge_weights(source=source, target=target) + * inverse_outdegree_edge_weights(source=source, target=target) + ).sqrt() + + +class RGCNRepresentations(RepresentationModule): + """Representations enriched by R-GCN.""" + + def __init__( + self, + triples_factory: TriplesFactory, + embedding_dim: int = 500, + num_bases_or_blocks: int = 5, + num_layers: int = 2, + use_bias: bool = True, + use_batch_norm: bool = False, + activation_cls: Optional[Type[nn.Module]] = None, + activation_kwargs: Optional[Mapping[str, Any]] = None, + sparse_messages_slcwa: bool = True, + edge_dropout: float = 0.4, + self_loop_dropout: float = 0.2, + edge_weighting: Callable[ + [torch.LongTensor, torch.LongTensor], + torch.FloatTensor, + ] = inverse_indegree_edge_weights, + decomposition: str = 'basis', + buffer_messages: bool = True, + base_representations: Optional[RepresentationModule] = None, + ): + super().__init__(shape=(embedding_dim,), max_id=triples_factory.num_entities) + + self.triples_factory = triples_factory + + # normalize representations + if base_representations is None: + base_representations = Embedding( + num_embeddings=triples_factory.num_entities, + embedding_dim=embedding_dim, + # https://github.com/MichSchli/RelationPrediction/blob/c77b094fe5c17685ed138dae9ae49b304e0d8d89/code/encoders/affine_transform.py#L24-L28 + initializer=nn.init.xavier_uniform_, + ) + self.base_embeddings = base_representations + self.embedding_dim = embedding_dim + + # check decomposition + self.decomposition = decomposition + if self.decomposition == 'basis': + if num_bases_or_blocks is None: + logging.info('Using a heuristic to determine the number of bases.') + num_bases_or_blocks = triples_factory.num_relations // 2 + 1 + if num_bases_or_blocks > triples_factory.num_relations: + raise ValueError('The number of bases should not exceed the number of relations.') + elif self.decomposition == 'block': + if num_bases_or_blocks is None: + logging.info('Using a heuristic to determine the number of blocks.') + num_bases_or_blocks = 2 + if embedding_dim % num_bases_or_blocks != 0: + raise ValueError( + 'With block decomposition, the embedding dimension has to be divisible by the number of' + f' blocks, but {embedding_dim} % {num_bases_or_blocks} != 0.', + ) + else: + raise ValueError(f'Unknown decomposition: "{decomposition}". Please use either "basis" or "block".') + + self.num_bases = num_bases_or_blocks + self.edge_weighting = edge_weighting + self.edge_dropout = edge_dropout + if self_loop_dropout is None: + self_loop_dropout = edge_dropout + self.self_loop_dropout = self_loop_dropout + self.use_batch_norm = use_batch_norm + if activation_cls is None: + activation_cls = nn.ReLU + self.activation_cls = activation_cls + self.activation_kwargs = activation_kwargs + if use_batch_norm: + if use_bias: + logger.warning('Disabling bias because batch normalization was used.') + use_bias = False + self.use_bias = use_bias + self.num_layers = num_layers + self.sparse_messages_slcwa = sparse_messages_slcwa + + # Save graph using buffers, such that the tensors are moved together with the model + h, r, t = self.triples_factory.mapped_triples.t() + self.register_buffer('sources', h) + self.register_buffer('targets', t) + self.register_buffer('edge_types', r) + + self.activations = nn.ModuleList([ + self.activation_cls(**(self.activation_kwargs or {})) for _ in range(self.num_layers) + ]) + + # Weights + self.bases = nn.ParameterList() + if self.decomposition == 'basis': + self.att = nn.ParameterList() + for _ in range(self.num_layers): + self.bases.append(nn.Parameter( + data=torch.empty( + self.num_bases, + self.embedding_dim, + self.embedding_dim, + ), + requires_grad=True, + )) + self.att.append(nn.Parameter( + data=torch.empty( + self.triples_factory.num_relations + 1, + self.num_bases, + ), + requires_grad=True, + )) + elif self.decomposition == 'block': + block_size = self.embedding_dim // self.num_bases + for _ in range(self.num_layers): + self.bases.append(nn.Parameter( + data=torch.empty( + self.triples_factory.num_relations + 1, + self.num_bases, + block_size, + block_size, + ), + requires_grad=True, + )) + + self.att = None + else: + raise NotImplementedError + if self.use_bias: + self.biases = nn.ParameterList([ + nn.Parameter(torch.empty(self.embedding_dim), requires_grad=True) + for _ in range(self.num_layers) + ]) + else: + self.biases = None + if self.use_batch_norm: + self.batch_norms = nn.ModuleList([ + nn.BatchNorm1d(num_features=self.embedding_dim) + for _ in range(self.num_layers) + ]) + else: + self.batch_norms = None + + # buffering of messages + self.buffer_messages = buffer_messages + self.enriched_embeddings = None + + def _get_relation_weights(self, i_layer: int, r: int) -> torch.FloatTensor: + if self.decomposition == 'block': + # allocate weight + w = torch.zeros(self.embedding_dim, self.embedding_dim, device=self.bases[i_layer].device) + + # Get blocks + this_layer_blocks = self.bases[i_layer] + + # self.bases[i_layer].shape (num_relations, num_blocks, embedding_dim/num_blocks, embedding_dim/num_blocks) + # note: embedding_dim is guaranteed to be divisible by num_bases in the constructor + block_size = self.embedding_dim // self.num_bases + for b, start in enumerate(range(0, self.embedding_dim, block_size)): + stop = start + block_size + w[start:stop, start:stop] = this_layer_blocks[r, b, :, :] + + elif self.decomposition == 'basis': + # The current basis weights, shape: (num_bases) + att = self.att[i_layer][r, :] + # the current bases, shape: (num_bases, embedding_dim, embedding_dim) + b = self.bases[i_layer] + # compute the current relation weights, shape: (embedding_dim, embedding_dim) + w = torch.sum(att[:, None, None] * b, dim=0) + + else: + raise AssertionError(f'Unknown decomposition: {self.decomposition}') + + return w + + def forward( + self, + indices: Optional[torch.LongTensor] = None, + ) -> torch.FloatTensor: + # use buffered messages if applicable + if indices is None and self.enriched_embeddings is not None: + return self.enriched_embeddings + + # Bind fields + # shape: (num_entities, embedding_dim) + x = self.base_embeddings(indices=None) + sources = self.sources + targets = self.targets + edge_types = self.edge_types + + # Edge dropout: drop the same edges on all layers (only in training mode) + if self.training and self.edge_dropout is not None: + # Get random dropout mask + edge_keep_mask = torch.rand(self.sources.shape[0], device=x.device) > self.edge_dropout + + # Apply to edges + sources = sources[edge_keep_mask] + targets = targets[edge_keep_mask] + edge_types = edge_types[edge_keep_mask] + + # Different dropout for self-loops (only in training mode) + if self.training and self.self_loop_dropout is not None: + node_keep_mask = torch.rand(self.triples_factory.num_entities, device=x.device) > self.self_loop_dropout + else: + node_keep_mask = None + + for i in range(self.num_layers): + # Initialize embeddings in the next layer for all nodes + new_x = torch.zeros_like(x) + + # TODO: Can we vectorize this loop? + for r in range(self.triples_factory.num_relations): + # Choose the edges which are of the specific relation + mask = (edge_types == r) + + # No edges available? Skip rest of inner loop + if not mask.any(): + continue + + # Get source and target node indices + sources_r = sources[mask] + targets_r = targets[mask] + + # send messages in both directions + sources_r, targets_r = torch.cat([sources_r, targets_r]), torch.cat([targets_r, sources_r]) + + # Select source node embeddings + x_s = x[sources_r] + + # get relation weights + w = self._get_relation_weights(i_layer=i, r=r) + + # Compute message (b x d) * (d x d) = (b x d) + m_r = x_s @ w + + # Normalize messages by relation-specific in-degree + if self.edge_weighting is not None: + m_r *= self.edge_weighting(sources_r, targets_r).unsqueeze(dim=-1) + + # Aggregate messages in target + new_x.index_add_(dim=0, index=targets_r, source=m_r) + + # Self-loop + self_w = self._get_relation_weights(i_layer=i, r=self.triples_factory.num_relations) + if node_keep_mask is None: + new_x += x @ self_w + else: + new_x[node_keep_mask] += x[node_keep_mask] @ self_w + + # Apply bias, if requested + if self.use_bias: + bias = self.biases[i] + new_x += bias + + # Apply batch normalization, if requested + if self.use_batch_norm: + batch_norm = self.batch_norms[i] + new_x = batch_norm(new_x) + + # Apply non-linearity + if self.activations is not None: + activation = self.activations[i] + new_x = activation(new_x) + + x = new_x + + if indices is None and self.buffer_messages: + self.enriched_embeddings = x + if indices is not None: + x = x[indices] + + return x + + def post_parameter_update(self) -> None: # noqa: D102 + super().post_parameter_update() + + # invalidate enriched embeddings + self.enriched_embeddings = None + + def reset_parameters(self): + self.base_embeddings.reset_parameters() + + gain = nn.init.calculate_gain(nonlinearity=self.activation_cls.__name__.lower()) + if self.decomposition == 'basis': + for base in self.bases: + nn.init.xavier_normal_(base, gain=gain) + for att in self.att: + # Random convex-combination of bases for initialization (guarantees that initial weight matrices are + # initialized properly) + # We have one additional relation for self-loops + nn.init.uniform_(att) + functional.normalize(att.data, p=1, dim=1, out=att.data) + elif self.decomposition == 'block': + for base in self.bases: + block_size = base.shape[-1] + # Xavier Glorot initialization of each block + std = torch.sqrt(torch.as_tensor(2.)) * gain / (2 * block_size) + nn.init.normal_(base, std=std) + + # Reset biases + if self.biases is not None: + for bias in self.biases: + nn.init.zeros_(bias) + + # Reset batch norm parameters + if self.batch_norms is not None: + for bn in self.batch_norms: + bn.reset_parameters() + + # Reset activation parameters, if any + for act in self.activations: + if hasattr(act, 'reset_parameters'): + act.reset_parameters() diff --git a/tests/test_models.py b/tests/test_models.py index ff59366772..3a7e3948b7 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -31,11 +31,7 @@ get_novelty_mask, ) from pykeen.models.cli import build_cli_from_cls -from pykeen.models.unimodal.rgcn import ( - RGCNRepresentations, inverse_indegree_edge_weights, - inverse_outdegree_edge_weights, - symmetric_edge_weights, -) +from pykeen.nn.representation import RGCNRepresentations, inverse_indegree_edge_weights, inverse_outdegree_edge_weights, symmetric_edge_weights from pykeen.regularizers import LpRegularizer, collect_regularization_terms from pykeen.testing.mocks import MockRepresentations from pykeen.training import LCWATrainingLoop, SLCWATrainingLoop, TrainingLoop From 32a500bd0a5c1ddfded2a9604d8e56d3dd29a250 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:10:18 +0100 Subject: [PATCH 582/690] Add tests for RGCN representations --- tests/test_nn.py | 34 +++++++++++++++++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index 2d60ee753f..9ef63b9306 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -12,9 +12,10 @@ from torch.nn import functional from pykeen.nn import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule -from pykeen.nn.representation import DIMS, get_expected_canonical_shape +from pykeen.nn.representation import DIMS, RGCNRepresentations, get_expected_canonical_shape from pykeen.nn.sim import kullback_leibler_similarity from pykeen.testing.base import GenericTests, TestsTest +from pykeen.triples import TriplesFactory from pykeen.typing import GaussianDistribution @@ -312,6 +313,37 @@ def _verify_content(self, x, indices): # noqa: D102 assert torch.allclose(x, exp_x) +class RGCNRepresentationTests(RepresentationModuleTests, unittest.TestCase): + """Test RGCN representations.""" + + cls = RGCNRepresentations + kwargs = dict( + num_bases_or_blocks=2, + ) + num_relations: int = 7 + num_triples: int = 31 + num_bases: int = 2 + + def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102 + kwargs = super()._pre_instantiation_hook(kwargs=kwargs) + # generate random triples + mapped_triples = numpy.stack([ + numpy.random.randint(max_id, size=(self.num_triples,)) + for max_id in (self.num, self.num_relations, self.num) + ], axis=-1) + entity_names = [f"e_{i}" for i in range(self.num)] + relation_names = [f"r_{i}" for i in range(self.num_relations)] + triples = numpy.stack([ + [names[i] for i in col.tolist()] + for col, names in zip( + mapped_triples.T, + (entity_names, relation_names, entity_names), + ) + ]) + kwargs["triples_factory"] = TriplesFactory(triples=triples) + return kwargs + + class RepresentationModuleTestsTest(TestsTest[RepresentationModule], unittest.TestCase): """Test that there are tests for all representation modules.""" From fcb3c92ac14c9a88089f70a2634923052ce2253c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:12:20 +0100 Subject: [PATCH 583/690] Fix instantiation of RGCN test instance --- tests/test_nn.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_nn.py b/tests/test_nn.py index 9ef63b9306..33b901d895 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -319,6 +319,7 @@ class RGCNRepresentationTests(RepresentationModuleTests, unittest.TestCase): cls = RGCNRepresentations kwargs = dict( num_bases_or_blocks=2, + embedding_dim=RepresentationModuleTests.exp_shape[0], ) num_relations: int = 7 num_triples: int = 31 From d93599f09b5bf7c2c366eccc3027c46eb5cd06fd Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:15:03 +0100 Subject: [PATCH 584/690] Add exception to RGCN representations for indices in wrong shape --- src/pykeen/nn/representation.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index 53322cd2f0..a48afdb5a2 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -630,6 +630,8 @@ def forward( # use buffered messages if applicable if indices is None and self.enriched_embeddings is not None: return self.enriched_embeddings + if indices is not None and indices.ndimension() > 1: + raise RuntimeError("indices must be None, or 1-dimensional.") # Bind fields # shape: (num_entities, embedding_dim) From fd181e8967102ceacc7bf8ac6affcb5b8cd5e62d Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:15:58 +0100 Subject: [PATCH 585/690] Inline constants --- src/pykeen/nn/representation.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index a48afdb5a2..0e0e03dd8c 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -29,10 +29,7 @@ logger = logging.getLogger(__name__) -HEAD_DIM = 1 -RELATION_DIM = 2 -TAIL_DIM = 3 -DIMS = dict(h=HEAD_DIM, r=RELATION_DIM, t=TAIL_DIM) +DIMS = dict(h=1, r=2, t=3) def _normalize_dim(dim: Union[int, str]) -> int: From 1a9e40c9f458200a55f139335eb98dd60738611e Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:16:19 +0100 Subject: [PATCH 586/690] Rename constant --- src/pykeen/nn/modules.py | 8 ++++---- src/pykeen/nn/representation.py | 4 ++-- tests/test_nn.py | 4 ++-- 3 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 0f856a44bc..948aa0b1b5 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -11,7 +11,7 @@ from torch import FloatTensor, nn from . import functional as pkf -from .representation import DIMS, convert_to_canonical_shape +from .representation import CANONICAL_DIMENSIONS, convert_to_canonical_shape from ..typing import HeadRepresentation, RelationRepresentation, Representation, TailRepresentation from ..utils import upgrade_to_sequence @@ -240,17 +240,17 @@ def _forward_slicing_wrapper( scores = torch.cat([ self(h=h_batch, r=r, t=t) for h_batch in _get_batches(h, slice_size) - ], dim=DIMS[slice_dim]) + ], dim=CANONICAL_DIMENSIONS[slice_dim]) elif slice_dim == "r": scores = torch.cat([ self(h=h, r=r_batch, t=t) for r_batch in _get_batches(r, slice_size) - ], dim=DIMS[slice_dim]) + ], dim=CANONICAL_DIMENSIONS[slice_dim]) elif slice_dim == "t": scores = torch.cat([ self(h=h, r=r, t=t_batch) for t_batch in _get_batches(t, slice_size) - ], dim=DIMS[slice_dim]) + ], dim=CANONICAL_DIMENSIONS[slice_dim]) else: raise ValueError(f'Invalid slice_dim: {slice_dim}') return scores diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index 0e0e03dd8c..9b933270ea 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -29,14 +29,14 @@ logger = logging.getLogger(__name__) -DIMS = dict(h=1, r=2, t=3) +CANONICAL_DIMENSIONS = dict(h=1, r=2, t=3) def _normalize_dim(dim: Union[int, str]) -> int: """Normalize the dimension selection.""" if isinstance(dim, int): return dim - return DIMS[dim.lower()[0]] + return CANONICAL_DIMENSIONS[dim.lower()[0]] def get_expected_canonical_shape( diff --git a/tests/test_nn.py b/tests/test_nn.py index 33b901d895..34bded06b6 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -12,7 +12,7 @@ from torch.nn import functional from pykeen.nn import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule -from pykeen.nn.representation import DIMS, RGCNRepresentations, get_expected_canonical_shape +from pykeen.nn.representation import CANONICAL_DIMENSIONS, RGCNRepresentations, get_expected_canonical_shape from pykeen.nn.sim import kullback_leibler_similarity from pykeen.testing.base import GenericTests, TestsTest from pykeen.triples import TriplesFactory @@ -86,7 +86,7 @@ def test_forward_with_invalid_indices(self): def _test_in_canonical_shape(self, indices: Optional[torch.LongTensor]): """Test get_in_canonical_shape with the given indices.""" # test both, using the actual dimension, and its name - for dim in itertools.chain(DIMS.keys(), DIMS.values()): + for dim in itertools.chain(CANONICAL_DIMENSIONS.keys(), CANONICAL_DIMENSIONS.values()): # batch_size, d1, d2, d3, * x = self.instance.get_in_canonical_shape(dim=dim, indices=indices) From bc8bd6617ca76e640011139a5a6284fb46bdc7cb Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:16:36 +0100 Subject: [PATCH 587/690] Add upgrade_to_sequence to export of utils.py --- src/pykeen/utils.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 4304709048..55f666c90e 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -46,6 +46,7 @@ 'set_random_seed', 'split_complex', 'split_list_in_batches_iter', + 'upgrade_to_sequence', 'view_complex', 'NoRandomSeedNecessary', 'Result', From bc0c663e1ad5a2d15ef5892b15d51d017a098207 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:23:12 +0100 Subject: [PATCH 588/690] Use complex dtype for ComplEx interaction --- src/pykeen/nn/functional.py | 17 +++-------------- 1 file changed, 3 insertions(+), 14 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 1ce51f15cf..170e7226da 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -13,7 +13,7 @@ from ..typing import GaussianDistribution from ..utils import ( broadcast_cat, clamp_norm, extended_einsum, is_cudnn_error, negative_norm_of_sum, project_entity, - split_complex, tensor_product, tensor_sum, view_complex, + tensor_product, tensor_sum, view_complex, ) __all__ = [ @@ -104,19 +104,8 @@ def complex_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # Re() - # = - - - - # = + + - - (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] - return tensor_sum(*( - factor * tensor_product(hh, rr, tt).sum(dim=-1) - for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] - )) + h, r, t = [view_complex(x=x) for x in (h, r, t)] + return torch.real((h * r * torch.conj(t)).sum(dim=-1)) @_add_cuda_warning From a15d328d515aed3ea9fdb0665649a36d83061ad5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:25:01 +0100 Subject: [PATCH 589/690] use tensor product --- src/pykeen/nn/functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 170e7226da..8a04066087 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -105,7 +105,7 @@ def complex_interaction( The scores. """ h, r, t = [view_complex(x=x) for x in (h, r, t)] - return torch.real((h * r * torch.conj(t)).sum(dim=-1)) + return torch.real(tensor_product(h, r, torch.conj(t)).sum(dim=-1)) @_add_cuda_warning From 19c54c25e4cadd65d1a360aa67459eb1d90b9cc5 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:28:00 +0100 Subject: [PATCH 590/690] Extract common utility for elementwise tensor combinations --- src/pykeen/utils.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 55f666c90e..05a4066a62 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -502,16 +502,25 @@ def complex_normalize(x: torch.Tensor) -> torch.Tensor: return x +def _tensor_elementwise( + op: Callable[[torch.FloatTensor, torch.FloatTensor], torch.FloatTensor], + *x: torch.FloatTensor, +) -> torch.FloatTensor: + """Combine tensors elementwise in optimal order.""" + if len(x) < 2: + return x[0] + # TODO: Optimize order + return functools.reduce(op, x[1:], x[0]) + + def tensor_sum(*x: torch.FloatTensor) -> torch.FloatTensor: """Compute elementwise sum of tensors in brodcastable shape.""" - # TODO: Optimize order - return sum(x) + return _tensor_elementwise(operator.add, *x) def tensor_product(*x: torch.FloatTensor) -> torch.FloatTensor: """Compute elementwise product of tensors in broadcastable shape.""" - # TODO: Optimize order - return functools.reduce(operator.mul, x[1:], x[0]) + return _tensor_elementwise(operator.mul, *x) def negative_norm_of_sum( From 856b841482f99140d5b8b45c1a8aca9dca21b501 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:34:26 +0100 Subject: [PATCH 591/690] Add elementwise combination optimization --- src/pykeen/utils.py | 96 ++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 87 insertions(+), 9 deletions(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 05a4066a62..e17017ba11 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -4,6 +4,7 @@ import ftplib import functools +import itertools import json import logging import operator @@ -12,6 +13,7 @@ from io import BytesIO from typing import Any, Callable, Dict, Iterable, List, Mapping, Optional, Sequence, Tuple, Type, TypeVar, Union +import numpy import numpy as np import pandas as pd import torch @@ -502,25 +504,101 @@ def complex_normalize(x: torch.Tensor) -> torch.Tensor: return x -def _tensor_elementwise( +def calculate_broadcasted_elementwise_result_shape( + first: Tuple[int, ...], + second: Tuple[int, ...], +) -> Tuple[int, ...]: + """Determine the return shape of a broadcasted elementwise operation.""" + return tuple(max(a, b) for a, b in zip(first, second)) + + +def estimate_cost_of_sequence( + shape: Tuple[int, ...], + *other_shapes: Tuple[int, ...], +) -> int: + """Cost of a sequence of broadcasted element-wise operations of tensors, given their shapes.""" + return sum( + map( + numpy.prod, + itertools.islice( + itertools.accumulate( + other_shapes, + calculate_broadcasted_elementwise_result_shape, + initial=shape, + ), + 1, + None, + ) + ) + ) + + +@functools.lru_cache(maxsize=32) +def _get_optimal_sequence( + *sorted_shapes: Tuple[int, ...], +) -> Tuple[int, Tuple[int, ...]]: + """Find the optimal sequence in which to combine tensors element-wise based on the shapes. + The shapes should be sorted to enable efficient caching. + :param sorted_shapes: + The shapes of the tensors to combine. + :return: + The optimal execution order (as indices), and the cost. + """ + return min( + (estimate_cost_of_sequence(*(sorted_shapes[i] for i in p)), p) + for p in itertools.permutations(list(range(len(sorted_shapes)))) + ) + + +def get_optimal_sequence(*shapes: Tuple[int, ...]) -> Tuple[int, Tuple[int, ...]]: + """Find the optimal sequence in which to combine tensors elementwise based on the shapes. + :param shapes: + The shapes of the tensors to combine. + :return: + The optimal execution order (as indices), and the cost. + """ + # create sorted list of shapes to allow utilization of lru cache (optimal execution order does not depend on the + # input sorting, as the order is determined by re-ordering the sequence anyway) + arg_sort = sorted(range(len(shapes)), key=shapes.__getitem__) + + # Determine optimal order and cost + cost, optimal_order = _get_optimal_sequence(*(shapes[new_index] for new_index in arg_sort)) + + # translate back to original order + optimal_order = tuple(arg_sort[i] for i in optimal_order) + + return cost, optimal_order + + +def _multi_combine( + tensors: Tuple[torch.FloatTensor, ...], op: Callable[[torch.FloatTensor, torch.FloatTensor], torch.FloatTensor], - *x: torch.FloatTensor, ) -> torch.FloatTensor: - """Combine tensors elementwise in optimal order.""" - if len(x) < 2: - return x[0] - # TODO: Optimize order - return functools.reduce(op, x[1:], x[0]) + """Broadcasted element-wise combination of tensors. + The optimal execution plan gets cached so that the optimization is only performed once for a fixed set of shapes. + + :param tensors: + The tensors, in broadcastable shape. + :param op: + The elementwise operator. + + :return: + The elementwise combination evaluated in optimal processing order. + """ + # determine optimal processing order + order = get_optimal_sequence(*(t.shape for t in tensors))[1] + tensors = [tensors[i] for i in order] + return functools.reduce(op, tensors[1:], tensors[0]) def tensor_sum(*x: torch.FloatTensor) -> torch.FloatTensor: """Compute elementwise sum of tensors in brodcastable shape.""" - return _tensor_elementwise(operator.add, *x) + return _multi_combine(tensors=x, op=operator.add) def tensor_product(*x: torch.FloatTensor) -> torch.FloatTensor: """Compute elementwise product of tensors in broadcastable shape.""" - return _tensor_elementwise(operator.mul, *x) + return _multi_combine(tensors=x, op=operator.mul) def negative_norm_of_sum( From 74870b541486a6537e3ca4582425a8ea4b6e3650 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:38:25 +0100 Subject: [PATCH 592/690] Add tests for broadcasted combination --- tests/test_utils.py | 107 ++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 104 insertions(+), 3 deletions(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index 876ce365e9..327993834a 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,15 +1,21 @@ # -*- coding: utf-8 -*- """Tests for the :mod:`pykeen.utils` module.""" - +import functools +import itertools +import operator +import random import string +import timeit import unittest +from typing import Iterable, Tuple +import numpy import torch from pykeen.utils import ( - _CUDA_OOM_ERROR, _CUDNN_ERROR, clamp_norm, combine_complex, compact_mapping, flatten_dictionary, - get_until_first_blank, is_cuda_oom_error, is_cudnn_error, project_entity, split_complex, + _CUDA_OOM_ERROR, _CUDNN_ERROR, calculate_broadcasted_elementwise_result_shape, clamp_norm, combine_complex, compact_mapping, estimate_cost_of_sequence, flatten_dictionary, + get_optimal_sequence, get_until_first_blank, is_cuda_oom_error, is_cudnn_error, project_entity, split_complex, tensor_product, tensor_sum, ) @@ -197,3 +203,98 @@ def test_is_cudnn_error(self): self.assertFalse(is_cuda_oom_error(runtime_error=error)) self.assertFalse(is_cudnn_error(runtime_error=self.not_cuda_error)) + + +def test_calculate_broadcasted_elementwise_result_shape(): + """Test calculate_broadcasted_elementwise_result_shape.""" + max_dim = 64 + for n_dim, i in itertools.product(range(2, 5), range(10)): + a_shape = [1 for _ in range(n_dim)] + b_shape = [1 for _ in range(n_dim)] + for j in range(n_dim): + dim = 2 + random.randrange(max_dim) + mod = random.randrange(3) + if mod % 2 == 0: + a_shape[j] = dim + if mod > 0: + b_shape[j] = dim + a = torch.empty(*a_shape) + b = torch.empty(*b_shape) + shape = calculate_broadcasted_elementwise_result_shape(first=a.shape, second=b.shape) + c = a + b + exp_shape = c.shape + assert shape == exp_shape + + +def _generate_shapes( + n_dim: int = 5, + n_terms: int = 4, + iterations: int = 64, +) -> Iterable[Tuple[Tuple[int, ...], ...]]: + """Generate shapes.""" + max_shape = torch.randint(low=2, high=32, size=(128,)) + for _i in range(iterations): + # create broadcastable shapes + idx = torch.randperm(max_shape.shape[0])[:n_dim] + this_max_shape = max_shape[idx] + this_min_shape = torch.ones_like(this_max_shape) + shapes = [] + for _j in range(n_terms): + mask = this_min_shape + while not (1 < mask.sum() < n_dim): + mask = torch.as_tensor(torch.rand(size=(n_dim,)) < 0.3, dtype=max_shape.dtype) + this_array_shape = this_max_shape * mask + this_min_shape * (1 - mask) + shapes.append(tuple(this_array_shape.tolist())) + yield tuple(shapes) + + +def test_estimate_cost_of_add_sequence(): + """Test ``estimate_cost_of_add_sequence()``.""" + # create random array, estimate the costs of addition, and measure some execution times. + # then, compute correlation between the estimated cost, and the measured time. + data = [] + for shapes in _generate_shapes(): + arrays = [torch.empty(*shape) for shape in shapes] + cost = estimate_cost_of_sequence(*(a.shape for a in arrays)) + consumption = timeit.timeit(stmt='sum(arrays)', globals=locals(), number=25) + data.append((cost, consumption)) + a = numpy.asarray(data) + + # check for strong correlation between estimated costs and measured execution time + assert (numpy.corrcoef(x=a[:, 0], y=a[:, 1])[0, 1]) > 0.8 + + +def test_get_optimal_add_sequence(): + """Test ``get_optimal_add_sequence()``.""" + for shapes in _generate_shapes(): + # get optimal sequence + opt_cost, opt_seq = get_optimal_sequence(*shapes) + + # check correct cost + exp_opt_cost = estimate_cost_of_sequence(*(shapes[i] for i in opt_seq)) + assert exp_opt_cost == opt_cost + + # check optimality + for perm in itertools.permutations(list(range(len(shapes)))): + cost = estimate_cost_of_sequence(*(shapes[i] for i in perm)) + assert cost >= opt_cost + + +def test_tensor_sum(): + """Test tensor_sum.""" + for shapes in _generate_shapes(): + tensors = [torch.rand(*shape) for shape in shapes] + result = tensor_sum(*tensors) + + # compare result to sequential addition + assert torch.allclose(result, sum(tensors)) + + +def test_tensor_product(): + """Test tensor_product.""" + for shapes in _generate_shapes(): + tensors = [torch.rand(*shape) for shape in shapes] + result = tensor_product(*tensors) + + # compare result to sequential addition + assert torch.allclose(result, functools.reduce(operator.mul, tensors[1:], tensors[0])) From a97c84b1b2ad36ab575e5d4be1ea018fbf6e0fd4 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:39:23 +0100 Subject: [PATCH 593/690] Add mark.slow --- tests/test_utils.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index 327993834a..c384d4c90c 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -11,11 +11,12 @@ from typing import Iterable, Tuple import numpy +import pytest import torch from pykeen.utils import ( _CUDA_OOM_ERROR, _CUDNN_ERROR, calculate_broadcasted_elementwise_result_shape, clamp_norm, combine_complex, compact_mapping, estimate_cost_of_sequence, flatten_dictionary, - get_optimal_sequence, get_until_first_blank, is_cuda_oom_error, is_cudnn_error, project_entity, split_complex, tensor_product, tensor_sum, + get_optimal_sequence, get_until_first_blank, is_cuda_oom_error, is_cudnn_error, project_entity, set_random_seed, split_complex, tensor_product, tensor_sum, ) @@ -248,8 +249,10 @@ def _generate_shapes( yield tuple(shapes) +@pytest.mark.slow def test_estimate_cost_of_add_sequence(): """Test ``estimate_cost_of_add_sequence()``.""" + set_random_seed(seed=42) # create random array, estimate the costs of addition, and measure some execution times. # then, compute correlation between the estimated cost, and the measured time. data = [] From 15b368c7683e504afb24ac6d92d89cba07911d7b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 12:44:10 +0100 Subject: [PATCH 594/690] Add two implementation variants of complex to functional --- src/pykeen/nn/functional.py | 33 ++++++++++++++++++++++++++++++--- 1 file changed, 30 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 8a04066087..3574a93626 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -13,7 +13,7 @@ from ..typing import GaussianDistribution from ..utils import ( broadcast_cat, clamp_norm, extended_einsum, is_cudnn_error, negative_norm_of_sum, project_entity, - tensor_product, tensor_sum, view_complex, + split_complex, tensor_product, tensor_sum, view_complex, ) __all__ = [ @@ -83,6 +83,34 @@ def wrapped(*args, **kwargs): return wrapped +def _complex_interaction_complex_native( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Use torch built-ins for computation with complex numbers.""" + h, r, t = [view_complex(x=x) for x in (h, r, t)] + return torch.real(tensor_product(h, r, torch.conj(t)).sum(dim=-1)) + + +def _complex_interaction_optimized_broadcasted( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Manually split into real/imag, and used optimized broadcasted combination.""" + (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] + return tensor_sum(*( + factor * tensor_product(hh, rr, tt).sum(dim=-1) + for factor, hh, rr, tt in [ + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] + )) + + def complex_interaction( h: torch.FloatTensor, r: torch.FloatTensor, @@ -104,8 +132,7 @@ def complex_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - h, r, t = [view_complex(x=x) for x in (h, r, t)] - return torch.real(tensor_product(h, r, torch.conj(t)).sum(dim=-1)) + return _complex_interaction_complex_native(h, r, t) @_add_cuda_warning From 6428ab5b8385f18d0a78f5f8e4f749d494af683a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 13:50:01 +0100 Subject: [PATCH 595/690] Add another complex variant --- src/pykeen/nn/functional.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 3574a93626..836c319f97 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -111,6 +111,19 @@ def _complex_interaction_optimized_broadcasted( )) +def _complex_interaction_direct( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Manually split into real/imag, and directly evaluate interaction.""" + (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] + return (h_re * r_re * t_re).sum(dim=-1) + \ + (h_re * r_im * t_im).sum(dim=-1) + \ + (h_im * r_re * t_im).sum(dim=-1) - \ + (h_im * r_im * t_re).sum(dim=-1) + + def complex_interaction( h: torch.FloatTensor, r: torch.FloatTensor, From 06627257a3ef66c6c909faf74d660e6c1ba3bec7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 13:54:35 +0100 Subject: [PATCH 596/690] Add cost estimation to rotate forward --- src/pykeen/nn/functional.py | 35 ++++++++++++++++++----------------- 1 file changed, 18 insertions(+), 17 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 836c319f97..a0f0183998 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -12,7 +12,7 @@ from .sim import KG2E_SIMILARITIES from ..typing import GaussianDistribution from ..utils import ( - broadcast_cat, clamp_norm, extended_einsum, is_cudnn_error, negative_norm_of_sum, project_entity, + broadcast_cat, clamp_norm, estimate_cost_of_sequence, extended_einsum, is_cudnn_error, negative_norm_of_sum, project_entity, split_complex, tensor_product, tensor_sum, view_complex, ) @@ -103,11 +103,11 @@ def _complex_interaction_optimized_broadcasted( return tensor_sum(*( factor * tensor_product(hh, rr, tt).sum(dim=-1) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) @@ -591,21 +591,22 @@ def rotate_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # TODO: rotate by inverse relation, if the cost(dist(h, t * r_inv)) < cost(dist(h * r, t)) - # # r expresses a rotation in complex plane. - # # The inverse rotation is expressed by the complex conjugate of r. - # # The score is computed as the distance of the relation-rotated head to the tail. - # # Equivalently, we can rotate the tail by the inverse relation, and measure the distance to the head, i.e. - # # |h * r - t| = |h - conj(r) * t| - # r_inv = torch.stack([r[:, :, :, 0], -r[:, :, :, 1]], dim=-1) + # r expresses a rotation in complex plane. h, r, t = [view_complex(x) for x in (h, r, t)] - - # Rotate (=Hadamard product in complex space). - hr = h * r + if estimate_cost_of_sequence(h.shape, r.shape) < estimate_cost_of_sequence(r.shape, t.shape): + # rotate head by relation (=Hadamard product in complex space) + h = h * r + else: + # rotate tail by inverse of relation + # The inverse rotation is expressed by the complex conjugate of r. + # The score is computed as the distance of the relation-rotated head to the tail. + # Equivalently, we can rotate the tail by the inverse relation, and measure the distance to the head, i.e. + # |h * r - t| = |h - conj(r) * t| + t = t * torch.conj(r) # Workaround until https://github.com/pytorch/pytorch/issues/30704 is fixed return negative_norm_of_sum( - hr, + h, -t, p=2, power_norm=False, From f2cbd0a11621883467430ae6cd46d49eda450721 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 14:18:37 +0100 Subject: [PATCH 597/690] Add draft of functional interaction function benchmark --- benchmarking/interactions.py | 107 +++++++++++++++++++++++++++++++++++ 1 file changed, 107 insertions(+) create mode 100644 benchmarking/interactions.py diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py new file mode 100644 index 0000000000..230304a4d5 --- /dev/null +++ b/benchmarking/interactions.py @@ -0,0 +1,107 @@ +import timeit +from typing import Mapping, Optional, Tuple + +import pandas +import torch +import tqdm + +from pykeen.nn import Interaction +from pykeen.nn.functional import _complex_interaction_complex_native, _complex_interaction_direct, _complex_interaction_optimized_broadcasted +from pykeen.nn.modules import ComplExInteraction, _unpack_singletons +from pykeen.typing import HeadRepresentation, RelationRepresentation, TailRepresentation + + +def _use_case_to_shape(use_case: str, b: int, n: int) -> Tuple[ + Tuple[int, int], + Tuple[int, int], + Tuple[int, int], +]: + if use_case == "hrt": + return (b, 1), (b, 1), (b, 1) + elif use_case == "t": + return (b, 1), (b, 1), (1, n) + elif use_case == "h": + return (1, n), (b, 1), (b, 1) + else: + raise ValueError + + +def _generate_hrt( + prefix_shapes: Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]], + interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], + dim: int, + additional_dims: Optional[Mapping[str, int]] = None, +) -> Tuple[HeadRepresentation, RelationRepresentation, TailRepresentation]: + additional_dims = additional_dims or dict() + additional_dims.setdefault("d", dim) + return _unpack_singletons(*( + torch.rand(*prefix_shape, *(additional_dims[s] for s in suffix_shape), requires_grad=True) + for prefix_shape, suffix_shape in zip( + prefix_shapes, + ( + interaction.entity_shape, + interaction.relation_shape, + interaction.tail_entity_shape or interaction.entity_shape + ) + ) + )) + + +def main(): + base_interaction: Interaction = ComplExInteraction() + variants = [ + _complex_interaction_complex_native, + _complex_interaction_optimized_broadcasted, + _complex_interaction_direct, + ] + use_case_labels = ["hrt", "t", "h"] + batch_sizes = [2 ** i for i in range(5, 10 + 1)] + # batch_sizes = [2 ** i for i in range(5, 7)] + num_entities = (100, 15_000) + # num_entities = (100,) + vector_dimensions = [2 ** i for i in range(5, 10 + 1)] + # vector_dimensions = [2 ** i for i in range(5, 7)] + data = [] + tasks = [ + (b, n, d, ul, _use_case_to_shape(use_case=ul, b=b, n=n)) + for ul in use_case_labels + for b in batch_sizes + for n in num_entities + for d in vector_dimensions + ] + for variant in tqdm.tqdm(variants, unit="variant", unit_scale=True): + # create variant + base_interaction.func = variant + for (b, n, d, ul, prefix_shapes) in tqdm.tqdm(tasks, unit="task", unit_scale=True): + h, r, t = _generate_hrt(prefix_shapes=prefix_shapes, interaction=base_interaction, dim=d) + try: + timer = timeit.Timer( + stmt="interaction(h=h, r=r, t=t)", + globals=dict(interaction=base_interaction, h=h, r=r, t=t), + ) + n_samples, total_time = timer.autorange() + time_per_sample = total_time / n_samples + except Exception as error: + n_samples, total_time, time_per_sample = 0, float('nan'), float('nan') + data.append((b, n, d, prefix_shapes, ul, variant.__name__, total_time, n_samples, time_per_sample)) + df = pandas.DataFrame(data=data, columns=[ + "batch_size", + "num_entities", + "dimension", + "prefix_shapes", + "use_case", + "variant", + "total_time", + "n_samples", + "time_per_sample", + ]) + df.to_csv("measurement.tsv", sep="\t", index=False) + print( + df.groupby( + by=["batch_size", "num_entities", "dimension", "use_case", "variant"] + ).agg({"time_per_sample": "mean"}).unstack() + ) + + +if __name__ == '__main__': + main() From 3f6a1e3d30828e089bc9dd6b6e9e5d1e0a9cff37 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 14:28:08 +0100 Subject: [PATCH 598/690] change benchmark script --- benchmarking/interactions.py | 33 ++++++++++++++++++++++----------- 1 file changed, 22 insertions(+), 11 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index 230304a4d5..e1c509355f 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -1,6 +1,7 @@ import timeit from typing import Mapping, Optional, Tuple +import numpy import pandas import torch import tqdm @@ -47,6 +48,10 @@ def _generate_hrt( )) +def _get_result_shape(prefix_shapes) -> Tuple[int, int, int, int]: + return (max(s[0] for s in prefix_shapes),) + tuple([s[1] for s in prefix_shapes]) + + def main(): base_interaction: Interaction = ComplExInteraction() variants = [ @@ -59,6 +64,7 @@ def main(): # batch_sizes = [2 ** i for i in range(5, 7)] num_entities = (100, 15_000) # num_entities = (100,) + max_result_elements = 2 ** 30 vector_dimensions = [2 ** i for i in range(5, 10 + 1)] # vector_dimensions = [2 ** i for i in range(5, 7)] data = [] @@ -69,20 +75,25 @@ def main(): for n in num_entities for d in vector_dimensions ] - for variant in tqdm.tqdm(variants, unit="variant", unit_scale=True): + progress = tqdm.tqdm(variants, unit="variant", unit_scale=True) + for variant in progress: # create variant base_interaction.func = variant for (b, n, d, ul, prefix_shapes) in tqdm.tqdm(tasks, unit="task", unit_scale=True): - h, r, t = _generate_hrt(prefix_shapes=prefix_shapes, interaction=base_interaction, dim=d) - try: - timer = timeit.Timer( - stmt="interaction(h=h, r=r, t=t)", - globals=dict(interaction=base_interaction, h=h, r=r, t=t), - ) - n_samples, total_time = timer.autorange() - time_per_sample = total_time / n_samples - except Exception as error: - n_samples, total_time, time_per_sample = 0, float('nan'), float('nan') + result_shape = _get_result_shape(prefix_shapes) + n_samples, total_time, time_per_sample = 0, float('nan'), float('nan') + if max_result_elements is None or numpy.prod(result_shape) < max_result_elements: + h, r, t = _generate_hrt(prefix_shapes=prefix_shapes, interaction=base_interaction, dim=d) + try: + # TODO: cuda sync + timer = timeit.Timer( + stmt="interaction(h=h, r=r, t=t)", + globals=dict(interaction=base_interaction, h=h, r=r, t=t), + ) + n_samples, total_time = timer.autorange() + time_per_sample = total_time / n_samples + except Exception as error: + progress.write(f"ERROR: {error}") data.append((b, n, d, prefix_shapes, ul, variant.__name__, total_time, n_samples, time_per_sample)) df = pandas.DataFrame(data=data, columns=[ "batch_size", From d924126e9fd0cb34d16b7053b401a01c05dfd389 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 14:46:27 +0100 Subject: [PATCH 599/690] Fix variant creation --- benchmarking/interactions.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index e1c509355f..7a0efed6f8 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -78,7 +78,7 @@ def main(): progress = tqdm.tqdm(variants, unit="variant", unit_scale=True) for variant in progress: # create variant - base_interaction.func = variant + base_interaction.__class__.func = variant for (b, n, d, ul, prefix_shapes) in tqdm.tqdm(tasks, unit="task", unit_scale=True): result_shape = _get_result_shape(prefix_shapes) n_samples, total_time, time_per_sample = 0, float('nan'), float('nan') @@ -94,6 +94,7 @@ def main(): time_per_sample = total_time / n_samples except Exception as error: progress.write(f"ERROR: {error}") + progress.set_postfix(dict(shape=prefix_shapes, time=time_per_sample)) data.append((b, n, d, prefix_shapes, ul, variant.__name__, total_time, n_samples, time_per_sample)) df = pandas.DataFrame(data=data, columns=[ "batch_size", From bf819faf075cd632244419878b5319b45c818cd7 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 24 Nov 2020 14:52:29 +0100 Subject: [PATCH 600/690] Store results better --- benchmarking/interactions.py | 19 ++++++++++++------- 1 file changed, 12 insertions(+), 7 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index 7a0efed6f8..dd9bb403ed 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -7,9 +7,13 @@ import tqdm from pykeen.nn import Interaction -from pykeen.nn.functional import _complex_interaction_complex_native, _complex_interaction_direct, _complex_interaction_optimized_broadcasted +from pykeen.nn.functional import ( + _complex_interaction_complex_native, _complex_interaction_direct, + _complex_interaction_optimized_broadcasted, +) from pykeen.nn.modules import ComplExInteraction, _unpack_singletons from pykeen.typing import HeadRepresentation, RelationRepresentation, TailRepresentation +from pykeen.version import get_git_hash def _use_case_to_shape(use_case: str, b: int, n: int) -> Tuple[ @@ -107,12 +111,13 @@ def main(): "n_samples", "time_per_sample", ]) - df.to_csv("measurement.tsv", sep="\t", index=False) - print( - df.groupby( - by=["batch_size", "num_entities", "dimension", "use_case", "variant"] - ).agg({"time_per_sample": "mean"}).unstack() - ) + git_hash = get_git_hash() + df.to_csv(f"{git_hash}_measurement.tsv", sep="\t", index=False) + df_agg = df.groupby( + by=["batch_size", "num_entities", "dimension", "use_case", "variant"] + ).agg({"time_per_sample": "mean"}).unstack() + df_agg.to_csv(f"{git_hash}_measurement_agg.tsv", sep="\t", index=False) + print(df_agg) if __name__ == '__main__': From ce69e86a909f66bf2a29e8dc35e8c6f26f5c561e Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 24 Nov 2020 14:54:14 +0100 Subject: [PATCH 601/690] Cleaner generation of stateless interaction class --- benchmarking/interactions.py | 7 +++---- src/pykeen/nn/modules.py | 14 +++++++++++++- 2 files changed, 16 insertions(+), 5 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index dd9bb403ed..453a34c2b4 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -57,7 +57,6 @@ def _get_result_shape(prefix_shapes) -> Tuple[int, int, int, int]: def main(): - base_interaction: Interaction = ComplExInteraction() variants = [ _complex_interaction_complex_native, _complex_interaction_optimized_broadcasted, @@ -82,17 +81,17 @@ def main(): progress = tqdm.tqdm(variants, unit="variant", unit_scale=True) for variant in progress: # create variant - base_interaction.__class__.func = variant + interaction = Interaction.from_func(variant) for (b, n, d, ul, prefix_shapes) in tqdm.tqdm(tasks, unit="task", unit_scale=True): result_shape = _get_result_shape(prefix_shapes) n_samples, total_time, time_per_sample = 0, float('nan'), float('nan') if max_result_elements is None or numpy.prod(result_shape) < max_result_elements: - h, r, t = _generate_hrt(prefix_shapes=prefix_shapes, interaction=base_interaction, dim=d) + h, r, t = _generate_hrt(prefix_shapes=prefix_shapes, interaction=interaction, dim=d) try: # TODO: cuda sync timer = timeit.Timer( stmt="interaction(h=h, r=r, t=t)", - globals=dict(interaction=base_interaction, h=h, r=r, t=t), + globals=dict(interaction=interaction, h=h, r=r, t=t), ) n_samples, total_time = timer.autorange() time_per_sample = total_time / n_samples diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 948aa0b1b5..7a2079cd81 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -5,7 +5,7 @@ import logging import math from abc import ABC -from typing import Any, Callable, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Union +from typing import Any, Callable, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, Union import torch from torch import FloatTensor, nn @@ -85,6 +85,18 @@ class Interaction(nn.Module, Generic[HeadRepresentation, RelationRepresentation, #: The functional interaction form func: Callable[..., torch.FloatTensor] + @classmethod + def from_func(cls, f) -> 'Interaction': + """Create an instance of a stateless interaction class.""" + return cls.cls_from_func(f)() + + @classmethod + def cls_from_func(cls, f) -> Type['Interaction']: + """Create a stateless interaction class.""" + class StatelessInteraction(cls): + func = f + return StatelessInteraction + @staticmethod def _prepare_hrt_for_functional( h: HeadRepresentation, From a3100a90648031f342f5a2269cb71b1727d52df9 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 24 Nov 2020 15:27:04 +0100 Subject: [PATCH 602/690] Update interaction benchmark --- benchmarking/interactions.py | 57 +++++++++++++++++++++++------------- 1 file changed, 36 insertions(+), 21 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index 453a34c2b4..7f14601f78 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -1,17 +1,20 @@ import timeit from typing import Mapping, Optional, Tuple +import click +import matplotlib.pyplot as plt import numpy import pandas +import seaborn as sns import torch -import tqdm +from tqdm import tqdm from pykeen.nn import Interaction from pykeen.nn.functional import ( _complex_interaction_complex_native, _complex_interaction_direct, _complex_interaction_optimized_broadcasted, ) -from pykeen.nn.modules import ComplExInteraction, _unpack_singletons +from pykeen.nn.modules import _unpack_singletons from pykeen.typing import HeadRepresentation, RelationRepresentation, TailRepresentation from pykeen.version import get_git_hash @@ -56,7 +59,9 @@ def _get_result_shape(prefix_shapes) -> Tuple[int, int, int, int]: return (max(s[0] for s in prefix_shapes),) + tuple([s[1] for s in prefix_shapes]) -def main(): +@click.command() +@click.option('-m', '--max-result-elements-power', type=int, default=30, show_default=True) +def main(max_result_elements_power: int): variants = [ _complex_interaction_complex_native, _complex_interaction_optimized_broadcasted, @@ -67,7 +72,7 @@ def main(): # batch_sizes = [2 ** i for i in range(5, 7)] num_entities = (100, 15_000) # num_entities = (100,) - max_result_elements = 2 ** 30 + max_result_elements = 2 ** max_result_elements_power vector_dimensions = [2 ** i for i in range(5, 10 + 1)] # vector_dimensions = [2 ** i for i in range(5, 7)] data = [] @@ -78,28 +83,32 @@ def main(): for n in num_entities for d in vector_dimensions ] - progress = tqdm.tqdm(variants, unit="variant", unit_scale=True) + progress = tqdm(variants, unit="variant") for variant in progress: # create variant interaction = Interaction.from_func(variant) - for (b, n, d, ul, prefix_shapes) in tqdm.tqdm(tasks, unit="task", unit_scale=True): + for i, (b, n, d, ul, prefix_shapes) in enumerate(tqdm(tasks, unit="task"), start=1): result_shape = _get_result_shape(prefix_shapes) n_samples, total_time, time_per_sample = 0, float('nan'), float('nan') - if max_result_elements is None or numpy.prod(result_shape) < max_result_elements: - h, r, t = _generate_hrt(prefix_shapes=prefix_shapes, interaction=interaction, dim=d) - try: - # TODO: cuda sync - timer = timeit.Timer( - stmt="interaction(h=h, r=r, t=t)", - globals=dict(interaction=interaction, h=h, r=r, t=t), - ) - n_samples, total_time = timer.autorange() - time_per_sample = total_time / n_samples - except Exception as error: - progress.write(f"ERROR: {error}") + if max_result_elements is not None and max_result_elements < numpy.prod(result_shape): + continue + h, r, t = _generate_hrt(prefix_shapes=prefix_shapes, interaction=interaction, dim=d) + try: + # TODO: cuda sync + timer = timeit.Timer( + stmt="interaction(h=h, r=r, t=t)", + globals=dict(interaction=interaction, h=h, r=r, t=t), + ) + n_samples, total_time = timer.autorange() + time_per_sample = total_time / n_samples + except Exception as error: + progress.write(f"ERROR: {error}") progress.set_postfix(dict(shape=prefix_shapes, time=time_per_sample)) - data.append((b, n, d, prefix_shapes, ul, variant.__name__, total_time, n_samples, time_per_sample)) + data.append((i, b, n, d, prefix_shapes, ul, variant.__name__, total_time, n_samples, time_per_sample)) + + git_hash = get_git_hash() df = pandas.DataFrame(data=data, columns=[ + "experiment_number", "batch_size", "num_entities", "dimension", @@ -110,14 +119,20 @@ def main(): "n_samples", "time_per_sample", ]) - git_hash = get_git_hash() df.to_csv(f"{git_hash}_measurement.tsv", sep="\t", index=False) + df_agg = df.groupby( by=["batch_size", "num_entities", "dimension", "use_case", "variant"] - ).agg({"time_per_sample": "mean"}).unstack() + ).agg({"time_per_sample": "mean"}).unstack().reset_index().dropna() df_agg.to_csv(f"{git_hash}_measurement_agg.tsv", sep="\t", index=False) print(df_agg) + viz_df = df[df['total_time'].notna()] + viz_df['variant'] = viz_df['variant'].map(lambda s: ' '.join(s.split('_')[3:]).capitalize()) + plt.figure(figsize=(14, 6)) + sns.lineplot(data=viz_df, x='experiment_number', y='total_time', hue='variant') + plt.savefig(f"{git_hash}_measurement.png", dpi=300) + if __name__ == '__main__': main() From 86875ab9c2228a72086d4cb825d2dc219302c465 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 16:01:55 +0100 Subject: [PATCH 603/690] Add more options --- benchmarking/interactions.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index 7f14601f78..3fdd31cd7f 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -61,20 +61,24 @@ def _get_result_shape(prefix_shapes) -> Tuple[int, int, int, int]: @click.command() @click.option('-m', '--max-result-elements-power', type=int, default=30, show_default=True) -def main(max_result_elements_power: int): +@click.option('-b', '--max-batch-size-power', type=int, default=10, show_default=True) +@click.option('-d', '--max-vector-dimension-power', type=int, default=10, show_default=True) +def main( + max_result_elements_power: int, + max_batch_size_power: int, + max_vector_dimension_power: int, +): variants = [ _complex_interaction_complex_native, _complex_interaction_optimized_broadcasted, _complex_interaction_direct, ] use_case_labels = ["hrt", "t", "h"] - batch_sizes = [2 ** i for i in range(5, 10 + 1)] - # batch_sizes = [2 ** i for i in range(5, 7)] + batch_sizes = [2 ** i for i in range(5, max_batch_size_power + 1)] num_entities = (100, 15_000) # num_entities = (100,) max_result_elements = 2 ** max_result_elements_power - vector_dimensions = [2 ** i for i in range(5, 10 + 1)] - # vector_dimensions = [2 ** i for i in range(5, 7)] + vector_dimensions = [2 ** i for i in range(5, max_vector_dimension_power + 1)] data = [] tasks = [ (b, n, d, ul, _use_case_to_shape(use_case=ul, b=b, n=n)) From b8d7b2b56ec1daf139ab095b31c32dfa5875b9f8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 24 Nov 2020 16:14:33 +0100 Subject: [PATCH 604/690] Add gpu support to kernel benchmark --- benchmarking/interactions.py | 18 +++++++++++++----- 1 file changed, 13 insertions(+), 5 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index 3fdd31cd7f..7927e93b23 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -1,4 +1,3 @@ -import timeit from typing import Mapping, Optional, Tuple import click @@ -7,6 +6,7 @@ import pandas import seaborn as sns import torch +from torch.utils.benchmark import Timer from tqdm import tqdm from pykeen.nn import Interaction @@ -38,12 +38,13 @@ def _generate_hrt( prefix_shapes: Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]], interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], dim: int, + device: torch.device, additional_dims: Optional[Mapping[str, int]] = None, ) -> Tuple[HeadRepresentation, RelationRepresentation, TailRepresentation]: additional_dims = additional_dims or dict() additional_dims.setdefault("d", dim) return _unpack_singletons(*( - torch.rand(*prefix_shape, *(additional_dims[s] for s in suffix_shape), requires_grad=True) + torch.rand(*prefix_shape, *(additional_dims[s] for s in suffix_shape), requires_grad=True, device=device) for prefix_shape, suffix_shape in zip( prefix_shapes, ( @@ -68,6 +69,8 @@ def main( max_batch_size_power: int, max_vector_dimension_power: int, ): + device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") + print(f"Running on {device}.") variants = [ _complex_interaction_complex_native, _complex_interaction_optimized_broadcasted, @@ -96,10 +99,14 @@ def main( n_samples, total_time, time_per_sample = 0, float('nan'), float('nan') if max_result_elements is not None and max_result_elements < numpy.prod(result_shape): continue - h, r, t = _generate_hrt(prefix_shapes=prefix_shapes, interaction=interaction, dim=d) + h, r, t = _generate_hrt( + prefix_shapes=prefix_shapes, + interaction=interaction, + dim=d, + device=device, + ) try: - # TODO: cuda sync - timer = timeit.Timer( + timer = Timer( stmt="interaction(h=h, r=r, t=t)", globals=dict(interaction=interaction, h=h, r=r, t=t), ) @@ -123,6 +130,7 @@ def main( "n_samples", "time_per_sample", ]) + df["device"] = device.type df.to_csv(f"{git_hash}_measurement.tsv", sep="\t", index=False) df_agg = df.groupby( From 7290163f3dde6bb068533edfdf8277c9ca24688f Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Wed, 25 Nov 2020 15:52:12 +0100 Subject: [PATCH 605/690] Pass flake8 --- src/pykeen/models/unimodal/rgcn.py | 5 ++++- src/pykeen/nn/functional.py | 24 +++++++++++++----------- src/pykeen/nn/representation.py | 4 ++-- src/pykeen/utils.py | 29 +++++++++++++++-------------- tests/test_models.py | 5 ++++- tests/test_utils.py | 10 ++++++---- 6 files changed, 44 insertions(+), 33 deletions(-) diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index 1e9ed5229a..d71a6963c7 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -12,7 +12,10 @@ from ...losses import Loss from ...nn import EmbeddingSpecification, Interaction from ...nn.modules import DistMultInteraction -from ...nn.representation import RGCNRepresentations, inverse_indegree_edge_weights, inverse_outdegree_edge_weights, symmetric_edge_weights +from ...nn.representation import ( + RGCNRepresentations, inverse_indegree_edge_weights, inverse_outdegree_edge_weights, + symmetric_edge_weights, +) from ...triples import TriplesFactory from ...typing import DeviceHint diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index a0f0183998..275f98c010 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -12,8 +12,8 @@ from .sim import KG2E_SIMILARITIES from ..typing import GaussianDistribution from ..utils import ( - broadcast_cat, clamp_norm, estimate_cost_of_sequence, extended_einsum, is_cudnn_error, negative_norm_of_sum, project_entity, - split_complex, tensor_product, tensor_sum, view_complex, + broadcast_cat, clamp_norm, estimate_cost_of_sequence, extended_einsum, is_cudnn_error, negative_norm_of_sum, + project_entity, split_complex, tensor_product, tensor_sum, view_complex, ) __all__ = [ @@ -103,11 +103,11 @@ def _complex_interaction_optimized_broadcasted( return tensor_sum(*( factor * tensor_product(hh, rr, tt).sum(dim=-1) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) @@ -118,10 +118,12 @@ def _complex_interaction_direct( ) -> torch.FloatTensor: """Manually split into real/imag, and directly evaluate interaction.""" (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] - return (h_re * r_re * t_re).sum(dim=-1) + \ - (h_re * r_im * t_im).sum(dim=-1) + \ - (h_im * r_re * t_im).sum(dim=-1) - \ - (h_im * r_im * t_re).sum(dim=-1) + return ( + (h_re * r_re * t_re).sum(dim=-1) + + (h_re * r_im * t_im).sum(dim=-1) + + (h_im * r_re * t_im).sum(dim=-1) + - (h_im * r_im * t_re).sum(dim=-1) + ) def complex_interaction( diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index 9b933270ea..dd113d49ce 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -623,7 +623,7 @@ def _get_relation_weights(self, i_layer: int, r: int) -> torch.FloatTensor: def forward( self, indices: Optional[torch.LongTensor] = None, - ) -> torch.FloatTensor: + ) -> torch.FloatTensor: # noqa:D102 # use buffered messages if applicable if indices is None and self.enriched_embeddings is not None: return self.enriched_embeddings @@ -726,7 +726,7 @@ def post_parameter_update(self) -> None: # noqa: D102 # invalidate enriched embeddings self.enriched_embeddings = None - def reset_parameters(self): + def reset_parameters(self): # noqa:D102 self.base_embeddings.reset_parameters() gain = nn.init.calculate_gain(nonlinearity=self.activation_cls.__name__.lower()) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index e17017ba11..8e6e615801 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -517,20 +517,18 @@ def estimate_cost_of_sequence( *other_shapes: Tuple[int, ...], ) -> int: """Cost of a sequence of broadcasted element-wise operations of tensors, given their shapes.""" - return sum( - map( - numpy.prod, - itertools.islice( - itertools.accumulate( - other_shapes, - calculate_broadcasted_elementwise_result_shape, - initial=shape, - ), - 1, - None, - ) - ) - ) + return sum(map( + numpy.prod, + itertools.islice( + itertools.accumulate( + other_shapes, + calculate_broadcasted_elementwise_result_shape, + initial=shape, + ), + 1, + None, + ), + )) @functools.lru_cache(maxsize=32) @@ -538,6 +536,7 @@ def _get_optimal_sequence( *sorted_shapes: Tuple[int, ...], ) -> Tuple[int, Tuple[int, ...]]: """Find the optimal sequence in which to combine tensors element-wise based on the shapes. + The shapes should be sorted to enable efficient caching. :param sorted_shapes: The shapes of the tensors to combine. @@ -552,6 +551,7 @@ def _get_optimal_sequence( def get_optimal_sequence(*shapes: Tuple[int, ...]) -> Tuple[int, Tuple[int, ...]]: """Find the optimal sequence in which to combine tensors elementwise based on the shapes. + :param shapes: The shapes of the tensors to combine. :return: @@ -575,6 +575,7 @@ def _multi_combine( op: Callable[[torch.FloatTensor, torch.FloatTensor], torch.FloatTensor], ) -> torch.FloatTensor: """Broadcasted element-wise combination of tensors. + The optimal execution plan gets cached so that the optimization is only performed once for a fixed set of shapes. :param tensors: diff --git a/tests/test_models.py b/tests/test_models.py index 3a7e3948b7..2566765fa5 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -31,7 +31,10 @@ get_novelty_mask, ) from pykeen.models.cli import build_cli_from_cls -from pykeen.nn.representation import RGCNRepresentations, inverse_indegree_edge_weights, inverse_outdegree_edge_weights, symmetric_edge_weights +from pykeen.nn.representation import ( + RGCNRepresentations, inverse_indegree_edge_weights, inverse_outdegree_edge_weights, + symmetric_edge_weights, +) from pykeen.regularizers import LpRegularizer, collect_regularization_terms from pykeen.testing.mocks import MockRepresentations from pykeen.training import LCWATrainingLoop, SLCWATrainingLoop, TrainingLoop diff --git a/tests/test_utils.py b/tests/test_utils.py index c384d4c90c..60de875009 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- """Tests for the :mod:`pykeen.utils` module.""" + import functools import itertools import operator @@ -15,8 +16,9 @@ import torch from pykeen.utils import ( - _CUDA_OOM_ERROR, _CUDNN_ERROR, calculate_broadcasted_elementwise_result_shape, clamp_norm, combine_complex, compact_mapping, estimate_cost_of_sequence, flatten_dictionary, - get_optimal_sequence, get_until_first_blank, is_cuda_oom_error, is_cudnn_error, project_entity, set_random_seed, split_complex, tensor_product, tensor_sum, + _CUDA_OOM_ERROR, _CUDNN_ERROR, calculate_broadcasted_elementwise_result_shape, clamp_norm, combine_complex, + compact_mapping, estimate_cost_of_sequence, flatten_dictionary, get_optimal_sequence, get_until_first_blank, + is_cuda_oom_error, is_cudnn_error, project_entity, set_random_seed, split_complex, tensor_product, tensor_sum, ) @@ -209,7 +211,7 @@ def test_is_cudnn_error(self): def test_calculate_broadcasted_elementwise_result_shape(): """Test calculate_broadcasted_elementwise_result_shape.""" max_dim = 64 - for n_dim, i in itertools.product(range(2, 5), range(10)): + for n_dim, _ in itertools.product(range(2, 5), range(10)): a_shape = [1 for _ in range(n_dim)] b_shape = [1 for _ in range(n_dim)] for j in range(n_dim): @@ -234,7 +236,7 @@ def _generate_shapes( ) -> Iterable[Tuple[Tuple[int, ...], ...]]: """Generate shapes.""" max_shape = torch.randint(low=2, high=32, size=(128,)) - for _i in range(iterations): + for _ in range(iterations): # create broadcastable shapes idx = torch.randperm(max_shape.shape[0])[:n_dim] this_max_shape = max_shape[idx] From 90b55fcaa0882c912955191b27d073741840576a Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Wed, 25 Nov 2020 16:50:32 +0100 Subject: [PATCH 606/690] Add docs stubs --- .../extending/extending_interactors.rst | 25 +++++++++++++++++++ docs/source/extending/index.rst | 6 +++++ docs/source/index.rst | 1 + 3 files changed, 32 insertions(+) create mode 100644 docs/source/extending/extending_interactors.rst create mode 100644 docs/source/extending/index.rst diff --git a/docs/source/extending/extending_interactors.rst b/docs/source/extending/extending_interactors.rst new file mode 100644 index 0000000000..4dce2d909f --- /dev/null +++ b/docs/source/extending/extending_interactors.rst @@ -0,0 +1,25 @@ +Extending the Interaction Functions +=================================== +In [ali2020]_, we argued that a knowledge graph embedding model (KGEM) consists of +several components: an interaction function, a loss function, a training approach, etc. + +Let's assume you have invented a new interaction model, +e.g. this variant of :class:`pykeen.models.DistMult` + +.. math:: + + f(h, r, t) = + +where :math:`h,r,t \in \mathbb{R}^d`, and :math:`\sigma` denotes the logistic sigmoid. + +.. code-block:: python + + from pykeen.nn import Interaction + + class ModifiedDistMultInteraction(Interaction): + def forward(self, h, r, t): + return h * r.sigmoid() * t + + +.. [ali2020] Mehdi, A., *et al.* (2020) `PyKEEN 1.0: A Python Library for Training and + Evaluating Knowledge Graph Embeddings `_ *arXiv*, 2007.14175. diff --git a/docs/source/extending/index.rst b/docs/source/extending/index.rst new file mode 100644 index 0000000000..6e37cfcff5 --- /dev/null +++ b/docs/source/extending/index.rst @@ -0,0 +1,6 @@ +Extending PyKEEN +================ +.. toctree:: + :name: extending + + extending_interactors diff --git a/docs/source/index.rst b/docs/source/index.rst index ee433108b9..013bf50dbc 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -17,6 +17,7 @@ PyKEEN tutorial/trackers/index tutorial/making_predictions tutorial/performance + extending/index .. toctree:: :caption: Reference From 305321781ad3ad57dbc31361cd6d83907334ba2a Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Wed, 25 Nov 2020 17:01:09 +0100 Subject: [PATCH 607/690] Pass mypy --- src/pykeen/nn/modules.py | 2 +- src/pykeen/nn/representation.py | 11 +++++++---- src/pykeen/utils.py | 6 +++--- 3 files changed, 11 insertions(+), 8 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 7a2079cd81..ea9c6e9937 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -93,7 +93,7 @@ def from_func(cls, f) -> 'Interaction': @classmethod def cls_from_func(cls, f) -> Type['Interaction']: """Create a stateless interaction class.""" - class StatelessInteraction(cls): + class StatelessInteraction(cls): # type: ignore func = f return StatelessInteraction diff --git a/src/pykeen/nn/representation.py b/src/pykeen/nn/representation.py index dd113d49ce..efec14888d 100644 --- a/src/pykeen/nn/representation.py +++ b/src/pykeen/nn/representation.py @@ -60,13 +60,16 @@ def get_expected_canonical_shape( :return: (batch_size, num_heads, num_relations, num_tails, ``*``). The expected shape, a tuple of at least 5 positive integers. """ - if torch.is_tensor(indices): - indices = indices.shape - exp_shape = [1, 1, 1, 1] + list(suffix_shape) + if isinstance(suffix_shape, int): + exp_shape = [1, 1, 1, 1, suffix_shape] + else: + exp_shape = [1, 1, 1, 1, *suffix_shape] dim = _normalize_dim(dim=dim) if indices is None: # 1-n scoring - exp_shape[dim] = num + exp_shape[dim] = num # type: ignore else: # batch dimension + if isinstance(indices, torch.Tensor): + indices = indices.shape exp_shape[0] = indices[0] if len(indices) > 1: # multi-target batching exp_shape[dim] = indices[1] diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 8e6e615801..1a950c3d8d 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -587,9 +587,9 @@ def _multi_combine( The elementwise combination evaluated in optimal processing order. """ # determine optimal processing order - order = get_optimal_sequence(*(t.shape for t in tensors))[1] - tensors = [tensors[i] for i in order] - return functools.reduce(op, tensors[1:], tensors[0]) + _, order = get_optimal_sequence(*(t.shape for t in tensors)) + head, *rest = [tensors[i] for i in order] + return functools.reduce(op, rest, head) def tensor_sum(*x: torch.FloatTensor) -> torch.FloatTensor: From b9000a3aafc318e5f71131b34a5bd0eb3c29f258 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 25 Nov 2020 17:39:25 +0100 Subject: [PATCH 608/690] Delete unused variable --- src/pykeen/models/unimodal/kg2e.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index ff89882a22..b7d5875e43 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -2,7 +2,6 @@ """Implementation of KG2E.""" -import math from typing import Optional import torch @@ -20,8 +19,6 @@ 'KG2E', ] -_LOG_2_PI = math.log(math.tau) - class KG2E(ERModel): r"""An implementation of KG2E from [he2015]_. From ecf12a9b8439713e0185aab12b87dc25a67a95d7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 25 Nov 2020 18:30:57 +0100 Subject: [PATCH 609/690] extract common computation for KG2E similarities --- src/pykeen/nn/compute_kernel.py | 29 +++++++++++++++++++++++++++++ src/pykeen/nn/sim.py | 12 ++++++++---- 2 files changed, 37 insertions(+), 4 deletions(-) create mode 100644 src/pykeen/nn/compute_kernel.py diff --git a/src/pykeen/nn/compute_kernel.py b/src/pykeen/nn/compute_kernel.py new file mode 100644 index 0000000000..f057e9553f --- /dev/null +++ b/src/pykeen/nn/compute_kernel.py @@ -0,0 +1,29 @@ +import torch + + +def _batched_dot_manual( + a: torch.FloatTensor, + b: torch.FloatTensor, +) -> torch.FloatTensor: + return (a * b).sum(dim=-1) + + +def _batched_dot_matmul( + a: torch.FloatTensor, + b: torch.FloatTensor, +) -> torch.FloatTensor: + return (a.unsqueeze(dim=-2) @ b.unsqueeze(dim=-1)).view(a.shape[:-1]) + + +def _batched_dot_einsum( + a: torch.FloatTensor, + b: torch.FloatTensor, +) -> torch.FloatTensor: + return torch.einsum("...i,...i->...", a, b) + + +def batched_dot( + a: torch.FloatTensor, + b: torch.FloatTensor, +) -> torch.FloatTensor: + return _batched_dot_manual(a, b) diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index 01c261fca3..ed44aa634e 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -6,6 +6,7 @@ import torch +from .compute_kernel import batched_dot from ..typing import GaussianDistribution from ..utils import tensor_sum @@ -54,8 +55,10 @@ def expected_likelihood( #: a = \mu^T\Sigma^{-1}\mu safe_sigma = torch.clamp_min(var, min=epsilon) - sigma_inv = torch.reciprocal(safe_sigma) - sim = torch.sum(sigma_inv * mean ** 2, dim=-1) + sim = batched_dot( + a=safe_sigma.reciprocal(), + b=(mean ** 2), + ) #: b = \log \det \Sigma sim = sim + safe_sigma.log().sum(dim=-1) @@ -116,7 +119,8 @@ def kullback_leibler_similarity( # since sigma_0, sigma_1 are diagonal matrices: # = sum (sigma_1^-1[i] sigma_0[i]) = sum (sigma_0[i] / sigma_1[i]) r_var_safe = r_var.clamp_min(min=epsilon) - terms.append((e_var / r_var_safe).sum(dim=-1)) + var_safe_reciprocal = r_var_safe.reciprocal() + terms.append(batched_dot(e_var, var_safe_reciprocal)) # 2. Component # (mu_1 - mu_0) * Sigma_1^-1 (mu_1 - mu_0) @@ -125,7 +129,7 @@ def kullback_leibler_similarity( # since Sigma_1 is diagonal # = mu**2 / sigma_1 mu = r_mean - e_mean - terms.append((mu.pow(2) / r_var_safe).sum(dim=-1)) + terms.append(batched_dot(mu.pow(2), r_var_safe)) # 3. Component if exact: From b51adea5b05fd45a8a0c22208fed0a4e7f1ae697 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 25 Nov 2020 18:39:26 +0100 Subject: [PATCH 610/690] Pass h/t down to KG2E similarity functions --- src/pykeen/nn/functional.py | 23 ++++++++++---------- src/pykeen/nn/sim.py | 42 +++++++++++++++++++++---------------- 2 files changed, 35 insertions(+), 30 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 275f98c010..456b60b825 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -103,11 +103,11 @@ def _complex_interaction_optimized_broadcasted( return tensor_sum(*( factor * tensor_product(hh, rr, tt).sum(dim=-1) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) @@ -450,13 +450,12 @@ def kg2e_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - similarity_fn = KG2E_SIMILARITIES[similarity] - # Compute entity distribution - e_mean = h_mean - t_mean - e_var = h_var + t_var - e = GaussianDistribution(mean=e_mean, diagonal_covariance=e_var) - r = GaussianDistribution(mean=r_mean, diagonal_covariance=r_var) - return similarity_fn(e=e, r=r, exact=exact) + return KG2E_SIMILARITIES[similarity]( + h=GaussianDistribution(mean=h_mean, diagonal_covariance=h_var), + r=GaussianDistribution(mean=r_mean, diagonal_covariance=r_var), + t=GaussianDistribution(mean=t_mean, diagonal_covariance=t_var), + exact=exact, + ) def ntn_interaction( diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index ed44aa634e..cd89655ca0 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -18,8 +18,9 @@ def expected_likelihood( - e: GaussianDistribution, + h: GaussianDistribution, r: GaussianDistribution, + t: GaussianDistribution, epsilon: float = 1.0e-10, exact: bool = True, ) -> torch.FloatTensor: @@ -37,10 +38,14 @@ def expected_likelihood( + \log \det \Sigma + d \log (2 \pi) \right) - :param e: shape: (batch_size, num_heads, 1, num_tails, d) - The entity Gaussian distribution. + with :math:`\mu_e = \mu_h - \mu_t` and :math:`\Sigma_e = \Sigma_h + \Sigma_t`. + + :param h: shape: (batch_size, num_heads, 1, 1, d) + The head entity Gaussian distribution. :param r: shape: (batch_size, 1, num_relations, 1, d) The relation Gaussian distribution. + :param t: shape: (batch_size, 1, 1, num_tails, d) + The tail entity Gaussian distribution. :param epsilon: float (default=1.0) Small constant used to avoid numerical issues when dividing. :param exact: @@ -50,8 +55,8 @@ def expected_likelihood( The similarity. """ # subtract, shape: (batch_size, num_heads, num_relations, num_tails, dim) - var = r.diagonal_covariance + e.diagonal_covariance - mean = e.mean - r.mean + var = tensor_sum(*(d.diagonal_covariance for d in (h, r, t))) + mean = tensor_sum(h.mean, -t.mean, -r.mean) #: a = \mu^T\Sigma^{-1}\mu safe_sigma = torch.clamp_min(var, min=epsilon) @@ -68,8 +73,9 @@ def expected_likelihood( def kullback_leibler_similarity( - e: GaussianDistribution, + h: GaussianDistribution, r: GaussianDistribution, + t: GaussianDistribution, epsilon: float = 1.0e-10, exact: bool = True, ) -> torch.FloatTensor: @@ -86,16 +92,20 @@ def kullback_leibler_similarity( + ln (det(\Sigma_1) / det(\Sigma_0)) ) + with :math:`\mu_e = \mu_h - \mu_t` and :math:`\Sigma_e = \Sigma_h + \Sigma_t`. + .. note :: This methods assumes diagonal covariance matrices :math:`\Sigma`. .. seealso :: https://en.wikipedia.org/wiki/Multivariate_normal_distribution#Kullback%E2%80%93Leibler_divergence - :param e: shape: (batch_size, num_heads, 1, num_tails, d) - The entity Gaussian distributions, as mean/diagonal covariance pairs. + :param h: shape: (batch_size, num_heads, 1, 1, d) + The head entity Gaussian distribution. :param r: shape: (batch_size, 1, num_relations, 1, d) - The relation Gaussian distributions, as mean/diagonal covariance pairs. + The relation Gaussian distribution. + :param t: shape: (batch_size, 1, 1, num_tails, d) + The tail entity Gaussian distribution. :param epsilon: float (default=1.0) Small constant used to avoid numerical issues when dividing. :param exact: @@ -104,13 +114,10 @@ def kullback_leibler_similarity( :return: torch.Tensor, shape: (s_1, ..., s_k) The similarity. """ - assert (e.diagonal_covariance > 0).all() and (r.diagonal_covariance > 0).all() - - e_mean = e.mean - e_var = e.diagonal_covariance + assert all((d.diagonal_covariance > 0).all() for d in (h, r, t)) - r_mean = r.mean - r_var = r.diagonal_covariance + e_var = (h.diagonal_covariance + t.diagonal_covariance) + r_var_safe = r.diagonal_covariance.clamp_min(min=epsilon) terms = [] @@ -118,7 +125,6 @@ def kullback_leibler_similarity( # tr(sigma_1^-1 sigma_0) = sum (sigma_1^-1 sigma_0)[i, i] # since sigma_0, sigma_1 are diagonal matrices: # = sum (sigma_1^-1[i] sigma_0[i]) = sum (sigma_0[i] / sigma_1[i]) - r_var_safe = r_var.clamp_min(min=epsilon) var_safe_reciprocal = r_var_safe.reciprocal() terms.append(batched_dot(e_var, var_safe_reciprocal)) @@ -128,12 +134,12 @@ def kullback_leibler_similarity( # = mu * Sigma_1^-1 mu # since Sigma_1 is diagonal # = mu**2 / sigma_1 - mu = r_mean - e_mean + mu = tensor_sum(r.mean, -h.mean, t.mean) terms.append(batched_dot(mu.pow(2), r_var_safe)) # 3. Component if exact: - terms.append(-e_mean.shape[-1]) + terms.append(-h.mean.shape[-1]) # 4. Component # ln (det(\Sigma_1) / det(\Sigma_0)) From 051e24d40c25a3a218a8b856960ff042448b8af7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 25 Nov 2020 18:41:34 +0100 Subject: [PATCH 611/690] Add docstring --- src/pykeen/nn/compute_kernel.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pykeen/nn/compute_kernel.py b/src/pykeen/nn/compute_kernel.py index f057e9553f..93835e2ce9 100644 --- a/src/pykeen/nn/compute_kernel.py +++ b/src/pykeen/nn/compute_kernel.py @@ -1,3 +1,4 @@ +"""Compute kernels for common sub-tasks.""" import torch @@ -26,4 +27,5 @@ def batched_dot( a: torch.FloatTensor, b: torch.FloatTensor, ) -> torch.FloatTensor: + """Compute "element-wise" dot-product between batched vectors.""" return _batched_dot_manual(a, b) From 10798604bceb3d8357d3a8169fcb3fe423849140 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 25 Nov 2020 18:41:45 +0100 Subject: [PATCH 612/690] Add tensor_sum / tensor_product to __all__ --- src/pykeen/utils.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 594da9bcbe..a6db6cc4ae 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -51,6 +51,8 @@ 'set_random_seed', 'split_complex', 'split_list_in_batches_iter', + 'tensor_sum', + 'tensor_product', 'upgrade_to_sequence', 'view_complex', 'NoRandomSeedNecessary', From 2db70f68919dac71d6992d011645524d9ee5c4e8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 25 Nov 2020 19:07:57 +0100 Subject: [PATCH 613/690] Do not use tensor_sum when it is clear that all summands are in the same shape --- src/pykeen/nn/functional.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 456b60b825..ac14afb2d9 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -100,14 +100,14 @@ def _complex_interaction_optimized_broadcasted( ) -> torch.FloatTensor: """Manually split into real/imag, and used optimized broadcasted combination.""" (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] - return tensor_sum(*( + return sum(*( factor * tensor_product(hh, rr, tt).sum(dim=-1) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) From 6e5eeccb1a591423898d1c4b30ec4a3d09caadc6 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 25 Nov 2020 19:08:21 +0100 Subject: [PATCH 614/690] Do not use tensor_sum when it is clear that all summands are in the same shape --- src/pykeen/nn/functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index ac14afb2d9..7699956091 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -210,7 +210,7 @@ def conve_interaction( x = (x @ t).squeeze(dim=-2) # add bias term - return tensor_sum(x, t_bias.squeeze(dim=-1)) + return x + t_bias.squeeze(dim=-1) def convkb_interaction( From f704595fb14e5ed1fe14482e96def5709c69197c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 25 Nov 2020 19:24:24 +0100 Subject: [PATCH 615/690] short-cut tensor_sum for less than 3 tensors --- src/pykeen/utils.py | 29 +++++++++++++++++------------ 1 file changed, 17 insertions(+), 12 deletions(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index a6db6cc4ae..42587ccf2e 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -555,6 +555,7 @@ def _get_optimal_sequence( ) +@functools.lru_cache(maxsize=64) def get_optimal_sequence(*shapes: Tuple[int, ...]) -> Tuple[int, Tuple[int, ...]]: """Find the optimal sequence in which to combine tensors elementwise based on the shapes. @@ -576,36 +577,40 @@ def get_optimal_sequence(*shapes: Tuple[int, ...]) -> Tuple[int, Tuple[int, ...] return cost, optimal_order -def _multi_combine( +def _reorder( tensors: Tuple[torch.FloatTensor, ...], - op: Callable[[torch.FloatTensor, torch.FloatTensor], torch.FloatTensor], -) -> torch.FloatTensor: - """Broadcasted element-wise combination of tensors. +) -> Tuple[torch.FloatTensor, ...]: + """Re-order tensors for broadcasted element-wise combination of tensors. The optimal execution plan gets cached so that the optimization is only performed once for a fixed set of shapes. :param tensors: The tensors, in broadcastable shape. - :param op: - The elementwise operator. :return: - The elementwise combination evaluated in optimal processing order. + The re-ordered tensors in optimal processing order. """ + if len(tensors) < 3: + return tensors # determine optimal processing order - _, order = get_optimal_sequence(*(t.shape for t in tensors)) - head, *rest = [tensors[i] for i in order] - return functools.reduce(op, rest, head) + order = get_optimal_sequence(*(t.shape for t in tensors))[1] + return tuple(tensors[i] for i in order) def tensor_sum(*x: torch.FloatTensor) -> torch.FloatTensor: """Compute elementwise sum of tensors in brodcastable shape.""" - return _multi_combine(tensors=x, op=operator.add) + return sum(_reorder(tensors=x)) + + +def tensor_sum_(*x: torch.FloatTensor) -> torch.FloatTensor: + """Compute elementwise sum of tensors in brodcastable shape.""" + return sum(x[i] for i in get_optimal_sequence(*(t.shape for t in x))[1]) def tensor_product(*x: torch.FloatTensor) -> torch.FloatTensor: """Compute elementwise product of tensors in broadcastable shape.""" - return _multi_combine(tensors=x, op=operator.mul) + head, *rest = _reorder(tensors=x) + return functools.reduce(operator.mul, rest, head) def negative_norm_of_sum( From 88e02dc3b5cb5034ac8f257e39fdae7dd28adc17 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 25 Nov 2020 19:43:31 +0100 Subject: [PATCH 616/690] Add benchmarking script for tensor_sum --- benchmarking/utils.py | 105 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 105 insertions(+) create mode 100644 benchmarking/utils.py diff --git a/benchmarking/utils.py b/benchmarking/utils.py new file mode 100644 index 0000000000..328432547b --- /dev/null +++ b/benchmarking/utils.py @@ -0,0 +1,105 @@ +"""Benchmark utility methods.""" +import functools +import itertools +import operator +import timeit +from typing import Sequence + +import pandas +import torch +from tqdm.auto import tqdm + +from pykeen.utils import tensor_sum + + +def _generate_tensors( + batch_size: int, + num: int, + dim: int, + use_case: str, + shapes: Sequence[str], +) -> Sequence[torch.FloatTensor]: + dims = dict(b=batch_size, h=num, r=num, t=num, d=dim) + canonical = "bhrtd" + return [ + torch.rand( + *( + dims[c] if ( + c in use_case and c in shape + ) else 1 + for c in canonical + ), + dtype=torch.float32, + requires_grad=True, + ) + for shape in shapes + ] + + +def tqdm_itertools_product(*args, **kwargs): + return tqdm(itertools.product(*args), **kwargs, total=functools.reduce(operator.mul, map(len, args), 1)) + + +def check_tensor_sum_performance(): + """Test whether tensor_sum actually offers any performance benefits.""" + batch_sizes = [2 ** i for i in range(5, 10 + 1)] + nums = [2 ** i for i in range(10, 15)] # 2**15 ~ 15k + dims = [2 ** i for i in range(5, 10 + 1)] + use_cases = ["b", "bh", "br", "bt"] # score_hrt, score_h, score_t, score_t + data = [] + progress = tqdm_itertools_product( + batch_sizes, + nums, + dims, + use_cases, + ( + ("TransE/TransD", "bh", "br", "bt"), # h, r, t + ("TransH", "bh", "bhr", "br", "bt", "brt"), # h, - w_r, d_r, -t,, w_r + ), + unit="configuration", + unit_scale=True, + ) + for batch_size, num, dim, use_case, (models, *shapes) in progress: + tensors = _generate_tensors( + batch_size=batch_size, + num=num, + dim=dim, + use_case=use_case, + shapes=shapes, + ) + result_shape = [max(ds) for ds in zip(*(t.shape for t in tensors))] + + # using normal sum + n_samples, time_baseline = timeit.Timer( + stmt="sum(tensors)", + globals=dict(tensors=tensors) + ).autorange() + time_baseline /= n_samples + + # use tensor_sum + n_samples, time = timeit.Timer( + setup="tensor_sum(*tensors)", + stmt="tensor_sum(*tensors)", + globals=dict(tensor_sum=tensor_sum, tensors=tensors) + ).autorange() + time /= n_samples + + data.append((batch_size, num, dim, use_case, shapes, time_baseline, time)) + progress.set_postfix(shape=result_shape, delta=time_baseline - time) + df = pandas.DataFrame(data=data, columns=[ + "batch_size", + "num", + "dim", + "use_case", + "shapes", + "time_sum", + "time_tensor_sum", + ]) + df.to_csv("tensor_sum.perf.tsv", sep="\t", index=False) + print("tensor_sum is better than sum in {percent:2.2%} of all cases.".format( + percent=(df["time_sum"] > df["time_tensor_sum"]).mean()) + ) + + +if __name__ == '__main__': + check_tensor_sum_performance() From 328fb596245cdd6b95d3954977de8fdc280f557c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 25 Nov 2020 19:43:57 +0100 Subject: [PATCH 617/690] Remove unused method --- src/pykeen/utils.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 42587ccf2e..847962cb94 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -602,11 +602,6 @@ def tensor_sum(*x: torch.FloatTensor) -> torch.FloatTensor: return sum(_reorder(tensors=x)) -def tensor_sum_(*x: torch.FloatTensor) -> torch.FloatTensor: - """Compute elementwise sum of tensors in brodcastable shape.""" - return sum(x[i] for i in get_optimal_sequence(*(t.shape for t in x))[1]) - - def tensor_product(*x: torch.FloatTensor) -> torch.FloatTensor: """Compute elementwise product of tensors in broadcastable shape.""" head, *rest = _reorder(tensors=x) From 3eeb5415230412d76fe1f0ff8302230747125d8b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 25 Nov 2020 19:52:13 +0100 Subject: [PATCH 618/690] Update utils benchmark with all shape combinations of tensorsum --- benchmarking/utils.py | 22 +++++++++++++--------- 1 file changed, 13 insertions(+), 9 deletions(-) diff --git a/benchmarking/utils.py b/benchmarking/utils.py index 328432547b..594817a43a 100644 --- a/benchmarking/utils.py +++ b/benchmarking/utils.py @@ -42,19 +42,23 @@ def tqdm_itertools_product(*args, **kwargs): def check_tensor_sum_performance(): """Test whether tensor_sum actually offers any performance benefits.""" - batch_sizes = [2 ** i for i in range(5, 10 + 1)] - nums = [2 ** i for i in range(10, 15)] # 2**15 ~ 15k - dims = [2 ** i for i in range(5, 10 + 1)] - use_cases = ["b", "bh", "br", "bt"] # score_hrt, score_h, score_t, score_t data = [] progress = tqdm_itertools_product( - batch_sizes, - nums, - dims, - use_cases, + [2 ** i for i in range(5, 10 + 1)], + [2 ** i for i in range(10, 15)], # 2**15 ~ 15k + [2 ** i for i in range(5, 10 + 1)], + ["b", "bh", "br", "bt"], # score_hrt, score_h, score_t, score_t ( - ("TransE/TransD", "bh", "br", "bt"), # h, r, t + ("ConvKB/ERMLP", "", "bh", "br", "bt"), # conv_bias, h, r, t + ("NTN", "bhrt", "bhr", "bht", "br"), # h w t, vh h, vt t, b + ("ProjE", "bh", "br", ""), # h, r, b_c + ("RotatE", "bhr", "bt"), # hr, -t, + ("RotatE-inv", "bh", "brt"), # h, -(r_inv)t, + ("StructuredEmbedding", "bhr", "brt"), # r h, r t + ("TransE/TransD/KG2E", "bh", "br", "bt"), # h, r, t ("TransH", "bh", "bhr", "br", "bt", "brt"), # h, - w_r, d_r, -t,, w_r + ("TransR", "bhr", "br", "brt"), # h m_r, r, -t m_r + ("UnstructuredModel", "bh", "bt"), # h, r ), unit="configuration", unit_scale=True, From a8bfd734850c878f2d0d48a7ce3e1428121268a0 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 25 Nov 2020 20:00:49 +0100 Subject: [PATCH 619/690] Add click --- benchmarking/utils.py | 60 ++++++++++++++++++++++++++++--------------- 1 file changed, 40 insertions(+), 20 deletions(-) diff --git a/benchmarking/utils.py b/benchmarking/utils.py index 594817a43a..2d9e13b719 100644 --- a/benchmarking/utils.py +++ b/benchmarking/utils.py @@ -3,8 +3,9 @@ import itertools import operator import timeit -from typing import Sequence +from typing import Sequence, Tuple +import click import pandas import torch from tqdm.auto import tqdm @@ -12,41 +13,55 @@ from pykeen.utils import tensor_sum -def _generate_tensors( +def _generate_shapes( batch_size: int, num: int, dim: int, use_case: str, shapes: Sequence[str], -) -> Sequence[torch.FloatTensor]: +) -> Sequence[Tuple[int, ...]]: dims = dict(b=batch_size, h=num, r=num, t=num, d=dim) canonical = "bhrtd" return [ - torch.rand( - *( - dims[c] if ( - c in use_case and c in shape - ) else 1 - for c in canonical - ), - dtype=torch.float32, - requires_grad=True, + tuple( + dims[c] if ( + c in use_case and c in shape + ) else 1 + for c in canonical ) for shape in shapes ] +def _generate_tensors(*shapes: Tuple[int, ...]) -> Sequence[torch.FloatTensor]: + return [ + torch.rand(shape, requires_grad=True, dtype=torch.float32) + for shape in shapes + ] + + def tqdm_itertools_product(*args, **kwargs): return tqdm(itertools.product(*args), **kwargs, total=functools.reduce(operator.mul, map(len, args), 1)) -def check_tensor_sum_performance(): +@click.command() +@click.option('-m', '--max-result-power', type=int, default=30, show_default=True) +@click.option('-b', '--max-batch-size-power', type=int, default=10, show_default=True) +@click.option('-n', '--max-num-power', type=int, default=14, show_default=True) +@click.option('-d', '--max-dim-power', type=int, default=10, show_default=True) +def main( + max_result_power: int, + max_batch_size_power: int, + max_dim_power: int, + max_num_power: int, +): """Test whether tensor_sum actually offers any performance benefits.""" + max_size = 2 ** max_result_power data = [] progress = tqdm_itertools_product( - [2 ** i for i in range(5, 10 + 1)], - [2 ** i for i in range(10, 15)], # 2**15 ~ 15k - [2 ** i for i in range(5, 10 + 1)], + [2 ** i for i in range(5, max_batch_size_power + 1)], + [2 ** i for i in range(10, max_num_power + 1)], # 2**15 ~ 15k + [2 ** i for i in range(5, max_dim_power + 1)], ["b", "bh", "br", "bt"], # score_hrt, score_h, score_t, score_t ( ("ConvKB/ERMLP", "", "bh", "br", "bt"), # conv_bias, h, r, t @@ -64,14 +79,18 @@ def check_tensor_sum_performance(): unit_scale=True, ) for batch_size, num, dim, use_case, (models, *shapes) in progress: - tensors = _generate_tensors( + this_shapes = _generate_shapes( batch_size=batch_size, num=num, dim=dim, use_case=use_case, shapes=shapes, ) - result_shape = [max(ds) for ds in zip(*(t.shape for t in tensors))] + result_shape = _get_result_shape(tensors) + num_result_elements = functools.reduce(operator.mul, result_shape, 1) + if num_result_elements > max_size: + continue + tensors = _generate_tensors(*this_shapes) # using normal sum n_samples, time_baseline = timeit.Timer( @@ -105,5 +124,6 @@ def check_tensor_sum_performance(): ) -if __name__ == '__main__': - check_tensor_sum_performance() +def _get_result_shape(tensors): + result_shape = [max(ds) for ds in zip(*(t.shape for t in tensors))] + return result_shape From bcc4a349e13f8b897ca14d630b32aead638f1917 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 25 Nov 2020 20:03:41 +0100 Subject: [PATCH 620/690] fixes --- benchmarking/utils.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/benchmarking/utils.py b/benchmarking/utils.py index 2d9e13b719..ff50d2046e 100644 --- a/benchmarking/utils.py +++ b/benchmarking/utils.py @@ -44,6 +44,10 @@ def tqdm_itertools_product(*args, **kwargs): return tqdm(itertools.product(*args), **kwargs, total=functools.reduce(operator.mul, map(len, args), 1)) +def _get_result_shape(shapes: Sequence[Tuple[int, ...]]) -> Tuple[int, ...]: + return tuple(max(ds) for ds in zip(*shapes)) + + @click.command() @click.option('-m', '--max-result-power', type=int, default=30, show_default=True) @click.option('-b', '--max-batch-size-power', type=int, default=10, show_default=True) @@ -86,7 +90,7 @@ def main( use_case=use_case, shapes=shapes, ) - result_shape = _get_result_shape(tensors) + result_shape = _get_result_shape(this_shapes) num_result_elements = functools.reduce(operator.mul, result_shape, 1) if num_result_elements > max_size: continue @@ -124,6 +128,5 @@ def main( ) -def _get_result_shape(tensors): - result_shape = [max(ds) for ds in zip(*(t.shape for t in tensors))] - return result_shape +if __name__ == '__main__': + main() From 02e2293ba9a4b50aa6c87037d0f03d9550cc4006 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 26 Nov 2020 16:53:22 +0100 Subject: [PATCH 621/690] Do not use tensor_sum for 2-element list --- src/pykeen/nn/functional.py | 26 ++++++++++---------------- src/pykeen/utils.py | 29 ++++++++++++++++++++++++----- 2 files changed, 34 insertions(+), 21 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 7699956091..53ca41784e 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -12,7 +12,7 @@ from .sim import KG2E_SIMILARITIES from ..typing import GaussianDistribution from ..utils import ( - broadcast_cat, clamp_norm, estimate_cost_of_sequence, extended_einsum, is_cudnn_error, negative_norm_of_sum, + broadcast_cat, clamp_norm, estimate_cost_of_sequence, extended_einsum, is_cudnn_error, negative_norm, negative_norm_of_sum, project_entity, split_complex, tensor_product, tensor_sum, view_complex, ) @@ -103,11 +103,11 @@ def _complex_interaction_optimized_broadcasted( return sum(*( factor * tensor_product(hh, rr, tt).sum(dim=-1) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) @@ -606,12 +606,7 @@ def rotate_interaction( t = t * torch.conj(r) # Workaround until https://github.com/pytorch/pytorch/issues/30704 is fixed - return negative_norm_of_sum( - h, - -t, - p=2, - power_norm=False, - ) + return negative_norm(h - t, p=2, power_norm=False) def simple_interaction( @@ -683,9 +678,8 @@ def structured_embedding_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return negative_norm_of_sum( - (r_h @ h.unsqueeze(dim=-1)).squeeze(dim=-1), - -(r_t @ t.unsqueeze(dim=-1)).squeeze(dim=-1), + return negative_norm( + (r_h @ h.unsqueeze(dim=-1) - r_t @ t.unsqueeze(dim=-1)).squeeze(dim=-1), p=p, power_norm=power_norm, ) @@ -931,4 +925,4 @@ def unstructured_model_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return negative_norm_of_sum(h, -t, p=p, power_norm=power_norm) + return negative_norm(h - t, p=p, power_norm=power_norm) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 847962cb94..906d63c48f 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -625,17 +625,36 @@ def negative_norm_of_sum( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - d: torch.FloatTensor = tensor_sum(*x) + return negative_norm(tensor_sum(*x), p=p, power_norm=power_norm) + + +def negative_norm( + x: torch.FloatTensor, + p: Union[str, int] = 2, + power_norm: bool = False, +) -> torch.FloatTensor: + """Evaluate negative norm of a vector. + + :param x: shape: (batch_size, num_heads, num_relations, num_tails, dim) + The vectors. + :param p: + The p for the norm. cf. torch.norm. + :param power_norm: + Whether to return $|x-y|_p^p$, cf. https://github.com/pytorch/pytorch/issues/28119 + + :return: shape: (batch_size, num_heads, num_relations, num_tails) + The scores. + """ if power_norm: assert not isinstance(p, str) - return -(d.abs() ** p).sum(dim=-1) + return -(x.abs() ** p).sum(dim=-1) - if torch.is_complex(d): + if torch.is_complex(x): assert not isinstance(p, str) # workaround for complex numbers: manually compute norm - return -(d.abs() ** p).sum(dim=-1) ** (1 / p) + return -(x.abs() ** p).sum(dim=-1) ** (1 / p) - return -d.norm(p=p, dim=-1) + return -x.norm(p=p, dim=-1) def extended_einsum( From 02fb26a526eb122c2e3beef5bc455e2a2e95f18a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 26 Nov 2020 17:05:02 +0100 Subject: [PATCH 622/690] Add caching test for get_optimal_sequence --- tests/test_utils.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/test_utils.py b/tests/test_utils.py index 3cb85e9957..0a51400715 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -289,6 +289,25 @@ def test_estimate_cost_of_add_sequence(): assert (numpy.corrcoef(x=a[:, 0], y=a[:, 1])[0, 1]) > 0.8 +@pytest.mark.slow +def test_get_optimal_add_sequence_caching(): + """Test caching of ``get_optimal_add_sequence()``.""" + for shapes in _generate_shapes(iterations=10): + # get optimal sequence + first_time = timeit.default_timer() + get_optimal_sequence(*shapes) + first_time = timeit.default_timer() - first_time + + # check caching + samples, second_time = timeit.Timer(stmt="get_optimal_sequence(*shapes)", globals=dict( + get_optimal_sequence=get_optimal_sequence, + shapes=shapes, + )).autorange() + second_time /= samples + + assert second_time < first_time + + def test_get_optimal_add_sequence(): """Test ``get_optimal_add_sequence()``.""" for shapes in _generate_shapes(): From 0d6c7df874a3489cf36cff4b91ded0c5014859b6 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 26 Nov 2020 17:05:35 +0100 Subject: [PATCH 623/690] Fix test name --- tests/test_utils.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index 0a51400715..ae9c6e54a4 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -290,8 +290,8 @@ def test_estimate_cost_of_add_sequence(): @pytest.mark.slow -def test_get_optimal_add_sequence_caching(): - """Test caching of ``get_optimal_add_sequence()``.""" +def test_get_optimal_sequence_caching(): + """Test caching of ``get_optimal_sequence()``.""" for shapes in _generate_shapes(iterations=10): # get optimal sequence first_time = timeit.default_timer() @@ -304,12 +304,12 @@ def test_get_optimal_add_sequence_caching(): shapes=shapes, )).autorange() second_time /= samples - + assert second_time < first_time -def test_get_optimal_add_sequence(): - """Test ``get_optimal_add_sequence()``.""" +def test_get_optimal_sequence(): + """Test ``get_optimal_sequence()``.""" for shapes in _generate_shapes(): # get optimal sequence opt_cost, opt_seq = get_optimal_sequence(*shapes) From bd04e55057f303a720418b3505de2b75fa51a316 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 26 Nov 2020 17:08:27 +0100 Subject: [PATCH 624/690] Add additional heuristic for score_hrt case --- src/pykeen/utils.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 906d63c48f..ac1103d244 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -593,7 +593,11 @@ def _reorder( if len(tensors) < 3: return tensors # determine optimal processing order - order = get_optimal_sequence(*(t.shape for t in tensors))[1] + shapes = tuple(tuple(t.shape) for t in tensors) + if len(set(s[0] for s in shapes)) < 2: + # heuristic + return tensors + order = get_optimal_sequence(*shapes)[1] return tuple(tensors[i] for i in order) From 041a8f0c10cb012f5e2f82e7547398ca69f5828f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 26 Nov 2020 17:10:11 +0100 Subject: [PATCH 625/690] do not benchmark two-element tensor_sums --- benchmarking/utils.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/benchmarking/utils.py b/benchmarking/utils.py index ff50d2046e..8d00eb9fb9 100644 --- a/benchmarking/utils.py +++ b/benchmarking/utils.py @@ -71,13 +71,13 @@ def main( ("ConvKB/ERMLP", "", "bh", "br", "bt"), # conv_bias, h, r, t ("NTN", "bhrt", "bhr", "bht", "br"), # h w t, vh h, vt t, b ("ProjE", "bh", "br", ""), # h, r, b_c - ("RotatE", "bhr", "bt"), # hr, -t, - ("RotatE-inv", "bh", "brt"), # h, -(r_inv)t, + # ("RotatE", "bhr", "bt"), # hr, -t, + # ("RotatE-inv", "bh", "brt"), # h, -(r_inv)t, ("StructuredEmbedding", "bhr", "brt"), # r h, r t ("TransE/TransD/KG2E", "bh", "br", "bt"), # h, r, t ("TransH", "bh", "bhr", "br", "bt", "brt"), # h, - w_r, d_r, -t,, w_r ("TransR", "bhr", "br", "brt"), # h m_r, r, -t m_r - ("UnstructuredModel", "bh", "bt"), # h, r + # ("UnstructuredModel", "bh", "bt"), # h, r ), unit="configuration", unit_scale=True, From f6c07e7e19f21844398f5e39c93e6b83aade7499 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 30 Nov 2020 12:15:31 +0100 Subject: [PATCH 626/690] Remove amo from model --- src/pykeen/models/base.py | 17 ----------------- src/pykeen/models/multimodal/base.py | 2 -- src/pykeen/models/multimodal/complex_literal.py | 2 -- .../models/multimodal/distmult_literal.py | 2 -- src/pykeen/models/unimodal/complex.py | 5 ----- src/pykeen/models/unimodal/conv_e.py | 2 -- src/pykeen/models/unimodal/conv_kb.py | 2 -- src/pykeen/models/unimodal/distmult.py | 2 -- src/pykeen/models/unimodal/ermlp.py | 2 -- src/pykeen/models/unimodal/ermlpe.py | 2 -- src/pykeen/models/unimodal/hole.py | 2 -- src/pykeen/models/unimodal/kg2e.py | 2 -- src/pykeen/models/unimodal/ntn.py | 2 -- src/pykeen/models/unimodal/proj_e.py | 13 +------------ src/pykeen/models/unimodal/rescal.py | 2 -- src/pykeen/models/unimodal/rgcn.py | 2 -- src/pykeen/models/unimodal/rotate.py | 2 -- src/pykeen/models/unimodal/simple.py | 2 -- .../models/unimodal/structured_embedding.py | 2 -- src/pykeen/models/unimodal/trans_d.py | 2 -- src/pykeen/models/unimodal/trans_e.py | 2 -- src/pykeen/models/unimodal/trans_h.py | 2 -- src/pykeen/models/unimodal/trans_r.py | 2 -- src/pykeen/models/unimodal/tucker.py | 2 -- .../models/unimodal/unstructured_model.py | 2 -- src/pykeen/testing/mocks.py | 3 +-- 26 files changed, 2 insertions(+), 80 deletions(-) diff --git a/src/pykeen/models/base.py b/src/pykeen/models/base.py index 15f470daa4..ac68b2d5a2 100644 --- a/src/pykeen/models/base.py +++ b/src/pykeen/models/base.py @@ -254,7 +254,6 @@ def __init__( triples_factory: TriplesFactory, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, ) -> None: @@ -268,10 +267,6 @@ def __init__( Whether to apply sigmoid onto the scores when predicting scores. Applying sigmoid at prediction time may lead to exactly equal scores for certain triples with very high, or very low score. When not trained with applying sigmoid (or using BCEWithLogitsLoss), the scores are not calibrated to perform well with sigmoid. - :param automatic_memory_optimization: - If set to `True`, the model derives the maximum possible batch sizes for the scoring of triples during - evaluation and also training (if no batch size was given). This allows to fully utilize the hardware at hand - and achieves the fastest calculations possible. :param preferred_device: The preferred device for model training and inference. :param random_seed: @@ -288,9 +283,6 @@ def __init__( elif random_seed is not NoRandomSeedNecessary: set_random_seed(random_seed) - if automatic_memory_optimization is None: - automatic_memory_optimization = True - # Loss if loss is None: self.loss = self.loss_default(**(self.loss_default_kwargs or {})) @@ -310,9 +302,6 @@ def __init__( ''' self.predict_with_sigmoid = predict_with_sigmoid - # This allows to store the optimized parameters - self.automatic_memory_optimization = automatic_memory_optimization - def __init_subclass__(cls, autoreset: bool = True, **kwargs): # noqa:D105 cls._is_base_model = not autoreset if not cls._is_base_model: @@ -1166,7 +1155,6 @@ def __init__( ] = None, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, ) -> None: @@ -1180,10 +1168,6 @@ def __init__( Whether to apply sigmoid onto the scores when predicting scores. Applying sigmoid at prediction time may lead to exactly equal scores for certain triples with very high, or very low score. When not trained with applying sigmoid (or using BCEWithLogitsLoss), the scores are not calibrated to perform well with sigmoid. - :param automatic_memory_optimization: - If set to `True`, the model derives the maximum possible batch sizes for the scoring of triples during - evaluation and also training (if no batch size was given). This allows to fully utilize the hardware at hand - and achieves the fastest calculations possible. :param preferred_device: The preferred device for model training and inference. :param random_seed: @@ -1191,7 +1175,6 @@ def __init__( """ super().__init__( triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/multimodal/base.py b/src/pykeen/models/multimodal/base.py index 23ac535ab6..628b66d79b 100644 --- a/src/pykeen/models/multimodal/base.py +++ b/src/pykeen/models/multimodal/base.py @@ -33,7 +33,6 @@ def __init__( relation_specification: Optional[EmbeddingSpecification] = None, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, - automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, ): @@ -42,7 +41,6 @@ def __init__( interaction=interaction, loss=loss, predict_with_sigmoid=predict_with_sigmoid, - automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, entity_representations=[ diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index 008310fc74..030e050cee 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -76,7 +76,6 @@ def __init__( self, triples_factory: TriplesNumericLiteralsFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, input_dropout: float = 0.2, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, @@ -104,7 +103,6 @@ def __init__( ), loss=loss, predict_with_sigmoid=predict_with_sigmoid, - automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, ) diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 48f265e0a1..ef35d25f70 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -38,7 +38,6 @@ def __init__( self, triples_factory: TriplesNumericLiteralsFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, input_dropout: float = 0.0, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -62,7 +61,6 @@ def __init__( ), loss=loss, predict_with_sigmoid=predict_with_sigmoid, - automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, ) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 6d4814a1cd..7d35d3a385 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -71,7 +71,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -84,9 +83,6 @@ def __init__( The triple factory connected to the model. :param embedding_dim: The embedding dimensionality of the entity embeddings. - :param automatic_memory_optimization: bool - Whether to automatically optimize the sub-batch size during training and batch size during evaluation with - regards to the hardware at hand. :param loss: OptionalLoss (optional) The loss to use. Defaults to SoftplusLoss. :param preferred_device: str (optional) @@ -116,7 +112,6 @@ def __init__( interaction=ComplExInteraction(), entity_representations=embedding_specification, relation_representations=relation_embedding_specification, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 5a29b01b24..000092ca91 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -118,7 +118,6 @@ def __init__( output_dropout: float = 0.3, feature_map_dropout: float = 0.2, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -162,7 +161,6 @@ def __init__( embedding_dim=embedding_dim, initializer=xavier_normal_, ), - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 4d4a20c925..c95769e432 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -74,7 +74,6 @@ def __init__( triples_factory: TriplesFactory, hidden_dropout_rate: float = 0., embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, num_filters: int = 400, @@ -98,7 +97,6 @@ def __init__( embedding_dim=embedding_dim, ), loss=loss, - automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, ) diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index e00a6871b4..5865d6939e 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -71,7 +71,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -101,7 +100,6 @@ def __init__( # Only relation embeddings are regularized regularizer=self._instantiate_default_regularizer(), ), - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 3d0c8cd2be..70f9140822 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -42,7 +42,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -64,7 +63,6 @@ def __init__( relation_representations=EmbeddingSpecification( embedding_dim=embedding_dim, ), - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index 00da44382c..63b7c74c00 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -58,7 +58,6 @@ def __init__( input_dropout: float = 0.2, hidden_dropout: float = 0.3, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -77,7 +76,6 @@ def __init__( relation_representations=EmbeddingSpecification( embedding_dim=embedding_dim, ), - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index dbf70612fa..c082e2d001 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -55,7 +55,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -76,7 +75,6 @@ def __init__( initializer=xavier_uniform_, ), loss=loss, - automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, ) diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index b7d5875e43..8b26fd720f 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -62,7 +62,6 @@ def __init__( c_min: float = 0.05, c_max: float = 5., loss: Optional[Loss] = None, - automatic_memory_optimization: Optional[bool] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, ) -> None: @@ -97,7 +96,6 @@ def __init__( ), entity_representations=representation_spec, relation_representations=representation_spec, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index 59dd97d11c..cea3c0c689 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -55,7 +55,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 100, - automatic_memory_optimization: Optional[bool] = None, num_slices: int = 4, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -89,7 +88,6 @@ def __init__( # u: (k,) EmbeddingSpecification(shape=(num_slices,)), ], - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index 1e2cae2457..b20ac90c82 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -63,19 +63,9 @@ def __init__( loss: Optional[Loss] = None, # Model parameters preferred_device: DeviceHint = None, - automatic_memory_optimization: Optional[bool] = None, random_seed: Optional[int] = None, ) -> None: - """Initialize :class:`ERModel` using :class:`ProjEInteraction`. - - :param triples_factory: - :param embedding_dim: - :param inner_non_linearity: - :param loss: - :param automatic_memory_optimization: - :param preferred_device: - :param random_seed: - """ + """Initialize :class:`ERModel` using :class:`ProjEInteraction`.""" super().__init__( triples_factory=triples_factory, interaction=ProjEInteraction( @@ -90,7 +80,6 @@ def __init__( embedding_dim=embedding_dim, initializer=xavier_uniform_, ), - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index 3ac048c9fa..af92064e16 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -54,7 +54,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -79,7 +78,6 @@ def __init__( shape=(embedding_dim, embedding_dim), regularizer=regularizer, ), - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index d71a6963c7..0fbde1e358 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -62,7 +62,6 @@ def __init__( triples_factory: TriplesFactory, interaction: Optional[Interaction] = None, embedding_dim: int = 500, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, preferred_device: DeviceHint = None, @@ -107,7 +106,6 @@ def __init__( ) super().__init__( triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index 36d83afa3b..f54e102163 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -54,7 +54,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -74,7 +73,6 @@ def __init__( dtype=torch.complex64, ), loss=loss, - automatic_memory_optimization=automatic_memory_optimization, preferred_device=preferred_device, random_seed=random_seed, ) diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index 0adf5cb3f8..b1cae069fe 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -67,7 +67,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -97,7 +96,6 @@ def __init__( regularizer=regularizer, ), ], - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index 73749bf6b5..f834be093a 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -49,7 +49,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -88,7 +87,6 @@ def __init__( initializer=relation_initializer, ), ], - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index dffd7fbd37..a1bed5e588 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -60,7 +60,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, relation_dim: int = 30, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -87,7 +86,6 @@ def __init__( ) for _ in range(2) ], - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 279493dc56..23a4237516 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -50,7 +50,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -80,7 +79,6 @@ def __init__( functional.normalize, ), ), - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index eb957a7c5d..aaed6ce6eb 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -67,7 +67,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 2, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, @@ -98,7 +97,6 @@ def __init__( constrainer=functional.normalize, ), ], - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 76777336bc..81317e55eb 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -63,7 +63,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, relation_dim: int = 30, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, @@ -100,7 +99,6 @@ def __init__( initializer=xavier_uniform_, ), ], - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index e7733a41b8..5740492b2b 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -66,7 +66,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 200, - automatic_memory_optimization: Optional[bool] = None, relation_dim: Optional[int] = None, loss: Optional[Loss] = None, preferred_device: DeviceHint = None, @@ -103,7 +102,6 @@ def __init__( embedding_dim=relation_dim, initializer=xavier_normal_, ), - automatic_memory_optimization=automatic_memory_optimization, loss=loss, preferred_device=preferred_device, random_seed=random_seed, diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index afd299e036..041be13a8b 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -45,7 +45,6 @@ def __init__( self, triples_factory: TriplesFactory, embedding_dim: int = 50, - automatic_memory_optimization: Optional[bool] = None, scoring_fct_norm: int = 1, loss: Optional[Loss] = None, predict_with_sigmoid: bool = False, @@ -60,7 +59,6 @@ def __init__( self.embedding_dim = embedding_dim super().__init__( triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, loss=loss, predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, diff --git a/src/pykeen/testing/mocks.py b/src/pykeen/testing/mocks.py index e67f460600..7a7f670d6f 100644 --- a/src/pykeen/testing/mocks.py +++ b/src/pykeen/testing/mocks.py @@ -21,10 +21,9 @@ class MockModel(Model): """A mock model returning fake scores.""" - def __init__(self, triples_factory: TriplesFactory, automatic_memory_optimization: bool = True): + def __init__(self, triples_factory: TriplesFactory): super().__init__( triples_factory=triples_factory, - automatic_memory_optimization=automatic_memory_optimization, ) def forward( From 41e9dfb359af7017711de6196f5fba197eb9dde3 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 30 Nov 2020 12:31:30 +0100 Subject: [PATCH 627/690] More manual amo merge --- src/pykeen/evaluation/evaluator.py | 4 ++- src/pykeen/evaluation/rank_based_evaluator.py | 3 +- src/pykeen/evaluation/sklearn.py | 9 +++-- src/pykeen/models/cli/builders.py | 3 ++ src/pykeen/models/cli/options.py | 11 +++--- src/pykeen/pipeline.py | 14 ++++++-- src/pykeen/training/slcwa.py | 5 +++ src/pykeen/training/training_loop.py | 9 +++-- tests/test_early_stopping.py | 13 +++---- tests/test_evaluators.py | 6 ++-- tests/test_models.py | 33 +++++++++++++---- tests/test_training.py | 35 +++++++++++++------ 12 files changed, 105 insertions(+), 40 deletions(-) diff --git a/src/pykeen/evaluation/evaluator.py b/src/pykeen/evaluation/evaluator.py index ab562ec3cf..2b1f3e9ad2 100644 --- a/src/pykeen/evaluation/evaluator.py +++ b/src/pykeen/evaluation/evaluator.py @@ -67,11 +67,13 @@ def __init__( requires_positive_mask: bool = False, batch_size: int = None, slice_size: int = None, + automatic_memory_optimization: bool = True, ): self.filtered = filtered self.requires_positive_mask = requires_positive_mask self.batch_size = batch_size self.slice_size = slice_size + self.automatic_memory_optimization = automatic_memory_optimization @classmethod def get_normalized_name(cls) -> str: @@ -135,7 +137,7 @@ def evaluate( if mapped_triples is None: mapped_triples = model.triples_factory.mapped_triples - if batch_size is None and model.automatic_memory_optimization: + if batch_size is None and self.automatic_memory_optimization: batch_size, slice_size = self.batch_and_slice( model=model, mapped_triples=mapped_triples, diff --git a/src/pykeen/evaluation/rank_based_evaluator.py b/src/pykeen/evaluation/rank_based_evaluator.py index 0093d7ca6c..d4e51fa21b 100644 --- a/src/pykeen/evaluation/rank_based_evaluator.py +++ b/src/pykeen/evaluation/rank_based_evaluator.py @@ -218,6 +218,7 @@ def __init__( self, ks: Optional[Iterable[Union[int, float]]] = None, filtered: bool = True, + automatic_memory_optimization: bool = True, **kwargs, ): """Initialize rank-based evaluator. @@ -228,7 +229,7 @@ def __init__( Whether to use the filtered evaluation protocol. If enabled, ranking another true triple higher than the currently considered one will not decrease the score. """ - super().__init__(filtered=filtered, **kwargs) + super().__init__(filtered=filtered, automatic_memory_optimization=automatic_memory_optimization, **kwargs) self.ks = tuple(ks) if ks is not None else (1, 3, 5, 10) for k in self.ks: if isinstance(k, float) and not (0 < k < 1): diff --git a/src/pykeen/evaluation/sklearn.py b/src/pykeen/evaluation/sklearn.py index e93abb9694..0a65a5f1ce 100644 --- a/src/pykeen/evaluation/sklearn.py +++ b/src/pykeen/evaluation/sklearn.py @@ -68,8 +68,13 @@ def get_metric(self, name: str) -> float: # noqa: D102 class SklearnEvaluator(Evaluator): """An evaluator that uses a Scikit-learn metric.""" - def __init__(self, **kwargs): - super().__init__(filtered=False, requires_positive_mask=True, **kwargs) + def __init__(self, automatic_memory_optimization: bool = True, **kwargs): + super().__init__( + filtered=False, + requires_positive_mask=True, + automatic_memory_optimization=automatic_memory_optimization, + **kwargs, + ) self.all_scores = {} self.all_positives = {} diff --git a/src/pykeen/models/cli/builders.py b/src/pykeen/models/cli/builders.py index 30a1dea349..e9eae27265 100644 --- a/src/pykeen/models/cli/builders.py +++ b/src/pykeen/models/cli/builders.py @@ -81,6 +81,7 @@ def _decorate_model_kwargs(command: click.Command) -> click.Command: @options.valiadation_option @options.optimizer_option @options.training_loop_option + @options.automatic_memory_optimization_option @options.number_epochs_option @options.batch_size_option @options.learning_rate_option @@ -107,6 +108,7 @@ def main( mlflow_tracking_uri, title, dataset, + automatic_memory_optimization, training_triples_factory, testing_triples_factory, validation_triples_factory, @@ -158,6 +160,7 @@ def main( title=title, ), random_seed=random_seed, + automatic_memory_optimization=automatic_memory_optimization, ) if not silent: diff --git a/src/pykeen/models/cli/options.py b/src/pykeen/models/cli/options.py index 788fb0d8f9..7572f2b843 100644 --- a/src/pykeen/models/cli/options.py +++ b/src/pykeen/models/cli/options.py @@ -103,12 +103,6 @@ def triples_factory_callback(_, __, path: Optional[str]) -> Optional[TriplesFact default='distmult', show_default=True, ), - 'automatic_memory_optimization': click.option( - '--automatic-memory-optimization', - type=bool, - default=True, - show_default=True, - ), } device_option = click.option( @@ -137,6 +131,11 @@ def triples_factory_callback(_, __, path: Optional[str]) -> Optional[TriplesFact default=_get_default(get_training_loop_cls, suffix=_TRAINING_LOOP_SUFFIX), show_default=True, ) +automatic_memory_optimization_option = click.option( + '--automatic-memory-optimization/--no-automatic-memory-optimization', + default=True, + show_default=True, +) stopper_option = click.option( '--stopper', type=click.Choice(list(stoppers)), diff --git a/src/pykeen/pipeline.py b/src/pykeen/pipeline.py index bddc96a13a..19bcc91ce3 100644 --- a/src/pykeen/pipeline.py +++ b/src/pykeen/pipeline.py @@ -738,6 +738,7 @@ def pipeline( # noqa: C901 result_tracker: Union[None, str, Type[ResultTracker]] = None, result_tracker_kwargs: Optional[Mapping[str, Any]] = None, # Misc + automatic_memory_optimization: bool = True, metadata: Optional[Dict[str, Any]] = None, device: Union[None, str, torch.device] = None, random_seed: Optional[int] = None, @@ -913,6 +914,7 @@ def pipeline( # noqa: C901 training_loop_instance: TrainingLoop = training_loop( model=model_instance, optimizer=optimizer_instance, + automatic_memory_optimization=automatic_memory_optimization, ) elif training_loop is not SLCWATrainingLoop: raise ValueError('Can not specify negative sampler with LCWA') @@ -925,14 +927,20 @@ def pipeline( # noqa: C901 training_loop_instance: TrainingLoop = SLCWATrainingLoop( model=model_instance, optimizer=optimizer_instance, + automatic_memory_optimization=automatic_memory_optimization, negative_sampler_cls=negative_sampler, negative_sampler_kwargs=negative_sampler_kwargs, ) evaluator = get_evaluator_cls(evaluator) - evaluator_instance: Evaluator = evaluator( - **(evaluator_kwargs or {}), - ) + # TODO @mehdi is setting the automatic memory optimization as an attribute + # of the class appropriate, since it doesn't cause any state to be stored? + # I think it might be better to have this as an argument to the + # Evaluator.evaluate() function instead + if evaluator_kwargs is None: + evaluator_kwargs = {} + evaluator_kwargs.setdefault('automatic_memory_optimization', automatic_memory_optimization) + evaluator_instance: Evaluator = evaluator(**evaluator_kwargs) if evaluation_kwargs is None: evaluation_kwargs = {} diff --git a/src/pykeen/training/slcwa.py b/src/pykeen/training/slcwa.py index e3b9727359..e878a1688e 100644 --- a/src/pykeen/training/slcwa.py +++ b/src/pykeen/training/slcwa.py @@ -35,6 +35,7 @@ def __init__( optimizer: Optional[Optimizer] = None, negative_sampler_cls: Optional[Type[NegativeSampler]] = None, negative_sampler_kwargs: Optional[Mapping[str, Any]] = None, + automatic_memory_optimization: bool = True, ): """Initialize the training loop. @@ -43,10 +44,14 @@ def __init__( :param negative_sampler_cls: The class of the negative sampler :param negative_sampler_kwargs: Keyword arguments to pass to the negative sampler class on instantiation for every positive one + :param automatic_memory_optimization: + Whether to automatically optimize the sub-batch size during + training and batch size during evaluation with regards to the hardware at hand. """ super().__init__( model=model, optimizer=optimizer, + automatic_memory_optimization=automatic_memory_optimization, ) if negative_sampler_cls is None: diff --git a/src/pykeen/training/training_loop.py b/src/pykeen/training/training_loop.py index c60a71d95a..ff40daa849 100644 --- a/src/pykeen/training/training_loop.py +++ b/src/pykeen/training/training_loop.py @@ -79,16 +79,21 @@ def __init__( self, model: Model, optimizer: Optional[Optimizer] = None, + automatic_memory_optimization: bool = True, ) -> None: """Initialize the training loop. :param model: The model to train :param optimizer: The optimizer to use while training the model + :param automatic_memory_optimization: bool + Whether to automatically optimize the sub-batch size during + training and batch size during evaluation with regards to the hardware at hand. """ self.model = model self.optimizer = optimizer self.training_instances = None self.losses_per_epochs = [] + self.automatic_memory_optimization = automatic_memory_optimization if self.loss_blacklist and isinstance(self.model.loss, tuple(self.loss_blacklist)): raise TrainingApproachLossMismatchError( @@ -264,13 +269,13 @@ def _train( # noqa: C901 # Take the biggest possible training batch_size, if batch_size not set batch_size_sufficient = False if batch_size is None: - if self.model.automatic_memory_optimization: + if self.automatic_memory_optimization: batch_size, batch_size_sufficient = self.batch_size_search() else: batch_size = 256 # This will find necessary parameters to optimize the use of the hardware at hand - if not only_size_probing and self.model.automatic_memory_optimization and not batch_size_sufficient: + if not only_size_probing and self.automatic_memory_optimization and not batch_size_sufficient: # return the relevant parameters slice_size and batch_size sub_batch_size, slice_size = self.sub_batch_and_slice(batch_size) diff --git a/tests/test_early_stopping.py b/tests/test_early_stopping.py index 2a8395e731..8a336de760 100644 --- a/tests/test_early_stopping.py +++ b/tests/test_early_stopping.py @@ -55,8 +55,8 @@ def test_is_improvement(self): class MockEvaluator(Evaluator): """A mock evaluator for testing early stopping.""" - def __init__(self, losses: Iterable[float]) -> None: - super().__init__() + def __init__(self, losses: Iterable[float], automatic_memory_optimization: bool = True) -> None: + super().__init__(automatic_memory_optimization=automatic_memory_optimization) self.losses = tuple(losses) self.losses_iter = iter(self.losses) @@ -150,10 +150,10 @@ class TestEarlyStopping(unittest.TestCase): def setUp(self): """Prepare for testing the early stopper.""" - self.mock_evaluator = MockEvaluator(self.mock_losses) + self.mock_evaluator = MockEvaluator(self.mock_losses, automatic_memory_optimization=True) # Set automatic_memory_optimization to false for tests nations = Nations() - self.model = MockModel(triples_factory=nations.training, automatic_memory_optimization=False) + self.model = MockModel(triples_factory=nations.training) self.stopper = EarlyStopper( model=self.model, evaluator=self.mock_evaluator, @@ -237,8 +237,8 @@ def test_early_stopping(self): """Tests early stopping.""" # Set automatic_memory_optimization to false during testing nations = Nations() - model: Model = TransE(triples_factory=nations.training, automatic_memory_optimization=False) - evaluator = RankBasedEvaluator() + model: Model = TransE(triples_factory=nations.training) + evaluator = RankBasedEvaluator(automatic_memory_optimization=False) stopper = EarlyStopper( model=model, evaluator=evaluator, @@ -250,6 +250,7 @@ def test_early_stopping(self): training_loop = SLCWATrainingLoop( model=model, optimizer=Adam(params=model.get_grad_params()), + automatic_memory_optimization=False, ) losses = training_loop.train( num_epochs=self.max_num_epochs, diff --git a/tests/test_evaluators.py b/tests/test_evaluators.py index 1309d0e2e0..f3b8ee696e 100644 --- a/tests/test_evaluators.py +++ b/tests/test_evaluators.py @@ -395,8 +395,8 @@ def test_filter_corrupted_triples(self): class DummyEvaluator(Evaluator): """A dummy evaluator for testing the structure of the evaluation function.""" - def __init__(self, *, counter: int, filtered: bool) -> None: - super().__init__(filtered=filtered) + def __init__(self, *, counter: int, filtered: bool, automatic_memory_optimization: bool = True) -> None: + super().__init__(filtered=filtered, automatic_memory_optimization=automatic_memory_optimization) self.counter = counter def process_tail_scores_( @@ -437,7 +437,7 @@ def setUp(self): self.counter = 1337 self.evaluator = DummyEvaluator(counter=self.counter, filtered=True) self.triples_factory = Nations().training - self.model = MockModel(triples_factory=self.triples_factory, automatic_memory_optimization=False) + self.model = MockModel(triples_factory=self.triples_factory) def test_evaluation_structure(self): """Test if the evaluator has a balanced call of head and tail processors.""" diff --git a/tests/test_models.py b/tests/test_models.py index 3f48252901..2449dec1a0 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -102,6 +102,9 @@ class _ModelTestCase: #: Additional arguments passed to the model's constructor method model_kwargs: ClassVar[Optional[Mapping[str, Any]]] = None + #: Additional arguments passed to the training loop's constructor method + training_loop_kwargs: ClassVar[Optional[Mapping[str, Any]]] = None + #: The triples factory instance factory: TriplesFactory @@ -255,6 +258,7 @@ def test_train_slcwa(self) -> None: loop = SLCWATrainingLoop( model=self.model, optimizer=Adagrad(params=self.model.get_grad_params(), lr=0.001), + **(self.training_loop_kwargs or {}), ) losses = self._safe_train_loop( loop, @@ -270,6 +274,7 @@ def test_train_lcwa(self) -> None: loop = LCWATrainingLoop( model=self.model, optimizer=Adagrad(params=self.model.get_grad_params(), lr=0.001), + **(self.training_loop_kwargs or {}), ) losses = self._safe_train_loop( loop, @@ -335,6 +340,16 @@ def cli_extras(self): for k, v in kwargs.items(): extras.append('--' + k.replace('_', '-')) extras.append(str(v)) + + # For the high/low memory test cases of NTN, SE, etc. + if self.training_loop_kwargs and 'automatic_memory_optimization' in self.training_loop_kwargs: + automatic_memory_optimization = self.training_loop_kwargs.get('automatic_memory_optimization') + if automatic_memory_optimization is True: + extras.append('--automatic-memory-optimization') + elif automatic_memory_optimization is False: + extras.append('--no-automatic-memory-optimization') + # else, leave to default + extras += [ '--number-epochs', self.train_num_epochs, '--embedding-dim', self.embedding_dim, @@ -711,6 +726,9 @@ class TestNTNLowMemory(_BaseNTNTest): model_kwargs = { 'num_slices': 2, + } + + training_loop_kwargs = { 'automatic_memory_optimization': True, } @@ -720,6 +738,9 @@ class TestNTNHighMemory(_BaseNTNTest): model_kwargs = { 'num_slices': 2, + } + + training_loop_kwargs = { 'automatic_memory_optimization': False, } @@ -816,17 +837,17 @@ def _check_constraints(self): class TestSELowMemory(_BaseTestSE): """Tests SE with low memory.""" - model_kwargs = dict( - automatic_memory_optimization=True, - ) + training_loop_kwargs = { + 'automatic_memory_optimization': True, + } class TestSEHighMemory(_BaseTestSE): """Tests SE with low memory.""" - model_kwargs = dict( - automatic_memory_optimization=False, - ) + training_loop_kwargs = { + 'automatic_memory_optimization': False, + } class TestTransD(_DistanceModelTestCase, unittest.TestCase): diff --git a/tests/test_training.py b/tests/test_training.py index 179148ea16..0a2186669a 100644 --- a/tests/test_training.py +++ b/tests/test_training.py @@ -20,8 +20,12 @@ class DummyTrainingLoop(SLCWATrainingLoop): """A wrapper around SLCWATrainingLoop.""" - def __init__(self, model: Model, sub_batch_size: int): - super().__init__(model=model, optimizer=optim.Adam(lr=1.0, params=model.parameters())) + def __init__(self, model: Model, sub_batch_size: int, automatic_memory_optimization: bool = False): + super().__init__( + model=model, + optimizer=optim.Adam(lr=1.0, params=model.parameters()), + automatic_memory_optimization=automatic_memory_optimization, + ) self.sub_batch_size = sub_batch_size def _process_batch( @@ -50,8 +54,12 @@ def _process_batch( class NaNTrainingLoop(SLCWATrainingLoop): """A wrapper around SLCWATrainingLoop returning NaN losses.""" - def __init__(self, model: Model, patience: int): - super().__init__(model=model, optimizer=optim.Adam(lr=1.0, params=model.parameters())) + def __init__(self, model: Model, patience: int, automatic_memory_optimization: bool = False): + super().__init__( + model=model, + optimizer=optim.Adam(lr=1.0, params=model.parameters()), + automatic_memory_optimization=automatic_memory_optimization, + ) self.patience = patience def _process_batch( @@ -89,14 +97,22 @@ def setUp(self) -> None: def test_sub_batching(self): """Test if sub-batching works as expected.""" - model = TransE(triples_factory=self.triples_factory, automatic_memory_optimization=False) - training_loop = DummyTrainingLoop(model=model, sub_batch_size=self.sub_batch_size) + model = TransE(triples_factory=self.triples_factory) + training_loop = DummyTrainingLoop( + model=model, + sub_batch_size=self.sub_batch_size, + automatic_memory_optimization=False, + ) training_loop.train(num_epochs=1, batch_size=self.batch_size, sub_batch_size=self.sub_batch_size) def test_sub_batching_support(self): """Test if sub-batching works as expected.""" - model = ConvE(triples_factory=self.triples_factory, automatic_memory_optimization=False) - training_loop = DummyTrainingLoop(model=model, sub_batch_size=self.sub_batch_size) + model = ConvE(triples_factory=self.triples_factory) + training_loop = DummyTrainingLoop( + model=model, + sub_batch_size=self.sub_batch_size, + automatic_memory_optimization=False, + ) def _try_train(): """Call train method.""" @@ -117,7 +133,6 @@ def test_blacklist_loss_on_slcwa(self): model = TransE( triples_factory=self.triples_factory, loss=CrossEntropyLoss(), - automatic_memory_optimization=False, ) with self.assertRaises(TrainingApproachLossMismatchError): - NaNTrainingLoop(model=model, patience=2) + NaNTrainingLoop(model=model, patience=2, automatic_memory_optimization=False) From bc7b1165d61e73ae004d2476ed6c5b75f786e56e Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 30 Nov 2020 12:59:14 +0100 Subject: [PATCH 628/690] Cleanup --- src/pykeen/evaluation/evaluator.py | 4 ++-- src/pykeen/nn/functional.py | 4 ++-- src/pykeen/training/training_loop.py | 3 +++ tests/test_early_stopping.py | 2 +- tests/test_evaluators.py | 2 +- 5 files changed, 9 insertions(+), 6 deletions(-) diff --git a/src/pykeen/evaluation/evaluator.py b/src/pykeen/evaluation/evaluator.py index 2b1f3e9ad2..6aa9bc39a7 100644 --- a/src/pykeen/evaluation/evaluator.py +++ b/src/pykeen/evaluation/evaluator.py @@ -65,8 +65,8 @@ def __init__( self, filtered: bool = False, requires_positive_mask: bool = False, - batch_size: int = None, - slice_size: int = None, + batch_size: Optional[int] = None, + slice_size: Optional[int] = None, automatic_memory_optimization: bool = True, ): self.filtered = filtered diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 53ca41784e..08152cc40c 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -12,8 +12,8 @@ from .sim import KG2E_SIMILARITIES from ..typing import GaussianDistribution from ..utils import ( - broadcast_cat, clamp_norm, estimate_cost_of_sequence, extended_einsum, is_cudnn_error, negative_norm, negative_norm_of_sum, - project_entity, split_complex, tensor_product, tensor_sum, view_complex, + broadcast_cat, clamp_norm, estimate_cost_of_sequence, extended_einsum, is_cudnn_error, negative_norm, + negative_norm_of_sum, project_entity, split_complex, tensor_product, tensor_sum, view_complex, ) __all__ = [ diff --git a/src/pykeen/training/training_loop.py b/src/pykeen/training/training_loop.py index ff40daa849..906976ebad 100644 --- a/src/pykeen/training/training_loop.py +++ b/src/pykeen/training/training_loop.py @@ -239,6 +239,9 @@ def _train( # noqa: C901 :param slice_size: >0 The divisor for the scoring function when using slicing. This is only possible for LCWA training loops in general and only for models that have the slicing capability implemented. + :param automatic_memory_optimization: bool + Whether to automatically optimize the sub-batch size during training and batch size during evaluation with + regards to the hardware at hand. :param label_smoothing: (0 <= label_smoothing < 1) If larger than zero, use label smoothing. :param sampler: (None or 'schlichtkrull') diff --git a/tests/test_early_stopping.py b/tests/test_early_stopping.py index 8a336de760..df28f71507 100644 --- a/tests/test_early_stopping.py +++ b/tests/test_early_stopping.py @@ -150,7 +150,7 @@ class TestEarlyStopping(unittest.TestCase): def setUp(self): """Prepare for testing the early stopper.""" - self.mock_evaluator = MockEvaluator(self.mock_losses, automatic_memory_optimization=True) + self.mock_evaluator = MockEvaluator(self.mock_losses, automatic_memory_optimization=False) # Set automatic_memory_optimization to false for tests nations = Nations() self.model = MockModel(triples_factory=nations.training) diff --git a/tests/test_evaluators.py b/tests/test_evaluators.py index f3b8ee696e..3175c0a050 100644 --- a/tests/test_evaluators.py +++ b/tests/test_evaluators.py @@ -435,7 +435,7 @@ class TestEvaluationStructure(unittest.TestCase): def setUp(self): """Prepare for testing the evaluation structure.""" self.counter = 1337 - self.evaluator = DummyEvaluator(counter=self.counter, filtered=True) + self.evaluator = DummyEvaluator(counter=self.counter, filtered=True, automatic_memory_optimization=False) self.triples_factory = Nations().training self.model = MockModel(triples_factory=self.triples_factory) From 8719441448fcbbc69f207e33b80af9e40b9cb1ec Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 30 Nov 2020 13:12:07 +0100 Subject: [PATCH 629/690] Pass lint --- MANIFEST.in | 2 ++ src/pykeen/nn/functional.py | 10 +++++----- 2 files changed, 7 insertions(+), 5 deletions(-) diff --git a/MANIFEST.in b/MANIFEST.in index 7f1b1fa299..604bb32838 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -3,6 +3,8 @@ graft tests prune docs/source/api prune notebooks +prune benchmarking +prune tests/mlruns recursive-include docs/source *.py recursive-include docs/source *.rst diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 08152cc40c..08d87d0502 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -103,11 +103,11 @@ def _complex_interaction_optimized_broadcasted( return sum(*( factor * tensor_product(hh, rr, tt).sum(dim=-1) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) From eea971e8df5effc86560e46271f9db5ffa11fbd2 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 30 Nov 2020 13:24:29 +0100 Subject: [PATCH 630/690] Fix ER --- src/pykeen/nn/sim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index cd89655ca0..0c8fc13bf6 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -139,7 +139,7 @@ def kullback_leibler_similarity( # 3. Component if exact: - terms.append(-h.mean.shape[-1]) + terms.append(-torch.as_tensor(data=[h.mean.shape[-1]])) # 4. Component # ln (det(\Sigma_1) / det(\Sigma_0)) From 3791a358be4a70d3fb7e8500e2f0060bb5b2cf61 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 30 Nov 2020 13:54:06 +0100 Subject: [PATCH 631/690] Adapt unittests for KL div similarity --- tests/test_nn.py | 93 +++++++++++++++++++++++++++++++----------------- 1 file changed, 60 insertions(+), 33 deletions(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index 34bded06b6..d3b307e3f0 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -12,7 +12,7 @@ from torch.nn import functional from pykeen.nn import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule -from pykeen.nn.representation import CANONICAL_DIMENSIONS, RGCNRepresentations, get_expected_canonical_shape +from pykeen.nn.representation import CANONICAL_DIMENSIONS, RGCNRepresentations, convert_to_canonical_shape, get_expected_canonical_shape from pykeen.nn.sim import kullback_leibler_similarity from pykeen.testing.base import GenericTests, TestsTest from pykeen.triples import TriplesFactory @@ -394,54 +394,81 @@ def test_make(self): class KullbackLeiblerTests(unittest.TestCase): """Tests for the vectorized computation of KL divergences.""" - d: int = 3 + batch_size: int = 2 + num_heads: int = 3 + num_relations: int = 5 + num_tails: int = 7 + d: int = 11 def setUp(self) -> None: # noqa: D102 - self.e_mean = torch.rand(self.d) - self.e_var = torch.rand(self.d).exp() - self.r_mean = torch.rand(self.d) - self.r_var = torch.rand(self.d).exp() - - def _get_e(self, pre_shape=(1, 1, 1)): - return GaussianDistribution( - mean=self.e_mean.view(*pre_shape, self.d), - diagonal_covariance=self.e_var.view(*pre_shape, self.d), - ) - - def _get_r(self, pre_shape=(1, 1)): - return GaussianDistribution( - mean=self.r_mean.view(*pre_shape, self.d), - diagonal_covariance=self.r_var.view(*pre_shape, self.d), - ) + dims = dict(h=self.num_heads, r=self.num_relations, t=self.num_tails) + (self.h_mean, self.r_mean, self.t_mean), (self.h_var, self.r_var, self.t_var) = [ + [ + convert_to_canonical_shape( + x=torch.rand(self.batch_size, num, self.d), + dim=dim, + num=num, + batch_size=self.batch_size, + ) + for dim, num in dims.items() + ] + for _ in ("mean", "diagonal_covariance") + ] + # ensure positivity + self.h_var, self.r_var, self.t_var = [x.exp() for x in (self.h_var, self.r_var, self.t_var)] + + def _get(self, name: str): + if name == "h": + mean, var = self.h_mean, self.h_var + elif name == "r": + mean, var = self.r_mean, self.r_var + elif name == "t": + mean, var = self.t_mean, self.t_var + elif name == "e": + mean, var = self.h_mean - self.t_mean, self.h_var + self.t_var + else: + raise ValueError + return GaussianDistribution(mean=mean, diagonal_covariance=var) def test_against_torch_builtin(self): """Compare value against torch.distributions.""" - # r: (batch_size, num_heads, num_tails, d) - e = self._get_e() - # r: (batch_size, num_relations, d) - r = self._get_r() - sim = kullback_leibler_similarity(e=e, r=r, exact=True).view(-1) + # compute using pykeen + sim = kullback_leibler_similarity( + h=self._get(name="h"), + r=self._get(name="r"), + t=self._get(name="t"), + exact=True, + ) - p = torch.distributions.MultivariateNormal(loc=self.e_mean, covariance_matrix=torch.diag(self.e_var)) - q = torch.distributions.MultivariateNormal(loc=self.r_mean, covariance_matrix=torch.diag(self.r_var)) - sim2 = -torch.distributions.kl_divergence(p=p, q=q).view(-1) - assert torch.allclose(sim, sim2) + # compute using pytorch + sim2 = torch.empty_like(sim) + for bi, hi, ri, ti in itertools.product(range(self.batch_size), range(self.num_heads), range(self.num_relations), range(self.num_tails)): + # prepare distributions + e_mean = self.h_mean[bi, hi, 0, 0, :] - self.t_mean[bi, 0, 0, ti, :] + e_var = torch.diag(self.h_var[bi, hi, 0, 0, :] + self.t_var[bi, 0, 0, ti, :]) + r_mean = self.r_mean[bi, 0, ri, 0, :] + r_var = torch.diag(self.r_var[bi, 0, ri, 0, :]) + p = torch.distributions.MultivariateNormal(loc=e_mean, covariance_matrix=e_var) + q = torch.distributions.MultivariateNormal(loc=r_mean, covariance_matrix=r_var) + sim2[bi, hi, ri, ti] = -torch.distributions.kl_divergence(p=p, q=q).view(-1) + assert torch.allclose(sim, sim2), (sim - sim2).abs() def test_self_similarity(self): """Check value of similarity to self.""" # e: (batch_size, num_heads, num_tails, d) # https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence#Properties # divergence = 0 => similarity = -divergence = 0 - e = self._get_e() - r = self._get_e(pre_shape=(1, 1)) - sim = kullback_leibler_similarity(e=e, r=r, exact=True) + # (h - t), r + r = self._get(name="r") + h = GaussianDistribution(mean=2 * r.mean, diagonal_covariance=0.5 * r.diagonal_covariance) + t = GaussianDistribution(mean=r.mean, diagonal_covariance=0.5 * r.diagonal_covariance) + sim = kullback_leibler_similarity(h=h, r=r, t=t, exact=True) assert torch.allclose(sim, torch.zeros_like(sim)) def test_value_range(self): """Check the value range.""" # https://en.wikipedia.org/wiki/Kullback%E2%80%93Leibler_divergence#Properties # divergence >= 0 => similarity = -divergence <= 0 - e = self._get_e() - r = self._get_r() - sim = kullback_leibler_similarity(e=e, r=r, exact=True) + h, r, t = [self._get(name=name) for name in "hrt"] + sim = kullback_leibler_similarity(h=h, r=r, t=t, exact=True) assert (sim <= 0).all() From 46d80966f5148a29c7d688e9be773ac71d5ef24a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 30 Nov 2020 14:04:45 +0100 Subject: [PATCH 632/690] add test in test --- tests/test_nn.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/test_nn.py b/tests/test_nn.py index d3b307e3f0..cfc6191b99 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -448,6 +448,11 @@ def test_against_torch_builtin(self): e_var = torch.diag(self.h_var[bi, hi, 0, 0, :] + self.t_var[bi, 0, 0, ti, :]) r_mean = self.r_mean[bi, 0, ri, 0, :] r_var = torch.diag(self.r_var[bi, 0, ri, 0, :]) + # check for correct slicing + assert e_mean.shape == (self.d,) + assert r_mean.shape == (self.d,) + assert e_var.shape == (self.d, self.d) + assert r_var.shape == (self.d, self.d) p = torch.distributions.MultivariateNormal(loc=e_mean, covariance_matrix=e_var) q = torch.distributions.MultivariateNormal(loc=r_mean, covariance_matrix=r_var) sim2[bi, hi, ri, ti] = -torch.distributions.kl_divergence(p=p, q=q).view(-1) From 58128f82d035fd3312c8b18ef4b2e4b9c4eb30a8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 30 Nov 2020 14:10:49 +0100 Subject: [PATCH 633/690] reformat --- tests/test_nn.py | 27 ++++++++++++++++----------- 1 file changed, 16 insertions(+), 11 deletions(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index cfc6191b99..5a08f5b22c 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -441,20 +441,25 @@ def test_against_torch_builtin(self): ) # compute using pytorch + e_mean = self.h_mean - self.t_mean + e_var = self.h_var + self.t_var + r_mean, r_var = self.r_var, self.r_mean + assert (e_var > 0).all() sim2 = torch.empty_like(sim) for bi, hi, ri, ti in itertools.product(range(self.batch_size), range(self.num_heads), range(self.num_relations), range(self.num_tails)): # prepare distributions - e_mean = self.h_mean[bi, hi, 0, 0, :] - self.t_mean[bi, 0, 0, ti, :] - e_var = torch.diag(self.h_var[bi, hi, 0, 0, :] + self.t_var[bi, 0, 0, ti, :]) - r_mean = self.r_mean[bi, 0, ri, 0, :] - r_var = torch.diag(self.r_var[bi, 0, ri, 0, :]) - # check for correct slicing - assert e_mean.shape == (self.d,) - assert r_mean.shape == (self.d,) - assert e_var.shape == (self.d, self.d) - assert r_var.shape == (self.d, self.d) - p = torch.distributions.MultivariateNormal(loc=e_mean, covariance_matrix=e_var) - q = torch.distributions.MultivariateNormal(loc=r_mean, covariance_matrix=r_var) + e_loc = e_mean[bi, hi, 0, ti, :] + r_loc = r_mean[bi, 0, ri, 0, :] + e_cov = torch.diag(e_var[bi, hi, 0, ti, :]) + r_cov = torch.diag(r_var[bi, 0, ri, 0, :]) + p = torch.distributions.MultivariateNormal( + loc=e_loc, + covariance_matrix=e_cov, + ) + q = torch.distributions.MultivariateNormal( + loc=r_loc, + covariance_matrix=r_cov, + ) sim2[bi, hi, ri, ti] = -torch.distributions.kl_divergence(p=p, q=q).view(-1) assert torch.allclose(sim, sim2), (sim - sim2).abs() From 83363aa27d7e8b62076ab8f8e4b33b539448be9a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 30 Nov 2020 14:11:02 +0100 Subject: [PATCH 634/690] reformat --- tests/test_nn.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index 5a08f5b22c..246d5351cd 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -446,7 +446,12 @@ def test_against_torch_builtin(self): r_mean, r_var = self.r_var, self.r_mean assert (e_var > 0).all() sim2 = torch.empty_like(sim) - for bi, hi, ri, ti in itertools.product(range(self.batch_size), range(self.num_heads), range(self.num_relations), range(self.num_tails)): + for bi, hi, ri, ti in itertools.product( + range(self.batch_size), + range(self.num_heads), + range(self.num_relations), + range(self.num_tails), + ): # prepare distributions e_loc = e_mean[bi, hi, 0, ti, :] r_loc = r_mean[bi, 0, ri, 0, :] From 2023a889b848a2988cdd52432bf30e816ae7df44 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 30 Nov 2020 14:22:29 +0100 Subject: [PATCH 635/690] Move debug implementation to sim.py --- src/pykeen/nn/sim.py | 62 ++++++++++++++++++++++++++++++++++++++------ tests/test_nn.py | 23 ++++++++-------- 2 files changed, 65 insertions(+), 20 deletions(-) diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index 0c8fc13bf6..a598062dba 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -1,14 +1,14 @@ # -*- coding: utf-8 -*- """Similarity functions.""" - +import itertools import math import torch from .compute_kernel import batched_dot from ..typing import GaussianDistribution -from ..utils import tensor_sum +from ..utils import calculate_broadcasted_elementwise_result_shape, tensor_sum __all__ = [ 'expected_likelihood', @@ -115,19 +115,26 @@ def kullback_leibler_similarity( The similarity. """ assert all((d.diagonal_covariance > 0).all() for d in (h, r, t)) + return _vectorized_kl_similarity(h=h, r=r, t=t, epsilon=epsilon, exact=exact) + +def _vectorized_kl_similarity( + h: GaussianDistribution, + r: GaussianDistribution, + t: GaussianDistribution, + epsilon: float = 1.0e-10, + exact: bool = True, +) -> torch.FloatTensor: + """Vectorized implementation.""" e_var = (h.diagonal_covariance + t.diagonal_covariance) r_var_safe = r.diagonal_covariance.clamp_min(min=epsilon) - terms = [] - # 1. Component # tr(sigma_1^-1 sigma_0) = sum (sigma_1^-1 sigma_0)[i, i] # since sigma_0, sigma_1 are diagonal matrices: # = sum (sigma_1^-1[i] sigma_0[i]) = sum (sigma_0[i] / sigma_1[i]) var_safe_reciprocal = r_var_safe.reciprocal() terms.append(batched_dot(e_var, var_safe_reciprocal)) - # 2. Component # (mu_1 - mu_0) * Sigma_1^-1 (mu_1 - mu_0) # with mu = (mu_1 - mu_0) @@ -136,11 +143,9 @@ def kullback_leibler_similarity( # = mu**2 / sigma_1 mu = tensor_sum(r.mean, -h.mean, t.mean) terms.append(batched_dot(mu.pow(2), r_var_safe)) - # 3. Component if exact: terms.append(-torch.as_tensor(data=[h.mean.shape[-1]])) - # 4. Component # ln (det(\Sigma_1) / det(\Sigma_0)) # = ln det Sigma_1 - ln det Sigma_0 @@ -152,11 +157,52 @@ def kullback_leibler_similarity( r_var_safe.log().sum(dim=-1), -e_var_safe.log().sum(dim=-1), )) - result = tensor_sum(*terms) if exact: result = 0.5 * result + return -result + +def _torch_kl_similarity( + h: GaussianDistribution, + r: GaussianDistribution, + t: GaussianDistribution, +) -> torch.FloatTensor: + """ + Implementation delegating to torch.distributions. + + .. note :: + Do not use this method in production code. + """ + e_mean = h.mean - t.mean + e_var = h.diagonal_covariance + t.diagonal_covariance + + # allocate result + batch_size, num_heads, num_relations, num_tails = calculate_broadcasted_elementwise_result_shape( + e_mean.shape, + r.mean.shape + )[:-1] + result = h.mean.new_empty(batch_size, num_heads, num_relations, num_tails) + for bi, hi, ri, ti in itertools.product( + range(batch_size), + range(num_heads), + range(num_relations), + range(num_tails), + ): + # prepare distributions + e_loc = e_mean[bi, hi, 0, ti, :] + r_loc = r.mean[bi, 0, ri, 0, :] + e_cov = torch.diag(e_var[bi, hi, 0, ti, :]) + r_cov = torch.diag(r.diagonal_covariance[bi, 0, ri, 0, :]) + p = torch.distributions.MultivariateNormal( + loc=e_loc, + covariance_matrix=e_cov, + ) + q = torch.distributions.MultivariateNormal( + loc=r_loc, + covariance_matrix=r_cov, + ) + result[bi, hi, ri, ti] = torch.distributions.kl_divergence(p=p, q=q).view(-1) return -result diff --git a/tests/test_nn.py b/tests/test_nn.py index 246d5351cd..1f8c04568f 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -13,7 +13,7 @@ from pykeen.nn import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule from pykeen.nn.representation import CANONICAL_DIMENSIONS, RGCNRepresentations, convert_to_canonical_shape, get_expected_canonical_shape -from pykeen.nn.sim import kullback_leibler_similarity +from pykeen.nn.sim import _torch_kl_similarity, kullback_leibler_similarity from pykeen.testing.base import GenericTests, TestsTest from pykeen.triples import TriplesFactory from pykeen.typing import GaussianDistribution @@ -430,22 +430,13 @@ def _get(self, name: str): raise ValueError return GaussianDistribution(mean=mean, diagonal_covariance=var) - def test_against_torch_builtin(self): - """Compare value against torch.distributions.""" - # compute using pykeen - sim = kullback_leibler_similarity( - h=self._get(name="h"), - r=self._get(name="r"), - t=self._get(name="t"), - exact=True, - ) - + def _get_kl_similarity_torch(self): # compute using pytorch e_mean = self.h_mean - self.t_mean e_var = self.h_var + self.t_var r_mean, r_var = self.r_var, self.r_mean assert (e_var > 0).all() - sim2 = torch.empty_like(sim) + sim2 = torch.empty(self.batch_size, self.num_heads, self.num_relations, self.num_tails) for bi, hi, ri, ti in itertools.product( range(self.batch_size), range(self.num_heads), @@ -466,6 +457,14 @@ def test_against_torch_builtin(self): covariance_matrix=r_cov, ) sim2[bi, hi, ri, ti] = -torch.distributions.kl_divergence(p=p, q=q).view(-1) + return sim2 + + def test_against_torch_builtin(self): + """Compare value against torch.distributions.""" + # compute using pykeen + h, r, t = [self._get(name=name) for name in "hrt"] + sim = kullback_leibler_similarity(h=h, r=r, t=t, exact=True) + sim2 = _torch_kl_similarity(h=h, r=r, t=t) assert torch.allclose(sim, sim2), (sim - sim2).abs() def test_self_similarity(self): From 95b655d37ddf318fa11066666d2a89ab2285f838 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 30 Nov 2020 14:47:30 +0100 Subject: [PATCH 636/690] Fix bug in KL --- src/pykeen/nn/sim.py | 70 +++++++++++++++++++++++++++++++++++++------- 1 file changed, 60 insertions(+), 10 deletions(-) diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index a598062dba..a81704215d 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -115,26 +115,76 @@ def kullback_leibler_similarity( The similarity. """ assert all((d.diagonal_covariance > 0).all() for d in (h, r, t)) - return _vectorized_kl_similarity(h=h, r=r, t=t, epsilon=epsilon, exact=exact) + return -_vectorized_kl_divergence( + h=h, + r=r, + t=t, + epsilon=epsilon, + exact=exact, + ) -def _vectorized_kl_similarity( +def _vectorized_kl_divergence( h: GaussianDistribution, r: GaussianDistribution, t: GaussianDistribution, epsilon: float = 1.0e-10, exact: bool = True, ) -> torch.FloatTensor: - """Vectorized implementation.""" + r""" + Vectorized implementation of KL-divergence. + + Computes the divergence between :math:`\mathcal{N}(\mu_e, \Sigma_e)` and :math:`\mathcal{N}(\mu_r, \Sigma_r)` + given by + + .. math :: + \mu_e = \mu_h - \mu_t + + \Sigma_e = \Sigma_h + \Sigma_t + + where all covariance matrices are diagonal. Hence we can simplify + + .. math :: + D(\mathcal{N}(\mu_e, \Sigma_e), \mathcal{N}(\mu_r, \Sigma_r)) + = + 0.5 * ( + \trace(\Sigma_r^-1 \Sigma_e) + + (\mu_r - \mu_e) * \Sigma_r^-1 (\mu_r - \mu_e) + - k + + \ln (\det(\Sigma_r) / \det(\Sigma_e)) + ) + = 0.5 * ( + \sum_i \Sigma_e[i] / Sigma_r[i] + + \sum_i \mu[i]^2 / \Sigma_r[i] + + \sum_i \ln Sigma_r[i] + - \sum_i \ln Sigma_e[i] + - k + ) + + where :math:`\mu = \mu_r - \mu_e = \mu_r - \mu_h + \mu_t` + + :param h: shape: (batch_size, num_heads, 1, 1, d) + The head entity Gaussian distribution. + :param r: shape: (batch_size, 1, num_relations, 1, d) + The relation Gaussian distribution. + :param t: shape: (batch_size, 1, 1, num_tails, d) + The tail entity Gaussian distribution. + :param epsilon: float (default=1.0) + Small constant used to avoid numerical issues when dividing. + :param exact: + Whether to return the exact similarity, or leave out constant offsets. + + :return: torch.Tensor, shape: (s_1, ..., s_k) + The KL-divergence. + """ + num_heads, num_relations, num_tails = [x.mean.shape[d] for d, x in enumerate((h, r, t))] e_var = (h.diagonal_covariance + t.diagonal_covariance) r_var_safe = r.diagonal_covariance.clamp_min(min=epsilon) terms = [] # 1. Component - # tr(sigma_1^-1 sigma_0) = sum (sigma_1^-1 sigma_0)[i, i] - # since sigma_0, sigma_1 are diagonal matrices: - # = sum (sigma_1^-1[i] sigma_0[i]) = sum (sigma_0[i] / sigma_1[i]) - var_safe_reciprocal = r_var_safe.reciprocal() - terms.append(batched_dot(e_var, var_safe_reciprocal)) + # \sum_i \Sigma_e[i] / Sigma_r[i] + r_var_safe_reciprocal = r_var_safe.reciprocal() + terms.append(batched_dot(e_var, r_var_safe_reciprocal)) # 2. Component # (mu_1 - mu_0) * Sigma_1^-1 (mu_1 - mu_0) # with mu = (mu_1 - mu_0) @@ -142,7 +192,7 @@ def _vectorized_kl_similarity( # since Sigma_1 is diagonal # = mu**2 / sigma_1 mu = tensor_sum(r.mean, -h.mean, t.mean) - terms.append(batched_dot(mu.pow(2), r_var_safe)) + terms.append(batched_dot(mu.pow(2), r_var_safe_reciprocal)) # 3. Component if exact: terms.append(-torch.as_tensor(data=[h.mean.shape[-1]])) @@ -160,7 +210,7 @@ def _vectorized_kl_similarity( result = tensor_sum(*terms) if exact: result = 0.5 * result - return -result + return result def _torch_kl_similarity( From 01af212b52fc10dc359622b0e519625e45ae0efa Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 30 Nov 2020 14:47:48 +0100 Subject: [PATCH 637/690] Remove unused variables --- src/pykeen/nn/sim.py | 1 - 1 file changed, 1 deletion(-) diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index a81704215d..452a8f5fb0 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -177,7 +177,6 @@ def _vectorized_kl_divergence( :return: torch.Tensor, shape: (s_1, ..., s_k) The KL-divergence. """ - num_heads, num_relations, num_tails = [x.mean.shape[d] for d, x in enumerate((h, r, t))] e_var = (h.diagonal_covariance + t.diagonal_covariance) r_var_safe = r.diagonal_covariance.clamp_min(min=epsilon) terms = [] From aac4bf592356b41e3eab5866e35d420c6a040ea7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 30 Nov 2020 14:51:09 +0100 Subject: [PATCH 638/690] Fix flake8 --- src/pykeen/nn/sim.py | 4 ++-- tests/test_nn.py | 4 +++- 2 files changed, 5 insertions(+), 3 deletions(-) diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index 452a8f5fb0..5d9d66442b 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -218,7 +218,7 @@ def _torch_kl_similarity( t: GaussianDistribution, ) -> torch.FloatTensor: """ - Implementation delegating to torch.distributions. + Compute KL similarity using torch.distributions. .. note :: Do not use this method in production code. @@ -229,7 +229,7 @@ def _torch_kl_similarity( # allocate result batch_size, num_heads, num_relations, num_tails = calculate_broadcasted_elementwise_result_shape( e_mean.shape, - r.mean.shape + r.mean.shape, )[:-1] result = h.mean.new_empty(batch_size, num_heads, num_relations, num_tails) for bi, hi, ri, ti in itertools.product( diff --git a/tests/test_nn.py b/tests/test_nn.py index 1f8c04568f..5f32589cdb 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -12,7 +12,9 @@ from torch.nn import functional from pykeen.nn import Embedding, EmbeddingSpecification, LiteralRepresentations, RepresentationModule -from pykeen.nn.representation import CANONICAL_DIMENSIONS, RGCNRepresentations, convert_to_canonical_shape, get_expected_canonical_shape +from pykeen.nn.representation import ( + CANONICAL_DIMENSIONS, RGCNRepresentations, convert_to_canonical_shape, get_expected_canonical_shape, +) from pykeen.nn.sim import _torch_kl_similarity, kullback_leibler_similarity from pykeen.testing.base import GenericTests, TestsTest from pykeen.triples import TriplesFactory From 3851f04af80f85e17fd863872756b6c380eca95a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 30 Nov 2020 17:46:33 +0100 Subject: [PATCH 639/690] Fix unittest --- tests/test_regularizers.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/tests/test_regularizers.py b/tests/test_regularizers.py index a5f1732898..e252821bfd 100644 --- a/tests/test_regularizers.py +++ b/tests/test_regularizers.py @@ -322,14 +322,16 @@ def test_collect_regularization_terms(self): PowerSumRegularizer(normalize=True), ] model = ERModel( - triples_factory=MagicMock(), - interaction=MagicMock(), + triples_factory=MagicMock(num_entities=3, num_relations=2), + interaction=MagicMock(relation_shape=("d",), entity_shape=("d",)), entity_representations=EmbeddingSpecification( regularizer=regularizers[0], - ).make(num_embeddings=3, embedding_dim=2, shape=None), + embedding_dim=2, + ), relation_representations=EmbeddingSpecification( regularizer=regularizers[1], - ).make(num_embeddings=3, embedding_dim=2, shape=None), + embedding_dim=2, + ), ) # add weighted modules From 08c68e6605d7172d2f9e5208d2fa258eae2301f7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 30 Nov 2020 17:48:20 +0100 Subject: [PATCH 640/690] Add skip_cls --- tests/test_nn.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/test_nn.py b/tests/test_nn.py index 5f32589cdb..bbd9e9593f 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -17,6 +17,7 @@ ) from pykeen.nn.sim import _torch_kl_similarity, kullback_leibler_similarity from pykeen.testing.base import GenericTests, TestsTest +from pykeen.testing.mocks import MockRepresentations from pykeen.triples import TriplesFactory from pykeen.typing import GaussianDistribution @@ -352,6 +353,7 @@ class RepresentationModuleTestsTest(TestsTest[RepresentationModule], unittest.Te base_cls = RepresentationModule base_test = RepresentationModuleTests + skip_cls = {MockRepresentations} class EmbeddingSpecificationTests(unittest.TestCase): From acda8e2638fbf87278108ee417c25bf1b12b9543 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 30 Nov 2020 18:25:00 +0100 Subject: [PATCH 641/690] Fix unittest for RotatEInteraction --- tests/test_interactions.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 6d02043a35..9547429d55 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -449,11 +449,20 @@ class RotatETests(InteractionTests, unittest.TestCase): cls = pykeen.nn.modules.RotatEInteraction + def _get_hrt(self, *shapes): # noqa: D102 + # normalize length of r + h, r, t = super()._get_hrt(*shapes) + rc = view_complex(r) + rl = (rc.abs() ** 2).sum(dim=-1).sqrt() + r = r / rl.unsqueeze(dim=-1) + return h, r, t + def _exp_score(self, h, r, t) -> torch.FloatTensor: # noqa: D102 h, r, t = strip_dim(*(view_complex(x) for x in (h, r, t))) - hr = h * r - d = hr - t - return -(d.abs() ** 2).sum().sqrt() + # check for unit length + assert torch.allclose((r.abs() ** 2).sum(dim=-1).sqrt(), torch.ones(1)) + d = h * r - t + return -(d.abs() ** 2).sum(dim=-1).sqrt() class TranslationalInteractionTests(InteractionTests): From 79a241038dd6f70e5e972ee4b26bef945d528362 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Mon, 30 Nov 2020 18:43:09 +0100 Subject: [PATCH 642/690] circumvent accumulate(..., initial=x) since it is unavailable in Python 3.7 --- src/pykeen/utils.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index ac1103d244..cb79f27e7c 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -527,9 +527,8 @@ def estimate_cost_of_sequence( numpy.prod, itertools.islice( itertools.accumulate( - other_shapes, + (shape,) + other_shapes, calculate_broadcasted_elementwise_result_shape, - initial=shape, ), 1, None, From a42ce997f90d937db0ac28380b77ca23e9760c96 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Mon, 30 Nov 2020 19:54:58 +0100 Subject: [PATCH 643/690] More unittesting hopefully this gives some better logging --- src/pykeen/models/__init__.py | 11 +++------ tests/test_interactions.py | 6 ++--- tests/test_models.py | 18 ++++++++------ tests/test_nn.py | 45 ++++++++++++++++++++--------------- 4 files changed, 43 insertions(+), 37 deletions(-) diff --git a/src/pykeen/models/__init__.py b/src/pykeen/models/__init__.py index 1763737405..7189a695c8 100644 --- a/src/pykeen/models/__init__.py +++ b/src/pykeen/models/__init__.py @@ -8,8 +8,8 @@ from typing import Mapping, Set, Type, Union -from .base import ERModel, Model -from .multimodal import ComplExLiteral, DistMultLiteral, LiteralModel +from .base import ERModel, Model # noqa:F401 +from .multimodal import ComplExLiteral, DistMultLiteral, LiteralModel # noqa:F401 from .unimodal import ( ComplEx, ConvE, @@ -63,15 +63,10 @@ 'get_model_cls', ] -_BASE_MODELS = { - ERModel, - LiteralModel, -} - def _concrete_subclasses(cls: Type[Model]): for subcls in cls.__subclasses__(): - if not subcls._is_base_model and subcls not in _BASE_MODELS: + if not subcls._is_base_model: yield subcls yield from _concrete_subclasses(subcls) diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 9547429d55..fba5bc26d0 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -14,14 +14,14 @@ import pykeen.nn.modules from pykeen.nn.functional import distmult_interaction from pykeen.nn.modules import Interaction, TranslationalInteraction -from pykeen.testing.base import GenericTests, TestsTest +from pykeen.testing import base as ptb from pykeen.typing import Representation from pykeen.utils import clamp_norm, project_entity, strip_dim, view_complex logger = logging.getLogger(__name__) -class InteractionTests(GenericTests[pykeen.nn.modules.Interaction]): +class InteractionTests(ptb.GenericTests[pykeen.nn.modules.Interaction]): """Generic test for interaction functions.""" dim: int = 2 @@ -605,7 +605,7 @@ def _exp_score(self, h, r, t, h_inv, r_inv, t_inv, clamp) -> torch.FloatTensor: return 0.5 * distmult_interaction(h, r, t) + 0.5 * distmult_interaction(h_inv, r_inv, t_inv) -class InteractionTestsTest(TestsTest[Interaction], unittest.TestCase): +class InteractionTestsTest(ptb.TestsTest[Interaction], unittest.TestCase): """Test for tests for all interaction functions.""" base_cls = Interaction diff --git a/tests/test_models.py b/tests/test_models.py index 2449dec1a0..eacd6e665a 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -3,7 +3,6 @@ """Test that models can be executed.""" import importlib -import itertools as itt import os import tempfile import traceback @@ -23,7 +22,7 @@ import pykeen.models from pykeen.datasets.kinships import KINSHIPS_TRAIN_PATH from pykeen.datasets.nations import NATIONS_TEST_PATH, NATIONS_TRAIN_PATH, Nations -from pykeen.models import LiteralModel, _BASE_MODELS, _MODELS +from pykeen.models import _MODELS from pykeen.models.base import ( ERModel, Model, @@ -31,6 +30,7 @@ get_novelty_mask, ) from pykeen.models.cli import build_cli_from_cls +from pykeen.models.multimodal.base import LiteralModel from pykeen.nn.representation import ( RGCNRepresentations, inverse_indegree_edge_weights, inverse_outdegree_edge_weights, symmetric_edge_weights, @@ -43,13 +43,17 @@ SKIP_MODULES = { Model.__name__, + ERModel.__name__, + LiteralModel.__name__, 'DummyModel', 'MockModel', 'models', 'get_model_cls', } -for cls in itt.chain(_BASE_MODELS, LiteralModel.__subclasses__()): - SKIP_MODULES.add(cls.__name__) +SKIP_MODULES.update({ + cls.__name__ + for cls in LiteralModel.__subclasses__() +}) _EPSILON = 1.0e-07 @@ -935,7 +939,7 @@ class TestTesting(unittest.TestCase): def test_testing(self): """Check that there's a test for all models. - For now, this is excluding multimodel models. Not sure how to test those yet. + For now, this is excluding multimodal models. Not sure how to test those yet. """ model_names = { cls.__name__ @@ -1060,8 +1064,8 @@ class TestModelUtilities(unittest.TestCase): def test_abstract(self): """Test that classes are checked as abstract properly.""" - for model_cls in _BASE_MODELS: - self.assertTrue(model_cls._is_base_model) + self.assertTrue(ERModel._is_base_model) + self.assertTrue(LiteralModel._is_base_model) for model_cls in _MODELS: self.assertFalse( model_cls._is_base_model, diff --git a/tests/test_nn.py b/tests/test_nn.py index bbd9e9593f..e0f3519b28 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- """Unittest for the :mod:`pykeen.nn` module.""" + import itertools import unittest from typing import Any, Iterable, Mapping, MutableMapping, Optional, Sequence @@ -16,13 +17,13 @@ CANONICAL_DIMENSIONS, RGCNRepresentations, convert_to_canonical_shape, get_expected_canonical_shape, ) from pykeen.nn.sim import _torch_kl_similarity, kullback_leibler_similarity -from pykeen.testing.base import GenericTests, TestsTest +from pykeen.testing import base as ptb from pykeen.testing.mocks import MockRepresentations from pykeen.triples import TriplesFactory from pykeen.typing import GaussianDistribution -class RepresentationModuleTests(GenericTests[RepresentationModule]): +class RepresentationModuleTests(ptb.GenericTests[RepresentationModule]): """Tests for RepresentationModule.""" #: The batch size @@ -123,6 +124,7 @@ def test_get_in_canonical_shape_with_2d_indices(self): def _check_call( + self: unittest.TestCase, call_count: int, should_be_called: bool, wrapped: MagicMock, @@ -131,6 +133,8 @@ def _check_call( """ Check whether a wrapped method is called. + :param self: + The test cas calling the check :param call_count: The previous call count. :param should_be_called: @@ -146,15 +150,15 @@ def _check_call( if should_be_called: call_count += 1 - assert wrapped.call_count == call_count + self.assertEqual(wrapped.call_count, call_count) # called with one positional argument ... - assert len(wrapped.call_args.args) == 1 + self.assertEqual(1, len(wrapped.call_args.args)) # .. and additional key-word based arguments. - assert len(wrapped.call_args.kwargs) == len(kwargs or {}) + self.assertEqual(len(wrapped.call_args.kwargs), len(kwargs or {})) else: - assert wrapped.call_count == call_count + self.assertEqual(wrapped.call_count, call_count) return call_count @@ -210,6 +214,7 @@ def _test_func_with_kwargs( # check call in reset_parameters embedding.reset_parameters() call_count = _check_call( + self, call_count=call_count, should_be_called=reset_parameters_call, wrapped=wrapped, @@ -219,6 +224,7 @@ def _test_func_with_kwargs( # check call in forward embedding.forward(indices=None) call_count = _check_call( + self, call_count=call_count, should_be_called=forward_call, wrapped=wrapped, @@ -228,6 +234,7 @@ def _test_func_with_kwargs( # check call in post_parameter_update embedding.post_parameter_update() _check_call( + self, call_count=call_count, should_be_called=post_parameter_update_call, wrapped=wrapped, @@ -313,7 +320,7 @@ def _verify_content(self, x, indices): # noqa: D102 exp_x = self.numeric_literals if indices is not None: exp_x = exp_x[indices] - assert torch.allclose(x, exp_x) + self.assertTrue(torch.allclose(x, exp_x)) class RGCNRepresentationTests(RepresentationModuleTests, unittest.TestCase): @@ -348,7 +355,7 @@ def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMa return kwargs -class RepresentationModuleTestsTest(TestsTest[RepresentationModule], unittest.TestCase): +class RepresentationModuleTestsTest(ptb.TestsTest[RepresentationModule], unittest.TestCase): """Test that there are tests for all representation modules.""" base_cls = RepresentationModule @@ -384,15 +391,15 @@ def test_make(self): emb = spec.make(num_embeddings=self.num) # check shape - assert emb.embedding_dim == (embedding_dim or int(numpy.prod(shape))) - assert emb.shape == (shape or (embedding_dim,)) - assert emb.num_embeddings == self.num + self.assertEqual(emb.embedding_dim, (embedding_dim or int(numpy.prod(shape)))) + self.assertEqual(emb.shape, (shape or (embedding_dim,))) + self.assertEqual(emb.num_embeddings, self.num) # check attributes - assert emb.initializer is initializer - assert emb.normalizer is normalizer - assert emb.constrainer is constrainer - assert emb.regularizer is regularizer + self.assertIs(emb.initializer, initializer) + self.assertIs(emb.normalizer, normalizer) + self.assertIs(emb.constrainer, constrainer) + self.assertIs(emb.regularizer, regularizer) class KullbackLeiblerTests(unittest.TestCase): @@ -439,7 +446,7 @@ def _get_kl_similarity_torch(self): e_mean = self.h_mean - self.t_mean e_var = self.h_var + self.t_var r_mean, r_var = self.r_var, self.r_mean - assert (e_var > 0).all() + self.assertTrue((e_var > 0).all()) sim2 = torch.empty(self.batch_size, self.num_heads, self.num_relations, self.num_tails) for bi, hi, ri, ti in itertools.product( range(self.batch_size), @@ -469,7 +476,7 @@ def test_against_torch_builtin(self): h, r, t = [self._get(name=name) for name in "hrt"] sim = kullback_leibler_similarity(h=h, r=r, t=t, exact=True) sim2 = _torch_kl_similarity(h=h, r=r, t=t) - assert torch.allclose(sim, sim2), (sim - sim2).abs() + self.assertTrue(torch.allclose(sim, sim2), msg=f'Difference: {(sim - sim2).abs()}') def test_self_similarity(self): """Check value of similarity to self.""" @@ -481,7 +488,7 @@ def test_self_similarity(self): h = GaussianDistribution(mean=2 * r.mean, diagonal_covariance=0.5 * r.diagonal_covariance) t = GaussianDistribution(mean=r.mean, diagonal_covariance=0.5 * r.diagonal_covariance) sim = kullback_leibler_similarity(h=h, r=r, t=t, exact=True) - assert torch.allclose(sim, torch.zeros_like(sim)) + self.assertTrue(torch.allclose(sim, torch.zeros_like(sim)), msg=f'Sim: {sim}') def test_value_range(self): """Check the value range.""" @@ -489,4 +496,4 @@ def test_value_range(self): # divergence >= 0 => similarity = -divergence <= 0 h, r, t = [self._get(name=name) for name in "hrt"] sim = kullback_leibler_similarity(h=h, r=r, t=t, exact=True) - assert (sim <= 0).all() + self.assertTrue((sim <= 0).all()) From 4b8b65ab8146239cc245ca219b0fcda5b7d85ab1 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 1 Dec 2020 02:02:58 +0100 Subject: [PATCH 644/690] Fix mocking in test Apparently this relied on an unstable API and the structure of the unittest.mock._Call class, which subclasses tuple. I switched it to using index based nomenclature instead --- tests/test_nn.py | 22 +++++++++++++++++----- 1 file changed, 17 insertions(+), 5 deletions(-) diff --git a/tests/test_nn.py b/tests/test_nn.py index e0f3519b28..7cd89ed41c 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -150,15 +150,27 @@ def _check_call( if should_be_called: call_count += 1 - self.assertEqual(wrapped.call_count, call_count) + self.assertEqual(call_count, wrapped.call_count) - # called with one positional argument ... - self.assertEqual(1, len(wrapped.call_args.args)) + # Lets check the tuple + self.assertIsInstance(wrapped.call_args, tuple) + + call_size = len(wrapped.call_args) + # Make sure tuple at least has positional arguments, could be 3 if kwargs available + self.assertLessEqual(2, call_size) + if call_size == 2: + args_idx, kwargs_idx = 0, 1 + else: # call_size == 3: + args_idx, kwargs_idx = 1, 2 + + # called with one positional argument ... + self.assertEqual(1, len(wrapped.call_args[args_idx]), + msg=f'Args: {wrapped.call_args[args_idx]} Kwargs: {wrapped.call_args[kwargs_idx]}') # .. and additional key-word based arguments. - self.assertEqual(len(wrapped.call_args.kwargs), len(kwargs or {})) + self.assertEqual(len(kwargs or {}), len(wrapped.call_args[kwargs_idx])) else: - self.assertEqual(wrapped.call_count, call_count) + self.assertEqual(call_count, wrapped.call_count) return call_count From 33dbe66d403834f773b9559d5c99b1957873be1c Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 1 Dec 2020 11:04:04 +0100 Subject: [PATCH 645/690] Use autorange instead of fixed number of samples --- tests/test_utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index ae9c6e54a4..2e76b0471e 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -281,7 +281,8 @@ def test_estimate_cost_of_add_sequence(): for shapes in _generate_shapes(): arrays = [torch.empty(*shape) for shape in shapes] cost = estimate_cost_of_sequence(*(a.shape for a in arrays)) - consumption = timeit.timeit(stmt='sum(arrays)', globals=locals(), number=25) + n_samples, time = timeit.Timer(stmt='sum(arrays)', globals=dict(arrays=arrays)).autorange() + consumption = time / n_samples data.append((cost, consumption)) a = numpy.asarray(data) From 2c66dfd406e6e38c0085653bead2b558cafdd6aa Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 10 Dec 2020 17:40:27 +0100 Subject: [PATCH 646/690] Fix imports for triples_factory.py --- src/pykeen/triples/triples_factory.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/src/pykeen/triples/triples_factory.py b/src/pykeen/triples/triples_factory.py index af936c4c13..c16be43158 100644 --- a/src/pykeen/triples/triples_factory.py +++ b/src/pykeen/triples/triples_factory.py @@ -7,9 +7,6 @@ import logging import os import re -import typing -from collections import Counter, defaultdict -from typing import Collection, Dict, Iterable, List, Mapping, Optional, Sequence, Set, TextIO, Tuple, Union from typing import Callable, Collection, List, Mapping, Optional, Sequence, Set, TextIO, Union import numpy as np @@ -19,10 +16,8 @@ from .instances import Instances, LCWAInstances, SLCWAInstances from .splitting import split from .utils import load_triples -from ..typing import EntityMapping, LabeledTriples, MappedTriples, RelationMapping -from ..utils import compact_mapping, invert_mapping, random_non_negative_int from ..typing import EntityMapping, LabeledTriples, MappedTriples, RelationMapping, TorchRandomHint -from ..utils import compact_mapping, format_relative_comparison, invert_mapping, slice_triples, torch_is_in_1d +from ..utils import compact_mapping, format_relative_comparison, invert_mapping, torch_is_in_1d __all__ = [ 'TriplesFactory', From 23715378671845cfd43b0328eab758a8b12c477f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 10 Dec 2020 17:42:17 +0100 Subject: [PATCH 647/690] fix more imports --- src/pykeen/models/multimodal/complex_literal.py | 1 + src/pykeen/models/multimodal/distmult_literal.py | 1 + src/pykeen/utils.py | 5 +---- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index 568aaae2ec..1788f06c32 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -8,6 +8,7 @@ import torch.nn as nn from .base import LiteralModel +from ...constants import DEFAULT_DROPOUT_HPO_RANGE, DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE from ...losses import BCEWithLogitsLoss, Loss from ...nn import EmbeddingSpecification from ...nn.modules import ComplExInteraction diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 2ac65eaea8..227453f690 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -7,6 +7,7 @@ import torch.nn as nn from .base import LiteralModel +from ...constants import DEFAULT_DROPOUT_HPO_RANGE, DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.modules import DistMultInteraction diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index 1badca441b..d0d386df57 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -27,7 +27,6 @@ from .typing import DeviceHint, RandomHint, TorchRandomHint from .version import get_git_hash -# get_subclasses, project_entity, set_random_seed, strip_dim, view_complex __all__ = [ 'broadcast_cat', 'compose', @@ -37,7 +36,6 @@ 'compact_mapping', 'ensure_torch_random_state', 'format_relative_comparison', - 'imag_part', 'complex_normalize', 'fix_dataclass_init_docs', 'flatten_dictionary', @@ -56,7 +54,6 @@ 'set_random_seed', 'split_complex', 'split_list_in_batches_iter', - 'split_list_in_batches', 'torch_is_in_1d', 'normalize_string', 'normalized_lookup', @@ -619,7 +616,7 @@ def estimate_cost_of_sequence( ) -> int: """Cost of a sequence of broadcasted element-wise operations of tensors, given their shapes.""" return sum(map( - numpy.prod, + np.prod, itertools.islice( itertools.accumulate( (shape,) + other_shapes, From d551b1e62b673ce4083f731efc6991bd4974f9c3 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 10 Dec 2020 17:43:15 +0100 Subject: [PATCH 648/690] Fix ComplEx constructor merge --- src/pykeen/models/unimodal/complex.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 2db29f5788..76b19b791e 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -116,7 +116,4 @@ def __init__( loss=loss, preferred_device=preferred_device, random_seed=random_seed, - regularizer=regularizer, - # initialize with entity and relation embeddings with standard normal distribution, cf. - # https://github.com/ttrouill/complex/blob/dc4eb93408d9a5288c986695b58488ac80b1cc17/efe/models.py#L481-L487 ) From 9a18c9d1401ce18e8db174b84dbaf027891d0dae Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 10 Dec 2020 17:46:13 +0100 Subject: [PATCH 649/690] More post-merge fixes --- src/pykeen/models/multimodal/complex_literal.py | 2 -- tests/test_nn.py | 3 ++- tests/test_pipeline.py | 5 ++++- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index 1788f06c32..9b8c01b7e2 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -102,6 +102,4 @@ def __init__( predict_with_sigmoid=predict_with_sigmoid, preferred_device=preferred_device, random_seed=random_seed, - entity_initializer=xavier_normal_, - relation_initializer=xavier_normal_, ) diff --git a/tests/test_nn.py b/tests/test_nn.py index 7cd89ed41c..ff6981e580 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -349,6 +349,7 @@ class RGCNRepresentationTests(RepresentationModuleTests, unittest.TestCase): def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMapping[str, Any]: # noqa: D102 kwargs = super()._pre_instantiation_hook(kwargs=kwargs) + # TODO: use triple generation # generate random triples mapped_triples = numpy.stack([ numpy.random.randint(max_id, size=(self.num_triples,)) @@ -363,7 +364,7 @@ def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMa (entity_names, relation_names, entity_names), ) ]) - kwargs["triples_factory"] = TriplesFactory(triples=triples) + kwargs["triples_factory"] = TriplesFactory.from_labeled_triples(triples=triples) return kwargs diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index d120899a6c..752d53feaa 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -7,8 +7,11 @@ import pandas as pd +import pykeen from pykeen.datasets import Nations -from pykeen.pipeline import pipeline +from pykeen.models import Model +from pykeen.pipeline import PipelineResult, pipeline +from pykeen.regularizers import NoRegularizer class TestPipeline(unittest.TestCase): From e94bdc439a6f437cfdf1ba89d781bb33f4fdd32f Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 10 Dec 2020 17:51:55 +0100 Subject: [PATCH 650/690] Extract literal interaction wrapper --- src/pykeen/models/multimodal/base.py | 53 +++++++++++-------- .../models/multimodal/complex_literal.py | 14 ++--- .../models/multimodal/distmult_literal.py | 12 +++-- 3 files changed, 47 insertions(+), 32 deletions(-) diff --git a/src/pykeen/models/multimodal/base.py b/src/pykeen/models/multimodal/base.py index 628b66d79b..c911c38fd6 100644 --- a/src/pykeen/models/multimodal/base.py +++ b/src/pykeen/models/multimodal/base.py @@ -2,7 +2,7 @@ """Base classes for multi-modal models.""" -from typing import Optional, TYPE_CHECKING +from typing import Optional, TYPE_CHECKING, Tuple import torch from torch import nn @@ -21,14 +21,43 @@ from ...typing import Representation # noqa +class LiteralInteraction( + Interaction[ + Tuple[Representation, Representation], + Representation, + Tuple[Representation, Representation], + ]): + + def __init__( + self, + base: Interaction[Representation, Representation, Representation], + combination: nn.Module, + ): + super().__init__() + self.base = base + self.combination = combination + + def forward( + self, + h: Tuple[Representation, Representation], + r: Representation, + t: Tuple[Representation, Representation], + ) -> torch.FloatTensor: + # combine entity embeddings + literals + h, t = [ + self.combination(torch.cat(x, dim=-1)) + for x in (h, t) + ] + return self.base(h=h, r=r, t=t) + + class LiteralModel(ERModel[HeadRepresentation, RelationRepresentation, TailRepresentation], autoreset=False): """Base class for models with entity literals.""" def __init__( self, triples_factory: TriplesNumericLiteralsFactory, - interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], - combination: nn.Module, + interaction: LiteralInteraction, entity_specification: Optional[EmbeddingSpecification] = None, relation_specification: Optional[EmbeddingSpecification] = None, loss: Optional[Loss] = None, @@ -59,21 +88,3 @@ def __init__( specification=relation_specification, ), ) - self.combination = combination - - def forward( - self, - h_indices: Optional[torch.LongTensor], - r_indices: Optional[torch.LongTensor], - t_indices: Optional[torch.LongTensor], - slice_size: Optional[int] = None, - slice_dim: Optional[str] = None, - ) -> torch.FloatTensor: # noqa: D102 - h, r, t = self._get_representations(h_indices, r_indices, t_indices) - # combine entity embeddings + literals - h, t = [ - self.combination(torch.cat(x, dim=-1)) - for x in (h, t) - ] - scores = self.interaction.score(h=h, r=r, t=t, slice_size=slice_size, slice_dim=slice_dim) - return self._repeat_scores_if_necessary(scores, h_indices, r_indices, t_indices) diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index 9b8c01b7e2..ab2f5893c9 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -7,7 +7,7 @@ import torch import torch.nn as nn -from .base import LiteralModel +from .base import LiteralInteraction, LiteralModel from ...constants import DEFAULT_DROPOUT_HPO_RANGE, DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE from ...losses import BCEWithLogitsLoss, Loss from ...nn import EmbeddingSpecification @@ -82,11 +82,13 @@ def __init__( """Initialize the model.""" super().__init__( triples_factory=triples_factory, - interaction=ComplExInteraction(), - combination=ComplExLiteralCombination( - embedding_dim=embedding_dim, - num_of_literals=triples_factory.numeric_literals.shape[-1], - dropout=input_dropout, + interaction=LiteralInteraction( + base=ComplExInteraction(), + combination=ComplExLiteralCombination( + embedding_dim=embedding_dim, + num_of_literals=triples_factory.numeric_literals.shape[-1], + dropout=input_dropout, + ), ), entity_specification=EmbeddingSpecification( embedding_dim=embedding_dim, diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 227453f690..1bf1960c4f 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -6,7 +6,7 @@ import torch.nn as nn -from .base import LiteralModel +from .base import LiteralInteraction, LiteralModel from ...constants import DEFAULT_DROPOUT_HPO_RANGE, DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE from ...losses import Loss from ...nn import EmbeddingSpecification @@ -47,10 +47,12 @@ def __init__( ) -> None: super().__init__( triples_factory=triples_factory, - interaction=DistMultInteraction(), - combination=nn.Sequential( - nn.Linear(embedding_dim + triples_factory.numeric_literals.shape[1], embedding_dim), - nn.Dropout(input_dropout), + interaction=LiteralInteraction( + base=DistMultInteraction(), + combination=nn.Sequential( + nn.Linear(embedding_dim + triples_factory.numeric_literals.shape[1], embedding_dim), + nn.Dropout(input_dropout), + ), ), entity_specification=EmbeddingSpecification( embedding_dim=embedding_dim, From 60ecdb636abaf3339480d7000885d7600d38e11a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 10 Dec 2020 18:07:47 +0100 Subject: [PATCH 651/690] fix complex --- src/pykeen/models/multimodal/base.py | 11 ++++++----- src/pykeen/models/multimodal/complex_literal.py | 3 ++- tests/test_interactions.py | 2 ++ 3 files changed, 10 insertions(+), 6 deletions(-) diff --git a/src/pykeen/models/multimodal/base.py b/src/pykeen/models/multimodal/base.py index c911c38fd6..d8caa9c809 100644 --- a/src/pykeen/models/multimodal/base.py +++ b/src/pykeen/models/multimodal/base.py @@ -11,7 +11,7 @@ from ...losses import Loss from ...nn import Embedding, EmbeddingSpecification, Interaction, LiteralRepresentations from ...triples import TriplesNumericLiteralsFactory -from ...typing import DeviceHint, HeadRepresentation, RelationRepresentation, TailRepresentation +from ...typing import DeviceHint, HeadRepresentation, RelationRepresentation, Representation, TailRepresentation __all__ = [ "LiteralModel", @@ -36,6 +36,7 @@ def __init__( super().__init__() self.base = base self.combination = combination + self.entity_shape = tuple(self.base.entity_shape) + ("e",) def forward( self, @@ -44,10 +45,10 @@ def forward( t: Tuple[Representation, Representation], ) -> torch.FloatTensor: # combine entity embeddings + literals - h, t = [ - self.combination(torch.cat(x, dim=-1)) - for x in (h, t) - ] + h = torch.cat(h, dim=-1) + h = self.combination(h.view(-1, h.shape[-1])).view(*h.shape[:-1], -1) + t = torch.cat(t, dim=-1) + t = self.combination(t.view(-1, t.shape[-1])).view(*t.shape[:-1], -1) return self.base(h=h, r=r, t=t) diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index ab2f5893c9..e5b555c72e 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -41,7 +41,8 @@ def __init__( nn.Linear(embedding_dim + num_of_literals, embedding_dim), torch.nn.Tanh(), ) - self.embedding_dim = embedding_dim + # TODO: Determine this automatically + self.embedding_dim = 2 * embedding_dim def forward( self, diff --git a/tests/test_interactions.py b/tests/test_interactions.py index fba5bc26d0..99bdb920a7 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -12,6 +12,7 @@ import torch import pykeen.nn.modules +from pykeen.models.multimodal.base import LiteralInteraction from pykeen.nn.functional import distmult_interaction from pykeen.nn.modules import Interaction, TranslationalInteraction from pykeen.testing import base as ptb @@ -612,4 +613,5 @@ class InteractionTestsTest(ptb.TestsTest[Interaction], unittest.TestCase): base_test = InteractionTests skip_cls = { TranslationalInteraction, + LiteralInteraction, } From d282f98ec9a97cfc93737732518b99664839c677 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 07:07:41 +0100 Subject: [PATCH 652/690] Add work load generation for fast SLCWA --- benchmarking/interactions.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index 7927e93b23..e70b761a08 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -19,13 +19,23 @@ from pykeen.version import get_git_hash -def _use_case_to_shape(use_case: str, b: int, n: int) -> Tuple[ +def _use_case_to_shape( + use_case: str, + b: int, + n: int, + num_neg_samples: int, +) -> Tuple[ Tuple[int, int], Tuple[int, int], Tuple[int, int], ]: if use_case == "hrt": + b = b * num_neg_samples return (b, 1), (b, 1), (b, 1) + elif use_case == "hrt+": + return (b, 1), (b, 1), (b, num_neg_samples) + elif use_case == "h+rt": + return (b, num_neg_samples), (b, 1), (b, 1) elif use_case == "t": return (b, 1), (b, 1), (1, n) elif use_case == "h": From c6a60fadd11fc7bb653f5c64032c2d182dfe3afe Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 07:10:02 +0100 Subject: [PATCH 653/690] Add docstring --- benchmarking/interactions.py | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index e70b761a08..6de247ea2c 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -23,19 +23,41 @@ def _use_case_to_shape( use_case: str, b: int, n: int, - num_neg_samples: int, + s: int, ) -> Tuple[ Tuple[int, int], Tuple[int, int], Tuple[int, int], ]: + """ + Generate prefix shapes for various use cases. + + :param use_case: + The use case. + + - "hrt": score_hrt naive + - "hrt+": score_hrt fast SCLWA with tail corruption + - "h+rt": score_hrt fast SCLWA with head corruption + - "t": score_t + - "h": score_t + + :param b: + The batch size. + :param n: + The number of entities. + :param s: + The number of negative samples. + + :return: + A 3-tuple, (head_prefix, relation_prefix, tail_prefix), each a 2-tuple of integers. + """ if use_case == "hrt": - b = b * num_neg_samples + b = b * s return (b, 1), (b, 1), (b, 1) elif use_case == "hrt+": - return (b, 1), (b, 1), (b, num_neg_samples) + return (b, 1), (b, 1), (b, s) elif use_case == "h+rt": - return (b, num_neg_samples), (b, 1), (b, 1) + return (b, s), (b, 1), (b, 1) elif use_case == "t": return (b, 1), (b, 1), (1, n) elif use_case == "h": From 9abaecc5b6c969d4af2ee1750a19b7783755d4b6 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 07:23:02 +0100 Subject: [PATCH 654/690] Move utility method --- src/pykeen/nn/modules.py | 17 +++-------------- src/pykeen/utils.py | 13 ++++++++++++- tests/test_interactions.py | 3 ++- 3 files changed, 17 insertions(+), 16 deletions(-) diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index ea9c6e9937..637229896c 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -12,8 +12,8 @@ from . import functional as pkf from .representation import CANONICAL_DIMENSIONS, convert_to_canonical_shape -from ..typing import HeadRepresentation, RelationRepresentation, Representation, TailRepresentation -from ..utils import upgrade_to_sequence +from ..typing import HeadRepresentation, RelationRepresentation, TailRepresentation +from ..utils import ensure_tuple, upgrade_to_sequence __all__ = [ # Base Classes @@ -45,17 +45,6 @@ logger = logging.getLogger(__name__) -def _ensure_tuple(*x: Union[Representation, Sequence[Representation]]) -> Sequence[Sequence[Representation]]: - return tuple(upgrade_to_sequence(xx) for xx in x) - - -def _unpack_singletons(*xs: Tuple) -> Sequence[Tuple]: - return [ - x[0] if len(x) == 1 else x - for x in xs - ] - - def _get_prefix(slice_size, slice_dim, d) -> str: if slice_size is None or slice_dim != d: return 'b' @@ -64,7 +53,7 @@ def _get_prefix(slice_size, slice_dim, d) -> str: def _get_batches(z, slice_size): - for batch in zip(*(hh.split(slice_size, dim=1) for hh in _ensure_tuple(z)[0])): + for batch in zip(*(hh.split(slice_size, dim=1) for hh in ensure_tuple(z)[0])): if len(batch) == 1: batch = batch[0] yield batch diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index d0d386df57..a21b5fab34 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -24,7 +24,7 @@ from torch.nn import functional from .constants import PYKEEN_BENCHMARKS -from .typing import DeviceHint, RandomHint, TorchRandomHint +from .typing import DeviceHint, RandomHint, Representation, TorchRandomHint from .version import get_git_hash __all__ = [ @@ -846,3 +846,14 @@ def strip_dim(*x, num: int = 4): def upgrade_to_sequence(x: Union[X, Sequence[X]]) -> Sequence[X]: """Ensure that the input is a sequence.""" return x if isinstance(x, Sequence) else (x,) + + +def ensure_tuple(*x: Union[X, Sequence[X]]) -> Sequence[Sequence[X]]: + return tuple(upgrade_to_sequence(xx) for xx in x) + + +def unpack_singletons(*xs: Tuple[X]) -> Sequence[Union[X, Tuple[X]]]: + return [ + x[0] if len(x) == 1 else x + for x in xs + ] diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 99bdb920a7..1f9920c60f 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -12,6 +12,7 @@ import torch import pykeen.nn.modules +import pykeen.utils from pykeen.models.multimodal.base import LiteralInteraction from pykeen.nn.functional import distmult_interaction from pykeen.nn.modules import Interaction, TranslationalInteraction @@ -51,7 +52,7 @@ def _get_hrt( [self.cls.entity_shape, self.cls.relation_shape, self.cls.entity_shape], ) ) - return tuple(pykeen.nn.modules._unpack_singletons(*result)) + return tuple(pykeen.utils.unpack_singletons(*result)) def _check_scores(self, scores: torch.FloatTensor, exp_shape: Tuple[int, ...]): """Check shape, dtype and gradients of scores.""" From bc07b12bffc5adfbc48bfcfd044a15a505c1065a Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 07:29:08 +0100 Subject: [PATCH 655/690] Cleanup interaction benchmark script --- benchmarking/interactions.py | 87 ++++++++++++++++++++++-------------- 1 file changed, 53 insertions(+), 34 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index 6de247ea2c..c45512bce6 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -14,8 +14,8 @@ _complex_interaction_complex_native, _complex_interaction_direct, _complex_interaction_optimized_broadcasted, ) -from pykeen.nn.modules import _unpack_singletons from pykeen.typing import HeadRepresentation, RelationRepresentation, TailRepresentation +from pykeen.utils import unpack_singletons from pykeen.version import get_git_hash @@ -66,25 +66,38 @@ def _use_case_to_shape( raise ValueError -def _generate_hrt( +def _resolve_shapes( prefix_shapes: Tuple[Tuple[int, int], Tuple[int, int], Tuple[int, int]], interaction: Interaction[HeadRepresentation, RelationRepresentation, TailRepresentation], dim: int, - device: torch.device, additional_dims: Optional[Mapping[str, int]] = None, -) -> Tuple[HeadRepresentation, RelationRepresentation, TailRepresentation]: +) -> Tuple[Tuple[Tuple[int, ...], ...], ...]: additional_dims = additional_dims or dict() additional_dims.setdefault("d", dim) - return _unpack_singletons(*( - torch.rand(*prefix_shape, *(additional_dims[s] for s in suffix_shape), requires_grad=True, device=device) + + return [ + tuple(*prefix_shape, *(additional_dims[s] for s in suffix_shape)) for prefix_shape, suffix_shape in zip( - prefix_shapes, - ( - interaction.entity_shape, - interaction.relation_shape, - interaction.tail_entity_shape or interaction.entity_shape + prefix_shapes, + ( + interaction.entity_shape, + interaction.relation_shape, + interaction.tail_entity_shape or interaction.entity_shape + ) ) - ) + ] + + +def _generate_hrt( + shapes: Tuple[Tuple[Tuple[int, ...], ...], ...], + device: torch.device, +) -> Tuple[HeadRepresentation, RelationRepresentation, TailRepresentation]: + return unpack_singletons(*( + [ + torch.rand(*shape, requires_grad=True, device=device) + for shape in single_shapes + ] + for single_shapes in shapes )) @@ -94,12 +107,16 @@ def _get_result_shape(prefix_shapes) -> Tuple[int, int, int, int]: @click.command() @click.option('-m', '--max-result-elements-power', type=int, default=30, show_default=True) +@click.option('-n', '--max-num-entities-power', type=int, default=15, show_default=True) @click.option('-b', '--max-batch-size-power', type=int, default=10, show_default=True) @click.option('-d', '--max-vector-dimension-power', type=int, default=10, show_default=True) +@click.option('-s', '--max-sample-power', type=int, default=10, show_default=True) def main( max_result_elements_power: int, + max_num_entities_power: int, max_batch_size_power: int, max_vector_dimension_power: int, + max_sample_power: int, ): device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") print(f"Running on {device}.") @@ -108,73 +125,75 @@ def main( _complex_interaction_optimized_broadcasted, _complex_interaction_direct, ] - use_case_labels = ["hrt", "t", "h"] + use_case_labels = ["hrt", "hrt+", "h+rt", "t", "h"] batch_sizes = [2 ** i for i in range(5, max_batch_size_power + 1)] - num_entities = (100, 15_000) - # num_entities = (100,) + negative_samples = [2 ** i for i in range(5, max_sample_power + 1)] + num_entities = [2 ** i for i in range(7, max_num_entities_power)] max_result_elements = 2 ** max_result_elements_power vector_dimensions = [2 ** i for i in range(5, max_vector_dimension_power + 1)] data = [] tasks = [ - (b, n, d, ul, _use_case_to_shape(use_case=ul, b=b, n=n)) - for ul in use_case_labels + (b, s, n, d, ul, _use_case_to_shape(use_case=ul, b=b, n=n, s=s)) for b in batch_sizes + for s in negative_samples for n in num_entities for d in vector_dimensions + for ul in use_case_labels ] progress = tqdm(variants, unit="variant") for variant in progress: # create variant interaction = Interaction.from_func(variant) - for i, (b, n, d, ul, prefix_shapes) in enumerate(tqdm(tasks, unit="task"), start=1): + for i, (b, s, n, d, ul, prefix_shapes) in enumerate(tqdm(tasks, unit="task"), start=1): result_shape = _get_result_shape(prefix_shapes) - n_samples, total_time, time_per_sample = 0, float('nan'), float('nan') + median = iqr = float('nan') if max_result_elements is not None and max_result_elements < numpy.prod(result_shape): continue - h, r, t = _generate_hrt( + shapes = _resolve_shapes( prefix_shapes=prefix_shapes, interaction=interaction, dim=d, - device=device, ) try: timer = Timer( stmt="interaction(h=h, r=r, t=t)", - globals=dict(interaction=interaction, h=h, r=r, t=t), + globals=dict(interaction=interaction, shapes=shapes, device=device), + setup="h, r, t = _generate_hrt(shapes=shapes, device=device)" ) - n_samples, total_time = timer.autorange() - time_per_sample = total_time / n_samples + time = timer.blocked_autorange() + median = time.median + iqr = time.iqr except Exception as error: progress.write(f"ERROR: {error}") - progress.set_postfix(dict(shape=prefix_shapes, time=time_per_sample)) - data.append((i, b, n, d, prefix_shapes, ul, variant.__name__, total_time, n_samples, time_per_sample)) + progress.set_postfix(dict(shape=prefix_shapes, time=median)) + data.append((i, b, s, n, d, ul, prefix_shapes, variant.__name__, median, iqr)) git_hash = get_git_hash() df = pandas.DataFrame(data=data, columns=[ "experiment_number", "batch_size", + "num_negative_samples", "num_entities", "dimension", - "prefix_shapes", "use_case", + "prefix_shapes", "variant", - "total_time", - "n_samples", - "time_per_sample", + "time_median", + "time_inter_quartile_range", ]) df["device"] = device.type df.to_csv(f"{git_hash}_measurement.tsv", sep="\t", index=False) df_agg = df.groupby( by=["batch_size", "num_entities", "dimension", "use_case", "variant"] - ).agg({"time_per_sample": "mean"}).unstack().reset_index().dropna() + ).agg({"time_median": "mean"}).unstack().reset_index().dropna() df_agg.to_csv(f"{git_hash}_measurement_agg.tsv", sep="\t", index=False) print(df_agg) - viz_df = df[df['total_time'].notna()] - viz_df['variant'] = viz_df['variant'].map(lambda s: ' '.join(s.split('_')[3:]).capitalize()) + viz_df = df[df["time_median"].notna()] + viz_df["variant"] = viz_df["variant"].map(lambda s: ' '.join(s.split('_')[3:]).capitalize()) plt.figure(figsize=(14, 6)) - sns.lineplot(data=viz_df, x='experiment_number', y='total_time', hue='variant') + sns.lineplot(data=viz_df, x='experiment_number', y='time_median', hue='variant') plt.savefig(f"{git_hash}_measurement.png", dpi=300) From 4d21e798135f573d8787dc3c65fc1227ea0e1e09 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 07:32:58 +0100 Subject: [PATCH 656/690] add another complex variant --- benchmarking/interactions.py | 3 ++- src/pykeen/nn/functional.py | 26 +++++++++++++++++++++----- 2 files changed, 23 insertions(+), 6 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index c45512bce6..9f027125d6 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -12,7 +12,7 @@ from pykeen.nn import Interaction from pykeen.nn.functional import ( _complex_interaction_complex_native, _complex_interaction_direct, - _complex_interaction_optimized_broadcasted, + _complex_interaction_optimized_broadcasted, _complex_select, ) from pykeen.typing import HeadRepresentation, RelationRepresentation, TailRepresentation from pykeen.utils import unpack_singletons @@ -121,6 +121,7 @@ def main( device = torch.device("cuda") if torch.cuda.is_available() else torch.device("cpu") print(f"Running on {device}.") variants = [ + _complex_select, _complex_interaction_complex_native, _complex_interaction_optimized_broadcasted, _complex_interaction_direct, diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 08d87d0502..1aaa22a8e6 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -83,6 +83,22 @@ def wrapped(*args, **kwargs): return wrapped +def _complex_select( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Decide based on result shape whether to combine hr or ht first.""" + hr_cost = numpy.prod([max(hs, rs) for hs, rs in zip(h.shape, r.shape)]) + rt_cost = numpy.prod([max(ts, rs) for ts, rs in zip(t.shape, r.shape)]) + (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] + if hr_cost < rt_cost: + h_re, h_im = (h_re * r_re - h_im * r_im), (h_re * r_im + h_im * r_re) + else: + t_re, t_im = (t_re * r_re - t_im * r_im), (t_re * r_im + t_im * r_re) + return h_re @ t_re.transpose(-2, -1) - h_im @ t_im.transpose(-2, -1) + + def _complex_interaction_complex_native( h: torch.FloatTensor, r: torch.FloatTensor, @@ -103,11 +119,11 @@ def _complex_interaction_optimized_broadcasted( return sum(*( factor * tensor_product(hh, rr, tt).sum(dim=-1) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] )) From 8a612869ff038a02d513118ce4b11d7e3b411beb Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 07:39:34 +0100 Subject: [PATCH 657/690] add fast flag --- benchmarking/interactions.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index 9f027125d6..0443236605 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -74,9 +74,8 @@ def _resolve_shapes( ) -> Tuple[Tuple[Tuple[int, ...], ...], ...]: additional_dims = additional_dims or dict() additional_dims.setdefault("d", dim) - return [ - tuple(*prefix_shape, *(additional_dims[s] for s in suffix_shape)) + tuple((*prefix_shape, *(additional_dims[s] for ss in s)) for s in suffix_shape) for prefix_shape, suffix_shape in zip( prefix_shapes, ( @@ -106,12 +105,14 @@ def _get_result_shape(prefix_shapes) -> Tuple[int, int, int, int]: @click.command() +@click.option('--fast/--no-fast', default=False) @click.option('-m', '--max-result-elements-power', type=int, default=30, show_default=True) @click.option('-n', '--max-num-entities-power', type=int, default=15, show_default=True) @click.option('-b', '--max-batch-size-power', type=int, default=10, show_default=True) @click.option('-d', '--max-vector-dimension-power', type=int, default=10, show_default=True) @click.option('-s', '--max-sample-power', type=int, default=10, show_default=True) def main( + fast: bool, max_result_elements_power: int, max_num_entities_power: int, max_batch_size_power: int, @@ -141,6 +142,8 @@ def main( for d in vector_dimensions for ul in use_case_labels ] + if fast: + tasks = tasks[:5] progress = tqdm(variants, unit="variant") for variant in progress: # create variant @@ -158,7 +161,7 @@ def main( try: timer = Timer( stmt="interaction(h=h, r=r, t=t)", - globals=dict(interaction=interaction, shapes=shapes, device=device), + globals=dict(interaction=interaction, shapes=shapes, device=device, _generate_hrt=_generate_hrt), setup="h, r, t = _generate_hrt(shapes=shapes, device=device)" ) time = timer.blocked_autorange() From 5d5d0568d68d94b5cec93217c1a9985d08eea478 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 07:41:46 +0100 Subject: [PATCH 658/690] add shuffle option for more reliable meta time estimates --- benchmarking/interactions.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index 0443236605..51f3ec93fe 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -1,3 +1,4 @@ +import random from typing import Mapping, Optional, Tuple import click @@ -106,6 +107,7 @@ def _get_result_shape(prefix_shapes) -> Tuple[int, int, int, int]: @click.command() @click.option('--fast/--no-fast', default=False) +@click.option('--shuffle/--no-shuffle', default=False) @click.option('-m', '--max-result-elements-power', type=int, default=30, show_default=True) @click.option('-n', '--max-num-entities-power', type=int, default=15, show_default=True) @click.option('-b', '--max-batch-size-power', type=int, default=10, show_default=True) @@ -113,6 +115,7 @@ def _get_result_shape(prefix_shapes) -> Tuple[int, int, int, int]: @click.option('-s', '--max-sample-power', type=int, default=10, show_default=True) def main( fast: bool, + shuffle: bool, max_result_elements_power: int, max_num_entities_power: int, max_batch_size_power: int, @@ -144,6 +147,8 @@ def main( ] if fast: tasks = tasks[:5] + if shuffle: + random.shuffle(tasks) progress = tqdm(variants, unit="variant") for variant in progress: # create variant From c3c697b6f83b2020bb433550b8e6491d9ae7a3e9 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 07:47:11 +0100 Subject: [PATCH 659/690] Also measure memory --- benchmarking/interactions.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index 51f3ec93fe..f05478e785 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -105,6 +105,16 @@ def _get_result_shape(prefix_shapes) -> Tuple[int, int, int, int]: return (max(s[0] for s in prefix_shapes),) + tuple([s[1] for s in prefix_shapes]) +def _get_memory(interaction, shapes, device) -> int: + torch.cuda.reset_accumulated_memory_stats() + torch.cuda.reset_peak_memory_stats() + h, r, t = _generate_hrt(shapes=shapes, device=device) + interaction(h=h, r=r, t=t) + stats = torch.cuda.memory_stats() + return stats["active_bytes.all.peak"] + + + @click.command() @click.option('--fast/--no-fast', default=False) @click.option('--shuffle/--no-shuffle', default=False) @@ -155,7 +165,7 @@ def main( interaction = Interaction.from_func(variant) for i, (b, s, n, d, ul, prefix_shapes) in enumerate(tqdm(tasks, unit="task"), start=1): result_shape = _get_result_shape(prefix_shapes) - median = iqr = float('nan') + max_memory = median = iqr = float('nan') if max_result_elements is not None and max_result_elements < numpy.prod(result_shape): continue shapes = _resolve_shapes( @@ -172,10 +182,12 @@ def main( time = timer.blocked_autorange() median = time.median iqr = time.iqr + max_memory = _get_memory(interaction, shapes, device) + except Exception as error: progress.write(f"ERROR: {error}") - progress.set_postfix(dict(shape=prefix_shapes, time=median)) - data.append((i, b, s, n, d, ul, prefix_shapes, variant.__name__, median, iqr)) + progress.set_postfix(dict(s=prefix_shapes, t=median, mem=max_memory)) + data.append((i, b, s, n, d, ul, prefix_shapes, variant.__name__, median, iqr, max_memory)) git_hash = get_git_hash() df = pandas.DataFrame(data=data, columns=[ @@ -189,6 +201,7 @@ def main( "variant", "time_median", "time_inter_quartile_range", + "max_memory", ]) df["device"] = device.type df.to_csv(f"{git_hash}_measurement.tsv", sep="\t", index=False) From abe6b43628a66195324528ac887f32b3f69d47f7 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 07:49:19 +0100 Subject: [PATCH 660/690] Remove viz --- benchmarking/interactions.py | 14 +++----------- 1 file changed, 3 insertions(+), 11 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index f05478e785..3ad0c6ccdf 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -2,10 +2,8 @@ from typing import Mapping, Optional, Tuple import click -import matplotlib.pyplot as plt import numpy import pandas -import seaborn as sns import torch from torch.utils.benchmark import Timer from tqdm import tqdm @@ -114,7 +112,6 @@ def _get_memory(interaction, shapes, device) -> int: return stats["active_bytes.all.peak"] - @click.command() @click.option('--fast/--no-fast', default=False) @click.option('--shuffle/--no-shuffle', default=False) @@ -163,7 +160,8 @@ def main( for variant in progress: # create variant interaction = Interaction.from_func(variant) - for i, (b, s, n, d, ul, prefix_shapes) in enumerate(tqdm(tasks, unit="task"), start=1): + for i, config in enumerate(tqdm(tasks, unit="task"), start=1): + b, s, n, d, ul, prefix_shapes = config result_shape = _get_result_shape(prefix_shapes) max_memory = median = iqr = float('nan') if max_result_elements is not None and max_result_elements < numpy.prod(result_shape): @@ -185,7 +183,7 @@ def main( max_memory = _get_memory(interaction, shapes, device) except Exception as error: - progress.write(f"ERROR: {error}") + progress.write(f"ERROR: {error} for {variant}:{config}") progress.set_postfix(dict(s=prefix_shapes, t=median, mem=max_memory)) data.append((i, b, s, n, d, ul, prefix_shapes, variant.__name__, median, iqr, max_memory)) @@ -212,12 +210,6 @@ def main( df_agg.to_csv(f"{git_hash}_measurement_agg.tsv", sep="\t", index=False) print(df_agg) - viz_df = df[df["time_median"].notna()] - viz_df["variant"] = viz_df["variant"].map(lambda s: ' '.join(s.split('_')[3:]).capitalize()) - plt.figure(figsize=(14, 6)) - sns.lineplot(data=viz_df, x='experiment_number', y='time_median', hue='variant') - plt.savefig(f"{git_hash}_measurement.png", dpi=300) - if __name__ == '__main__': main() From db48aac150c94ecf75fdc00f4f605ebffb986cbb Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 07:49:23 +0100 Subject: [PATCH 661/690] fix variant --- src/pykeen/nn/functional.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 1aaa22a8e6..b3a0331609 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -116,7 +116,7 @@ def _complex_interaction_optimized_broadcasted( ) -> torch.FloatTensor: """Manually split into real/imag, and used optimized broadcasted combination.""" (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] - return sum(*( + return sum(( factor * tensor_product(hh, rr, tt).sum(dim=-1) for factor, hh, rr, tt in [ (+1, h_re, r_re, t_re), From 010bc7b4b3bdbccaed7ef4f6b99dee141a0677a2 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 08:03:25 +0100 Subject: [PATCH 662/690] Use single progress bar --- benchmarking/interactions.py | 69 ++++++++++++++++++------------------ src/pykeen/nn/functional.py | 33 +++++++++++++++++ 2 files changed, 68 insertions(+), 34 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index 3ad0c6ccdf..3eb9ef3a05 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -11,7 +11,7 @@ from pykeen.nn import Interaction from pykeen.nn.functional import ( _complex_interaction_complex_native, _complex_interaction_direct, - _complex_interaction_optimized_broadcasted, _complex_select, + _complex_interaction_optimized_broadcasted, _complex_select, _complex_stacked, _complex_stacked_select, ) from pykeen.typing import HeadRepresentation, RelationRepresentation, TailRepresentation from pykeen.utils import unpack_singletons @@ -136,6 +136,8 @@ def main( _complex_interaction_complex_native, _complex_interaction_optimized_broadcasted, _complex_interaction_direct, + _complex_stacked, + _complex_stacked_select, ] use_case_labels = ["hrt", "hrt+", "h+rt", "t", "h"] batch_sizes = [2 ** i for i in range(5, max_batch_size_power + 1)] @@ -145,47 +147,46 @@ def main( vector_dimensions = [2 ** i for i in range(5, max_vector_dimension_power + 1)] data = [] tasks = [ - (b, s, n, d, ul, _use_case_to_shape(use_case=ul, b=b, n=n, s=s)) + (v, b, s, n, d, ul, _use_case_to_shape(use_case=ul, b=b, n=n, s=s)) + for v in variants for b in batch_sizes for s in negative_samples for n in num_entities for d in vector_dimensions for ul in use_case_labels ] - if fast: - tasks = tasks[:5] if shuffle: random.shuffle(tasks) - progress = tqdm(variants, unit="variant") - for variant in progress: - # create variant - interaction = Interaction.from_func(variant) - for i, config in enumerate(tqdm(tasks, unit="task"), start=1): - b, s, n, d, ul, prefix_shapes = config - result_shape = _get_result_shape(prefix_shapes) - max_memory = median = iqr = float('nan') - if max_result_elements is not None and max_result_elements < numpy.prod(result_shape): - continue - shapes = _resolve_shapes( - prefix_shapes=prefix_shapes, - interaction=interaction, - dim=d, + if fast: + tasks = tasks[:5] + progress = tqdm(tasks, unit="task") + for i, config in enumerate(progress, start=1): + v, b, s, n, d, ul, prefix_shapes = config + interaction = Interaction.from_func(v) + result_shape = _get_result_shape(prefix_shapes) + max_memory = median = iqr = float('nan') + if max_result_elements is not None and max_result_elements < numpy.prod(result_shape): + continue + shapes = _resolve_shapes( + prefix_shapes=prefix_shapes, + interaction=interaction, + dim=d, + ) + try: + timer = Timer( + stmt="interaction(h=h, r=r, t=t)", + globals=dict(interaction=interaction, shapes=shapes, device=device, _generate_hrt=_generate_hrt), + setup="h, r, t = _generate_hrt(shapes=shapes, device=device)" ) - try: - timer = Timer( - stmt="interaction(h=h, r=r, t=t)", - globals=dict(interaction=interaction, shapes=shapes, device=device, _generate_hrt=_generate_hrt), - setup="h, r, t = _generate_hrt(shapes=shapes, device=device)" - ) - time = timer.blocked_autorange() - median = time.median - iqr = time.iqr - max_memory = _get_memory(interaction, shapes, device) - - except Exception as error: - progress.write(f"ERROR: {error} for {variant}:{config}") - progress.set_postfix(dict(s=prefix_shapes, t=median, mem=max_memory)) - data.append((i, b, s, n, d, ul, prefix_shapes, variant.__name__, median, iqr, max_memory)) + time = timer.blocked_autorange() + median = time.median + iqr = time.iqr + max_memory = _get_memory(interaction, shapes, device) + + except Exception as error: + progress.write(f"ERROR: {error} for {v}:{config}") + progress.set_postfix(dict(s=prefix_shapes, t=median, mem=max_memory)) + data.append((i, b, s, n, d, ul, prefix_shapes, v.__name__, median, iqr, max_memory)) git_hash = get_git_hash() df = pandas.DataFrame(data=data, columns=[ @@ -206,7 +207,7 @@ def main( df_agg = df.groupby( by=["batch_size", "num_entities", "dimension", "use_case", "variant"] - ).agg({"time_median": "mean"}).unstack().reset_index().dropna() + ).agg({"time_median": "mean"}).unstack().reset_index()#.dropna() df_agg.to_csv(f"{git_hash}_measurement_agg.tsv", sep="\t", index=False) print(df_agg) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index b3a0331609..f1200150a9 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -142,6 +142,39 @@ def _complex_interaction_direct( ) +def _complex_stacked( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Stack vectors.""" + (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (r, t)] + h = torch.cat([h, h], dim=-1) # re im re im + r = torch.cat([r_re, r_re, r_im, r_im], dim=-1) # re re im im + t = torch.cat([t_re, t_im, t_im, t_re], dim=-1) # re im im re + return (h * r * t).sum(dim=-1) + + +def _complex_stacked_select( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Stack vectors and select order.""" + (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (r, t)] + h = torch.cat([h, h], dim=-1) # re im re im + r = torch.cat([r_re, r_re, r_im, r_im], dim=-1) # re re im im + t = torch.cat([t_re, t_im, t_im, t_re], dim=-1) # re im im re + hr_cost = numpy.prod([max(hs, rs) for hs, rs in zip(h.shape, r.shape)]) + rt_cost = numpy.prod([max(ts, rs) for ts, rs in zip(t.shape, r.shape)]) + if hr_cost < rt_cost: + # h = h_re, -h_im + h = h * r + else: + t = r * t + return h @ t.transpose(-2, -1) + + def complex_interaction( h: torch.FloatTensor, r: torch.FloatTensor, From c51c62134c7fadef96ff85e5336b2b70bf5dcfdb Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 11:42:24 +0100 Subject: [PATCH 663/690] Extend docstring --- src/pykeen/nn/functional.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index f1200150a9..bb9f677725 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -1,6 +1,12 @@ # -*- coding: utf-8 -*- -"""Functional forms of interaction methods.""" +""" +Functional forms of interaction methods. + +The functional forms always assume the general form of the interaction function, where head, relation and tail +representations are provided in shape (batch_size, num_heads, *), (batch_size, num_relations, *), and +(batch_size, num_tails, *), and return a score tensor of shape (batch_size, num_heads, num_relations, num_tails). +""" from typing import Optional, Tuple, Union From 447e185b597fddec649100d13ba8e146b3f6e742 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 11:46:08 +0100 Subject: [PATCH 664/690] Add einsum variant --- src/pykeen/nn/functional.py | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index bb9f677725..63139256be 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -181,6 +181,26 @@ def _complex_stacked_select( return h @ t.transpose(-2, -1) +def _complex_einsum( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Use einsum.""" + x = h.new_zeros(2, 2, 2) + x[0, 0, 0] = 1 + x[0, 1, 1] = 1 + x[1, 0, 1] = 1 + x[1, 1, 0] = -1 + return extended_einsum( + "ijk,bhdi,brdj,btdk->bhrt", + x, + h.view(*h.shape[:-1], -1, 2), + r.view(*r.shape[:-1], -1, 2), + t.view(*t.shape[:-1], -1, 2), + ) + + def complex_interaction( h: torch.FloatTensor, r: torch.FloatTensor, From bcf20d09d09773c08b2a75bc2c0289ca304f6feb Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 14:44:57 +0100 Subject: [PATCH 665/690] Fix device error in KL divergence --- src/pykeen/nn/sim.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index 5d9d66442b..21df717abe 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -194,7 +194,7 @@ def _vectorized_kl_divergence( terms.append(batched_dot(mu.pow(2), r_var_safe_reciprocal)) # 3. Component if exact: - terms.append(-torch.as_tensor(data=[h.mean.shape[-1]])) + terms.append(-torch.as_tensor(data=[h.mean.shape[-1]], device=mu.device)) # 4. Component # ln (det(\Sigma_1) / det(\Sigma_0)) # = ln det Sigma_1 - ln det Sigma_0 From e2863de6c283639aff2a48182080aef4cb411be3 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 14:48:43 +0100 Subject: [PATCH 666/690] Do not use root logger --- src/pykeen/models/unimodal/trans_r.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 537990c48e..789682753c 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -1,7 +1,6 @@ # -*- coding: utf-8 -*- """Implementation of TransR.""" - import logging from typing import Optional @@ -21,6 +20,8 @@ 'TransR', ] +logger = logging.getLogger(__name__) + class TransR(ERModel): r"""An implementation of TransR from [lin2015]_. @@ -104,4 +105,4 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, ) - logging.warning("Initialize from TransE") + logger.warning("Initialize from TransE") From 03d8b2a05f43c45813d04b183a6cf4c1c6bcd681 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 15 Dec 2020 14:52:14 +0100 Subject: [PATCH 667/690] Update tox.ini --- tox.ini | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tox.ini b/tox.ini index d5e033448c..d1a0b419b3 100644 --- a/tox.ini +++ b/tox.ini @@ -118,7 +118,7 @@ commands = src/pykeen/nn \ src/pykeen/regularizers.py \ src/pykeen/losses.py \ - src/pykeen/trackers.py \ + src/pykeen/trackers \ src/pykeen/models/base.py \ src/pykeen/triples/triples_factory.py From a59fc7b8e6d608da0a16f0bdfaf647fb2fd266ba Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 15 Dec 2020 14:52:19 +0100 Subject: [PATCH 668/690] Pass flake8 --- src/pykeen/models/multimodal/base.py | 3 ++- src/pykeen/models/unimodal/simple.py | 6 +++--- src/pykeen/nn/functional.py | 14 +++++++------- src/pykeen/utils.py | 2 +- 4 files changed, 13 insertions(+), 12 deletions(-) diff --git a/src/pykeen/models/multimodal/base.py b/src/pykeen/models/multimodal/base.py index d8caa9c809..90a4dd973e 100644 --- a/src/pykeen/models/multimodal/base.py +++ b/src/pykeen/models/multimodal/base.py @@ -26,7 +26,8 @@ class LiteralInteraction( Tuple[Representation, Representation], Representation, Tuple[Representation, Representation], - ]): + ], +): def __init__( self, diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index 83500aad15..fa6da00030 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -117,9 +117,9 @@ def forward( t_source.get_in_canonical_shape(dim="t", indices=t_indices), ) for h_source, r_source, t_source in ( - (self.entity_representations[0], self.relation_representations[0], self.entity_representations[1]), - (self.entity_representations[1], self.relation_representations[1], self.entity_representations[0]), - ) + (self.entity_representations[0], self.relation_representations[0], self.entity_representations[1]), + (self.entity_representations[1], self.relation_representations[1], self.entity_representations[0]), + ) )) scores = self.interaction.score(h=h, r=r, t=t, slice_size=slice_size, slice_dim=slice_dim) return self._repeat_scores_if_necessary(scores, h_indices, r_indices, t_indices) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 63139256be..8a7bea0a1d 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -122,15 +122,15 @@ def _complex_interaction_optimized_broadcasted( ) -> torch.FloatTensor: """Manually split into real/imag, and used optimized broadcasted combination.""" (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] - return sum(( + return sum( factor * tensor_product(hh, rr, tt).sum(dim=-1) for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] - )) + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] + ) def _complex_interaction_direct( diff --git a/src/pykeen/utils.py b/src/pykeen/utils.py index a21b5fab34..a06e247560 100644 --- a/src/pykeen/utils.py +++ b/src/pykeen/utils.py @@ -24,7 +24,7 @@ from torch.nn import functional from .constants import PYKEEN_BENCHMARKS -from .typing import DeviceHint, RandomHint, Representation, TorchRandomHint +from .typing import DeviceHint, RandomHint, TorchRandomHint from .version import get_git_hash __all__ = [ From b3502e7b50394b4c9915d4c680c13cc0a4ff315f Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 15 Dec 2020 15:19:43 +0100 Subject: [PATCH 669/690] Update functional.py --- src/pykeen/nn/functional.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 8a7bea0a1d..0f77d9c092 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -4,8 +4,8 @@ Functional forms of interaction methods. The functional forms always assume the general form of the interaction function, where head, relation and tail -representations are provided in shape (batch_size, num_heads, *), (batch_size, num_relations, *), and -(batch_size, num_tails, *), and return a score tensor of shape (batch_size, num_heads, num_relations, num_tails). +representations are provided in shape (batch_size, num_heads, ``*``), (batch_size, num_relations, ``*``), and +(batch_size, num_tails, ``*``), and return a score tensor of shape (batch_size, num_heads, num_relations, num_tails). """ from typing import Optional, Tuple, Union From 3a8c9a7a6db5e617df2627a2e6cd20e81e811c45 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Tue, 15 Dec 2020 15:32:33 +0100 Subject: [PATCH 670/690] Fix tucker without explicit relation dim --- src/pykeen/models/unimodal/tucker.py | 1 + 1 file changed, 1 insertion(+) diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index aee9e24b09..c1be22f0da 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -85,6 +85,7 @@ def __init__( where h,r,t are the head, relation, and tail embedding, W is the core tensor, x_i denotes the tensor product along the i-th mode, BN denotes batch normalization, and DO dropout. """ + relation_dim = relation_dim or embedding_dim super().__init__( triples_factory=triples_factory, interaction=TuckerInteraction( From 392480a097c780f6cb906a451b555f97f65f1ada Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Tue, 15 Dec 2020 22:13:35 +0100 Subject: [PATCH 671/690] Pass mypy Trigger CI --- src/pykeen/models/multimodal/base.py | 4 ++-- src/pykeen/nn/emb.py | 13 ++++++++----- src/pykeen/nn/modules.py | 10 +++++++++- src/pykeen/triples/splitting.py | 5 +++-- src/pykeen/triples/triples_factory.py | 6 +++--- 5 files changed, 25 insertions(+), 13 deletions(-) diff --git a/src/pykeen/models/multimodal/base.py b/src/pykeen/models/multimodal/base.py index 90a4dd973e..496b8743ca 100644 --- a/src/pykeen/models/multimodal/base.py +++ b/src/pykeen/models/multimodal/base.py @@ -47,9 +47,9 @@ def forward( ) -> torch.FloatTensor: # combine entity embeddings + literals h = torch.cat(h, dim=-1) - h = self.combination(h.view(-1, h.shape[-1])).view(*h.shape[:-1], -1) + h = self.combination(h.view(-1, h.shape[-1])).view(*h.shape[:-1], -1) # type: ignore t = torch.cat(t, dim=-1) - t = self.combination(t.view(-1, t.shape[-1])).view(*t.shape[:-1], -1) + t = self.combination(t.view(-1, t.shape[-1])).view(*t.shape[:-1], -1) # type: ignore return self.base(h=h, r=r, t=t) diff --git a/src/pykeen/nn/emb.py b/src/pykeen/nn/emb.py index bded90e6d8..52b7def6eb 100644 --- a/src/pykeen/nn/emb.py +++ b/src/pykeen/nn/emb.py @@ -90,15 +90,18 @@ def __init__( if initializer_kwargs: self.initializer = functools.partial(initializer, **initializer_kwargs) else: - self.initializer = initializer - if constrainer_kwargs: + self.initializer = initializer # type: ignore + + if constrainer is not None and constrainer_kwargs: self.constrainer = functools.partial(constrainer, **constrainer_kwargs) else: - self.constrainer = constrainer - if normalizer_kwargs: + self.constrainer = constrainer # type: ignore + + if normalizer is not None and normalizer_kwargs: self.normalizer = functools.partial(normalizer, **normalizer_kwargs) else: - self.normalizer = normalizer + self.normalizer = normalizer # type: ignore + self._embeddings = torch.nn.Embedding( num_embeddings=num_embeddings, embedding_dim=embedding_dim, diff --git a/src/pykeen/nn/modules.py b/src/pykeen/nn/modules.py index 637229896c..91eef5a97a 100644 --- a/src/pykeen/nn/modules.py +++ b/src/pykeen/nn/modules.py @@ -5,7 +5,10 @@ import logging import math from abc import ABC -from typing import Any, Callable, Generic, Mapping, MutableMapping, Optional, Sequence, Tuple, Type, Union +from typing import ( + Any, Callable, Generic, Mapping, MutableMapping, Optional, Sequence, TYPE_CHECKING, Tuple, Type, + Union, +) import torch from torch import FloatTensor, nn @@ -15,6 +18,9 @@ from ..typing import HeadRepresentation, RelationRepresentation, TailRepresentation from ..utils import ensure_tuple, upgrade_to_sequence +if TYPE_CHECKING: + from ..typing import Representation # noqa + __all__ = [ # Base Classes 'Interaction', @@ -82,8 +88,10 @@ def from_func(cls, f) -> 'Interaction': @classmethod def cls_from_func(cls, f) -> Type['Interaction']: """Create a stateless interaction class.""" + class StatelessInteraction(cls): # type: ignore func = f + return StatelessInteraction @staticmethod diff --git a/src/pykeen/triples/splitting.py b/src/pykeen/triples/splitting.py index c14d7e798f..d04e7dd228 100644 --- a/src/pykeen/triples/splitting.py +++ b/src/pykeen/triples/splitting.py @@ -3,6 +3,7 @@ """Implementation of triples splitting functions.""" import logging +import typing from typing import Optional, Sequence, Tuple, Union import numpy @@ -244,10 +245,10 @@ def _prepare_cleanup( columns = [[0, 2], [1]] to_move_mask = torch.zeros(1, dtype=torch.bool) if max_ids is None: - max_ids = [ + max_ids = typing.cast(Tuple[int, int], tuple( max(training[:, col].max().item(), testing[:, col].max().item()) + 1 for col in columns - ] + )) for col, max_id in zip(columns, max_ids): # IDs not in training not_in_training_mask = torch.ones(max_id, dtype=torch.bool) diff --git a/src/pykeen/triples/triples_factory.py b/src/pykeen/triples/triples_factory.py index 8cd8c2583d..dbb1a66682 100644 --- a/src/pykeen/triples/triples_factory.py +++ b/src/pykeen/triples/triples_factory.py @@ -363,7 +363,7 @@ def clone_and_exchange_triples( create_inverse_triples=self.create_inverse_triples, metadata={ **(extra_metadata or {}), - **(self.metadata if keep_metadata else {}), + **(self.metadata if keep_metadata else {}), # type: ignore }, ) @@ -403,7 +403,7 @@ def extra_repr(self) -> str: ('num_triples', self.num_triples), ('inverse_triples', self.create_inverse_triples), ] - d.extend(sorted(self.metadata.items())) + d.extend(sorted(self.metadata.items())) # type: ignore return ', '.join( f'{k}="{v}"' if isinstance(v, str) else f'{k}={v}' for k, v in d @@ -416,7 +416,7 @@ def get_inverse_relation_id(self, relation: Union[str, int]) -> int: """Get the inverse relation identifier for the given relation.""" if not self.create_inverse_triples: raise ValueError('Can not get inverse triple, they have not been created.') - relation = next(iter(self.relations_to_ids(relations=[relation]))) + relation = next(iter(self.relations_to_ids(relations=[relation]))) # type:ignore return self._get_inverse_relation_id(relation) @staticmethod From 19a2ec151368ee3323b09abbc6bbe4b745a75ce9 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Wed, 16 Dec 2020 13:24:55 +0100 Subject: [PATCH 672/690] Bring back configurable regularization Trigger CI --- src/pykeen/hpo/hpo.py | 23 +++++++++++++---------- src/pykeen/models/unimodal/complex.py | 16 ++++++++++------ src/pykeen/models/unimodal/conv_kb.py | 15 ++++++++------- src/pykeen/models/unimodal/distmult.py | 7 +++++-- src/pykeen/models/unimodal/rescal.py | 6 ++++-- src/pykeen/models/unimodal/simple.py | 6 ++++-- src/pykeen/models/unimodal/trans_h.py | 18 +++++++++++------- src/pykeen/pipeline.py | 13 ++++++------- tests/test_pipeline.py | 18 ++++++++++++------ 9 files changed, 73 insertions(+), 49 deletions(-) diff --git a/src/pykeen/hpo/hpo.py b/src/pykeen/hpo/hpo.py index 4b0a087fcd..d90ecc116b 100644 --- a/src/pykeen/hpo/hpo.py +++ b/src/pykeen/hpo/hpo.py @@ -57,7 +57,7 @@ class Objective: dataset: Union[None, str, Type[Dataset]] # 1. model: Type[Model] # 2. loss: Type[Loss] # 3. - regularizer: Type[Regularizer] # 4. + regularizer: Optional[Type[Regularizer]] # 4. optimizer: Type[Optimizer] # 5. training_loop: Type[TrainingLoop] # 6. evaluator: Type[Evaluator] # 8. @@ -149,13 +149,16 @@ def __call__(self, trial: Trial) -> Optional[float]: kwargs_ranges=self.loss_kwargs_ranges, ) # 4. Regularizer - # _regularizer_kwargs = _get_kwargs( - # trial=trial, - # prefix='regularizer', - # default_kwargs_ranges=self.regularizer.hpo_default, - # kwargs=self.regularizer_kwargs, - # kwargs_ranges=self.regularizer_kwargs_ranges, - # ) + if self.regularizer is not None: + _regularizer_kwargs = _get_kwargs( + trial=trial, + prefix='regularizer', + default_kwargs_ranges=self.regularizer.hpo_default, + kwargs=self.regularizer_kwargs, + kwargs_ranges=self.regularizer_kwargs_ranges, + ) + else: + _regularizer_kwargs = None # 5. Optimizer _optimizer_kwargs = _get_kwargs( trial=trial, @@ -205,8 +208,8 @@ def __call__(self, trial: Trial) -> Optional[float]: loss=self.loss, loss_kwargs=_loss_kwargs, # 4. Regularizer - # regularizer=self.regularizer, - # regularizer_kwargs=_regularizer_kwargs, + regularizer=self.regularizer, + regularizer_kwargs=_regularizer_kwargs, clear_optimizer=True, # 5. Optimizer optimizer=self.optimizer, diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 76b19b791e..79d89efbd9 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -12,7 +12,7 @@ from ...losses import Loss, SoftplusLoss from ...nn import EmbeddingSpecification from ...nn.modules import ComplExInteraction -from ...regularizers import LpRegularizer +from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -73,6 +73,7 @@ def __init__( triples_factory: TriplesFactory, embedding_dim: int = 200, loss: Optional[Loss] = None, + regularizer: Optional[Regularizer] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, embedding_specification: Optional[EmbeddingSpecification] = None, @@ -84,14 +85,17 @@ def __init__( The triple factory connected to the model. :param embedding_dim: The embedding dimensionality of the entity embeddings. - :param loss: OptionalLoss (optional) - The loss to use. Defaults to SoftplusLoss. - :param preferred_device: str (optional) + :param loss: + The loss to use. Defaults to :data:`loss_default`. + :param regularizer: + The regularizer to use. Defaults to :data:`regularizer_default`. + :param preferred_device: The default device where to model is located. - :param random_seed: int (optional) + :param random_seed: An optional random seed to set before the initialization of weights. """ - regularizer = self._instantiate_default_regularizer() + if regularizer is None: + regularizer = self._instantiate_default_regularizer() # initialize with entity and relation embeddings with standard normal distribution, cf. # https://github.com/ttrouill/complex/blob/dc4eb93408d9a5288c986695b58488ac80b1cc17/efe/models.py#L481-L487 if embedding_specification is None: diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 1095d918ec..85f392bbb8 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -10,7 +10,7 @@ from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.modules import ConvKBInteraction -from ...regularizers import LpRegularizer +from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -76,6 +76,7 @@ def __init__( hidden_dropout_rate: float = 0., embedding_dim: int = 200, loss: Optional[Loss] = None, + regularizer: Optional[Regularizer] = None, preferred_device: DeviceHint = None, num_filters: int = 400, random_seed: Optional[int] = None, @@ -101,12 +102,12 @@ def __init__( preferred_device=preferred_device, random_seed=random_seed, ) - regularizer = self._instantiate_default_regularizer() + if regularizer is None: + regularizer = self._instantiate_default_regularizer() # In the code base only the weights of the output layer are used for regularization # c.f. https://github.com/daiquocnguyen/ConvKB/blob/73a22bfa672f690e217b5c18536647c7cf5667f1/model.py#L60-L66 - if regularizer is not None: - self.append_weight_regularizer( - parameter=self.interaction.parameters(), - regularizer=regularizer, - ) + self.append_weight_regularizer( + parameter=self.interaction.parameters(), + regularizer=regularizer, + ) logger.warning('To be consistent with the paper, initialize entity and relation embeddings from TransE.') diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 7ebd75b872..87ec5216e8 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -12,7 +12,7 @@ from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.modules import DistMultInteraction -from ...regularizers import LpRegularizer +from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint from ...utils import compose @@ -73,6 +73,7 @@ def __init__( triples_factory: TriplesFactory, embedding_dim: int = 50, loss: Optional[Loss] = None, + regularizer: Optional[Regularizer] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, ) -> None: @@ -80,6 +81,8 @@ def __init__( :param embedding_dim: The entity embedding dimension $d$. Is usually $d \in [50, 300]$. """ + if regularizer is None: + regularizer = self._instantiate_default_regularizer() super().__init__( triples_factory=triples_factory, interaction=DistMultInteraction(), @@ -99,7 +102,7 @@ def __init__( functional.normalize, ), # Only relation embeddings are regularized - regularizer=self._instantiate_default_regularizer(), + regularizer=regularizer, ), loss=loss, preferred_device=preferred_device, diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index eac216fd0b..b42b469489 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -9,7 +9,7 @@ from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.modules import RESCALInteraction -from ...regularizers import LpRegularizer +from ...regularizers import LpRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -56,6 +56,7 @@ def __init__( triples_factory: TriplesFactory, embedding_dim: int = 50, loss: Optional[Loss] = None, + regularizer: Optional[Regularizer] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, ) -> None: @@ -67,7 +68,8 @@ def __init__( - OpenKE `implementation of RESCAL `_ """ - regularizer = self._instantiate_default_regularizer() + if regularizer is None: + regularizer = self._instantiate_default_regularizer() super().__init__( triples_factory=triples_factory, interaction=RESCALInteraction(), diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index fa6da00030..bfd58328f5 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -11,7 +11,7 @@ from ...losses import Loss, SoftplusLoss from ...nn import EmbeddingSpecification from ...nn.modules import SimplEInteraction -from ...regularizers import PowerSumRegularizer +from ...regularizers import PowerSumRegularizer, Regularizer from ...triples import TriplesFactory from ...typing import DeviceHint @@ -69,11 +69,13 @@ def __init__( triples_factory: TriplesFactory, embedding_dim: int = 200, loss: Optional[Loss] = None, + regularizer: Optional[Regularizer] = None, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, clamp_score: Optional[Union[float, Tuple[float, float]]] = None, ) -> None: - regularizer = self._instantiate_default_regularizer() + if regularizer is None: + regularizer = self._instantiate_default_regularizer() super().__init__( triples_factory=triples_factory, interaction=SimplEInteraction(clamp_score=clamp_score), diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index aef8adc7a1..47419045fd 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -11,7 +11,7 @@ from ...losses import Loss from ...nn import EmbeddingSpecification from ...nn.modules import TransHInteraction -from ...regularizers import TransHRegularizer +from ...regularizers import Regularizer, TransHRegularizer from ...triples import TriplesFactory from ...typing import DeviceHint from ...utils import pop_only @@ -70,6 +70,7 @@ def __init__( embedding_dim: int = 50, scoring_fct_norm: int = 2, loss: Optional[Loss] = None, + regularizer: Optional[Regularizer] = None, predict_with_sigmoid: bool = False, preferred_device: DeviceHint = None, random_seed: Optional[int] = None, @@ -103,9 +104,12 @@ def __init__( random_seed=random_seed, predict_with_sigmoid=predict_with_sigmoid, ) - # Note that the TransH regularizer has a different interface - self.regularizer = self._instantiate_default_regularizer( - entity_embeddings=pop_only(self.entity_representations[0].parameters()), - relation_embeddings=pop_only(self.relation_representations[0].parameters()), - normal_vector_embeddings=pop_only(self.relation_representations[1].parameters()), - ) + if regularizer is None: + # Note that the TransH regularizer has a different interface + self.regularizer = self._instantiate_default_regularizer( + entity_embeddings=pop_only(self.entity_representations[0].parameters()), + relation_embeddings=pop_only(self.relation_representations[0].parameters()), + normal_vector_embeddings=pop_only(self.relation_representations[1].parameters()), + ) + else: + self.regularizer = regularizer diff --git a/src/pykeen/pipeline.py b/src/pykeen/pipeline.py index 03302c8913..786c5badb8 100644 --- a/src/pykeen/pipeline.py +++ b/src/pykeen/pipeline.py @@ -188,7 +188,7 @@ from .models.base import Model from .nn import Embedding from .optimizers import get_optimizer_cls -from .regularizers import Regularizer +from .regularizers import Regularizer, get_regularizer_cls from .sampling import NegativeSampler, get_negative_sampler_cls from .stoppers import EarlyStopper, Stopper, get_stopper_cls from .trackers import ResultTracker, get_result_tracker_cls @@ -899,13 +899,12 @@ def pipeline( # noqa: C901 model_kwargs.setdefault('random_seed', random_seed) if regularizer is not None: - logger.warning('Specification of the regularizer from the pipeline() is currently under maitenance') # FIXME this should never happen. - # if 'regularizer' in model_kwargs: - # logger.warning('Can not specify regularizer in kwargs and model_kwargs. removing from model_kwargs') - # del model_kwargs['regularizer'] - # regularizer_cls: Type[Regularizer] = get_regularizer_cls(regularizer) - # model_kwargs['regularizer'] = regularizer_cls(**(regularizer_kwargs or {})) + if 'regularizer' in model_kwargs: + logger.warning('Can not specify regularizer in kwargs and model_kwargs. removing from model_kwargs') + del model_kwargs['regularizer'] + regularizer_cls: Type[Regularizer] = get_regularizer_cls(regularizer) + model_kwargs['regularizer'] = regularizer_cls(**(regularizer_kwargs or {})) if loss is not None: if 'loss' in model_kwargs: # FIXME diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py index 752d53feaa..da05d878a3 100644 --- a/tests/test_pipeline.py +++ b/tests/test_pipeline.py @@ -9,7 +9,8 @@ import pykeen from pykeen.datasets import Nations -from pykeen.models import Model +from pykeen.models import DistMult +from pykeen.nn import Embedding from pykeen.pipeline import PipelineResult, pipeline from pykeen.regularizers import NoRegularizer @@ -185,7 +186,7 @@ class TestAttributes(unittest.TestCase): def test_specify_regularizer(self): """Test a pipeline that uses a regularizer.""" for regularizer, cls in [ - (None, pykeen.regularizers.NoRegularizer), + (None, pykeen.regularizers.LpRegularizer), # if none, goes to default. ('no', pykeen.regularizers.NoRegularizer), (NoRegularizer, pykeen.regularizers.NoRegularizer), ('powersum', pykeen.regularizers.PowerSumRegularizer), @@ -193,11 +194,16 @@ def test_specify_regularizer(self): ]: with self.subTest(regularizer=regularizer): pipeline_result = pipeline( - model='TransE', + model='DistMult', dataset='Nations', regularizer=regularizer, - training_kwargs=dict(num_epochs=1), + training_kwargs=dict(num_epochs=1, use_tqdm=False), ) self.assertIsInstance(pipeline_result, PipelineResult) - self.assertIsInstance(pipeline_result.model, Model) - self.assertIsInstance(pipeline_result.model.regularizer, cls) + self.assertIsInstance(pipeline_result.model, DistMult) + self.assertEqual(1, len(pipeline_result.model.entity_representations)) + self.assertIsInstance(pipeline_result.model.entity_representations[0], Embedding) + self.assertIsNone(pipeline_result.model.entity_representations[0].regularizer) + self.assertEqual(1, len(pipeline_result.model.relation_representations)) + self.assertIsInstance(pipeline_result.model.relation_representations[0], Embedding) + self.assertIsInstance(pipeline_result.model.relation_representations[0].regularizer, cls) From 44d156d5ecc6aa01dfc9f57a8630da22b45cec78 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Wed, 16 Dec 2020 14:55:09 +0100 Subject: [PATCH 673/690] Update pipeline.py --- src/pykeen/pipeline.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/src/pykeen/pipeline.py b/src/pykeen/pipeline.py index 786c5badb8..720ad381fe 100644 --- a/src/pykeen/pipeline.py +++ b/src/pykeen/pipeline.py @@ -904,7 +904,10 @@ def pipeline( # noqa: C901 logger.warning('Can not specify regularizer in kwargs and model_kwargs. removing from model_kwargs') del model_kwargs['regularizer'] regularizer_cls: Type[Regularizer] = get_regularizer_cls(regularizer) - model_kwargs['regularizer'] = regularizer_cls(**(regularizer_kwargs or {})) + model_kwargs['regularizer'] = regularizer_cls( + # device=device, + **(regularizer_kwargs or {}), + ) if loss is not None: if 'loss' in model_kwargs: # FIXME @@ -1012,8 +1015,8 @@ def pipeline( # noqa: C901 logging.debug(f"model_kwargs: {model_kwargs}") logging.debug(f"loss: {loss}") logging.debug(f"loss_kwargs: {loss_kwargs}") - # logging.debug(f"regularizer: {regularizer}") - # logging.debug(f"regularizer_kwargs: {regularizer_kwargs}") + logging.debug(f"regularizer: {regularizer}") + logging.debug(f"regularizer_kwargs: {regularizer_kwargs}") logging.debug(f"optimizer: {optimizer}") logging.debug(f"optimizer_kwargs: {optimizer_kwargs}") logging.debug(f"training_loop: {training_loop}") From 65f36be31789ebd3257771560cd7f1bd15a80200 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Wed, 16 Dec 2020 14:56:47 +0100 Subject: [PATCH 674/690] Now I remember why I did this :laughing: Gotta revert this Trigger CI --- src/pykeen/models/unimodal/conv_kb.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 85f392bbb8..31d5890425 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -106,8 +106,9 @@ def __init__( regularizer = self._instantiate_default_regularizer() # In the code base only the weights of the output layer are used for regularization # c.f. https://github.com/daiquocnguyen/ConvKB/blob/73a22bfa672f690e217b5c18536647c7cf5667f1/model.py#L60-L66 - self.append_weight_regularizer( - parameter=self.interaction.parameters(), - regularizer=regularizer, - ) + if regularizer is not None: + self.append_weight_regularizer( + parameter=self.interaction.parameters(), + regularizer=regularizer, + ) logger.warning('To be consistent with the paper, initialize entity and relation embeddings from TransE.') From 1904672751905c14168c0e9cdbaff6b1c350dbc9 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Wed, 16 Dec 2020 15:50:16 +0100 Subject: [PATCH 675/690] add another complex variant --- src/pykeen/nn/functional.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 0f77d9c092..89f981d10c 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -105,6 +105,23 @@ def _complex_select( return h_re @ t_re.transpose(-2, -1) - h_im @ t_im.transpose(-2, -1) +def _complex_native_select( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Use torch built-ins for computation with complex numbers and select whether to combine hr or ht first.""" + h, r, t = [view_complex(x=x) for x in (h, r, t)] + hr_cost = numpy.prod([max(hs, rs) for hs, rs in zip(h.shape, r.shape)]) + rt_cost = numpy.prod([max(ts, rs) for ts, rs in zip(t.shape, r.shape)]) + t = torch.conj(t) + if hr_cost < rt_cost: + h = h * r + else: + t = r * t + return torch.real((h * t).sum(dim=-1)) + + def _complex_interaction_complex_native( h: torch.FloatTensor, r: torch.FloatTensor, From 2bc1becac83811f90ec95a369d9e59ab2ba5182d Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Wed, 16 Dec 2020 19:17:31 +0100 Subject: [PATCH 676/690] Update classvars --- src/pykeen/models/unimodal/complex.py | 12 ++++++------ src/pykeen/models/unimodal/conv_e.py | 2 +- src/pykeen/models/unimodal/conv_kb.py | 6 +++--- src/pykeen/models/unimodal/distmult.py | 6 +++--- src/pykeen/models/unimodal/ermlp.py | 4 ++-- src/pykeen/models/unimodal/ermlpe.py | 2 +- src/pykeen/models/unimodal/hole.py | 4 ++-- src/pykeen/models/unimodal/kg2e.py | 4 ++-- src/pykeen/models/unimodal/ntn.py | 4 ++-- src/pykeen/models/unimodal/proj_e.py | 10 +++++----- src/pykeen/models/unimodal/rescal.py | 6 +++--- src/pykeen/models/unimodal/rgcn.py | 4 ++-- src/pykeen/models/unimodal/rotate.py | 4 ++-- src/pykeen/models/unimodal/simple.py | 4 ++-- src/pykeen/models/unimodal/structured_embedding.py | 4 ++-- src/pykeen/models/unimodal/trans_d.py | 4 ++-- src/pykeen/models/unimodal/trans_e.py | 4 ++-- src/pykeen/models/unimodal/trans_h.py | 6 +++--- src/pykeen/models/unimodal/trans_r.py | 5 +++-- src/pykeen/models/unimodal/tucker.py | 2 +- src/pykeen/models/unimodal/unstructured_model.py | 4 ++-- 21 files changed, 51 insertions(+), 50 deletions(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 79d89efbd9..979905c8d1 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -2,7 +2,7 @@ """Implementation of the ComplEx model.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional, Type import torch import torch.nn as nn @@ -52,17 +52,17 @@ class ComplEx(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, ) #: The default loss function class - loss_default = SoftplusLoss + loss_default: Type[Loss] = SoftplusLoss #: The default parameters for the default loss function class - loss_default_kwargs = dict(reduction='mean') + loss_default_kwargs: ClassVar[Mapping[str, Any]] = dict(reduction='mean') #: The regularizer used by [trouillon2016]_ for ComplEx. - regularizer_default = LpRegularizer + regularizer_default: Type[Regularizer] = LpRegularizer #: The LP settings used by [trouillon2016]_ for ComplEx. - regularizer_default_kwargs = dict( + regularizer_default_kwargs: ClassVar[Mapping[str, Any]] = dict( weight=0.01, p=2.0, normalize=True, diff --git a/src/pykeen/models/unimodal/conv_e.py b/src/pykeen/models/unimodal/conv_e.py index 9fddaadd44..9a08a9effa 100644 --- a/src/pykeen/models/unimodal/conv_e.py +++ b/src/pykeen/models/unimodal/conv_e.py @@ -95,7 +95,7 @@ class ConvE(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( output_channels=dict(type=int, low=4, high=6, scale='power_two'), input_dropout=DEFAULT_DROPOUT_HPO_RANGE, output_dropout=DEFAULT_DROPOUT_HPO_RANGE, diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index 31d5890425..ff0204d742 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -3,7 +3,7 @@ """Implementation of the ConvKB model.""" import logging -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional from ..base import ERModel from ...constants import DEFAULT_DROPOUT_HPO_RANGE, DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE @@ -55,7 +55,7 @@ class ConvKB(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, hidden_dropout_rate=DEFAULT_DROPOUT_HPO_RANGE, num_filters=dict(type=int, low=7, high=9, scale='power_two'), @@ -63,7 +63,7 @@ class ConvKB(ERModel): #: The regularizer used by [nguyen2018]_ for ConvKB. regularizer_default = LpRegularizer #: The LP settings used by [nguyen2018]_ for ConvKB. - regularizer_default_kwargs = dict( + regularizer_default_kwargs: ClassVar[Mapping[str, Any]] = dict( weight=0.001 / 2, p=2.0, normalize=True, diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index 87ec5216e8..c9b9ac7b96 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -2,7 +2,7 @@ """Implementation of DistMult.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional from torch import nn from torch.nn import functional @@ -53,7 +53,7 @@ class DistMult(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, ) #: The regularizer used by [yang2014]_ for DistMult @@ -62,7 +62,7 @@ class DistMult(ERModel): #: why the weight has to be increased by a factor of 100 to have the same configuration as in the paper. regularizer_default = LpRegularizer #: The LP settings used by [yang2014]_ for DistMult - regularizer_default_kwargs = dict( + regularizer_default_kwargs: ClassVar[Mapping[str, Any]] = dict( weight=0.1, p=2.0, normalize=True, diff --git a/src/pykeen/models/unimodal/ermlp.py b/src/pykeen/models/unimodal/ermlp.py index 9c3e6bdeb4..91a499cb79 100644 --- a/src/pykeen/models/unimodal/ermlp.py +++ b/src/pykeen/models/unimodal/ermlp.py @@ -2,7 +2,7 @@ """Implementation of ERMLP.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional from ..base import ERModel from ...constants import DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE @@ -35,7 +35,7 @@ class ERMLP(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, ) diff --git a/src/pykeen/models/unimodal/ermlpe.py b/src/pykeen/models/unimodal/ermlpe.py index fbc9358021..39d05080be 100644 --- a/src/pykeen/models/unimodal/ermlpe.py +++ b/src/pykeen/models/unimodal/ermlpe.py @@ -41,7 +41,7 @@ class ERMLPE(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, hidden_dim=dict(type=int, low=5, high=9, scale='power_two'), input_dropout=DEFAULT_DROPOUT_HPO_RANGE, diff --git a/src/pykeen/models/unimodal/hole.py b/src/pykeen/models/unimodal/hole.py index 9819d4b191..a80dbb4587 100644 --- a/src/pykeen/models/unimodal/hole.py +++ b/src/pykeen/models/unimodal/hole.py @@ -2,7 +2,7 @@ """Implementation of the HolE model.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional from ..base import ERModel from ...constants import DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE @@ -48,7 +48,7 @@ class HolE(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, ) diff --git a/src/pykeen/models/unimodal/kg2e.py b/src/pykeen/models/unimodal/kg2e.py index 3809f08e0e..f001b3621b 100644 --- a/src/pykeen/models/unimodal/kg2e.py +++ b/src/pykeen/models/unimodal/kg2e.py @@ -2,7 +2,7 @@ """Implementation of KG2E.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional import torch import torch.autograd @@ -49,7 +49,7 @@ class KG2E(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, c_min=dict(type=float, low=0.01, high=0.1, scale='log'), c_max=dict(type=float, low=1.0, high=10.0), diff --git a/src/pykeen/models/unimodal/ntn.py b/src/pykeen/models/unimodal/ntn.py index ad80ecf2df..fb91085814 100644 --- a/src/pykeen/models/unimodal/ntn.py +++ b/src/pykeen/models/unimodal/ntn.py @@ -2,7 +2,7 @@ """Implementation of NTN.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional from torch import nn @@ -47,7 +47,7 @@ class NTN(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, num_slices=dict(type=int, low=2, high=4), ) diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index b1910e5272..e1e4f69491 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -2,13 +2,13 @@ """Implementation of ProjE.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional from torch import nn from ..base import ERModel from ...constants import DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE -from ...losses import Loss +from ...losses import BCEWithLogitsLoss, Loss from ...nn import EmbeddingSpecification from ...nn.init import xavier_uniform_ from ...nn.modules import ProjEInteraction @@ -46,13 +46,13 @@ class ProjE(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, ) #: The default loss function class - loss_default = nn.BCEWithLogitsLoss + loss_default = BCEWithLogitsLoss #: The default parameters for the default loss function class - loss_default_kwargs = dict(reduction='mean') + loss_default_kwargs: ClassVar[Mapping[str, Any]] = dict(reduction='mean') def __init__( self, diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index b42b469489..4b6cae37b6 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -2,7 +2,7 @@ """Implementation of RESCAL.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional from ..base import ERModel from ...constants import DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE @@ -37,7 +37,7 @@ class RESCAL(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, ) #: The regularizer used by [nickel2011]_ for for RESCAL @@ -45,7 +45,7 @@ class RESCAL(ERModel): #: a normalized weight of 10 is used. regularizer_default = LpRegularizer #: The LP settings used by [nickel2011]_ for for RESCAL - regularizer_default_kwargs = dict( + regularizer_default_kwargs: ClassVar[Mapping[str, Any]] = dict( weight=10, p=2., normalize=True, diff --git a/src/pykeen/models/unimodal/rgcn.py b/src/pykeen/models/unimodal/rgcn.py index bfd137d51b..b8436e92e2 100644 --- a/src/pykeen/models/unimodal/rgcn.py +++ b/src/pykeen/models/unimodal/rgcn.py @@ -2,7 +2,7 @@ """Implementation of the R-GCN model.""" -from typing import Any, Callable, Mapping, Optional, Type +from typing import Any, Callable, ClassVar, Mapping, Optional, Type import torch from torch import nn @@ -39,7 +39,7 @@ class RGCN(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=dict(type=int, low=16, high=1024, q=16), num_bases_or_blocks=dict(type=int, low=2, high=20, q=1), num_layers=dict(type=int, low=1, high=5, q=1), diff --git a/src/pykeen/models/unimodal/rotate.py b/src/pykeen/models/unimodal/rotate.py index 926d1c2aa2..cb02f68386 100644 --- a/src/pykeen/models/unimodal/rotate.py +++ b/src/pykeen/models/unimodal/rotate.py @@ -2,7 +2,7 @@ """Implementation of the RotatE model.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional import torch @@ -46,7 +46,7 @@ class RotatE(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=dict(type=int, low=32, high=1024, q=16), ) diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index bfd58328f5..43e4296b2d 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -46,7 +46,7 @@ class SimplE(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, ) #: The default loss function class @@ -58,7 +58,7 @@ class SimplE(ERModel): #: regularization term by the number of elements, which is 200. regularizer_default = PowerSumRegularizer #: The power sum settings used by [trouillon2016]_ for SimplE - regularizer_default_kwargs = dict( + regularizer_default_kwargs: ClassVar[Mapping[str, Any]] = dict( weight=20, p=2.0, normalize=True, diff --git a/src/pykeen/models/unimodal/structured_embedding.py b/src/pykeen/models/unimodal/structured_embedding.py index 2eaa623b58..878a7b20bb 100644 --- a/src/pykeen/models/unimodal/structured_embedding.py +++ b/src/pykeen/models/unimodal/structured_embedding.py @@ -3,7 +3,7 @@ """Implementation of structured model (SE).""" import functools -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional import numpy as np from torch import nn @@ -41,7 +41,7 @@ class StructuredEmbedding(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, scoring_fct_norm=dict(type=int, low=1, high=2), ) diff --git a/src/pykeen/models/unimodal/trans_d.py b/src/pykeen/models/unimodal/trans_d.py index 84ba171823..4ae8d264bc 100644 --- a/src/pykeen/models/unimodal/trans_d.py +++ b/src/pykeen/models/unimodal/trans_d.py @@ -2,7 +2,7 @@ """Implementation of TransD.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional from ..base import ERModel from ...constants import DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE @@ -52,7 +52,7 @@ class TransD(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, relation_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, ) diff --git a/src/pykeen/models/unimodal/trans_e.py b/src/pykeen/models/unimodal/trans_e.py index 817b071fac..3b9c545709 100644 --- a/src/pykeen/models/unimodal/trans_e.py +++ b/src/pykeen/models/unimodal/trans_e.py @@ -2,7 +2,7 @@ """TransE.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional from torch.nn import functional @@ -42,7 +42,7 @@ class TransE(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, scoring_fct_norm=dict(type=int, low=1, high=2), ) diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 47419045fd..92de855f3b 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -2,7 +2,7 @@ """An implementation of TransH.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional from torch.nn import functional @@ -52,14 +52,14 @@ class TransH(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, scoring_fct_norm=dict(type=int, low=1, high=2), ) #: The custom regularizer used by [wang2014]_ for TransH regularizer_default = TransHRegularizer #: The settings used by [wang2014]_ for TransH - regularizer_default_kwargs = dict( + regularizer_default_kwargs: ClassVar[Mapping[str, Any]] = dict( weight=0.05, epsilon=1e-5, ) diff --git a/src/pykeen/models/unimodal/trans_r.py b/src/pykeen/models/unimodal/trans_r.py index 789682753c..88c1b93c09 100644 --- a/src/pykeen/models/unimodal/trans_r.py +++ b/src/pykeen/models/unimodal/trans_r.py @@ -1,8 +1,9 @@ # -*- coding: utf-8 -*- """Implementation of TransR.""" + import logging -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional from torch.nn import functional @@ -55,7 +56,7 @@ class TransR(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, relation_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, scoring_fct_norm=dict(type=int, low=1, high=2), diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index c1be22f0da..77f9c1aeb8 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -51,7 +51,7 @@ class TuckER(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, relation_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, dropout_0=DEFAULT_DROPOUT_HPO_RANGE, diff --git a/src/pykeen/models/unimodal/unstructured_model.py b/src/pykeen/models/unimodal/unstructured_model.py index 4d6b7b4c71..cc39532f02 100644 --- a/src/pykeen/models/unimodal/unstructured_model.py +++ b/src/pykeen/models/unimodal/unstructured_model.py @@ -2,7 +2,7 @@ """Implementation of UM.""" -from typing import Optional +from typing import Any, ClassVar, Mapping, Optional from ..base import ERModel from ...constants import DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE @@ -37,7 +37,7 @@ class UnstructuredModel(ERModel): """ #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, scoring_fct_norm=dict(type=int, low=1, high=2), ) From 14d463ebd6c23f8fab822500bc390ce473ed2198 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Wed, 16 Dec 2020 19:20:37 +0100 Subject: [PATCH 677/690] More types --- src/pykeen/models/unimodal/complex.py | 4 ++-- src/pykeen/models/unimodal/conv_kb.py | 4 ++-- src/pykeen/models/unimodal/distmult.py | 4 ++-- src/pykeen/models/unimodal/proj_e.py | 4 ++-- src/pykeen/models/unimodal/rescal.py | 4 ++-- src/pykeen/models/unimodal/simple.py | 6 +++--- src/pykeen/models/unimodal/trans_h.py | 4 ++-- src/pykeen/models/unimodal/tucker.py | 4 ++-- tox.ini | 19 +++++++++---------- 9 files changed, 26 insertions(+), 27 deletions(-) diff --git a/src/pykeen/models/unimodal/complex.py b/src/pykeen/models/unimodal/complex.py index 979905c8d1..42f3c826f4 100644 --- a/src/pykeen/models/unimodal/complex.py +++ b/src/pykeen/models/unimodal/complex.py @@ -56,11 +56,11 @@ class ComplEx(ERModel): embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, ) #: The default loss function class - loss_default: Type[Loss] = SoftplusLoss + loss_default: ClassVar[Type[Loss]] = SoftplusLoss #: The default parameters for the default loss function class loss_default_kwargs: ClassVar[Mapping[str, Any]] = dict(reduction='mean') #: The regularizer used by [trouillon2016]_ for ComplEx. - regularizer_default: Type[Regularizer] = LpRegularizer + regularizer_default: ClassVar[Type[Regularizer]] = LpRegularizer #: The LP settings used by [trouillon2016]_ for ComplEx. regularizer_default_kwargs: ClassVar[Mapping[str, Any]] = dict( weight=0.01, diff --git a/src/pykeen/models/unimodal/conv_kb.py b/src/pykeen/models/unimodal/conv_kb.py index ff0204d742..605ee00a21 100644 --- a/src/pykeen/models/unimodal/conv_kb.py +++ b/src/pykeen/models/unimodal/conv_kb.py @@ -3,7 +3,7 @@ """Implementation of the ConvKB model.""" import logging -from typing import Any, ClassVar, Mapping, Optional +from typing import Any, ClassVar, Mapping, Optional, Type from ..base import ERModel from ...constants import DEFAULT_DROPOUT_HPO_RANGE, DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE @@ -61,7 +61,7 @@ class ConvKB(ERModel): num_filters=dict(type=int, low=7, high=9, scale='power_two'), ) #: The regularizer used by [nguyen2018]_ for ConvKB. - regularizer_default = LpRegularizer + regularizer_default: ClassVar[Type[Regularizer]] = LpRegularizer #: The LP settings used by [nguyen2018]_ for ConvKB. regularizer_default_kwargs: ClassVar[Mapping[str, Any]] = dict( weight=0.001 / 2, diff --git a/src/pykeen/models/unimodal/distmult.py b/src/pykeen/models/unimodal/distmult.py index c9b9ac7b96..5ca4faa38c 100644 --- a/src/pykeen/models/unimodal/distmult.py +++ b/src/pykeen/models/unimodal/distmult.py @@ -2,7 +2,7 @@ """Implementation of DistMult.""" -from typing import Any, ClassVar, Mapping, Optional +from typing import Any, ClassVar, Mapping, Optional, Type from torch import nn from torch.nn import functional @@ -60,7 +60,7 @@ class DistMult(ERModel): #: In the paper, they use weight of 0.0001, mini-batch-size of 10, and dimensionality of vector 100 #: Thus, when we use normalized regularization weight, the normalization factor is 10*sqrt(100) = 100, which is #: why the weight has to be increased by a factor of 100 to have the same configuration as in the paper. - regularizer_default = LpRegularizer + regularizer_default: ClassVar[Type[Regularizer]] = LpRegularizer #: The LP settings used by [yang2014]_ for DistMult regularizer_default_kwargs: ClassVar[Mapping[str, Any]] = dict( weight=0.1, diff --git a/src/pykeen/models/unimodal/proj_e.py b/src/pykeen/models/unimodal/proj_e.py index e1e4f69491..c33881a6a6 100644 --- a/src/pykeen/models/unimodal/proj_e.py +++ b/src/pykeen/models/unimodal/proj_e.py @@ -2,7 +2,7 @@ """Implementation of ProjE.""" -from typing import Any, ClassVar, Mapping, Optional +from typing import Any, ClassVar, Mapping, Optional, Type from torch import nn @@ -50,7 +50,7 @@ class ProjE(ERModel): embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, ) #: The default loss function class - loss_default = BCEWithLogitsLoss + loss_default: ClassVar[Type[Loss]] = BCEWithLogitsLoss #: The default parameters for the default loss function class loss_default_kwargs: ClassVar[Mapping[str, Any]] = dict(reduction='mean') diff --git a/src/pykeen/models/unimodal/rescal.py b/src/pykeen/models/unimodal/rescal.py index 4b6cae37b6..aba19774b2 100644 --- a/src/pykeen/models/unimodal/rescal.py +++ b/src/pykeen/models/unimodal/rescal.py @@ -2,7 +2,7 @@ """Implementation of RESCAL.""" -from typing import Any, ClassVar, Mapping, Optional +from typing import Any, ClassVar, Mapping, Optional, Type from ..base import ERModel from ...constants import DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE @@ -43,7 +43,7 @@ class RESCAL(ERModel): #: The regularizer used by [nickel2011]_ for for RESCAL #: According to https://github.com/mnick/rescal.py/blob/master/examples/kinships.py #: a normalized weight of 10 is used. - regularizer_default = LpRegularizer + regularizer_default: ClassVar[Type[Regularizer]] = LpRegularizer #: The LP settings used by [nickel2011]_ for for RESCAL regularizer_default_kwargs: ClassVar[Mapping[str, Any]] = dict( weight=10, diff --git a/src/pykeen/models/unimodal/simple.py b/src/pykeen/models/unimodal/simple.py index 43e4296b2d..56feb1313d 100644 --- a/src/pykeen/models/unimodal/simple.py +++ b/src/pykeen/models/unimodal/simple.py @@ -2,7 +2,7 @@ """Implementation of SimplE.""" -from typing import Any, ClassVar, Mapping, Optional, Tuple, Union +from typing import Any, ClassVar, Mapping, Optional, Tuple, Type, Union import torch @@ -50,13 +50,13 @@ class SimplE(ERModel): embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, ) #: The default loss function class - loss_default = SoftplusLoss + loss_default: ClassVar[Type[Loss]] = SoftplusLoss #: The default parameters for the default loss function class loss_default_kwargs: ClassVar[Mapping[str, Any]] = {} #: The regularizer used by [trouillon2016]_ for SimplE #: In the paper, they use weight of 0.1, and do not normalize the #: regularization term by the number of elements, which is 200. - regularizer_default = PowerSumRegularizer + regularizer_default: ClassVar[Type[Regularizer]] = PowerSumRegularizer #: The power sum settings used by [trouillon2016]_ for SimplE regularizer_default_kwargs: ClassVar[Mapping[str, Any]] = dict( weight=20, diff --git a/src/pykeen/models/unimodal/trans_h.py b/src/pykeen/models/unimodal/trans_h.py index 92de855f3b..f2cf2163e8 100644 --- a/src/pykeen/models/unimodal/trans_h.py +++ b/src/pykeen/models/unimodal/trans_h.py @@ -2,7 +2,7 @@ """An implementation of TransH.""" -from typing import Any, ClassVar, Mapping, Optional +from typing import Any, ClassVar, Mapping, Optional, Type from torch.nn import functional @@ -57,7 +57,7 @@ class TransH(ERModel): scoring_fct_norm=dict(type=int, low=1, high=2), ) #: The custom regularizer used by [wang2014]_ for TransH - regularizer_default = TransHRegularizer + regularizer_default: ClassVar[Type[Regularizer]] = TransHRegularizer #: The settings used by [wang2014]_ for TransH regularizer_default_kwargs: ClassVar[Mapping[str, Any]] = dict( weight=0.05, diff --git a/src/pykeen/models/unimodal/tucker.py b/src/pykeen/models/unimodal/tucker.py index 77f9c1aeb8..3bdd2bbcf3 100644 --- a/src/pykeen/models/unimodal/tucker.py +++ b/src/pykeen/models/unimodal/tucker.py @@ -2,7 +2,7 @@ """Implementation of TuckEr.""" -from typing import Any, ClassVar, Mapping, Optional +from typing import Any, ClassVar, Mapping, Optional, Type from ..base import ERModel from ...constants import DEFAULT_DROPOUT_HPO_RANGE, DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE @@ -59,7 +59,7 @@ class TuckER(ERModel): dropout_2=DEFAULT_DROPOUT_HPO_RANGE, ) #: The default loss function class - loss_default = BCEAfterSigmoidLoss + loss_default: ClassVar[Type[Loss]] = BCEAfterSigmoidLoss #: The default parameters for the default loss function class loss_default_kwargs: ClassVar[Mapping[str, Any]] = {} diff --git a/tox.ini b/tox.ini index d1a0b419b3..702e1dc71c 100644 --- a/tox.ini +++ b/tox.ini @@ -111,16 +111,15 @@ description = Check all python files do not have mistaken trailing commas [testenv:mypy] deps = mypy skip_install = true -commands = - mypy --ignore-missing-imports \ - src/pykeen/typing.py \ - src/pykeen/utils.py \ - src/pykeen/nn \ - src/pykeen/regularizers.py \ - src/pykeen/losses.py \ - src/pykeen/trackers \ - src/pykeen/models/base.py \ - src/pykeen/triples/triples_factory.py +commands = mypy --ignore-missing-imports \ + src/pykeen/typing.py \ + src/pykeen/utils.py \ + src/pykeen/nn \ + src/pykeen/regularizers.py \ + src/pykeen/losses.py \ + src/pykeen/trackers \ + src/pykeen/models/base.py \ + src/pykeen/triples/triples_factory.py description = Run the mypy tool to check static typing on the project. From 8982553352ef8c96df08bd9e17141ebdf87fa93f Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Wed, 16 Dec 2020 20:54:24 +0100 Subject: [PATCH 678/690] More typing --- src/pykeen/models/multimodal/complex_literal.py | 6 +++--- src/pykeen/models/multimodal/distmult_literal.py | 6 +++--- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/src/pykeen/models/multimodal/complex_literal.py b/src/pykeen/models/multimodal/complex_literal.py index e5b555c72e..d1324f1620 100644 --- a/src/pykeen/models/multimodal/complex_literal.py +++ b/src/pykeen/models/multimodal/complex_literal.py @@ -2,7 +2,7 @@ """Implementation of the ComplexLiteral model.""" -from typing import Any, ClassVar, Mapping, Optional +from typing import Any, ClassVar, Mapping, Optional, Type import torch import torch.nn as nn @@ -61,12 +61,12 @@ class ComplExLiteral(LiteralModel): """An implementation of ComplexLiteral from [agustinus2018]_ based on the LCWA training approach.""" #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, input_dropout=DEFAULT_DROPOUT_HPO_RANGE, ) #: The default loss function class - loss_default = BCEWithLogitsLoss + loss_default: ClassVar[Type[Loss]] = BCEWithLogitsLoss #: The default parameters for the default loss function class loss_default_kwargs: ClassVar[Mapping[str, Any]] = {} diff --git a/src/pykeen/models/multimodal/distmult_literal.py b/src/pykeen/models/multimodal/distmult_literal.py index 1bf1960c4f..76dccb93f8 100644 --- a/src/pykeen/models/multimodal/distmult_literal.py +++ b/src/pykeen/models/multimodal/distmult_literal.py @@ -2,7 +2,7 @@ """Implementation of the DistMultLiteral model.""" -from typing import Optional, TYPE_CHECKING +from typing import Any, ClassVar, Mapping, Optional, TYPE_CHECKING import torch.nn as nn @@ -28,12 +28,12 @@ class DistMultLiteral(LiteralModel): """An implementation of DistMultLiteral from [agustinus2018]_.""" #: The default strategy for optimizing the model's hyper-parameters - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( embedding_dim=DEFAULT_EMBEDDING_HPO_EMBEDDING_DIM_RANGE, input_dropout=DEFAULT_DROPOUT_HPO_RANGE, ) #: The default parameters for the default loss function class - loss_default_kwargs = dict(margin=0.0) + loss_default_kwargs: ClassVar[Mapping[str, Any]] = dict(margin=0.0) def __init__( self, From b404b84c751e8c316665614e0b22c479dfb7fb9b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 17 Dec 2020 07:55:27 +0100 Subject: [PATCH 679/690] Add size information data class --- src/pykeen/nn/functional.py | 38 ++++++++++++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 89f981d10c..4603eba960 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -7,7 +7,8 @@ representations are provided in shape (batch_size, num_heads, ``*``), (batch_size, num_relations, ``*``), and (batch_size, num_tails, ``*``), and return a score tensor of shape (batch_size, num_heads, num_relations, num_tails). """ - +import dataclasses +from dataclasses import dataclass from typing import Optional, Tuple, Union import numpy @@ -46,6 +47,41 @@ ] +@dataclasses.dataclass +class SizeInformation: + """Size information of generic score function.""" + + #: The batch size of the head representations. + bh: int + + #: The number of head representations per batch + nh: int + + #: The batch size of the relation representations. + br: int + + #: The number of relation representations per batch + nr: int + + #: The batch size of the tail representations. + bt: int + + #: The number of tail representations per batch + nt: int + + +def _extract_size_information( + h: torch.Tensor, + r: torch.Tensor, + t: torch.Tensor, +) -> SizeInformation: + """Extract size information from tensors.""" + bh, nh = h.shape[:2] + br, nr = r.shape[:2] + bt, nt = t.shape[:2] + return SizeInformation(bh=bh, nh=nh, br=br, nr=nr, bt=bt, nt=nt) + + def _extract_sizes( h: torch.Tensor, r: torch.Tensor, From 3d941a90baefcdecf4c28ea47057b44a3f959461 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 17 Dec 2020 08:28:28 +0100 Subject: [PATCH 680/690] use utility for ER-MLP --- src/pykeen/nn/functional.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 4603eba960..708f1af4dd 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -8,7 +8,6 @@ (batch_size, num_tails, ``*``), and return a score tensor of shape (batch_size, num_heads, num_relations, num_tails). """ import dataclasses -from dataclasses import dataclass from typing import Optional, Tuple, Union import numpy @@ -69,6 +68,11 @@ class SizeInformation: #: The number of tail representations per batch nt: int + @property + def same(self) -> bool: + """Whether all representations have the same shape.""" + return self.bh == self.br and self.bh == self.bt and self.nh == self.nr and self.nh == self.nt + def _extract_size_information( h: torch.Tensor, @@ -373,9 +377,6 @@ def convkb_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - # bind sizes - num_heads, num_relations, num_tails, embedding_dim, _ = _extract_sizes(h, r, t) - # decompose convolution for faster computation in 1-n case num_filters = conv.weight.shape[0] assert conv.weight.shape == (num_filters, 1, 1, 3) @@ -408,7 +409,7 @@ def convkb_interaction( x = hidden_dropout(x) # Linear layer for final scores; use flattened representations, shape: (b, h, r, t, d * f) - x = x.view(*x.shape[:-2], embedding_dim * num_filters) + x = x.view(*x.shape[:-2], -1) x = linear(x) return x.squeeze(dim=-1) @@ -461,14 +462,21 @@ def ermlp_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - num_heads, num_relations, num_tails, embedding_dim, _ = _extract_sizes(h, r, t) + sizes = _extract_size_information(h, r, t) + + # same shape + if sizes.same: + return final(activation( + hidden(torch.cat([h, r, t], dim=-1).view(-1, 3 * h.shape[-1]))) + ).view(sizes.bh, sizes.nh, sizes.nr, sizes.nt) + hidden_dim = hidden.weight.shape[0] # split, shape: (embedding_dim, hidden_dim) - head_to_hidden, rel_to_hidden, tail_to_hidden = hidden.weight.t().split(embedding_dim) + head_to_hidden, rel_to_hidden, tail_to_hidden = hidden.weight.t().split(h.shape[-1]) bias = hidden.bias.view(1, 1, 1, 1, -1) - h = h @ head_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) - r = r @ rel_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) - t = t @ tail_to_hidden.view(1, 1, 1, embedding_dim, hidden_dim) + h = h @ head_to_hidden.view(1, 1, 1, h.shape[-1], hidden_dim) + r = r @ rel_to_hidden.view(1, 1, 1, r.shape[-1], hidden_dim) + t = t @ tail_to_hidden.view(1, 1, 1, t.shape[-1], hidden_dim) return final(activation(tensor_sum(bias, h, r, t))).squeeze(dim=-1) From ca234afc5c226473dd42fdc54788484bb25f8ee8 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 17 Dec 2020 08:28:58 +0100 Subject: [PATCH 681/690] Update docstring --- src/pykeen/nn/functional.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 708f1af4dd..23fd4b1130 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -4,8 +4,9 @@ Functional forms of interaction methods. The functional forms always assume the general form of the interaction function, where head, relation and tail -representations are provided in shape (batch_size, num_heads, ``*``), (batch_size, num_relations, ``*``), and -(batch_size, num_tails, ``*``), and return a score tensor of shape (batch_size, num_heads, num_relations, num_tails). +representations are provided in shape (batch_size, num_heads, 1, 1, ``*``), (batch_size, 1, num_relations, 1, ``*``), +and (batch_size, 1, 1, num_tails, ``*``), and return a score tensor of shape +(batch_size, num_heads, num_relations, num_tails). """ import dataclasses from typing import Optional, Tuple, Union From 7d137c790a3dfade75b3b26290980a676f962750 Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 17 Dec 2020 08:33:05 +0100 Subject: [PATCH 682/690] Move variants to module --- benchmarking/interactions.py | 11 +-- src/pykeen/nn/compute_kernel.py | 132 +++++++++++++++++++++++++++++++ src/pykeen/nn/functional.py | 134 +------------------------------- 3 files changed, 139 insertions(+), 138 deletions(-) diff --git a/benchmarking/interactions.py b/benchmarking/interactions.py index 3eb9ef3a05..b3e8b24013 100644 --- a/benchmarking/interactions.py +++ b/benchmarking/interactions.py @@ -9,10 +9,7 @@ from tqdm import tqdm from pykeen.nn import Interaction -from pykeen.nn.functional import ( - _complex_interaction_complex_native, _complex_interaction_direct, - _complex_interaction_optimized_broadcasted, _complex_select, _complex_stacked, _complex_stacked_select, -) +from pykeen.nn.compute_kernel import _complex_native_complex, _complex_direct, _complex_broadcast_optimized, _complex_select, _complex_stacked, _complex_stacked_select from pykeen.typing import HeadRepresentation, RelationRepresentation, TailRepresentation from pykeen.utils import unpack_singletons from pykeen.version import get_git_hash @@ -133,9 +130,9 @@ def main( print(f"Running on {device}.") variants = [ _complex_select, - _complex_interaction_complex_native, - _complex_interaction_optimized_broadcasted, - _complex_interaction_direct, + _complex_native_complex, + _complex_broadcast_optimized, + _complex_direct, _complex_stacked, _complex_stacked_select, ] diff --git a/src/pykeen/nn/compute_kernel.py b/src/pykeen/nn/compute_kernel.py index 93835e2ce9..c9e5b1d5ad 100644 --- a/src/pykeen/nn/compute_kernel.py +++ b/src/pykeen/nn/compute_kernel.py @@ -1,6 +1,9 @@ """Compute kernels for common sub-tasks.""" +import numpy import torch +from pykeen.utils import extended_einsum, split_complex, tensor_product, view_complex + def _batched_dot_manual( a: torch.FloatTensor, @@ -29,3 +32,132 @@ def batched_dot( ) -> torch.FloatTensor: """Compute "element-wise" dot-product between batched vectors.""" return _batched_dot_manual(a, b) + + +def _complex_broadcast_optimized( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Manually split into real/imag, and used optimized broadcasted combination.""" + (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] + return sum( + factor * tensor_product(hh, rr, tt).sum(dim=-1) + for factor, hh, rr, tt in [ + (+1, h_re, r_re, t_re), + (+1, h_re, r_im, t_im), + (+1, h_im, r_re, t_im), + (-1, h_im, r_im, t_re), + ] + ) + + +def _complex_direct( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Manually split into real/imag, and directly evaluate interaction.""" + (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] + return ( + (h_re * r_re * t_re).sum(dim=-1) + + (h_re * r_im * t_im).sum(dim=-1) + + (h_im * r_re * t_im).sum(dim=-1) + - (h_im * r_im * t_re).sum(dim=-1) + ) + + +def _complex_einsum( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Use einsum.""" + x = h.new_zeros(2, 2, 2) + x[0, 0, 0] = 1 + x[0, 1, 1] = 1 + x[1, 0, 1] = 1 + x[1, 1, 0] = -1 + return extended_einsum( + "ijk,bhdi,brdj,btdk->bhrt", + x, + h.view(*h.shape[:-1], -1, 2), + r.view(*r.shape[:-1], -1, 2), + t.view(*t.shape[:-1], -1, 2), + ) + + +def _complex_native_complex( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Use torch built-ins for computation with complex numbers.""" + h, r, t = [view_complex(x=x) for x in (h, r, t)] + return torch.real(tensor_product(h, r, torch.conj(t)).sum(dim=-1)) + + +def _complex_native_complex_select( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Use torch built-ins for computation with complex numbers and select whether to combine hr or ht first.""" + h, r, t = [view_complex(x=x) for x in (h, r, t)] + hr_cost = numpy.prod([max(hs, rs) for hs, rs in zip(h.shape, r.shape)]) + rt_cost = numpy.prod([max(ts, rs) for ts, rs in zip(t.shape, r.shape)]) + t = torch.conj(t) + if hr_cost < rt_cost: + h = h * r + else: + t = r * t + return torch.real((h * t).sum(dim=-1)) + + +def _complex_select( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Decide based on result shape whether to combine hr or ht first.""" + hr_cost = numpy.prod([max(hs, rs) for hs, rs in zip(h.shape, r.shape)]) + rt_cost = numpy.prod([max(ts, rs) for ts, rs in zip(t.shape, r.shape)]) + (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] + if hr_cost < rt_cost: + h_re, h_im = (h_re * r_re - h_im * r_im), (h_re * r_im + h_im * r_re) + else: + t_re, t_im = (t_re * r_re - t_im * r_im), (t_re * r_im + t_im * r_re) + return h_re @ t_re.transpose(-2, -1) - h_im @ t_im.transpose(-2, -1) + + +def _complex_stacked( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Stack vectors.""" + (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (r, t)] + h = torch.cat([h, h], dim=-1) # re im re im + r = torch.cat([r_re, r_re, r_im, r_im], dim=-1) # re re im im + t = torch.cat([t_re, t_im, t_im, t_re], dim=-1) # re im im re + return (h * r * t).sum(dim=-1) + + +def _complex_stacked_select( + h: torch.FloatTensor, + r: torch.FloatTensor, + t: torch.FloatTensor, +) -> torch.FloatTensor: + """Stack vectors and select order.""" + (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (r, t)] + h = torch.cat([h, h], dim=-1) # re im re im + r = torch.cat([r_re, r_re, r_im, r_im], dim=-1) # re re im im + t = torch.cat([t_re, t_im, t_im, t_re], dim=-1) # re im im re + hr_cost = numpy.prod([max(hs, rs) for hs, rs in zip(h.shape, r.shape)]) + rt_cost = numpy.prod([max(ts, rs) for ts, rs in zip(t.shape, r.shape)]) + if hr_cost < rt_cost: + # h = h_re, -h_im + h = h * r + else: + t = r * t + return h @ t.transpose(-2, -1) diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 23fd4b1130..434706a01a 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -16,11 +16,12 @@ import torch.fft from torch import nn +from .compute_kernel import _complex_native_complex from .sim import KG2E_SIMILARITIES from ..typing import GaussianDistribution from ..utils import ( broadcast_cat, clamp_norm, estimate_cost_of_sequence, extended_einsum, is_cudnn_error, negative_norm, - negative_norm_of_sum, project_entity, split_complex, tensor_product, tensor_sum, view_complex, + negative_norm_of_sum, project_entity, tensor_product, tensor_sum, view_complex, ) __all__ = [ @@ -130,135 +131,6 @@ def wrapped(*args, **kwargs): return wrapped -def _complex_select( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, -) -> torch.FloatTensor: - """Decide based on result shape whether to combine hr or ht first.""" - hr_cost = numpy.prod([max(hs, rs) for hs, rs in zip(h.shape, r.shape)]) - rt_cost = numpy.prod([max(ts, rs) for ts, rs in zip(t.shape, r.shape)]) - (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] - if hr_cost < rt_cost: - h_re, h_im = (h_re * r_re - h_im * r_im), (h_re * r_im + h_im * r_re) - else: - t_re, t_im = (t_re * r_re - t_im * r_im), (t_re * r_im + t_im * r_re) - return h_re @ t_re.transpose(-2, -1) - h_im @ t_im.transpose(-2, -1) - - -def _complex_native_select( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, -) -> torch.FloatTensor: - """Use torch built-ins for computation with complex numbers and select whether to combine hr or ht first.""" - h, r, t = [view_complex(x=x) for x in (h, r, t)] - hr_cost = numpy.prod([max(hs, rs) for hs, rs in zip(h.shape, r.shape)]) - rt_cost = numpy.prod([max(ts, rs) for ts, rs in zip(t.shape, r.shape)]) - t = torch.conj(t) - if hr_cost < rt_cost: - h = h * r - else: - t = r * t - return torch.real((h * t).sum(dim=-1)) - - -def _complex_interaction_complex_native( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, -) -> torch.FloatTensor: - """Use torch built-ins for computation with complex numbers.""" - h, r, t = [view_complex(x=x) for x in (h, r, t)] - return torch.real(tensor_product(h, r, torch.conj(t)).sum(dim=-1)) - - -def _complex_interaction_optimized_broadcasted( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, -) -> torch.FloatTensor: - """Manually split into real/imag, and used optimized broadcasted combination.""" - (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] - return sum( - factor * tensor_product(hh, rr, tt).sum(dim=-1) - for factor, hh, rr, tt in [ - (+1, h_re, r_re, t_re), - (+1, h_re, r_im, t_im), - (+1, h_im, r_re, t_im), - (-1, h_im, r_im, t_re), - ] - ) - - -def _complex_interaction_direct( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, -) -> torch.FloatTensor: - """Manually split into real/imag, and directly evaluate interaction.""" - (h_re, h_im), (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (h, r, t)] - return ( - (h_re * r_re * t_re).sum(dim=-1) - + (h_re * r_im * t_im).sum(dim=-1) - + (h_im * r_re * t_im).sum(dim=-1) - - (h_im * r_im * t_re).sum(dim=-1) - ) - - -def _complex_stacked( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, -) -> torch.FloatTensor: - """Stack vectors.""" - (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (r, t)] - h = torch.cat([h, h], dim=-1) # re im re im - r = torch.cat([r_re, r_re, r_im, r_im], dim=-1) # re re im im - t = torch.cat([t_re, t_im, t_im, t_re], dim=-1) # re im im re - return (h * r * t).sum(dim=-1) - - -def _complex_stacked_select( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, -) -> torch.FloatTensor: - """Stack vectors and select order.""" - (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (r, t)] - h = torch.cat([h, h], dim=-1) # re im re im - r = torch.cat([r_re, r_re, r_im, r_im], dim=-1) # re re im im - t = torch.cat([t_re, t_im, t_im, t_re], dim=-1) # re im im re - hr_cost = numpy.prod([max(hs, rs) for hs, rs in zip(h.shape, r.shape)]) - rt_cost = numpy.prod([max(ts, rs) for ts, rs in zip(t.shape, r.shape)]) - if hr_cost < rt_cost: - # h = h_re, -h_im - h = h * r - else: - t = r * t - return h @ t.transpose(-2, -1) - - -def _complex_einsum( - h: torch.FloatTensor, - r: torch.FloatTensor, - t: torch.FloatTensor, -) -> torch.FloatTensor: - """Use einsum.""" - x = h.new_zeros(2, 2, 2) - x[0, 0, 0] = 1 - x[0, 1, 1] = 1 - x[1, 0, 1] = 1 - x[1, 1, 0] = -1 - return extended_einsum( - "ijk,bhdi,brdj,btdk->bhrt", - x, - h.view(*h.shape[:-1], -1, 2), - r.view(*r.shape[:-1], -1, 2), - t.view(*t.shape[:-1], -1, 2), - ) - - def complex_interaction( h: torch.FloatTensor, r: torch.FloatTensor, @@ -280,7 +152,7 @@ def complex_interaction( :return: shape: (batch_size, num_heads, num_relations, num_tails) The scores. """ - return _complex_interaction_complex_native(h, r, t) + return _complex_native_complex(h, r, t) @_add_cuda_warning From c9b240e89f666400c0ce4dfbc125a8980013887b Mon Sep 17 00:00:00 2001 From: Max Berrendorf Date: Thu, 17 Dec 2020 08:39:38 +0100 Subject: [PATCH 683/690] extract common code --- src/pykeen/nn/compute_kernel.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/pykeen/nn/compute_kernel.py b/src/pykeen/nn/compute_kernel.py index c9e5b1d5ad..2635998381 100644 --- a/src/pykeen/nn/compute_kernel.py +++ b/src/pykeen/nn/compute_kernel.py @@ -130,16 +130,21 @@ def _complex_select( return h_re @ t_re.transpose(-2, -1) - h_im @ t_im.transpose(-2, -1) +def _complex_to_stacked(h, r, t): + (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (r, t)] + h = torch.cat([h, h], dim=-1) # re im re im + r = torch.cat([r_re, r_re, r_im, r_im], dim=-1) # re re im im + t = torch.cat([t_re, t_im, t_im, t_re], dim=-1) # re im im re + return h, r, t + + def _complex_stacked( h: torch.FloatTensor, r: torch.FloatTensor, t: torch.FloatTensor, ) -> torch.FloatTensor: """Stack vectors.""" - (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (r, t)] - h = torch.cat([h, h], dim=-1) # re im re im - r = torch.cat([r_re, r_re, r_im, r_im], dim=-1) # re re im im - t = torch.cat([t_re, t_im, t_im, t_re], dim=-1) # re im im re + h, r, t = _complex_to_stacked(h, r, t) return (h * r * t).sum(dim=-1) @@ -149,10 +154,7 @@ def _complex_stacked_select( t: torch.FloatTensor, ) -> torch.FloatTensor: """Stack vectors and select order.""" - (r_re, r_im), (t_re, t_im) = [split_complex(x=x) for x in (r, t)] - h = torch.cat([h, h], dim=-1) # re im re im - r = torch.cat([r_re, r_re, r_im, r_im], dim=-1) # re re im im - t = torch.cat([t_re, t_im, t_im, t_re], dim=-1) # re im im re + h, r, t = _complex_to_stacked(h, r, t) hr_cost = numpy.prod([max(hs, rs) for hs, rs in zip(h.shape, r.shape)]) rt_cost = numpy.prod([max(ts, rs) for ts, rs in zip(t.shape, r.shape)]) if hr_cost < rt_cost: From 3bf1ab9e199a6aaed9c1831cfdb443365f81d773 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 17 Dec 2020 14:18:04 +0100 Subject: [PATCH 684/690] Code cleanup --- src/pykeen/nn/compute_kernel.py | 3 ++ src/pykeen/nn/functional.py | 8 ++++- src/pykeen/nn/sim.py | 1 + src/pykeen/testing/base.py | 17 ++++++++--- tests/test_interactions.py | 54 ++++++++++++++++----------------- tests/test_nn.py | 20 ++++++------ 6 files changed, 61 insertions(+), 42 deletions(-) diff --git a/src/pykeen/nn/compute_kernel.py b/src/pykeen/nn/compute_kernel.py index 2635998381..f56fe99c12 100644 --- a/src/pykeen/nn/compute_kernel.py +++ b/src/pykeen/nn/compute_kernel.py @@ -1,4 +1,7 @@ +# -*- coding: utf-8 -*- + """Compute kernels for common sub-tasks.""" + import numpy import torch diff --git a/src/pykeen/nn/functional.py b/src/pykeen/nn/functional.py index 434706a01a..968c0eaf13 100644 --- a/src/pykeen/nn/functional.py +++ b/src/pykeen/nn/functional.py @@ -8,6 +8,7 @@ and (batch_size, 1, 1, num_tails, ``*``), and return a score tensor of shape (batch_size, num_heads, num_relations, num_tails). """ + import dataclasses from typing import Optional, Tuple, Union @@ -73,7 +74,12 @@ class SizeInformation: @property def same(self) -> bool: """Whether all representations have the same shape.""" - return self.bh == self.br and self.bh == self.bt and self.nh == self.nr and self.nh == self.nt + return ( + self.bh == self.br + and self.bh == self.bt + and self.nh == self.nr + and self.nh == self.nt + ) def _extract_size_information( diff --git a/src/pykeen/nn/sim.py b/src/pykeen/nn/sim.py index 21df717abe..949d759283 100644 --- a/src/pykeen/nn/sim.py +++ b/src/pykeen/nn/sim.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- """Similarity functions.""" + import itertools import math diff --git a/src/pykeen/testing/base.py b/src/pykeen/testing/base.py index e34c94493c..8c037a2e61 100644 --- a/src/pykeen/testing/base.py +++ b/src/pykeen/testing/base.py @@ -1,12 +1,21 @@ +# -*- coding: utf-8 -*- + """Base classes for simplified testing.""" + +import unittest from typing import Any, Collection, Generic, Mapping, MutableMapping, Optional, Type, TypeVar -from pykeen.utils import get_subclasses, set_random_seed +from ..utils import get_subclasses, set_random_seed + +__all__ = [ + 'GenericTestCase', + 'TestsTestCase', +] T = TypeVar("T") -class GenericTests(Generic[T]): +class GenericTestCase(Generic[T]): """Generic tests.""" cls: Type[T] @@ -30,11 +39,11 @@ def post_instantiation_hook(self) -> None: """Perform actions after instantiation.""" -class TestsTest(Generic[T]): +class TestsTestCase(Generic[T], unittest.TestCase): """A generic test for tests.""" base_cls: Type[T] - base_test: Type[GenericTests[T]] + base_test: Type[GenericTestCase[T]] skip_cls: Collection[T] = tuple() def test_testing(self): diff --git a/tests/test_interactions.py b/tests/test_interactions.py index 1f9920c60f..1e9e294b97 100644 --- a/tests/test_interactions.py +++ b/tests/test_interactions.py @@ -23,7 +23,7 @@ logger = logging.getLogger(__name__) -class InteractionTests(ptb.GenericTests[pykeen.nn.modules.Interaction]): +class InteractionTestCase(ptb.GenericTestCase[pykeen.nn.modules.Interaction]): """Generic test for interaction functions.""" dim: int = 2 @@ -254,7 +254,7 @@ def _exp_score(self, **kwargs) -> torch.FloatTensor: raise NotImplementedError(f"{self.cls.__name__}({sorted(kwargs.keys())})") -class ComplExTests(InteractionTests, unittest.TestCase): +class ComplExTests(InteractionTestCase, unittest.TestCase): """Tests for ComplEx interaction function.""" cls = pykeen.nn.modules.ComplExInteraction @@ -264,7 +264,7 @@ def _exp_score(self, h, r, t) -> torch.FloatTensor: # noqa: D102 return (h * r * torch.conj(t)).sum().real -class ConvETests(InteractionTests, unittest.TestCase): +class ConvETests(InteractionTestCase, unittest.TestCase): """Tests for ConvE interaction function.""" cls = pykeen.nn.modules.ConvEInteraction @@ -273,7 +273,7 @@ class ConvETests(InteractionTests, unittest.TestCase): embedding_width=2, kernel_height=2, kernel_width=1, - embedding_dim=InteractionTests.dim, + embedding_dim=InteractionTestCase.dim, ) def _get_hrt( @@ -298,13 +298,13 @@ def _exp_score( return (x.view(1, -1) * t.view(1, -1)).sum() + t_bias -class ConvKBTests(InteractionTests, unittest.TestCase): +class ConvKBTests(InteractionTestCase, unittest.TestCase): """Tests for ConvKB interaction function.""" cls = pykeen.nn.modules.ConvKBInteraction kwargs = dict( - embedding_dim=InteractionTests.dim, - num_filters=2 * InteractionTests.dim - 1, + embedding_dim=InteractionTestCase.dim, + num_filters=2 * InteractionTestCase.dim - 1, ) def _exp_score(self, h, r, t, conv, activation, hidden_dropout, linear) -> torch.FloatTensor: # noqa: D102 @@ -316,7 +316,7 @@ def _exp_score(self, h, r, t, conv, activation, hidden_dropout, linear) -> torch return linear(x.view(1, -1)) -class DistMultTests(InteractionTests, unittest.TestCase): +class DistMultTests(InteractionTestCase, unittest.TestCase): """Tests for DistMult interaction function.""" cls = pykeen.nn.modules.DistMultInteraction @@ -325,13 +325,13 @@ def _exp_score(self, h, r, t) -> torch.FloatTensor: return (h * r * t).sum(dim=-1) -class ERMLPTests(InteractionTests, unittest.TestCase): +class ERMLPTests(InteractionTestCase, unittest.TestCase): """Tests for ERMLP interaction function.""" cls = pykeen.nn.modules.ERMLPInteraction kwargs = dict( - embedding_dim=InteractionTests.dim, - hidden_dim=2 * InteractionTests.dim - 1, + embedding_dim=InteractionTestCase.dim, + hidden_dim=2 * InteractionTestCase.dim - 1, ) def _exp_score(self, h, r, t, hidden, activation, final) -> torch.FloatTensor: @@ -339,13 +339,13 @@ def _exp_score(self, h, r, t, hidden, activation, final) -> torch.FloatTensor: return final(activation(hidden(x))) -class ERMLPETests(InteractionTests, unittest.TestCase): +class ERMLPETests(InteractionTestCase, unittest.TestCase): """Tests for ERMLP-E interaction function.""" cls = pykeen.nn.modules.ERMLPEInteraction kwargs = dict( - embedding_dim=InteractionTests.dim, - hidden_dim=2 * InteractionTests.dim - 1, + embedding_dim=InteractionTestCase.dim, + hidden_dim=2 * InteractionTestCase.dim - 1, ) def _exp_score(self, h, r, t, mlp) -> torch.FloatTensor: # noqa: D102 @@ -353,7 +353,7 @@ def _exp_score(self, h, r, t, mlp) -> torch.FloatTensor: # noqa: D102 return mlp(x).view(1, -1) @ t.view(-1, 1) -class HolETests(InteractionTests, unittest.TestCase): +class HolETests(InteractionTestCase, unittest.TestCase): """Tests for HolE interaction function.""" cls = pykeen.nn.modules.HolEInteraction @@ -365,7 +365,7 @@ def _exp_score(self, h, r, t) -> torch.FloatTensor: # noqa: D102 return (c * r).sum() -class NTNTests(InteractionTests, unittest.TestCase): +class NTNTests(InteractionTestCase, unittest.TestCase): """Tests for NTN interaction function.""" cls = pykeen.nn.modules.NTNInteraction @@ -389,12 +389,12 @@ def _exp_score(self, h, t, w, vt, vh, b, u, activation) -> torch.FloatTensor: return score -class ProjETests(InteractionTests, unittest.TestCase): +class ProjETests(InteractionTestCase, unittest.TestCase): """Tests for ProjE interaction function.""" cls = pykeen.nn.modules.ProjEInteraction kwargs = dict( - embedding_dim=InteractionTests.dim, + embedding_dim=InteractionTestCase.dim, ) def _exp_score(self, h, r, t, d_e, d_r, b_c, b_p, activation) -> torch.FloatTensor: @@ -403,7 +403,7 @@ def _exp_score(self, h, r, t, d_e, d_r, b_c, b_p, activation) -> torch.FloatTens return (t * activation((d_e * h) + (d_r * r) + b_c)).sum() + b_p -class RESCALTests(InteractionTests, unittest.TestCase): +class RESCALTests(InteractionTestCase, unittest.TestCase): """Tests for RESCAL interaction function.""" cls = pykeen.nn.modules.RESCALInteraction @@ -414,7 +414,7 @@ def _exp_score(self, h, r, t) -> torch.FloatTensor: return h.view(1, -1) @ r @ t.view(-1, 1) -class KG2ETests(InteractionTests, unittest.TestCase): +class KG2ETests(InteractionTestCase, unittest.TestCase): """Tests for KG2E interaction function.""" cls = pykeen.nn.modules.KG2EInteraction @@ -428,12 +428,12 @@ def _exp_score(self, exact, h_mean, h_var, r_mean, r_var, similarity, t_mean, t_ return -torch.distributions.kl.kl_divergence(p, q) -class TuckerTests(InteractionTests, unittest.TestCase): +class TuckerTests(InteractionTestCase, unittest.TestCase): """Tests for Tucker interaction function.""" cls = pykeen.nn.modules.TuckerInteraction kwargs = dict( - embedding_dim=InteractionTests.dim, + embedding_dim=InteractionTestCase.dim, ) def _exp_score(self, bn_h, bn_hr, core_tensor, do_h, do_r, do_hr, h, r, t) -> torch.FloatTensor: @@ -446,7 +446,7 @@ def _exp_score(self, bn_h, bn_hr, core_tensor, do_h, do_r, do_hr, h, r, t) -> to return (d * t[None, None, :]).sum() -class RotatETests(InteractionTests, unittest.TestCase): +class RotatETests(InteractionTestCase, unittest.TestCase): """Tests for RotatE interaction function.""" cls = pykeen.nn.modules.RotatEInteraction @@ -467,7 +467,7 @@ def _exp_score(self, h, r, t) -> torch.FloatTensor: # noqa: D102 return -(d.abs() ** 2).sum(dim=-1).sqrt() -class TranslationalInteractionTests(InteractionTests): +class TranslationalInteractionTests(InteractionTestCase): """Common tests for translational interaction.""" kwargs = dict( @@ -596,7 +596,7 @@ def _exp_score(self, h, t, p, power_norm) -> torch.FloatTensor: return -(h - t).pow(p).sum() -class SimplEInteractionTests(InteractionTests, unittest.TestCase): +class SimplEInteractionTests(InteractionTestCase, unittest.TestCase): """Tests for SimplE interaction function.""" cls = pykeen.nn.modules.SimplEInteraction @@ -607,11 +607,11 @@ def _exp_score(self, h, r, t, h_inv, r_inv, t_inv, clamp) -> torch.FloatTensor: return 0.5 * distmult_interaction(h, r, t) + 0.5 * distmult_interaction(h_inv, r_inv, t_inv) -class InteractionTestsTest(ptb.TestsTest[Interaction], unittest.TestCase): +class InteractionTestsTestCase(ptb.TestsTestCase[Interaction]): """Test for tests for all interaction functions.""" base_cls = Interaction - base_test = InteractionTests + base_test = InteractionTestCase skip_cls = { TranslationalInteraction, LiteralInteraction, diff --git a/tests/test_nn.py b/tests/test_nn.py index ff6981e580..0ed640ff96 100644 --- a/tests/test_nn.py +++ b/tests/test_nn.py @@ -23,7 +23,7 @@ from pykeen.typing import GaussianDistribution -class RepresentationModuleTests(ptb.GenericTests[RepresentationModule]): +class RepresentationModuleTestCase(ptb.GenericTestCase[RepresentationModule]): """Tests for RepresentationModule.""" #: The batch size @@ -174,13 +174,13 @@ def _check_call( return call_count -class EmbeddingTests(RepresentationModuleTests, unittest.TestCase): +class EmbeddingTests(RepresentationModuleTestCase, unittest.TestCase): """Tests for Embedding.""" cls = Embedding kwargs = dict( - num_embeddings=RepresentationModuleTests.num, - shape=RepresentationModuleTests.exp_shape, + num_embeddings=RepresentationModuleTestCase.num, + shape=RepresentationModuleTestCase.exp_shape, ) def test_constructor_errors(self): @@ -305,13 +305,13 @@ def test_constrainer_kwargs(self): ) -class TensorEmbeddingTests(RepresentationModuleTests, unittest.TestCase): +class TensorEmbeddingTests(RepresentationModuleTestCase, unittest.TestCase): """Tests for Embedding with 2-dimensional shape.""" cls = Embedding exp_shape = (3, 7) kwargs = dict( - num_embeddings=RepresentationModuleTests.num, + num_embeddings=RepresentationModuleTestCase.num, shape=(3, 7), ) @@ -335,13 +335,13 @@ def _verify_content(self, x, indices): # noqa: D102 self.assertTrue(torch.allclose(x, exp_x)) -class RGCNRepresentationTests(RepresentationModuleTests, unittest.TestCase): +class RGCNRepresentationTests(RepresentationModuleTestCase, unittest.TestCase): """Test RGCN representations.""" cls = RGCNRepresentations kwargs = dict( num_bases_or_blocks=2, - embedding_dim=RepresentationModuleTests.exp_shape[0], + embedding_dim=RepresentationModuleTestCase.exp_shape[0], ) num_relations: int = 7 num_triples: int = 31 @@ -368,11 +368,11 @@ def _pre_instantiation_hook(self, kwargs: MutableMapping[str, Any]) -> MutableMa return kwargs -class RepresentationModuleTestsTest(ptb.TestsTest[RepresentationModule], unittest.TestCase): +class RepresentationModuleTestsTestCase(ptb.TestsTestCase[RepresentationModule]): """Test that there are tests for all representation modules.""" base_cls = RepresentationModule - base_test = RepresentationModuleTests + base_test = RepresentationModuleTestCase skip_cls = {MockRepresentations} From e3ca9bfc0bc8b1d8357395d44271573b6ef44c7a Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 17 Dec 2020 14:25:51 +0100 Subject: [PATCH 685/690] Update test_models.py --- tests/test_models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_models.py b/tests/test_models.py index eacd6e665a..907f5809f1 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1060,7 +1060,7 @@ def test_symmetric_edge_weights(self): class TestModelUtilities(unittest.TestCase): - """Extra tests for utilit functions.""" + """Extra tests for utility functions.""" def test_abstract(self): """Test that classes are checked as abstract properly.""" From ff9a9dbae1984b0061a7be14e974478608602112 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 17 Dec 2020 14:41:02 +0100 Subject: [PATCH 686/690] Update losses.py --- src/pykeen/losses.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pykeen/losses.py b/src/pykeen/losses.py index d373074de9..e5a179ff92 100644 --- a/src/pykeen/losses.py +++ b/src/pykeen/losses.py @@ -98,7 +98,7 @@ class MarginRankingLoss(nn.MarginRankingLoss, PairwiseLoss): synonyms = {"Pairwise Hinge Loss"} - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( margin=dict(type=int, low=0, high=3, q=1), ) @@ -160,7 +160,7 @@ class NSSALoss(SetwiseLoss): synonyms = {'Self-Adversarial Negative Sampling Loss', 'Negative Sampling Self-Adversarial Loss'} - hpo_default = dict( + hpo_default: ClassVar[Mapping[str, Any]] = dict( margin=dict(type=int, low=3, high=30, q=3), adversarial_temperature=dict(type=float, low=0.5, high=1.0), ) From ebdddbdea0ef4f95aab42edc9dd6e87464df416e Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 17 Dec 2020 14:43:22 +0100 Subject: [PATCH 687/690] Update typing.py --- src/pykeen/typing.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/src/pykeen/typing.py b/src/pykeen/typing.py index 8634bd11ad..e834dfacb1 100644 --- a/src/pykeen/typing.py +++ b/src/pykeen/typing.py @@ -16,7 +16,6 @@ 'Normalizer', 'Constrainer', 'DeviceHint', - 'RandomHint', 'TorchRandomHint', 'HeadRepresentation', 'RelationRepresentation', @@ -37,7 +36,6 @@ Constrainer = Callable[[TensorType], TensorType] DeviceHint = Union[None, str, torch.device] -RandomHint = Union[None, int, np.random.RandomState] TorchRandomHint = Union[None, int, torch.Generator] Representation = torch.FloatTensor From bf63db37ddd917fbdaa985cb563d440f338483a0 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 17 Dec 2020 14:44:17 +0100 Subject: [PATCH 688/690] Update losses.py --- src/pykeen/losses.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/src/pykeen/losses.py b/src/pykeen/losses.py index e5a179ff92..19f5c7e80c 100644 --- a/src/pykeen/losses.py +++ b/src/pykeen/losses.py @@ -86,6 +86,8 @@ class BCEWithLogitsLoss(PointwiseLoss, nn.BCEWithLogitsLoss): a negative distance as score and cannot produce positive model outputs. """ + synonyms = {'Negative Log Likelihood Loss'} + class MSELoss(PointwiseLoss, nn.MSELoss): """A wrapper around the PyTorch mean square error loss.""" From c988b4984230f68a7ac878c7f718c9e78138fcc0 Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 17 Dec 2020 14:45:52 +0100 Subject: [PATCH 689/690] Update tests.yml --- .github/workflows/tests.yml | 3 --- 1 file changed, 3 deletions(-) diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml index 0bf64f7354..bcdd86f2b7 100644 --- a/.github/workflows/tests.yml +++ b/.github/workflows/tests.yml @@ -29,9 +29,6 @@ jobs: - name: Check code quality with flake8 run: tox -e flake8 - - name: Check typing with MyPy - run: tox -e mypy - - name: Check package metadata with Pyroma run: tox -e pyroma From 79b21871af7f56570678f420a518b64d6eaa66dc Mon Sep 17 00:00:00 2001 From: Charles Tapley Hoyt Date: Thu, 21 Jan 2021 14:11:47 +0100 Subject: [PATCH 690/690] Update test_utils.py --- tests/test_utils.py | 86 +++++++++++++++++++++++---------------------- 1 file changed, 44 insertions(+), 42 deletions(-) diff --git a/tests/test_utils.py b/tests/test_utils.py index 2375216a08..e27773d455 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -198,48 +198,50 @@ def test_torch_is_in_1d(): ) assert (result == expected_result).all() - def test_complex_utils(self): - """Test complex tensor utilities.""" - re = torch.rand(20, 10) - im = torch.rand(20, 10) - x = combine_complex(x_re=re, x_im=im) - re2, im2 = split_complex(x) - assert (re2 == re).all() - assert (im2 == im).all() - - def test_project_entity(self): - """Test _project_entity.""" - batch_size = 2 - embedding_dim = 3 - relation_dim = 5 - num_entities = 7 - - # random entity embeddings & projections - e = torch.rand(1, num_entities, embedding_dim) - e = clamp_norm(e, maxnorm=1, p=2, dim=-1) - e_p = torch.rand(1, num_entities, embedding_dim) - - # random relation embeddings & projections - r_p = torch.rand(batch_size, 1, relation_dim) - - # project - e_bot = project_entity(e=e, e_p=e_p, r_p=r_p) - - # check shape: - assert e_bot.shape == (batch_size, num_entities, relation_dim) - - # check normalization - assert (torch.norm(e_bot, dim=-1, p=2) <= 1.0 + 1.0e-06).all() - - # check equivalence of re-formulation - # e_{\bot} = M_{re} e = (r_p e_p^T + I^{d_r \times d_e}) e - # = r_p (e_p^T e) + e' - m_re = r_p.unsqueeze(dim=-1) @ e_p.unsqueeze(dim=-2) - m_re = m_re + torch.eye(relation_dim, embedding_dim).view(1, 1, relation_dim, embedding_dim) - assert m_re.shape == (batch_size, num_entities, relation_dim, embedding_dim) - e_vanilla = (m_re @ e.unsqueeze(dim=-1)).squeeze(dim=-1) - e_vanilla = clamp_norm(e_vanilla, p=2, dim=-1, maxnorm=1) - assert torch.allclose(e_vanilla, e_bot) + +def test_complex_utils(): + """Test complex tensor utilities.""" + re = torch.rand(20, 10) + im = torch.rand(20, 10) + x = combine_complex(x_re=re, x_im=im) + re2, im2 = split_complex(x) + assert (re2 == re).all() + assert (im2 == im).all() + + +def test_project_entity(): + """Test _project_entity.""" + batch_size = 2 + embedding_dim = 3 + relation_dim = 5 + num_entities = 7 + + # random entity embeddings & projections + e = torch.rand(1, num_entities, embedding_dim) + e = clamp_norm(e, maxnorm=1, p=2, dim=-1) + e_p = torch.rand(1, num_entities, embedding_dim) + + # random relation embeddings & projections + r_p = torch.rand(batch_size, 1, relation_dim) + + # project + e_bot = project_entity(e=e, e_p=e_p, r_p=r_p) + + # check shape: + assert e_bot.shape == (batch_size, num_entities, relation_dim) + + # check normalization + assert (torch.norm(e_bot, dim=-1, p=2) <= 1.0 + 1.0e-06).all() + + # check equivalence of re-formulation + # e_{\bot} = M_{re} e = (r_p e_p^T + I^{d_r \times d_e}) e + # = r_p (e_p^T e) + e' + m_re = r_p.unsqueeze(dim=-1) @ e_p.unsqueeze(dim=-2) + m_re = m_re + torch.eye(relation_dim, embedding_dim).view(1, 1, relation_dim, embedding_dim) + assert m_re.shape == (batch_size, num_entities, relation_dim, embedding_dim) + e_vanilla = (m_re @ e.unsqueeze(dim=-1)).squeeze(dim=-1) + e_vanilla = clamp_norm(e_vanilla, p=2, dim=-1, maxnorm=1) + assert torch.allclose(e_vanilla, e_bot) class TestCudaExceptionsHandling(unittest.TestCase):