def test_to_hetero_with_bases_and_rgcn_equal_output():
torch.manual_seed(1234)
# Run `RGCN` with basis decomposition:
x = torch.randn(10, 16) # 6 paper nodes, 4 author nodes
adj = (torch.rand(10, 10) > 0.5)
adj[6:, 6:] = False
edge_index = adj.nonzero(as_tuple=False).t().contiguous()
row, col = edge_index
# # 0 = paper<->paper, 1 = author->paper, 2 = paper->author
edge_type = torch.full((edge_index.size(1), ), -1, dtype=torch.long)
edge_type[(row < 6) & (col < 6)] = 0
edge_type[(row < 6) & (col >= 6)] = 1
edge_type[(row >= 6) & (col < 6)] = 2
assert edge_type.min() == 0
num_bases = 4
conv = RGCNConv(16, 32, num_relations=3, num_bases=num_bases, aggr='add')
out1 = conv(x, edge_index, edge_type)
# Run `to_hetero_with_bases`:
x_dict = {
'paper': x[:6],
'author': x[6:],
}
edge_index_dict = {
('paper', '_', 'paper'):
edge_index[:, edge_type == 0],
('paper', '_', 'author'):
edge_index[:, edge_type == 1] - torch.tensor([[0], [6]]),
('author', '_', 'paper'):
edge_index[:, edge_type == 2] - torch.tensor([[6], [0]]),
}
adj_t_dict = {
key: SparseTensor.from_edge_index(edge_index).t()
for key, edge_index in edge_index_dict.items()
}
metadata = (list(x_dict.keys()), list(edge_index_dict.keys()))
model = to_hetero_with_bases(RGCN(16, 32), metadata, num_bases=num_bases,
debug=False)
# Set model weights:
for i in range(num_bases):
model.conv.convs[i].lin.weight.data = conv.weight[i].data.t()
model.conv.convs[i].edge_type_weight.data = conv.comp[:, i].data.t()
model.lin.weight.data = conv.root.data.t()
model.lin.bias.data = conv.bias.data
out2 = model(x_dict, edge_index_dict)
out2 = torch.cat([out2['paper'], out2['author']], dim=0)
assert torch.allclose(out1, out2, atol=1e-6)
out3 = model(x_dict, adj_t_dict)
out3 = torch.cat([out3['paper'], out3['author']], dim=0)
assert torch.allclose(out1, out3, atol=1e-6)
def test_to_hetero_with_bases():
metadata = (['paper', 'author'], [('paper', 'cites', 'paper'),
('paper', 'written_by', 'author'),
('author', 'writes', 'paper')])
x_dict = {'paper': torch.randn(100, 16), 'author': torch.randn(100, 8)}
edge_index_dict = {
('paper', 'cites', 'paper'):
torch.randint(100, (2, 200), dtype=torch.long),
('paper', 'written_by', 'author'):
torch.randint(100, (2, 200), dtype=torch.long),
('author', 'writes', 'paper'):
torch.randint(100, (2, 200), dtype=torch.long),
}
edge_attr_dict = {
('paper', 'cites', 'paper'): torch.randn(200, 8),
('paper', 'written_by', 'author'): torch.randn(200, 8),
('author', 'writes', 'paper'): torch.randn(200, 8),
}
model = Net1()
in_channels = {'x': 16, 'edge_attr': 8}
model = to_hetero_with_bases(model,
metadata,
num_bases=4,
in_channels=in_channels,
debug=False)
out = model(x_dict, edge_attr_dict)
assert isinstance(out, tuple) and len(out) == 2
assert isinstance(out[0], dict) and len(out[0]) == 2
assert out[0]['paper'].size() == (100, 32)
assert out[0]['author'].size() == (100, 32)
assert isinstance(out[1], dict) and len(out[1]) == 3
assert out[1][('paper', 'cites', 'paper')].size() == (200, 16)
assert out[1][('paper', 'written_by', 'author')].size() == (200, 16)
assert out[1][('author', 'writes', 'paper')].size() == (200, 16)
model = Net2()
in_channels = {'x': 16}
model = to_hetero_with_bases(model,
metadata,
num_bases=4,
in_channels=in_channels,
debug=False)
out = model(x_dict, edge_index_dict)
assert isinstance(out, dict) and len(out) == 2
assert out['paper'].size() == (100, 32)
assert out['author'].size() == (100, 32)
model = Net3()
in_channels = {'x': 16, 'edge_attr': 8}
model = to_hetero_with_bases(model,
metadata,
num_bases=4,
in_channels=in_channels,
debug=False)
out = model(x_dict, edge_index_dict, edge_attr_dict)
assert isinstance(out, dict) and len(out) == 2
assert out['paper'].size() == (100, 32)
assert out['author'].size() == (100, 32)
model = Net4()
in_channels = {'x': 16}
model = to_hetero_with_bases(model,
metadata,
num_bases=4,
in_channels=in_channels,
debug=False)
out = model(x_dict, edge_index_dict)
assert isinstance(out, dict) and len(out) == 2
assert out['paper'].size() == (100, 32)
assert out['author'].size() == (100, 32)
model = Net5(num_layers=2)
in_channels = {'x': 16}
model = to_hetero_with_bases(model,
metadata,
num_bases=4,
in_channels=in_channels,
debug=False)
out = model(x_dict, edge_index_dict)
assert isinstance(out, dict) and len(out) == 2
assert out['paper'].size() == (100, 16)
assert out['author'].size() == (100, 16)
model = Net6(num_layers=2)
in_channels = {'x': 16}
model = to_hetero_with_bases(model,
metadata,
num_bases=4,
in_channels=in_channels,
debug=False)
out = model(x_dict, edge_index_dict)
assert isinstance(out, dict) and len(out) == 2
assert out['paper'].size() == (100, 16)
assert out['author'].size() == (100, 16)
model = Net7()
in_channels = {'x': 16}
model = to_hetero_with_bases(model,
metadata,
num_bases=4,
in_channels=in_channels,
debug=False)
out = model(x_dict, edge_index_dict)
assert isinstance(out, dict) and len(out) == 2
assert out['paper'].size() == (100, 32)
assert out['author'].size() == (100, 32)