Implement Base.show for LRP analyzers (#89)

* Implement `Base.show` for LRP analyzers

* Add tests

* Update "Advanced LRP" docs

Author: adrhill

docs/literate/advanced_lrp.jl

@@ -20,9 +20,10 @@ x, _ = MNIST(Float32, :test)[10]
 input = reshape(x, 28, 28, 1, :);
 
 # ## LRP composites
-# ### Custom composites
+# ### Assigning individual rules
 # 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:
+# The array of rules has to match the length of the Flux chain.
+# The `LRP` analyzer will show a summary of how layers and rules got matched:
 rules = [
     ZBoxRule(0.0f0, 1.0f0),
     EpsilonRule(),
@@ -35,6 +36,8 @@ rules = [
 ]
 
 analyzer = LRP(model, rules)
+
+#
 heatmap(input, analyzer)
 
 # Since some Flux Chains contain other Flux Chains, ExplainableAI provides
@@ -44,8 +47,10 @@ heatmap(input, analyzer)
 #md #     Not all models can be flattened, e.g. those using
 #md #     `Parallel` and `SkipConnection` layers.
 
+# ### Custom composites
 # 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.
+# A composite contructs a list of LRP-rules by sequentially applying
+# [Composite primitives](@ref composite_primitive_api) it contains.
 #
 # To obtain the same set of rules as in the previous example, we can define
 composite = Composite(
@@ -57,11 +62,12 @@ composite = Composite(
     FirstLayerRule(ZBoxRule(0.0f0, 1.0f0)),  # apply ZBoxRule on the first layer
 )
 
-analyzer = LRP(model, composite)             # construct LRP analyzer from composite
-heatmap(input, analyzer)
+# We now construct an LRP analyzer from `composite`
+analyzer = LRP(model, composite)
 
-# This analyzer contains the same rules as our previous one:
-analyzer.rules # show rules
+# As you can see, this analyzer contains the same rules as our previous one
+# and therefore also produces the same heatmaps:
+heatmap(input, analyzer)
 
 # ### Composite primitives
 # The following [Composite primitives](@ref composite_primitive_api) can used to construct a [`Composite`](@ref).
@@ -87,7 +93,9 @@ analyzer.rules # show rules
 # ### Default composites
 # A list of implemented default composites can be found under
 # [Default composites](@ref default_composite_api) in the API reference, e.g. [`EpsilonPlusFlat`](@ref):
-EpsilonPlusFlat()
+composite = EpsilonPlusFlat()
+#
+analyzer = LRP(model, composite)
 
 # ## Custom LRP rules
 # Let's define a rule that modifies the weights and biases of our layer on the forward pass.
@@ -101,7 +109,7 @@ struct MyGammaRule <: AbstractLRPRule end
 import ExplainableAI: modify_param!
 
 function modify_param!(::MyGammaRule, param)
-    param .+= 0.25 * relu.(param)
+    param .+= 0.25f0 * relu.(param)
     return nothing
 end
 

src/lrp/show.jl

@@ -1,4 +1,33 @@
-# Composites
+const COLOR_COMMENT = :light_black
+const COLOR_RULE = :yellow
+const COLOR_TYPE = :blue
+const COLOR_RANGE = :green
+
+typename(x) = string(nameof(typeof(x)))
+
+################
+# LRP analyzer #
+################
+
+_print_layer(io::IO, l) = string(sprint(show, l; context=io))
+function Base.show(io::IO, m::MIME"text/plain", analyzer::LRP)
+    layer_names = [_print_layer(io, l) for l in analyzer.model]
+    npad = maximum(length.(layer_names)) + 1 # padding to align rules with rpad
+
+    println(io, "LRP", "(")
+    for (l, r) in zip(layer_names, analyzer.rules)
+        print(io, "  ", rpad(l, npad), " => ")
+        printstyled(io, r; color=COLOR_RULE)
+        println(io, ",")
+    end
+    println(io, ")")
+    return nothing
+end
+
+#############
+# Composite #
+#############
+
 _range_string(r::LayerRule) = "layer $(r.n)"
 _range_string(::GlobalRule) = "all layers"
 _range_string(r::RangeRule) = "layers $(r.range)"
@@ -11,13 +40,6 @@ _range_string(::LastLayerTypeRule) = "last layer"
 _range_string(r::FirstNTypeRule) = "layers $(1:r.n)"
 _range_string(r::LastNTypeRule) = "last $(r.n) layers"
 
-const COLOR_COMMENT = :light_black
-const COLOR_RULE = :yellow
-const COLOR_TYPE = :blue
-const COLOR_RANGE = :green
-
-typename(x) = string(nameof(typeof(x)))
-
 function Base.show(io::IO, m::MIME"text/plain", c::Composite)
     println(io, "Composite", "(")
     for p in c.primitives

test/references/show/lrp1.txt

@@ -0,0 +1,21 @@
+LRP(
+  Conv((3, 3), 3 => 64, relu, pad=1)     => ZBoxRule{Float32}(-3.0f0, 3.0f0),
+  MaxPool((2, 2))                        => EpsilonRule{Float32}(1.0f-6),
+  Conv((3, 3), 64 => 128, relu, pad=1)   => FlatRule(),
+  MaxPool((2, 2))                        => EpsilonRule{Float32}(1.0f-5),
+  Conv((3, 3), 128 => 256, relu, pad=1)  => FlatRule(),
+  Conv((3, 3), 256 => 256, relu, pad=1)  => FlatRule(),
+  MaxPool((2, 2))                        => EpsilonRule{Float32}(1.0f-5),
+  Conv((3, 3), 256 => 512, relu, pad=1)  => AlphaBetaRule{Float32}(2.0f0, 1.0f0),
+  Conv((3, 3), 512 => 512, relu, pad=1)  => AlphaBetaRule{Float32}(1.0f0, 0.0f0),
+  MaxPool((2, 2))                        => EpsilonRule{Float32}(1.0f-5),
+  Conv((3, 3), 512 => 512, relu, pad=1)  => AlphaBetaRule{Float32}(2.0f0, 1.0f0),
+  Conv((3, 3), 512 => 512, relu, pad=1)  => AlphaBetaRule{Float32}(2.0f0, 1.0f0),
+  MaxPool((2, 2))                        => EpsilonRule{Float32}(1.0f-6),
+  Flux.flatten                           => PassRule(),
+  Dense(25088 => 4096, relu)             => EpsilonRule{Float32}(1.0f-6),
+  Dropout(0.5)                           => PassRule(),
+  Dense(4096 => 4096, relu)              => EpsilonRule{Float32}(1.0f-7),
+  Dropout(0.5)                           => ZeroRule(),
+  Dense(4096 => 1000)                    => PassRule(),
+)

test/references/show/lrp2.txt

@@ -0,0 +1,10 @@
+LRP(
+  Conv((5, 5), 1 => 6, relu)   => AlphaBetaRule{Float32}(2.0f0, 1.0f0),
+  MaxPool((2, 2))              => ZeroRule(),
+  Conv((5, 5), 6 => 16, relu)  => ZeroRule(),
+  MaxPool((2, 2))              => ZeroRule(),
+  Flux.flatten                 => ZeroRule(),
+  Dense(256 => 120, relu)      => ZeroRule(),
+  Dense(120 => 84, relu)       => ZeroRule(),
+  Dense(84 => 10)              => EpsilonRule{Float32}(2.0f-5),
+)

test/test_composite.jl

@@ -20,9 +20,7 @@ for (name, c) in DEFAULT_COMPOSITES
     @test_reference "references/show/$name.txt" repr("text/plain", c)
 end
 
-# Test printing
-
-# This composite is non-sensical, but covers as many composite primitives as possible
+# This composite is non-sensical, but covers many composite primitives
 composite1 = Composite(
     ZeroRule(), # default rule
     GlobalRule(PassRule()), # override default rule
@@ -62,6 +60,7 @@ analyzer1 = LRP(model, composite1)
     ZeroRule()
     PassRule()
 ]
+@test_reference "references/show/lrp1.txt" repr("text/plain", analyzer1)
 @test_reference "references/show/composite1.txt" repr("text/plain", composite1)
 
 model = Chain(
@@ -91,4 +90,5 @@ analyzer2 = LRP(model, composite2)
     ZeroRule()
     EpsilonRule(2.0f-5)
 ]
+@test_reference "references/show/lrp2.txt" repr("text/plain", analyzer2)
 @test_reference "references/show/composite2.txt" repr("text/plain", composite2)