PyTorch: Image Classification using Pre-Trained Models¶

Image classification is an active area of research in computer vision where we look at an image and assign a label to it. Image classification tasks generally have images with one object present in them which they classify into a particular category. Over the years, many competitions are organized by Imagenet for producing a high-quality image classifier. The winner of these competitions are complicated models that are quite good at the task. Imagenet dataset has nearly 1000 target categories which include different kinds of animals, sea life, daily use objects (clock, paper, etc), etc. These models are so good at image classification tasks that nowadays rarely does someone design a new network for their purpose. We can directly use one of the pre-trained networks available from deep learning libraries like PyTorch.

What can you learn from this Article?¶

As a part of this tutorial, we have covered how to use pre-trained networks/models available from Python library PyTorch (torchvision) to solve image classification task. PyTorch has a helper module named torchvision specifically designed for computer vision tasks, providing these pre-trained networks with weights. We have an option to load these network architectures without weights if we have enough data and computing resources to train the network. Below, we have listed some of the models provided by torchvision module.

• VGG
• AlexNet
• ConvNeXt
• DenseNet
• EfficientNet
• EfficientNetV2
• GoogLeNet
• Inception V3
• MNASNet
• MobileNet V2
• MobileNet V3
• RegNet
• ResNet
• ResNeXt
• ShuffleNet V2
• SqueezeNet

Below, we have listed essential sections of the Tutorial to give an overview of the material covered.

Important Sections Of Tutorial¶

1. Load Images
   • 1.1 Download Images from the Internet
   • 1.2 Load Images in Memory using Pillow Library
   • 1.3 Convert Pillow Images to Torch Tensors
2. Load Model with Pre-Trained Weights
3. Preprocess Images and Make Predictions
4. Rerieve Target Labels
5. Visualize Results
   • 5.1 RestNet Predictions Visualization
   • 5.2 MobileNet Predictions Visualization
6. Try Other Models

Below, we have imported the necessary libraries that we have used in our tutorial and printed the versions of them.

Install Latest torchvision¶

• !pip install -U torchvision

1. Load Images ¶

In this section, we have simply downloaded a few random images from the internet and loaded them in memory. We have converted them to torch tensors as well for image classification through Pytorch networks.

1.1 Download Images from Internet¶

Below, we have downloaded 6 images from the internet. The images are of panda, koala, lion, sea lion, wall_clock, and digital clock. The images are downloaded using the shell command wget. Please feel free to download other images if you want to try them. Just make sure that images have just one object in them else it can confuse the classifier.

1.2 Load Images in Memory using Pillow Library¶

Below, we have simply loaded images in memory using the Python library Pillow. The Pillow provides us with lots of functionalities to work with images like loading, cropping, resizing, filtering, etc. Do check the below link in your free time.

• Pillow

1.3 Convert Pillow Images to Torch Tensors¶

Here, we have converted our Pillow images to torch tensors. The torchvision library provides a functional API to convert images to tensors and vice-versa. The function pil_to_tensor() is used to convert Pillow image to torch tensor and to_pil_image() is used to convert torch tensor to pillow image.

2. Load Model with Pre-Trained Weights ¶

In this section, we have loaded pre-trained PyTorch image classifiers available from torchvision. The models are available from "models" sub-module of torchvision.

We have loaded two models for our experimentation purpose.

1. ResNet101
2. MobileNet V3

We just need to create an instance of these models by loading them. By default, weights parameter of the constructor is None which means that no weights will be loaded (only architecture). In order to load weights, we need to import model_name_Weights object. It has an attribute named DEFAULT which will load default weights as there is more than one version of weights for the model (E.g., for Resnet there are V1 and V2 weights as well).

After loading models, we have set them in evaluation mode by calling eval() on them. This will deactivate batch normalization and dropout layers.

As we had said earlier, if you have enough images then you can train the model as well. In that case, you have options like training the whole network (do not load pre-trained weights) and fine-tuning existing weights (transfer learning).

3. Preprocess Images and Make Predictions ¶

In this section, we are making predictions using our loaded models.

All these models are trained with specific image sizes and can make predictions on that size only. So, we'll need to resize images and also normalize them.

To resize and normalize the image, the weights object that we loaded earlier provides us with transform() method. It returns a PyTorch transformation that can be applied to the Pillow image. It'll resize the image as well as normalize it. We can give this processed image directly to the network for prediction.

After transforming images, we have made predictions on all 6 images using both models. As both models are trained on ImageNet dataset, they have 1000 output probabilities per example. The category with the highest probability will be the predicted target label.

Below, we have listed transformation applied by RestNet weights for reference purposes.

crop_size=[224]
resize_size=[256]
mean=[0.485, 0.456, 0.406]
std=[0.229, 0.224, 0.225]
interpolation=InterpolationMode.BILINEAR

4. Rerieve Target Labels ¶

Here, we have retrieved the predicted target label for images by our models.

As we said earlier, the prediction made by networks is a thousand probabilities per image. We need to retrieve target labels using these probabilities.

First, we have sorted probabilities from maximum to minimum and taken indexes of the first 3 probabilities.

Then, we retrieved the target label using these 3 probability indexes. The actual target label for all thousand classes is available through meta dictionaries' "categories" key.

After retrieving the predicted target labels, we have printed them as well.

5. Visualize Results ¶

At last, in this section, we have visualized the prediction of our models.

5.1 RestNet Predictions Visualization¶

Below, we have visualized prediction made by ResNet101 model using matplotlib. We can notice from the results that the first label is the correct label for all images. Hence, we can conclude that ResNet101's results are quite good.

5.2 MobileNet Predictions Visualization¶

Here, we have visualized predictions of MobileNet V3 predictions. We can notice from the results that all images are correctly identified by the classifier. It seems that both classifiers are quite good at the job.

6. Try Other Models ¶

Here, we have listed many other pre-trained image classifiers available from PyTorch. If you are not getting good results using the above models then you should try one of the below.

References¶

• PyTorch
• Transfer Learning Using PyTorch
• Torchvision: Object Detection using Pre-Trained PyTorch Models