Add PyTorch SparseTensor support for GINConv, SAGEConv, and GraphConv (#6532)

This PR aims to add PyTorch SparseTensor support for `GINConv`,
`SAGEConv`, and `GraphConv`, ~~and also update type hints in
`message_and_aggregate` correspondingly~~.

Author: EdisonLeeeee

CHANGELOG.md

@@ -60,7 +60,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `Data.edge_subgraph` and `HeteroData.edge_subgraph` functionalities ([#6193](https://github.com/pyg-team/pytorch_geometric/pull/6193))
 - Added `input_time` option to `LightningNodeData` and `transform_sampler_output` to `NodeLoader` and `LinkLoader` ([#6187](https://github.com/pyg-team/pytorch_geometric/pull/6187))
 - Added `summary` for PyG/PyTorch models ([#5859](https://github.com/pyg-team/pytorch_geometric/pull/5859), [#6161](https://github.com/pyg-team/pytorch_geometric/pull/6161))
-- Started adding `torch.sparse` support to PyG ([#5906](https://github.com/pyg-team/pytorch_geometric/pull/5906), [#5944](https://github.com/pyg-team/pytorch_geometric/pull/5944), [#6003](https://github.com/pyg-team/pytorch_geometric/pull/6003), [#6033](https://github.com/pyg-team/pytorch_geometric/pull/6033), [#6514](https://github.com/pyg-team/pytorch_geometric/pull/6514))
+- Started adding `torch.sparse` support to PyG ([#5906](https://github.com/pyg-team/pytorch_geometric/pull/5906), [#5944](https://github.com/pyg-team/pytorch_geometric/pull/5944), [#6003](https://github.com/pyg-team/pytorch_geometric/pull/6003), [#6033](https://github.com/pyg-team/pytorch_geometric/pull/6033), [#6514](https://github.com/pyg-team/pytorch_geometric/pull/6514), [#6532](https://github.com/pyg-team/pytorch_geometric/pull/6532))
 - Add `inputs_channels` back in training benchmark ([#6154](https://github.com/pyg-team/pytorch_geometric/pull/6154))
 - Added support for dropping nodes in `utils.to_dense_batch` in case `max_num_nodes` is smaller than the number of nodes ([#6124](https://github.com/pyg-team/pytorch_geometric/pull/6124))
 - Added the RandLA-Net architecture as an example ([#5117](https://github.com/pyg-team/pytorch_geometric/pull/5117))

test/nn/conv/test_gin_conv.py

@@ -14,6 +14,7 @@ def test_gin_conv():
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
     row, col = edge_index
     adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
+    adj2 = adj.to_torch_sparse_coo_tensor()
 
     nn = Seq(Lin(16, 32), ReLU(), Lin(32, 32))
     conv = GINConv(nn, train_eps=True)
@@ -25,8 +26,9 @@ def test_gin_conv():
         '))')
     out = conv(x1, edge_index)
     assert out.size() == (4, 32)
-    assert conv(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
-    assert conv(x1, adj.t()).tolist() == out.tolist()
+    assert torch.allclose(conv(x1, edge_index, size=(4, 4)), out, atol=1e-6)
+    assert torch.allclose(conv(x1, adj.t()), out, atol=1e-6)
+    assert torch.allclose(conv(x1, adj2.t()), out, atol=1e-6)
 
     if is_full_test():
         t = '(Tensor, Tensor, Size) -> Tensor'
@@ -39,13 +41,16 @@ def test_gin_conv():
         assert jit(x1, adj.t()).tolist() == out.tolist()
 
     adj = adj.sparse_resize((4, 2))
+    adj2 = adj.to_torch_sparse_coo_tensor()
     out1 = conv((x1, x2), edge_index)
     out2 = conv((x1, None), edge_index, (4, 2))
     assert out1.size() == (2, 32)
     assert out2.size() == (2, 32)
-    assert conv((x1, x2), edge_index, (4, 2)).tolist() == out1.tolist()
-    assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
-    assert conv((x1, None), adj.t()).tolist() == out2.tolist()
+    assert torch.allclose(conv((x1, x2), edge_index, (4, 2)), out1, atol=1e-6)
+    assert torch.allclose(conv((x1, x2), adj.t()), out1, atol=1e-6)
+    assert torch.allclose(conv((x1, None), adj.t()), out2, atol=1e-6)
+    assert torch.allclose(conv((x1, x2), adj2.t()), out1, atol=1e-6)
+    assert torch.allclose(conv((x1, None), adj2.t()), out2, atol=1e-6)
 
     if is_full_test():
         t = '(OptPairTensor, Tensor, Size) -> Tensor'

test/nn/conv/test_graph_conv.py

@@ -13,18 +13,23 @@ def test_graph_conv():
     value = torch.randn(edge_index.size(1))
     adj1 = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
     adj2 = SparseTensor(row=row, col=col, value=value, sparse_sizes=(4, 4))
+    adj3 = adj1.to_torch_sparse_coo_tensor()
+    adj4 = adj2.to_torch_sparse_coo_tensor()
 
     conv = GraphConv(8, 32)
     assert conv.__repr__() == 'GraphConv(8, 32)'
     out11 = conv(x1, edge_index)
     assert out11.size() == (4, 32)
-    assert conv(x1, edge_index, size=(4, 4)).tolist() == out11.tolist()
-    assert conv(x1, adj1.t()).tolist() == out11.tolist()
+    assert torch.allclose(conv(x1, edge_index, size=(4, 4)), out11, atol=1e-6)
+    assert torch.allclose(conv(x1, adj1.t()), out11, atol=1e-6)
+    assert torch.allclose(conv(x1, adj3.t()), out11, atol=1e-6)
 
     out12 = conv(x1, edge_index, value)
     assert out12.size() == (4, 32)
-    assert conv(x1, edge_index, value, size=(4, 4)).tolist() == out12.tolist()
-    assert conv(x1, adj2.t()).tolist() == out12.tolist()
+    assert torch.allclose(conv(x1, edge_index, value, size=(4, 4)), out12,
+                          atol=1e-6)
+    assert torch.allclose(conv(x1, adj2.t()), out12, atol=1e-6)
+    assert torch.allclose(conv(x1, adj4.t()), out12, atol=1e-6)
 
     if is_full_test():
         t = '(Tensor, Tensor, OptTensor, Size) -> Tensor'
@@ -42,6 +47,8 @@ def test_graph_conv():
 
     adj1 = adj1.sparse_resize((4, 2))
     adj2 = adj2.sparse_resize((4, 2))
+    adj3 = adj1.to_torch_sparse_coo_tensor()
+    adj4 = adj2.to_torch_sparse_coo_tensor()
     conv = GraphConv((8, 16), 32)
     assert conv.__repr__() == 'GraphConv((8, 16), 32)'
     out21 = conv((x1, x2), edge_index)
@@ -52,10 +59,14 @@ def test_graph_conv():
     assert out22.size() == (2, 32)
     assert out23.size() == (2, 32)
     assert out24.size() == (2, 32)
-    assert conv((x1, x2), edge_index, size=(4, 2)).tolist() == out21.tolist()
-    assert conv((x1, x2), edge_index, value, (4, 2)).tolist() == out22.tolist()
-    assert conv((x1, x2), adj1.t()).tolist() == out21.tolist()
-    assert conv((x1, x2), adj2.t()).tolist() == out22.tolist()
+    assert torch.allclose(conv((x1, x2), edge_index, size=(4, 2)), out21,
+                          atol=1e-6)
+    assert torch.allclose(conv((x1, x2), edge_index, value, (4, 2)), out22,
+                          atol=1e-6)
+    assert torch.allclose(conv((x1, x2), adj1.t()), out21, atol=1e-6)
+    assert torch.allclose(conv((x1, x2), adj2.t()), out22, atol=1e-6)
+    assert torch.allclose(conv((x1, x2), adj3.t()), out21, atol=1e-6)
+    assert torch.allclose(conv((x1, x2), adj4.t()), out22, atol=1e-6)
 
     if is_full_test():
         t = '(OptPairTensor, Tensor, OptTensor, Size) -> Tensor'

test/nn/conv/test_sage_conv.py

@@ -7,53 +7,62 @@
 
 
 @pytest.mark.parametrize('project', [False, True])
-def test_sage_conv(project):
+@pytest.mark.parametrize('aggr', ['mean', 'sum'])
+def test_sage_conv(project, aggr):
     x1 = torch.randn(4, 8)
     x2 = torch.randn(2, 16)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
     row, col = edge_index
     adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
+    adj2 = adj.to_torch_sparse_coo_tensor()
 
-    conv = SAGEConv(8, 32, project=project)
-    assert str(conv) == 'SAGEConv(8, 32, aggr=mean)'
+    conv = SAGEConv(8, 32, project=project, aggr=aggr)
+    assert str(conv) == f'SAGEConv(8, 32, aggr={aggr})'
     out = conv(x1, edge_index)
     assert out.size() == (4, 32)
-    assert conv(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
-    assert conv(x1, adj.t()).tolist() == out.tolist()
+    assert torch.allclose(conv(x1, edge_index, size=(4, 4)), out, atol=1e-6)
+    assert torch.allclose(conv(x1, adj.t()), out, atol=1e-6)
+    if aggr == 'sum':
+        assert torch.allclose(conv(x1, adj2.t()), out, atol=1e-6)
 
     if is_full_test():
         t = '(Tensor, Tensor, Size) -> Tensor'
         jit = torch.jit.script(conv.jittable(t))
-        assert jit(x1, edge_index).tolist() == out.tolist()
-        assert jit(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
+        assert torch.allclose(jit(x1, edge_index), out, atol=1e-6)
+        assert torch.allclose(jit(x1, edge_index, size=(4, 4)), out, atol=1e-6)
 
         t = '(Tensor, SparseTensor, Size) -> Tensor'
         jit = torch.jit.script(conv.jittable(t))
-        assert jit(x1, adj.t()).tolist() == out.tolist()
+        assert torch.allclose(jit(x1, adj.t()), out, atol=1e-6)
 
     adj = adj.sparse_resize((4, 2))
-    conv = SAGEConv((8, 16), 32, project=project)
-    assert str(conv) == 'SAGEConv((8, 16), 32, aggr=mean)'
+    adj2 = adj.to_torch_sparse_coo_tensor()
+    conv = SAGEConv((8, 16), 32, project=project, aggr=aggr)
+    assert str(conv) == f'SAGEConv((8, 16), 32, aggr={aggr})'
     out1 = conv((x1, x2), edge_index)
     out2 = conv((x1, None), edge_index, (4, 2))
     assert out1.size() == (2, 32)
     assert out2.size() == (2, 32)
-    assert conv((x1, x2), edge_index, (4, 2)).tolist() == out1.tolist()
-    assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
-    assert conv((x1, None), adj.t()).tolist() == out2.tolist()
+    assert torch.allclose(conv((x1, x2), edge_index, (4, 2)), out1, atol=1e-6)
+    assert torch.allclose(conv((x1, x2), adj.t()), out1, atol=1e-6)
+    assert torch.allclose(conv((x1, None), adj.t()), out2, atol=1e-6)
+    if aggr == 'sum':
+        assert torch.allclose(conv((x1, x2), adj2.t()), out1, atol=1e-6)
+        assert torch.allclose(conv((x1, None), adj2.t()), out2, atol=1e-6)
 
     if is_full_test():
         t = '(OptPairTensor, Tensor, Size) -> Tensor'
         jit = torch.jit.script(conv.jittable(t))
-        assert jit((x1, x2), edge_index).tolist() == out1.tolist()
-        assert jit((x1, x2), edge_index, size=(4, 2)).tolist() == out1.tolist()
-        assert jit((x1, None), edge_index,
-                   size=(4, 2)).tolist() == out2.tolist()
+        assert torch.allclose(jit((x1, x2), edge_index), out1, atol=1e-6)
+        assert torch.allclose(jit((x1, x2), edge_index, size=(4, 2)), out1,
+                              atol=1e-6)
+        assert torch.allclose(jit((x1, None), edge_index, size=(4, 2)), out2,
+                              atol=1e-6)
 
         t = '(OptPairTensor, SparseTensor, Size) -> Tensor'
         jit = torch.jit.script(conv.jittable(t))
-        assert jit((x1, x2), adj.t()).tolist() == out1.tolist()
-        assert jit((x1, None), adj.t()).tolist() == out2.tolist()
+        assert torch.allclose(jit((x1, x2), adj.t()), out1, atol=1e-6)
+        assert torch.allclose(jit((x1, None), adj.t()), out2, atol=1e-6)
 
 
 def test_lstm_aggr_sage_conv():

torch_geometric/nn/conv/gin_conv.py

@@ -2,11 +2,12 @@
 
 import torch
 from torch import Tensor
-from torch_sparse import SparseTensor, matmul
+from torch_sparse import SparseTensor
 
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.dense.linear import Linear
 from torch_geometric.typing import Adj, OptPairTensor, OptTensor, Size
+from torch_geometric.utils import spmm
 
 from ..inits import reset
 
@@ -84,8 +85,9 @@ def message(self, x_j: Tensor) -> Tensor:
 
     def message_and_aggregate(self, adj_t: SparseTensor,
                               x: OptPairTensor) -> Tensor:
-        adj_t = adj_t.set_value(None, layout=None)
-        return matmul(adj_t, x[0], reduce=self.aggr)
+        if isinstance(adj_t, SparseTensor):
+            adj_t = adj_t.set_value(None, layout=None)
+        return spmm(adj_t, x[0], reduce=self.aggr)
 
     def __repr__(self) -> str:
         return f'{self.__class__.__name__}(nn={self.nn})'

torch_geometric/nn/conv/graph_conv.py

@@ -1,11 +1,12 @@
 from typing import Tuple, Union
 
 from torch import Tensor
-from torch_sparse import SparseTensor, matmul
+from torch_sparse import SparseTensor
 
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.dense.linear import Linear
 from torch_geometric.typing import Adj, OptPairTensor, OptTensor, Size
+from torch_geometric.utils import spmm
 
 
 class GraphConv(MessagePassing):
@@ -91,4 +92,4 @@ def message(self, x_j: Tensor, edge_weight: OptTensor) -> Tensor:
 
     def message_and_aggregate(self, adj_t: SparseTensor,
                               x: OptPairTensor) -> Tensor:
-        return matmul(adj_t, x[0], reduce=self.aggr)
+        return spmm(adj_t, x[0], reduce=self.aggr)

torch_geometric/nn/conv/sage_conv.py

@@ -3,12 +3,13 @@
 import torch.nn.functional as F
 from torch import Tensor
 from torch.nn import LSTM
-from torch_sparse import SparseTensor, matmul
+from torch_sparse import SparseTensor
 
 from torch_geometric.nn.aggr import Aggregation, MultiAggregation
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.dense.linear import Linear
 from torch_geometric.typing import Adj, OptPairTensor, Size
+from torch_geometric.utils import spmm
 
 
 class SAGEConv(MessagePassing):
@@ -145,8 +146,9 @@ def message(self, x_j: Tensor) -> Tensor:
 
     def message_and_aggregate(self, adj_t: SparseTensor,
                               x: OptPairTensor) -> Tensor:
-        adj_t = adj_t.set_value(None, layout=None)
-        return matmul(adj_t, x[0], reduce=self.aggr)
+        if isinstance(adj_t, SparseTensor):
+            adj_t = adj_t.set_value(None, layout=None)
+        return spmm(adj_t, x[0], reduce=self.aggr)
 
     def __repr__(self) -> str:
         return (f'{self.__class__.__name__}({self.in_channels}, '