Captum: Interpret Predictions Of PyTorch Image Classification Networks¶

With the rise of deep learning, neural networks have become deeper and deeper with different types of layers (Convolution, Recurrent, LSTM, ConvLSTM, etc.) to transform data and find out patterns. Though this deep neural network gives the best results, especially for unstructured data (image, text, audio, etc), they are super hard to interpret. The predictions made by traditional ML models (decision trees, random forests, gradient boosting machines, etc) which are generally considered white-box models are fairly simple to interpret. Still, they do not perform well for unstructured data. Interpretability is a very important aspect of a network's performance evaluation as we want to know why our model is making a particular prediction and what parts of the data are contributing to that prediction. If our image classification network classifying cat vs dog is giving 95%+ accuracy then we want to know that it is using pixels of an image where either cat or dog are present to make predictions and not using some background pixels for predictions. Though interpreting predictions of deep neural networks (black-box models) is hard, over time, many algorithms have been developed to tackle the task.

As a part of this tutorial, we have explained a library called Captum which is specifically designed to interpret/explain predictions of neural networks designed using PyTorch. It provides an implementation for the majority of interpretation algorithms that have been invented till now. It'll keep adding more algorithms as they are introduced. In this tutorial, we have designed a simple convolutional neural network using PyTorch for image classification tasks. We'll use captum to explain predictions made by this network.

We have already covered the basics of captum and listed all available algorithms in a separate started tutorial. Please feel free to check it in your free time as it'll get you started with the library if you are new to it. It explains the library usage over tabular datasets.

• Captum: Interpret Predictions Of PyTorch Networks (Tabular Data)

Below, we have listed important sections of tutorial to give an overview of the material covered.

Important Sections Of Tutorial¶

1. Load Dataset
2. Define Network
3. Train Network
4. Evaluate Network Performance
5. Explain Network Predictions using CAPTUM
   1. Primary Attribution
      • Integrated Gradients
      • Gradient SHAP
      • Guided Grad-CAM
      • DeepLIFT
      • Kernel SHAP
   2. Layer Attribution
      • Grad-CAM
      • Layer DeepLIFT
      • Layer Integrated Gradients
   3. Neuron Attribution
      • Neuron Gradient
      • Neuron Integrated Gradients
      • Neuron Gradient SHAP

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

1. Load Dataset ¶

In this section, we have loaded Fashion MNIST dataset that we are going to use for our classification task. The dataset has grayscale images of shape (28,28) pixels for 10 different fashion items. The dataset is available from datasets module of torchvision Python library. The dataset is already divided into the train (60k images) and test (10k images) sets. After loading datasets, we have created data loaders from them which will be used to go through data in batches during the training process. The batch size of 64 is used. Below, we have included a table of mapping from target class index to target class name.

After loading datasets and creating data loaders, we have plotted a few sample images for reference purposes below.

2. Define Network ¶

In this section, we have defined a network that we'll use for our image classification task. The network consists of 2 convolution layers and one linear layer. The first convolution layer has 32 output channels and applies kernel of shape (3,3) on input data. The second convolution layer has 16 output channels and applies kernel of shape (3,3) on input data. The relu activation is applied to the output of both convolution layers. The output of the second convolution layer is given to the linear layer that has 10 output units. The output of the linear layer is a prediction of network. We have defined our network using Sequential API of PyTorch which is almost the same as Keras Sequential API.

After defining the network, we initialized it and performed a forward pass using a few examples for verification purposes.

If you are new to PyTorch and looking for guidance on creating neural networks then please feel free to check the below links. It'll get you started with the library.

• PyTorch: Simple Guide to Create Neural Networks using "torch.nn" Module
• PyTorch: Convolutional Neural Networks

3. Train Network ¶

In this section, we are training the network that we designed in the previous cell. In order to train the network, we have defined a function. The function takes model, loss function, optimizer, train data loader, validation data loader, and a number of epochs as input. The function executes training loop number of epochs times. For each epoch, it loops through training data in batches. For each batch of data, it performs a forward pass to make predictions, calculates loss, calculates gradients, and updates network parameters using gradients. It records loss for each batch and prints the average training loss at the end of each epoch. We have also created other helper functions that calculate validation loss and accuracy.

