chr5tphr/zennit

Zennit is a high-level framework in Python using PyTorch for explaining/exploring neural networks using attribution methods like LRP.

Zennit

Zennit (Zennit explains neural networks in torch) is a high-level framework in Python using Pytorch for explaining/exploring neural networks. Its design philosophy is intended to provide high customizability and integration as a standardized solution for applying rule-based attribution methods in research, with a strong focus on Layerwise Relevance Propagation (LRP). Zennit strictly requires models to use Pytorch's torch.nn.Module structure (including activation functions).

Zennit is currently under active development, but should be mostly stable.

If you find Zennit useful for your research, please consider citing our related paper:

@article{anders2021software,
      author  = {Anders, Christopher J. and
                 Neumann, David and
                 Samek, Wojciech and
                 Müller, Klaus-Robert and
                 Lapuschkin, Sebastian},
      title   = {Software for Dataset-wide XAI: From Local Explanations to Global Insights with {Zennit}, {CoRelAy}, and {ViRelAy}},
      journal = {CoRR},
      volume  = {abs/2106.13200},
      year    = {2021},
}

Documentation

The latest documentation is hosted at zennit.readthedocs.io.

Install

To install directly from PyPI using pip, use:

$ pip install zennit

Alternatively, install from a manually cloned repository to try out the examples:

$ git clone https://github.com/chr5tphr/zennit.git
$ pip install ./zennit

Usage

At its heart, Zennit registers hooks at Pytorch's Module level, to modify the backward pass to produce rule-based attributions like LRP (instead of the usual gradient). All rules are implemented as hooks (zennit/rules.py) and most use the LRP basis BasicHook (zennit/core.py).

Composites (zennit/composites.py) are a way of choosing the right hook for the right layer. In addition to the abstract NameMapComposite, which assigns hooks to layers by name, and LayerMapComposite, which assigns hooks to layers based on their Type, there exist explicit Composites, some of which are EpsilonGammaBox (ZBox in input, Epsilon in dense, Gamma in convolutions) or EpsilonPlus (Epsilon in dense, ZPlus in convolutions). All composites may be used by directly importing from zennit.composites, or by using their snake-case name as key for zennit.composites.COMPOSITES.

Canonizers (zennit/canonizers.py) temporarily transform models into a canonical form, if required, like SequentialMergeBatchNorm, which automatically detects and merges BatchNorm layers followed by linear layers in sequential networks, or AttributeCanonizer, which temporarily overwrites attributes of applicable modules, e.g. to handle the residual connection in ResNet-Bottleneck modules.

Attributors (zennit/attribution.py) directly execute the necessary steps to apply certain attribution methods, like the simple Gradient, SmoothGrad or Occlusion. An optional Composite may be passed, which will be applied during the Attributor's execution to compute the modified gradient, or hybrid methods.

Using all of these components, an LRP-type attribution for VGG16 with batch-norm layers with respect to label 0 may be computed using:

import torch
from torchvision.models import vgg16_bn

from zennit.composites import EpsilonGammaBox
from zennit.canonizers import SequentialMergeBatchNorm
from zennit.attribution import Gradient


data = torch.randn(1, 3, 224, 224)
model = vgg16_bn()

canonizers = [SequentialMergeBatchNorm()]
composite = EpsilonGammaBox(low=-3., high=3., canonizers=canonizers)

with Gradient(model=model, composite=composite) as attributor:
    out, relevance = attributor(data, torch.eye(1000)[[0]])

A similar setup using the example script produces the following attribution heatmaps:

For more details and examples, have a look at our documentation.

More Example Heatmaps

More heatmaps of various attribution methods for VGG16 and ResNet50, all generated using share/example/feed_forward.py, can be found below.

Heatmaps for VGG16

Heatmaps for ResNet50

Contributing

See CONTRIBUTING.md for detailed instructions on how to contribute.

License

Zennit is licensed under the GNU LESSER GENERAL PUBLIC LICENSE VERSION 3 OR LATER -- see the LICENSE, COPYING and COPYING.LESSER files for details.

zennit's Issues

Core/Composite: Simplify Composite.context

Composite.context can be implemented slimmer/simpler using contextlib.contextmanager.
Furthermore, instead of calling Composite.context, the same functionality could be implemented as Composite.__call__, as the context is the main functionality, and this would remove a little bit of overhead, although .context may be more descriptive.

Potential issue with learning BatchNorm parameters

Hi, in my project I have encountered and issue that I'm not sure if it's caused by invalid usage of the library or there is some bug in the library code. I cannot provide minimal code for reproduction because bug occurs during training so I will describe it as best as I can.
Pseudo-code for my training looks something like:

for i in range(number_epochs):
    # train loop
    model.train()
    for input, gt in train_dataloader:
        with composite.context(model) as modified_model:
            model_out = modified_model(input)
            task_loss = get_task_loss(model_out, gt)
            gt_maps = torch.autograd.grad(model_out, input, torch.ones_like(labels), retain_graph=True)[0].sum(1)
            salience_loss = get_saliance_loss(model_out, gt)
            total_loss = task_loss + salience_loss
        
        total_loss.backward()
        self.optimizer.step()
        self.optimizer.zero_grad()
            
    #validation loop
    model.eval()
    for input, gt in train_dataloader:
        with composite.context(model) as modified_model:
            model_out = modified_model(input)
            ...

For model I have tested VGG16, Resnet34 and VGG16_bn with appropriate canonizers, and for composite I have used EpsilonPlusFlat. All models have their heads changed to have 20 outputs, and are randomly initialized. I have noticed that models with BatchNorm have significant difference between output when in train mode and when in eval.

