Update nn.aggr documentation (#5099)

rusty1s · web-flow · commit 4f392ace2072 · 2022-07-31T03:01:57.000+02:00
* initial commit

* changelog

* update
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -43,7 +43,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added the `bias` vector to the `GCN` model definition in the "Create Message Passing Networks" tutorial ([#4755](https://github.com/pyg-team/pytorch_geometric/pull/4755))
 - Added `transforms.RootedSubgraph` interface with two implementations: `RootedEgoNets` and `RootedRWSubgraph` ([#3926](https://github.com/pyg-team/pytorch_geometric/pull/3926))
 - Added `ptr` vectors for `follow_batch` attributes within `Batch.from_data_list` ([#4723](https://github.com/pyg-team/pytorch_geometric/pull/4723))
-- Added `torch_geometric.nn.aggr` package ([#4687](https://github.com/pyg-team/pytorch_geometric/pull/4687), [#4721](https://github.com/pyg-team/pytorch_geometric/pull/4721), [#4731](https://github.com/pyg-team/pytorch_geometric/pull/4731), [#4762](https://github.com/pyg-team/pytorch_geometric/pull/4762), [#4749](https://github.com/pyg-team/pytorch_geometric/pull/4749), [#4779](https://github.com/pyg-team/pytorch_geometric/pull/4779), [#4863](https://github.com/pyg-team/pytorch_geometric/pull/4863), [#4864](https://github.com/pyg-team/pytorch_geometric/pull/4864), [#4865](https://github.com/pyg-team/pytorch_geometric/pull/4865), [#4866](https://github.com/pyg-team/pytorch_geometric/pull/4866), [#4872](https://github.com/pyg-team/pytorch_geometric/pull/4872), [#4934](https://github.com/pyg-team/pytorch_geometric/pull/4934), [#4935](https://github.com/pyg-team/pytorch_geometric/pull/4935), [#4957](https://github.com/pyg-team/pytorch_geometric/pull/4957), [#4973](https://github.com/pyg-team/pytorch_geometric/pull/4973), [#4973](https://github.com/pyg-team/pytorch_geometric/pull/4973), [#4986](https://github.com/pyg-team/pytorch_geometric/pull/4986), [#4995](https://github.com/pyg-team/pytorch_geometric/pull/4995), [#5000](https://github.com/pyg-team/pytorch_geometric/pull/5000), [#5034](https://github.com/pyg-team/pytorch_geometric/pull/5034), [#5036](https://github.com/pyg-team/pytorch_geometric/pull/5036), [#5039](https://github.com/pyg-team/pytorch_geometric/issues/5039), [#4522](https://github.com/pyg-team/pytorch_geometric/pull/4522), [#5033](https://github.com/pyg-team/pytorch_geometric/pull/5033]), [#5085](https://github.com/pyg-team/pytorch_geometric/pull/5085), [#5097](https://github.com/pyg-team/pytorch_geometric/pull/5097))
+- Added `torch_geometric.nn.aggr` package ([#4687](https://github.com/pyg-team/pytorch_geometric/pull/4687), [#4721](https://github.com/pyg-team/pytorch_geometric/pull/4721), [#4731](https://github.com/pyg-team/pytorch_geometric/pull/4731), [#4762](https://github.com/pyg-team/pytorch_geometric/pull/4762), [#4749](https://github.com/pyg-team/pytorch_geometric/pull/4749), [#4779](https://github.com/pyg-team/pytorch_geometric/pull/4779), [#4863](https://github.com/pyg-team/pytorch_geometric/pull/4863), [#4864](https://github.com/pyg-team/pytorch_geometric/pull/4864), [#4865](https://github.com/pyg-team/pytorch_geometric/pull/4865), [#4866](https://github.com/pyg-team/pytorch_geometric/pull/4866), [#4872](https://github.com/pyg-team/pytorch_geometric/pull/4872), [#4934](https://github.com/pyg-team/pytorch_geometric/pull/4934), [#4935](https://github.com/pyg-team/pytorch_geometric/pull/4935), [#4957](https://github.com/pyg-team/pytorch_geometric/pull/4957), [#4973](https://github.com/pyg-team/pytorch_geometric/pull/4973), [#4973](https://github.com/pyg-team/pytorch_geometric/pull/4973), [#4986](https://github.com/pyg-team/pytorch_geometric/pull/4986), [#4995](https://github.com/pyg-team/pytorch_geometric/pull/4995), [#5000](https://github.com/pyg-team/pytorch_geometric/pull/5000), [#5034](https://github.com/pyg-team/pytorch_geometric/pull/5034), [#5036](https://github.com/pyg-team/pytorch_geometric/pull/5036), [#5039](https://github.com/pyg-team/pytorch_geometric/issues/5039), [#4522](https://github.com/pyg-team/pytorch_geometric/pull/4522), [#5033](https://github.com/pyg-team/pytorch_geometric/pull/5033]), [#5085](https://github.com/pyg-team/pytorch_geometric/pull/5085), [#5097](https://github.com/pyg-team/pytorch_geometric/pull/5097), [#5099](https://github.com/pyg-team/pytorch_geometric/pull/5099))
 - Added the `DimeNet++` model ([#4432](https://github.com/pyg-team/pytorch_geometric/pull/4432), [#4699](https://github.com/pyg-team/pytorch_geometric/pull/4699), [#4700](https://github.com/pyg-team/pytorch_geometric/pull/4700), [#4800](https://github.com/pyg-team/pytorch_geometric/pull/4800))
 - Added an example of using PyG with PyTorch Ignite ([#4487](https://github.com/pyg-team/pytorch_geometric/pull/4487))
 - Added `GroupAddRev` module with support for reducing training GPU memory ([#4671](https://github.com/pyg-team/pytorch_geometric/pull/4671), [#4701](https://github.com/pyg-team/pytorch_geometric/pull/4701), [#4715](https://github.com/pyg-team/pytorch_geometric/pull/4715), [#4730](https://github.com/pyg-team/pytorch_geometric/pull/4730))
diff --git a/docs/source/modules/nn.rst b/docs/source/modules/nn.rst
@@ -25,9 +25,8 @@ Convolutional Layers
      {{ cls }}
    {% endfor %}
 
-.. automodule:: torch_geometric.nn.conv.message_passing
+.. autoclass:: torch_geometric.nn.conv.MessagePassing
    :members:
-   :exclude-members: training
 
 .. automodule:: torch_geometric.nn.conv
    :members:
@@ -49,7 +48,7 @@ Aggregation Operators
      {{ cls }}
    {% endfor %}
 
-.. automodule:: torch_geometric.nn.aggr.base
+.. autoclass:: torch_geometric.nn.aggr.Aggregation
    :members:
 
 .. automodule:: torch_geometric.nn.aggr
diff --git a/examples/equilibrium_median.py b/examples/equilibrium_median.py
@@ -11,7 +11,7 @@
 import numpy as np
 import torch
 
-from torch_geometric.nn.aggr import EquilibriumAggregation
+from torch_geometric.nn import EquilibriumAggregation
 
 input_size = 100
 steps = 10000000
diff --git a/test/nn/aggr/test_equilibrium.py b/test/nn/aggr/test_equilibrium.py
@@ -7,7 +7,6 @@
 @pytest.mark.parametrize('iter', [0, 1, 5])
 @pytest.mark.parametrize('alpha', [0, .1, 5])
 def test_equilibrium(iter, alpha):
-
     batch_size = 10
     feature_channels = 3
     output_channels = 2
@@ -30,7 +29,6 @@ def test_equilibrium(iter, alpha):
 @pytest.mark.parametrize('iter', [0, 1, 5])
 @pytest.mark.parametrize('alpha', [0, .1, 5])
 def test_equilibrium_batch(iter, alpha):
-
     batch_1, batch_2 = 4, 6
     feature_channels = 3
     output_channels = 2
diff --git a/torch_geometric/nn/aggr/attention.py b/torch_geometric/nn/aggr/attention.py
@@ -32,14 +32,6 @@ class AttentionalAggregation(Aggregation):
             shape :obj:`[-1, in_channels]` to shape :obj:`[-1, out_channels]`
             before combining them with the attention scores, *e.g.*, defined by
             :class:`torch.nn.Sequential`. (default: :obj:`None`)
-
-    Shapes:
-        - **input:**
-          node features :math:`(|\mathcal{V}|, F)`,
-          batch vector :math:`(|\mathcal{V}|)` *(optional)*
-        - **output:**
-          graph features :math:`(|\mathcal{G}|, F)` where
-          :math:`|\mathcal{G}|` denotes the number of graphs in the batch
     """
     def __init__(self, gate_nn: torch.nn.Module,
                  nn: Optional[torch.nn.Module] = None):
diff --git a/torch_geometric/nn/aggr/base.py b/torch_geometric/nn/aggr/base.py
@@ -8,7 +8,16 @@
 
 
 class Aggregation(torch.nn.Module):
-    r"""An abstract base class for implementing custom aggregations."""
+    r"""An abstract base class for implementing custom aggregations.
+
+    Shapes:
+        - **input:**
+          node features :math:`(|\mathcal{V}|, F_{in})` or edge features
+          :math:`(|\mathcal{E}|, F_{in})`,
+          index vector :math:`(|\mathcal{V}|)` or :math:`(|\mathcal{E}|)`,
+        - **output:** graph features :math:`(|\mathcal{G}|, F_{out})` or node
+          features :math:`(|\mathcal{V}|, F_{out})`
+    """
 
     # @abstractmethod
     def forward(self, x: Tensor, index: Optional[Tensor] = None,
diff --git a/torch_geometric/nn/aggr/equilibrium.py b/torch_geometric/nn/aggr/equilibrium.py
@@ -102,31 +102,30 @@ def forward(
 
 
 class EquilibriumAggregation(Aggregation):
-    r"""
-    The graph global pooling layer from the
-    `"Equilibrium Aggregation: Encoding Sets via Optimization"
-    <https://arxiv.org/abs/2202.12795>`_ paper.
-    This output of this layer :math:`\mathbf{y}` is defined implicitly by
-    defining a potential function :math:`F(\mathbf{x}, \mathbf{y})`
-    and regulatization function :math:`R(\mathbf{y})` and the condition
+    r"""The equilibrium aggregation layer from the `"Equilibrium Aggregation:
+    Encoding Sets via Optimization" <https://arxiv.org/abs/2202.12795>`_ paper.
+    The output of this layer :math:`\mathbf{y}` is defined implicitly via a
+    potential function :math:`F(\mathbf{x}, \mathbf{y})`, a regularization term
+    :math:`R(\mathbf{y})`, and the condition
 
     .. math::
-        \mathbf{y} = \min_\mathbf{y} R(\mathbf{y}) +
-        \sum_{i} F(\mathbf{x}_i, \mathbf{y})
+        \mathbf{y} = \min_\mathbf{y} R(\mathbf{y}) + \sum_{i}
+        F(\mathbf{x}_i, \mathbf{y}).
 
-    This implementation uses a ResNet Like model for the potential function
-    and a simple L2 norm for the regularizer with learnable weight
-    :math:`\lambda`.
+    The given implementation uses a ResNet-like model for the potential
+    function and a simple :math:`L_2` norm :math:`R(\mathbf{y}) =
+    \textrm{softplus}(\lambda) \cdot {\| \mathbf{y} \|}^2_2` for the
+    regularizer with learnable weight :math:`\lambda`.
 
     Args:
-        in_channels (int): The number of channels in the input to the layer.
-        out_channels (float): The number of channels in the ouput.
-        num_layers (List[int): A list of the number of hidden units in the
-            potential function.
+        in_channels (int): Size of each input sample.
+        out_channels (int): Size of each output sample.
+        num_layers (List[int): List of hidden channels in the potential
+            function.
         grad_iter (int): The number of steps to take in the internal gradient
             descent. (default: :obj:`5`)
-        lamb (float): The initial regularization constant. Is learnable.
-            descent. (default: :obj:`0.1`)
+        lamb (float): The initial regularization constant.
+            (default: :obj:`0.1`)
     """
     def __init__(self, in_channels: int, out_channels: int,
                  num_layers: List[int], grad_iter: int = 5, lamb: float = 0.1):
@@ -135,29 +134,25 @@ def __init__(self, in_channels: int, out_channels: int,
         self.potential = ResNetPotential(in_channels + out_channels, 1,
                                          num_layers)
         self.optimizer = MomentumOptimizer()
-        self._initial_lambda = lamb
-        self._labmda = torch.nn.Parameter(Tensor([lamb]), requires_grad=True)
+        self.initial_lamb = lamb
+        self.lamb = torch.nn.Parameter(Tensor(1), requires_grad=True)
         self.softplus = torch.nn.Softplus()
         self.grad_iter = grad_iter
         self.output_dim = out_channels
         self.reset_parameters()
 
     def reset_parameters(self):
-        self.lamb.data.fill_(self._initial_lambda)
+        self.lamb.data.fill_(self.initial_lamb)
         reset(self.optimizer)
         reset(self.potential)
 
-    @property
-    def lamb(self):
-        return self.softplus(self._labmda)
-
     def init_output(self, index: Optional[Tensor] = None) -> Tensor:
         index_size = 1 if index is None else int(index.max().item() + 1)
         return torch.zeros(index_size, self.output_dim,
                            requires_grad=True).float()
 
-    def reg(self, y: Tensor) -> float:
-        return self.lamb * y.square().mean(dim=-2).sum(dim=0)
+    def reg(self, y: Tensor) -> Tensor:
+        return self.softplus(self.lamb) * y.square().sum(dim=-1).mean()
 
     def energy(self, x: Tensor, y: Tensor, index: Optional[Tensor]):
         return self.potential(x, y, index) + self.reg(y)
diff --git a/torch_geometric/nn/aggr/gmt.py b/torch_geometric/nn/aggr/gmt.py
@@ -171,14 +171,6 @@ class GraphMultisetTransformer(Aggregation):
             (default: :obj:`4`)
         layer_norm (bool, optional): If set to :obj:`True`, will make use of
             layer normalization. (default: :obj:`False`)
-
-    Shapes:
-        - **input:**
-          node features :math:`(|\mathcal{V}|, F_{in})`,
-          batch vector :math:`(|\mathcal{V}|)`,
-          edge indices :math:`(2, |\mathcal{E}|)` *(optional)*
-        - **output:** graph features :math:`(|\mathcal{G}|, F_{out})` where
-          :math:`|\mathcal{G}|` denotes the number of graphs in the batch
     """
     def __init__(
         self,
diff --git a/torch_geometric/nn/aggr/lstm.py b/torch_geometric/nn/aggr/lstm.py
@@ -8,14 +8,12 @@
 
 class LSTMAggregation(Aggregation):
     r"""Performs LSTM-style aggregation in which the elements to aggregate are
-    interpreted as a sequence.
+    interpreted as a sequence, as described in the `"Inductive Representation
+    Learning on Large Graphs" <https://arxiv.org/abs/1706.02216>`_ paper.
 
     .. warning::
         :class:`LSTMAggregation` is not a permutation-invariant operator.
 
-    .. note::
-        :class:`LSTMAggregation` requires sorted indices as input.
-
     Args:
         in_channels (int): Size of each input sample.
         out_channels (int): Size of each output sample.
diff --git a/torch_geometric/nn/aggr/multi.py b/torch_geometric/nn/aggr/multi.py
@@ -10,7 +10,10 @@
 
 class MultiAggregation(Aggregation):
     r"""Performs aggregations with one or more aggregators and combines
-    aggregated results.
+    aggregated results, as described in the `"Principal Neighbourhood
+    Aggregation for Graph Nets" <https://arxiv.org/abs/2004.05718>`_ and
+    `"Adaptive Filters and Aggregator Fusion for Efficient Graph Convolutions"
+    <https://arxiv.org/abs/2104.01481>`_ papers.
 
     Args:
         aggrs (list): The list of aggregation schemes to use.
