Define manually sorted MNIST dataset with batch size = 1 in PyTorch

user_14052350 0
Viewed 538

[] : this indicates a batch. For example, if the batch size is 5, then the batch will look something like this [1,4,7,4,2]. The length of [] indicates the batch size.

What I want to make a training set something looks like this:

[1] -> [1] -> [1] -> [1] -> [1] -> [7] -> [7] -> [7] -> [7] -> [7] -> [3] -> [3] -> [3] -> [3] -> [3] -> ... and so on

Which means that firstly five 1s (batch size = 1), secondly five 7s (batch size = 1), thirdly five 3s (batch size = 1) and so on...

Can someone please provide me an idea?

It will be very helpful if someone can explain how to implement this with codes.

Thank you! :)

pytorch-dataloader batchsize mnist pytorch
2 Answers

If you want a DataLoader where you just want to define the class label for each sample then you can make use of the torch.data.utils.Subset class. Despite its name it doesn't necessarily need to define a subset of dataset. For example

import torch
import torchvision
import torchvision.transforms as T
from itertools import cycle

mnist = torchvision.datasets.MNIST(root='./', train=True, transform=T.ToTensor())

# not sure what "...and so on" implies, but define this list however you like
target_classes = [1, 1, 1, 1, 1, 7, 7, 7, 7, 7, 3, 3, 3, 3, 3]

# create cyclic iterators of indices for each class in MNIST
indices = dict()
for label in torch.unique(mnist.targets).tolist():
    indices[label] = cycle(torch.nonzero(mnist.targets == label).flatten().tolist())

# define the order of indices in the new mnist subset based on target_classes
new_indices = []
for t in target_classes:
    new_indices.append(next(indices[t]))

# create a Subset of MNIST based on new_indices
mnist_modified = torch.utils.data.Subset(mnist, new_indices)
dataloader = torch.utils.data.DataLoader(mnist_modified, batch_size=1, shuffle=False)

for idx, (x, y) in enumerate(dataloader):
    # training loop
    print(f'Batch {idx+1} labels: {y.tolist()}')
user_27900470

If you want a DataLoader that returns five samples in a row of the same class, but you don't want to define the class for each index manually then you can create a custom sampler. For example

import torch
import torchvision
import torchvision.transforms as T
from itertools import cycle

class RepeatClassSampler(torch.utils.data.Sampler):
    def __init__(self, targets, repeat_count, length, shuffle=False):
        if not torch.is_tensor(targets):
            targets = torch.tensor(targets)

        self.targets = targets
        self.repeat_count = repeat_count
        self.length = length
        self.shuffle = shuffle

        self.classes = torch.unique(targets).tolist()
        self.class_indices = dict()
        for label in self.classes:
            self.class_indices[label] = torch.nonzero(targets == label).flatten() 

    def __iter__(self):
        class_index_iters = dict()
        for label in self.classes:
            if self.shuffle:
                class_index_iters[label] = cycle(self.class_indices[label][torch.randperm(len(self.class_indices))].tolist())
            else:
                class_index_iters[label] = cycle(self.class_indices[label].tolist())

        if self.shuffle:
            target_iter = cycle(self.targets[torch.randperm(len(self.targets))].tolist())
        else:
            target_iter = cycle(self.targets.tolist())

        def index_generator():
            for i in range(self.length):
                if i % self.repeat_count == 0:
                    current_class = next(target_iter)
                yield next(class_index_iters[current_class])
    
        return index_generator()

    def __len__(self):
        return self.length


mnist = torchvision.datasets.MNIST(root='./', train=True, transform=T.ToTensor())
dataloader = torch.utils.data.DataLoader(
        mnist,
        batch_size=1,
        sampler=RepeatClassSampler(
            targets=mnist.targets,
            repeat_count=5,
            length=15,      # How many total to pick from your dataset
            shuffle=True))

for idx, (x, y) in enumerate(dataloader):
    # training loop
    print(f'Batch {idx+1} labels: {y.tolist()}')
user_27900470
Related