Interpretability in PyTorch, Integrated Gradient | Towards Data Science

Neural networks have been taking the world by storm. Not a week passes without great news about how GPT-3 supposedly automates yet another language task. Or how AI is helping doctors during the current pandemic.
But an increasing share of the population is becoming more and more skeptical about the turn AI is taking. Did it pick this treatment method because it is the right treatment method? Or because it happened to be the best one for that other 5.29-year-old patient in the training set?
Entire population groups warn us of the implications of bad designs in AI, and we should better listen. After all, racial and gender-related features have been used in the recent hiring process algorithms in discriminating ways!
How can we, as an AI-community,
Make sure we build inclusive tools.Make sure our models do what we think they are doing.Explain to others what our models are doing.
Make sure we build inclusive tools.
Make sure our models do what we think they are doing.
Explain to others what our models are doing.
The answer to 1,2 and 3 is clear Interpretability! It has been taking in the world of Deep Learning by storm. Let’s learn today what it is and how we can easily integrate it into our process.
We will build a Neural Network to predict the survival chance of a titanic passenger. And finally, understand what we did, and interpret our model using Captum. All using the mighty Lightning framework.
We will start by importing our dependencies.
# Our ML thingsimport pytorch_lightning as plimport torchfrom torch.utils.data import DataLoader, Datasetfrom captum.attr import IntegratedGradientsfrom pytorch_lightning import seed_everythingfrom pytorch_lightning import Trainer# Visualizationimport pandas as pdimport seaborn as snsimport matplotlib.pyplot as plt# Utilsfrom enum import Enumimport pandas as pdfrom sklearn.metrics import classification_report
To make this minimal example, we will work with the Titanic dataset (one small file). It can be downloaded here. Once you downloaded it, put it into a new folder ‘data/train.csv..’ This dataset is probably the most popular Kaggle dataset ever and therefore allows us to verify our results.
It contains the 890 different passengers of the Titanic. For each passenger, we are going to predict if they survived!
To do so, we are given several features.
PassengerId  Survived  Pclass  ...     Fare Cabin  Embarked0              1         0     ...   7.2500   NaN         S
We will only work with a small subset of these, namely:
Survived  Pclass     Sex   Age  SibSp     Fare0           0       3    male  22.0      1   7.2500
Now that we got an overview, let’s define the Dataset. A Dataset is a class from PyTorch. It needs to define 3 private functions, init(), len(), getitem().
If you are not familiar with this, you can find an easy explanation here (1. The Dataset). The goal here is to make a minimal preprocessing to show our interpretation later.
# Simple enum that helps organizing# Not really importantclass TrainType(Enum):    train = 1    test = 2# Our Custom Dataset classclass RMSTitanic(Dataset):    def __init__(self, type: TrainType):        file = "data/train.csv"        # Preprocessing        df = pd.read_csv(file)        # Select relevant fields        df = df[['Survived', 'Pclass',                 "Sex", "Age", "SibSp",                 "Fare"]]        # Convert Sex -> 0/1        m = {'male': 1, 'female': 0}        df['Sex'] = df['Sex'].str.lower().map(m)        # Fix the non available vals        # Also normalize        df = df.fillna(df.median())        df -= df.min()        df /= df.max()        # The first 80% are Train        if type == TrainType.train:            self.df = df.loc[:int(0.8 * len(df))]        if type == TrainType.test:            self.df = df.loc[int(0.8 * len(df)):]        # We will use this later for interpretation        self.base = torch.tensor(df[                                     ['Pclass', "Sex",                                      "Age", "SibSp",                                      "Fare"]                                 ].mean()).cuda().float()    def __len__(self):        return len(self.df)    def __getitem__(self, item):        # This function return the i-th elem        row = self.df.iloc[item]        label = row['Survived']        features = row[['Pclass', "Sex",                        "Age", "SibSp",                        "Fare"]]        # return the (label,features)        return (            torch.tensor(features).cuda().float(),            torch.tensor(label).cuda().float()        )
And that’s already our minimal Dataset. Keep in mind that this is only a small helper class. It helps Pytorch Lightning do its magic, so we save more code down the line.
The features we got and what their interpretation is.
Pclass   "What class is person traveling in? First being best"Sex      "Male or Female"Age      "How old in years?"SibSp    "How many Sibiling/Spouses are there onboard"Fare     "Amount in money spend"
We will now build a short but powerful model. It will call the following layers in order.
  | Name      | Type     |   ------------------------0 | input     | Linear   |   1 | r1        | PReLU    |  2 | l1        | Linear   |3 | r2        | PReLU    |    4 | out       | Linear   |   5 | sigmoid   | Sigmoid  |
