Building an Image Classification Model From Scratch Using PyTorch
By: Benedict Neo
Tags: Image Classification, Data Science, Artificial Intelligence, Machine Learning, Deep Learning

Deep Learning

Building an Image Classification Model From Scratch Using PyTorch

Look at the image above. What do you see? You might think this is a trivial question — clearly, it’s a yellow banana casting a shadow on a pink background. But how might computers recognize this image without having eyes and brains like humans do?

Computers aren’t able to infer about the world intuitively like we do, so for computers to “see” and recognize objects, scientists have to crack the complex system that is the human brain and implement it onto a computer. As children, we have an innate curiosity to explore and experiment with the world. We stumble upon unfamiliar objects, interact with and learn about them, and over time make inferences about things like whether that object is dangerous or harmless.

For computers to learn and extrapolate information similarly, we either program a set of rules for them to follow in a process called “supervised learning,” or we give them the answers and they try to understand what was the initial question or goal, or “unsupervised learning.” With both methods, it’s a major challenge to encode the right rules or give the right format of answers for computers to see images.

However, with recent advancements in deep learning, computers can now recognize and classify images and even videos with impressive accuracy. So what’s the magic behind this revolutionary change?

To answer that, we’ll take a dive into the field of image recognition.

Image Recognition

Image recognition is essentially a computer vision technique that gives “eyes” to computers for them to “see” and understand the world through images and videos. But as mentioned at the start of this article, computers don’t see the world the way we do. What they can “see” when they are given image data are numbers that relate to the intensity, brightness, color, shape, outline, etc. of an image.

Since similar objects will have the same information in brightness, color, etc. computers can learn those patterns, and remember what that object is the next time it “sees” it. This is called “image recognition” — a supervised ML technique where computers learn and predict image contents.

Training the model

Image recognition models are trained to take in an image as input, deconstruct it down to its basic form, then produce labels that categorize the image via a neural network (NN).

Images (input) → NN (layers) → Labels (output)

Example

As an example, let’s train a model to recognize if an image is of the Eiffel Tower.

Here’s an example of what the model does in practice:

Input: Image of Eiffel Tower Layers in NN: The model will first see the image as pixels, then detect the edges and contours of its content. Finally, it will look at the whole object before producing a final guess about what the model “sees.” Output: Eiffel Tower (label)

Types of image recognition

The 3 classes of image recognition are:

Single class — one label per image (our example) Multiclass — several labels per image (dog and cat in an image) Binary classifiers — two classes (i.e. “Eiffel Tower” or “Not Eiffel Tower”)

Alongside with labels, multiclass image recognition models will produce confidence scores that tell us how confident they are with their prediction. Say there’s a picture of a red apple fed into the network. The model might give a score of 97% for the prediction of an apple and 3% for a red ball, meaning that the model is 97% sure it is an apple.

If you want to go more into detail about the technical aspects of the different architectures of image recognition models, and how the work at the fundamental level, check out this series by fritz.ai on image recognition.

Building an image recognition model

We can build an image recognition model using traditional statistical approaches such as using Support Vector Machines or Decision Trees, but the state-of-the-art method is with Neural Networks.

Today, the de facto algorithm used for image recognition are convolutional neural networks (CNNs).

Based on the ImageNet Large Scale Visual Recognition Challenge, a CNN model made predictions on millions of images with 1000 classes and its performance is now close to that of humans.

Convolutional Neural Networks (CNNs)

CNNs are a class of Neural Network which work incredibly well on image data. They work similarly to how we humans recognize objects. As mentioned previously, the network first looks at the pixels of an image, then it gradually extracts the important features of the images using convolutions.

The term convolution here refers to the mathematical combination of two functions, thus producing a third function. CNNs perform convolutions on the image data using a filter or kernel, then produce a feature map.

The example above uses a robot as the input image and multiple feature maps for processing. You can see that the model produces an output via convolutions and subsampling.

An example of feature maps (source)

The above images show the process of feature mapping. While to human eyes they just look like weird cat pictures, these filters allows computers to pinpoint important features of the image.

Now that we’ve covered the basics, let’s get to the fun part : building our CNN model. To do so, we’ll use an image recognition framework called PyTorch.

PyTorch

PyTorch is an open source machine learning framework that speeds up the path from research prototyping to production deployment. Its two primary purposes are:

Replacing Numpy to use the power of GPUs for faster computation Calculating gradients to perform backpropagation on neural networks

Because PyTorch is easy to start and learn, it’s excellent for anyone already familiar with Python and looking to get started with deep learning.