How to compute the Cosine Similarity between two tensors in PyTorch?

To compute the cosine similarity between two tensors, we use the CosineSimilarity() function provided by the torch.nn module. It returns the cosine similarity value computed along dim.

dim is an optional parameter to this function along which cosine similarity is computed.

• For 1D tensors, we can compute the cosine similarity along dim=0 only.

• For 2D tensors, we can compute cosine similarity along dim=0 or 1.

• The size of both tensors must be the same to compute the cosine similarity. Both tensors must be real-valued. Cosine similarity is often used to measure document similarity in text analysis.

Syntax

torch.nn.CosineSimilarity(dim=1)

The default dim is set to 1. But if you measure the cosine similarity between 1D tensors, then we set dim to 0.

Steps

• Import the required library. In all the following examples, the required Python library is torch. Make sure you have already installed it.

import torch

• Create two tensors and print them. Both tensors must be real-valued.

tensor1 = torch.randn(3,4)
tensor2 = torch.randn(3,4)

• Define a method to measure cosine similarity along dimension dim.

cos = torch.nn.CosineSimilarity(dim=0)

• Compute the Cosine Similarity using the above defined method.

output = cos(tensor1, tensor2)

• Print the computed tensor with cosine similarity values.

print("Cosine Similarity:",output)

Example 1

The following Python program computes the Cosine Similarity between two 1D tensors.

# Import the required library
import torch

# define two input tensors
tensor1 = torch.tensor([0.1, 0.3, 2.3, 0.45])
tensor2 = torch.tensor([0.13, 0.23, 2.33, 0.45])

# print above defined two tensors
print("Tensor 1:
", tensor1)
print("Tensor 2:
", tensor2)

# define a method to measure cosine similarity
cos = torch.nn.CosineSimilarity(dim=0)
output = cos(tensor1, tensor2)

# display the output tensor
print("Cosine Similarity:",output)

Output

Tensor 1:
   tensor([0.1000, 0.3000, 2.3000, 0.4500])
Tensor 2:
   tensor([0.1300, 0.2300, 2.3300, 0.4500])
Cosine Similarity: tensor(0.9995)

Example 2

In this Python program, we compute the Cosine Similarity between two 2D tensors along different dim.

# Import the required library
import torch

# define two input tensors
tensor1 = torch.randn(3,4)
tensor2 = torch.randn(3,4)

# print above defined two tensors
print("Tensor 1:
", tensor1)
print("Tensor 2:
", tensor2)

# define a method to measure cosine similarity in dim 0
cos0 = torch.nn.CosineSimilarity(dim=0)
output0 = cos0(tensor1, tensor2)
print("Cosine Similarity in dim 0:
",output0)

# define a method to measure cosine similarity in dim 1
cos1 = torch.nn.CosineSimilarity(dim=1)
output1 = cos1(tensor1, tensor2)
print("Cosine Similarity in dim 1:
",output1)

Output

Tensor 1:
   tensor([[ 0.2714, 1.1430, 1.3997, 0.8788],
      [-2.2268, 1.9799, 1.5682, 0.5850],
      [ 1.2289, 0.5043, -0.1625, 1.1403]])
Tensor 2:
   tensor([[-0.3299, 0.6360, -0.2014, 0.5989],
      [-0.6679, 0.0793, -2.5842, -1.5123],
      [ 1.1110, -0.1212, 0.0324, 1.1277]])
Cosine Similarity in dim 0:
   tensor([ 0.8076, 0.5388, -0.7941, 0.3016])
Cosine Similarity in dim 1:
   tensor([ 0.4553, -0.3140, 0.9258])