I have logged the sum of output during training to show this for different models.

For VGG16 we can see that output sums have around the same order of magnitude which is expected:

For ResNet34 we see drastic change in output sums, around 4 orders of magnitudes difference

For VGG16_bn we again see difference in output sums but difference is "only" around 1 order of magnitude:

I see that this behaviour is very strange but it all points to something being wrong with BatchNorm.
Version of Zennit I'm using is 0.5.2.dev5.
I would really appreciate your help regarding this one.
Thanks in advance.

Composite: bug in BetaSmooth constructor?

Hi,

I noticed that the beta_smooth parameter is not passed to the ReLUBetaSmooth hook when initializing the BetaSmooth composite in composites.py. This results in the default use of beta_smooth=10 at all times. I assume that's not intended, or am I missing sth?

Best

RuntimeError: Mismatch in shape: grad_output[0] has a shape of torch.Size([1, 1, 80, 80]) and output[0] has a shape of torch.Size([1, 1, 96, 96]).

I encountered a problem below when I applied the lrp method using the zennit to U-net model ResNet-based backend with the input shape (1, 32, 96, 96) (e.g. (batch size, channels, width, height)) and output shape (1, 1, 80, 80). Do the input and output shapes have to be the same? If not, do you have any solution to solve this problem? Any help would be appreciated!

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-30-5d15a2db92a9> in <module>
     55         output_relevance[:, -1, :, :] = output[:, -1, :, :]
     56         # this will compute the modified gradient of model, with the on
---> 57         output, relevance = attributor(data.cuda(), output_relevance.cuda())
     58 
     59         # sum over the color channel for visualization

~/tools/miniconda3/envs/pytorch/lib/python3.8/site-packages/zennit/attribution.py in __call__(self, input, attr_output)
    130 
    131         if self.composite is None or self.composite.handles:
--> 132             return self.forward(input, attr_output_fn)
    133 
    134         with self:

~/tools/miniconda3/envs/pytorch/lib/python3.8/site-packages/zennit/attribution.py in forward(self, input, attr_output_fn)
    175         input = input.detach().requires_grad_(True)
    176         output = self.model(input)
--> 177         gradient, = torch.autograd.grad((output,), (input,), grad_outputs=(attr_output_fn(output.detach()),))
    178         return output, gradient
    179 

~/tools/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/autograd/__init__.py in grad(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused)
    195 
    196     grad_outputs_ = _tensor_or_tensors_to_tuple(grad_outputs, len(outputs))
--> 197     grad_outputs_ = _make_grads(outputs, grad_outputs_)
    198 
    199     if retain_graph is None:

~/tools/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/autograd/__init__.py in _make_grads(outputs, grads)
     31         if isinstance(grad, torch.Tensor):
     32             if not out.shape == grad.shape:
