PyTorch MNIST data augmentation

I’m trying to learn to use PyTorch and the first thing I wanted to do was MNIST predictor. I got pretty good results (something like 99%) but I want to use some tricks which Jeremy have been taught. The code is below so can someone explain to me how I can rotate the numbers a little to get more training data (aka data augmentation).

MNIST dataset

train_dataset = torchvision.datasets.MNIST(root='../../data/',
                                       train=True, 
                                       transform=transforms.ToTensor(),
                                       download=True)

test_dataset = torchvision.datasets.MNIST(root='../../data/',
                                      train=False, 
                                      transform=transforms.ToTensor())

Data loader

train_loader = torch.utils.data.DataLoader(dataset=train_dataset,
                                       batch_size=batch_size, 
                                       shuffle=True)

test_loader = torch.utils.data.DataLoader(dataset=test_dataset,
                                      batch_size=batch_size, 
                                      shuffle=False)

Hi, adding transforms.RandomRotation will do what you want.

So the transform will be

transform=transforms.Compose([
    transforms.RandomRotaion(max_degree),
    transforms.ToTensor(),
])

Thanks for reply! Can you show how to use this with my code. I’m not sure am I just stupid or is this wrong code. Where I should add train_loader object?

Now I got it! Thanks Masaki Kozuki!

Oh, sorry for my late reply.

But glad to hear that you solved by yourself! Hooray.