FastAI With TPU In PyTorch For Multiclass Image Classification

FastAI With TPU In PyTorch

Computer vision is one of the most trending subfields in Artificial Intelligence because of its wide variety of applications. In some domains, they even suppress human intelligence in recognising images with speed and accuracy. In this article, we will demonstrate one of the most popular computer vision applications — multiclass image classification problems using fastAI library and TPU as the hardware accelerator. TPU, or Tensor Processing Unit accelerates the training process of heavy deep learning models.

Topics covered in this article:

  • Multiclass Image classification
  • Popular models for image classification
  • Hands-on implementation in PyTorch using TPU

Multiclass Image classification

We use image classification for recognising objects in the image and can be used in detecting brand names using logo, classifying objects, etc. But these solutions have a limitation that they can just recognise the objects but they cannot find the location of the object. But for recognising the objects these image classification models are easy to implement compared to object localisation.

AIM Daily XO

Join our editors every weekday evening as they steer you through the most significant news of the day, introduce you to fresh perspectives, and provide unexpected moments of joy
Your newsletter subscriptions are subject to AIM Privacy Policy and Terms and Conditions.

Popular Models for the Image Classification

We can use VGG -16,19, Resnet, Inception v1,v2,v3, Wideresnt, Resnext, DenseNet etc, which are the advanced variants of a convolutional neural network. These are the popular image classification networks and used as the backbone for many state-of-the-art object detection and segmentation algorithms.

Image Classification using FasAI library and TPU hardware

We will implement this work in the following steps:-

Download our Mobile App

1. Selecting Hardware Accelerator

Here we are using Google Colab for our implementation. To use TPU in Google Colab, we need to open the edit option then notebook settings, and change the hardware accelerator to TPU.

By running the below code snippet you can check whether your notebook is using TPU or not.

import os

assert os.environ['COLAB_TPU_ADDR']

Path = 'grpc://'+os.environ['COLAB_TPU_ADDR']

print('TPU Address:', Path)

2. Loading FastAI library

In the below code snippet, we will import the fastAI library. 

from import *

from fastai.metrics import error_rate, accuracy

3. Customised Dataset

In the below code snippet, you can also try with your customised dataset.

PATH = '/content/images/dataset'


tfms = get_transforms(do_flip=True)

data = ImageDataBunch.from_folder(PATH, valid_pct=0.2, ds_tfms=tfms, size=299, bs=16).normalize(imagenet_stats)

data.show_batch(rows=4, figsize=(8, 8))

4. Loading Pre-Trained Deep Learning Model

In the below code snippet, we are going to import the VGG-19 batch_normalisation model. We will use this as an instance of the computer vision learner module of fastAI.

learn = cnn_learner(data, models.vgg19_bn, metrics=accuracy)

5. Training the Model

 In the below code snippet, we tried with one epoch. 


In the output, we can see that we have got an accuracy of 0.99 and it took 1 minute, 2 second.

In the below code snippet, we are displaying our results using the confusion matrix.

con_matrix = ClassificationInterpretation.from_learner(learn)


FastAI With TPU In PyTorch

6. Making Prediction using the Model

In the below code snippet, we can test our own image by giving the path of the image in test_your_image.

test_your_image='/content/images (3).jpg'

test = open_image(test_your_image)

FastAI With TPU In PyTorch

In the below code snippet we can get the output tensor and which class it belongs.


FastAI With TPU In PyTorch

As we can see in the above output, the model has predicted the class label for the input image that it belongs to the ‘flower’ category.


In the above demonstration, we implemented a multiclass image classification with few lines of code using the fastAI library with TPU and we used the pre-trained VGG-19 model. We achieved 0.99 accuracy in classifying the validation dataset in this task.

Sign up for The Deep Learning Podcast

by Vijayalakshmi Anandan

The Deep Learning Curve is a technology-based podcast hosted by Vijayalakshmi Anandan - Video Presenter and Podcaster at Analytics India Magazine. This podcast is the narrator's journey of curiosity and discovery in the world of technology.

Prudhvi varma
AI enthusiast, Currently working with Analytics India Magazine. I have experience of working with Machine learning, Deep learning real-time problems, Neural networks, structuring and machine learning projects. I am a Computer Vision researcher and I am Interested in solving real-time computer vision problems.

Our Upcoming Events

24th Mar, 2023 | Webinar
Women-in-Tech: Are you ready for the Techade

27-28th Apr, 2023 I Bangalore
Data Engineering Summit (DES) 2023

23 Jun, 2023 | Bangalore
MachineCon India 2023 [AI100 Awards]

21 Jul, 2023 | New York
MachineCon USA 2023 [AI100 Awards]

3 Ways to Join our Community

Telegram group

Discover special offers, top stories, upcoming events, and more.

Discord Server

Stay Connected with a larger ecosystem of data science and ML Professionals

Subscribe to our Daily newsletter

Get our daily awesome stories & videos in your inbox

Council Post: Evolution of Data Science: Skillset, Toolset, and Mindset

In my opinion, there will be considerable disorder and disarray in the near future concerning the emerging fields of data and analytics. The proliferation of platforms such as ChatGPT or Bard has generated a lot of buzz. While some users are enthusiastic about the potential benefits of generative AI and its extensive use in business and daily life, others have raised concerns regarding the accuracy, ethics, and related issues.