---> 33                 raise RuntimeError("Mismatch in shape: grad_output["
     34                                    + str(grads.index(grad)) + "] has a shape of "
     35                                    + str(grad.shape) + " and output["

RuntimeError: Mismatch in shape: grad_output[0] has a shape of torch.Size([1, 1, 80, 80]) and output[0] has a shape of torch.Size([1, 1, 96, 96]).

RuntimeError: Mismatch in shape: grad_output[0] has a shape of torch.Size([1, 1, 80, 80]) and output[0] has a shape of torch.Size([1, 1, 96, 96]).

I encountered a problem below when I applied the lrp method using the zennit to U-net model ResNet-based backend with the input shape (1, 32, 96, 96) (e.g. (batch size, channels, width, height)) and output shape (1, 1, 80, 80). Do the input and output shapes have to be the same? If not, If not, do you have a solution to this problem? Any help would be appreciated!

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-30-5d15a2db92a9> in <module>
     55         output_relevance[:, -1, :, :] = output[:, -1, :, :]
     56         # this will compute the modified gradient of model, with the on
---> 57         output, relevance = attributor(data.cuda(), output_relevance.cuda())
     58 
     59         # sum over the color channel for visualization

~/tools/miniconda3/envs/pytorch/lib/python3.8/site-packages/zennit/attribution.py in __call__(self, input, attr_output)
    130 
    131         if self.composite is None or self.composite.handles:
--> 132             return self.forward(input, attr_output_fn)
    133 
    134         with self:

~/tools/miniconda3/envs/pytorch/lib/python3.8/site-packages/zennit/attribution.py in forward(self, input, attr_output_fn)
    175         input = input.detach().requires_grad_(True)
    176         output = self.model(input)
--> 177         gradient, = torch.autograd.grad((output,), (input,), grad_outputs=(attr_output_fn(output.detach()),))
    178         return output, gradient
    179 

~/tools/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/autograd/__init__.py in grad(outputs, inputs, grad_outputs, retain_graph, create_graph, only_inputs, allow_unused)
    195 
    196     grad_outputs_ = _tensor_or_tensors_to_tuple(grad_outputs, len(outputs))
--> 197     grad_outputs_ = _make_grads(outputs, grad_outputs_)
    198 
    199     if retain_graph is None:

~/tools/miniconda3/envs/pytorch/lib/python3.8/site-packages/torch/autograd/__init__.py in _make_grads(outputs, grads)
     31         if isinstance(grad, torch.Tensor):
     32             if not out.shape == grad.shape:
---> 33                 raise RuntimeError("Mismatch in shape: grad_output["
     34                                    + str(grads.index(grad)) + "] has a shape of "
     35                                    + str(grad.shape) + " and output["

RuntimeError: Mismatch in shape: grad_output[0] has a shape of torch.Size([1, 1, 80, 80]) and output[0] has a shape of torch.Size([1, 1, 96, 96]).

Changing The Value of epsilon in the composite Epsilonplusflat creates a bug

Hi to everyone
I wanted to describe a problem that I encounter when I use the code of the tutorial.

• create a composite

composite = EpsilonPlusFlat()

• choose a target class for the attribution (label 437 is lighthouse)

target = torch.eye(1000)[[437]]

• create the attributor, specifying model and composite

with Gradient(model=model, composite=composite) as attributor:
    # compute the model output and attribution
    output, attribution = attributor(data, target)

print(f'Prediction: {output.argmax(1)[0].item()}')

This is the code i'm talking about, when i put a new value of epsilon inside the parenthesis of epsilonplusflat, the new heatmap is literally the same even if i change the epsilon

Custom rules for vision transformer

Hello
I'm trying to use this method on a vision transformer model(model = torchvision.models.vit_b_16(), first several layers in below image). I read the document, And I think I need to write and use new rules?(I see that there are some new types of layers that doesn't have an existing class . And also submodules are a little complex. So I have to use new rules right?). I read the document of how to write a custom rule, but I can't think of an idea which rules use on which layer in this VIT model.( I want to get images like original epsilonplusflat.) I run the code below and got error show below. Do you have any recommendations on how to run lrp method on this model?
Thank you!

composite = EpsilonPlusFlat()
with Gradient(model=model, composite=composite) as attributor:
     output, attribution = attributor(data, target)

Docs: Build Tutorials on ReadTheDocs with trained weights

• add a flag to disable/enable downloading and use of trained weights in tutorials
• add this flag to the notebooks
• this will result in better visualizations in the documentation at readthedocs.org

LinearAttention Module

Hi Christopher,

hope you're fine and I'm really glad that the zennit community grows, congratulation!
With a growing community, more nn.Modules desire to be explained and that's why I'm writing this issue.
A student in our department tries to explain a LinearAttention module. (The implementation is below for reference).

It contains a series of
torch.einsum
and
torch.transpose
operations.

It uses the rearrange function of the einops library, a new syntax to write basic torch code like transpose, reshape etc.

I think, zennit should be able to analyse a series of reshaping and transposing operations. However, I am not completely sure.
I'd be glad, if you could give your opinion on analyzing such a linear attention module. If you don't know, that's also no problem (: Then, it's the beginning of a new research topic.

(And the softmax function is also a problem, but maybe Arras et. al has a solution to this which the student could implement... )

Best,
Reduan

class LinearAttention(nn.Module):
    def __init__(self, dim, heads=4, dim_head=32):
        super().__init__()
        self.scale = dim_head**-0.5
        self.heads = heads
        hidden_dim = dim_head * heads
        self.to_qkv = nn.Conv2d(dim, hidden_dim * 3, 1, bias=False)

        self.to_out = nn.Sequential(nn.Conv2d(hidden_dim, dim, 1),
                                    nn.GroupNorm(1, dim))

    def forward(self, x):
        b, c, h, w = x.shape
        qkv = self.to_qkv(x).chunk(3, dim=1)
        q, k, v = map(
            lambda t: rearrange(t, "b (h c) x y -> b h c (x y)", h=self.heads), qkv
        )

        q = q.softmax(dim=-2)
        k = k.softmax(dim=-1)

        q = q * self.scale
        context = torch.einsum("b h d n, b h e n -> b h d e", k, v)

        out = torch.einsum("b h d e, b h d n -> b h e n", context, q)
        out = rearrange(out, "b h c (x y) -> b (h c) x y", h=self.heads, x=h, y=w)
        return self.to_out(out)

Numerical instability in ResNet50 heatmaps

Calculating the relevance on ResNet50 seems to be prone to a numerical instability, producing heatmaps where all attribution is concentrated in a few spots because the values in those spots have become larger than the rest. See heatmap in bug reproduction section. I can confirm that this unexpected behavior also happens using different composites and different images.

I have also seen this issue on VGG16 in my own LRP implementation depending on the heuristic used in the stabilize function.

Bug reproduction

Code based on snippet provided in #76 (comment).

Minimal reproducible example:

import cv2
import numpy
import torch
from matplotlib import pyplot as plt
from torchvision.models import resnet50
from zennit.composites import EpsilonGammaBox
from zennit.image import imgify
from zennit.torchvision import ResNetCanonizer


# 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=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 = resnet50(pretrained=True)
model.eval()

# create the composite from the name map
composite = EpsilonGammaBox(low=-high, high=high, canonizers=[ResNetCanonizer()])

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())

    # relevance scores
    R = data.grad

    # show maximum and minimum attribution
    print(torch.aminmax(R))

    heatmap = imgify(
        R.detach().cpu().sum(1),
        symmetric=True,
        grid=True,
        cmap='seismic',
    )
    
    plt.imshow(heatmap)

Input(s):

• Input image:

  

Outputs:

• Text:

  755226.0625
  torch.return_types.aminmax(
  min=tensor(-7540.4922),
  max=tensor(2985.0886))

• Heatmap:

  

Additional information

The bug is not limited to the castle.jpg image, it can also be reproduced using the following image. See the corresponding heatmap below.

• Input image:

  

• Heatmap:

  

Smooth MaxPool2D rule

Hey,

we'd like to add a new rule that smooths the MaxPool2D operation by replacing it by an AveragePool2D backward pass:

