/
test_neuron_integrated_gradients.py
136 lines (120 loc) · 4.88 KB
/
test_neuron_integrated_gradients.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
#!/usr/bin/env python3
import unittest
import torch
from captum.attr._core.integrated_gradients import IntegratedGradients
from captum.attr._core.neuron.neuron_integrated_gradients import (
NeuronIntegratedGradients,
)
from ..helpers.basic_models import (
BasicModel_ConvNet,
BasicModel_MultiLayer,
BasicModel_MultiLayer_MultiInput,
)
from ..helpers.utils import assertArraysAlmostEqual, BaseTest
class Test(BaseTest):
def test_simple_ig_input_linear2(self):
net = BasicModel_MultiLayer()
inp = torch.tensor([[0.0, 100.0, 0.0]])
self._ig_input_test_assert(net, net.linear2, inp, 0, [0.0, 390.0, 0.0])
def test_simple_ig_input_linear1(self):
net = BasicModel_MultiLayer()
inp = torch.tensor([[0.0, 100.0, 0.0]], requires_grad=True)
self._ig_input_test_assert(net, net.linear1, inp, (0,), [0.0, 100.0, 0.0])
def test_simple_ig_input_relu(self):
net = BasicModel_MultiLayer()
inp = torch.tensor([[0.0, 6.0, 14.0]], requires_grad=True)
self._ig_input_test_assert(net, net.relu, inp, (0,), [0.0, 3.0, 7.0])
def test_simple_ig_input_relu2(self):
net = BasicModel_MultiLayer()
inp = torch.tensor([[0.0, 5.0, 4.0]])
self._ig_input_test_assert(net, net.relu, inp, 1, [0.0, 5.0, 4.0])
def test_simple_ig_multi_input_linear2(self):
net = BasicModel_MultiLayer_MultiInput()
inp1 = torch.tensor([[0.0, 10.0, 0.0]])
inp2 = torch.tensor([[0.0, 10.0, 0.0]])
inp3 = torch.tensor([[0.0, 5.0, 0.0]])
self._ig_input_test_assert(
net,
net.model.linear2,
(inp1, inp2, inp3),
(0,),
([[0.0, 156.0, 0.0]], [[0.0, 156.0, 0.0]], [[0.0, 78.0, 0.0]]),
(4,),
)
def test_simple_ig_multi_input_relu(self):
net = BasicModel_MultiLayer_MultiInput()
inp1 = torch.tensor([[0.0, 6.0, 14.0]])
inp2 = torch.tensor([[0.0, 6.0, 14.0]])
inp3 = torch.tensor([[0.0, 0.0, 0.0]])
self._ig_input_test_assert(
net,
net.model.relu,
(inp1, inp2),
(0,),
([[0.0, 1.5, 3.5]], [[0.0, 1.5, 3.5]]),
(inp3, 0.5),
)
def test_simple_ig_multi_input_relu_batch(self):
net = BasicModel_MultiLayer_MultiInput()
inp1 = torch.tensor([[0.0, 6.0, 14.0], [0.0, 80.0, 0.0]])
inp2 = torch.tensor([[0.0, 6.0, 14.0], [0.0, 20.0, 0.0]])
inp3 = torch.tensor([[0.0, 0.0, 0.0], [0.0, 20.0, 0.0]])
self._ig_input_test_assert(
net,
net.model.relu,
(inp1, inp2),
(0,),
([[0.0, 1.5, 3.5], [0.0, 40.0, 0.0]], [[0.0, 1.5, 3.5], [0.0, 10.0, 0.0]]),
(inp3, 0.5),
)
def test_matching_output_gradient(self):
net = BasicModel_ConvNet()
inp = 100 * torch.randn(2, 1, 10, 10, requires_grad=True)
baseline = 20 * torch.randn(2, 1, 10, 10, requires_grad=True)
self._ig_matching_test_assert(net, net.softmax, inp, baseline)
def _ig_input_test_assert(
self,
model,
target_layer,
test_input,
test_neuron,
expected_input_ig,
additional_input=None,
):
for internal_batch_size in [None, 1, 20]:
grad = NeuronIntegratedGradients(model, target_layer)
attributions = grad.attribute(
test_input,
test_neuron,
n_steps=200,
method="gausslegendre",
additional_forward_args=additional_input,
internal_batch_size=internal_batch_size,
)
if isinstance(expected_input_ig, tuple):
for i in range(len(expected_input_ig)):
for j in range(attributions[i].shape[0]):
assertArraysAlmostEqual(
attributions[i][j].squeeze(0).tolist(),
expected_input_ig[i][j],
delta=0.1,
)
else:
assertArraysAlmostEqual(
attributions.squeeze(0).tolist(), expected_input_ig, delta=0.1
)
def _ig_matching_test_assert(self, model, output_layer, test_input, baseline=None):
out = model(test_input)
input_attrib = IntegratedGradients(model)
ig_attrib = NeuronIntegratedGradients(model, output_layer)
for i in range(out.shape[1]):
ig_vals = input_attrib.attribute(test_input, target=i, baselines=baseline)
neuron_ig_vals = ig_attrib.attribute(test_input, (i,), baselines=baseline)
assertArraysAlmostEqual(
ig_vals.reshape(-1).tolist(),
neuron_ig_vals.reshape(-1).tolist(),
delta=0.001,
)
self.assertEqual(neuron_ig_vals.shape, test_input.shape)
if __name__ == "__main__":
unittest.main()