From 183d6534fe3144b2f5aca55a4c3ad5b2f7e9d1f4 Mon Sep 17 00:00:00 2001 From: gerhagj Date: Tue, 20 Aug 2024 15:15:10 +0200 Subject: [PATCH 1/3] Added Global AttentionPooling --- tensorframes/nn/pooling.py | 58 ++++++++++++++++++++++++++++++++++++++ 1 file changed, 58 insertions(+) create mode 100644 tensorframes/nn/pooling.py diff --git a/tensorframes/nn/pooling.py b/tensorframes/nn/pooling.py new file mode 100644 index 0000000..c3151b0 --- /dev/null +++ b/tensorframes/nn/pooling.py @@ -0,0 +1,58 @@ +from typing import Union + +import torch +import torch_geometric +import torch_geometric.utils +from torch import Tensor + +from tensorframes.reps.irreps import Irreps +from tensorframes.reps.tensorreps import TensorReps + + +class GlobalAttentionPooling(torch.nn.Module): + """The GlobalAggregationPooling module. + + This module takes a tensor and aggregates the values in each batch using a global attention mechanism. + TODO: at the moment this only works if one only wants a scalar output in the network. + """ + + def __init__( + self, in_reps: Union[TensorReps, Irreps], out_reps: Union[TensorReps, Irreps] + ) -> None: + """Initialize the GlobalAggregationPooling module. + + Args: + in_reps (list): List of input representations. + out_reps (list): List of output representations. + """ + super().__init__() + self.in_reps = in_reps + self.out_reps = out_reps + + self.query = torch.nn.Parameter(torch.randn(in_reps.dim)) + self.key = torch.nn.Linear(in_reps.dim, in_reps.dim) + self.value = torch.nn.Linear(in_reps.dim, out_reps.dim) + + self.lin = torch.nn.Linear(in_reps.dim, out_reps.dim) + + def forward(self, x: Tensor, batch: Tensor) -> Tensor: + """Applies the GlobalAggregationPooling module. + + Args: + x (torch.Tensor): The input tensor. + batch (torch.Tensor): The batch tensor. + + Returns: + torch.Tensor: The output tensor. + """ + q = self.query + k = self.key(x) + v = self.value(x) + + softmax = torch_geometric.utils.softmax(q @ k.transpose(-1, -2), batch, dim=-1) + + x = torch.einsum("i,ij->ij", softmax, v) + + out = torch_geometric.nn.pool.global_add_pool(x, batch) + + return out From 508cd4fdc2c7bb1423627a3ac675ee2e88fddae7 Mon Sep 17 00:00:00 2001 From: gerhagj Date: Thu, 22 Aug 2024 10:26:39 +0200 Subject: [PATCH 2/3] added the possibility to control bias in GlobalAttentionPooling --- tensorframes/nn/pooling.py | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/tensorframes/nn/pooling.py b/tensorframes/nn/pooling.py index c3151b0..a2531dc 100644 --- a/tensorframes/nn/pooling.py +++ b/tensorframes/nn/pooling.py @@ -10,33 +10,35 @@ class GlobalAttentionPooling(torch.nn.Module): - """The GlobalAggregationPooling module. + """The GlobalAttentionPooling module. This module takes a tensor and aggregates the values in each batch using a global attention mechanism. TODO: at the moment this only works if one only wants a scalar output in the network. """ def __init__( - self, in_reps: Union[TensorReps, Irreps], out_reps: Union[TensorReps, Irreps] + self, + in_reps: Union[TensorReps, Irreps], + out_reps: Union[TensorReps, Irreps], + bias: bool = False, ) -> None: - """Initialize the GlobalAggregationPooling module. + """Initialize the GlobalAttentionPooling module. Args: in_reps (list): List of input representations. out_reps (list): List of output representations. + bias (bool, optional): Whether to include bias terms. Defaults to False. """ super().__init__() self.in_reps = in_reps self.out_reps = out_reps self.query = torch.nn.Parameter(torch.randn(in_reps.dim)) - self.key = torch.nn.Linear(in_reps.dim, in_reps.dim) - self.value = torch.nn.Linear(in_reps.dim, out_reps.dim) - - self.lin = torch.nn.Linear(in_reps.dim, out_reps.dim) + self.key = torch.nn.Linear(in_reps.dim, in_reps.dim, bias=bias) + self.value = torch.nn.Linear(in_reps.dim, out_reps.dim, bias=bias) def forward(self, x: Tensor, batch: Tensor) -> Tensor: - """Applies the GlobalAggregationPooling module. + """Applies the GlobalAttentionPooling module. Args: x (torch.Tensor): The input tensor. From 5fbf32208f7d99bdb9df444528ea5ff899f4e435 Mon Sep 17 00:00:00 2001 From: gerhagj Date: Wed, 28 Aug 2024 11:42:33 +0200 Subject: [PATCH 3/3] now the new Reps typehinting is used --- tensorframes/nn/pooling.py | 12 +++++------- 1 file changed, 5 insertions(+), 7 deletions(-) diff --git a/tensorframes/nn/pooling.py b/tensorframes/nn/pooling.py index a2531dc..8f0fdc3 100644 --- a/tensorframes/nn/pooling.py +++ b/tensorframes/nn/pooling.py @@ -1,12 +1,10 @@ -from typing import Union - import torch import torch_geometric import torch_geometric.utils from torch import Tensor +from torch_geometric.nn import global_add_pool -from tensorframes.reps.irreps import Irreps -from tensorframes.reps.tensorreps import TensorReps +from tensorframes.reps.reps import Reps class GlobalAttentionPooling(torch.nn.Module): @@ -18,8 +16,8 @@ class GlobalAttentionPooling(torch.nn.Module): def __init__( self, - in_reps: Union[TensorReps, Irreps], - out_reps: Union[TensorReps, Irreps], + in_reps: Reps, + out_reps: Reps, bias: bool = False, ) -> None: """Initialize the GlobalAttentionPooling module. @@ -55,6 +53,6 @@ def forward(self, x: Tensor, batch: Tensor) -> Tensor: x = torch.einsum("i,ij->ij", softmax, v) - out = torch_geometric.nn.pool.global_add_pool(x, batch) + out = global_add_pool(x, batch) return out