class SmoothMaxPool2dRule(BasicHook):

    def __init__(self, epsilon=1e-6, zero_params=None):
        stabilizer_fn = Stabilizer.ensure(epsilon)
        super().__init__(
            gradient_mapper=(lambda out_grad, outputs: out_grad / stabilizer_fn(outputs[0])),
            reducer=(lambda inputs, gradients: inputs[0] * gradients[0]),
        )

    def backward(self, module, grad_input, grad_output):
        '''Backward hook to compute LRP based on the class attributes.'''
        original_input = self.stored_tensors['input'][0].clone()
        inputs, outputs = [], []
        kernel_size = module.kernel_size
        stride = module.stride
        padding = module.padding
        
        input = original_input.requires_grad_()
        with torch.autograd.enable_grad():
            output = F.avg_pool2d(input, kernel_size, stride, padding, ceil_mode=False, count_include_pad=True, divisor_override=None)
        inputs.append(input)
        outputs.append(output)
        
        grad_outputs = self.gradient_mapper(grad_output[0], outputs)
        gradients = torch.autograd.grad(
            outputs,
            inputs,
            grad_outputs=grad_outputs,
            create_graph=grad_output[0].requires_grad
        )
        relevance = self.reducer(inputs, gradients)
        return tuple(relevance if original.shape == relevance.shape else None for original in grad_input)

You can test the code with

import torch.nn as nn
from zennit.rules import *
from zennit.core import BasicHook
import torch.nn.functional as F

if __name__ == "__main__":
    
    input = torch.linspace(0, 35, 36).view(1, 1, 6, 6).requires_grad_()

    layer = nn.MaxPool2d(2, 2, 0)
    norm_rule = Norm()
    h = norm_rule.register(layer)

    output = layer(input)
    grad, = torch.autograd.grad(output, input, torch.ones_like(output))
    h.remove()

    print(input)
    print(output)
    print(grad)

    print("###")

    rule = SmoothMaxPool2dRule()
    h = rule.register(layer)

    output = layer(input)
    grad, = torch.autograd.grad(output, input, torch.ones_like(output))
    h.remove()

    print(input)
    print(output)
    print(grad)

Do you think that's fine? I can create a pull request if you want.

Best,
Reduan

Feature request: Assign rule by layer index

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.

Core: Second Order Gradients

I am currently working on supporting second-order gradients, i.e. gradients of the modified gradients, which is used for example to compute adversarial explanations.
The current issue which prevents second order gradients is that the gradient modification introduced by rules will also be applied when in the second-order backward pass.
This will prevented by disabling the modification temporarily when computing the second-order gradient, likely using something like a no_modification context for composites/attributors/rules.

As also pointed out in #125, handles are not stored for the backward hooks for tensors. Storing and removing the hooks before the second-order backward pass would correctly compute the modified-gradient derivatives, although then the same graph cannot be used to compute the modified gradient for a different gradient output. By adding the context, I am considering to also enable complete removal of the tensor backward hooks.

ResNet canonizer misses canonizer for BasicBlock module

Smaller ResNets such as ResNet18 or ResNet34 consist of BasicBlock modules instead of Bottleneck blocks.

In order to ensure canonization of small ResNets, a BasicBlock should be added to the ResNetCanonizer.

Model inference outside of composite context

Hi @chr5tphr ,

I'm creating a new issue regarding the inference of a model.
I noticed that when I infer a model outside of (before) composite context (which has appropriate model cannonizer) I do not obtain the same attribution as when the inference is done inside of the context. This has me concerned because in order to properly learn batch norm's parameters inference should be done outside of the context because context effectively creates identity out of batch norm so inference inside of the context would never update it's parameter values. Is there something I'm not understanding here correctly?

Here is a code snippet I used to validate that attribution is not the same. Same test was also conducted on vgg16 model and yielded the same result.

import torch
from torchvision.models import resnet18

from zennit.composites import EpsilonPlusFlat
from zennit.torchvision import ResNetCanonizer

from PIL import Image
from torchvision.transforms import Compose, Resize, CenterCrop
from torchvision.transforms import ToTensor, Normalize

import matplotlib.pyplot as plt

# define the base image transform
transform_img = Compose([
    Resize(256),
    CenterCrop(224),
])
# define the normalization transform
transform_norm = Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
# define the full tensor transform
transform = Compose([
    transform_img,
    ToTensor(),
    transform_norm,
])
# load the image
image = Image.open('dornbusch-lighthouse.jpg')
# transform the PIL image and insert a batch-dimension
data = transform(image)[None]
data.requires_grad = True
# define target
target = torch.eye(1000)[[437]]


model = resnet18()

canonizers = [ResNetCanonizer()]
composite = EpsilonPlusFlat(canonizers=canonizers)

# Inference before context
model.eval()  # Put model in eval so batch-norm is frozen
model_out_before = model(data)
with composite.context(model) as modified_model:
    attribution_before, = torch.autograd.grad(model_out_before, data, target)

# Inference inside context
with composite.context(model) as modified_model:
    model_out_in = modified_model(data)
    attribution_in, = torch.autograd.grad(model_out_in, data, target)

relevance_before = attribution_before.cpu().sum(1).squeeze(0).numpy()
relevance_in = attribution_in.cpu().sum(1).squeeze(0).numpy()

plt.figure(figsize=(15, 5))
plt.subplot(1,3,1)
plt.imshow(transform_img(image))
plt.axis('off')
plt.subplot(1,3,2)
plt.imshow(relevance_before)
plt.axis('off')
plt.subplot(1,3,3)
plt.imshow(relevance_in)
plt.axis('off')
plt.show()

Obtanining graident of LRP otput w.r.t. network parameters

Hi, first of all thank you for all the hard work that was put into developing this framework and then making it available to everyone.
I was wondering if there is a way to obtain gradient of the explanation obtained using LRP with respect to the network parameters in order to optimize it.
I stumbled across your overview paper and would like to use the framework in my own EGL research.

Binary classification model with one output for output layer

Hi,

