From ec85bd6c0b6703b208c992cdcfe086f94c31ed15 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Sun, 26 Mar 2023 19:25:07 +0000 Subject: [PATCH 1/3] update --- .github/workflows/latest_testing.yml | 2 +- test/data/test_batch.py | 175 +++++++++++++++------------ test/data/test_data.py | 12 +- test/data/test_dataset.py | 7 +- test/data/test_graph_store.py | 16 ++- test/data/test_hetero_data.py | 14 +-- 6 files changed, 128 insertions(+), 98 deletions(-) diff --git a/.github/workflows/latest_testing.yml b/.github/workflows/latest_testing.yml index aed4ca6b2195..d17dd6100c53 100644 --- a/.github/workflows/latest_testing.yml +++ b/.github/workflows/latest_testing.yml @@ -52,7 +52,7 @@ jobs: pytest test/test_seed.py pytest test/test_typing.py pytest test/contrib/ - # pytest test/data/ + pytest test/data/ pytest test/datasets/ pytest test/explain/ pytest test/graphgym/ diff --git a/test/data/test_batch.py b/test/data/test_batch.py index 62b836e3f0a5..5a521c03e00b 100644 --- a/test/data/test_batch.py +++ b/test/data/test_batch.py @@ -2,97 +2,70 @@ import numpy as np import torch -from torch_sparse import SparseTensor import torch_geometric from torch_geometric.data import Batch, Data, HeteroData -from torch_geometric.testing import get_random_edge_index +from torch_geometric.testing import get_random_edge_index, withPackage +from torch_geometric.typing import SparseTensor -def test_batch(): +def test_batch_basic(): torch_geometric.set_debug(True) - x1 = torch.tensor([1, 2, 3], dtype=torch.float) - y1 = 1 - x1_sp = SparseTensor.from_dense(x1.view(-1, 1)) - e1 = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]]) - adj1 = SparseTensor.from_edge_index(e1) - s1 = '1' - array1 = ['1', '2'] - x2 = torch.tensor([1, 2], dtype=torch.float) - y2 = 2 - x2_sp = SparseTensor.from_dense(x2.view(-1, 1)) - e2 = torch.tensor([[0, 1], [1, 0]]) - adj2 = SparseTensor.from_edge_index(e2) - s2 = '2' - array2 = ['3', '4', '5'] - x3 = torch.tensor([1, 2, 3, 4], dtype=torch.float) - y3 = 3 - x3_sp = SparseTensor.from_dense(x3.view(-1, 1)) - e3 = torch.tensor([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]]) - adj3 = SparseTensor.from_edge_index(e3) - s3 = '3' - array3 = ['6', '7', '8', '9'] - - data1 = Data(x=x1, y=y1, x_sp=x1_sp, edge_index=e1, adj=adj1, s=s1, - array=array1, num_nodes=3) - data2 = Data(x=x2, y=y2, x_sp=x2_sp, edge_index=e2, adj=adj2, s=s2, - array=array2, num_nodes=2) - data3 = Data(x=x3, y=y3, x_sp=x3_sp, edge_index=e3, adj=adj3, s=s3, - array=array3, num_nodes=4) + x = torch.tensor([1.0, 2.0, 3.0]) + edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]]) + data1 = Data(x=x, y=1, edge_index=edge_index, string='1', array=['1', '2'], + num_nodes=3) + + x = torch.tensor([1.0, 2.0]) + edge_index = torch.tensor([[0, 1], [1, 0]]) + data2 = Data(x=x, y=2, edge_index=edge_index, string='2', + array=['3', '4', '5'], num_nodes=2) + + x = torch.tensor([1.0, 2.0, 3.0, 4.0]) + edge_index = torch.tensor([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]]) + data3 = Data(x=x, y=3, edge_index=edge_index, string='3', + array=['6', '7', '8', '9'], num_nodes=4) batch = Batch.from_data_list([data1]) assert str(batch) == ('DataBatch(x=[3], edge_index=[2, 4], y=[1], ' - 'x_sp=[3, 1, nnz=3], adj=[3, 3, nnz=4], s=[1], ' - 'array=[1], num_nodes=3, batch=[3], ptr=[2])') + 'string=[1], array=[1], num_nodes=3, batch=[3], ' + 'ptr=[2])') assert batch.num_graphs == len(batch) == 1 assert batch.x.tolist() == [1, 2, 3] assert batch.y.tolist() == [1] - assert batch.x_sp.to_dense().view(-1).tolist() == batch.x.tolist() assert batch.edge_index.tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]] - edge_index = torch.stack(batch.adj.coo()[:2], dim=0) - assert edge_index.tolist() == batch.edge_index.tolist() - assert batch.s == ['1'] + assert batch.string == ['1'] assert batch.array == [['1', '2']] assert batch.num_nodes == 3 assert batch.batch.tolist() == [0, 0, 0] assert batch.ptr.tolist() == [0, 3] - batch = Batch.from_data_list([data1, data2, data3], follow_batch=['s']) + batch = Batch.from_data_list([data1, data2, data3], + follow_batch=['string']) assert str(batch) == ('DataBatch(x=[9], edge_index=[2, 12], y=[3], ' - 'x_sp=[9, 1, nnz=9], adj=[9, 9, nnz=12], s=[3], ' - 's_batch=[3], s_ptr=[4], array=[3], num_nodes=9, ' - 'batch=[9], ptr=[4])') + 'string=[3], string_batch=[3], string_ptr=[4], ' + 'array=[3], num_nodes=9, batch=[9], ptr=[4])') assert batch.num_graphs == len(batch) == 3 assert batch.x.tolist() == [1, 2, 3, 1, 2, 1, 2, 3, 4] assert batch.y.tolist() == [1, 2, 3] - assert batch.x_sp.to_dense().view(-1).tolist() == batch.x.tolist() assert batch.edge_index.tolist() == [[0, 1, 1, 2, 3, 4, 5, 6, 6, 7, 7, 8], [1, 0, 2, 1, 4, 3, 6, 5, 7, 6, 8, 7]] - edge_index = torch.stack(batch.adj.coo()[:2], dim=0) - assert edge_index.tolist() == batch.edge_index.tolist() - assert batch.s == ['1', '2', '3'] - assert batch.s_batch.tolist() == [0, 1, 2] - assert batch.s_ptr.tolist() == [0, 1, 2, 3] + assert batch.string == ['1', '2', '3'] + assert batch.string_batch.tolist() == [0, 1, 2] + assert batch.string_ptr.tolist() == [0, 1, 2, 3] assert batch.array == [['1', '2'], ['3', '4', '5'], ['6', '7', '8', '9']] assert batch.num_nodes == 9 assert batch.batch.tolist() == [0, 0, 0, 1, 1, 2, 2, 2, 2] assert batch.ptr.tolist() == [0, 3, 5, 9] - data = batch[0] - assert str(data) == ("Data(x=[3], edge_index=[2, 4], y=[1], " - "x_sp=[3, 1, nnz=3], adj=[3, 3, nnz=4], s='1', " - "array=[2], num_nodes=3)") - data = batch[1] - assert str(data) == ("Data(x=[2], edge_index=[2, 2], y=[1], " - "x_sp=[2, 1, nnz=2], adj=[2, 2, nnz=2], s='2', " - "array=[3], num_nodes=2)") - - data = batch[2] - assert str(data) == ("Data(x=[4], edge_index=[2, 6], y=[1], " - "x_sp=[4, 1, nnz=4], adj=[4, 4, nnz=6], s='3', " - "array=[4], num_nodes=4)") + assert str(batch[0]) == ("Data(x=[3], edge_index=[2, 4], y=[1], " + "string='1', array=[2], num_nodes=3)") + assert str(batch[1]) == ("Data(x=[2], edge_index=[2, 2], y=[1], " + "string='2', array=[3], num_nodes=2)") + assert str(batch[2]) == ("Data(x=[4], edge_index=[2, 6], y=[1], " + "string='3', array=[4], num_nodes=4)") assert len(batch.index_select([1, 0])) == 2 assert len(batch.index_select(torch.tensor([1, 0]))) == 2 @@ -104,43 +77,95 @@ def test_batch(): data_list = batch.to_data_list() assert len(data_list) == 3 - assert len(data_list[0]) == 8 + assert len(data_list[0]) == 6 assert data_list[0].x.tolist() == [1, 2, 3] assert data_list[0].y.tolist() == [1] - assert data_list[0].x_sp.to_dense().view(-1).tolist() == [1, 2, 3] assert data_list[0].edge_index.tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]] - edge_index = torch.stack(data_list[0].adj.coo()[:2], dim=0) - assert edge_index.tolist() == data_list[0].edge_index.tolist() - assert data_list[0].s == '1' + assert data_list[0].string == '1' assert data_list[0].array == ['1', '2'] assert data_list[0].num_nodes == 3 - assert len(data_list[1]) == 8 + assert len(data_list[1]) == 6 assert data_list[1].x.tolist() == [1, 2] assert data_list[1].y.tolist() == [2] - assert data_list[1].x_sp.to_dense().view(-1).tolist() == [1, 2] assert data_list[1].edge_index.tolist() == [[0, 1], [1, 0]] - edge_index = torch.stack(data_list[1].adj.coo()[:2], dim=0) - assert edge_index.tolist() == data_list[1].edge_index.tolist() - assert data_list[1].s == '2' + assert data_list[1].string == '2' assert data_list[1].array == ['3', '4', '5'] assert data_list[1].num_nodes == 2 - assert len(data_list[2]) == 8 + assert len(data_list[2]) == 6 assert data_list[2].x.tolist() == [1, 2, 3, 4] assert data_list[2].y.tolist() == [3] - assert data_list[2].x_sp.to_dense().view(-1).tolist() == [1, 2, 3, 4] assert data_list[2].edge_index.tolist() == [[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]] - edge_index = torch.stack(data_list[2].adj.coo()[:2], dim=0) - assert edge_index.tolist() == data_list[2].edge_index.tolist() - assert data_list[2].s == '3' + assert data_list[2].string == '3' assert data_list[2].array == ['6', '7', '8', '9'] assert data_list[2].num_nodes == 4 torch_geometric.set_debug(True) +@withPackage('torch_sparse') +def test_batch_with_sparse_tensor(): + x = SparseTensor.from_dense(torch.tensor([[1.0], [2.0], [3.0]])) + edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]]) + adj = SparseTensor.from_edge_index(edge_index) + data1 = Data(x=x, adj=adj) + + x = SparseTensor.from_dense(torch.tensor([[1.0], [2.0]])) + edge_index = torch.tensor([[0, 1], [1, 0]]) + adj = SparseTensor.from_edge_index(edge_index) + data2 = Data(x=x, adj=adj) + + x = SparseTensor.from_dense(torch.tensor([[1.0], [2.0], [3.0], [4.0]])) + edge_index = torch.tensor([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]]) + adj = SparseTensor.from_edge_index(edge_index) + data3 = Data(x=x, adj=adj) + + batch = Batch.from_data_list([data1]) + assert str(batch) == ('DataBatch(x=[3, 1, nnz=3], adj=[3, 3, nnz=4], ' + 'batch=[3], ptr=[2])') + assert batch.num_graphs == len(batch) == 1 + assert batch.x.to_dense().tolist() == [[1], [2], [3]] + assert batch.adj.coo()[0].tolist() == [0, 1, 1, 2] + assert batch.adj.coo()[1].tolist() == [1, 0, 2, 1] + assert batch.batch.tolist() == [0, 0, 0] + assert batch.ptr.tolist() == [0, 3] + + batch = Batch.from_data_list([data1, data2, data3]) + + assert str(batch) == ('DataBatch(x=[9, 1, nnz=9], adj=[9, 9, nnz=12], ' + 'batch=[9], ptr=[4])') + assert batch.num_graphs == len(batch) == 3 + assert batch.x.to_dense().view(-1).tolist() == [1, 2, 3, 1, 2, 1, 2, 3, 4] + assert batch.adj.coo()[0].tolist() == [0, 1, 1, 2, 3, 4, 5, 6, 6, 7, 7, 8] + assert batch.adj.coo()[1].tolist() == [1, 0, 2, 1, 4, 3, 6, 5, 7, 6, 8, 7] + assert batch.batch.tolist() == [0, 0, 0, 1, 1, 2, 2, 2, 2] + assert batch.ptr.tolist() == [0, 3, 5, 9] + + assert str(batch[0]) == ("Data(x=[3, 1, nnz=3], adj=[3, 3, nnz=4])") + assert str(batch[1]) == ("Data(x=[2, 1, nnz=2], adj=[2, 2, nnz=2])") + assert str(batch[2]) == ("Data(x=[4, 1, nnz=4], adj=[4, 4, nnz=6])") + + data_list = batch.to_data_list() + assert len(data_list) == 3 + + assert len(data_list[0]) == 2 + assert data_list[0].x.to_dense().tolist() == [[1], [2], [3]] + assert data_list[0].adj.coo()[0].tolist() == [0, 1, 1, 2] + assert data_list[0].adj.coo()[1].tolist() == [1, 0, 2, 1] + + assert len(data_list[1]) == 2 + assert data_list[1].x.to_dense().tolist() == [[1], [2]] + assert data_list[1].adj.coo()[0].tolist() == [0, 1] + assert data_list[1].adj.coo()[1].tolist() == [1, 0] + + assert len(data_list[2]) == 2 + assert data_list[2].x.to_dense().tolist() == [[1], [2], [3], [4]] + assert data_list[2].adj.coo()[0].tolist() == [0, 1, 1, 2, 2, 3] + assert data_list[2].adj.coo()[1].tolist() == [1, 0, 2, 1, 3, 2] + + def test_batching_with_new_dimension(): torch_geometric.set_debug(True) diff --git a/test/data/test_data.py b/test/data/test_data.py index 09fd9d95e0b0..5badd09eb055 100644 --- a/test/data/test_data.py +++ b/test/data/test_data.py @@ -3,11 +3,11 @@ import pytest import torch import torch.multiprocessing as mp -import torch_sparse import torch_geometric from torch_geometric.data import Data from torch_geometric.data.storage import AttrType +from torch_geometric.testing import withPackage def test_data(): @@ -392,13 +392,11 @@ def test_basic_feature_store(): # Graph Store ################################################################# +@withPackage('torch_sparse') def test_basic_graph_store(): r"""Test the core graph store API.""" data = Data() - edge_index = torch.LongTensor([[0, 1], [1, 2]]) - adj = torch_sparse.SparseTensor(row=edge_index[0], col=edge_index[1]) - def assert_equal_tensor_tuple(expected, actual): assert len(expected) == len(actual) for i in range(len(expected)): @@ -406,9 +404,9 @@ def assert_equal_tensor_tuple(expected, actual): # We put all three tensor types: COO, CSR, and CSC, and we get them back # to confirm that `GraphStore` works as intended. - coo = adj.coo()[:-1] - csr = adj.csr()[:-1] - csc = adj.csc()[-2::-1] # (row, colptr) + coo = (torch.tensor([0, 1]), torch.tensor([1, 2])) + csr = (torch.tensor([0, 1, 2, 2]), torch.tensor([1, 2])) + csc = (torch.tensor([0, 1]), torch.tensor([0, 0, 1, 2])) # Put: data.put_edge_index(coo, layout='coo', size=(3, 3)) diff --git a/test/data/test_dataset.py b/test/data/test_dataset.py index abe0e06b2d2b..067a055a12f1 100644 --- a/test/data/test_dataset.py +++ b/test/data/test_dataset.py @@ -2,9 +2,10 @@ import pytest import torch -from torch_sparse import SparseTensor from torch_geometric.data import Data, HeteroData, InMemoryDataset +from torch_geometric.testing import withPackage +from torch_geometric.typing import SparseTensor class MyTestDataset(InMemoryDataset): @@ -117,6 +118,7 @@ def test_to_datapipe(): assert torch.equal(dataset[1].edge_index, list(dp)[1].edge_index) +@withPackage('torch_sparse') def test_in_memory_sparse_tensor_dataset(): x = torch.randn(11, 16) adj = SparseTensor( @@ -279,7 +281,8 @@ def tr(n, m): assert dataset[3].xs[1].size() == (16, 4) -def test_lists_of_SparseTensors(): +@withPackage('torch_sparse') +def test_lists_of_sparse_tensors(): e1 = torch.tensor([[4, 1, 3, 2, 2, 3], [1, 3, 2, 3, 3, 2]]) e2 = torch.tensor([[0, 1, 4, 7, 2, 9], [7, 2, 2, 1, 4, 7]]) e3 = torch.tensor([[3, 5, 1, 2, 3, 3], [5, 0, 2, 1, 3, 7]]) diff --git a/test/data/test_graph_store.py b/test/data/test_graph_store.py index 960280a62a2c..8628bf699c6e 100644 --- a/test/data/test_graph_store.py +++ b/test/data/test_graph_store.py @@ -1,9 +1,13 @@ import pytest import torch -from torch_sparse import SparseTensor from torch_geometric.data.graph_store import EdgeAttr, EdgeLayout from torch_geometric.testing import MyGraphStore, get_random_edge_index +from torch_geometric.utils import ( + to_torch_coo_tensor, + to_torch_csc_tensor, + to_torch_csr_tensor, +) def test_graph_store(): @@ -39,9 +43,13 @@ def test_graph_store(): def test_graph_store_conversion(): graph_store = MyGraphStore() - coo = (row, col) = get_random_edge_index(100, 100, 300) - adj = SparseTensor(row=row, col=col, sparse_sizes=(100, 100)) - csr, csc = adj.csr()[:2], adj.csc()[:2][::-1] + edge_index = get_random_edge_index(100, 100, 300) + adj = to_torch_coo_tensor(edge_index, size=(100, 100)) + coo = (adj.indices()[0], adj.indices()[1]) + adj = to_torch_csr_tensor(edge_index, size=(100, 100)) + csr = (adj.crow_indices(), adj.col_indices()) + adj = to_torch_csc_tensor(edge_index, size=(100, 100)) + csc = (adj.row_indices(), adj.ccol_indices()) graph_store.put_edge_index(coo, ('v', '1', 'v'), 'coo', size=(100, 100)) graph_store.put_edge_index(csr, ('v', '2', 'v'), 'csr', size=(100, 100)) diff --git a/test/data/test_hetero_data.py b/test/data/test_hetero_data.py index ee345709d925..1ea3c79bbe99 100644 --- a/test/data/test_hetero_data.py +++ b/test/data/test_hetero_data.py @@ -2,11 +2,10 @@ import pytest import torch -import torch_sparse from torch_geometric.data import HeteroData from torch_geometric.data.storage import EdgeStorage -from torch_geometric.testing import get_random_edge_index +from torch_geometric.testing import get_random_edge_index, withPackage x_paper = torch.randn(10, 16) x_author = torch.randn(5, 32) @@ -562,13 +561,10 @@ def test_basic_feature_store(): # Graph Store ################################################################# +@withPackage('torch_sparse') def test_basic_graph_store(): data = HeteroData() - edge_index = torch.LongTensor([[0, 1], [1, 2]]) - adj = torch_sparse.SparseTensor(row=edge_index[0], col=edge_index[1], - sparse_sizes=(3, 3)) - def assert_equal_tensor_tuple(expected, actual): assert len(expected) == len(actual) for i in range(len(expected)): @@ -576,9 +572,9 @@ def assert_equal_tensor_tuple(expected, actual): # We put all three tensor types: COO, CSR, and CSC, and we get them back # to confirm that `GraphStore` works as intended. - coo = adj.coo()[:-1] - csr = adj.csr()[:-1] - csc = adj.csc()[-2::-1] # (row, colptr) + coo = (torch.tensor([0, 1]), torch.tensor([1, 2])) + csr = (torch.tensor([0, 1, 2, 2]), torch.tensor([1, 2])) + csc = (torch.tensor([0, 1]), torch.tensor([0, 0, 1, 2])) # Put: data.put_edge_index(coo, layout='coo', edge_type=('a', 'to', 'b'), From a6235c4d3f9b8bc84061f2684f8d94b4a4e6ee46 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Sun, 26 Mar 2023 19:26:17 +0000 Subject: [PATCH 2/3] update --- CHANGELOG.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 57433f7c0b93..ec7477948bf6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,7 +9,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/). ### Changed -- Accelerated sparse tensor conversion routiens ([#7042](https://github.com/pyg-team/pytorch_geometric/pull/7042)) +- Accelerated sparse tensor conversion routiens ([#7042](https://github.com/pyg-team/pytorch_geometric/pull/7042), [#7043](https://github.com/pyg-team/pytorch_geometric/pull/7043)) - Change `torch_sparse.SparseTensor` logic to utilize `torch.sparse_csr` instead ([#7041](https://github.com/pyg-team/pytorch_geometric/pull/7041)) ### Removed From dae4b824c7a299e2f77ef058692761cf311f0d22 Mon Sep 17 00:00:00 2001 From: rusty1s Date: Sun, 26 Mar 2023 19:35:02 +0000 Subject: [PATCH 3/3] update --- test/data/lightning/test_datamodule.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/test/data/lightning/test_datamodule.py b/test/data/lightning/test_datamodule.py index e190988bcb87..55b7803ab52d 100644 --- a/test/data/lightning/test_datamodule.py +++ b/test/data/lightning/test_datamodule.py @@ -177,6 +177,7 @@ def configure_optimizers(self): @onlyCUDA @onlyFullTest +@withPackage('pyg_lib') @withPackage('pytorch_lightning>=2.0.0') @withPackage('torchmetrics>=0.11.0') @pytest.mark.parametrize('loader', ['full', 'neighbor']) @@ -275,6 +276,7 @@ def configure_optimizers(self): @onlyCUDA @onlyFullTest +@withPackage('pyg_lib') @withPackage('pytorch_lightning>=2.0.0') @withPackage('torchmetrics>=0.11.0') def test_lightning_hetero_node_data(get_dataset): @@ -324,6 +326,7 @@ def sample_from_nodes(self, *args, **kwargs): @onlyCUDA @onlyFullTest +@withPackage('pyg_lib') @withPackage('pytorch_lightning') def test_lightning_hetero_link_data(): torch.manual_seed(12345) @@ -383,6 +386,7 @@ def test_lightning_hetero_link_data(): assert 'edge_label_time' in batch['author', 'paper'] +@withPackage('pyg_lib') @withPackage('pytorch_lightning') def test_lightning_hetero_link_data_custom_store(): torch.manual_seed(12345) @@ -419,6 +423,7 @@ def test_lightning_hetero_link_data_custom_store(): assert 'edge_label_index' in batch['author', 'paper'] +@withPackage('pyg_lib') @withPackage('pytorch_lightning') def test_eval_loader_kwargs(get_dataset): data = get_dataset(name='Cora')[0]