Feature request: Assign rule by layer index

[rodrigobdz]

I'd like to assign LRP-rules by layer index, like shown in the screenshots. Please correct me if this is already possible, I've taken a look at the code and the paper but it seems currently it's only possible by type.

Source: gmontavon/lrp-tutorial

Source: Layer-Wise Relevance Propagation: An Overview

Thanks for the great framework! I especially like its architecture.

[chr5tphr]

A note about the code from the tutorial: The last line if l >= 31: will actually never happen for vgg16, since there are no AvgPool2d or Conv2d layers beyond index 30. Also, the LRP-0 above is more like the conventional LRP-Epsilon rule, and the LRP-Epsilon rule above depends on the output, and as such is not really the conventional LRP-Epsilon rule

While it would be straight forward to write a composite which does take the layer index into account, I would suggest to rather go with the NameMapComposite here, since the rules are very model-dependent and it's much more transparent than any index-based approach.

If you want to automatically create a name-map as the composite from the tutorial/ paper, here's some code on how to do it:

Click to unfold code

```python import torch from torch.nn import Conv2d, AvgPool2d from torchvision.models import vgg16 from zennit.composites import NameMapComposite from zennit.core import BasicHook, collect_leaves, stabilize from zennit.rules import Gamma, Epsilon # the LRP-Epsilon from the tutorial class GMontavonEpsilon(BasicHook): def __init__(self, epsilon=1e-6, delta=0.25): super().__init__( input_modifiers=[lambda input: input], param_modifiers=[lambda param, _: param], output_modifiers=[lambda output: output], gradient_mapper=(lambda out_grad, outputs: out_grad / stabilize(outputs[0] + delta * (outputs[0] ** 2).mean() ** .5, epsilon)), reducer=(lambda inputs, gradients: inputs[0] * gradients[0]) ) model = vgg16() # only these get rules, linear layers will be attributed by the gradient alone target_types = (Conv2d, AvgPool2d) # lookup module -> name child_name = {module: name for name, module in model.named_modules()} # the layers in sequential order without any containers etc. layers = list(enumerate(collect_leaves(model))) # list of tuples [([names..], rule)] as used by NameMapComposite name_map = [ ([child_name[module] for n, module in layers if n <= 16 and isinstance(module, target_types)], Gamma(0.25)), ([child_name[module] for n, module in layers if 17 <= n <= 30 and isinstance(module, target_types)], GMontavonEpsilon(1e-9, 0.25)), ([child_name[module] for n, module in layers if 30 <= n and isinstance(module, target_types)], Epsilon(1e-9)), ] # look at the name_map and you will see that there is no layer for which the last condition holds print(name_map) # create the composite from the name map composite = NameMapComposite(name_map) with composite.context(model) as modified_model: # compute attribution data = torch.randn(1, 3, 224, 224, requires_grad=True) output = modified_model(data) output.backward(torch.eye(1000)[[0]]) # print absolute sum of attribution print(data.grad.abs().sum().item()) ```

Note that doing model.named_modules() alone will give you all modules, eg. Sequential, and thus not count the layers correctly for vgg16. However, when manually constructing a name-map, this will show you the names of all layers.

[rodrigobdz]

@chr5tphr Awesome! Thank you for your time, the snippet has been of great help.

---

Clarification about the tutorial:

• The case if l >= 31: is indeed evaluated because dense layers in the classifier are converted to convolutional layers in this LOC:

  layers = list(model.features) + utils.toconv(list(model.classifier))

---

With your help, I've been able to reproduce the results from the LRP tutorial; here are the changes I made:

