CoolFace
Apppublic

adil22jaleel/CIFAR_ResnetModel

sourceHugging Facemitupdated 3y agoView on Hugging Face
0likes
App README

CIFAR-10 Classification using Lightning Module

This repository contains code for training a ResNet-based model for CIFAR-10 classification using PyTorch Lightning. The model architecture includes custom ResBlocks for feature extraction. Albumentations library is used for data augmentation. The code also includes functionality for GradCAM visualization of misclassified images.

Understanding the different functions

ResBlock

The ResBlock is a custom PyTorch module designed to create a residual block, which is a fundamental building block in the ResNet architecture. The ResNet architecture is widely used for deep learning tasks, particularly for image classification, due to its ability to effectively train very deep neural networks.

Model Architecture

The model architecture is based on the ResNet design, which includes custom ResBlocks for feature extraction. The model is composed of the following components:

  • preplayer: The initial 2D convolutional layer followed by batch normalization and ReLU activation.
  • layer1_X: The first part of the first residual block, which includes a 2D convolutional layer, max-pooling, batch normalization, and ReLU activation.
  • layer1_R: The second part of the first residual block, which is an instance of the custom ResBlock class.
  • layer2_X: The first part of the second residual block, which includes a 2D convolutional layer, max-pooling, batch normalization, and ReLU activation.
  • layer3_X: The first part of the third residual block, which includes a 2D convolutional layer, max-pooling, batch normalization, and ReLU activation.
  • layer3_R: The second part of the third residual block, which is another instance of the custom ResBlock class.
  • maxpool_4: A max-pooling layer with a kernel size of 4.
  • linear: The final fully connected layer that produces the output logits.

The detail information about the model and the lightning module is explained here

Gradio

Gradio is a Python library that simplifies creating web-based interfaces for machine learning models. It allows researchers, developers, and data scientists to quickly share their machine learning models and provide an interactive experience for end-users. Gradio's simple and intuitive API makes it easy to generate interfaces for a wide range of models and tasks.

Usage

  • Make sure to have the model weights file model.pth and the images for the examples in the specified folder images_folder.
  • The provided code assumes that the model architecture is LitCIFAR_ResNet and the GradCAM is implemented in GradCAM class.
  • The interface allows you to input an image and visualize its prediction along with the top classes and GradCAM output. You can also explore misclassified images from the test set.
  • The GradCAM output helps in understanding which parts of the image contributed most to the model's decision.

Steps to resuse the model

  1. 1.Load the pre-trained model using the provided LitCIFAR_ResNet class and load its weights from the model.pth file.
  1. 1.Use the inference function to predict the classes and visualize the GradCAM output for an input image.
  1. 1.To show misclassified images from the test set, use the showmisclassifiedimageswrap function, which internally uses the model's showmisclassified_images method.
  1. 1.Create a Gradio interface for image prediction and visualization of GradCAM output.
  1. 1.Create a separate Gradio interface for displaying misclassified images.
  1. 1.Create a tabbed Gradio interface that includes both image prediction and misclassified images.

Inferences

Model Test Accuracy: 90% Model Test Loss: 0.30

Outcome Gradio UI

input_image

misclassified_image