Computing Cosine Distance with Differently shaped tensors

user_4160007 0
Viewed 2275

I have the following tensor representing a word vector

A = (2, 500)

Where the first dimension is the BATCH dimension (i.e. A contains two word vectors each with 500 elements)

I also have the following tensor

B = (10, 500)

I want to compute the cosine distance between A and B such that I get

C = (2, 10, 1)

i.e for each row in A compute the cosine distance with each row in B

I looked at using torch.nn.functional.F.cosine_similarity however this doesn't work as the dimensions must be the same.

Whats the best efficient way of achieving this in pytorch?

cosine-similarity pytorch
2 Answers

The accepted solution appears to be inefficient --it took forever on my machine and finally crushed the kernel due to insufficient memory-- while this solution took a few milliseconds:

import torch.nn.functional as F

# cosine similarity = normalize the vectors & multiply
C = F.normalize(A) @ F.normalize(B).t()

This is the implementation in sentence-transformers

user_10541540

Use broadcasting technique with unsqueeze

import torch.nn.functional as F

C = F.cosine_similarity(A.unsqueeze(1), B, dim=-1)

print(C.shape)
# torch.size([2,10])
user_42384080
Related