We will now define a PytorchModel called MyHeartWillGoOn. Basically, only the forward function is crucial.
class MyHeartWillGoOn(pl.LightningModule):    def __init__(self):        # Setting up our model        super().__init__()        # way to fancy model        self.lr = 0.01        self.batch_size = 512        l1 = 128        # We send our 5 features into first layer        self.input = torch.nn.Linear(5, l1)        # PRELU is just a fancy activation function        self.r1 = torch.nn.PReLU()        # More Layers        self.l1 = torch.nn.Linear(l1, l1)        self.r2 = torch.nn.PReLU()        self.out = torch.nn.Linear(l1, 1)        # Befor the Output use a sigmoid        self.sigmoid = torch.nn.Sigmoid()        # Define loss        self.criterion = torch.nn.BCELoss()    def forward(self, x):        # Heart of our model        x = self.input(x)        x = self.l1(x)        x = self.r1(x)        x = self.out(x)        x = self.r2(x)        x = self.sigmoid(x)        return x    def train_dataloader(self):        # Load our Dataset: TRAIN        return DataLoader(            RMSTitanic(type=TrainType.train),            batch_size=self.batch_size,            shuffle=True)    def val_dataloader(self):        # Load our Dataset: TEST        # Simplification: TEST=VAL        return DataLoader(            RMSTitanic(type=TrainType.test),            batch_size=self.batch_size,            shuffle=False)    def configure_optimizers(self):        return torch.optim.Adam(self.parameters(), lr=self.lr)    def training_step(self, batch, batch_idx):        # Here we just log some basics        # We can look at them later in tensorboard        x, y = batch        y_hat = self(x)        y = torch.reshape(y, (y.size()[0], 1))        loss = self.criterion(y_hat, y)        tensorboard_logs = {'loss': loss}        return {'loss': loss, 'log': tensorboard_logs}
Training with PyTorch Lightning is done with a nice little Trainer.
if __name__ == '__main__':    # Seed so we can reproduce    seed_everything(404)    # Load model to GPU    device = torch.device("cuda")    model = MyHeartWillGoOn().to(device)    # Make the Trainer do the work    trainer = Trainer(max_epochs=20, min_epochs=1, auto_lr_find=True, progress_bar_refresh_rate=10)    trainer.fit(model)    # Accuracy #ToBeMeasured    ts = RMSTitanic(TrainType.test)    # Bit of Hacking    x = torch.stack(        [ts.__getitem__(i)[0] for i in range(ts.__len__())]    )    y = torch.stack(        [ts.__getitem__(i)[1] for i in range(ts.__len__())]    ).cpu().detach().numpy()    y_hat = model.forward(x).cpu().detach().numpy()    y_hat = (y_hat > 0.5)    # Accuracy and other metrics    print("REPORT:\n",          classification_report(y, y_hat))
Training our model on a GPU for 20 epochs will give us the following results.
SURVIVED   precision recall  f1-score   supportFALSE      0.84      0.94      0.89       115TRUE       0.86      0.67      0.75        64accuracy                               0.84       179macro avg          0.85      0.81      0.82       179weighted avg       0.85      0.84      0.84       179
As we can see, we get an easy 84% accuracy, which is acceptable for this task. We are more interested in interpreting this time. If you want to know how to get a perfect score on this task, have a look at Kaggle.
Since Pytorch Lightning is really new, and Captum is even newer ;) We are pretty cutting edge here.
We need to get around a few little problems, which yet have not been solved properly. To get this right, we need to add a smaller wrapper function to present later to Captum. Basically, Captum is a Library that has several methods available. These methods predict the importance of input features to Neural models. They do so in several ways. We will use the probably most popular method IntegratedGradient. I will give an intuition of this method below.
# Let's start with the interpretationSTEP_AMOUNT = 50SAMPLE_DIM = 5# Since Captum is not directly# made for Lightning,# we need this wrapperdef modified_f(in_vec):    # Changes the shapes correctly    # X:Shape([SAMPLE_DIM*STEP_AMOUNT]=250)    x = torch.reshape(in_vec,        (int(in_vec.size()[0] / SAMPLE_DIM), SAMPLE_DIM)    )    # RES:Shape([50,5])    # Since we have 50ti Batches and 5 features    res = model.forward(x)    # Again reshape to correct dims    res = torch.reshape(res, (res.size()[0], 1))    return resig = IntegratedGradients(modified_f)
Now we can start learning how our model thinks. Give this object IntegratedGradients. We can use its main method attribute()!
attribute() is a function that takes as an input a tensor and returns us a tensor of the same shape
This means that when we tell our model, our case study object:
Pclass  Sex   Age  SibSp     Fare3       male  22.0      1       7.2500 (I am not normalized)
It will make a prediction using these 5 features. Let’s say 0.3, which means 0.3% survival chance, for this 22-year-old man paying 7.25 in the fare. After predicting, we will send this 30% Survival rate ->0 %, meaning he died.
Now Integrated gradient returns us a tensor, also having 5 values. Each of these values will tell us how important the respective feature is. This may be a difference in y. E.g.
Pclass  Sex   Age  SibSp     Fare0.1619, -0.1594,  0.0196, -0.0024,  0.0068
A little theory lesson in the middle. Integrated Gradients has been proposed in “Axiomatic Attribution for Deep Networks.” Integrated Gradients is used to understand how our Neural Network works. Integrated Gradients is a so-called interpretability algorithm.
There are several interpretability algorithms used in today's AI community, but IG is one of the first and most established ones. The scientific community has not clearly decided which algorithm is the best. Feel free to test some of the other ones, like DeepLift.
This is a simplified explanation, please refer to the paper for more details “Axiomatic Attribution for Deep Networks”
To optimize a loss function in a Neural Network fashion, we build a gradient. To understand how much each input feature contributes to our loss function, we reverse this process (called Integration), so we Integrate the Gradient.
Now, calculating the Integral is hard, and the answer is not clear. This is why we have to approximate it. We approximate this Integral using a Riemann Sum.
We approximate the value by using a hyperparameter n_steps and a baseline. The baseline is what we compare our input to (we look at this in detail below). The n_steps is the number of times we backpropagate to estimate the integral. N_steps basically controls how precise our estimate of the integral is.
We now have 3 things:
A Trained modelA Test setIntegrated Gradient
A Trained model
A Test set
Integrated Gradient
Using these things, we will understand our model
# Test to understand the basics# First get the 6th test exampleval = ts.__getitem__(6)[0]print("IN X1:", val)# Predict the importance of the features# for the male exampleimp_m = ig.attribute(inputs=val,                     baselines=ts.base,                     n_steps=STEP_AMOUNT)print("IMPORTANCE X_m:", imp_m)print("Probability SURVIVAl X_m:", modified_f(val))# Predict the importance of the features# for the female exampleprint("Let's Change the gender ->X2")val_f = valval_f[1] = 0imp_f = ig.attribute(inputs=val_f,                     baselines=ts.base,                     n_steps=STEP_AMOUNT)print("IMPORTANCE X_f:", imp_f)print("Probability SURVIVAl X_f:", modified_f(val_f))
This code will print:
IN X1: [1.0000, 1.0000, 0.3466, 0.0000, 0.0303]IMPORTANCE X_m:[-0.0593, -0.1121, -0.0010,  0.0019, -0.0040]Probability SURVIVAl X_m:0.1265IMPORTANCE X_f:[-0.1178,  0.4049, -0.0020,  0.0037, -0.0080]Probability SURVIVAl X_f:0.5817
As you can see, when we give Captum, these two identical examples. With the exception of Gender. We get drastically different survival rates.
# MALEPclass  Sex   Age  SibSp     Fare3       male  22.0      1       7.2500 (I am not normalized)# FEMALEPclass  Sex     Age  SibSp     Fare3       female  22.0      1       7.2500 (I am not normalized)
And the Survival guess for the male is 12% and 58% for the female! The two Importance values we got are:
                  Pclass  Sex         Age    SibSp     Fare  IMPORTANCE X_m:[-0.0593, -0.1121, -0.0010,  0.0019, -0.0040]IMPORTANCE X_f:[-0.1178,  0.4049, -0.0020,  0.0037, -0.0080]
