Image Classification on the CIFAR-10 Dataset using Convolutional Neural Networks
Image classification is a fundamental problem in computer vision that involves assigning a label or class to an input image. Some real-world applications include facial recognition, medical image diagnosis, self-driving cars, and more. In recent years, deep learning approaches, especially convolutional neural networks (CNNs), have achieved state-of-the-art performance on image classification tasks.
In this blog post, we‘ll walk through how to build an image classifier using a CNN and the popular CIFAR-10 dataset. By the end, you‘ll understand the key components of a CNN architecture and how to train a model to recognize 10 different image categories. Let‘s dive in!
The CIFAR-10 Dataset
CIFAR-10 is a widely used dataset for benchmarking image classification models. It consists of 60,000 32×32 color images across 10 classes, with 6,000 images per class. The classes include common objects like airplanes, cars, birds, cats, deer, dogs, frogs, horses, ships, and trucks.
Here are some key characteristics of the CIFAR-10 dataset:
- 50,000 training images and 10,000 test images
- 10 classes with 6,000 images per class total
- 5,000 training images and 1,000 test images per class
- Low resolution 32×32 images (smaller than typical photos)
- Color images with 3 channels (RGB)
The relatively small size and low resolution of the images makes CIFAR-10 a good starter dataset for experimenting with CNN architectures without needing a costly GPU or long training times. At the same time, achieving high accuracy on CIFAR-10 is still challenging due to the diversity of the image classes.
Why Use a CNN for Image Classification?
CNNs are currently the go-to neural network architecture for computer vision tasks. But what makes them so effective for working with image data? The key is how CNNs are able to capture spatial patterns and hierarchies of visual features.
The main building blocks of a CNN are:
- Convolutional layers that learn visual features by scanning a sliding window over the input
- Pooling layers that downsample the spatial dimensions to reduce computation
- Dense layers at the end for classification
By stacking multiple convolutional and pooling layers, a CNN can learn a rich hierarchy of features, from basic edges and shapes to more complex visual concepts. This allows CNNs to effectively distill the relevant information from raw images into a form suitable for classification.
Another advantage of CNNs is that they are translation invariant, meaning they can recognize visual patterns regardless of their spatial location in the image. This is achieved by the sliding window approach that applies the same convolutional filters across the entire input.
With these properties, CNNs have been highly successful not just for image classification but also related tasks like object detection, semantic segmentation, and more. In the next section, we‘ll see how to implement a CNN for CIFAR-10 classification in code.
Building the CNN Model
We‘ll use the Keras deep learning framework to build our CNN image classifier. Keras provides a simple, intuitive API for defining neural network architectures as a stack of layers. Here‘s the step-by-step process:
1. Load and preprocess the data
First, we load the CIFAR-10 dataset from Keras using cifar10.load_data(). This gives us pre-split training and test sets. We then rescale the pixel values from [0, 255] to [0, 1] as a simple normalization step. We also one-hot encode the labels to match the output of the classifier.
from tensorflow.keras.datasets import cifar10
from tensorflow.keras.utils import to_categorical
(X_train, y_train), (X_test, y_test) = cifar10.load_data()
X_train = X_train.astype("float32") / 255.0
X_test = X_test.astype("float32") / 255.0
y_train = to_categorical(y_train, 10)
y_test = to_categorical(y_test, 10)
2. Define the model architecture
Next, we use the Keras Sequential model to stack layers into a CNN architecture:
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Conv2D, MaxPooling2D, Dense, Flatten
model = Sequential()
model.add(Conv2D(32, (3, 3), activation=‘relu‘, padding=‘same‘, input_shape=X_train.shape[1:]))
model.add(Conv2D(32, (3, 3), activation=‘relu‘, padding=‘same‘))
model.add(MaxPooling2D((2, 2)))
model.add(Conv2D(64, (3, 3), activation=‘relu‘, padding=‘same‘))
model.add(Conv2D(64, (3, 3), activation=‘relu‘, padding=‘same‘))
model.add(MaxPooling2D((2, 2)))
model.add(Conv2D(128, (3, 3), activation=‘relu‘, padding=‘same‘))
model.add(Conv2D(128, (3, 3), activation=‘relu‘, padding=‘same‘))
model.add(MaxPooling2D((2, 2)))
model.add(Flatten())
model.add(Dense(128, activation=‘relu‘))
model.add(Dense(10, activation=‘softmax‘))
This architecture consists of three blocks of two convolutional layers each followed by max pooling. The number of filters doubles at each block, from 32 to 64 to 128. ReLU activation adds non-linearity after each conv layer.
After the conv layers, the 2D feature maps are flattened into a 1D vector. The final dense layers learn the classification mapping from the features to the 10 classes. Softmax activation gives the final class probabilities.
3. Train the model
We compile the model with categorical cross-entropy loss and Adam optimization. Accuracy is a useful metric to monitor in addition to loss during training:
model.compile(loss=‘categorical_crossentropy‘, optimizer=‘adam‘, metrics=[‘accuracy‘])
history = model.fit(X_train, y_train, validation_data=(X_test, y_test),
epochs=50, batch_size=64)
We train for 50 epochs with a batch size of 64. Setting aside a validation set (in this case the test set) allows us to monitor overfitting during training.
4. Evaluate performance
Once the model is trained, we can evaluate its final loss and accuracy on the test set:
loss, accuracy = model.evaluate(X_test, y_test)
print(f"Test loss: {loss:.3f}")
print(f"Test accuracy: {accuracy:.3f}")
We can also plot the training curves to see how the model‘s performance evolved over time:
%matplotlib inline
import matplotlib.pyplot as plt
plt.plot(history.history[‘accuracy‘], label=‘Training accuracy‘)
plt.plot(history.history[‘val_accuracy‘], label=‘Validation accuracy‘)
plt.xlabel(‘Epoch‘)
plt.ylabel(‘Accuracy‘)
plt.ylim([0.5, 1])
plt.legend(loc=‘lower right‘)
plt.show()
The training curves can help diagnose issues like overfitting (when validation performance starts to degrade) or underfitting (when validation performance plateaus at a low value).
Experimenting with Model Design
There are many ways to tweak the CNN architecture that can impact its performance on the CIFAR-10 task. Some key design choices include:
- Number of convolutional layers and filters per layer
- Kernel size of the convolutional filters (e.g. 3×3, 5×5)
- Pooling size (e.g. 2×2, 4×4) and type (e.g. max pooling, average pooling)
- Depth of the network (number of layers)
- Addition of fully connected layers at the end
- Regularization (e.g. dropout, L2 weight decay)
In general, making the network deeper with more layers and filters will increase its representational power, but also requires more training data and computation to avoid overfitting.
Pooling provides an important means of downsampling and summarizing the spatial features between convolutional layers. Max pooling tends to work better than average pooling for extracting the most salient features. Very large pooling windows may lose too much spatial information.
Adding one or more fully connected layers before the final classification layer can improve accuracy by learning a non-linear mapping from the spatial features to the classes. However, additional FC layers increase the parameter count and risk of overfitting.
Regularization techniques help prevent overfitting by constraining the model complexity in some way. Some common approaches in CNNs include adding dropout layers and L2 weight decay to the convolutional and dense layers.
The best architecture will depend on the particular dataset and task. It‘s good practice to experiment with a few different designs to get intuition for what works well.
Conclusion and Next Steps
In this blog post, we explored how to train a CNN to classify images from the CIFAR-10 dataset. Some key takeaways:
- CIFAR-10 is a good benchmark dataset for image classification models with 60,000 32×32 color images across 10 classes
- CNNs are effective for image classification due to their ability to learn spatial hierarchies of visual features
- Keras makes it easy to build CNNs as a stack of convolutional, pooling, and dense layers
- Techniques like data normalization, data augmentation, and regularization are important for training CNNs effectively
- Model architecture and hyperparameters can be tuned to improve performance on a task
To dive deeper into this topic, I recommend the following:
- Explore transfer learning by starting with a pre-trained CNN and fine-tuning it on CIFAR-10
- Visualize the learned weights and activations of the CNN to gain insight into what features it‘s learning
- Apply data augmentation techniques like random cropping, flipping, and rotating to expand the training set and increase model robustness
- Experiment with the latest CNN architectures like ResNet, DenseNet, and EfficientNet to see how they perform on CIFAR-10
At a high level, the same process of data loading, model definition, training, and evaluation can be applied to other image classification tasks and datasets beyond CIFAR-10. I encourage you to try building a CNN classifier on your own image data as well!