1. Set data to castle image, then I plotted the heatmap but it lacked the attributions from the classifier layers and the ZBox rule for the pixel layer.

   

   Code diff

   ```diff import torch from torch.nn import Conv2d, AvgPool2d from torchvision.models import vgg16 from zennit.composites import NameMapComposite from zennit.core import BasicHook, collect_leaves, stabilize from zennit.rules import Gamma, Epsilon +import cv2 +import numpy +import utils # the LRP-Epsilon from the tutorial class GMontavonEpsilon(BasicHook): def __init__(self, epsilon=1e-6, delta=0.25): super().__init__( input_modifiers=[lambda input: input], param_modifiers=[lambda param, _: param], output_modifiers=[lambda output: output], gradient_mapper=(lambda out_grad, outputs: out_grad / stabilize(outputs[0] + delta * (outputs[0] ** 2).mean() ** .5, epsilon)), reducer=(lambda inputs, gradients: inputs[0] * gradients[0]) ) model = vgg16() +class BatchNormalize: + def __init__(self, mean, std, device=None): + self.mean = torch.tensor(mean, device=device)[None, :, None, None] + self.std = torch.tensor(std, device=device)[None, :, None, None] + + def __call__(self, tensor): + return (tensor - self.mean) / self.std + + +# mean and std of ILSVRC2012 as computed for the torchvision models +norm_fn = BatchNormalize((0.485, 0.456, 0.406), + (0.229, 0.224, 0.225), device='cpu') + # only these get rules, linear layers will be attributed by the gradient alone target_types = (Conv2d, AvgPool2d) # lookup module -> name child_name = {module: name for name, module in model.named_modules()} # the layers in sequential order without any containers etc. layers = list(enumerate(collect_leaves(model))) # list of tuples [([names..], rule)] as used by NameMapComposite name_map = [ ([child_name[module] for n, module in layers if n <= 16 and isinstance(module, target_types)], Gamma(0.25)), ([child_name[module] for n, module in layers if 17 <= n <= 30 and isinstance(module, target_types)], GMontavonEpsilon(1e-9, 0.25)), ([child_name[module] for n, module in layers if 30 <= n and isinstance(module, target_types)], Epsilon(1e-9)), ] # look at the name_map and you will see that there is no layer for which the last condition holds print(name_map) # create the composite from the name map composite = NameMapComposite(name_map) +R = None with composite.context(model) as modified_model: # compute attribution - data = torch.randn(1, 3, 224, 224, requires_grad=True) + # Returns a numpy array in BGR color space, not RGB + img = cv2.imread('castle.jpg') + + # Convert from BGR to RGB color space + img = img[..., ::-1] + + # img.shape is (224, 224, 3), where 3 corresponds to RGB channels + # Divide by 255 (max. RGB value) to normalize pixel values to [0,1] + img = img/255.0 + + data = norm_fn( + torch.FloatTensor( + img[numpy.newaxis].transpose([0, 3, 1, 2])*1 + ) + ) + data.requires_grad = True + output = modified_model(data) - output.backward(torch.eye(1000)[[0]]) + output[0].max().backward() + # print absolute sum of attribution print(data.grad.abs().sum().item()) + + R = data.grad + + utils.heatmap(R[0].sum(dim=0).detach().numpy(), 4,4) ```

   

