def test_batch(): torch_geometric.set_debug(True) x1 = torch.tensor([1, 2, 3], dtype=torch.float) e1 = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]]) s1 = '1' x2 = torch.tensor([1, 2], dtype=torch.float) e2 = torch.tensor([[0, 1], [1, 0]]) s2 = '2' data = Batch.from_data_list([Data(x1, e1, s=s1), Data(x2, e2, s=s2)]) assert data.__repr__() == ( 'Batch(batch=[5], edge_index=[2, 6], ptr=[3], s=[2], x=[5])') assert len(data) == 5 assert data.x.tolist() == [1, 2, 3, 1, 2] assert data.edge_index.tolist() == [[0, 1, 1, 2, 3, 4], [1, 0, 2, 1, 4, 3]] assert data.s == ['1', '2'] assert data.batch.tolist() == [0, 0, 0, 1, 1] assert data.ptr.tolist() == [0, 3, 5] assert data.num_graphs == 2 data_list = data.to_data_list() assert len(data_list) == 2 assert len(data_list[0]) == 3 assert data_list[0].x.tolist() == [1, 2, 3] assert data_list[0].edge_index.tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]] assert data_list[0].s == '1' assert len(data_list[1]) == 3 assert data_list[1].x.tolist() == [1, 2] assert data_list[1].edge_index.tolist() == [[0, 1], [1, 0]] assert data_list[1].s == '2' torch_geometric.set_debug(True)
def test_batch(): torch_geometric.set_debug(True) x1 = torch.tensor([1, 2, 3], dtype=torch.float) e1 = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]]) x2 = torch.tensor([1, 2], dtype=torch.float) e2 = torch.tensor([[0, 1], [1, 0]]) data = Batch.from_data_list([Data(x1, e1), Data(x2, e2)]) assert data.__repr__() == 'Batch(batch=[5], edge_index=[2, 6], x=[5])' assert len(data) == 3 assert data.x.tolist() == [1, 2, 3, 1, 2] assert data.edge_index.tolist() == [[0, 1, 1, 2, 3, 4], [1, 0, 2, 1, 4, 3]] assert data.batch.tolist() == [0, 0, 0, 1, 1] assert data.num_graphs == 2 data_list = data.to_data_list() assert len(data_list) == 2 assert len(data_list[0]) == 2 assert data_list[0].x.tolist() == [1, 2, 3] assert data_list[0].edge_index.tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]] assert len(data_list[1]) == 2 assert data_list[1].x.tolist() == [1, 2] assert data_list[1].edge_index.tolist() == [[0, 1], [1, 0]] torch_geometric.set_debug(True)
def test_debug(): assert is_debug_enabled() is False set_debug(True) assert is_debug_enabled() is True set_debug(False) assert is_debug_enabled() is False assert is_debug_enabled() is False with set_debug(True): assert is_debug_enabled() is True assert is_debug_enabled() is False assert is_debug_enabled() is False set_debug(True) assert is_debug_enabled() is True with set_debug(False): assert is_debug_enabled() is False assert is_debug_enabled() is True set_debug(False) assert is_debug_enabled() is False assert is_debug_enabled() is False with debug(): assert is_debug_enabled() is True assert is_debug_enabled() is False
def test_debug_data(): torch_geometric.set_debug(True) Data() Data(edge_index=torch.zeros((2, 0), dtype=torch.long), num_nodes=10) Data(face=torch.zeros((3, 0), dtype=torch.long), num_nodes=10) Data(edge_index=torch.tensor([[0, 1], [1, 0]]), edge_attr=torch.randn(2)) Data(x=torch.torch.randn(5, 3), num_nodes=5) Data(pos=torch.torch.randn(5, 3), num_nodes=5) Data(norm=torch.torch.randn(5, 3), num_nodes=5) torch_geometric.set_debug(False)
def test_batch(): torch_geometric.set_debug(True) x1 = torch.tensor([1, 2, 3], dtype=torch.float) e1 = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]]) s1 = '1' x2 = torch.tensor([1, 2], dtype=torch.float) e2 = torch.tensor([[0, 1], [1, 0]]) s2 = '2' batch = Batch.from_data_list([Data(x1, e1, s=s1), Data(x2, e2, s=s2)]) assert batch.__repr__() == ( 'Batch(batch=[5], edge_index=[2, 6], ptr=[3], s=[2], x=[5])') assert len(batch) == 5 assert batch.x.tolist() == [1, 2, 3, 1, 2] assert batch.edge_index.tolist() == [[0, 1, 1, 2, 3, 4], [1, 0, 2, 1, 4, 3]] assert batch.s == ['1', '2'] assert batch.batch.tolist() == [0, 0, 0, 1, 1] assert batch.ptr.tolist() == [0, 3, 5] assert batch.num_graphs == 2 data = batch[0] assert data.__repr__() == ('Data(edge_index=[2, 4], s="1", x=[3])') data = batch[1] assert data.__repr__() == ('Data(edge_index=[2, 2], s="2", x=[2])') assert len(batch.index_select([1, 0])) == 2 assert len(batch.index_select(torch.tensor([1, 0]))) == 2 assert len(batch.index_select(torch.tensor([True, False]))) == 1 assert len(batch.index_select(np.array([1, 0], dtype=np.int64))) == 2 assert len(batch.index_select(np.array([True, False]))) == 1 assert len(batch[:2]) == 2 data_list = batch.to_data_list() assert len(data_list) == 2 assert len(data_list[0]) == 3 assert data_list[0].x.tolist() == [1, 2, 3] assert data_list[0].edge_index.tolist() == [[0, 1, 1, 2], [1, 0, 2, 1]] assert data_list[0].s == '1' assert len(data_list[1]) == 3 assert data_list[1].x.tolist() == [1, 2] assert data_list[1].edge_index.tolist() == [[0, 1], [1, 0]] assert data_list[1].s == '2' torch_geometric.set_debug(True)
def test_batching_with_new_dimension(): torch_geometric.set_debug(True) class MyData(Data): def __cat_dim__(self, key, value, *args, **kwargs): if key == 'foo': return None else: return super().__cat_dim__(key, value, *args, **kwargs) x1 = torch.tensor([1, 2, 3], dtype=torch.float) foo1 = torch.randn(4) y1 = torch.tensor(1) x2 = torch.tensor([1, 2], dtype=torch.float) foo2 = torch.randn(4) y2 = torch.tensor(2) batch = Batch.from_data_list( [MyData(x=x1, foo=foo1, y=y1), MyData(x=x2, foo=foo2, y=y2)]) assert batch.__repr__() == ( 'Batch(x=[5], y=[2], foo=[2, 4], batch=[5], ptr=[3])') assert len(batch) == 5 assert batch.x.tolist() == [1, 2, 3, 1, 2] assert batch.foo.size() == (2, 4) assert batch.foo[0].tolist() == foo1.tolist() assert batch.foo[1].tolist() == foo2.tolist() assert batch.y.tolist() == [1, 2] assert batch.batch.tolist() == [0, 0, 0, 1, 1] assert batch.ptr.tolist() == [0, 3, 5] assert batch.num_graphs == 2 data = batch[0] assert str(data) == ('MyData(x=[3], y=[1], foo=[4])') data = batch[1] assert str(data) == ('MyData(x=[2], y=[1], foo=[4])') torch_geometric.set_debug(True)
def test_batch(): torch_geometric.set_debug(True) x1 = torch.tensor([1, 2, 3], dtype=torch.float) 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) 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) 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, x_sp=x1_sp, edge_index=e1, adj=adj1, s=s1, array=array1, num_nodes=3) data2 = Data(x=x2, x_sp=x2_sp, edge_index=e2, adj=adj2, s=s2, array=array2, num_nodes=2) data3 = Data(x=x3, x_sp=x3_sp, edge_index=e3, adj=adj3, s=s3, array=array3, num_nodes=4) batch = Batch.from_data_list([data1]) assert str(batch) == ('Batch(x=[3], edge_index=[2, 4], ' 'x_sp=[3, 1, nnz=3], adj=[3, 3, nnz=4], s=[1], ' 'array=[1], num_nodes=3, batch=[3], ptr=[2])') assert batch.num_graphs == 1 assert len(batch) == 9 assert batch.x.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], [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.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']) assert str(batch) == ('Batch(x=[9], edge_index=[2, 12], ' 'x_sp=[9, 1, nnz=9], adj=[9, 9, nnz=12], s=[3], ' 's_batch=[3], array=[3], num_nodes=9, batch=[9], ' 'ptr=[4])') assert batch.num_graphs == 3 assert len(batch) == 10 assert batch.x.tolist() == [1, 2, 3, 1, 2, 1, 2, 3, 4] 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.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], 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], 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], x_sp=[4, 1, nnz=4], " "adj=[4, 4, nnz=6], s='3', array=[4], num_nodes=4)") assert len(batch.index_select([1, 0])) == 2 assert len(batch.index_select(torch.tensor([1, 0]))) == 2 assert len(batch.index_select(torch.tensor([True, False]))) == 1 assert len(batch.index_select(np.array([1, 0], dtype=np.int64))) == 2 assert len(batch.index_select(np.array([True, False]))) == 1 assert len(batch[:2]) == 2 data_list = batch.to_data_list() assert len(data_list) == 3 assert len(data_list[0]) == 7 assert data_list[0].x.tolist() == [1, 2, 3] 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].array == ['1', '2'] assert data_list[0].num_nodes == 3 assert len(data_list[1]) == 7 assert data_list[1].x.tolist() == [1, 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].array == ['3', '4', '5'] assert data_list[1].num_nodes == 2 assert len(data_list[2]) == 7 assert data_list[2].x.tolist() == [1, 2, 3, 4] 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].array == ['6', '7', '8', '9'] assert data_list[2].num_nodes == 4 torch_geometric.set_debug(True)
def test_data(): torch_geometric.set_debug(True) x = torch.tensor([[1, 3, 5], [2, 4, 6]], dtype=torch.float).t() edge_index = torch.tensor([[0, 0, 1, 1, 2], [1, 1, 0, 2, 1]]) data = Data(x=x, edge_index=edge_index).to(torch.device('cpu')) N = data.num_nodes assert N == 3 assert data.x.tolist() == x.tolist() assert data['x'].tolist() == x.tolist() assert sorted(data.keys) == ['edge_index', 'x'] assert len(data) == 2 assert 'x' in data and 'edge_index' in data and 'pos' not in data D = data.to_dict() assert len(D) == 2 assert 'x' in D and 'edge_index' in D D = data.to_namedtuple() assert len(D) == 2 assert D.x is not None and D.edge_index is not None assert data.__cat_dim__('x', data.x) == 0 assert data.__cat_dim__('edge_index', data.edge_index) == -1 assert data.__inc__('x', data.x) == 0 assert data.__inc__('edge_index', data.edge_index) == data.num_nodes assert not data.x.is_contiguous() data.contiguous() assert data.x.is_contiguous() assert not data.is_coalesced() data = data.coalesce() assert data.is_coalesced() clone = data.clone() assert clone != data assert len(clone) == len(data) assert clone.x.tolist() == data.x.tolist() assert clone.edge_index.tolist() == data.edge_index.tolist() # Test `data.to_heterogenous()`: out = data.to_heterogeneous() assert torch.allclose(data.x, out['0'].x) assert torch.allclose(data.edge_index, out['0', '0'].edge_index) data.edge_type = torch.tensor([0, 0, 1, 0]) out = data.to_heterogeneous() assert torch.allclose(data.x, out['0'].x) assert [store.num_edges for store in out.edge_stores] == [3, 1] data.edge_type = None data['x'] = x + 1 assert data.x.tolist() == (x + 1).tolist() assert str(data) == 'Data(x=[3, 2], edge_index=[2, 4])' dictionary = {'x': data.x, 'edge_index': data.edge_index} data = Data.from_dict(dictionary) assert sorted(data.keys) == ['edge_index', 'x'] assert not data.has_isolated_nodes() assert not data.has_self_loops() assert data.is_undirected() assert not data.is_directed() assert data.num_nodes == 3 assert data.num_edges == 4 assert data.num_faces is None assert data.num_node_features == 2 assert data.num_features == 2 data.edge_attr = torch.randn(data.num_edges, 2) assert data.num_edge_features == 2 data.edge_attr = None data.x = None assert data.num_nodes == 3 data.edge_index = None assert data.num_nodes is None assert data.num_edges == 0 data.num_nodes = 4 assert data.num_nodes == 4 data = Data(x=x, attribute=x) assert len(data) == 2 assert data.x.tolist() == x.tolist() assert data.attribute.tolist() == x.tolist() face = torch.tensor([[0, 1], [1, 2], [2, 3]]) data = Data(num_nodes=4, face=face) assert data.num_faces == 2 assert data.num_nodes == 4 data = Data(title='test') assert str(data) == "Data(title='test')" assert data.num_node_features == 0 assert data.num_edge_features == 0 key = value = 'test_value' data[key] = value assert data[key] == value del data[value] del data[value] # Deleting unset attributes should work as well. assert data.get(key) is None assert data.get('title') == 'test' torch_geometric.set_debug(False)
def test_data(): torch_geometric.set_debug(True) x = torch.tensor([[1, 3, 5], [2, 4, 6]], dtype=torch.float).t() edge_index = torch.tensor([[0, 0, 1, 1, 2], [1, 1, 0, 2, 1]]) data = Data(x=x, edge_index=edge_index).to(torch.device('cpu')) N = data.num_nodes assert N == 3 assert data.x.tolist() == x.tolist() assert data['x'].tolist() == x.tolist() assert sorted(data.keys) == ['edge_index', 'x'] assert len(data) == 2 assert 'x' in data and 'edge_index' in data and 'pos' not in data D = data.to_dict() assert len(D) == 2 assert 'x' in D and 'edge_index' in D D = data.to_namedtuple() assert len(D) == 2 assert D.x is not None and D.edge_index is not None assert data.__cat_dim__('x', data.x) == 0 assert data.__cat_dim__('edge_index', data.edge_index) == -1 assert data.__inc__('x', data.x) == 0 assert data.__inc__('edge_index', data.edge_index) == data.num_nodes assert not data.x.is_contiguous() data.contiguous() assert data.x.is_contiguous() assert not data.is_coalesced() data = data.coalesce() assert data.is_coalesced() clone = data.clone() assert clone != data assert len(clone) == len(data) assert clone.x.tolist() == data.x.tolist() assert clone.edge_index.tolist() == data.edge_index.tolist() data['x'] = x + 1 assert data.x.tolist() == (x + 1).tolist() assert str(data) == 'Data(x=[3, 2], edge_index=[2, 4])' dictionary = {'x': data.x, 'edge_index': data.edge_index} data = Data.from_dict(dictionary) assert sorted(data.keys) == ['edge_index', 'x'] assert not data.has_isolated_nodes() assert not data.has_self_loops() assert data.is_undirected() assert not data.is_directed() assert data.num_nodes == 3 assert data.num_edges == 4 assert data.num_faces is None assert data.num_node_features == 2 assert data.num_features == 2 data.edge_attr = torch.randn(data.num_edges, 2) assert data.num_edge_features == 2 data.edge_attr = None data.x = None assert data.num_nodes == 3 data.edge_index = None assert data.num_nodes is None assert data.num_edges == 0 data.num_nodes = 4 assert data.num_nodes == 4 data = Data(x=x, attribute=x) assert len(data) == 2 assert data.x.tolist() == x.tolist() assert data.attribute.tolist() == x.tolist() face = torch.tensor([[0, 1], [1, 2], [2, 3]]) data = Data(num_nodes=4, face=face) assert data.num_faces == 2 assert data.num_nodes == 4 data = Data(title='test') assert str(data) == "Data(title='test')" assert data.num_node_features == 0 assert data.num_edge_features == 0 torch_geometric.set_debug(False)
from torch_geometric.nn import SAGEConv import torch import torch_geometric import torch.nn.functional as F torch_geometric.set_debug(True) class GraphSAGE(torch.nn.Module): def __init__(self, hidden_channels, seed): super().__init__() torch.manual_seed(seed) self.conv1 = SAGEConv((-1, -1), hidden_channels, normalize=True) self.conv2 = SAGEConv((-1, -1), 2, normalize=True) def forward(self, x, edge_index): x = self.conv1(x, edge_index) x = x.relu() x = F.dropout(x, p=0.2, training=self.training) x = self.conv2(x, edge_index) return F.log_softmax(x, dim=1)