Add LRP composites (#84)

* Add Composite and primitives

* Add pretty-printing of composites

* Add tests and docstrings

* Add "Advanced LRP" docs example

* Update simple example

Author: adrhill

docs/literate/advanced_lrp.jl

@@ -5,7 +5,7 @@
 #
 # This example will show you how to implement custom LRP rules and register custom layers
 # and activation functions.
-# For this purpose, we will quickly load our model from the previous section:
+# For this purpose, we will quickly load our model from the previous section
 using ExplainableAI
 using Flux
 using MLDatasets
@@ -14,14 +14,14 @@ using BSON
 
 model = BSON.load("../model.bson", @__MODULE__)[:model]
 
+# and data from the MNIST dataset
 index = 10
 x, _ = MNIST(Float32, :test)[10]
 input = reshape(x, 28, 28, 1, :);
 
 # ## Custom LRP composites
-# Instead of creating an LRP-analyzer from a single rule (e.g. `LRP(model, GammaRule())`),
-# we can also assign rules to each layer individually.
-# For this purpose, we create an array of rules that matches the length of the Flux chain:
+# When creating an LRP-analyzer, we can assign individual rules to each layer.
+# The array of rules has to match the length of the Flux chain:
 rules = [
     ZBoxRule(0.0f0, 1.0f0),
     EpsilonRule(),
@@ -43,6 +43,44 @@ heatmap(input, analyzer)
 #md #     Not all models can be flattened, e.g. those using
 #md #     `Parallel` and `SkipConnection` layers.
 
+# Instead of manually defining a list of rules, we can also use a [`Composite`](@ref).
+# A composite contructs a list of LRP-rules by sequentially applying composite primitives.
+#
+# To obtain the same set of rules as in the previous example, we can define
+composite = Composite(
+    ZeroRule(),                         # default rule
+    GlobalRuleMap(
+        Conv => GammaRule(),       # apply GammaRule on all convolutional layers
+        MaxPool => EpsilonRule(),  # apply EpsilonRule on all pooling-layers
+    ),
+    FirstRule(ZBoxRule(0.0f0, 1.0f0)),  # apply ZBoxRule on the first layer
+)
+
+analyzer = LRP(model, composite)        # construct LRP analyzer from composite
+heatmap(input, analyzer)
+
+# This analyzer contains the same rules as our previous one:
+analyzer.rules # show rules
+
+# ### Composite primitives
+# The following sets of primitives can used to construct a [`Composite`](@ref).
+#
+# To apply a single rule, use:
+# * [`LayerRule`](@ref) to apply a rule to the `n`-th layer of a model
+# * [`GlobalRule`](@ref) to apply a rule to all layers of a model
+# * [`RangeRule`](@ref) to apply a rule to a positional range of layers of a model
+# * [`FirstRule`](@ref) to apply a rule to the first layer of a model
+# * [`LastRule`](@ref) to apply a rule to the last layer of a model
+#
+# To apply a set of rules to multiple layers, use:
+# * [`GlobalRuleMap`](@ref) to apply a dictionary that maps layer types to LRP-rules
+# * [`RangeRuleMap`](@ref) for a `RuleMap` on generalized ranges
+# * [`FirstNRuleMap`](@ref) for a `RuleMap` on the first `n` layers of a model
+# * [`LastNRuleMap`](@ref) for a `RuleMap` on the last `n` layers
+#
+# Primitives are called sequentially in the order the `Composite` was created with
+# and overwrite rules specified by previous primitives.
+
 # ## Custom LRP rules
 # Let's define a rule that modifies the weights and biases of our layer on the forward pass.
 # The rule has to be of type `AbstractLRPRule`.

docs/literate/example.jl

@@ -20,6 +20,7 @@ model = BSON.load("../model.bson", @__MODULE__)[:model]
 #md # !!! warning "Strip softmax"
 #md #
 #md #     For models with softmax activations on the output, it is necessary to call
+#md #     [`strip_softmax`](@ref)
 #md #     ```julia
 #md #     model = strip_softmax(model)
 #md #     ```
@@ -59,14 +60,14 @@ expl = analyze(input, analyzer);
 # * `expl.neuron_selection`: the neuron index of used for the attribution
 # * `expl.analyzer`: a symbol corresponding the used analyzer, e.g. `:LRP`
 
-# Finally, we can visualize the `Explanation` through heatmapping:
+# Finally, we can visualize the `Explanation` through [`heatmap`](@ref):
 heatmap(expl)
 
 # Or get the same result by combining both analysis and heatmapping into one step:
 heatmap(input, analyzer)
 
 # ## Neuron selection
-# By passing an additional index to our call to `analyze`, we can compute the attribution
+# By passing an additional index to our call to [`analyze`](@ref), we can compute the attribution
 # with respect to a specific output neuron.
 # Let's see why the output wasn't interpreted as a 4 (output neuron at index 5)
 heatmap(input, analyzer, 5)
@@ -76,7 +77,7 @@ heatmap(input, analyzer, 5)
 
 #md # !!! note
 #md #
-#md #     The ouput neuron can also be specified when calling `analyze`:
+#md #     The ouput neuron can also be specified when calling [`analyze`](@ref):
 #md #     ```julia
 #md #     expl = analyze(img, analyzer, 5)
 #md #     ```
@@ -110,26 +111,31 @@ mosaic(heatmap(batch, analyzer, 1); nrow=10)
 # ├── SmoothGrad
 # ├── IntegratedGradients
 # └── LRP
-#     ├── LRPZero
-#     ├── LRPEpsilon
-#     └── LRPGamma
+#     ├── ZeroRule
+#     ├── EpsilonRule
+#     ├── GammaRule
+#     ├── WSquareRule
+#     ├── FlatRule
+#     ├── ZBoxRule
+#     ├── AlphaBetaRule
+#     └── PassRule
 # ```
 #
-# Let's try `InputTimesGradient`
+# Let's try [`InputTimesGradient`](@ref)
 analyzer = InputTimesGradient(model)
 heatmap(input, analyzer)
 
-# and `Gradient`
+# and [`Gradient`](@ref)
 analyzer = Gradient(model)
 heatmap(input, analyzer)
 
-# As you can see, the function `heatmap` automatically applies common presets for each method.
+# As you can see, the function [`heatmap`](@ref) automatically applies common presets for each method.
 #
-# Since `InputTimesGradient` and LRP both compute attributions, their presets are similar.
+# Since [`InputTimesGradient`](@ref) and [`LRP`](@ref) both compute attributions, their presets are similar.
 # Gradient methods however are typically shown in grayscale.
 
 # ## Custom heatmap settings
-# We can partially or fully override presets by passing keyword arguments to `heatmap`:
+# We can partially or fully override presets by passing keyword arguments to [`heatmap`](@ref):
 using ColorSchemes
 heatmap(expl; cs=ColorSchemes.jet)
 #

docs/src/api.md

@@ -45,6 +45,28 @@ LRP_CONFIG.supports_layer
 LRP_CONFIG.supports_activation
 ```
 
+## Composites
+```@docs
+Composite
+```
+
+Composite primitives that apply a single rule:
+```@docs
+LayerRule
+GlobalRule
+RangeRule
+FirstRule
+LastRule
+```
+
+Composite primitives that apply a set of rules to multiple layers:
+```@docs
+GlobalRuleMap
+RangeRuleMap
+FirstNRuleMap
+LastNRuleMap
+```
+
 # Utilities
 ```@docs
 strip_softmax

src/ExplainableAI.jl

@@ -20,15 +20,17 @@ using PrettyTables
 include("compat.jl")
 include("neuron_selection.jl")
 include("analyze_api.jl")
-include("types.jl")
+include("flux_types.jl")
 include("flux_utils.jl")
 include("utils.jl")
 include("input_augmentation.jl")
 include("gradient.jl")
 include("lrp/canonize.jl")
 include("lrp/checks.jl")
 include("lrp/rules.jl")
+include("lrp/composite.jl")
 include("lrp/lrp.jl")
+include("lrp/show.jl")
 include("heatmap.jl")
 include("preprocessing.jl")
 export analyze
@@ -49,6 +51,12 @@ export modify_input, modify_denominator
 export modify_param!, modify_layer!
 export check_model
 
+# LRP composites
+export Composite, AbstractCompositePrimitive
+export LayerRule, GlobalRule, RangeRule, FirstRule, LastRule
+export GlobalRuleMap, RangeRuleMap, FirstNRuleMap, LastNRuleMap
+export ConvLayer, PoolingLayer, DropoutLayer, ReshapingLayer
+
 # heatmapping
 export heatmap