Add Explanation type and dispatch heatmap on it (#36)

* adds wrapper `Explanation` around analysis that contains meta-data such as the used analyzer
* adds `:norm` reducer
* dispatch `heatmap` on analyzer symbol in Explanation
* add option to directly call `heatmap` with analyzer

Author: adrhill

docs/literate/example.jl

@@ -41,15 +41,22 @@ input = permutedims(input, (2,1,3))[:,:,:,:] * 255; # flip X/Y axes, add batch d
 # We can now select an analyzer of our choice
 # and call [`analyze`](@ref) to get an explanation `expl`:
 analyzer = LRPZero(model)
-expl, out = analyze(input, analyzer);
+expl = analyze(input, analyzer);
 
 # Finally, we can visualize the explanation through heatmapping:
 heatmap(expl)
 
+# Or do both in one combined step:
+heatmap(input, analyzer)
+
 #md # !!! tip "Neuron selection"
 #md #     To get an explanation with respect to a specific output neuron (e.g. class 42) call
 #md #     ```julia
-#md #     expl, out = analyze(img, analyzer, 42)
+#md #     expl = analyze(img, analyzer, 42)
+#md #     ```
+#md #     or using `heatmap`
+#md #     ```julia
+#md #     heatmap(img, analyzer, 42)
 #md #     ```
 #
 # Currently, the following analyzers are implemented:
@@ -75,13 +82,11 @@ model = flatten_model(model)
 #
 # Now we set a rule for each layer
 rules = [
-    ZBoxRule(),
-    repeat([GammaRule()], 15)...,
-    repeat([ZeroRule()], length(model) - 16)...
+    ZBoxRule(), repeat([GammaRule()], 15)..., repeat([ZeroRule()], length(model) - 16)...
 ]
 # and define a custom LRP analyzer:
 analyzer = LRP(model, rules)
-expl, out = analyze(input, analyzer)
+expl = analyze(input, analyzer)
 heatmap(expl)
 
 # ## Custom rules
@@ -98,7 +103,7 @@ end
 
 # We can directly use this rule to make an analyzer!
 analyzer = LRP(model, MyCustomLRPRule())
-expl, out = analyze(input, analyzer)
+expl = analyze(input, analyzer)
 heatmap(expl)
 
 #md # !!! tip "Pull requests welcome"

src/analyze_api.jl

@@ -13,7 +13,7 @@ Otherwise, the output neuron with the highest activation is automatically chosen
 function analyze(
     input::AbstractArray{<:Real},
     method::AbstractXAIMethod,
-    neuron_selection::Integer,
+    neuron_selection::Integer;
     kwargs...,
 )
     return method(input, IndexNS(neuron_selection); kwargs...)
@@ -22,3 +22,13 @@ end
 function analyze(input::AbstractArray{<:Real}, method::AbstractXAIMethod; kwargs...)
     return method(input, MaxActivationNS(); kwargs...)
 end
+
+# Explanations and outputs are returned in a wrapper.
+# Metadata such as the analyzer allows dispatching on functions like `heatmap`.
+struct Explanation{A,O,L}
+    attribution::A
+    output::O
+    neuron_selection::Int
+    analyzer::Symbol
+    layerwise_relevances::L
+end

src/gradient.jl

@@ -10,8 +10,8 @@ end
 function (analyzer::Gradient)(input, ns::AbstractNeuronSelector)
     output = analyzer.model(input)
     output_neuron = ns(output)
-    expl = gradient((in) -> analyzer.model(in)[output_neuron], input)[1]
-    return expl, output
+    attr = gradient((in) -> analyzer.model(in)[output_neuron], input)[1]
+    return Explanation(attr, output, output_neuron, :Gradient, Nothing)
 end
 
 """
@@ -29,6 +29,6 @@ end
 function (analyzer::InputTimesGradient)(input, ns::AbstractNeuronSelector)
     output = analyzer.model(input)
     output_neuron = ns(output)
-    expl = input .* gradient((in) -> analyzer.model(in)[output_neuron], input)[1]
-    return expl, output
+    attr = input .* gradient((in) -> analyzer.model(in)[output_neuron], input)[1]
+    return Explanation(attr, output, output_neuron, :InputTimesGradient, Nothing)
 end

src/heatmap.jl

@@ -1,28 +1,50 @@
 # NOTE: Heatmapping assumes Flux's WHCN convention (width, height, color channels, batch size).
 
+const HEATMAPPING_PRESETS = Dict{Symbol,Tuple{ColorScheme,Symbol,Symbol}}(
+    # Analyzer => (colorscheme, reduce, normalize)
+    :LRP => (ColorSchemes.bwr, :sum, :centered),
+    :InputTimesGradient => (ColorSchemes.bwr, :sum, :centered), # same as LRP
+    :Gradient => (ColorSchemes.grays, :norm, :extrema),
+)
+
 """
-    heatmap(expl; kwargs...)
+    heatmap(expl::Explanation; kwargs...)
+    heatmap(attr::AbstractArray; kwargs...)
+    heatmap(input, analyzer::AbstractXAIMethod)
+    heatmap(input, analyzer::AbstractXAIMethod, neuron_selection::Int)
 
 Visualize explanation.
 Assumes the Flux's WHCN convention (width, height, color channels, batch size).
 
 ## Keyword arguments
--`cs::ColorScheme`: ColorScheme that is applied. Defaults to `ColorSchemes.bwr`.
+-`cs::ColorScheme`: ColorScheme that is applied.
+    When calling `heatmap` with an `Explanation` or analyzer, the method default is selected.
+    When calling `heatmap` with an array, the default is `ColorSchemes.bwr`.
 -`reduce::Symbol`: How the color channels are reduced to a single number to apply a colorscheme.
-    Can be either `:sum` or `:maxabs`. `:sum` sums up all color channels for each pixel.
-    `:maxabs` selects the `maximum(abs, x)` over the color channel in each pixel.
-    Default is `:sum`.
+    The following methods can be selected, which are then applied over the color channels
+    for each "pixel" in the attribution:
+    - `:sum`: sum up color channels
+    - `:norm`: compute 2-norm over the color channels
+    - `:maxabs`: compute `maximum(abs, x)` over the color channels in
+    When calling `heatmap` with an `Explanation` or analyzer, the method default is selected.
+    When calling `heatmap` with an array, the default is `:sum`.
 -`normalize::Symbol`: How the color channel reduced heatmap is normalized before the colorscheme is applied.
-    Can be either `:extrema` or `:centered`. Default for use with colorscheme `bwr` is `:centered`.
+    Can be either `:extrema` or `:centered`.
+    When calling `heatmap` with an `Explanation` or analyzer, the method default is selected.
+    When calling `heatmap` with an array, the default for use with the `bwr` colorscheme is `:centered`.
+-`permute::Bool`: Whether to flip W&H input channels. Default is `true`.
+
+**Note:** these keyword arguments can't be used when calling `heatmap` with an analyzer.
 """
+
 function heatmap(
-    expl::AbstractArray;
+    attr::AbstractArray;
     cs::ColorScheme=ColorSchemes.bwr,
     reduce::Symbol=:sum,
     normalize::Symbol=:centered,
     permute::Bool=true,
 )
-    _size = size(expl)
+    _size = size(attr)
     length(_size) != 4 && throw(
         DomainError(
             _size,
@@ -36,11 +58,27 @@ function heatmap(
             """heatmap is only applicable to a single attribution, got a batch dimension of $(_size[end]).""",
         ),
     )
-    # drop batch dim -> reduce color channels -> normalize image -> apply color scheme
-    img = _normalize(_reduce(dropdims(expl; dims=4), reduce), normalize)
+
+    img = _normalize(dropdims(_reduce(dropdims(attr; dims=4), reduce); dims=3), normalize)
     permute && (img = permutedims(img))
     return ColorSchemes.get(cs, img)
 end
+# Use HEATMAPPING_PRESETS for default kwargs when dispatching on Explanation
+function heatmap(expl::Explanation; permute::Bool=true, kwargs...)
+    _cs, _reduce, _normalize = HEATMAPPING_PRESETS[expl.analyzer]
+    return heatmap(
+        expl.attribution;
+        reduce=get(kwargs, :reduce, _reduce),
+        normalize=get(kwargs, :normalize, _normalize),
+        cs=get(kwargs, :cs, _cs),
+        permute=permute,
+    )
+end
+# Analyze & heatmap in one go
+function heatmap(input, analyzer::AbstractXAIMethod, args...; kwargs...)
+    return heatmap(analyze(input, analyzer, args...; kwargs...))
+end
+
 
 # Normalize activations across pixels
 function _normalize(attr, method::Symbol)
@@ -58,16 +96,20 @@ function _normalize(attr, method::Symbol)
     return (attr .- min) / (max - min)
 end
 
-# Reduces activation in a pixel with multiple color channels into a single activation
-function _reduce(attr, method::Symbol)
-    if method == :maxabs
-        return dropdims(maximum(abs, attr; dims=3); dims=3)
+# Reduce attributions across color channels into a single scalar – assumes WHCN convention
+function _reduce(attr::T, method::Symbol) where {T}
+    if size(attr, 3) == 1 # nothing need to reduce
+        return attr
+    elseif method == :maxabs
+        return maximum(abs, attr; dims=3)
+    elseif method == :norm
+        return mapslices(norm, attr; dims=3)::T
     elseif method == :sum
-        return dropdims(sum(attr; dims=3); dims=3)
+        return sum(attr; dims=3)
     end
     throw(
         ArgumentError(
-            "Color channel reducer :$method not supported, `reduce` should be :maxabs or :sum",
+            "Color channel reducer :$method not supported, `reduce` should be :maxabs, :sum or :norm",
         ),
     )
 end

src/lrp.jl

@@ -66,9 +66,11 @@ function (analyzer::LRP)(input, ns::AbstractNeuronSelector; layerwise_relevances
         rels[i] .= lrp(rule, layers[i], acts[i], rels[i + 1])
     end
 
-    if layerwise_relevances
-        return rels, acts
-    end
-
-    return rels[1], acts[end] # expl, output
+    return Explanation(
+        first(rels),
+        last(acts),
+        output_neuron,
+        :LRP,
+        ifelse(layerwise_relevances, rels, Nothing),
+    )
 end

test/references/heatmaps/reduce_norm_normalize_centered.txt

@@ -0,0 +1 @@
+▀▀

test/references/heatmaps/reduce_norm_normalize_extrema.txt

@@ -0,0 +1 @@
+▀▀

test/references/heatmaps/vgg11_Gradient.txt

@@ -0,0 +1,15 @@
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀
+▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀▀