captum_integrated_gradients_linear_model_attribution_quickstart.py

1import torch
2from captum.attr import IntegratedGradients
3
4# A simple network with one linear layer
5class ToyModel(torch.nn.Module):
6    def __init__(self):
7        super().__init__()
8        self.lin = torch.nn.Linear(3, 1)
9        self.lin.weight = torch.nn.Parameter(torch.tensor([[3.0, 1.0, 0.0]]))
10        self.lin.bias = torch.nn.Parameter(torch.tensor([0.0]))
11
12    def forward(self, input):
13        return self.lin(input)
14
15# Initialize model and input tensor
16model = ToyModel()
17input = torch.tensor([[2.0, 5.0, 7.0]], requires_grad=True)
18
19# Initialize Integrated Gradients and compute attributions
20ig = IntegratedGradients(model)
21attributions, delta = ig.attribute(input, return_convergence_delta=True)
22
23print('IG Attributions:', attributions)
24print('Convergence Delta:', delta)