captum_integrated_gradients_attribution_simple_neural_network.py

1import torch
2import torch.nn as nn
3
4from captum.attr import IntegratedGradients
5
6class ToyModel(nn.Module):
7    def __init__(self):
8        super().__init__()
9        self.lin1 = nn.Linear(3, 3)
10        self.relu = nn.ReLU()
11        self.lin2 = nn.Linear(3, 2)
12
13    def forward(self, input):
14        return self.lin2(self.relu(self.lin1(input)))
15
16net = ToyModel()
17input = torch.randn(2, 3, requires_grad=True)
18baseline = torch.zeros(2, 3)
19
20ig = IntegratedGradients(net)
21attributions, delta = ig.attribute(input, baseline, target=0, return_convergence_delta=True)
22
23print('IG Attributions:', attributions)
24print('Convergence Delta:', delta)