2. Finally, added Linear to target_types, fixed LRP-0 rule, and assigned ZBox-rule to pixel layer.

   Code diff

   ```diff import torch -from torch.nn import Conv2d, AvgPool2d +from torch.nn import Conv2d, AvgPool2d, Linear from torchvision.models import vgg16 from zennit.composites import NameMapComposite from zennit.core import BasicHook, collect_leaves, stabilize -from zennit.rules import Gamma, Epsilon +from zennit.rules import Gamma, Epsilon, ZBox import cv2 import numpy import utils # the LRP-Epsilon from the tutorial class GMontavonEpsilon(BasicHook): - def __init__(self, epsilon=1e-6, delta=0.25): + def __init__(self, stabilize_epsilon=1e-6, epsilon=0.25): super().__init__( input_modifiers=[lambda input: input], param_modifiers=[lambda param, _: param], output_modifiers=[lambda output: output], - gradient_mapper=(lambda out_grad, outputs: out_grad / stabilize(outputs[0] + delta * (outputs[0] ** 2).mean() ** .5, epsilon)), + gradient_mapper=(lambda out_grad, outputs: out_grad / stabilize( + outputs[0] + epsilon * (outputs[0] ** 2).mean() ** .5, stabilize_epsilon)), reducer=(lambda inputs, gradients: inputs[0] * gradients[0]) ) -model = vgg16() + +# use the gpu if requested and available, else use the cpu +device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') + class BatchNormalize: def __init__(self, mean, std, device=None): self.mean = torch.tensor(mean, device=device)[None, :, None, None] self.std = torch.tensor(std, device=device)[None, :, None, None] def __call__(self, tensor): return (tensor - self.mean) / self.std # mean and std of ILSVRC2012 as computed for the torchvision models norm_fn = BatchNormalize((0.485, 0.456, 0.406), - (0.229, 0.224, 0.225), device='cpu') + (0.229, 0.224, 0.225), device=device) +batch_size = 1 +# the maximal input shape, needed for the ZBox rule +shape = (batch_size, 3, 224, 224) + +# the highest and lowest pixel values for the ZBox rule +low = norm_fn(torch.zeros(*shape, device=device)) +high = norm_fn(torch.ones(*shape, device=device)) + + +model = vgg16(pretrained=True) +model.eval() # only these get rules, linear layers will be attributed by the gradient alone -target_types = (Conv2d, AvgPool2d) +target_types = (Conv2d, AvgPool2d, Linear) # lookup module -> name child_name = {module: name for name, module in model.named_modules()} # the layers in sequential order without any containers etc. layers = list(enumerate(collect_leaves(model))) # list of tuples [([names..], rule)] as used by NameMapComposite name_map = [ - ([child_name[module] for n, module in layers if n <= 16 and isinstance(module, target_types)], Gamma(0.25)), - ([child_name[module] for n, module in layers if 17 <= n <= 30 and isinstance(module, target_types)], GMontavonEpsilon(1e-9, 0.25)), - ([child_name[module] for n, module in layers if 30 <= n and isinstance(module, target_types)], Epsilon(1e-9)), + ([child_name[module] for n, module in layers if n == 0 and isinstance(module, target_types)], ZBox(low=low, high=high)), + ([child_name[module] for n, module in layers if 1 <= n <= 16 and isinstance(module, target_types)], Gamma(0.25)), + ([child_name[module] for n, module in layers if 17 <= n <= 30 and isinstance(module, target_types)], GMontavonEpsilon(stabilize_epsilon=0, epsilon=0.25)), + ([child_name[module] for n, module in layers if 31 <= n and isinstance(module, target_types)], Epsilon(0)), ] + # look at the name_map and you will see that there is no layer for which the last condition holds print(name_map) # create the composite from the name map composite = NameMapComposite(name_map) R = None with composite.context(model) as modified_model: # compute attribution # Returns a numpy array in BGR color space, not RGB img = cv2.imread('castle.jpg') # Convert from BGR to RGB color space img = img[..., ::-1] # img.shape is (224, 224, 3), where 3 corresponds to RGB channels # Divide by 255 (max. RGB value) to normalize pixel values to [0,1] img = img/255.0 data = norm_fn( torch.FloatTensor( img[numpy.newaxis].transpose([0, 3, 1, 2])*1 ) ) data.requires_grad = True output = modified_model(data) output[0].max().backward() # print absolute sum of attribution print(data.grad.abs().sum().item()) R = data.grad utils.heatmap(R[0].sum(dim=0).detach().numpy(), 4,4) ```

   

   Python code

   ```python import torch from torch.nn import Conv2d, AvgPool2d, Linear from torchvision.models import vgg16 from zennit.composites import NameMapComposite from zennit.core import BasicHook, collect_leaves, stabilize from zennit.rules import Gamma, Epsilon, ZBox import cv2 import numpy import utils # the LRP-Epsilon from the tutorial class GMontavonEpsilon(BasicHook): def __init__(self, stabilize_epsilon=1e-6, epsilon=0.25): super().__init__( input_modifiers=[lambda input: input], param_modifiers=[lambda param, _: param], output_modifiers=[lambda output: output], gradient_mapper=(lambda out_grad, outputs: out_grad / stabilize( outputs[0] + epsilon * (outputs[0] ** 2).mean() ** .5, stabilize_epsilon)), reducer=(lambda inputs, gradients: inputs[0] * gradients[0]) ) # use the gpu if requested and available, else use the cpu device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu') # Source: https://github.com/chr5tphr/zennit/blob/6251a9e17aa31c3381799de92f92b1d259b392b2/share/example/feed_forward.py#L32-L38 class BatchNormalize: def __init__(self, mean, std, device=None): self.mean = torch.tensor(mean, device=device)[None, :, None, None] self.std = torch.tensor(std, device=device)[None, :, None, None] def __call__(self, tensor): return (tensor - self.mean) / self.std # mean and std of ILSVRC2012 as computed for the torchvision models norm_fn = BatchNormalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225), device=device) batch_size = 1 # the maximal input shape, needed for the ZBox rule shape = (batch_size, 3, 224, 224) # the highest and lowest pixel values for the ZBox rule low = norm_fn(torch.zeros(*shape, device=device)) high = norm_fn(torch.ones(*shape, device=device)) model = vgg16(pretrained=True) model.eval() # only these get rules, linear layers will be attributed by the gradient alone # target_types = (Conv2d, AvgPool2d) target_types = (Conv2d, AvgPool2d, Linear) # lookup module -> name child_name = {module: name for name, module in model.named_modules()} # the layers in sequential order without any containers etc. layers = list(enumerate(collect_leaves(model))) # list of tuples [([names..], rule)] as used by NameMapComposite name_map = [ ([child_name[module] for n, module in layers if n == 0 and isinstance(module, target_types)], ZBox(low=low, high=high)), ([child_name[module] for n, module in layers if 1 <= n <= 16 and isinstance(module, target_types)], Gamma(0.25)), ([child_name[module] for n, module in layers if 17 <= n <= 30 and isinstance(module, target_types)], GMontavonEpsilon(stabilize_epsilon=0, epsilon=0.25)), ([child_name[module] for n, module in layers if 31 <= n and isinstance(module, target_types)], Epsilon(0)), ] # look at the name_map and you will see that there is no layer for which the last condition holds display(name_map) # create the composite from the name map composite = NameMapComposite(name_map) R = None with composite.context(model) as modified_model: # compute attribution # Returns a numpy array in BGR color space, not RGB img = cv2.imread('castle.jpg') # Convert from BGR to RGB color space img = img[..., ::-1] # img.shape is (224, 224, 3), where 3 corresponds to RGB channels # Divide by 255 (max. RGB value) to normalize pixel values to [0,1] img = img/255.0 data = norm_fn( torch.FloatTensor( img[numpy.newaxis].transpose([0, 3, 1, 2])*1 ) ) data.requires_grad = True output = modified_model(data) output[0].max().backward() # print absolute sum of attribution print(data.grad.abs().sum().item()) R = data.grad utils.heatmap(R[0].sum(dim=0).detach().numpy(), 4,4) ```