🎓 A1Learner
🔍

PyTorch Sorting and Topk

Sorting puts a tensor's values in order, smallest first. topk returns only the k largest values and their positions.

Sort values and positions

sort returns two tensors. The first holds the values in order and the second holds their old positions.

import torch

scores = torch.tensor([72, 95, 61, 88])
sorted_scores, positions = torch.sort(scores)
print(sorted_scores)
print(positions)

Get the order with argsort

argsort returns the same positions on their own. Reading the scores in that order gives the sorted list.

scores = torch.tensor([72, 95, 61, 88])
print(scores.argsort())

Take the top two

topk skips the full sort. It returns the largest values first, with their positions.

scores = torch.tensor([72, 95, 61, 88])
values, indices = scores.topk(2)
print(values)
print(indices)

Example

import torch

scores = torch.tensor([72, 95, 61, 88, 79])
top_scores, top_students = torch.topk(scores, 3)
print("Top 3 scores:")
print(top_scores)
print("Student indices:")
print(top_students)

Key points

  • Use torch.sort(tensor) to get the sorted values plus each value's original position
  • Use argsort to get the positions in sorted order, without the values
  • k: how many of the largest values topk should return
  • Use topk(k) to get the k largest values and their positions
  • In short, sort orders everything and topk goes straight to the leaders
  • Use descending=True with sort to run from the largest value down

In this example, scores holds one score for each of five students. torch.topk(scores, 3) asks for the three largest scores. The values print largest first, 95, then 88, then 79. The indices print the student each score came from. Student 1 scored 95, student 3 scored 88, and student 4 scored 79.

Try it Yourself

Edit the code, then press Run.

Running is simulated · Python Playground

Exercise

What does scores.topk(3) return?