I want to use zennit to calculate LRP relevance for Binary classification model which only one neuron in output layer, how to set the second para of attributor?

class test_model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.lin1 = torch.nn.Linear(3 * 512 + 512, 1024)
        self.lin2 = torch.nn.Linear(1024, 512)
        self.lin3 = torch.nn.Linear(512, 1)

    def forward(self, x):
        out = torch.relu(self.lin1(x))
        out = torch.relu(self.lin2(out))
        out = self.lin3(out)
        return out
model = test_model()

from zennit.composites import EpsilonPlusFlat
composite = EpsilonPlusFlat()
test_input.requires_grad = True
from zennit.attribution import Gradient
attributor = Gradient(model , composite)
with attributor:
    output, relevance = attributor(test_input, torch.tensor([1]))

Core: GradOutHooks

I am currently working on GradOutHook, which is different from the current zennit.core.Hook in that instead of overwriting the full gradient of the module, it only changes the gradient output. For Composites using zennit.core.Hook, only a single Hook can be attached at a time, because it will change the full gradient of the module. The GradOutHook can modify the output gradient multiple times, and can be used together with zennit.core.Hook. This can lead to using multiple Composites at a time. Another way to enable multiple hooks would be to let the module_map function of Composites allow to return a tuple of Hooks to be applied.

The main use case for this is to mask or re-weight neurons, mainly to support LRP for GNNs. Another use-case is to mask certain neurons to get LRP for a subset of features/concepts.

This will somewhat change the Hook-inheritance, where a HookBase will be added to specify the interface necessary for all Hooks. Also, I am considering to add a Mask rule to zennit/rule.py which takes a function or a tensor to mask the gradient output, which can be used without subclassing the planned GradOutHook.

error occurred when run tutorial on custom model

Hi
I'm trying to use zennit to draw a lrp heatmap for my self trained vgg11 model. I followed the tutorial on the document, but when I use my models(which just trained the last layer and changed output dim to 6) will have errors. It will say that input type and weight type should be the same. Then I changed my data to cuda mode, the second error "tuple has no attribute detach" occurred. Is this tutorial only for pure pretrained model? Can you tell me how to implement it on my slightly customed vgg11 model?
Thank you!

Image fetching in example is broken

image-net.org removed their API, breaking the download procedure used in the example with subimagenet.sh.

Default Rule

Hi,

I'd like to know which rule Zennit applies if the layer is unknown?
More precisely: which rule is applied to BatchNorm2d when not using a canonizer?

I would be able to figure this out on my own with a debugger, but unfortunately VS Code and Pycharm won't let me set a breakpoint in the backward pass of Zennit, but in the forward pass it works.

Thanks!

[How to use] heatmap variants [...] such as the guided_backprop or smoothgrad [...] ?

Hi, thanks again for the response. After looking through the linked original example. I'm wondering why I can't use other heatmap variants other than epsilon ? such as the guided_backprop or smoothgrad which seem to produce a better heatmap then epsilon_plus or epsilon_alpha2_beta1_flat

Originally posted by @linghu258 in #24 (comment)

support unet with Upsample or ConvTranspose2d layer?

I want to use LRP to explain the semantic segmentation task using Unet model (Pytorch). I tested the LRP in captum but not support nn.Upsample and nn.ConvTranspose2d. I would like to know if the semantic segmentation model like Unet can be supported, and if not, how should it be implemented? Any help would be appreciated!

Switch over to Module.register_full_backward_hook

The most recent version of PyTorch seems to have introduced a proper module backward hook behavior.
In favor of this new behavior, the old backward hook behavior will be deprecated.
At some point, we should use the pytorch full backward hook instead of our work-around of doing a full backward hook.

Memory Leak

Hi Chris,

unfortunately, there is a memory leak when you run zennit several times in a row.
You can reproduce the results on a GPU using:

from torchvision.models.vgg import vgg16
import torch
import zennit
from zennit.composites import COMPOSITES

model = vgg16(pretrained=True).to("cuda")
model.eval()

data = torch.randn((1,3,224,224)).to("cuda")
data.requires_grad = True
target = 0

eye = torch.eye(1000, device="cuda")
output_relevance = eye[[target]]

for i in range(100):
    print(i)

    try:
        composite = COMPOSITES["epsilon_plus_flat"]()
    except:
        raise ValueError(f"Method not defined. Available are {list(COMPOSITES.keys())}")

    with composite.context(model) as modified:
        # one-hot tensor of target

        pred = modified(data)

        torch.autograd.backward((pred,), (output_relevance,))

    # the attribution will be stored in the gradient's place
    heatmaps = data.grad.sum(1).detach().cpu().numpy()

    r = torch.cuda.memory_reserved(0)
    a = torch.cuda.memory_allocated(0)
    print("memory reserved ", r, "memory allocated ", a)

Do you suspect where the leak is coming from?

Feature Request: Debugger

Hey,

I think it would be really nice, if we could check if the composite actually attached to all the modules we named and to see if there are any modules left that have no rules attached to.
In the long run - but much more difficult to realize - it would be nice to see, if there are any non-Module elements that unexpectedly changed the Relevance values.

Best

feed_forward is broken for torchvision 0.10.0

Due to torchvision PR 3496, torchvision.datasets.folder.make_dataset now raises a FileNotFoundError if there are empty classes. Therefore, we cannot simply make a dataset with a set amount of classes and label files accordingly by moving them to their respective label folder, if there are classes without members.

Notebook for feed_forward.py

Hi,

I'm playing with zennit at the moment; usually, I like to start with a simple notebook. So, I made such a notebook for myself based on share/examples/feed_forward.py.

Perhaps, we can convert it to be a tutorial of the package.

