Explaining Custom CNN CIFAR-10 Classification Using the Attributions Explainer

import matplotlib.pyplot as plt
import matplotlib
import numpy as np

%matplotlib inline

# PyTorch
import torch
import torchvision
import torchvision.transforms as transforms
import torchvision.transforms.functional as TF
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models
import torch.optim as optim

transform = transforms.Compose(
    [transforms.ToTensor(),
     transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])

trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
                                        download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False,
                                       download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                         shuffle=False, num_workers=2)

classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

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 x


net = Net()

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

USE_PRETRAINED_MODEL = False # if the model is not saved, set false

if USE_PRETRAINED_MODEL:
    print("Using existing trained model")
    from urllib.request import urlopen
    import os.path
    if os.path.isfile("models/cifar_torchvision.pt"):
        print("File found, will be loaded")
        net.load_state_dict(torch.load('models/cifar_torchvision.pt'))
    else:
        print("Please train the model first by setting USE_PRETRAINED_MODEL to False")
else:
    for epoch in range(1):  # loop over the dataset multiple times

        running_loss = 0.0
        for i, data in enumerate(trainloader, 0):
            # get the inputs
            inputs, labels = data
            # zero the parameter gradients
            optimizer.zero_grad()

            # forward + backward + optimize
            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            # print statistics
            running_loss += loss.item()
            if i % 2000 == 1999:    # print every 2000 mini-batches
                print('[%d, %5d] loss: %.3f' %
                      (epoch + 1, i + 1, running_loss / 2000))
                running_loss = 0.0

    print('Finished Training')
    torch.save(net.state_dict(), 'cifar_torchvision.pt')

def imshow(img, transpose = True):
    img = img / 2 + 0.5     # unnormalize
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()

dataiter = iter(testloader)
images, labels = next(dataiter)

# print images
imshow(torchvision.utils.make_grid(images))
print('GroundTruth: ', ' '.join('%5s' % classes[labels[j]] for j in range(4)))


outputs = net(images)

_, predicted = torch.max(outputs, 1)

print('Predicted: ', ' '.join('%5s' % classes[predicted[j]]
                              for j in range(4)))

ind = 3

input = images[ind].unsqueeze(0)
input.requires_grad = True

net.eval()

from intel_ai_safety.explainer.attributions import pt_attributions as attributions
from captum.attr import visualization as viz

# handeling Original Image
print('Original Image')
print('Predicted:', classes[predicted[ind]], ' Probability:', torch.max(F.softmax(outputs, 1)).item())
original_image = np.transpose((images[ind].cpu().detach().numpy() / 2) + 0.5, (1, 2, 0))
viz.visualize_image_attr(None, original_image, method="original_image", title="Original Image")

# Entry Points
attributions.saliency(net).visualize(input,labels[ind],original_image,"Saliency")
attributions.integratedgradients(net).visualize(input,labels[ind],original_image,"Integrated Gradients")
attributions.deeplift(net).visualize(input,labels[ind],original_image,"Deep Lift")
attributions.smoothgrad(net).visualize(input,labels[ind],original_image,"Smooth Grad")
attributions.featureablation(net).visualize(input,labels[ind],original_image,"Feature Ablation")