Below, we have trained our network using a function defined in the previous cell. We have first initialized a number of epochs to 8 and the learning rate to 0.001. Then, we have initialized the classification model, cross entropy loss, and Adam optimizer. At last, we have called our training routine with the necessary parameters to perform the training process. We can notice from the loss and accuracy values getting printed after each epoch that our network is doing a good job at the image classification task.

4. Evaluate Network Performance ¶

In this section, we have evaluated the performance of our trained network by calculating accuracy score, classification report (precision, recall, and f1-score per target class) and confusion matrix ML metrics on test predictions. We can notice from the accuracy score that the network is doing an above-average job at the task. We have calculated these metrics using functions available from scikit-learn.

Please feel free to check the below link if you are interested in learning about various ML metrics available from sklearn.

• Scikit-Learn - Model Evaluation & Scoring Metrics

Apart from calculations, we have also created a visualization of the confusion matrix using Python library scikit-plot. The visualization highlights that our network is not good at classifying images of target categories Pullover and Shirt. Many shirts are confused with T-shirts. Pullovers are confused with coats very often.

Scikit-plot provides visualization for many other ML metrics. Please check the below link if you are interested in it.

• Scikit-Plot: Visualizing Machine Learning Algorithm Results & Performance Metrics

5. Explain Network Predictions using CAPTUM ¶

Now, we'll explain predictions made by our network using Python library captum. It provides an implementation for many famous interpretation algorithms. Captum divides algorithms into 3 categories.

1. Primary Attribution - The algorithms in this category let us find out the contribution of our input data features towards prediction. (Input Features --> Prediction)
2. Layer Attribution - The algorithms in this category let us find out the contribution of activations of the selected layer towards prediction. (Layer Activations --> Prediction)
3. Neuron Attribution - The algorithms in this category let us find out the contributions of input data features towards the activation of a particular neuron. (Input Features --> Neuron Activation)

We'll explain the usage of a few algorithms from each category. All algorithms are available from 'attr' sub-module of captum. If the reader is interested in the theory of algorithms then we recommend going through the below link that covers them all in detail.

• Captum Algorithm

1. Primary Attribution¶

In this section, we'll explain a few algorithms from Primary Attribution category. Below, we have retrieved a few images and their target labels from the test data loader. Then, we made predictions on the first image using our trained model. We have printed the actual label and predicted label for reference purposes. Our network correctly predicts the target label as Ankle Boot. Next, we'll interpret this prediction and find out which features contributed to it.

Integrated Gradients¶

In this section, we have explained the usage of integrated gradients algorithm. We have created an instance of algorithm using IntegratedGradients() constructor available from 'attr' sub-module of captum. We have given our PyTorch network to the constructor. In order to generate an explanation for prediction, we need to call attribute() method. We have called the method with the first image and its target label. The output is an interpretation. Here, we have given only one grayscale image hence shape is (1,1,28,28) but you can give as many as you want. This returned image is a kind of mast that highlights pixels that contributed to predicting the target label as Ankle Boot.

Next, we have visualized original image and returned explanation using visualize_image_attr() function available from 'visualization' sub-module of 'attr'. The method uses Python library matplotlib to generate visualizations. We can notice from the explanation image that pixels present in the ankle boot section are contributing to prediction which is good. If this returns pixels other than object pixels (for example background pixels) then we can consider that our model has not generalized well.

Please make a NOTE that we have reshaped images before giving them to the visualization method as matplotlib requires channels at last and we had channels at the beginning.

Gradient SHAP¶

In this section, we have explained the usage of gradient shap algorithm. We have selected a different image this time. The selected image is Pullover and our network correctly predicts the target label. Our model has little low prediction probability for this prediction.

We have created an algorithm using the constructor GradientShap(). We have generated an explanation for the second image from our data and then visualized it. We can notice from the visualization that pixels inside of the pullover are contributing to prediction.

Guided Grad-CAM¶

In this section, we have explained the usage of Guide Grad-CAM algorithm. This algorithm requires reference to the convolution layer of the network and produces an explanation with respect to that layer.

We have created an algorithm using GuidedGradCam() constructor. We have provided reference to the second convolution layer. After generating an explanation, we have plotted it as usual. The visualization indicates that the second convolution layer seems to be detecting an edge of the pullover.

