Classify Passenger Jets Using PyTorch
This is the first ever image classifier I’ve ever built. This is the first ever data science related blog I’ve written.
I love being at airports, not only because I like traveling, but more importantly, I like to see those passenger jets, the most advanced and comprehensive technology products that a normal person can actually touch and use on a daily basis.
But there always be questions for me and my friends: “How can we tell, at a single glance, whether the jet is Boeing or Airbus”. Although experts and enthusiasts can easily distinguish different jets, my friends and I often have hard time correctly identify them.
With the help of machine learning algorithms, neural networks in particular, we can actually build a classifier to help us distinguish. We can now easily access so many existing architectures that can classify images very well. That being said, we don’t need to spend a lot of time studying those mathematical formulas and coding an algorithm from scratch. Rather, we can utilize packages to do the classification in just lines of codes.
Below is a typical neural network layout. Trust me, it’s a simple one. You can find almost everything about a neural net or deep learning on the internet, so I’ll just link a wiki page here.
OK, that’s enough, let’s dive straight into the codes.
To build the classifier, that’s all you need to import into your file.
To load images used for training, validation and testing, we can use
torchvision.datasets. Firstly, we need to transform the image into the format that the pre-trained architectures was trained on. To avoid potential overfitting, we need to do some image augmentation on the training images. But we will keep the validation and test images as what they were except for basic resize and crop.
Noticed that we have to keep size and the normalization params the same as the pre-trained architecture was trained on. 224 is the size ,
[0.485, 0.456, 0.406] are the means used, and
[0.229, 0.224, 0.225] are the standard deviation used.
Next, we need to extract the model, and apply our own classifier. In the meantime, we need to specify the
optimizer of the model, and train it using a
for loop. It is also easy in several lines of codes
Until this time, your model can predict the kind of a passenger jet.
Follow the codes above will give you a one time prediction tool, but obviously you don’t want to train your model every time you want a prediction. That’s where the
torch.save comes in. You can save the
state_dict of your trained model, and load it next time for usage.
Alright! Let’s see how our classifier works!
Our model has an accuracy of 76%, which is not bad, given it only spent around 30 minutes, 2 epochs in training. Also the images actually contain a lot of noise, and are not well formatted for better training result.
Below are some of the predictions.
❌A320 into B767 ⭕️A321 ⭕️A330 but with noise.
⭕️A340 ⭕️A350 ⭕️A380, all with high confidence.
⭕️B737 ⭕️B747 ⭕️B757, all with high confidence.
⭕️B767 ⭕️B777 ⭕️B787, all with high confidence.
- As we can see, the images used for training and prediction are not showing the whole jets, rather, only middle parts were contained in the image. Some features maybe lost because of this.
- There are some noise in the training set, such as the image of the control panel or one engine of the jet. Although there are only a small number of such images, it may lead to inaccurate classification.
- The majority of the images are showing side views, the front/tail or top/bottom views seldom appear.
Overall, our simple model works good with 76% accuracy. It almost correctly classified all shown examples above.