Skip to content

Commit 1e8b4cc

Browse files
craymichaelfacebook-github-bot
authored andcommitted
Add recurrent model attribution tutorial (#1205) (#1834)
Summary: Fixes #1205. Issue: #1205 Issue summary: Captum lacked a recurrent-model tutorial showing sequence and recurrent-layer attribution. - Add an LSTM sequence-classification tutorial with input attribution and time/feature aggregation. - Include LayerIntegratedGradients on the recurrent layer with attribute_to_layer_input=True. Test Plan: - venv/uv/bin/python -m json.tool tutorials/Recurrent_Model_Interpret.ipynb - Pyre: not applicable; tutorial-only change. Reviewed By: Ayush-Warikoo Differential Revision: D106053939 Pulled By: craymichael
1 parent 00fb699 commit 1e8b4cc

2 files changed

Lines changed: 242 additions & 0 deletions

File tree

Lines changed: 238 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,238 @@
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+
}

website/tutorials.json

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,10 @@
3030
"id": "House_Prices_Regression_Interpret",
3131
"title": "Interpreting a regression model of California house prices"
3232
},
33+
{
34+
"id": "Recurrent_Model_Interpret",
35+
"title": "Interpreting recurrent models"
36+
},
3337
{
3438
"id": "Segmentation_Interpret",
3539
"title": "Interpreting semantic segmentation models"

0 commit comments

Comments
 (0)