DeepLIFT¶

In this section, we have generated an explanation using DeepLIFT algorithm. The algorithm is created using DeepLift() constructor. The visualization shows that pixels in the pullover seem to be contributing to prediction.

Kernel SHAP¶

In this section, we have generated an explanation using kernel SHAP algorithm. The algorithm is initialized using KernelShap() constructor. The visualization shows that other than pullover pixels some background pixels are also contributing to prediction which is not good. This can happen because our model is already weak at classifying pullovers and it predicted our selected pullover correct with less probability.

2. Layer Attribution¶

In this section, we'll explain network predictions using Layer Attribution algorithms. The algorithms in this section require us to provide a reference layer from which to generate an explanation. All algorithms from this category start with the string 'Layer' to differentiate them from other categories.

Grad-CAM¶

In this section, we have explained the usage of Grad-CAM algorithm. We have instantiated algorithm using LayerGradCam() constructor. We have provided a classification model and second convolution layer to the constructor. Then, we have created an explanation using attribute() method and visualized it along with the original image for comparison purposes. We can notice from the visualization that the second convolution layer seems to detect the edges of images.

Below, we have generated an explanation with respect to the first convolution layer. We can notice from the visualization that it also seems to detect edges along with internal pixels.

Layer DeepLIFT¶

In this section, we have created an explanation using Layer DeepLIFT algorithm with respect to the first convolution layer. The visualization highlights that the results are almost the same as the previous Grad-CAM algorithm.

Layer Integrated Gradients¶

In this section, we have created an explanation using Layer integrated algorithms with respect to the first convolution layer. We can notice from the visualization that the second convolution layer is making a few edge pixels and a few middle pixels contribute towards prediction.

3. Neuron Attribution¶

In this section, we have explained how to use algorithms from Neuron Attribution category. The algorithms in this category require us to provide a reference layer and neuron index as it lets us know input data features contribute to the activation of that neuron. All algorithms from this category start with the string 'Neuron' to differentiate them from other categories.

Neuron Gradient¶

In this section, we have generated an explanation using Neuron gradient algorithm. We have initialized the algorithm using NeuronGradient() constructor. We have provided a classification network and the first convolution layer to the constructor. When calling attribute() method for algorithms in this section, we need to provide a neuron index. The output of second convolution layer has shape of (batch_size,32,28,28) hence per example there are (32,28,28) = (channels, height, width) activated neurons. We need to provide a tuple of 3 values to index a neuron. We have selected neuron at (0,15,15) which points at the first channel, 16th pixel in height dimension, and 16th pixel in the width dimension.

After generating an explanation, we have visualized it as well. We can notice from visualization the contribution of features towards the activation of that neuron.

Neuron Integrated Gradients¶

In this section, we have explained the usage of neuron integrated gradients. The algorithm is initiated using NeuronIntegratedGradients() constructor. The explanation is created with respect to neuron at position (0,15,15) in the first convolution layer. After generating an explanation, we have visualized it as well.

Neuron Gradient SHAP¶

In this section, we have explained the usage of neuron gradient SHAP algorithm. The algorithm is initiated using NeuronGradientShap() constructor. We have generated input data features contribution with respect to the neuron at position (15,15,15) from the second convolution layer. After generating an explanation, we have visualized it as well.

This ends our small tutorial explaining how we can use various interpretation algorithms available from Captum to interpret predictions of PyTorch Image classification networks.

References¶

LIME Tutorials¶

• How to Use LIME to Understand sklearn Models Predictions?
• LIME: Explain Keras Image Classification Network (CNN) Predictions

SHAP Values Tutorials¶

• SHAP - Explain Machine Learning Model Predictions using Game Theoretic Approach
• SHAP Values for Image Classification Tasks (Keras)
• Explain Flax (JAX) Image Classification Network Predictions using SHAP Values

Eli5 Tutorials¶

• How to Use eli5 to Understand sklearn Models, their Performance, and their Predictions?
• Eli5: Explain Image Classifier Predictions Using Grad-CAM (Keras)

Grad-CAM Implementations¶

• PyTorch: Grad-CAM
• Flax (JAX): Grad-CAM