CODEWITHSUNDEEP
pytorchgpupytorch-geometricgraph-neural-network
Adrien Lagesse
Adrien Lagesse

Reputation: 21

Batch computation of similarity matrices of different sizes (Pytorch)

I am computing node features of dimension D for a B different graphs where the graph i has N_i nodes, hence I have a batch representation as a tensor of dimension (N_1 + ... + N_B)xD. I want to compute the similarity of each node to each other within the same graph, if X_i is the node features of graph i (of dimension N_i x D), I need to compute X_i@X_i.T (dimension N_i x D) for each i.

Afterwards, I will compute a the cross entropy loss between X_i@X_i.T and Id_{N_i}.

I need to have a fast way to compute this loss for a batch (batch size of 100) in parallel and on the GPU (with Pytorch and Pytorch Geometric). I already tried to reduce the batch size to one and to pad the tensors so that the batch can be reshaped into B x N_max x F, however, the batch size of 1 is way to slow and padding the tensors doesn't work because torch_geometric expect contiguous tensors for the message passing part to compute the node features.

Upvotes: 0

Views: 101

Answers (1)

MinhNH
MinhNH

Reputation: 564

Consider taking the similarity of the whole batch and masked out the similarity between nodes of different graphs. The easiest way to construct such mask is to first create a vector and then do a equal operator like this:

batch_vector = torch.cat([[1]*N_1, [2]*N_2, ..., [B]*N_B], dim = 0) # shape (N_1 + ... + N_B,)

mask = batch_vector.unsqueeze(0) == batch_vector.unsqueeze(1) # shape (N_1 + ... + N_B, N_1 + ... + N_B)

This would not be memory efficient though, so that's a trade-off.

Upvotes: 0

Related Questions