Hello, I’ve been trying to implement the Mahalanobis distance between multiple nodes. So far I’ve started with an input of size (batch, time_stamps, num_nodes, embeding_size) and I will like to have an output of size (batch_size, time_stamps, num_nodes, num_nodes). The computation is quite simple for each pair of nodes (x_i, x_j, where the batch and the time_stamp matches) I need to compute: (x_i - x_j)^T * M * (x_j - x_i). I managed to implement by looping, but it was too time consuming, I also implemented it by using torch.repeat and some reshapes but i’m running out of memory.
So far I have this:
sub1 = x_perm.repeat([1, 1, num_nodes, 1]).reshape(batch_size, ts, num_nodes, num_nodes,
num_channels).transpose(
3, 2) - \
x_perm.repeat([1, 1, num_nodes, 1]).reshape(batch_size, ts, num_nodes,
num_nodes, num_channels)
sub2 = x_perm.repeat([1, 1, num_nodes, 1]).reshape(batch_size, ts, num_nodes,
num_nodes, num_channels) - \
x_perm.repeat([1, 1, num_nodes, 1]).reshape(batch_size, ts, num_nodes,
num_nodes, num_channels).transpose(
3, 2)
sub1 = sub1.unsqueeze(4)
sub2 = sub2.unsqueeze(4)
score = torch.matmul(sub1, M])
score = torch.matmul(score, sub2.transpose(5, 4))
score = score.reshape(batch_size, ts, num_nodes, num_nodes)
Any suggestions are more than welcome.