create DegreeScalerAggregation

Author: Padarn

torch_geometric/nn/aggr/__init__.py

@@ -13,19 +13,11 @@
 )
 from .lstm import LSTMAggregation
 from .set2set import Set2Set
+from .scaler import DegreeScalerAggregation
 
 __all__ = classes = [
-    'Aggregation',
-    'MultiAggregation',
-    'MeanAggregation',
-    'SumAggregation',
-    'MaxAggregation',
-    'MinAggregation',
-    'MulAggregation',
-    'VarAggregation',
-    'StdAggregation',
-    'SoftmaxAggregation',
-    'PowerMeanAggregation',
-    'LSTMAggregation',
-    'Set2Set',
+    'Aggregation', 'MultiAggregation', 'MeanAggregation', 'SumAggregation',
+    'MaxAggregation', 'MinAggregation', 'MulAggregation', 'VarAggregation',
+    'StdAggregation', 'SoftmaxAggregation', 'PowerMeanAggregation',
+    'LSTMAggregation', 'Set2Set', 'DegreeScalerAggregation'
 ]

torch_geometric/nn/aggr/scaler.py

@@ -0,0 +1,72 @@
+from typing import Any, Dict, List, Optional, Union
+
+import torch
+from torch import Tensor
+
+from torch_geometric.nn.aggr import Aggregation, MultiAggregation
+from torch_geometric.utils import degree
+
+
+class DegreeScalerAggregation(Aggregation):
+    """
+    Class that combines together one or more aggregators and then transforms
+    the result with one or more scalers. The scalers are normalised by the
+    in-degree of the training set and so must be provided at construction.
+
+    Args:
+        aggr (string or list or Aggregation, optional): The list of
+            aggregations given as :class:`~torch_geometric.nn.aggr.Aggregation`
+            (or any string that automatically resolves to it).
+        scalers (list of str): Set of scaling function identifiers, namely
+                :obj:`"identity"`, :obj:`"amplification"`,
+                :obj:`"attenuation"`, :obj:`"linear"` and
+                :obj:`"inverse_linear"`.
+        deg (Tensor): Histogram of in-degrees of nodes in the training set,
+            used by scalers to normalize.
+        aggr_kwargs (List[Dict[str, Any]], optional): Arguments passed to the
+            respective aggregation functions in case it gets automatically
+            resolved. (default: :obj:`None`)
+    """
+    def __init__(self, aggrs: List[Union[Aggregation, str]],
+                 scalers: List[str], deg: Tensor,
+                 aggrs_kwargs: Optional[List[Dict[str, Any]]] = None):
+
+        super().__init__()
+
+        self.agg = MultiAggregation(aggrs, aggrs_kwargs)
+        self.scalersz = scalers
+
+        deg = deg.to(torch.float)
+        num_nodes = int(deg.sum())
+        bin_degrees = torch.arange(deg.numel())
+        self.avg_deg: Dict[str, float] = {
+            'lin': float((bin_degrees * deg).sum()) / num_nodes,
+            'log': float(((bin_degrees + 1).log() * deg).sum()) / num_nodes,
+            'exp': float((bin_degrees.exp() * deg).sum()) / num_nodes,
+        }
+
+    def forward(self, x: Tensor, index: Optional[Tensor] = None,
+                ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
+                dim: int = -2) -> Tensor:
+
+        out = self.agg(x, index, ptr, dim_size, dim)
+
+        deg = degree(index, dtype=out.dtype)
+        deg = deg.clamp_(1).view(-1, 1, 1)
+
+        outs = []
+        for scaler in self.scalers:
+            if scaler == 'identity':
+                pass
+            elif scaler == 'amplification':
+                out = out * (torch.log(deg + 1) / self.avg_deg['log'])
+            elif scaler == 'attenuation':
+                out = out * (self.avg_deg['log'] / torch.log(deg + 1))
+            elif scaler == 'linear':
+                out = out * (deg / self.avg_deg['lin'])
+            elif scaler == 'inverse_linear':
+                out = out * (self.avg_deg['lin'] / deg)
+            else:
+                raise ValueError(f'Unknown scaler "{scaler}".')
+            outs.append(out)
+        return torch.cat(outs, dim=-1)

torch_geometric/nn/conv/pna_conv.py

@@ -4,10 +4,10 @@
 from torch import Tensor
 from torch.nn import ModuleList, ReLU, Sequential
 
+from torch_geometric.nn.aggr import DegreeScalerAggregation
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.dense.linear import Linear
 from torch_geometric.typing import Adj, OptTensor
-from torch_geometric.utils import degree
 
 from ..inits import reset
 
@@ -85,7 +85,9 @@ def __init__(self, in_channels: int, out_channels: int,
                  pre_layers: int = 1, post_layers: int = 1,
                  divide_input: bool = False, **kwargs):
 
-        kwargs.setdefault('aggr', aggregators)
+        aggr = DegreeScalerAggregation(aggregators, scalers, deg)
+        kwargs.setdefault('aggr', aggr)
+
         super().__init__(node_dim=0, **kwargs)
 
         if divide_input:
@@ -94,7 +96,6 @@ def __init__(self, in_channels: int, out_channels: int,
 
         self.in_channels = in_channels
         self.out_channels = out_channels
-        self.scalers = scalers
         self.edge_dim = edge_dim
         self.towers = towers
         self.divide_input = divide_input
@@ -132,9 +133,6 @@ def __init__(self, in_channels: int, out_channels: int,
 
         self.lin = Linear(out_channels, out_channels)
 
-        # scale hook
-        self.register_aggregate_forward_hook(self.scale_hook)
-
         self.reset_parameters()
 
     def reset_parameters(self):
@@ -179,29 +177,6 @@ def message(self, x_i: Tensor, x_j: Tensor,
         hs = [nn(h[:, i]) for i, nn in enumerate(self.pre_nns)]
         return torch.stack(hs, dim=1)
 
-    @staticmethod
-    def scale_hook(module, inputs, out: Tensor) -> Tensor:
-
-        deg = degree(inputs[0]['index'], dtype=out.dtype)
-        deg = deg.clamp_(1).view(-1, 1, 1)
-
-        outs = []
-        for scaler in module.scalers:
-            if scaler == 'identity':
-                pass
-            elif scaler == 'amplification':
-                out = out * (torch.log(deg + 1) / module.avg_deg['log'])
-            elif scaler == 'attenuation':
-                out = out * (module.avg_deg['log'] / torch.log(deg + 1))
-            elif scaler == 'linear':
-                out = out * (deg / module.avg_deg['lin'])
-            elif scaler == 'inverse_linear':
-                out = out * (module.avg_deg['lin'] / deg)
-            else:
-                raise ValueError(f'Unknown scaler "{scaler}".')
-            outs.append(out)
-        return torch.cat(outs, dim=-1)
-
     def __repr__(self):
         return (f'{self.__class__.__name__}({self.in_channels}, '
                 f'{self.out_channels}, towers={self.towers}, '