def test_hetero_data_functions():
data = HeteroData()
data['paper'].x = x_paper
data['author'].x = x_author
data['paper', 'paper'].edge_index = edge_index_paper_paper
data['paper', 'author'].edge_index = edge_index_paper_author
data['author', 'paper'].edge_index = edge_index_author_paper
data['paper', 'paper'].edge_attr = edge_attr_paper_paper
assert len(data) == 3
assert sorted(data.keys) == ['edge_attr', 'edge_index', 'x']
assert 'x' in data and 'edge_index' in data and 'edge_attr' in data
assert data.num_nodes == 15
assert data.num_edges == 110
assert data.num_node_features == {'paper': 16, 'author': 32}
assert data.num_edge_features == {
('paper', 'to', 'paper'): 8,
('paper', 'to', 'author'): 0,
('author', 'to', 'paper'): 0,
}
node_types, edge_types = data.metadata()
assert node_types == ['paper', 'author']
assert edge_types == [
('paper', 'to', 'paper'),
('paper', 'to', 'author'),
('author', 'to', 'paper'),
]
x_dict = data.collect('x')
assert len(x_dict) == 2
assert x_dict['paper'].tolist() == x_paper.tolist()
assert x_dict['author'].tolist() == x_author.tolist()
assert x_dict == data.x_dict
data.y = 0
assert data['y'] == 0 and data.y == 0
assert len(data) == 4
assert sorted(data.keys) == ['edge_attr', 'edge_index', 'x', 'y']
del data['paper', 'author']
node_types, edge_types = data.metadata()
assert node_types == ['paper', 'author']
assert edge_types == [('paper', 'to', 'paper'), ('author', 'to', 'paper')]
assert len(data.to_dict()) == 5
assert len(data.to_namedtuple()) == 5
assert data.to_namedtuple().y == 0
assert len(data.to_namedtuple().paper) == 1
def test_heterogeneous_neighbor_loader_on_cora(directed):
root = osp.join('/', 'tmp', str(random.randrange(sys.maxsize)))
dataset = Planetoid(root, 'Cora')
data = dataset[0]
data.edge_weight = torch.rand(data.num_edges)
hetero_data = HeteroData()
hetero_data['paper'].x = data.x
hetero_data['paper'].n_id = torch.arange(data.num_nodes)
hetero_data['paper', 'paper'].edge_index = data.edge_index
hetero_data['paper', 'paper'].edge_weight = data.edge_weight
split_idx = torch.arange(5, 8)
loader = NeighborLoader(hetero_data,
num_neighbors=[-1, -1],
batch_size=split_idx.numel(),
input_nodes=('paper', split_idx),
directed=directed)
assert len(loader) == 1
hetero_batch = next(iter(loader))
batch_size = hetero_batch['paper'].batch_size
if not directed:
n_id, _, _, e_mask = k_hop_subgraph(split_idx,
num_hops=2,
edge_index=data.edge_index,
num_nodes=data.num_nodes)
n_id = n_id.sort()[0]
assert n_id.tolist() == hetero_batch['paper'].n_id.sort()[0].tolist()
assert hetero_batch['paper', 'paper'].num_edges == int(e_mask.sum())
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).relu()
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, data.edge_weight)[split_idx]
out2 = hetero_model(hetero_batch.x_dict, hetero_batch.edge_index_dict,
hetero_batch.edge_weight_dict)['paper'][:batch_size]
assert torch.allclose(out1, out2, atol=1e-6)
try:
shutil.rmtree(root)
except PermissionError:
pass
def test_hgt_loader_on_cora(get_dataset):
dataset = get_dataset(name='Cora')
data = dataset[0]
data.edge_weight = torch.rand(data.num_edges)
hetero_data = HeteroData()
hetero_data['paper'].x = data.x
hetero_data['paper'].n_id = torch.arange(data.num_nodes)
hetero_data['paper', 'paper'].edge_index = data.edge_index
hetero_data['paper', 'paper'].edge_weight = data.edge_weight
split_idx = torch.arange(5, 8)
# Sample the complete two-hop neighborhood:
loader = HGTLoader(hetero_data,
num_samples=[data.num_nodes] * 2,
batch_size=split_idx.numel(),
input_nodes=('paper', split_idx))
assert len(loader) == 1
hetero_batch = next(iter(loader))
batch_size = hetero_batch['paper'].batch_size
n_id, _, _, e_mask = k_hop_subgraph(split_idx,
num_hops=2,
edge_index=data.edge_index,
num_nodes=data.num_nodes)
n_id = n_id.sort()[0]
assert n_id.tolist() == hetero_batch['paper'].n_id.sort()[0].tolist()
assert hetero_batch['paper', 'paper'].num_edges == int(e_mask.sum())
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).relu()
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, data.edge_weight)[split_idx]
out2 = hetero_model(hetero_batch.x_dict, hetero_batch.edge_index_dict,
hetero_batch.edge_weight_dict)['paper'][:batch_size]
assert torch.allclose(out1, out2, atol=1e-6)
def test_hgt_conv_out_of_place():
data = HeteroData()
data['author'].x = torch.randn(4, 16)
data['paper'].x = torch.randn(6, 32)
index1 = torch.randint(0, 4, (20, ), dtype=torch.long)
index2 = torch.randint(0, 6, (20, ), dtype=torch.long)
data['author', 'paper'].edge_index = torch.stack([index1, index2], dim=0)
data['paper', 'author'].edge_index = torch.stack([index2, index1], dim=0)
conv = HGTConv(-1, 64, data.metadata(), heads=1)
x_dict, edge_index_dict = data.x_dict, data.edge_index_dict
assert x_dict['author'].size() == (4, 16)
assert x_dict['paper'].size() == (6, 32)
_ = conv(x_dict, edge_index_dict)
assert x_dict['author'].size() == (4, 16)
assert x_dict['paper'].size() == (6, 32)
def create_hetero_mock_data(n_count, feature_dict):
_x_dict = {
'author':
torch.FloatTensor(
np.random.uniform(0, 1, (n_count, feature_dict['author']))),
'paper':
torch.FloatTensor(
np.random.uniform(0, 1, (n_count, feature_dict['paper'])))
}
_edge_index_dict = {
('author', 'writes', 'paper'):
torch.LongTensor(get_edge_array(n_count))
}
data = HeteroData()
data['author'].x = _x_dict['author']
data['paper'].x = _x_dict['paper']
data[('author', 'writes',
'paper')].edge_index = _edge_index_dict[('author', 'writes',
'paper')]
data = T.ToUndirected()(data)
return data.x_dict, data.edge_index_dict, data.metadata()
