[FIXED] Training set not showing in graph

Issue

For some reason my training set is not showing as a matplotlib graph, I am using the torchvision dataset MNIST.

The code is running without any errors however the graph of data won't be shown. I tried doing the .view() to reshape the data shape from a (1, 28, 28) to a (28, 28) that will be allowed, however the data won't be shown. Here is my code:

import torch
import torchvision
import matplotlib.pyplot as plt
from torchvision import transforms, datasets

train = datasets.MNIST("", train=True, download=True,
                       transform= transforms.Compose([transforms.ToTensor()]))

test = datasets.MNIST("", train=False, download=True,
                       transform= transforms.Compose([transforms.ToTensor()]))

trainset = torch.utils.data.DataLoader(train, batch_size=10, shuffle=True)
testset = torch.utils.data.DataLoader(test, batch_size=10, shuffle=True)


for data in trainset:
    print(data)
    break

x, y = data[0][0], data[1][0]


print(y)


plt.imshow(data[0][0].view(28,28))
plt.show()

Solution

Firstly, define x, y inside of the loop.

for data in trainset:
    x, y = data[0][0], data[1][0]
    break

Then use .reshape() to change from (1, 28, 28) to (28, 28)

x = x.reshape([28, 28])

Convert to a numpy array (better for matplotlib) and plot

x = x.numpy()
plt.imshow(x, cmap="Greys_r")
plt.show()

Below is the output

Output

Answered By - edgar_maddocks