Calculate Cosine Similarity Between Tensors in PyTorch

In this tutorial, we will introduce you how to calculate cosine similarity between two tensors in pytorch. We will list two methods.

Method 1: Create a function to compute

For example:

import torch
import torch.nn.functional as F
def cosine_similarity(x, y):
    cosine = torch.sum(F.normalize(x, dim = 1)* F.normalize(y, dim = 1), dim = 1)
    return cosine

Then, we will get cosine similarity between x and y.

In order to understand how to use F.normalize(), you can read:

Understand torch.nn.functional.normalize() with Examples – PyTorch Tutorial

x = torch.randn(5,200)
y = torch.randn(5,200)
cosine = cosine_similarity(x, y)
print(cosine)

Run this code, we will see:

tensor([ 0.0136, -0.1019,  0.0674,  0.0229, -0.0716])

Method 2: use torch.nn.CosineSimilarity()

It is defined as:

torch.nn.CosineSimilarity(dim=1, eps=1e-08)

We can use it as follows:

cosine = torch.nn.CosineSimilarity(dim=1)
print(cosine(x, y))

Run it, we also can get:

tensor([ 0.0136, -0.1019,  0.0674,  0.0229, -0.0716])