How to sort a dataset in pytorch

Question

I would like to sort my dataset by the numerical values in the labels.

Is there a function from pytorch to handle this efficiently?

my dataset type() is in this from:

 <class 'torchvision.datasets.mnist.MNIST'>

Answer 1

There is no generic way to do this efficiently, as the dataset class is only implements a __getitem__ and __len__ method, and don't necessarily have any "stored" information about the labels.

In the case of the MNIST dataset class however you can sort the dataset from the label list.

For example when you want to list the indices that have the label 5.

mnist = torchvision.datasets.mnist.MNIST("/")
labels = mnist.train_labels
fives = (labels == 5).nonzero()