How to Calculate Cosine Similarity Between Two Tensors in PyTorch?

by | Nov 22, 2023 | How-To/Tips

Cosine similarity measures how similar two vectors are by computing the cosine of the angle between them — returning a value between -1 (opposite) and 1 (identical direction). PyTorch provides torch.nn.CosineSimilarity() in the torch.nn module and torch.nn.functional.cosine_similarity() for a functional API approach. Both work on 1D and 2D tensors along a specified dimension.

Quick Answer: To calculate cosine similarity in PyTorch, use torch.nn.CosineSimilarity(dim=0)(tensor1, tensor2) for 1D vectors, or torch.nn.functional.cosine_similarity(t1, t2, dim=1) for row-wise comparison of 2D tensors.

CosineSimilarity: Class vs Functional API

PyTorch offers two ways to compute cosine similarity:

API Usage Best For
torch.nn.CosineSimilarity(dim) Create an object, then call it with two tensors Reusing the same similarity function in a model or loop
F.cosine_similarity(t1, t2, dim) Direct function call — no object creation needed One-off computations, cleaner scripts

Both return the same result. The functional form (import torch.nn.functional as F) is often preferred for its brevity.

Example 1: Calculate Cosine Similarity Between 1D Tensors

For 1D tensors, use dim=0 — similarity is computed across all elements.

Step 1: Import PyTorch

import torch
Importing the torch library in Python for cosine similarity calculation

Step 2: Define Two 1D Tensors

Create two 1D tensors of equal size with real values:

Tens1 = torch.tensor([0.7, 0.1, 1.4, 0.9])
Tens2 = torch.tensor([1.6, 0.2, 1.3, 1.8])
Defining two 1D PyTorch tensors Tens1 and Tens2 for cosine similarity

Note: Both tensors must have the same size and contain real values.

Step 3: Print the Tensor Elements

print(Tens1)
print(Tens2)
Printing elements of 1D tensors Tens1 and Tens2

Step 4: Compute Cosine Similarity

Use CosineSimilarity(dim=0) to compute similarity element-wise across a 1D tensor:

cos_sim = torch.nn.CosineSimilarity(dim=0)
Creating a CosineSimilarity object with dim=0 for 1D tensor comparison
output = cos_sim(Tens1, Tens2)
Calculating cosine similarity between two 1D PyTorch tensors
print("Cosine Similarity:", output)
Output showing the cosine similarity value between two 1D PyTorch tensors

Example 2: Calculate Cosine Similarity Between 2D Tensors

For 2D tensors, you can compute cosine similarity column-wise (dim=0) or row-wise (dim=1).

Step 1: Import PyTorch

import torch
Importing the torch library for 2D tensor cosine similarity example

Step 2: Define Two 2D Tensors

T1 = torch.randn(3, 4)
T2 = torch.randn(3, 4)
Defining two 2D PyTorch tensors T1 and T2 with random values using torch.randn

Step 3: Print the Tensor Elements

print(T1)
print(T2)
Printing elements of 2D tensors T1 and T2 showing their random values

Step 4: Compute Cosine Similarity Along dim=0 (Column-wise)

cos_sim0 = torch.nn.CosineSimilarity(dim=0)
output = cos_sim0(T1, T2)
print("Cosine Similarity in dim 0:", output)
Cosine similarity between 2D tensors computed column-wise along dim=0

Computing along dim=1 compares tensors row-wise instead:

cos_sim1 = torch.nn.CosineSimilarity(dim=1)
output = cos_sim1(T1, T2)
print("Cosine Similarity in dim 1:", output)
Cosine similarity between 2D tensors computed row-wise along dim=1

Note: For the complete working code, access the Google Colab Notebook.

Conclusion

To calculate cosine similarity in PyTorch, use torch.nn.CosineSimilarity(dim=0) for 1D tensors and specify dim=0 (column-wise) or dim=1 (row-wise) for 2D tensors. Both tensors must have the same size and contain real values. For a more concise alternative, use the functional API: torch.nn.functional.cosine_similarity(t1, t2, dim=1) — it produces the same result without creating a separate object.