import os.path as osp import numpy as np import pytest import torch import torch_geometric.typing from torch_geometric import EdgeIndex from torch_geometric.data import Data, HeteroData from torch_geometric.loader import NeighborLoader from torch_geometric.nn import GraphConv, to_hetero from torch_geometric.sampler.base import SubgraphType from torch_geometric.testing import ( MyFeatureStore, MyGraphStore, get_random_edge_index, get_random_tensor_frame, onlyLinux, onlyNeighborSampler, onlyOnline, withCUDA, withPackage, ) from torch_geometric.typing import ( WITH_EDGE_TIME_NEIGHBOR_SAMPLE, WITH_PYG_LIB, WITH_TORCH_SPARSE, WITH_WEIGHTED_NEIGHBOR_SAMPLE, TensorFrame, ) from torch_geometric.utils import ( is_undirected, sort_edge_index, to_torch_csr_tensor, to_undirected, ) DTYPES = [ pytest.param(torch.int64, id='int64'), pytest.param(torch.int32, id='int32'), ] SUBGRAPH_TYPES = [ pytest.param(SubgraphType.directional, id='directional'), pytest.param(SubgraphType.bidirectional, id='bidirectional'), pytest.param(SubgraphType.induced, id='induced'), ] FILTER_PER_WORKERS = [ pytest.param(None, id='auto_filter'), pytest.param(True, id='filter_per_worker'), pytest.param(False, id='filter_in_main'), ] def is_subset(subedge_index, edge_index, src_idx, dst_idx): num_nodes = int(edge_index.max()) + 1 idx = num_nodes * edge_index[0] + edge_index[1] subidx = num_nodes * src_idx[subedge_index[0]] + dst_idx[subedge_index[1]] mask = torch.from_numpy(np.isin(subidx.cpu().numpy(), idx.cpu().numpy())) return int(mask.sum()) == mask.numel() @withCUDA @onlyNeighborSampler @pytest.mark.parametrize('dtype', DTYPES) @pytest.mark.parametrize('subgraph_type', SUBGRAPH_TYPES) @pytest.mark.parametrize('filter_per_worker', FILTER_PER_WORKERS) def test_homo_neighbor_loader_basic( device, subgraph_type, dtype, filter_per_worker, ): if dtype != torch.int64 and not torch_geometric.typing.WITH_PT20: return induced = SubgraphType.induced if subgraph_type == SubgraphType.induced and not WITH_TORCH_SPARSE: return if dtype != torch.int64 and (not WITH_PYG_LIB or subgraph_type == induced): return torch.manual_seed(12345) data = Data() data.x = torch.arange(100, device=device) data.edge_index = get_random_edge_index(100, 100, 500, dtype, device) data.edge_attr = torch.arange(500, device=device) loader = NeighborLoader( data, num_neighbors=[5] * 2, batch_size=20, subgraph_type=subgraph_type, filter_per_worker=filter_per_worker, ) assert str(loader) == 'NeighborLoader()' assert len(loader) == 5 batch = loader([0]) assert isinstance(batch, Data) assert batch.n_id[:1].tolist() == [0] for i, batch in enumerate(loader): assert isinstance(batch, Data) assert batch.x.device == device assert batch.x.size(0) <= 100 assert batch.n_id.size() == (batch.num_nodes, ) assert batch.input_id.numel() == batch.batch_size == 20 assert batch.x.min() >= 0 and batch.x.max() < 100 # TODO Re-enable once `EdgeIndex` is stable. assert not isinstance(batch.edge_index, EdgeIndex) # batch.edge_index.validate() # size = (batch.num_nodes, batch.num_nodes) # assert batch.edge_index.sparse_size() == size # assert batch.edge_index.sort_order == 'col' assert batch.edge_index.device == device assert batch.edge_index.min() >= 0 assert batch.edge_index.max() < batch.num_nodes assert batch.edge_attr.device == device assert batch.edge_attr.size(0) == batch.edge_index.size(1) # Input nodes are always sampled first: assert torch.equal( batch.x[:batch.batch_size], torch.arange(i * batch.batch_size, (i + 1) * batch.batch_size, device=device), ) if subgraph_type != SubgraphType.bidirectional: assert batch.e_id.size() == (batch.num_edges, ) assert batch.edge_attr.min() >= 0 assert batch.edge_attr.max() < 500 assert is_subset( batch.edge_index.to(torch.int64), data.edge_index.to(torch.int64), batch.x, batch.x, ) @onlyNeighborSampler @pytest.mark.parametrize('dtype', DTYPES) @pytest.mark.parametrize('subgraph_type', SUBGRAPH_TYPES) def test_hetero_neighbor_loader_basic(subgraph_type, dtype): if dtype != torch.int64 and not torch_geometric.typing.WITH_PT20: return induced = SubgraphType.induced if subgraph_type == SubgraphType.induced and not WITH_TORCH_SPARSE: return if dtype != torch.int64 and (not WITH_PYG_LIB or subgraph_type == induced): return torch.manual_seed(12345) data = HeteroData() data['paper'].x = torch.arange(100) data['author'].x = torch.arange(100, 300) edge_index = get_random_edge_index(100, 100, 500, dtype) data['paper', 'paper'].edge_index = edge_index data['paper', 'paper'].edge_attr = torch.arange(500) edge_index = get_random_edge_index(100, 200, 1000, dtype) data['paper', 'author'].edge_index = edge_index data['paper', 'author'].edge_attr = torch.arange(500, 1500) edge_index = get_random_edge_index(200, 100, 1000, dtype) data['author', 'paper'].edge_index = edge_index data['author', 'paper'].edge_attr = torch.arange(1500, 2500) r1, c1 = data['paper', 'paper'].edge_index r2, c2 = data['paper', 'author'].edge_index + torch.tensor([[0], [100]]) r3, c3 = data['author', 'paper'].edge_index + torch.tensor([[100], [0]]) batch_size = 20 with pytest.raises(ValueError, match="hops must be the same across all"): loader = NeighborLoader( data, num_neighbors={ ('paper', 'to', 'paper'): [-1], ('paper', 'to', 'author'): [-1, -1], ('author', 'to', 'paper'): [-1, -1], }, input_nodes='paper', batch_size=batch_size, subgraph_type=subgraph_type, ) next(iter(loader)) loader = NeighborLoader( data, num_neighbors=[10] * 2, input_nodes='paper', batch_size=batch_size, subgraph_type=subgraph_type, ) assert str(loader) == 'NeighborLoader()' assert len(loader) == (100 + batch_size - 1) // batch_size for batch in loader: assert isinstance(batch, HeteroData) # Test node type selection: assert set(batch.node_types) == {'paper', 'author'} assert batch['paper'].n_id.size() == (batch['paper'].num_nodes, ) assert batch['paper'].x.size(0) <= 100 assert batch['paper'].input_id.numel() == batch_size assert batch['paper'].batch_size == batch_size assert batch['paper'].x.min() >= 0 and batch['paper'].x.max() < 100 assert batch['author'].n_id.size() == (batch['author'].num_nodes, ) assert batch['author'].x.size(0) <= 200 assert batch['author'].x.min() >= 100 and batch['author'].x.max() < 300 # Test edge type selection: assert set(batch.edge_types) == {('paper', 'to', 'paper'), ('paper', 'to', 'author'), ('author', 'to', 'paper')} for edge_type, edge_index in batch.edge_index_dict.items(): src, _, dst = edge_type # TODO Re-enable once `EdgeIndex` is stable. assert not isinstance(edge_index, EdgeIndex) # edge_index.validate() # size = (batch[src].num_nodes, batch[dst].num_nodes) # assert edge_index.sparse_size() == size # assert edge_index.sort_order == 'col' row, col = batch['paper', 'paper'].edge_index assert row.min() >= 0 and row.max() < batch['paper'].num_nodes assert col.min() >= 0 and col.max() < batch['paper'].num_nodes if subgraph_type != SubgraphType.bidirectional: assert batch['paper', 'paper'].e_id.size() == (row.numel(), ) value = batch['paper', 'paper'].edge_attr assert value.min() >= 0 and value.max() < 500 assert is_subset( batch['paper', 'paper'].edge_index.to(torch.int64), data['paper', 'paper'].edge_index.to(torch.int64), batch['paper'].x, batch['paper'].x, ) elif subgraph_type != SubgraphType.directional: assert 'e_id' not in batch['paper', 'paper'] assert 'edge_attr' not in batch['paper', 'paper'] assert is_undirected(batch['paper', 'paper'].edge_index) row, col = batch['paper', 'author'].edge_index assert row.min() >= 0 and row.max() < batch['paper'].num_nodes assert col.min() >= 0 and col.max() < batch['author'].num_nodes if subgraph_type != SubgraphType.bidirectional: assert batch['paper', 'author'].e_id.size() == (row.numel(), ) value = batch['paper', 'author'].edge_attr assert value.min() >= 500 and value.max() < 1500 assert is_subset( batch['paper', 'author'].edge_index.to(torch.int64), data['paper', 'author'].edge_index.to(torch.int64), batch['paper'].x, batch['author'].x - 100, ) elif subgraph_type != SubgraphType.directional: assert 'e_id' not in batch['paper', 'author'] assert 'edge_attr' not in batch['paper', 'author'] edge_index1 = batch['paper', 'author'].edge_index edge_index2 = batch['author', 'paper'].edge_index assert torch.equal( edge_index1, sort_edge_index(edge_index2.flip([0]), sort_by_row=False), ) row, col = batch['author', 'paper'].edge_index assert row.min() >= 0 and row.max() < batch['author'].num_nodes assert col.min() >= 0 and col.max() < batch['paper'].num_nodes if subgraph_type != SubgraphType.bidirectional: assert batch['author', 'paper'].e_id.size() == (row.numel(), ) value = batch['author', 'paper'].edge_attr assert value.min() >= 1500 and value.max() < 2500 assert is_subset( batch['author', 'paper'].edge_index.to(torch.int64), data['author', 'paper'].edge_index.to(torch.int64), batch['author'].x - 100, batch['paper'].x, ) elif subgraph_type != SubgraphType.directional: assert 'e_id' not in batch['author', 'paper'] assert 'edge_attr' not in batch['author', 'paper'] edge_index1 = batch['author', 'paper'].edge_index edge_index2 = batch['paper', 'author'].edge_index assert torch.equal( edge_index1, sort_edge_index(edge_index2.flip([0]), sort_by_row=False), ) # Test for isolated nodes (there shouldn't exist any): assert not batch.has_isolated_nodes() @onlyOnline @onlyNeighborSampler @pytest.mark.parametrize('subgraph_type', SUBGRAPH_TYPES) def test_homo_neighbor_loader_on_karate(get_dataset, subgraph_type): if subgraph_type == SubgraphType.induced and not WITH_TORCH_SPARSE: return dataset = get_dataset(name='karate') data = dataset[0] mask = data.edge_index[0] < data.edge_index[1] edge_index = data.edge_index[:, mask] edge_weight = torch.rand(edge_index.size(1)) data.edge_index, data.edge_weight = to_undirected(edge_index, edge_weight) split_idx = torch.arange(5, 8) loader = NeighborLoader( data, num_neighbors=[-1, -1], batch_size=split_idx.numel(), input_nodes=split_idx, subgraph_type=subgraph_type, ) assert len(loader) == 1 batch = next(iter(loader)) batch_size = batch.batch_size class GNN(torch.nn.Module): def __init__(self, in_channels, hidden_channels, out_channels): super().__init__() self.conv1 = GraphConv(in_channels, hidden_channels) self.conv2 = GraphConv(hidden_channels, out_channels) def forward(self, x, edge_index, edge_weight): x = self.conv1(x, edge_index, edge_weight).relu() x = self.conv2(x, edge_index, edge_weight) return x model = GNN(dataset.num_features, 16, dataset.num_classes) out1 = model(data.x, data.edge_index, data.edge_weight)[split_idx] out2 = model(batch.x, batch.edge_index, batch.edge_weight)[:batch_size] assert torch.allclose(out1, out2, atol=1e-6) @onlyOnline @onlyNeighborSampler @pytest.mark.parametrize('subgraph_type', SUBGRAPH_TYPES) def test_hetero_neighbor_loader_on_karate(get_dataset, subgraph_type): if subgraph_type == SubgraphType.induced and not WITH_TORCH_SPARSE: return dataset = get_dataset(name='karate') data = dataset[0] hetero_data = HeteroData() hetero_data['v'].x = data.x hetero_data['v', 'v'].edge_index = data.edge_index split_idx = torch.arange(5, 8) loader = NeighborLoader( hetero_data, num_neighbors=[-1, -1], batch_size=split_idx.numel(), input_nodes=('v', split_idx), subgraph_type=subgraph_type, ) assert len(loader) == 1 hetero_batch = next(iter(loader)) batch_size = hetero_batch['v'].batch_size class GNN(torch.nn.Module): def __init__(self, in_channels, hidden_channels, out_channels): super().__init__() self.conv1 = GraphConv(in_channels, hidden_channels) self.conv2 = GraphConv(hidden_channels, out_channels) def forward(self, x, edge_index): x = self.conv1(x, edge_index).relu() x = self.conv2(x, edge_index) return x model = GNN(dataset.num_features, 16, dataset.num_classes) hetero_model = to_hetero(model, hetero_data.metadata()) out1 = model(data.x, data.edge_index)[split_idx] out2 = hetero_model(hetero_batch.x_dict, hetero_batch.edge_index_dict)['v'][:batch_size] assert torch.allclose(out1, out2, atol=1e-6) @onlyOnline @withPackage('pyg_lib') def test_temporal_hetero_neighbor_loader_on_karate(get_dataset): dataset = get_dataset(name='karate') data = dataset[0] hetero_data = HeteroData() hetero_data['v'].x = data.x hetero_data['v'].time = torch.arange(data.num_nodes, 0, -1) hetero_data['v', 'v'].edge_index = data.edge_index loader = NeighborLoader(hetero_data, num_neighbors=[-1, -1], input_nodes='v', time_attr='time', batch_size=1) for batch in loader: mask = batch['v'].time[0] >= batch['v'].time[1:] assert torch.all(mask) @onlyNeighborSampler def test_custom_neighbor_loader(): # Initialize feature store, graph store, and reference: feature_store = MyFeatureStore() graph_store = MyGraphStore() # Set up node features: x = torch.arange(100, 300) feature_store.put_tensor(x, group_name=None, attr_name='x', index=None) y = torch.arange(100, 300) feature_store.put_tensor(y, group_name=None, attr_name='y', index=None) # COO: edge_index = get_random_edge_index(100, 100, 500, coalesce=True) edge_index = edge_index[:, torch.randperm(edge_index.size(1))] coo = (edge_index[0], edge_index[1]) graph_store.put_edge_index(edge_index=coo, edge_type=None, layout='coo', size=(100, 100)) data = Data(x=x, edge_index=edge_index, y=y, num_nodes=200) # Construct neighbor loaders: loader1 = NeighborLoader(data, batch_size=20, input_nodes=torch.arange(100), num_neighbors=[-1] * 2) loader2 = NeighborLoader((feature_store, graph_store), batch_size=20, input_nodes=torch.arange(100), num_neighbors=[-1] * 2) assert str(loader1) == str(loader2) assert len(loader1) == len(loader2) for batch1, batch2 in zip(loader1, loader2): assert len(batch1) == len(batch2) assert batch1.num_nodes == batch2.num_nodes assert batch1.num_edges == batch2.num_edges assert batch1.batch_size == batch2.batch_size # Mapped indices of neighbors may be differently sorted ... assert torch.allclose(batch1.x.sort()[0], batch2.x.sort()[0]) assert torch.allclose(batch1.y.sort()[0], batch2.y.sort()[0]) @onlyNeighborSampler def test_custom_hetero_neighbor_loader(): # Initialize feature store, graph store, and reference: feature_store = MyFeatureStore() graph_store = MyGraphStore() data = HeteroData() # Set up node features: x = torch.arange(100) data['paper'].x = x feature_store.put_tensor(x, group_name='paper', attr_name='x', index=None) x = torch.arange(100, 300) data['author'].x = x feature_store.put_tensor(x, group_name='author', attr_name='x', index=None) # COO: edge_index = get_random_edge_index(100, 100, 500, coalesce=True) edge_index = edge_index[:, torch.randperm(edge_index.size(1))] data['paper', 'to', 'paper'].edge_index = edge_index coo = (edge_index[0], edge_index[1]) graph_store.put_edge_index(edge_index=coo, edge_type=('paper', 'to', 'paper'), layout='coo', size=(100, 100)) # CSR: edge_index = get_random_edge_index(100, 200, 1000, coalesce=True) data['paper', 'to', 'author'].edge_index = edge_index adj = to_torch_csr_tensor(edge_index, size=(100, 200)) csr = (adj.crow_indices(), adj.col_indices()) graph_store.put_edge_index(edge_index=csr, edge_type=('paper', 'to', 'author'), layout='csr', size=(100, 200)) # CSC: edge_index = get_random_edge_index(200, 100, 1000, coalesce=True) data['author', 'to', 'paper'].edge_index = edge_index adj = to_torch_csr_tensor(edge_index.flip([0]), size=(100, 200)) csc = (adj.col_indices(), adj.crow_indices()) graph_store.put_edge_index(edge_index=csc, edge_type=('author', 'to', 'paper'), layout='csc', size=(200, 100)) # COO (sorted): edge_index = get_random_edge_index(200, 200, 100, coalesce=True) edge_index = edge_index[:, edge_index[1].argsort()] data['author', 'to', 'author'].edge_index = edge_index coo = (edge_index[0], edge_index[1]) graph_store.put_edge_index(edge_index=coo, edge_type=('author', 'to', 'author'), layout='coo', size=(200, 200), is_sorted=True) # Construct neighbor loaders: loader1 = NeighborLoader(data, batch_size=20, input_nodes=('paper', range(100)), num_neighbors=[-1] * 2) loader2 = NeighborLoader((feature_store, graph_store), batch_size=20, input_nodes=('paper', range(100)), num_neighbors=[-1] * 2) assert str(loader1) == str(loader2) assert len(loader1) == len(loader2) for batch1, batch2 in zip(loader1, loader2): # `loader2` explicitly adds `num_nodes` to the batch: assert len(batch1) + 1 == len(batch2) assert batch1['paper'].batch_size == batch2['paper'].batch_size # Mapped indices of neighbors may be differently sorted ... for node_type in data.node_types: assert torch.allclose( batch1[node_type].x.sort()[0], batch2[node_type].x.sort()[0], ) # ... but should sample the exact same number of edges: for edge_type in data.edge_types: assert batch1[edge_type].num_edges == batch2[edge_type].num_edges @onlyOnline @withPackage('pyg_lib') def test_temporal_custom_neighbor_loader_on_karate(get_dataset): dataset = get_dataset(name='karate') data = dataset[0] data.time = torch.arange(data.num_nodes, 0, -1) # Initialize feature store, graph store, and reference: feature_store = MyFeatureStore() graph_store = MyGraphStore() hetero_data = HeteroData() feature_store.put_tensor( data.x, group_name='v', attr_name='x', index=None, ) hetero_data['v'].x = data.x feature_store.put_tensor( data.time, group_name='v', attr_name='time', index=None, ) hetero_data['v'].time = data.time # Sort according to time in local neighborhoods: row, col = data.edge_index perm = ((col * (data.num_nodes + 1)) + data.time[row]).argsort() edge_index = data.edge_index[:, perm] graph_store.put_edge_index( edge_index, edge_type=('v', 'to', 'v'), layout='coo', is_sorted=True, size=(data.num_nodes, data.num_nodes), ) hetero_data['v', 'to', 'v'].edge_index = data.edge_index loader1 = NeighborLoader( hetero_data, num_neighbors=[-1, -1], input_nodes='v', time_attr='time', batch_size=128, ) loader2 = NeighborLoader( (feature_store, graph_store), num_neighbors=[-1, -1], input_nodes='v', time_attr='time', batch_size=128, ) for batch1, batch2 in zip(loader1, loader2): assert torch.equal(batch1['v'].time, batch2['v'].time) @withPackage('pyg_lib', 'torch_sparse') def test_pyg_lib_and_torch_sparse_homo_equality(): edge_index = get_random_edge_index(20, 20, 100) adj = to_torch_csr_tensor(edge_index.flip([0]), size=(20, 20)) colptr, row = adj.crow_indices(), adj.col_indices() seed = torch.arange(10) sample = torch.ops.pyg.neighbor_sample out1 = sample(colptr, row, seed, [-1, -1], None, None, None, None, True) sample = torch.ops.torch_sparse.neighbor_sample out2 = sample(colptr, row, seed, [-1, -1], False, True) row1, col1, node_id1, edge_id1 = out1[:4] node_id2, row2, col2, edge_id2 = out2 assert torch.equal(node_id1, node_id2) assert torch.equal(row1, row2) assert torch.equal(col1, col2) assert torch.equal(edge_id1, edge_id2) @withPackage('pyg_lib', 'torch_sparse') def test_pyg_lib_and_torch_sparse_hetero_equality(): edge_index = get_random_edge_index(20, 10, 50) adj = to_torch_csr_tensor(edge_index.flip([0]), size=(10, 20)) colptr1, row1 = adj.crow_indices(), adj.col_indices() edge_index = get_random_edge_index(10, 20, 50) adj = to_torch_csr_tensor(edge_index.flip([0]), size=(20, 10)) colptr2, row2 = adj.crow_indices(), adj.col_indices() node_types = ['paper', 'author'] edge_types = [('paper', 'to', 'author'), ('author', 'to', 'paper')] colptr_dict = { 'paper__to__author': colptr1, 'author__to__paper': colptr2, } row_dict = { 'paper__to__author': row1, 'author__to__paper': row2, } seed_dict = {'paper': torch.arange(1)} num_neighbors_dict = { 'paper__to__author': [-1, -1], 'author__to__paper': [-1, -1], } sample = torch.ops.pyg.hetero_neighbor_sample out1 = sample(node_types, edge_types, colptr_dict, row_dict, seed_dict, num_neighbors_dict, None, None, None, None, True, False, True, False, "uniform", True) sample = torch.ops.torch_sparse.hetero_neighbor_sample out2 = sample(node_types, edge_types, colptr_dict, row_dict, seed_dict, num_neighbors_dict, 2, False, True) row1_dict, col1_dict, node_id1_dict, edge_id1_dict = out1[:4] node_id2_dict, row2_dict, col2_dict, edge_id2_dict = out2 assert len(node_id1_dict) == len(node_id2_dict) for key in node_id1_dict.keys(): assert torch.equal(node_id1_dict[key], node_id2_dict[key]) assert len(row1_dict) == len(row2_dict) for key in row1_dict.keys(): assert torch.equal(row1_dict[key], row2_dict[key]) assert len(col1_dict) == len(col2_dict) for key in col1_dict.keys(): assert torch.equal(col1_dict[key], col2_dict[key]) assert len(edge_id1_dict) == len(edge_id2_dict) for key in edge_id1_dict.keys(): assert torch.equal(edge_id1_dict[key], edge_id2_dict[key]) @onlyLinux @onlyNeighborSampler def test_memmap_neighbor_loader(tmp_path): path = osp.join(tmp_path, 'x.npy') x = np.memmap(path, dtype=np.float32, mode='w+', shape=(100, 32)) x[:] = np.random.randn(100, 32) data = Data() data.x = np.memmap(path, dtype=np.float32, mode='r', shape=(100, 32)) data.edge_index = get_random_edge_index(100, 100, 500) assert str(data) == 'Data(x=[100, 32], edge_index=[2, 500])' assert data.num_nodes == 100 loader = NeighborLoader(data, num_neighbors=[5] * 2, batch_size=20, num_workers=2) batch = next(iter(loader)) assert batch.num_nodes <= 100 assert isinstance(batch.x, torch.Tensor) assert batch.x.size() == (batch.num_nodes, 32) @withPackage('pyg_lib') def test_homo_neighbor_loader_sampled_info(): edge_index = torch.tensor([ [2, 3, 4, 5, 7, 7, 10, 11, 12, 13], [0, 1, 2, 3, 2, 3, 7, 7, 7, 7], ]) data = Data(edge_index=edge_index, num_nodes=14) loader = NeighborLoader( data, num_neighbors=[1, 2, 4], batch_size=2, shuffle=False, ) batch = next(iter(loader)) assert batch.num_sampled_nodes == [2, 2, 3, 4] assert batch.num_sampled_edges == [2, 4, 4] @withPackage('pyg_lib') def test_hetero_neighbor_loader_sampled_info(): edge_index = torch.tensor([ [2, 3, 4, 5, 7, 7, 10, 11, 12, 13], [0, 1, 2, 3, 2, 3, 7, 7, 7, 7], ]) data = HeteroData() data['paper'].num_nodes = data['author'].num_nodes = 14 data['paper', 'paper'].edge_index = edge_index data['paper', 'author'].edge_index = edge_index data['author', 'paper'].edge_index = edge_index loader = NeighborLoader( data, num_neighbors=[1, 2, 4], batch_size=2, input_nodes='paper', shuffle=False, ) batch = next(iter(loader)) expected_num_sampled_nodes = { 'paper': [2, 2, 3, 4], 'author': [0, 2, 3, 4], } expected_num_sampled_edges = { ('paper', 'to', 'paper'): [2, 4, 4], ('paper', 'to', 'author'): [0, 4, 4], ('author', 'to', 'paper'): [2, 4, 4], } for node_type in batch.node_types: assert (batch[node_type].num_sampled_nodes == expected_num_sampled_nodes[node_type]) for edge_type in batch.edge_types: assert (batch[edge_type].num_sampled_edges == expected_num_sampled_edges[edge_type]) @withPackage('pyg_lib') def test_neighbor_loader_mapping(): edge_index = torch.tensor([ [0, 0, 0, 0, 0, 1, 1, 1, 2, 2, 3, 5], [1, 2, 3, 4, 5, 8, 6, 7, 9, 10, 6, 11], ]) data = Data(edge_index=edge_index, num_nodes=12) loader = NeighborLoader( data, num_neighbors=[1], batch_size=2, shuffle=True, ) for batch in loader: assert torch.equal( batch.n_id[batch.edge_index], data.edge_index[:, batch.e_id], ) @pytest.mark.skipif( not WITH_WEIGHTED_NEIGHBOR_SAMPLE, reason="'pyg-lib' does not support weighted neighbor sampling", ) def test_weighted_homo_neighbor_loader(): edge_index = torch.tensor([ [1, 3, 0, 4], [2, 2, 1, 3], ]) edge_weight = torch.tensor([0.0, 1.0, 0.0, 1.0]) data = Data(num_nodes=5, edge_index=edge_index, edge_weight=edge_weight) loader = NeighborLoader( data, input_nodes=torch.tensor([2]), num_neighbors=[1] * 2, batch_size=1, weight_attr='edge_weight', ) assert len(loader) == 1 batch = next(iter(loader)) assert batch.num_nodes == 3 assert batch.n_id.tolist() == [2, 3, 4] assert batch.num_edges == 2 assert batch.n_id[batch.edge_index].tolist() == [[3, 4], [2, 3]] @pytest.mark.skipif( not WITH_WEIGHTED_NEIGHBOR_SAMPLE, reason="'pyg-lib' does not support weighted neighbor sampling", ) def test_weighted_hetero_neighbor_loader(): edge_index = torch.tensor([ [1, 3, 0, 4], [2, 2, 1, 3], ]) edge_weight = torch.tensor([0.0, 1.0, 0.0, 1.0]) data = HeteroData() data['paper'].num_nodes = 5 data['paper', 'to', 'paper'].edge_index = edge_index data['paper', 'to', 'paper'].edge_weight = edge_weight loader = NeighborLoader( data, input_nodes=('paper', torch.tensor([2])), num_neighbors=[1] * 2, batch_size=1, weight_attr='edge_weight', ) assert len(loader) == 1 batch = next(iter(loader)) assert batch['paper'].num_nodes == 3 assert batch['paper'].n_id.tolist() == [2, 3, 4] assert batch['paper', 'paper'].num_edges == 2 global_edge_index = batch['paper'].n_id[batch['paper', 'paper'].edge_index] assert global_edge_index.tolist() == [[3, 4], [2, 3]] @pytest.mark.skipif( not WITH_EDGE_TIME_NEIGHBOR_SAMPLE, reason="'pyg-lib' does not support weighted neighbor sampling", ) def test_edge_level_temporal_homo_neighbor_loader(): edge_index = torch.tensor([ [0, 1, 1, 2, 2, 3, 3, 4], [1, 0, 2, 1, 3, 2, 4, 3], ]) edge_time = torch.arange(edge_index.size(1)) data = Data(edge_index=edge_index, edge_time=edge_time, num_nodes=5) loader = NeighborLoader( data, num_neighbors=[-1, -1], input_time=torch.tensor([4, 4, 4, 4, 4]), time_attr='edge_time', batch_size=1, ) for batch in loader: assert batch.edge_time.numel() == batch.num_edges if batch.edge_time.numel() > 0: assert batch.edge_time.max() <= 4 @pytest.mark.skipif( not WITH_EDGE_TIME_NEIGHBOR_SAMPLE, reason="'pyg-lib' does not support weighted neighbor sampling", ) def test_edge_level_temporal_hetero_neighbor_loader(): edge_index = torch.tensor([ [0, 1, 1, 2, 2, 3, 3, 4], [1, 0, 2, 1, 3, 2, 4, 3], ]) edge_time = torch.arange(edge_index.size(1)) data = HeteroData() data['A'].num_nodes = 5 data['A', 'A'].edge_index = edge_index data['A', 'A'].edge_time = edge_time loader = NeighborLoader( data, num_neighbors=[-1, -1], input_nodes='A', input_time=torch.tensor([4, 4, 4, 4, 4]), time_attr='edge_time', batch_size=1, ) for batch in loader: assert batch['A', 'A'].edge_time.numel() == batch['A', 'A'].num_edges if batch['A', 'A'].edge_time.numel() > 0: assert batch['A', 'A'].edge_time.max() <= 4 @withCUDA @onlyNeighborSampler @withPackage('torch_frame') def test_neighbor_loader_with_tensor_frame(device): data = Data() data.tf = get_random_tensor_frame(num_rows=100, device=device) data.edge_index = get_random_edge_index(100, 100, 500, device=device) data.edge_attr = get_random_tensor_frame(500, device=device) data.global_tf = get_random_tensor_frame(num_rows=1, device=device) loader = NeighborLoader(data, num_neighbors=[5] * 2, batch_size=20) assert len(loader) == 5 for batch in loader: assert isinstance(batch.tf, TensorFrame) assert batch.tf.device == device assert batch.tf.num_rows == batch.n_id.numel() assert batch.tf == data.tf[batch.n_id] assert isinstance(batch.edge_attr, TensorFrame) assert batch.edge_attr.device == device assert batch.edge_attr.num_rows == batch.e_id.numel() assert batch.edge_attr == data.edge_attr[batch.e_id] assert isinstance(batch.global_tf, TensorFrame) assert batch.global_tf.device == device assert batch.global_tf.num_rows == 1 assert batch.global_tf == data.global_tf @onlyNeighborSampler def test_neighbor_loader_input_id(): data = HeteroData() data['a'].num_nodes = 10 data['b'].num_nodes = 12 row = torch.randint(0, data['a'].num_nodes, (40, )) col = torch.randint(0, data['b'].num_nodes, (40, )) data['a', 'b'].edge_index = torch.stack([row, col], dim=0) data['b', 'a'].edge_index = torch.stack([col, row], dim=0) mask = torch.ones(data['a'].num_nodes, dtype=torch.bool) mask[0] = False loader = NeighborLoader( data, input_nodes=('a', mask), batch_size=2, num_neighbors=[2, 2], ) for i, batch in enumerate(loader): if i < 4: expected = [(2 * i) + 1, (2 * i) + 2] else: expected = [(2 * i) + 1] assert batch['a'].input_id.tolist() == expected @withPackage('pyg_lib') def test_temporal_neighbor_loader_single_link(): data = HeteroData() data['a'].x = torch.arange(10) data['b'].x = torch.arange(10) data['c'].x = torch.arange(10) data['b'].time = torch.arange(0, 10) data['c'].time = torch.arange(1, 11) data['a', 'b'].edge_index = torch.arange(10).view(1, -1).repeat(2, 1) data['b', 'a'].edge_index = torch.arange(10).view(1, -1).repeat(2, 1) data['a', 'c'].edge_index = torch.arange(10).view(1, -1).repeat(2, 1) data['c', 'a'].edge_index = torch.arange(10).view(1, -1).repeat(2, 1) loader = NeighborLoader( data, num_neighbors=[-1], input_nodes='a', time_attr='time', input_time=torch.arange(0, 10), batch_size=10, ) batch = next(iter(loader)) assert batch['a'].num_nodes == 10 assert batch['b'].num_nodes == 10 assert batch['c'].num_nodes == 0