https://colab.research.google.com/drive/1MNLkfMi9DB9lZkkcOPOodQhXCatkwQgO#scrollTo=LIGiH5SP8IC4

Layer-wise LRP score

Hi,

Thanks a lot for the awesome work! May I know how I can extract layer-wise LRP attribution scores for ResNet-18?

How to add support to DenseNet from torchvision

I wanted to use LRP with DesneNet121 from torchvision. So far, in zennit.torchvision I found canonizers for ResNet and VGG and I wonder if I may use them (to get some results) or I need to write my own custom canonizer (because the network has some new layers that were not covered by previous canonizers?).

Thanks for your help!

Support for padding layers

Hi,

I am trying to analyze a model consisting of Padding Layers using LRP. I cant seem to find any mentions of padding layers in the code however. Are they implicitly supported already and if not what steps are necessary to make it work?

Thanks!

Module with Multiple Inputs

Hey Chris,

hope you're well. I noticed an implementation detail where I am unsure if this was programmed on purpose and why.

At line you take in the backward pass only the first input, while saving previously all inputs in line.

I see that you defined the summation layer using a concat operation at line, so I assume restricting the inputs is on purpose.

So do you think, it is possible to attribute a summation layer defined in the following way in the future? And why did you restrict the input layers to have only one input?

class Sum(torch.nn.Module):

    def forward(self, input_1, input_2):
        return input_1 + input_2

Thanks a lot!

Bug in AlphaBeta rule?

Hi,

I've been working very extensively with the LRP method lately and I also tried to implement the method with the most commonly used rules by myself. In order to check the correctness of my implementation, I compared some results with already existing implementations (like yours ;) ). Thereby I always get different relevances with the AlphaBeta-rule (I already opened an issue in innvestigate with the same example). Maybe you can explain the following behavior or confirm that this is really a bug on your side:

Let's suppose we have only one layer with two inputs, one output and no activation function. The layer has the following weights and bias vector: W = (1, -1) and b = -1. For the input x = (1,1), the formula of the AlphaBeta rule (Eq. (60)) in Bach et al. reduces to (in this case is r_out = -1)

• r1 = alpha * (1 * 1 / (1 * 1)) + beta * (0 / -1) * r_out = alpha * r_out
• r2 = alpha * (0 / 1) + beta * ((1 * -1) / (1 * -1 + (-1))) * r_out = beta * 0.5 * r_out

This yields a relevance of (-1, 0) for the Alpha1_Beta0-rule and (-2, 0.5) for the Alpha2_Beta1-rule. But with your implementation I get both times (0.5, -0.5). Also for other choices of alpha and beta I always get the same result. Here is my code snippet for the Alpha1_Beta0 rule:

import torch
import torch.nn as nn
from zennit.rules import AlphaBeta

input = torch.tensor([[1.,1.]], requires_grad = True)

model = nn.Sequential(
      nn.Linear(2, 1)
  )
model.get_submodule("0").weight.data = torch.tensor([[1., -1.]])
model.get_submodule("0").bias.data = torch.tensor([-1.])

rule = AlphaBeta(alpha = 1, beta = 0)
rule.register(model)
output = model(input)

grad_out = torch.ones_like(output) * output

attr, = torch.autograd.grad(
    output, input, grad_outputs=grad_out
)

attr
# tensor([[ 0.5000, -0.5000]])

I hope you can help me to clarify this behavior.

Best Niklas

Canonizers: BatchNorm loses bias relevance

With the implementation of #185 allowing for gradient computation wrt. the parameters given an attribution with canonizers, the BatchNorms seem to be leaking attribution even with zero_params='bias' on ResNet18.

Bug: MultiheadAttention in types

Hi,

here for some mysterious reason the nn.MultiheadAttention is part of the Activation types.

I am actually implementing a canonizer for this module and then it just didn't trigger and I was wondering what is happening until I found out that the layer_base automatically overwrites my rule with the Pass() one (;

Best,
Reduan

ResNet: Unstable Attributions

With the introduction of #185 , ResNet18 attributions result in negative attribution sums in the input layer, leading to bad attributions. Although #185 increased the stability of the attribution sums for ResNet, the previous instability seems to have inflated the positive parts of the attributions, circumventing this problem pre #185.

