Implementation of LRP for pytorch

PyTorch implementation of some of the Layer-Wise Relevance Propagation (LRP)
rules, [1, 2, 3], for linear layers and convolutional layers.

The modules decorates torch.nn.Sequential, torch.nn.Linear, and
torch.nn.Conv2d to be able to use autograd backprop algorithm to compute
explanations.

Installation

To install requirements, refer to the requirements.yml
file.

If you use conda, then you can install an environment called torchlrp by
executing the following command:

> conda env create -f requirements.yml

To be able to import lrp as below, make sure that the TorchLRP directory is
included in your path.

Usage

The code can be used as follows:

import torch
import lrp
model = Sequential(
    lrp.Conv2d(1, 32, 3, 1, 1),
    torch.nn.ReLU(),
    torch.nn.MaxPool2d(2, 2),
    torch.nn.Flatten(),
    lrp.Linear(14*14*32, 10)
)
x = ... # business as usual
y_hat = model.forward(x, explain=True, rule="alpha2beta1")
y_hat = y_hat[torch.arange(batch_size), y_hat.max(1)[1]] # Choose maximizing output neuron
y_hat = y_hat.sum()
# Backward pass (do explanation)
y_hat.backward()
explanation = x.grad

Implemented rules:
|Rule |Key | Note |
|:———————————————|:———————————|:—————————————————————————|
|epsilon-rule | “epsilon” | Implemented but epsilon fixed to 1e-1 |
|gamma-rule | “gamma” | Implemented but gamma fixed to 1e-1 |
|epsilon-rule | “epsilon” | gamma and epsilon fixed to 1e-1 |
|alpha=1 beta=0 | “alpha1beta0” | |
|alpha=2 beta=1 | “alpha2beta1” | |
|PatternAttribution (all) | “patternattribution” | Use additional argument pattern=patterns_all |
|PatternAttribution (positive) | “patternattribution” | Use additional argument pattern=patterns_pos |
|PatternNet (all) | “patternnet” | Use additional argument pattern=patterns_all |
|PatternNet (positive) | “patternnet” | Use additional argument pattern=patterns_pos |

To compute patterns for the two PatternAttribution methods, import
lrp.patterns and call

import lrp.patterns.*
patterns_all = fit_patternnet(model, train_loader)
patterns_pos = fit_patternnet_positive(model, train_loader)

Note: Biases are currently ignored in the alphabeta-rule implementations.

Trace intermediate relevances

Thanks to francescomalandrino, you can now also
trace the intermediate relevances by enabling traces:

... 
lrp.trace.enable_and_clean()
y_hat.backward()
all_relevances=lrp.trace.collect_and_disable()
for i,t in enumerate(all_relevances):
    print(i,t.shape)

MNIST

For a complete running example, please see examples/explain_mnist.py.
The code generates this plot:

To run the example code, simply activate the conda environment and execute the code from the root of the project:

> conda activate torchlrp
> python examples/explain_mnist.py

VGG / ImageNet

It is also possible to use this code for pretrained vgg models from torchvision,
by using the lrp.convert_vgg function to convert torch.nn.Conv2d and torch.nn.Linear layers to lrp.Conv2d and lrp.Linear, respectively.

It takes a bit to make the vgg example work. First, you need An imagenet dataloader.
In the code, we use the dataloader from the torch_imagenet repo.
You could also make your own.

The most interesting parts is converting the torch vgg models, such that they can be
explained. To do so, do as follows:

vgg = torchvision.models.vgg16(pretrained=True).to(device)
vgg.eval()
lrp_vgg = lrp.convert_vgg(vgg).to(device)

The lrp_vgg model will then have the same parameters as the original network.
Afterwards, explanations can be produced as the example above.

Note:

The code example reads a config.ini file from the root of this project. In
that file you can specify the parent of the torch_imagenet repo such that the
correct dataloader is loaded:

[DEFAULT]
ImageNetDir = /home/user/example/data

Possible bugs

Fixed - Description

• [ ] According to [3] Section 10.3.2, it is apparently a good idea to use gradient of average pooling for LRP backpropagation. I have started to implement this but not finished, as I didn’t need it so far.

• [x] Fixed in commit 4277098f4f37a81ae9a21154c8cba49cae918770_. Judging
  from the plot, something is probably wrong with the positive
  PatternAttribution and PatternNet, as it doesn’t compare visually to, e.g.,
  this implementation.

References

[1] Bach, S., Binder, A., Montavon, G., Klauschen, F., Müller, K.R. and Samek, W., 2015. On pixel-wise explanations for non-linear classifier decisions by layer-wise relevance propagation. PloS one, 10(7), p.e0130140.
[2] Kindermans, P.J., Schütt, K.T., Alber, M., Müller, K.R., Erhan, D., Kim, B. and Dähne, S., 2017. Learning how to explain neural networks: Patternnet and patternattribution. arXiv preprint arXiv:1705.05598.
[3] Montavon, G., Binder, A., Lapuschkin, S., Samek, W. and Müller, K.R., 2019. Layer-wise relevance propagation: an overview. In Explainable AI: interpreting, explaining and visualizing deep learning (pp. 193-209). Springer, Cham.