def test_gat_conv():
in_channels, out_channels = (16, 32)
edge_index = torch.tensor([[0, 0, 0, 1, 2, 3], [1, 2, 3, 0, 0, 0]])
num_nodes = edge_index.max().item() + 1
x = torch.randn((num_nodes, in_channels))
conv = GATConv(in_channels, out_channels, heads=2, dropout=0.0)
assert conv.__repr__() == 'GATConv(16, 32, heads=2)'
out = conv(x, edge_index)
assert out.size() == (num_nodes, 2 * out_channels)
assert conv((x, x), edge_index).tolist() == out.tolist()
model = Net1()
model.conv = conv.jittable(x=x, edge_index=edge_index)
model = torch.jit.script(model)
assert model(x, edge_index).tolist() == out.tolist()
result = conv(x, edge_index, return_attention_weights=True)
assert result[0].tolist() == out.tolist()
assert result[1][0][:, :-num_nodes].tolist() == edge_index.tolist()
assert result[1][1].size() == (edge_index.size(1) + num_nodes, 2)
assert conv._alpha is None
model = Net2()
model.conv = conv.jittable(x=x, edge_index=edge_index)
model = torch.jit.script(model)
assert model(x, edge_index)[0].tolist() == result[0].tolist()
assert model(x, edge_index)[1][0].tolist() == result[1][0].tolist()
assert model(x, edge_index)[1][1].tolist() == result[1][1].tolist()
conv = GATConv(in_channels, out_channels, heads=2, concat=False)
out = conv(x, edge_index)
assert out.size() == (num_nodes, out_channels)
assert conv((x, x), edge_index).tolist() == out.tolist()
model = Net1()
model.conv = conv.jittable(x=x, edge_index=edge_index)
model = torch.jit.script(model)
assert model(x, edge_index).tolist() == out.tolist()
result = conv(x, edge_index, return_attention_weights=True)
assert result[0].tolist() == out.tolist()
assert result[1][0][:, :-num_nodes].tolist() == edge_index.tolist()
assert result[1][1].size() == (edge_index.size(1) + num_nodes, 2)
assert conv._alpha is None
model = Net2()
model.conv = conv.jittable(x=x, edge_index=edge_index)
model = torch.jit.script(model)
assert model(x, edge_index)[0].tolist() == result[0].tolist()
assert model(x, edge_index)[1][0].tolist() == result[1][0].tolist()
assert model(x, edge_index)[1][1].tolist() == result[1][1].tolist()
class Net(torch.nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = GATConv(dataset.num_features, 8, heads=8, dropout=0.6)
self.conv1 = self.conv1.jittable(x=data.x, edge_index=data.edge_index)
self.conv2 = GATConv(64,
dataset.num_classes,
heads=1,
concat=True,
dropout=0.6)
self.conv2 = self.conv2.jittable(x=torch.randn(data.num_nodes, 64),
edge_index=data.edge_index)
def forward(self, x, edge_index):
x = F.dropout(x, p=0.6, training=self.training)
x = F.elu(self.conv1(x, edge_index))
x = F.dropout(x, p=0.6, training=self.training)
x = self.conv2(x, edge_index)
return F.log_softmax(x, dim=1)
def test_gat_conv():
x1 = torch.randn(4, 8)
x2 = torch.randn(2, 16)
edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
row, col = edge_index
adj = SparseTensor(row=row, col=col, sparse_sizes=(4, 4))
conv = GATConv(8, 32, heads=2)
assert conv.__repr__() == 'GATConv(8, 32, heads=2)'
out = conv(x1, edge_index)
assert out.size() == (4, 64)
assert conv(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
assert torch.allclose(conv(x1, adj.t()), out, atol=1e-6)
t = '(Tensor, Tensor, OptTensor, Size, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, edge_index).tolist() == out.tolist()
assert jit(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
t = '(Tensor, SparseTensor, OptTensor, Size, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit(x1, adj.t()), out, atol=1e-6)
# Test `return_attention_weights`.
result = conv(x1, edge_index, return_attention_weights=True)
assert result[0].tolist() == out.tolist()
assert result[1][0].size() == (2, 7)
assert result[1][1].size() == (7, 2)
assert result[1][1].min() >= 0 and result[1][1].max() <= 1
assert conv._alpha is None
result = conv(x1, adj.t(), return_attention_weights=True)
assert torch.allclose(result[0], out, atol=1e-6)
assert result[1].sizes() == [4, 4, 2] and result[1].nnz() == 7
assert conv._alpha is None
t = ('(Tensor, Tensor, OptTensor, Size, bool) -> '
'Tuple[Tensor, Tuple[Tensor, Tensor]]')
jit = torch.jit.script(conv.jittable(t))
result = jit(x1, edge_index, return_attention_weights=True)
assert result[0].tolist() == out.tolist()
assert result[1][0].size() == (2, 7)
assert result[1][1].size() == (7, 2)
assert result[1][1].min() >= 0 and result[1][1].max() <= 1
assert conv._alpha is None
t = ('(Tensor, SparseTensor, OptTensor, Size, bool) -> '
'Tuple[Tensor, SparseTensor]')
jit = torch.jit.script(conv.jittable(t))
result = jit(x1, adj.t(), return_attention_weights=True)
assert torch.allclose(result[0], out, atol=1e-6)
assert result[1].sizes() == [4, 4, 2] and result[1].nnz() == 7
assert conv._alpha is None
adj = adj.sparse_resize((4, 2))
conv = GATConv((8, 16), 32, heads=2)
assert conv.__repr__() == 'GATConv((8, 16), 32, heads=2)'
out1 = conv((x1, x2), edge_index)
out2 = conv((x1, None), edge_index, size=(4, 2))
assert out1.size() == (2, 64)
assert out2.size() == (2, 64)
assert conv((x1, x2), edge_index, size=(4, 2)).tolist() == out1.tolist()
assert torch.allclose(conv((x1, x2), adj.t()), out1, atol=1e-6)
assert torch.allclose(conv((x1, None), adj.t()), out2, atol=1e-6)
t = '(OptPairTensor, Tensor, OptTensor, Size, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), edge_index).tolist() == out1.tolist()
assert jit((x1, x2), edge_index, size=(4, 2)).tolist() == out1.tolist()
assert jit((x1, None), edge_index, size=(4, 2)).tolist() == out2.tolist()
t = '(OptPairTensor, SparseTensor, OptTensor, Size, NoneType) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert torch.allclose(jit((x1, x2), adj.t()), out1, atol=1e-6)
assert torch.allclose(jit((x1, None), adj.t()), out2, atol=1e-6)
