I keep coming back to the MNIST dataset when I need to quickly test image processing pipelines or validate that a model architecture is actually learning. It is the hello world of machine learning datasets, and after working with it dozens of times I have a workflow that works every time. Let me walk you through exactly how I load and plot the MNIST dataset in Python without running into the common pitfalls.
MNIST stands for Modified National Institute of Standards and Technology database. It contains 70,000 square images (28×28 pixels) of handwritten digits from 0 to 9, split into 60,000 training images and 10,000 test images. The images are grayscale and stored as pixel values between 0 and 255. Nothing complicated, which is exactly why I reach for it first whenever I am prototyping something new.
TLDR
- Load MNIST with
tensorflow.keras.datasets.mnist.load_data()– it returns tuples of (train_X, train_y) and (test_X, test_y) - Training data shape: (60000, 28, 28), labels shape: (60000,). Test data: (10000, 28, 28), labels: (10000,)
- Plot samples with Matplotlib using
pyplot.imshow()withcmap='gray' - No preprocessing needed – data loads in ready-to-use format
- Next step: feed to a simple neural network or CNN for digit classification
What Makes MNIST So Popular?
I get asked this a lot, so let me give you the short version. MNIST is popular for three straightforward reasons. First, it is publicly available and free to use. Second, the data requires almost no preprocessing – you get clean pixel arrays that you can feed directly into a model. Third, it is a proper benchmark, which means you can compare your results against published baselines and know if you are on the right track. That combination is rare, which is why you see MNIST used in tutorials, research papers, and quick prototypes alike.
Loading the MNIST Dataset in Python
The fastest way to load MNIST in Python is through TensorFlow Keras. I use this almost every time because it handles the download and caching for you – you call one function and the data is ready.
Step 1: Import and Load
from tensorflow.keras.datasets import mnist
(train_X, train_y), (test_X, test_y) = mnist.load_data()
The first time you run this, Keras downloads the dataset from the internet and caches it in your home directory. After that, it loads from cache and you will not notice any delay.
Step 2: Inspect the Shapes
I always print the shapes right after loading. It takes one line and catches most data issues before they cause confusing errors downstream. Here is what you should see.
print('X_train:', train_X.shape)
print('Y_train:', train_y.shape)
print('X_test: ', test_X.shape)
print('Y_test: ', test_y.shape)
Output:
X_train: (60000, 28, 28)
Y_train: (60000,)
X_test: (10000, 28, 28)
Y_test: (10000,)
Here is how I read this. The training input is 60,000 images of 28 by 28 pixels each. The training labels are 60,000 scalar values (the digit each image represents, 0 through 9). The test set mirrors this structure with 10,000 samples. Each pixel value in the array ranges from 0 to 255 – 0 is black, 255 is white, and values in between represent shades of gray.
Step 3: Plot Some Samples
Before I build any model, I always plot a few samples. Visual inspection catches things like corrupted images, unexpected label distributions, or preprocessing bugs that are hard to spot from numbers alone.
from matplotlib import pyplot as plt
for i in range(9):
plt.subplot(330 + 1 + i)
plt.imshow(train_X[i], cmap='gray')
plt.show()
The cmap='gray' argument is important. Without it, Matplotlib defaults to a colormap that maps low values to blue and high values to red, which is not what you want for grayscale images. I have seen people get confused by their MNIST data because they forgot this one argument.
Complete Code to Load and Plot MNIST
Here is the full script in one piece. Copy it into a notebook or a .py file and run it directly.
from tensorflow.keras.datasets import mnist
from matplotlib import pyplot as plt
# Load data
(train_X, train_y), (test_X, test_y) = mnist.load_data()
# Print shapes
print('X_train:', train_X.shape)
print('Y_train:', train_y.shape)
print('X_test: ', test_X.shape)
print('Y_test: ', test_y.shape)
# Plot 9 samples
for i in range(9):
plt.subplot(330 + 1 + i)
plt.imshow(train_X[i], cmap='gray')
plt.show()
What to Do After Loading
Now that you have the data loaded, you have a few natural next steps. The most common one is digit classification – teaching a model to predict which digit (0-9) each image represents. For this task, a simple fully-connected neural network works, but a Convolutional Neural Network (CNN) will give you noticeably better results. A CNN learns spatial patterns in the pixel grid, which is exactly what makes handwritten digits distinguishable.
A CNN for MNIST typically uses three types of layers. The convolutional layer applies small filters across the image to detect edges and shapes. The pooling layer reduces the spatial size of the data while keeping the most important information. The flattening layer converts the 2D feature maps into a 1D vector that a dense output layer can process. Stack these together and you have a classifier that routinely achieves over 99% accuracy on the test set.
FAQ
Q: Can I use MNIST without TensorFlow?
Yes. PyTorch has torchvision.datasets.MNIST and scikit-learn has a function in sklearn.datasets. All three options load the same underlying data.
Q: Are the pixel values normalized?
No, they load as raw integers from 0 to 255. Most practitioners normalize them to a 0-1 range by dividing by 255 before feeding them to a model.
Q: How many digits does MNIST contain?
MNIST contains 10 digit classes (0 through 9). The training set has 60,000 images and the test set has 10,000 images.
Q: What image size is MNIST?
Each image is 28 by 28 pixels, grayscale, with pixel values ranging from 0 to 255.
Q: Is MNIST still relevant for modern machine learning research?
For baselines and quick prototyping, yes. It is not challenging enough to represent the frontier of computer vision, but it remains a useful sanity check for new ideas and architectures.
