def test_gin_conv():
x1 = torch.randn(4, 16)
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))
nn = Seq(Lin(16, 32), ReLU(), Lin(32, 32))
conv = GINConv(nn, train_eps=True)
assert conv.__repr__() == (
'GINConv(nn=Sequential(\n'
' (0): Linear(in_features=16, out_features=32, bias=True)\n'
' (1): ReLU()\n'
' (2): Linear(in_features=32, out_features=32, bias=True)\n'
'))')
out = conv(x1, edge_index)
assert out.size() == (4, 32)
assert conv(x1, edge_index, size=(4, 4)).tolist() == out.tolist()
assert conv(x1, adj.t()).tolist() == out.tolist()
if is_full_test():
t = '(Tensor, Tensor, Size) -> 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, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit(x1, adj.t()).tolist() == out.tolist()
adj = adj.sparse_resize((4, 2))
out1 = conv((x1, x2), edge_index)
out2 = conv((x1, None), edge_index, (4, 2))
assert out1.size() == (2, 32)
assert out2.size() == (2, 32)
assert conv((x1, x2), edge_index, (4, 2)).tolist() == out1.tolist()
assert conv((x1, x2), adj.t()).tolist() == out1.tolist()
assert conv((x1, None), adj.t()).tolist() == out2.tolist()
if is_full_test():
t = '(OptPairTensor, Tensor, Size) -> 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, Size) -> Tensor'
jit = torch.jit.script(conv.jittable(t))
assert jit((x1, x2), adj.t()).tolist() == out1.tolist()
assert jit((x1, None), adj.t()).tolist() == out2.tolist()
def __init__(self, in_channels, hidden_channels, out_channels, num_layers):
super(Net, self).__init__()
self.convs = torch.nn.ModuleList()
self.batch_norms = torch.nn.ModuleList()
for i in range(num_layers):
mlp = Sequential(
Linear(in_channels, 2 * hidden_channels),
BatchNorm(2 * hidden_channels),
ReLU(),
Linear(2 * hidden_channels, hidden_channels),
)
conv = GINConv(mlp, train_eps=True)
conv = conv.jittable(x=torch.randn(data.num_nodes, in_channels),
edge_index=data.edge_index)
self.convs.append(conv)
self.batch_norms.append(BatchNorm(hidden_channels))
in_channels = hidden_channels
self.lin1 = Linear(hidden_channels, hidden_channels)
self.batch_norm1 = BatchNorm(hidden_channels)
self.lin2 = Linear(hidden_channels, out_channels)
def test_gin_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))
nn = Seq(Lin(in_channels, 32), ReLU(), Lin(32, out_channels))
conv = GINConv(nn, train_eps=True)
assert conv.__repr__() == (
'GINConv(nn=Sequential(\n'
' (0): Linear(in_features=16, out_features=32, bias=True)\n'
' (1): ReLU()\n'
' (2): Linear(in_features=32, out_features=32, bias=True)\n'
'))')
out = conv(x, edge_index)
assert out.size() == (num_nodes, out_channels)
jit_conv = conv.jittable(x=x, edge_index=edge_index)
jit_conv = torch.jit.script(jit_conv)
assert jit_conv(x, edge_index).tolist() == out.tolist()
