Pytorch Model Visual Interpretation | by Himanshu Sharma | Towards Data Science

Pytorch is an open-source python library that is used to create deep learning/machine learning models, it is one of the most used library after TensorFlow. It is based on the Torch library and is mostly used for Computer Vision and NLP applications.
There are different pre-trained PyTorch models that we can use according to the application we are working on. But the question is do we really understand what comes out of this black box. Model interpretation is important because we need to understand how the model works and how we can improve the model.
Captum is an open-source python library that is used to interpret PyTorch models. It creates transparency in terms of model understanding and interpretability. It helps understand the most important features that data have and how it contributes to the output of the model.
In this article, we will create a Captum insights dashboard for PyTorch model interpretation.
Let’s get started...
We will start by installing a flask_compress and Captum using pip. The command given below will do that.
!pip install flask_compress !pip install captum
In this step, we will import the required libraries for creating a PyTorch model and a Captum dashboard for the model interpretation. In this article, I will not be covering how to create a model.
import osimport torchimport torch.nn as nnimport torchvisionimport torchvision.transforms as transformsfrom captum.insights import AttributionVisualizer, Batchfrom captum.insights.attr_vis.features import ImageFeature
For this article, I will be using a pre-trained model for creating our model.
def get_classes():    classes = [        "Plane",        "Car",        "Bird",        "Cat",        "Deer",        "Dog",        "Frog",        "Horse",        "Ship",        "Truck",    ]    return classesdef get_pretrained_model():    class Net(nn.Module):        def __init__(self):            super(Net, self).__init__()            self.conv1 = nn.Conv2d(3, 6, 5)            self.pool1 = nn.MaxPool2d(2, 2)            self.pool2 = nn.MaxPool2d(2, 2)            self.conv2 = nn.Conv2d(6, 16, 5)            self.fc1 = nn.Linear(16 * 5 * 5, 120)            self.fc2 = nn.Linear(120, 84)            self.fc3 = nn.Linear(84, 10)            self.relu1 = nn.ReLU()            self.relu2 = nn.ReLU()            self.relu3 = nn.ReLU()            self.relu4 = nn.ReLU()def forward(self, x):            x = self.pool1(self.relu1(self.conv1(x)))            x = self.pool2(self.relu2(self.conv2(x)))            x = x.view(-1, 16 * 5 * 5)            x = self.relu3(self.fc1(x))            x = self.relu4(self.fc2(x))            x = self.fc3(x)            return xnet = Net()    net.load_state_dict(torch.load("cifar_torchvision.pt"))    return netdef baseline_func(input):    return input * 0def formatted_data_iter():    dataset = torchvision.datasets.CIFAR10(        root="data/test", train=False, download=True, transform=transforms.ToTensor()    )    dataloader = iter(        torch.utils.data.DataLoader(dataset, batch_size=4, shuffle=False, num_workers=2)    )    while True:        images, labels = next(dataloader)        yield Batch(inputs=images, labels=labels)
This is the final step where will create the dashboard for model understanding and interpretation.
normalize = transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))model = get_pretrained_model()visualizer = AttributionVisualizer(    models=[model],    score_func=lambda o: torch.nn.functional.softmax(o, 1),    classes=get_classes(),    features=[        ImageFeature(            "Photo",            baseline_transforms=[baseline_func],            input_transforms=[normalize],        )    ],    dataset=formatted_data_iter(),)#Visualize Modelvisualizer.render()
Here you can clearly visualize how the model is working, you can select the class that you want to interpret, similarly choose the attribution method and the arguments for that. It is clearly showing the attribution magnitude, labels, and prediction also.
Go ahead try this with different datasets and create beautiful Captum dashboards to interpret the model. In case you find any difficulty please let me know in the response section.
This article is in collaboration with Piyush Ingale.
Thanks for reading! If you want to get in touch with me, feel free to reach me at hmix13@gmail.com or my LinkedIn Profile. You can view my Github profile for different data science projects and packages tutorials. Also, feel free to explore my profile and read different articles I have written related to Data Science.