Rocket or Tower

I recently started working through the fast.ai course. This is my first attempt at training an image classifier.

Attempt One

My training set contained 10 rockets and 10 towers.

Examples from the rocket-and-tower training set

I used the fast.ai library to fine-tune a pretrained ResNet model with only a few lines of code.

A prediction near 0 is more rocket-like; one near 1 is more tower-like.

With only 10 examples of each class, it performed surprisingly well. It barely misclassified the CN Tower as a rocket.

Predictions from the first rocket-or-tower model

Attempt Two

I expanded the dataset to about 300 training images and 15 validation images per class.

This time, I added data augmentation, a changing learning rate, and a second training stage with the model unfrozen.

Data Augmentation

Data augmentation applies realistic transformations to the training images, such as shifts, zooms, and flips. The transformed versions are then used as additional training examples.

Augmented versions of the rocket-and-tower images

Stochastic Gradient Descent with Restarts

I used a cyclical learning-rate schedule that starts high, decreases during a cycle, and then restarts. Varying the learning rate this way can help the optimizer move out of narrow local regions of the loss surface.

I refined the schedule by using lower learning rates in the later epochs.

Learning-rate schedule with lower rates in later epochs

Training the Pretrained Model

Initially, I trained only the model's final layers. I then unfroze the earlier layers and trained them with progressively lower learning rates, using the smallest rate near the input.

Result

The model reached 96.6% accuracy, with one classification error.

The confusion matrix showed which class had been misclassified.

Confusion matrix for the rocket-or-tower classifier

Below are the images about which the model was least certain, including the incorrectly labelled rocket at the bottom left.

Images with the highest classification losses