How can I avoid this for loop in pytorch? Is there a function for efficient computation?

user_15705006 0
Viewed 89

I have the following code in my Pytorch neural net:

cos = nn.CosineSimilarity(dim=1)

d = torch.zeros(batch_sz, n, n).to(device="cuda")
    
for i in range(n):
   for j in range(n):
       d[:, i, j] = cos(q[:, i, :], k[:, j, :])

q and k are both of size (batch_sz, n, m). This piece of Code obviously slows down my program and I'm wondering if Pytorch offers any functions that might make this more efficient.

Thanks so much!

pytorch performance
1 Answers

I am not sure how to vectorize using nn.CosineSimilarity but you could use this vectorized implementation. It computes the cosine similarity in the same way as PyTorch's internal module.

import torch
import torch.nn as nn
import time

# some dummy inputs
n=20
m=30
batch_sz = 10

k = torch.rand(batch_sz, n, m)
q = torch.rand(batch_sz, n, m)
d = torch.zeros(batch_sz, n, n)


cos = nn.CosineSimilarity(dim=1)

for i in range(n):
    for j in range(n):
        d[:, i, j] = cos(q[:, i, :], k[:, j, :])



# dot product (numerator)
out = torch.bmm(q, k.transpose(1,2))

# computing the denominator in the next 5 steps

# compute the norm and restore dimensions
q_norm = q.norm(dim=2).unsqueeze(2)
k_norm = k.norm(dim=2).unsqueeze(1)

# This repeats the norms along dim 2 for q and dim 1 for k
q_norm_expanded = q_norm.expand(batch_sz, n, n)
k_norm_expanded = k_norm.expand(batch_sz, n, n)

# we compute the product. 
norms = q_norm_expanded* k_norm_expanded

# cosine similarity
out = out/(norms+1e-9)

print(torch.allclose(d, out))

The process of expanding and multiplying the norms is actually computing the outer product. So you could use also use this operation:

norms = torch.bmm(q_norm, k_norm)

instead of

q_norm_expanded = q_norm.expand(batch_sz, n, n)
k_norm_expanded = k_norm.expand(batch_sz, n, n)

norms = q_norm_expanded* k_norm_expanded

I just realized you could normalize the vectors before hand for a More concise and computationally stable version

q_norm = q.norm(dim=2)+1e-9
k_norm = k.norm(dim=2)+1e-9
q = q/q_norm.unsqueeze(2)
k = k/k_norm.unsqueeze(2)

out = torch.bmm(q, k.transpose(1,2))
user_32369250
Related