Pytorch cosine similarity NxN elements

user_15723356 0
Viewed 1312

I have 128 vectors of embeddings

image.shape = torch.Size([128, 512])
text.shape = torch.Size([128, 512])

And I want to calculate the tensor containing the cosine similarity between all elements (i.e:

cosine.shape = torch.Size([128, 128])

Where the first row is the cosine similarity between the 1st image and all text (128), etc.

At the moment I'm only doing this, but the result is a one-dimension array containing only N cosine similarities.

cosine_similarity = torch.nn.CosineSimilarity()

cosine = cosine_similarity(image, text)

How can I do it? I tried to transpose text but didn't work

pytorch python
2 Answers

The way pytorch computes cosine similarity internally is like this:

    def cos_sim(A, B, dim, eps=1e-08):
      numerator = torch.mul(A, B).sum(axis=dim, keepdims=True)
      A_l2 = torch.mul(A, A).sum(axis=dim, keepdims=True)
      B_l2 = torch.mul(B, B).sum(axis=dim, keepdims=True)
      denominator = torch.max(torch.sqrt(torch.mul(A_l2, B_l2)) torch.tensor(eps))
      return torch.div(numerator, denominator).squeeze()

In order to get NxN cosine similarity, you can instead use this function:

    def nxn_cos_sim(A, B, dim=1, eps=1e-8):
      numerator = A @ B.T
      A_l2 = torch.mul(A, A).sum(axis=dim)
      B_l2 = torch.mul(B, B).sum(axis=dim)
      denominator = torch.max(torch.sqrt(torch.outer(A_l2, B_l2)), torch.tensor(eps))
      return torch.div(numerator, denominator)

Ahmed's answer also works, but if the value for N is large you will face memory issues using it.

user_191794280

This should do the trick:

repeated_image = image.repeat_interleave(128, dim=0)
repeated_text = text.repeat(128, 1)

distances = torch.nn.functional.cosine_similarity(repeated_image, repeated_text, dim=1)
distances = distances.reshape(128, 128)

Explanation: In order to do the computations faster, it is better to copy the data and then use the GPU to do the steps in parallel, rather than for-loop over all possible pairs. The first two lines are copying the data in an appropriate way:

repeat_interleave will copy each element of the tensor (which is a 512-dimensional image) along dim=0, which will create a (128*128, 512) tensor whose first 128 elements belong to the first image, and so on.

repeat will repeat the entire tensor 128 times along the 0-th dimension, again creating a (128*128, 512) tensor, where the first element belongs to the first text, the second element belongs to the second text, and so on (the 129-th element belongs to the first text.)

We then compute all pairwise distances using the cosine_similarity function, and then reshape to get them into the appropriate shape.

user_59547410
Related