As we can see in both cases, the importance is the strongest (in absolute terms) in the feature SEX. This indicates that this is the most important feature!
Interpreting the sign: We can think of it as something that adds to the prediction (+) or deducted from the prediction (-). I.e. + makes it more likely to be (1) or less likely to be (1).
We can see that depending on the Sex(Male\Female), the resulting vector is quite different.
The baseline is what we compare our value. In order to understand its impact, we will compare the most obvious choices. In the initial paper, the ideas of noise and black images (zeros) were used.
We will compare 4 different baselines.
Compare to:
1.the average!
2. noise, compared to random
3. all ones
4. all zeros
To following, code will produce the desired interpretation.
# define a collectionto_be_df = []# Compare each element of the test set to out baselines# we will than use thisfor i in range(0,               1):  # ts.__len__()):    # load our test example    in_val = ts.__getitem__(i)[0]    # compare it to the 4 baselines    att_b = ig.attribute(        inputs=in_val,        baselines=ts.base,        n_steps=STEP_AMOUNT).detach().cpu().numpy()    att_r = ig.attribute(        inputs=in_val,        baselines=torch.rand(            5).cuda(),        n_steps=STEP_AMOUNT).detach().cpu().numpy()    att_z = ig.attribute(        inputs=in_val,        baselines=torch.zeros(            5).cuda(),        n_steps=STEP_AMOUNT).detach().cpu().numpy()    att_1 = ig.attribute(        inputs=in_val,        baselines=torch.ones(            5).cuda(),        n_steps=STEP_AMOUNT).detach().cpu().numpy()    # save result, this will produce a df    # you can skip the details    for base_type, vals in [        ("mean-base", att_b),        ("random-base", att_r),        ("zero-base", att_z),        ('one-base', att_1),    ]:        for i, name in enumerate(['Pclass',                                  "Sex",                                  "Age",                                  "SibSp",                                  "Fare"]):            to_be_df.append({                "base-type": base_type,                "feature": name,                "value": vals[i],            })# Convert our data to a pandasdf = pd.DataFrame(to_be_df)df.to_csv('data/interpretation_results.csv')print("OUR INTERPRETATION:\n\n",      df)
Our results look something like this.
      base-type feature     value0     mean-base  Pclass  0.1620041     mean-base     Sex -0.1594212     mean-base     Age  0.0196523     mean-base   SibSp -0.0024334     mean-base    Fare  0.0068215   random-base  Pclass  0.117517
