Load the dataset, using the following code:

mnist = input_data.read_data_sets("data/mnist", one_hot=True)

In the preceding code, data/mnist implies the location where we store the MNIST dataset, and one_hot=True implies that we are one-hot encoding the labels (0 to 9).

We will see what we have in our data by executing the following code:

print("No of images in training set {}".format(mnist.train.images.shape))
print("No of labels in training set {}".format(mnist.train.labels.shape))

print("No of images in test set {}".format(mnist.test.images.shape))
print("No of labels in test set {}".format(mnist.test.labels.shape))

No of images in training set (55000, 784)
No of labels in training set (55000, 10)
No of images in test set (10000, 784)
No of labels in test set (10000, 10)

We have 55000 images in the training set, each image is of size 784, and we have 10 labels, which are actually 0 to 9. Similarly, we have 10000 images in the test set.

Now, we'll plot an input image to see what it looks like:

img1 = mnist.train.images[0].reshape(28,28)
plt.imshow(img1, cmap='Greys')

Thus, our input image looks like the following: