Batch inner product decoder for VGAE

Hi,

Do you guys have any idea on how to implement an inner product decoder for architectures such as VGAE (Eq. 2) using mini-batches in torch_geometric?

Thank you!

[dskoda]

Thank you for your reply, @rusty1s !

After a few changes in your code, it was possible to make it. Here’s what worked for me:

row = torch.arange(0, data.num_nodes, dtype=torch.long)
col = data.batch
index = torch.stack([row, col], dim=0)
value = Z.new_ones(data.num_nodes)
size = torch.Size([data.num_nodes, data.batch[-1].item() + 1]) # added one to the index of last batch
assignment = torch.sparse_coo_tensor(index, value, size).to_dense() # converted sparse matrix to dense

mask = torch.matmul(assignment, assignment.t())

out = torch.matmul(Z, Z.t())
out = out * mask # masking with element-wise multiplication

I couldn’t make the line torch.matmul(assignment, assignment.t()) work for sparse matrices. It returned the following error: RuntimeError: Sparse tensors do not have strides.

I suppose this way is not the optimal, since performing dense matrix operations instead of sparse ones will take more time.

[rusty1s]

If I‘m not mistaken, masking could be achieved by:

row = torch.arange(0, data.num_nodes, dtype=torch.long)
col = data.batch
index = torch.stack([row, col], dim=0)
value = Z.new_ones(data.num_nodes)
size = torch.Size([data.num_nodes, data.batch[-1]])
assignment = torch.sparse_coo_tensor(index, value, size).to_dense()

mask = torch.matmul(assignment, assignment.t())

out = torch.matmul(Z, Z.t())
out = out * mask

You can then train against your densified input graph. However, this involves computing and comparing two N x N dense matrices, but this is more a problem of VGAE than with the proposed approach.