Merge branch 'master' into master

Author: rusty1s

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ## [2.0.5] - 2022-MM-DD
 ### Added
 - Let `ImbalancedSampler` accept `torch.Tensor` as input ([#5138](https://github.com/pyg-team/pytorch_geometric/pull/5138))
+- Added `flow` argument to `gcn_norm` to correctly normalize the adjacency matrix in `GCNConv` ([#5149](https://github.com/pyg-team/pytorch_geometric/pull/5149))
 - `NeighborSampler` supports graphs without edges ([#5072](https://github.com/pyg-team/pytorch_geometric/pull/5072))
 - Added the `MeanSubtractionNorm` layer ([#5068](https://github.com/pyg-team/pytorch_geometric/pull/5068))
 - Added `pyg_lib.segment_matmul` integration within `RGCNConv` ([#5052](https://github.com/pyg-team/pytorch_geometric/pull/5052), [#5096](https://github.com/pyg-team/pytorch_geometric/pull/5096))
@@ -66,6 +67,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added support for graph-level outputs in `to_hetero` ([#4582](https://github.com/pyg-team/pytorch_geometric/pull/4582))
 - Added `CHANGELOG.md` ([#4581](https://github.com/pyg-team/pytorch_geometric/pull/4581))
 ### Changed
+- Fixed `GINEConv` bug with non-sequential input ([#5154](https://github.com/pyg-team/pytorch_geometric/pull/5154)]
 - Improved error message ([#5095](https://github.com/pyg-team/pytorch_geometric/pull/5095))
 - Fixed `HGTLoader` bug which produced outputs with missing edge types ([#5067](https://github.com/pyg-team/pytorch_geometric/pull/5067))
 - Fixed dynamic inheritance issue in data batching ([#5051](https://github.com/pyg-team/pytorch_geometric/pull/5051))
@@ -101,7 +103,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Refactored reading molecular positions from sdf file for qm9 datasets ([4654](https://github.com/pyg-team/pytorch_geometric/pull/4654))
 - Fixed `MLP.jittable()` bug in case `return_emb=True` ([#4645](https://github.com/pyg-team/pytorch_geometric/pull/4645), [#4648](https://github.com/pyg-team/pytorch_geometric/pull/4648))
 - The generated node features of `StochasticBlockModelDataset` are now ordered with respect to their labels ([#4617](https://github.com/pyg-team/pytorch_geometric/pull/4617))
-- Fixed typos in the documentation ([#4616](https://github.com/pyg-team/pytorch_geometric/pull/4616), [#4824](https://github.com/pyg-team/pytorch_geometric/pull/4824), [#4895](https://github.com/pyg-team/pytorch_geometric/pull/4895))
+- Fixed typos in the documentation ([#4616](https://github.com/pyg-team/pytorch_geometric/pull/4616), [#4824](https://github.com/pyg-team/pytorch_geometric/pull/4824), [#4895](https://github.com/pyg-team/pytorch_geometric/pull/4895), [#5161](https://github.com/pyg-team/pytorch_geometric/pull/5161))
 - The `bias` argument in `TAGConv` is now actually applied ([#4597](https://github.com/pyg-team/pytorch_geometric/pull/4597))
 - Fixed subclass behaviour of `process` and `download` in `Datsaet` ([#4586](https://github.com/pyg-team/pytorch_geometric/pull/4586))
 - Fixed filtering of attributes for loaders in case `__cat_dim__ != 0` ([#4629](https://github.com/pyg-team/pytorch_geometric/pull/4629))

test/nn/conv/test_gcn_conv.py

@@ -78,3 +78,15 @@ def test_static_gcn_conv():
     conv = GCNConv(16, 32)
     out = conv(x, edge_index)
     assert out.size() == (3, 4, 32)
+
+
+def test_gcn_conv_norm():
+    x = torch.randn(4, 16)
+    edge_index = torch.tensor([[0, 0, 0], [1, 2, 3]])
+    row, col = edge_index
+
+    conv = GCNConv(16, 32, flow="source_to_target")
+    out1 = conv(x, edge_index)
+    conv.flow = "target_to_source"
+    out2 = conv(x, edge_index.flip(0))
+    assert torch.allclose(out1, out2, atol=1e-6)

test/nn/conv/test_gin_conv.py

@@ -126,6 +126,11 @@ def test_gine_conv_edge_dim():
     out = conv(x, edge_index, edge_attr)
     assert out.size() == (4, 32)
 
+    nn = Lin(16, 32)
+    conv = GINEConv(nn, train_eps=True, edge_dim=8)
+    out = conv(x, edge_index, edge_attr)
+    assert out.size() == (4, 32)
+
 
 def test_static_gin_conv():
     x = torch.randn(3, 4, 16)

torch_geometric/graphgym/utils/agg_runs.py

@@ -59,15 +59,18 @@ def agg_dict_list(dict_list):
 
 
 def name_to_dict(run):
-    cols = run.split('-')[1:]
+    run = run.split('-', 1)[-1]
+    cols = run.split('=')
     keys, vals = [], []
-    for col in cols:
+    keys.append(cols[0])
+    for col in cols[1:-1]:
         try:
-            key, val = col.split('=')
+            val, key = col.rsplit('-', 1)
         except Exception:
             print(col)
         keys.append(key)
         vals.append(string_to_python(val))
+    vals.append(cols[-1])
     return dict(zip(keys, vals))
 
 

torch_geometric/loader/hgt_loader.py

@@ -52,7 +52,7 @@ class HGTLoader(torch.utils.data.DataLoader):
             num_samples={key: [512] * 4 for key in hetero_data.node_types},
             # Use a batch size of 128 for sampling training nodes of type paper
             batch_size=128,
-            input_nodes=('paper': hetero_data['paper'].train_mask),
+            input_nodes=('paper', hetero_data['paper'].train_mask),
         )
 
         sampled_hetero_data = next(iter(loader))

torch_geometric/nn/aggr/equilibrium.py

@@ -41,7 +41,7 @@ def forward(self, x: Tensor, y: Tensor, index: Optional[Tensor]) -> Tensor:
             return h.mean()
 
         size = int(index.max().item() + 1)
-        return scatter(x, index, dim=0, dim_size=size, reduce='mean').sum()
+        return scatter(h, index, dim=0, dim_size=size, reduce='mean').sum()
 
 
 class MomentumOptimizer(torch.nn.Module):

torch_geometric/nn/conv/appnp.py

@@ -85,7 +85,7 @@ def forward(self, x: Tensor, edge_index: Adj,
                 if cache is None:
                     edge_index, edge_weight = gcn_norm(  # yapf: disable
                         edge_index, edge_weight, x.size(self.node_dim), False,
-                        self.add_self_loops, dtype=x.dtype)
+                        self.add_self_loops, self.flow, dtype=x.dtype)
                     if self.cached:
                         self._cached_edge_index = (edge_index, edge_weight)
                 else:
@@ -96,7 +96,7 @@ def forward(self, x: Tensor, edge_index: Adj,
                 if cache is None:
                     edge_index = gcn_norm(  # yapf: disable
                         edge_index, edge_weight, x.size(self.node_dim), False,
-                        self.add_self_loops, dtype=x.dtype)
+                        self.add_self_loops, self.flow, dtype=x.dtype)
                     if self.cached:
                         self._cached_adj_t = edge_index
                 else:

torch_geometric/nn/conv/arma_conv.py

@@ -108,12 +108,12 @@ def forward(self, x: Tensor, edge_index: Adj,
         if isinstance(edge_index, Tensor):
             edge_index, edge_weight = gcn_norm(  # yapf: disable
                 edge_index, edge_weight, x.size(self.node_dim),
-                add_self_loops=False, dtype=x.dtype)
+                add_self_loops=False, flow=self.flow, dtype=x.dtype)
 
         elif isinstance(edge_index, SparseTensor):
             edge_index = gcn_norm(  # yapf: disable
                 edge_index, edge_weight, x.size(self.node_dim),
-                add_self_loops=False, dtype=x.dtype)
+                add_self_loops=False, flow=self.flow, dtype=x.dtype)
 
         x = x.unsqueeze(-3)
         out = x

torch_geometric/nn/conv/dna_conv.py

@@ -275,7 +275,7 @@ def forward(self, x: Tensor, edge_index: Adj,
                 if cache is None:
                     edge_index, edge_weight = gcn_norm(  # yapf: disable
                         edge_index, edge_weight, x.size(self.node_dim), False,
-                        self.add_self_loops, dtype=x.dtype)
+                        self.add_self_loops, self.flow, dtype=x.dtype)
                     if self.cached:
                         self._cached_edge_index = (edge_index, edge_weight)
                 else:
@@ -286,7 +286,7 @@ def forward(self, x: Tensor, edge_index: Adj,
                 if cache is None:
                     edge_index = gcn_norm(  # yapf: disable
                         edge_index, edge_weight, x.size(self.node_dim), False,
-                        self.add_self_loops, dtype=x.dtype)
+                        self.add_self_loops, self.flow, dtype=x.dtype)
                     if self.cached:
                         self._cached_adj_t = edge_index
                 else:

torch_geometric/nn/conv/eg_conv.py

@@ -130,7 +130,8 @@ def forward(self, x: Tensor, edge_index: Adj) -> Tensor:
                 if cache is None:
                     edge_index, symnorm_weight = gcn_norm(  # yapf: disable
                         edge_index, None, num_nodes=x.size(self.node_dim),
-                        improved=False, add_self_loops=self.add_self_loops)
+                        improved=False, add_self_loops=self.add_self_loops,
+                        flow=self.flow)
                     if self.cached:
                         self._cached_edge_index = (edge_index, symnorm_weight)
                 else:
@@ -141,7 +142,8 @@ def forward(self, x: Tensor, edge_index: Adj) -> Tensor:
                 if cache is None:
                     edge_index = gcn_norm(  # yapf: disable
                         edge_index, None, num_nodes=x.size(self.node_dim),
-                        improved=False, add_self_loops=self.add_self_loops)
+                        improved=False, add_self_loops=self.add_self_loops,
+                        flow=self.flow)
                     if self.cached:
                         self._cached_adj_t = edge_index
                 else:

torch_geometric/nn/conv/fa_conv.py

@@ -115,7 +115,7 @@ def forward(self, x: Tensor, x_0: Tensor, edge_index: Adj,
                 if cache is None:
                     edge_index, edge_weight = gcn_norm(  # yapf: disable
                         edge_index, None, x.size(self.node_dim), False,
-                        self.add_self_loops, dtype=x.dtype)
+                        self.add_self_loops, self.flow, dtype=x.dtype)
                     if self.cached:
                         self._cached_edge_index = (edge_index, edge_weight)
                 else:
@@ -127,7 +127,7 @@ def forward(self, x: Tensor, x_0: Tensor, edge_index: Adj,
                 if cache is None:
                     edge_index = gcn_norm(  # yapf: disable
                         edge_index, None, x.size(self.node_dim), False,
-                        self.add_self_loops, dtype=x.dtype)
+                        self.add_self_loops, self.flow, dtype=x.dtype)
                     if self.cached:
                         self._cached_adj_t = edge_index
                 else:

torch_geometric/nn/conv/gcn2_conv.py

@@ -119,7 +119,7 @@ def forward(self, x: Tensor, x_0: Tensor, edge_index: Adj,
                 if cache is None:
                     edge_index, edge_weight = gcn_norm(  # yapf: disable
                         edge_index, edge_weight, x.size(self.node_dim), False,
-                        self.add_self_loops, dtype=x.dtype)
+                        self.add_self_loops, self.flow, dtype=x.dtype)
                     if self.cached:
                         self._cached_edge_index = (edge_index, edge_weight)
                 else:
@@ -130,7 +130,7 @@ def forward(self, x: Tensor, x_0: Tensor, edge_index: Adj,
                 if cache is None:
                     edge_index = gcn_norm(  # yapf: disable
                         edge_index, edge_weight, x.size(self.node_dim), False,
-                        self.add_self_loops, dtype=x.dtype)
+                        self.add_self_loops, self.flow, dtype=x.dtype)
                     if self.cached:
                         self._cached_adj_t = edge_index
                 else:

torch_geometric/nn/conv/gcn_conv.py

@@ -17,24 +17,25 @@
 
 @torch.jit._overload
 def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False,
-             add_self_loops=True, dtype=None):
-    # type: (Tensor, OptTensor, Optional[int], bool, bool, Optional[int]) -> PairTensor  # noqa
+             add_self_loops=True, flow="source_to_target", dtype=None):
+    # type: (Tensor, OptTensor, Optional[int], bool, bool, str, Optional[int]) -> PairTensor  # noqa
     pass
 
 
 @torch.jit._overload
 def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False,
-             add_self_loops=True, dtype=None):
-    # type: (SparseTensor, OptTensor, Optional[int], bool, bool, Optional[int]) -> SparseTensor  # noqa
+             add_self_loops=True, flow="source_to_target", dtype=None):
+    # type: (SparseTensor, OptTensor, Optional[int], bool, bool, str, Optional[int]) -> SparseTensor  # noqa
     pass
 
 
 def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False,
-             add_self_loops=True, dtype=None):
+             add_self_loops=True, flow="source_to_target", dtype=None):
 
     fill_value = 2. if improved else 1.
 
     if isinstance(edge_index, SparseTensor):
+        assert flow in ["source_to_target"]
         adj_t = edge_index
         if not adj_t.has_value():
             adj_t = adj_t.fill_value(1., dtype=dtype)
@@ -48,6 +49,7 @@ def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False,
         return adj_t
 
     else:
+        assert flow in ["source_to_target", "target_to_source"]
         num_nodes = maybe_num_nodes(edge_index, num_nodes)
 
         if edge_weight is None:
@@ -61,7 +63,8 @@ def gcn_norm(edge_index, edge_weight=None, num_nodes=None, improved=False,
             edge_weight = tmp_edge_weight
 
         row, col = edge_index[0], edge_index[1]
-        deg = scatter_add(edge_weight, col, dim=0, dim_size=num_nodes)
+        idx = col if flow == "source_to_target" else row
+        deg = scatter_add(edge_weight, idx, dim=0, dim_size=num_nodes)
         deg_inv_sqrt = deg.pow_(-0.5)
         deg_inv_sqrt.masked_fill_(deg_inv_sqrt == float('inf'), 0)
         return edge_index, deg_inv_sqrt[row] * edge_weight * deg_inv_sqrt[col]
@@ -171,7 +174,7 @@ def forward(self, x: Tensor, edge_index: Adj,
                 if cache is None:
                     edge_index, edge_weight = gcn_norm(  # yapf: disable
                         edge_index, edge_weight, x.size(self.node_dim),
-                        self.improved, self.add_self_loops)
+                        self.improved, self.add_self_loops, self.flow)
                     if self.cached:
                         self._cached_edge_index = (edge_index, edge_weight)
                 else:
@@ -182,7 +185,7 @@ def forward(self, x: Tensor, edge_index: Adj,
                 if cache is None:
                     edge_index = gcn_norm(  # yapf: disable
                         edge_index, edge_weight, x.size(self.node_dim),
-                        self.improved, self.add_self_loops)
+                        self.improved, self.add_self_loops, self.flow)
                     if self.cached:
                         self._cached_adj_t = edge_index
                 else:

torch_geometric/nn/conv/gin_conv.py

@@ -130,8 +130,9 @@ class GINEConv(MessagePassing):
         - **output:** node features :math:`(|\mathcal{V}|, F_{out})` or
           :math:`(|\mathcal{V}_t|, F_{out})` if bipartite
     """
-    def __init__(self, nn: Callable, eps: float = 0., train_eps: bool = False,
-                 edge_dim: Optional[int] = None, **kwargs):
+    def __init__(self, nn: torch.nn.Module, eps: float = 0.,
+                 train_eps: bool = False, edge_dim: Optional[int] = None,
+                 **kwargs):
         kwargs.setdefault('aggr', 'add')
         super().__init__(**kwargs)
         self.nn = nn
@@ -141,11 +142,16 @@ def __init__(self, nn: Callable, eps: float = 0., train_eps: bool = False,
         else:
             self.register_buffer('eps', torch.Tensor([eps]))
         if edge_dim is not None:
-            if hasattr(self.nn[0], 'in_features'):
-                in_channels = self.nn[0].in_features
+            if isinstance(self.nn, torch.nn.Sequential):
+                nn = self.nn[0]
+            if hasattr(nn, 'in_features'):
+                in_channels = nn.in_features
+            elif hasattr(nn, 'in_channels'):
+                in_channels = nn.in_channels
             else:
-                in_channels = self.nn[0].in_channels
+                raise ValueError("Could not infer input channels from `nn`.")
             self.lin = Linear(edge_dim, in_channels)
+
         else:
             self.lin = None
         self.reset_parameters()

torch_geometric/nn/conv/lg_conv.py

@@ -42,11 +42,12 @@ def forward(self, x: Tensor, edge_index: Adj,
         """"""
         if self.normalize and isinstance(edge_index, Tensor):
             out = gcn_norm(edge_index, edge_weight, x.size(self.node_dim),
-                           add_self_loops=False, dtype=x.dtype)
+                           add_self_loops=False, flow=self.flow, dtype=x.dtype)
             edge_index, edge_weight = out
         elif self.normalize and isinstance(edge_index, SparseTensor):
             edge_index = gcn_norm(edge_index, None, x.size(self.node_dim),
-                                  add_self_loops=False, dtype=x.dtype)
+                                  add_self_loops=False, flow=self.flow,
+                                  dtype=x.dtype)
 
         # propagate_type: (x: Tensor, edge_weight: OptTensor)
         return self.propagate(edge_index, x=x, edge_weight=edge_weight,

torch_geometric/nn/conv/pdn_conv.py

@@ -100,10 +100,11 @@ def forward(self, x: Tensor, edge_index: Adj,
             if isinstance(edge_index, Tensor):
                 edge_index, edge_attr = gcn_norm(edge_index, edge_attr,
                                                  x.size(self.node_dim), False,
-                                                 self.add_self_loops)
+                                                 self.add_self_loops,
+                                                 self.flow)
             elif isinstance(edge_index, SparseTensor):
                 edge_index = gcn_norm(edge_index, None, x.size(self.node_dim),
-                                      False, self.add_self_loops)
+                                      False, self.add_self_loops, self.flow)
 
         x = self.lin(x)
 

torch_geometric/nn/conv/sg_conv.py

@@ -83,11 +83,11 @@ def forward(self, x: Tensor, edge_index: Adj,
             if isinstance(edge_index, Tensor):
                 edge_index, edge_weight = gcn_norm(  # yapf: disable
                     edge_index, edge_weight, x.size(self.node_dim), False,
-                    self.add_self_loops, dtype=x.dtype)
+                    self.add_self_loops, self.flow, dtype=x.dtype)
             elif isinstance(edge_index, SparseTensor):
                 edge_index = gcn_norm(  # yapf: disable
                     edge_index, edge_weight, x.size(self.node_dim), False,
-                    self.add_self_loops, dtype=x.dtype)
+                    self.add_self_loops, self.flow, dtype=x.dtype)
 
             for k in range(self.K):
                 # propagate_type: (x: Tensor, edge_weight: OptTensor)

torch_geometric/nn/conv/tag_conv.py

@@ -75,12 +75,13 @@ def forward(self, x: Tensor, edge_index: Adj,
             if isinstance(edge_index, Tensor):
                 edge_index, edge_weight = gcn_norm(  # yapf: disable
                     edge_index, edge_weight, x.size(self.node_dim),
-                    improved=False, add_self_loops=False, dtype=x.dtype)
+                    improved=False, add_self_loops=False, flow=self.flow,
+                    dtype=x.dtype)
 
             elif isinstance(edge_index, SparseTensor):
                 edge_index = gcn_norm(  # yapf: disable
                     edge_index, edge_weight, x.size(self.node_dim),
-                    add_self_loops=False, dtype=x.dtype)
+                    add_self_loops=False, flow=self.flow, dtype=x.dtype)
 
         out = self.lins[0](x)
         for lin in self.lins[1:]: