Added support for torch.sparse.Tensor in DataLoader (pyg-team#7252)

this implementation isnt working yet, it currently fails with shape
mismatch on a Linear layer but passes the collate part

example repro:
`cd /opt/pyg; pip uninstall -y torch-geometric torch-sparse; rm -rf
pytorch_geometric; git clone -b collate_fix
https://github.com/pyg-team/pytorch_geometric.git; cd
/opt/pyg/pytorch_geometric; pip install .; python3 examples/gcn2_ppi.py`
```
e_idxs_to_cat.size()= [torch.Size([2, 48146]), torch.Size([2, 88335])]
value.size()= torch.Size([4693, 2815])
Traceback (most recent call last):
  File "examples/gcn2_ppi.py", line 93, in <module>
    loss = train()
  File "examples/gcn2_ppi.py", line 70, in train
    loss = criterion(model(data.x, data.adj_t), data.y)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1533, in _call_impl
    return forward_call(*args, **kwargs)
  File "examples/gcn2_ppi.py", line 46, in forward
    h = conv(h, x_0, adj_t)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1533, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric/nn/conv/gcn2_conv.py", line 138, in forward
    x = self.propagate(edge_index, x=x, edge_weight=edge_weight, size=None)
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric/nn/conv/message_passing.py", line 437, in propagate
    out = self.message_and_aggregate(edge_index, **msg_aggr_kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric/nn/conv/gcn2_conv.py", line 159, in message_and_aggregate
    return spmm(adj_t, x, reduce=self.aggr)
  File "/usr/local/lib/python3.8/dist-packages/torch_geometric/utils/spmm.py", line 80, in spmm
    return torch.sparse.mm(src, other)
```
(just remove the check that triggers `This example requires
'torch-sparse'`)

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: rusty1s <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added support for `torch.sparse.Tensor` in `DataLoader` ([#7252](https://github.com/pyg-team/pytorch_geometric/pull/7252))
 - Added `save` and `load` methods to `InMemoryDataset` ([#7250](https://github.com/pyg-team/pytorch_geometric/pull/7250))
 - Added an example for heterogeneous GNN explanation via `CaptumExplainer` ([#7096](https://github.com/pyg-team/pytorch_geometric/pull/7096))
 - Added `visualize_feature_importance` functionality to `HeteroExplanation` ([#7096](https://github.com/pyg-team/pytorch_geometric/pull/7096))

examples/gcn2_ppi.py

@@ -9,10 +9,6 @@
 from torch_geometric.datasets import PPI
 from torch_geometric.loader import DataLoader
 from torch_geometric.nn import GCN2Conv
-from torch_geometric.typing import WITH_TORCH_SPARSE
-
-if not WITH_TORCH_SPARSE:
-    quit("This example requires 'torch-sparse'")
 
 path = osp.join(osp.dirname(osp.realpath(__file__)), '..', 'data', 'GCN2_PPI')
 pre_transform = T.Compose([T.GCNNorm(), T.ToSparseTensor()])

test/data/test_batch.py

@@ -1,12 +1,14 @@
 import os.path as osp
 
 import numpy as np
+import pytest
 import torch
 
 import torch_geometric
 from torch_geometric.data import Batch, Data, HeteroData
 from torch_geometric.testing import get_random_edge_index, withPackage
 from torch_geometric.typing import SparseTensor
+from torch_geometric.utils import to_edge_index, to_torch_sparse_tensor
 
 
 def test_batch_basic():
@@ -466,12 +468,10 @@ def tr(n, m):
     d4 = Data(xs=[tr(4, 3), tr(16, 4), tr(1, 2)], a={"aa": tr(8, 3)},
               x=tr(8, 5))
 
-    # Dataset
     data_list = [d1, d2, d3, d4]
 
     batch = Batch.from_data_list(data_list, follow_batch=['xs', 'a'])
 
-    # assert shapes
     assert batch.xs[0].shape == (19, 3)
     assert batch.xs[1].shape == (56, 4)
     assert batch.xs[2].shape == (7, 2)
@@ -480,7 +480,6 @@ def tr(n, m):
     assert len(batch.xs_batch) == 3
     assert len(batch.a_batch) == 1
 
-    # assert _batch
     assert batch.xs_batch[0].tolist() == \
            [0, 0, 0, 0, 1, 1, 1, 1, 1, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3]
     assert batch.xs_batch[1].tolist() == \
@@ -490,3 +489,31 @@ def tr(n, m):
 
     assert batch.a_batch['aa'].tolist() == \
            [0] * 11 + [1] * 2 + [2] * 4 + [3] * 8
+
+
+@withPackage('torch>=2.0.0')
+@pytest.mark.parametrize('layout', [
+    torch.sparse_coo,
+    torch.sparse_csr,
+    torch.sparse_csc,
+])
+def test_torch_sparse_batch(layout):
+    x_dense = torch.randn(3, 4)
+    x = x_dense.to_sparse(layout=layout)
+    edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
+    edge_attr = torch.rand(4)
+    adj = to_torch_sparse_tensor(edge_index, edge_attr, layout=layout)
+
+    data = Data(x=x, adj=adj)
+
+    batch = Batch.from_data_list([data, data])
+
+    assert batch.x.size() == (6, 4)
+    assert batch.x.layout == layout
+    assert torch.equal(batch.x.to_dense(), torch.cat([x_dense, x_dense], 0))
+
+    assert batch.adj.size() == (6, 6)
+    assert batch.adj.layout == layout
+    out = to_edge_index(batch.adj.to_sparse(layout=torch.sparse_csr))
+    assert torch.equal(out[0], torch.cat([edge_index, edge_index + 3], 1))
+    assert torch.equal(out[1], torch.cat([edge_attr, edge_attr], 0))

torch_geometric/data/collate.py

@@ -9,6 +9,8 @@
 from torch_geometric.data.data import BaseData
 from torch_geometric.data.storage import BaseStorage, NodeStorage
 from torch_geometric.typing import SparseTensor, torch_sparse
+from torch_geometric.utils import is_sparse, is_torch_sparse_tensor
+from torch_geometric.utils.sparse import cat
 
 
 def collate(
@@ -122,7 +124,7 @@ def _collate(
 
     elem = values[0]
 
-    if isinstance(elem, Tensor):
+    if isinstance(elem, Tensor) and not is_sparse(elem):
         # Concatenate a list of `torch.Tensor` along the `cat_dim`.
         # NOTE: We need to take care of incrementing elements appropriately.
         key = str(key)
@@ -160,15 +162,18 @@ def _collate(
         value = torch.cat(values, dim=cat_dim or 0, out=out)
         return value, slices, incs
 
-    elif isinstance(elem, SparseTensor) and increment:
+    elif is_sparse(elem) and increment:
         # Concatenate a list of `SparseTensor` along the `cat_dim`.
         # NOTE: `cat_dim` may return a tuple to allow for diagonal stacking.
         key = str(key)
         cat_dim = data_list[0].__cat_dim__(key, elem, stores[0])
         cat_dims = (cat_dim, ) if isinstance(cat_dim, int) else cat_dim
         repeats = [[value.size(dim) for dim in cat_dims] for value in values]
         slices = cumsum(repeats)
-        value = torch_sparse.cat(values, dim=cat_dim)
+        if is_torch_sparse_tensor(elem):
+            value = cat(values, dim=cat_dim)
+        else:
+            value = torch_sparse.cat(values, dim=cat_dim)
         return value, slices, None
 
     elif isinstance(elem, (int, float)):

torch_geometric/data/data.py

@@ -37,7 +37,7 @@
     OptTensor,
     SparseTensor,
 )
-from torch_geometric.utils import select, subgraph
+from torch_geometric.utils import is_sparse, select, subgraph
 
 
 class BaseData:
@@ -518,7 +518,7 @@ def update(self, data: Union['Data', Dict[str, Any]]) -> 'Data':
         return self
 
     def __cat_dim__(self, key: str, value: Any, *args, **kwargs) -> Any:
-        if isinstance(value, SparseTensor) and 'adj' in key:
+        if is_sparse(value) and 'adj' in key:
             return (0, 1)
         elif 'index' in key or key == 'face':
             return -1

torch_geometric/data/separate.py

@@ -65,7 +65,8 @@ def _separate(
         start, end = int(slices[idx]), int(slices[idx + 1])
         value = narrow(value, cat_dim or 0, start, end - start)
         value = value.squeeze(0) if cat_dim is None else value
-        if decrement and (incs.dim() > 1 or int(incs[idx]) != 0):
+        if (decrement and incs is not None
+                and (incs.dim() > 1 or int(incs[idx]) != 0)):
             value = value - incs[idx].to(value.device)
         return value
 

torch_geometric/utils/__init__.py

@@ -27,7 +27,8 @@
 from .nested import to_nested_tensor, from_nested_tensor
 from .sparse import (dense_to_sparse, is_sparse, is_torch_sparse_tensor,
                      to_torch_coo_tensor, to_torch_csr_tensor,
-                     to_torch_csc_tensor, to_edge_index)
+                     to_torch_csc_tensor, to_torch_sparse_tensor,
+                     to_edge_index)
 from .spmm import spmm
 from .unbatch import unbatch, unbatch_edge_index
 from .one_hot import one_hot
@@ -99,6 +100,7 @@
     'to_torch_coo_tensor',
     'to_torch_csr_tensor',
     'to_torch_csc_tensor',
+    'to_torch_sparse_tensor',
     'to_edge_index',
     'spmm',
     'unbatch',

torch_geometric/utils/sparse.py

@@ -1,4 +1,4 @@
-from typing import Any, Optional, Tuple, Union
+from typing import Any, List, Optional, Tuple, Union
 
 import torch
 from torch import Tensor
@@ -239,6 +239,44 @@ def to_torch_csc_tensor(
     return adj
 
 
+def to_torch_sparse_tensor(
+    edge_index: Tensor,
+    edge_attr: Optional[Tensor] = None,
+    size: Optional[Union[int, Tuple[int, int]]] = None,
+    is_coalesced: bool = False,
+    layout: torch.layout = torch.sparse_coo,
+):
+    r"""Converts a sparse adjacency matrix defined by edge indices and edge
+    attributes to a :class:`torch.sparse.Tensor` with custom :obj:`layout`.
+    See :meth:`~torch_geometric.utils.to_edge_index` for the reverse operation.
+
+    Args:
+        edge_index (LongTensor): The edge indices.
+        edge_attr (Tensor, optional): The edge attributes.
+            (default: :obj:`None`)
+        size (int or (int, int), optional): The size of the sparse matrix.
+            If given as an integer, will create a quadratic sparse matrix.
+            If set to :obj:`None`, will infer a quadratic sparse matrix based
+            on :obj:`edge_index.max() + 1`. (default: :obj:`None`)
+        is_coalesced (bool): If set to :obj:`True`, will assume that
+            :obj:`edge_index` is already coalesced and thus avoids expensive
+            computation. (default: :obj:`False`)
+        layout (torch.layout, optional): The layout of the output sparse tensor
+            (:obj:`torch.sparse_coo`, :obj:`torch.sparse_csr`,
+            :obj:`torch.sparse_csc`). (default: :obj:`torch.sparse_coo`)
+
+    :rtype: :class:`torch.sparse.Tensor`
+    """
+    if layout == torch.sparse_coo:
+        return to_torch_coo_tensor(edge_index, edge_attr, size, is_coalesced)
+    if layout == torch.sparse_csr:
+        return to_torch_csr_tensor(edge_index, edge_attr, size, is_coalesced)
+    if layout == torch.sparse_csc:
+        return to_torch_csc_tensor(edge_index, edge_attr, size, is_coalesced)
+
+    raise ValueError(f"Unexpected sparse tensor layout (got '{layout}')")
+
+
 def to_edge_index(adj: Union[Tensor, SparseTensor]) -> Tuple[Tensor, Tensor]:
     r"""Converts a :class:`torch.sparse.Tensor` or a
     :class:`torch_sparse.SparseTensor` to edge indices and edge attributes.
@@ -341,3 +379,42 @@ def ptr2index(ptr: Tensor) -> Tensor:
 def index2ptr(index: Tensor, size: int) -> Tensor:
     return torch._convert_indices_from_coo_to_csr(
         index, size, out_int32=index.dtype == torch.int32)
+
+
+def cat(tensors: List[Tensor], dim: Union[int, Tuple[int, int]]) -> Tensor:
+    # TODO (matthias) We can make this more efficient by directly operating on
+    # the individual sparse tensor layouts.
+    assert dim in {0, 1, (0, 1)}
+
+    size = [0, 0]
+    edge_indices = []
+    edge_attrs = []
+    for tensor in tensors:
+        assert is_torch_sparse_tensor(tensor)
+        edge_index, edge_attr = to_edge_index(tensor)
+        edge_index = edge_index.clone()
+
+        if dim == 0:
+            edge_index[0] += size[0]
+            size[0] += tensor.size(0)
+            size[1] = max(size[1], tensor.size(1))
+        elif dim == 1:
+            edge_index[1] += size[1]
+            size[0] = max(size[0], tensor.size(0))
+            size[1] += tensor.size(1)
+        else:
+            edge_index[0] += size[0]
+            edge_index[1] += size[1]
+            size[0] += tensor.size(0)
+            size[1] += tensor.size(1)
+
+        edge_indices.append(edge_index)
+        edge_attrs.append(edge_attr)
+
+    return to_torch_sparse_tensor(
+        edge_index=torch.cat(edge_indices, dim=1),
+        edge_attr=torch.cat(edge_attrs, dim=0),
+        size=size,
+        is_coalesced=dim == (0, 1),
+        layout=tensors[0].layout,
+    )