Improving CNN performance using GradCAM (Collab)

Motivtion

Doubt: Unexpectedly high classical performance metrics (AUC, F1, Acc) of compact CNNs (< 150M parameters) on small (~5k X-rays) Pneumonia Dataset.
Hypothesis: Possible model overfitting on unwanted features rather than the target pathology, Pneumonia.

Problem Statement

Improving Convolutional Neural Network classification and feature localisation for Pneumonia Chest X-rays in the scenario of a lack of extensive annotated data and access to extensive GPU Training architectures.

Approach

• Developed a model-independent, self-sufficient and recursive optimisation framework for CNN models utilising Gradient Class Activation Maps (GradCAMs), to tackle challenges of limited data and suboptimal localisation by iteratively augmenting and re-using the original dataset.

  • GradCAM, like heatmaps, are used to visualise regions of the input image which strongly influence the model's classification output.
  • Trained a compact CNN (~100M parameters) on a small dataset (~4.5k X-rays) that achieved high accuracy and F1 score, but showed poor pathology localisation when analysed using Grad-CAMs.
  • Iteratively enhanced the dataset by overlaying Grad-CAMs from high-confidence predictions onto input images to highlight pneumonia-relevant regions, creating a refined dataset for model fine-tuning.
  • Repeated the process to progressively improve localisation, producing models that better identified disease-relevant regions with each cycle.

  Image depicting one aforementioned cycle

 

INSTRUCTIONS

Details

General

• Code: Pytorch | Data-set: paultimothymooney.
• After opening the notebook in Collab, please go to File > Save as copy in Google Drive to experiment with the code after reading the Data Handling section below.

Data Handling

• .zip file containing the Pneumonia Dataset must be uploaded to the My Drive folder of the Google Drive mounted to the collab notebook.
• .zip file is** automatically extracted** to drive\My Drive\ML Data Sets while removing any corrupted images.
• To use external data-sets, place them in drive\My Drive\ML Data Sets and change variable names.

CNN Models

• Baseline model architecture: TinyVGG
• Links to .pth weights of the trained models can be found in the models folder.
• popular pre-trained architectures (VGG-16, VGG-19, ResNet-50, ResNet-101, ResNet-152) can be used and corresponding models loaded.
• Model, optimizer & scheduler state_dicts can be saved to Google Drive (drive\MyDrive\ML Models\<xyz.pth>). These can be loaded in future for seamless resumption of progress.
  

mismatch in model parameters while loading models can be a result of modified architecture in the code

---

 

RESULTS FOR VARIOUS ARCHITECTURES

Details

Initial testing for Pneumonia localisation

• Implemented GradCAM functionality to check Pneumonia localisation of baseline and pre-trained models.
• Identified and optimised the optimal model architecture.

Different Architectures

TinyVGG CNN_GradCAM_v1.pth

• Test Loss: 0.305 | Test Acc: 88.28% | Sensitivity: 91.79% | Specificity: 81.62%

• Problem: Focus on unwanted features, extremely poor localisation of Pneumonia. Heavy overfitting.

  
  Pneumonia Class | Good here "Good" means prediction score >= 85%
  blank boxes indicate wrong predictions

VGG-16 VGG-16_GradCAM_v2.pth

• Test Loss: 0.362 | Test Acc: 83.91% | Sensitivity: 88.72% | Specificity: 75.07%

• Progress: Improved but unsatisfactory localisation.

  
  

ResNet-101 ResNet-101_GradCAM_v1.pth

• Test loss: 0.369 | Test acc: 83.59% | Sensitivity: 95.38% | Specificity: 62.82%

• Progress: Improved localisation after adding FC layer with 10k params.

  
  

ResNet-152 ResNet-152_GradCAM_v1.pth

• Test Loss: 0.379 | Test acc: 84.69% | Sensitivity: 93.33% | Specificity: 70.09%

• Progress: Similar/worse localisation compared to ResNet-101 with similar evaluation metrics.

  

Best Overall Localisation and Performance: ResNet-101 ResNet-101_GradCAM_v1.pth

• There is an unwanted focus on void regions surrounding the skeleton.

• To fix this, a small centre crop was added to training images before feeding them to the model, thus preventing the model from "learning" these void regions.

  
  

Optimizing ResNet-101_GradCAM_v1.pth

• Finally, ~17M parameters were re-trained** from scratch (~38% of ResNet-101) using above cropped train data** to achieve the current "best" model: ResNet-101_GradCAM_Cropped_v4.pth

• Test Loss: 0.264 | Test Acc: 91.56% | Sensitivity: 94.62% | Specificity: 85.90%

• Progress: Obtained noticeable improvement in Pneumonia localisation and evaluation metrics.
  The model is now ready to create an enhanced data set.

  
  GradCAMs showing improved localisation

---

 

ENHANCED DATASET RESULTS

Details

Enhanced Data Set

• Train Data enhanced by overlaying accurate GarCAMs.

  • "enhanced" - important features emphasised (in "white") while the rest of the image is suppressed.

  • "accurate" - The assumption is that GradCAMs localised to the lung are fairly accurate.

    
    Image showing 49 images from the enhanced train data
    post data transformations

Results using Enhanced Data Set ResNet-101_GradCAM_Cropped_v4_enh_v1.pth

• Re-trained previously "un-frozen" ~17M ResNet-101_GradCAM_Cropped_v4.pth params using above enhanced data set.

  • One cycle resulted in a ~15% improvement in the localization of pneumonia, with only a ~3% decrease in test accuracy with respect to ResNet-101_GradCAM_Cropped_v4.pth.

  • Test Loss: 0.397 | Test acc: 88.59% | Sensitivity: 86.41% | Specificity: 91.88%

    
    Improved GradCAMs post fine-tuning using enhanced data

---

 

CONCLUSION

Details

Summary

• We can successfully develop a workaround for the limited data constraint to iteratively obtain better localisation & identification of pneumonia.
• During this process, as expected, test accuracy takes a small hit as the model diverges away from unwanted "easy" to detect features.

Results on annotated images

• Model was not trained on annotated data; the presence of arrows pointing to Pneumonia in the X-rays makes no difference in the model choice

• Covid-19 Pneumonia X-ray
  
  

• Pneumonia X-ray
  
  

• Pneumonia X-ray
  
  

Post and Pre-Enhanced Data Fine-Tuning Comparison

• Post-fine-tuning the model on enhanced data, we observe significantly better localisation of Pneumonia to the lungs.

• The model also relies less on unwanted features and void regions while predicting.

• The results are limited to only one cycle shown in the schematic above; further process iterations are bound to yield gradual but better results.

  
  Final Comparison on the same 100 test images

---

 

REFERENCES

• GradCAM - JimEverest
• Lung Segmentation - IlliaOvcharenko's repo