Repository: pyg-team/pytorch_geometric
Pull request #7370 (2023-05-15 17:45:11)

`PMLP` implementation
I am not quite sure if the tests are sufficient. Also, at the moment, I have limited the MLP layers to have the same bias/dropouts etc as per the reference repo, but wondering if they should follow more the format of the MLP model where you can input a list?


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 83984b37f..3f3c7d399 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added the `PMLP` model ([#7370](https://github.com/pyg-team/pytorch_geometric/pull/7370))
 - Added padding capabilities to `HeteroData.to_homogeneous()` in case feature dimensionalities do not match ([#7374](https://github.com/pyg-team/pytorch_geometric/pull/7374))
 - Added an optional `batch_size` argument to `fps`, `knn`, `knn_graph`, `radius` and `radius_graph` ([#7368](https://github.com/pyg-team/pytorch_geometric/pull/7368))
 - Added `PrefetchLoader` capabilities ([#7376](https://github.com/pyg-team/pytorch_geometric/pull/7376), [#7378](https://github.com/pyg-team/pytorch_geometric/pull/7378), [#7383](https://github.com/pyg-team/pytorch_geometric/pull/7383))
diff --git a/README.md b/README.md
index e67c1572c..be467cc17 100644
--- a/README.md
+++ b/README.md
@@ -279,6 +279,7 @@ Unlike simple stacking of GNN layers, these models could involve pre-processing,
 * **[Deep Graph Infomax](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.models.DeepGraphInfomax.html)** from Veličković *et al.*: [Deep Graph Infomax](https://arxiv.org/abs/1809.10341) (ICLR 2019) [[**Example1**](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/infomax_transductive.py), [**Example2**](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/infomax_inductive.py)]
 * **Deep Multiplex Graph Infomax** from Park *et al.*: [Unsupervised Attributed Multiplex Network Embedding](https://arxiv.org/abs/1911.06750) (AAAI 2020) [[**Example**](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/hetero/dmgi_unsup.py)]
 * **[Masked Label Prediction](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.models.MaskLabel.html)** from Shi *et al.*: [Masked Label Prediction: Unified Message Passing Model for Semi-Supervised Classification](https://arxiv.org/abs/2009.03509) (CoRR 2020) [[**Example**](https://github.com/pyg-team/pytorch_geometric/blob/master/examples/unimp_arxiv.py)]
+* **[PMLP](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.models.PMLP.html)** from Yang *et al.*: [Graph Neural Networks are Inherently Good Generalizers: Insights by Bridging GNNs and MLPs](https://arxiv.org/abs/2212.09034) (ICLR 2023)
 
 <details>
 <summary><b>Expand to see all implemented GNN models...</b></summary>
diff --git a/torch_geometric/nn/conv/simple_conv.py b/torch_geometric/nn/conv/simple_conv.py
index 5a11a95e8..532f5ebc7 100644
--- a/torch_geometric/nn/conv/simple_conv.py
+++ b/torch_geometric/nn/conv/simple_conv.py
@@ -11,8 +11,9 @@ from torch_geometric.typing import (
     OptTensor,
     Size,
     SparseTensor,
+    torch_sparse,
 )
-from torch_geometric.utils import spmm
+from torch_geometric.utils import add_self_loops, spmm
 
 
 class SimpleConv(MessagePassing):
@@ -34,7 +35,7 @@ class SimpleConv(MessagePassing):
             that automatically resolves to it). (default: :obj:`"sum"`)
         combine_root (str, optional): Specifies whether or how to combine the
             central node representation (one of :obj:`"sum"`, :obj:`"cat"`,
-            :obj:`None`). (default: :obj:`None`)
+            :obj:`"self_loop"`, :obj:`None`). (default: :obj:`None`)
         **kwargs (optional): Additional arguments of
             :class:`torch_geometric.nn.conv.MessagePassing`.
 
@@ -53,7 +54,7 @@ class SimpleConv(MessagePassing):
         combine_root: Optional[str] = None,
         **kwargs,
     ):
-        if combine_root not in ['sum', 'cat', None]:
+        if combine_root not in ['sum', 'cat', 'self_loop', None]:
             raise ValueError(f"Received invalid value for 'combine_root' "
                              f"(got '{combine_root}')")
 
@@ -62,6 +63,19 @@ class SimpleConv(MessagePassing):
 
     def forward(self, x: Union[Tensor, OptPairTensor], edge_index: Adj,
                 edge_weight: OptTensor = None, size: Size = None) -> Tensor:
+
+        if self.combine_root is not None:
+            if self.combine_root == 'self_loop':
+                if not isinstance(x, Tensor) or (size is not None
+                                                 and size[0] != size[1]):
+                    raise ValueError("Cannot use `combine_root='self_loop'` "
+                                     "for bipartite message passing")
+                if isinstance(edge_index, Tensor):
+                    edge_index, edge_weight = add_self_loops(
+                        edge_index, edge_weight, num_nodes=x.size(0))
+                elif isinstance(edge_index, SparseTensor):
+                    edge_index = torch_sparse.set_diag(edge_index)
+
         if isinstance(x, Tensor):
             x: OptPairTensor = (x, x)
 
diff --git a/torch_geometric/nn/models/__init__.py b/torch_geometric/nn/models/__init__.py
index 46b9fd26f..9b716e21e 100644
--- a/torch_geometric/nn/models/__init__.py
+++ b/torch_geometric/nn/models/__init__.py
@@ -23,6 +23,7 @@ from .lightgcn import LightGCN
 from .mask_label import MaskLabel
 from .rev_gnn import GroupAddRev
 from .gnnff import GNNFF
+from .pmlp import PMLP
 
 __all__ = classes = [
     'MLP',
@@ -62,4 +63,5 @@ __all__ = classes = [
     'MaskLabel',
     'GroupAddRev',
     'GNNFF',
+    'PMLP',
 ]
diff --git a/torch_geometric/nn/models/pmlp.py b/torch_geometric/nn/models/pmlp.py
new file mode 100644
index 000000000..7677e0b34
--- /dev/null
+++ b/torch_geometric/nn/models/pmlp.py
@@ -0,0 +1,92 @@
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+from torch import Tensor
+
+from torch_geometric.nn import SimpleConv
+from torch_geometric.nn.dense.linear import Linear
+
+
+class PMLP(torch.nn.Module):
+    r"""The P(ropagational)MLP model from the `"Graph Neural Networks are
+    Inherently Good Generalizers: Insights by Bridging GNNs and MLPs"
+    <https://arxiv.org/abs/2212.09034>`_ paper.
+    :class:`PMLP` is identical to a standard MLP during training, but then
+    adopts a GNN architecture during testing.
+
+    Args:
+        in_channels (int): Size of each input sample.
+        hidden_channels (int): Size of each hidden sample.
+        out_channels (int): Size of each output sample.
+        num_layers (int): The number of layers.
+        dropout (float, optional): Dropout probability of each hidden
+            embedding. (default: :obj:`0.`)
+        bias (bool, optional): If set to :obj:`False`, the module
+            will not learn additive biases. (default: :obj:`True`)
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        hidden_channels: int,
+        out_channels: int,
+        num_layers: int,
+        dropout: float = 0.,
+        bias: bool = True,
+    ):
+        super().__init__()
+
+        self.in_channels = in_channels
+        self.hidden_channels = hidden_channels
+        self.out_channels = out_channels
+        self.num_layers = num_layers
+        self.dropout = dropout
+        self.bias = bias
+
+        self.lins = torch.nn.ModuleList()
+        self.lins.append(Linear(in_channels, hidden_channels, self.bias))
+        for _ in range(self.num_layers - 2):
+            lin = Linear(hidden_channels, hidden_channels, self.bias)
+            self.lins.append(lin)
+        self.lins.append(Linear(hidden_channels, out_channels, self.bias))
+
+        self.norm = torch.nn.BatchNorm1d(hidden_channels, affine=False,
+                                         track_running_stats=False)
+
+        self.conv = SimpleConv(aggr='mean', combine_root='self_loop')
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        r"""Resets all learnable parameters of the module."""
+        for lin in self.lins:
+            torch.nn.init.xavier_uniform_(lin.weight, gain=1.414)
+            if self.bias:
+                torch.nn.init.zeros_(lin.bias)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        edge_index: Optional[Tensor] = None,
+    ) -> torch.Tensor:
+        """"""
+        if not self.training and edge_index is None:
+            raise ValueError(f"'edge_index' needs to be present during "
+                             f"inference in '{self.__class__.__name__}'")
+
+        for i in range(self.num_layers):
+            x = x @ self.lins[i].weight.t()
+            if not self.training:
+                x = self.conv(x, edge_index)
+            if self.bias:
+                x = x + self.lins[i].bias
+            if i != self.num_layers - 1:
+                x = self.norm(x)
+                x = x.relu()
+                x = F.dropout(x, p=self.dropout, training=self.training)
+
+        return x
+
+    def __repr__(self) -> str:
+        return (f'{self.__class__.__name__}({self.in_channels}, '
+                f'{self.out_channels}, num_layers={self.num_layers})')


Test patch:
diff --git a/test/nn/conv/test_simple_conv.py b/test/nn/conv/test_simple_conv.py
index ed251cd15..28d7cffdb 100644
--- a/test/nn/conv/test_simple_conv.py
+++ b/test/nn/conv/test_simple_conv.py
@@ -12,11 +12,12 @@ from torch_geometric.utils import to_torch_csc_tensor
     ('mean', None),
     ('sum', 'sum'),
     (['mean', 'max'], 'cat'),
+    ('mean', 'self_loop'),
 ])
 def test_simple_conv(aggr, combine_root):
     x1 = torch.randn(4, 8)
     x2 = torch.randn(2, 8)
-    edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
+    edge_index = torch.tensor([[0, 1, 2, 3], [1, 0, 1, 1]])
     adj1 = to_torch_csc_tensor(edge_index, size=(4, 4))
 
     conv = SimpleConv(aggr, combine_root)
@@ -46,13 +47,15 @@ def test_simple_conv(aggr, combine_root):
         assert torch.allclose(jit(x1, adj2.t()), out)
 
     # Test bipartite message passing:
-    adj1 = to_torch_csc_tensor(edge_index, size=(4, 2))
-
-    out = conv((x1, x2), edge_index)
-    assert out.size() == (2, output_size)
-    assert torch.allclose(conv((x1, x2), edge_index, size=(4, 2)), out)
-    assert torch.allclose(conv((x1, x2), adj1.t()), out)
-
-    if torch_geometric.typing.WITH_TORCH_SPARSE:
-        adj2 = SparseTensor.from_edge_index(edge_index, sparse_sizes=(4, 2))
-        assert torch.allclose(conv((x1, x2), adj2.t()), out)
+    if combine_root != 'self_loop':
+        adj1 = to_torch_csc_tensor(edge_index, size=(4, 2))
+
+        out = conv((x1, x2), edge_index)
+        assert out.size() == (2, output_size)
+        assert torch.allclose(conv((x1, x2), edge_index, size=(4, 2)), out)
+        assert torch.allclose(conv((x1, x2), adj1.t()), out)
+
+        if torch_geometric.typing.WITH_TORCH_SPARSE:
+            adj2 = SparseTensor.from_edge_index(edge_index,
+                                                sparse_sizes=(4, 2))
+            assert torch.allclose(conv((x1, x2), adj2.t()), out)
diff --git a/test/nn/models/test_pmlp.py b/test/nn/models/test_pmlp.py
new file mode 100644
index 000000000..c7933e959
--- /dev/null
+++ b/test/nn/models/test_pmlp.py
@@ -0,0 +1,23 @@
+import pytest
+import torch
+
+from torch_geometric.nn.models import PMLP
+
+
+def test_pmlp():
+    x = torch.randn(4, 16)
+    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+
+    pmlp = PMLP(in_channels=16, hidden_channels=32, out_channels=2,
+                num_layers=4)
+    assert str(pmlp) == 'PMLP(16, 2, num_layers=4)'
+
+    pmlp.training = True
+    assert pmlp(x).size() == (4, 2)
+
+    pmlp.training = False
+    assert pmlp(x, edge_index).size() == (4, 2)
+
+    with pytest.raises(ValueError, match="'edge_index' needs to be present"):
+        pmlp.training = False
+        pmlp(x)
diff --git a/test/nn/models/test_rect.py b/test/nn/models/test_rect.py
index 74f2b0320..6364dbe21 100644
--- a/test/nn/models/test_rect.py
+++ b/test/nn/models/test_rect.py
@@ -19,13 +19,13 @@ def test_rect():
     assert out.size() == (6, 8)
     if torch_geometric.typing.WITH_TORCH_SPARSE:
         adj = SparseTensor.from_edge_index(edge_index, sparse_sizes=(6, 6))
-        assert torch.allclose(out, model(x, adj.t()))
+        assert torch.allclose(out, model(x, adj.t()), atol=1e-6)
 
     # Test `embed`:
     embed_out = model.embed(x, edge_index)
     assert embed_out.size() == (6, 16)
     if torch_geometric.typing.WITH_TORCH_SPARSE:
-        assert torch.allclose(embed_out, model.embed(x, adj.t()))
+        assert torch.allclose(embed_out, model.embed(x, adj.t()), atol=1e-6)
 
     # Test `get_semantic_labels`:
     labeds_out = model.get_semantic_labels(x, y, mask)