Now we found what we are after!
Once we have done the work, we can sit back and enjoy it!
# Aggregate and Visualize# Load Datadf = pd.read_csv('data/interpretation_results.csv')# Defined the color map for our heatmap to be red to greencmap = sns.diverging_palette(h_neg=10, h_pos=130, s=99,                             l=55, sep=3, as_cmap=True)# Aggregate the CSV by meandf = df.groupby(["base-type", 'feature', 'epoch'],                as_index=False).mean()df = df[df['epoch'] == max_epoch]# Make one plot per baseline to comparefor b in ["mean-base",          "random-base",          "zero-base",          'one-base']:    # Let's plot them isolated    tmp = df[df['base-type'] == b]    # Create a pivot frame    tmp = tmp.pivot(index='base-type', columns='feature',                    values='value')    print("We will plot:\n",tmp)
A pivot data frame is basically just taking a different view on the same data. Here we tell it to have the index (the rows on the base-type) and the columns on the ‘features.’ This looks something like this.
feature         Age      Fare    Pclass       Sex     SibSpmean-base  0.023202  0.004757  0.116368  0.225421  0.022177
Now all that is left to do is plot it using
# Create a pivot frame tmp = tmp.pivot(index='base-type', columns='feature',                    values='value')# Some code to make a heatmap using seabornfig, ax = plt.subplots()fig.set_size_inches(10, 2.5)plt.title("Feature Importance ", fontsize=15)sns.heatmap(tmp, ax=ax, cmap=cmap, annot=True)plt.text(0, 0,         'By Sandro Luck\nTwitter:@san_sluck',         horizontalalignment='center',         verticalalignment='bottom',         fontsize=10)plt.savefig(f'data/{b}.png')
Great, now I hope also the last skeptic is convinced Sex is Great! We can clearly see that the impact of the feature Sex is significant. When we look at the dataset, we could notice that being a man increases your chance of dying significantly. “Women and children first” is the way they go with ships, after all.
Here we see quite a difference between the baselines. If we compare the two baselines with the Mean-Base, we notice that the importance of SibSp changes drastically. This generally should illustrate that using baselines such as 1 and 0 is not advisable. The reason for this is that we calculate how much the prediction changes between baseline and value.
For Images for example using a 0-baseline might make sense. The reason is that a black square represents in this case the absence of information.
The surprise of the day comes for me in the form of the success of Random baselines. The authors of Integrated Gradients mentioned that using a random baseline is advisable. But i did not anticipate it working so well for this example.
GIFs are crucial when making people pay attention to your data on certain platforms. Especially on social media, animation can make the difference between a click and a quick glimpse.
Since this article is already pretty long, please refer to my article “How to make GIFS with Python and Seaborn From Google Trends Data” to learn at lightning speed to create the GIF above.
We have learned what Interpretability is and how we can use it in PyTorch. We have seen how easy we can use it in our Prototyping and how to connect it to PyTorch Lightning.
I hope this introduction to the world of Interpretability was easy and painless. Understanding what you are building can help our customers and us. It increases their confidence in our predictions significantly.
If you enjoyed this article, I would be excited to connect on Twitter or LinkedIn.
Make sure to check out my YouTube channel, where I will be publishing new videos every week.