This seems to be related to leaking attributions (#193 ) combined with skip connections that can cause negative attributions.

A quickfix for EpsilonGammaBox is to use a slightly higher gamma value.

Flat Rule for Pooling Layers

Hi Chris,

I am defining a new composite for instance:

@register_composite('all_flat')
class AllFlat(LayerMapComposite):

    def __init__(self, canonizers=None):
        layer_map = [
            (Linear, Flat()),
            (AvgPool, Flat()),
            (Activation, Pass()),
            (Sum, Norm()),
        ]
       
        super().__init__(layer_map, canonizers=canonizers)

The problem is, that the Flat() rule changes the parameter of a layer and the pooling layers do not define the "weight" parameter. As a consequence, there will be a RuntimeError saying, that zennit tries to access the parameter "weight" which is not available.
The solution would be to define a new rule that does not have a param_modifier, for instance:

class FlatPooling(LinearHook):
    '''This is the Flat LRP rule. It is essentially the same as the WSquare Rule, but with all parameters set to ones.
    '''
    def __init__(self):
        super().__init__(
            input_modifiers=[torch.ones_like],
            param_modifiers=[None],
            output_modifiers=[lambda output: output],
            gradient_mapper=(lambda out_grad, outputs: out_grad / stabilize(outputs[0])),
            reducer=(lambda inputs, gradients: gradients[0])
        )

What do you think?

Best

Feature request: Support for max-pooling layers

In the code there seems to be no mention of max-pooling layers, is LRP for this type of layers already (implicitly) supported?

Neuralized K-Means: make k-means amenable to neural network explanations

Description

K-Means finds some centroids $\mathbf{\mu}_1,\ldots,\mathbf{\mu}_K$ and assigns data points to a cluster as $c = {\rm argmin}_k \lbrace \Vert\boldsymbol{x} - \mathbf{\mu}_k\Vert^2\rbrace$, or in code:

c = torch.argmin(torch.cdist(x, centroids)**2)

Neither the assignment $c$ nor the distance $\Vert\boldsymbol{x} - \mathbf{\mu}_c\Vert^2$ are really explainable.
The assignment $c$ is not continuous, and the distance is in fact measuring a dissimilarity to the cluster, essentially the opposite of what we want to explain.

Fixes

From Clustering to Cluster Explanation via Neural Networks (2022) present a method based on neuralization. They show that the k-means cluster assignment can be exactly rewritten as a set of piecewise linear functions

$$f_c(\boldsymbol{x}) = \min_{k\neq c}\lbrace \boldsymbol{w}_{ck}^\top\boldsymbol{x} + b_{ck}\rbrace$$

and the original k-means cluster assignment can be recovered as $c = {\rm argmax}_k\lbrace f_k(\boldsymbol{x})\rbrace$.

The cluster discriminant $f_c$ is a measure of cluster membership (instead of a dissimilarity) and is also structurally amenable to LRP and friends. It can also be plugged on top of a neural network feature extractor to make deep cluster assignments explainable.

Additional Information

• There is no pairwise distance layer in PyTorch. Since zennit canonizers apply to torch.nn.Module modules, such module would need to be added
• Each discriminant $f_c$ has it's own weight matrix $W\in\mathbb{R}^{(K-1)\times D}$. In principle, a neuralized k-means layer could be written as a layer

torch.nn.Sequential(
    torch.nn.Linear(in_features=D, out_features=K*(K-1)), 
    torch.nn.Unflatten(1, torch.Size([K, K-1]))
)

or alternatively via torch.einsum, which I think is easier. Either way, a special layer is needed for that.

• The paper suggests a specific LRP rule for the $\min_{k\neq c}\lbrace s_k\rbrace$ to make the explanations continuous. In practice, we can replace it by $\beta^{-1}{\rm LogMeanExp}_k\lbrace \beta\cdot s_k\rbrace$ which approximates the min for $\beta < 0$, but has continuous gradients. Gradient propagation through this layer recovers exactly the LRP rule from the paper, but this also provides explanation continuity for other gradient-based explanation methods. I would suggest to implement such a LogMeanExpPool layer. Alternative could be to implement a MinPool layer with a custom LRP rule.
• A canonizer could be used to replace the distance/k-means layer by it's neuralized counterpart. Alternatively, a composite might also work, but it seems a bit more complicated (where are the weights computed? how to set LRP rules for the linear/einsum layer?)

Proper way of handling classifier heads in Transformers

Hi,

I am attempting to implement Conservative LRP rules as part of #184 since I need it for a university project. I am running into certain issues with the classifier heads and was hoping you could point out what I'm potentially doing wrong.

I have so far implemented a composite in the following manner:

@register_composite('transformer')
class Transformer(LayerMapComposite):

    def __init__(self, linear_layer_epsilon=1e-6, layer_norm_epsilon=1e-5, layer_map=None, canonizers=None, zero_params=None):
        if layer_map is None:
            layer_map = []

        rule_kwargs = {'zero_params': zero_params}
        layer_map = [(Activation, Pass()),
                     (Convolution, ZPlus()),
                     (AvgPool, Norm()),
                     (MultiheadAttention, AHConservative()),
                     (LayerNorm, LNConservative(epsilon=layer_norm_epsilon)),
                     (Linear, Epsilon(epsilon=linear_layer_epsilon, **rule_kwargs))]

        # layer_map.composites += general_layer_map
        # named_map = NameLayerMapComposite([(("classifier",), )])

        super().__init__(layer_map=layer_map, canonizers=canonizers)

Where AHConservative and LNConservative are Rules described in the CLRP Paper.

I have also implemented a custom attributor which calculates relevance scores with respect to the embeddings (since integer inputs are not differentiable). However the problem I am facing seems to appear also with a basic Gradient attributor. Here is a minimal example:

import torch
from torch import nn
from zennit.composites import Transformer
from zennit.attribution import Gradient

composite = Transformer()

model = nn.Linear(768, 2, bias=True)

with Gradient(model=model, composite=composite) as attributor:
    input = torch.randn((1, 768))
    out, relevance = attributor(input.float(), attr_output=torch.ones((1, 2)))

Working with BertForSequenceClassification from the Huggingface transformers library the classifier head is an nn.Linear module with the size (768, 2). I am however getting an error from the Epsilon rule, specifically from the gradient_mapper:

Exception has occurred: RuntimeError       (note: full exception trace is shown but execution is paused at: wrapper)
The size of tensor a (768) must match the size of tensor b (2) at non-singleton dimension 1
  File "/home/chris/zennit/src/zennit/rules.py", line 120, in <lambda>
    gradient_mapper=(lambda out_grad, outputs: out_grad / stabilizer_fn(outputs[0])),
  File "/home/chris/zennit/src/zennit/core.py", line 539, in backward
    grad_outputs = self.gradient_mapper(grad_output[0], outputs)
  File "/home/chris/zennit/src/zennit/core.py", line 388, in wrapper (Current frame)
    return hook.backward(module, grad_input, grad_output)
  File "/home/chris/zennit/.venv/lib/python3.10/site-packages/torch/autograd/__init__.py", line 303, in grad
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
  File "/home/chris/zennit/src/zennit/attribution.py", line 257, in grad
    gradient, = torch.autograd.grad(
  File "/home/chris/zennit/src/zennit/attribution.py", line 287, in forward
    return self.grad(input, attr_output_fn)
  File "/home/chris/zennit/src/zennit/attribution.py", line 181, in __call__
    return self.forward(input, attr_output_fn)
  File "/home/chris/zennit/t3.py", line 25, in <module>
    out, relevance = attributor(input.float(), attr_output=torch.ones((1, 2)))
  File "/home/chris/miniconda3/lib/python3.10/runpy.py", line 86, in _run_code
    exec(code, run_globals)
  File "/home/chris/miniconda3/lib/python3.10/runpy.py", line 196, in _run_module_as_main
    return _run_code(code, main_globals, None,
RuntimeError: The size of tensor a (768) must match the size of tensor b (2) at non-singleton dimension 1

I do understand that the tensor sizes must agree for division in the gradient_mapper. I am therefore suspecting I'm mishandling the classifier head but I am not sure how to proceed. Should I upsample the output to the size of input and just use Epsilon? Should I implement a custom rule? Any help would be largely appreciated! I'd love to get a rough implementation of the changes suggested in #184 working so we could push the needle on XAI for Transformers a bit.

SmoothGrad - potential typo in implementation of noise level

Hi,

Thanks for a great library!

There might be a discrepancy between the implementation and the original paper when it comes to how noise is calculated https://arxiv.org/pdf/1706.03825.pdf.

zennit/src/zennit/attribution.py

Line 356 in 60a2c08

std = self.noise_level * (input.amax(dims, keepdim=True) - input.amin(dims, keepdim=True))

i.e., it should be $std = \frac{sigma}{(xmax-xmin)}$ and not $std = sigma * (xmax-xmin)$.

Canonizers

Hi Chris,

hopefully you're fine.
There are two small problems I have noticed:

1. Applying for instance SequentialMergeBatchNorm, the code iterates trough all modules in a model and checks if two subsequent modules are of Conv and BatchNorm type. Unfortunately, this does not work if the modules are not ordered linearly in the model. For instance I defined a model using torch's ModuleList and shuffled the layers as I wished in the forward pass. As a consequence, Zennit was not able to identify which module follows another one. I see, that in torch it is really difficult. Some time ago, I draw a model graph using torch's jit. Do you think there is a way to find out which module follows which one in a forward pass and then adapt Zennit to consider this correct ordering? For now, I just rewrote my code but maybe someone else will not be able to... at least we need to write this in the documentation if not already.

2. SequentialMergeBatchNorm does not work for Conv1d layer unfortunately. The problem is following line in
   def merge_batch_norm(modules, batch_norm):

module.weight.data = (original_weight * scale[:, None, None, None])

There, you cast the weights to a fixed dimension of 4. But Conv1d layers have weight of dimension 3 i.e. we need
module.weight.data = (original_weight * scale[:, None, None])

Do you know an elegant way to achieve this for all dimensions?

Best

New version of zennit produces wrong LRP heatmaps

Hi,

it is a long time ago, I pulled all the new zennit updates. So I missed a lot of comments and I can not exactly point out where the LRP pipeline seems to have broken down the first time.
For test purposes, I always use the same LeNet model pretrained on FashionMNIST.
When I generate a heatmap with "epsilon_plus", I get with the old version of zennit the same heatmap as in Innvestigate. But the new version suddenly produces completely different heatmaps.

You can try it out with my code example on https://github.com/rachtibat/mnist_zennit I uploaded for you.
Just change the import "zennit_old" to "zennit_new" and you see the difference. I changed the name of zennit.torchvision to zennit.torchvision_2 making it work fast and dirty (; (see #6)

Do you think, I made a mistake or is zennit somehow computing the heatmaps differently?

Best

P.s. if this is really a bug, it might solve the other issues

Core: Warn the user when falling back to gradient

Some modules are implicitly mapped to the gradient.

We can explicitly map Module types to None in their respective module_map in composites and warn the user when no rule is found to prevent hidden incompatibilities.

An example for an implicitly mapped module is MaxPooling.

Alpha Beta still includes Bias two times

Hi,

if I am not wrong, the alpha beta rule still includes the bias two times.

param_modifiers=[
                lambda param, _: param.clamp(min=0),
                lambda param, name: param.clamp(max=0) if name != 'bias' else torch.zeros_like(param),
                lambda param, _: param.clamp(max=0),
                lambda param, _: param.clamp(min=0),
            ],

Maybe you can write

param_modifiers=[
                lambda param, _: param.clamp(min=0),
                lambda param, name: param.clamp(max=0) if name != 'bias' else torch.zeros_like(param),
                lambda param, _: param.clamp(max=0),
                lambda param, name: param.clamp(min=0) if name != 'bias' else torch.zeros_like(param),
            ],

Best

torch 1.10 version not working

Hi Chris,

it seems that the pytorch version 1.10 produces wrong heatmaps and maybe also newer versions of torch.
I updated my conda environment and suddenly the "epsilon_plus_flat" rule did not work anymore.
Returning to
pip install torch=1.9 torchvision=0.10
fixed the issue.

Best

Not removing backward hook handler *might* cause problems?

Hi Chris,

Thank you for your support.
I noticed in the Hook source code,
that you register a backward hook, but you don't save its handler.

As far as I understand, the hook is destroyed when the tensor is destroyed.
But if the tensor stays alive for whatever reason, then the hook will persist. This might cause a problem?
What do you think?

Best