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)
