|
| 1 | +{ |
| 2 | + "cells": [ |
| 3 | + { |
| 4 | + "cell_type": "markdown", |
| 5 | + "metadata": {}, |
| 6 | + "source": [ |
| 7 | + "# Interpreting Recurrent Models" |
| 8 | + ] |
| 9 | + }, |
| 10 | + { |
| 11 | + "cell_type": "markdown", |
| 12 | + "metadata": {}, |
| 13 | + "source": [ |
| 14 | + "This notebook demonstrates Captum attribution on an LSTM classifier whose input shape is `(batch, sequence_length, features)`. It shows how to compute input attributions, how to aggregate them by time step or by feature, and how to use layer attribution on the recurrent layer.\n", |
| 15 | + "\n", |
| 16 | + "**Note:** Before running this tutorial, please install matplotlib." |
| 17 | + ] |
| 18 | + }, |
| 19 | + { |
| 20 | + "cell_type": "code", |
| 21 | + "execution_count": null, |
| 22 | + "metadata": {}, |
| 23 | + "outputs": [], |
| 24 | + "source": [ |
| 25 | + "import matplotlib.pyplot as plt\n", |
| 26 | + "import torch\n", |
| 27 | + "import torch.nn as nn\n", |
| 28 | + "import torch.nn.functional as F\n", |
| 29 | + "from torch.utils.data import DataLoader, TensorDataset\n", |
| 30 | + "\n", |
| 31 | + "from captum.attr import IntegratedGradients, LayerIntegratedGradients" |
| 32 | + ] |
| 33 | + }, |
| 34 | + { |
| 35 | + "cell_type": "markdown", |
| 36 | + "metadata": {}, |
| 37 | + "source": [ |
| 38 | + "## Create a sequence classification task" |
| 39 | + ] |
| 40 | + }, |
| 41 | + { |
| 42 | + "cell_type": "markdown", |
| 43 | + "metadata": {}, |
| 44 | + "source": [ |
| 45 | + "The synthetic label depends on feature 0 near the end of the sequence, feature 1 near the beginning, and feature 2 in the middle. This gives us a known pattern to compare against the attribution summaries." |
| 46 | + ] |
| 47 | + }, |
| 48 | + { |
| 49 | + "cell_type": "code", |
| 50 | + "execution_count": null, |
| 51 | + "metadata": {}, |
| 52 | + "outputs": [], |
| 53 | + "source": [ |
| 54 | + "torch.manual_seed(7)\n", |
| 55 | + "\n", |
| 56 | + "def make_sequences(num_examples, seq_len=24, num_features=4):\n", |
| 57 | + " x = torch.randn(num_examples, seq_len, num_features)\n", |
| 58 | + " time_feature = torch.linspace(-1, 1, seq_len).view(1, seq_len)\n", |
| 59 | + " x[:, :, 3] = time_feature\n", |
| 60 | + "\n", |
| 61 | + " evidence = (\n", |
| 62 | + " x[:, -8:, 0].sum(dim=1)\n", |
| 63 | + " - x[:, :8, 1].sum(dim=1)\n", |
| 64 | + " + 0.5 * x[:, 8:16, 2].sum(dim=1)\n", |
| 65 | + " )\n", |
| 66 | + " y = (evidence > 0).long()\n", |
| 67 | + " return x, y\n", |
| 68 | + "\n", |
| 69 | + "train_x, train_y = make_sequences(2048)\n", |
| 70 | + "test_x, test_y = make_sequences(256)\n", |
| 71 | + "\n", |
| 72 | + "train_loader = DataLoader(\n", |
| 73 | + " TensorDataset(train_x, train_y),\n", |
| 74 | + " batch_size=64,\n", |
| 75 | + " shuffle=True,\n", |
| 76 | + ")" |
| 77 | + ] |
| 78 | + }, |
| 79 | + { |
| 80 | + "cell_type": "markdown", |
| 81 | + "metadata": {}, |
| 82 | + "source": [ |
| 83 | + "## Train a small LSTM classifier" |
| 84 | + ] |
| 85 | + }, |
| 86 | + { |
| 87 | + "cell_type": "code", |
| 88 | + "execution_count": null, |
| 89 | + "metadata": {}, |
| 90 | + "outputs": [], |
| 91 | + "source": [ |
| 92 | + "class LSTMClassifier(nn.Module):\n", |
| 93 | + " def __init__(self, input_size, hidden_size, output_size):\n", |
| 94 | + " super().__init__()\n", |
| 95 | + " self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True)\n", |
| 96 | + " self.fc = nn.Linear(hidden_size, output_size)\n", |
| 97 | + "\n", |
| 98 | + " def forward(self, x):\n", |
| 99 | + " output, _ = self.lstm(x)\n", |
| 100 | + " return self.fc(output[:, -1, :])\n", |
| 101 | + "\n", |
| 102 | + "model = LSTMClassifier(input_size=4, hidden_size=16, output_size=2)\n", |
| 103 | + "optimizer = torch.optim.Adam(model.parameters(), lr=0.01)\n", |
| 104 | + "\n", |
| 105 | + "for epoch in range(8):\n", |
| 106 | + " model.train()\n", |
| 107 | + " for batch_x, batch_y in train_loader:\n", |
| 108 | + " logits = model(batch_x)\n", |
| 109 | + " loss = F.cross_entropy(logits, batch_y)\n", |
| 110 | + " optimizer.zero_grad()\n", |
| 111 | + " loss.backward()\n", |
| 112 | + " optimizer.step()\n", |
| 113 | + "\n", |
| 114 | + "model.eval()\n", |
| 115 | + "with torch.no_grad():\n", |
| 116 | + " accuracy = (model(test_x).argmax(dim=1) == test_y).float().mean()\n", |
| 117 | + "\n", |
| 118 | + "print(\"Test accuracy:\", accuracy.item())" |
| 119 | + ] |
| 120 | + }, |
| 121 | + { |
| 122 | + "cell_type": "markdown", |
| 123 | + "metadata": {}, |
| 124 | + "source": [ |
| 125 | + "## Attribute a prediction to sequence inputs" |
| 126 | + ] |
| 127 | + }, |
| 128 | + { |
| 129 | + "cell_type": "markdown", |
| 130 | + "metadata": {}, |
| 131 | + "source": [ |
| 132 | + "`IntegratedGradients` returns attributions with the same shape as the input. For this model, that means one attribution score per `(time step, feature)` pair." |
| 133 | + ] |
| 134 | + }, |
| 135 | + { |
| 136 | + "cell_type": "code", |
| 137 | + "execution_count": null, |
| 138 | + "metadata": {}, |
| 139 | + "outputs": [], |
| 140 | + "source": [ |
| 141 | + "example = test_x[:1]\n", |
| 142 | + "baseline = torch.zeros_like(example)\n", |
| 143 | + "target = model(example).argmax(dim=1).item()\n", |
| 144 | + "\n", |
| 145 | + "ig = IntegratedGradients(model)\n", |
| 146 | + "\n", |
| 147 | + "# CuDNN RNN backward does not support eval mode on some GPU setups.\n", |
| 148 | + "with torch.backends.cudnn.flags(enabled=False):\n", |
| 149 | + " attributions, delta = ig.attribute(\n", |
| 150 | + " example,\n", |
| 151 | + " baselines=baseline,\n", |
| 152 | + " target=target,\n", |
| 153 | + " return_convergence_delta=True,\n", |
| 154 | + " )\n", |
| 155 | + "\n", |
| 156 | + "print(\"Attribution shape:\", tuple(attributions.shape))\n", |
| 157 | + "print(\"Convergence delta:\", delta.item())" |
| 158 | + ] |
| 159 | + }, |
| 160 | + { |
| 161 | + "cell_type": "markdown", |
| 162 | + "metadata": {}, |
| 163 | + "source": [ |
| 164 | + "Aggregate over the feature dimension to rank time steps, or aggregate over the time dimension to rank input features." |
| 165 | + ] |
| 166 | + }, |
| 167 | + { |
| 168 | + "cell_type": "code", |
| 169 | + "execution_count": null, |
| 170 | + "metadata": {}, |
| 171 | + "outputs": [], |
| 172 | + "source": [ |
| 173 | + "time_importance = attributions.abs().sum(dim=2).squeeze(0).detach()\n", |
| 174 | + "feature_importance = attributions.abs().sum(dim=1).squeeze(0).detach()\n", |
| 175 | + "\n", |
| 176 | + "fig, axes = plt.subplots(1, 2, figsize=(12, 4))\n", |
| 177 | + "\n", |
| 178 | + "axes[0].bar(range(example.shape[1]), time_importance)\n", |
| 179 | + "axes[0].set_xlabel(\"Time step\")\n", |
| 180 | + "axes[0].set_ylabel(\"Absolute attribution\")\n", |
| 181 | + "axes[0].set_title(\"Importance by time step\")\n", |
| 182 | + "\n", |
| 183 | + "axes[1].bar(range(example.shape[2]), feature_importance)\n", |
| 184 | + "axes[1].set_xlabel(\"Feature\")\n", |
| 185 | + "axes[1].set_title(\"Importance by feature\")\n", |
| 186 | + "\n", |
| 187 | + "plt.tight_layout()\n", |
| 188 | + "plt.show()" |
| 189 | + ] |
| 190 | + }, |
| 191 | + { |
| 192 | + "cell_type": "markdown", |
| 193 | + "metadata": {}, |
| 194 | + "source": [ |
| 195 | + "## Attribute through the recurrent layer" |
| 196 | + ] |
| 197 | + }, |
| 198 | + { |
| 199 | + "cell_type": "markdown", |
| 200 | + "metadata": {}, |
| 201 | + "source": [ |
| 202 | + "`LayerIntegratedGradients` can attribute to the recurrent layer inputs or outputs. Attributing to layer inputs keeps the result aligned with the original `(sequence_length, features)` axes." |
| 203 | + ] |
| 204 | + }, |
| 205 | + { |
| 206 | + "cell_type": "code", |
| 207 | + "execution_count": null, |
| 208 | + "metadata": {}, |
| 209 | + "outputs": [], |
| 210 | + "source": [ |
| 211 | + "lig = LayerIntegratedGradients(model, model.lstm)\n", |
| 212 | + "\n", |
| 213 | + "with torch.backends.cudnn.flags(enabled=False):\n", |
| 214 | + " layer_input_attr = lig.attribute(\n", |
| 215 | + " example,\n", |
| 216 | + " baselines=baseline,\n", |
| 217 | + " target=target,\n", |
| 218 | + " attribute_to_layer_input=True,\n", |
| 219 | + " )\n", |
| 220 | + "\n", |
| 221 | + "print(\"Layer input attribution shape:\", tuple(layer_input_attr.shape))" |
| 222 | + ] |
| 223 | + } |
| 224 | + ], |
| 225 | + "metadata": { |
| 226 | + "kernelspec": { |
| 227 | + "display_name": "python3", |
| 228 | + "language": "python", |
| 229 | + "name": "python3" |
| 230 | + }, |
| 231 | + "language_info": { |
| 232 | + "name": "python", |
| 233 | + "pygments_lexer": "ipython3" |
| 234 | + } |
| 235 | + }, |
| 236 | + "nbformat": 4, |
| 237 | + "nbformat_minor": 2 |
| 238 | +} |
0 commit comments