Feature: New LayerFilter class

[rodrigobdz]

Circling back on the thread about assigning LRP rules by layer index, I have implemented a LayerFilter class to reduce the boilerplate code required to specify LRP rules by layer index.
This class would be particularly useful for users taking advantage of NameMapComposite.

@chr5tphr Concretely, let me know if you consider this class interesting for zennit and I would be happy to submit a proper PR. Before putting in the work, I wanted to check with you first.

TL;DR

Instead of this:

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)], Gamma(0.15)),
    ([child_name[module] for n, module in layers if 31 <=
        n and isinstance(module, target_types)], Gamma(0.05)),
]

I propose:

name_map = [
    (filter_by_layer_index_type(lambda n: n == 0), ZBox(low=low, high=high)),
    (filter_by_layer_index_type(lambda n: 1 <= n <= 16), Gamma(0.25)),
    (filter_by_layer_index_type(lambda n: 17 <= n <= 30), Gamma(0.15)),
    (filter_by_layer_index_type(lambda n: 31 <= n), Gamma(0.05)),
]

Demo

Here's an example of how the new LayerFilter class would improve the user-friendliness when using NameMapComposite by using the LayerFilter class to elegantly filter layers by index and type:

Before

The definition of the name_map variable might be cumbersome and requires boilerplate code in this first snippet.

Without LayerFilter class

import torch
from torch.nn import Conv2d, AvgPool2d, Linear
from torchvision.models import vgg16

from zennit.composites import NameMapComposite
from zennit.core import collect_leaves
from zennit.rules import Gamma, ZBox

import cv2
import numpy

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

# only these get rules, linear layers will be attributed by the gradient alone
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)], Gamma(0.15)),
    ([child_name[module] for n, module in layers if 31 <=
        n and isinstance(module, target_types)], Gamma(0.05)),
]

# 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

After

LayerFilter class definition

r"""Filter layers before rules are mapped."""

from typing import Callable, Dict, List, Optional, Tuple

import torch

from zennit.core import collect_leaves


class LayerFilter:  # pylint: disable=too-few-public-methods
    r"""Filter the layers by layer index and type."""

    def __init__(self, model: torch.nn.Module, target_types: Tuple[type]) -> None:
        r"""Build a filter for the given model.

        :param model: Model to filter layers from
        :param target_types: Types of layers to filter
        """
        # Collect all layers by name
        self.name_lookup: Dict[str, torch.nn.Module] = {
            module: name for name, module in model.named_modules()}

        # Model layers in sequential order
        self.layers: List[Tuple[int, torch.nn.Module]] = list(
            enumerate(collect_leaves(model)))

        # Types of layers to filter
        self.target_types = target_types

    def __call__(self, cond: Callable[[int], bool],
                 target_types: Optional[Tuple[type]] = None) -> List[str]:
        r"""Filter layers by layer index and type.

        :param cond: Function to filter layers by layer index
        :param target_types: Types of layers to filter

        :return: List of layer names
        """
        # Use default target types if not specified
        allowed_types = target_types or self.target_types

        return [self.name_lookup[module] for n, module in self.layers
                if cond(n) and isinstance(module, allowed_types)]

The definition of the name_map variable is improved in this snippet by reducing the noise around the actual content.

With LayerFilter class

import torch
from torch.nn import Conv2d, AvgPool2d, Linear
from torchvision.models import vgg16

from zennit.composites import NameMapComposite
from zennit.core import collect_leaves
from zennit.rules import Gamma, ZBox

import cv2
import numpy

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

# Init layer filter
# only these get rules, linear layers will be attributed by the gradient alone
target_types: Tuple[type] = (Conv2d, AvgPool2d, Linear)
filter_by_layer_index_type: LayerFilter = LayerFilter(model=model,
                                                    target_types=target_types)

# list of tuples [([names..], rule)] as used by NameMapComposite
name_map = [
  (filter_by_layer_index_type(lambda n: n == 0), ZBox(low=low, high=high)),
  (filter_by_layer_index_type(lambda n: 1 <= n <= 16), Gamma(0.25)),
  (filter_by_layer_index_type(lambda n: 17 <= n <= 30), Gamma(0.15)),
  (filter_by_layer_index_type(lambda n: 31 <= n), Gamma(0.05)),
]

# 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

[chr5tphr]

Hey Rodrigo,

thank you for the idea and detailed description.
I really think that defining Composites over the layer-depth is not a good approach.
Nevertheless, if you think this is an important use-case, we could discuss a way to make quality of life improvements to building Composites in this way.

Designing Composites by Layer-Depth

I previously proposed to use a NameMapComposite for your specific use-case, since when using the layer-depth as a rule-indicator, the Composite is made with a specific architecture in mind.
Ideally, you would make the NameMap once and re-use it for the same model. The approach is demonstrated here where NameMapComposites are introduced.

While I gave you an example of how to map layers to names as a demonstration to how these name_maps can be constructed, I would not recommend this is as the best approach, but rather suggest implementing it with a custom Composite, skipping the intermediate mapping to the names.
In the documentation, the example I give for a custom Composite is exactly the one you are looking for, which is defined over the layer-depth and the type (I recommend the second code cell). Note that there you can supply tuples of types, or rewrite the layer to use a default tuple of types when only a 1-tuple is supplied.

LayerMapFilter

The vanilla Composite requires a module_map function, which acts as custom mapping from layer, name and context to a Rule. The LayerMapFilter you provide is a new construct, which tries to circumvent implementing a new Composite by adding an additional layer of abstraction using the layer name. Design-wise, I do not think this fits very well into the framework.

However, instead of adding this functionality as a new layer of abstraction, I could imagine implementing an abstract LayerFilterComposite (there may be a better name), similar to LayerMapComposite, which instead of defining a single module_map function, could require a list of 2-tuples, with a function taking the same arguments as the module_map, but with a boolean return type, and a rule template, which will be used if the function evaluates to True given a module.

I asked myself this question when I implemented the core Composite class, where I ultimately decided that having a single function return the Hook template is a better approach than supplying one function per hook template, since I felt this was clearer, easier to understand (only a single function as opposed to many), and more flexible.

So actually, I think it may be a cleaner approach to directly implement a Composite subclass with a custom module_map, instead of supplying many tuples of lambda-expressions and Rule-templates.

Conclusion

• Unless there is another use-case for your proposed LayerFilter, I do not think it should be implemented as a new utility-type.
• The LayerFilter concept could be implemented as a new abstract Composite type, but I think it is more convoluted than just using the core Composite and implementing a custom module_map.
• I would like to not promote choosing rule-maps based on the layer-depth. If you (or other users of this approach) can convince me (and I like to be convinced of things that make sense) that the layer-depth is a good, simple, widely-adapted approach to defining rule Composites (outside of Grégoire's tutorial), I would rather adopt the LeafNumberTypeComposite from the documentation, possibly with a few tweaks.

[rodrigobdz]

Hi, Christopher,

Perhaps we could let the userbase vote on this feature to see if it's worth introducing and maintining this additional LayerFilter class. I do understand that this would add some further complexity to the already feature-rich codebase.

Argument

Aside from the minutiae of the implementation of LayerFilter, which can certainly be improved, personally, my main argument in favor of the feature of assigning rules by layer index is:

Since I'm conducting a hyperparameter grid search for LRP rules to evaluate explanations, I need the flexibility to test as many combinations as possible—this includes assigning also rules by layer index.