How to make manual batch in pytorch with MNIST dataset - pytorch

Question

I am new to pytorch.
I am trying to use a different approach when training MNIST dataset in pytorch.
Normally, we put all shuffled numbers for each batch.
What I want to do is (assuming by batch size is 20 and there are 1000 training examples).
Make a batch look this this (I wrote only labels for simplicity)
batch_1 = [0, 0, 1, 0, ... ] only 0 and 1 for this batch
batch_2 = [6, 6, 7, 7, ... ] only 6 and 7 for this batch
batch_3 = [2, 3, 2, 2, ... ] only 2 and 3 for this batch
batch_4 = [4, 5, 5, 4, ... ] only 4 and 5 for this batch
batch_5 = [6, 6, 7, 7, ... ] only 6 and 7 for this batch
batch_6 = [8, 9, 9, 9, ... ] only 8 and 9 for this batch
batch_7 = [2, 3, 2, 2, ... ] only 2 and 3 for this batch
batch_8 = [4, 5, 5, 4, ... ] only 4 and 5 for this batch
...
batch_50 = [4, 5, 5, 4, ... ] only 4 and 5 for this batch
How can I do this in pytorch??
Thank you
So far this is my implementation for training / testing.
I will expand/manipulate this code below to develop the above.
import torch
from torch.utils.data import DataLoader, Subset
import torchvision.datasets as datasets
from torchvision import transforms
import numpy as np
mnist_trainset = datasets.MNIST(root='./data', train=True,
download=True,
transform=transforms.Compose([transforms.ToTensor()]))
mnist_testset = datasets.MNIST(root='./data',
train=False,
download=True,
transform=transforms.Compose([transforms.ToTensor()]))
class_inds = [torch.where(mnist_trainset.targets == class_idx)[0]
for class_idx in mnist_trainset.class_to_idx.values()]
train_dataloaders = [
DataLoader(dataset=Subset(mnist_trainset, inds),
batch_size=50,
shuffle=True,
drop_last=False
)
for inds in class_inds
]
test_dataloader = torch.utils.data.DataLoader(mnist_testset,
batch_size=50,
shuffle=False)
class Model(torch.nn.Module):
def __init__(self):
super(Model, self).__init__()
self.linear_1 = torch.nn.Linear(784, 256)
self.linear_2 = torch.nn.Linear(256, 10)
self.sigmoid = torch.nn.Sigmoid()
def forward(self, x):
x = x.reshape(x.size(0), -1)
x = self.linear_1(x)
x = self.sigmoid(x)
pred = self.linear_2(x)
return pred
model = Model()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
epochs = 20
criterion = torch.nn.CrossEntropyLoss()
test_acc = list()
for epoch in range(epochs):
print(f"epoch {epoch} started")
model.train()
itr = 0
iterators = list(map(iter, train_dataloaders))
while iterators:
itr = itr + 1
iterator = np.random.choice(iterators)
try:
image, label = next(iterator)
optimizer.zero_grad()
pred = model(image)
loss = criterion(pred, label)
loss.backward()
optimizer.step()
except StopIteration:
iterators.remove(iterator)
model.eval()
total = 0
for itr, (image, label) in enumerate(test_dataloader):
pred = model(image)
loss = criterion(pred, label)
# we now need softmax because we are testing.
pred = torch.nn.functional.softmax(pred, dim=1)
for i, p in enumerate(pred):
if label[i] == torch.max(p.data, 0)[1]:
total = total + 1
accuracy = total / len(mnist_testset)
# append accuracy here
test_acc.append(accuracy)