How to split MNIST images based on their labels?

1

Dear Altruists,

I am currently working with MNIST dataset. I am able to download and load training data. For my project, I need to train my model with images belonging to a subset of labels, not all of them. Is there a way to do it?

I am new to Pytorch. Thank you in advance!

2

You could filter out the unwanted labels by directly accessing the dataset.targets and/or wrap it in a custom Dataset.

3

Hello @ptrblck thanks a lot for your help. This is how I did it:

train_data = datasets.MNIST(root='~/.pytorch/MNIST_data/', train=True,
                                   download=True, transform=transform)
train_set = [] 

for i in range(0,len(train_data)):
    if train_data[i][1] < 7:
        train_set.append(train_data[i])