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

Question

[] : 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! :)

Answer 1

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()}')

Answer 2

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()}')