Add inplace updating lrp! rule calls and reuse gradient code (#38)

* Change `lrp` to `lrp!` that inplace updates Rₖ
* Faster LRP preallocation
* Reuse gradient method code with `gradient_wrt_input`
* Drop precompilation

Author: adrhill

benchmark/benchmarks.jl

@@ -1,7 +1,7 @@
 using BenchmarkTools
 using Flux
 using ExplainabilityMethods
-import ExplainabilityMethods: _modify_layer, lrp
+import ExplainabilityMethods: _modify_layer, lrp!
 
 on_CI = haskey(ENV, "GITHUB_ACTIONS")
 
@@ -44,7 +44,7 @@ struct TestWrapper{T}
 end
 (w::TestWrapper)(x) = w.layer(x)
 _modify_layer(r::AbstractLRPRule, w::TestWrapper) = _modify_layer(r, w.layer)
-lrp(rule::ZBoxRule, w::TestWrapper, aₖ, Rₖ₊₁) = lrp(rule, w.layer, aₖ, Rₖ₊₁)
+lrp!(rule::ZBoxRule, w::TestWrapper, Rₖ, aₖ, Rₖ₊₁) = lrp!(rule, w.layer, Rₖ, aₖ, Rₖ₊₁)
 
 # generate input for conv layers
 insize = (64, 64, 3, 1)
@@ -69,11 +69,11 @@ rules = Dict(
 SUITE["Layer"] = BenchmarkGroup([k for k in keys(layers)])
 for (layername, (layer, aₖ)) in layers
     SUITE["Layer"][layername] = BenchmarkGroup([k for k in keys(rules)])
-
+    Rₖ = similar(aₖ)
     Rₖ₊₁ = layer(aₖ)
     for (rulename, rule) in rules
-        SUITE["Layer"][layername][rulename] = @benchmarkable lrp(
-            $(rule), $(layer), $(aₖ), $(Rₖ₊₁)
+        SUITE["Layer"][layername][rulename] = @benchmarkable lrp!(
+            $(rule), $(layer), $(Rₖ), $(aₖ), $(Rₖ₊₁)
         )
     end
 end

src/ExplainabilityMethods.jl

@@ -1,13 +1,16 @@
 module ExplainabilityMethods
 
+using Base.Iterators
 using LinearAlgebra
 using Flux
 using Zygote
-using ColorSchemes
-using ImageCore
-using Base.Iterators
 using Tullio
 
+# Heatmapping:
+using ImageCore
+using ColorSchemes
+
+# Model checks:
 using Markdown
 using PrettyTables
 
@@ -20,8 +23,6 @@ include("lrp_checks.jl")
 include("lrp_rules.jl")
 include("lrp.jl")
 include("heatmap.jl")
-include("precompile.jl")
-_precompile_()
 
 export analyze
 
@@ -34,7 +35,7 @@ export LRP, LRPZero, LRPEpsilon, LRPGamma
 export AbstractLRPRule
 export LRP_CONFIG
 export ZeroRule, EpsilonRule, GammaRule, ZBoxRule
-export lrp, modify_params, modify_denominator
+export modify_params, modify_denominator
 export check_model
 
 # heatmapping

src/gradient.jl

@@ -1,3 +1,7 @@
+function gradient_wrt_input(model, input::T, output_neuron)::T where {T}
+    return only(gradient((in) -> model(in)[output_neuron], input))
+end
+
 """
     Gradient(model)
 
@@ -10,8 +14,8 @@ end
 function (analyzer::Gradient)(input, ns::AbstractNeuronSelector)
     output = analyzer.model(input)
     output_neuron = ns(output)
-    attr = gradient((in) -> analyzer.model(in)[output_neuron], input)[1]
-    return Explanation(attr, output, output_neuron, :Gradient, Nothing)
+    grad = gradient_wrt_input(analyzer.model, input, output_neuron)
+    return Explanation(grad, output, output_neuron, :Gradient, Nothing)
 end
 
 """
@@ -29,6 +33,6 @@ end
 function (analyzer::InputTimesGradient)(input, ns::AbstractNeuronSelector)
     output = analyzer.model(input)
     output_neuron = ns(output)
-    attr = input .* gradient((in) -> analyzer.model(in)[output_neuron], input)[1]
+    attr = input .* gradient_wrt_input(analyzer.model, input, output_neuron)
     return Explanation(attr, output, output_neuron, :InputTimesGradient, Nothing)
 end

src/lrp.jl

@@ -47,23 +47,24 @@ LRPEpsilon(model::Chain; kwargs...) = LRP(model, EpsilonRule(); kwargs...)
 LRPGamma(model::Chain; kwargs...) = LRP(model, GammaRule(); kwargs...)
 
 # The call to the LRP analyzer.
-function (analyzer::LRP)(input, ns::AbstractNeuronSelector; layerwise_relevances=false)
+function (analyzer::LRP)(
+    input::AbstractArray{T}, ns::AbstractNeuronSelector; layerwise_relevances=false
+) where {T}
     layers = analyzer.model.layers
-    acts = Vector{Any}([input])
-    # Forward pass through layers, keeping track of activations
-    for layer in layers
-        append!(acts, [layer(acts[end])])
-    end
-    rels = deepcopy(acts) # allocate arrays
+    # Compute layerwise activations on forward pass through model:
+    acts = [input, Flux.activations(analyzer.model, input)...]
+
+    # Allocate array for layerwise relevances:
+    rels = similar.(acts)
 
     # Mask output neuron
     output_neuron = ns(acts[end])
-    rels[end] *= 0
+    rels[end] .= zero(T)
     rels[end][output_neuron] = acts[end][output_neuron]
 
     # Backward pass through layers, applying LRP rules
     for (i, rule) in Iterators.reverse(enumerate(analyzer.rules))
-        rels[i] .= lrp(rule, layers[i], acts[i], rels[i + 1])
+        lrp!(rule, layers[i], rels[i], acts[i], rels[i + 1]) # inplace update rels[i]
     end
 
     return Explanation(

src/lrp_rules.jl

@@ -4,13 +4,13 @@
 # can be implemented by dispatching on the functions `modify_params` & `modify_denominator`,
 # which make use of the generalized LRP implementation shown in [1].
 #
-# If the relevance propagation falls outside of this scheme, custom functions
+# If the relevance propagation falls outside of this scheme, custom low-level functions
 # ```julia
-# (::MyLRPRule)(layer, aₖ, Rₖ₊₁) = ...
-# (::MyLRPRule)(layer::MyLayer, aₖ, Rₖ₊₁) = ...
-# (::AbstractLRPRule)(layer::MyLayer, aₖ, Rₖ₊₁) = ...
+# lrp!(::MyLRPRule, layer, Rₖ, aₖ, Rₖ₊₁) = ...
+# lrp!(::MyLRPRule, layer::MyLayer, Rₖ, aₖ, Rₖ₊₁) = ...
+# lrp!(::AbstractLRPRule, layer::MyLayer, Rₖ, aₖ, Rₖ₊₁) = ...
 # ```
-# that return `Rₖ` can be implemented.
+# that inplace-update `Rₖ` can be implemented.
 # This is used for the ZBoxRule and for faster computations on common layers.
 #
 # References:
@@ -22,12 +22,13 @@ abstract type AbstractLRPRule end
 # This is the generic relevance propagation rule which is used for the 0, γ and ϵ rules.
 # It can be extended for new rules via `modify_denominator` and `modify_params`.
 # Since it uses autodiff, it is used as a fallback for layer types without custom implementation.
-function lrp(rule::R, layer::L, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule,L}
-    return lrp_autodiff(rule, layer, aₖ, Rₖ₊₁)
+function lrp!(rule::R, layer::L, Rₖ, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule,L}
+    lrp_autodiff!(rule, layer, Rₖ, aₖ, Rₖ₊₁)
+    return nothing
 end
 
-function lrp_autodiff(
-    rule::R, layer::L, aₖ::T1, Rₖ₊₁::T2
+function lrp_autodiff!(
+     rule::R, layer::L, Rₖ::T1, aₖ::T1, Rₖ₊₁::T2
 ) where {R<:AbstractLRPRule,L,T1,T2}
     layerᵨ = _modify_layer(rule, layer)
     c::T1 = only(
@@ -37,23 +38,26 @@ function lrp_autodiff(
             z ⋅ s
         end,
     )
-    return aₖ .* c # Rₖ
+    Rₖ .= aₖ .* c
+    return nothing
 end
 
 # For linear layer types such as Dense layers, using autodiff is overkill.
-function lrp(rule::R, layer::Dense, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule}
-    return lrp_dense(rule, layer, aₖ, Rₖ₊₁)
+function lrp!(rule::R, layer::Dense, Rₖ,  aₖ, Rₖ₊₁) where {R<:AbstractLRPRule}
+    lrp_dense!(rule, layer, Rₖ, aₖ, Rₖ₊₁)
+    return nothing
 end
 
-function lrp_dense(rule::R, l, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule}
+function lrp_dense!(rule::R, l, Rₖ, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule}
     ρW, ρb = modify_params(rule, get_params(l)...)
     ãₖ₊₁ = modify_denominator(rule, ρW * aₖ + ρb)
-    return @tullio Rₖ[j] := aₖ[j] * ρW[k, j] / ãₖ₊₁[k] * Rₖ₊₁[k]
+    @tullio Rₖ[j] = aₖ[j] * ρW[k, j] / ãₖ₊₁[k] * Rₖ₊₁[k]
+    return nothing
 end
 
 # Other special cases that are dispatched on layer type:
-lrp(::AbstractLRPRule, ::DropoutLayer, aₖ, Rₖ₊₁) = Rₖ₊₁
-lrp(::AbstractLRPRule, ::ReshapingLayer, aₖ, Rₖ₊₁) = reshape(Rₖ₊₁, size(aₖ))
+lrp!(::AbstractLRPRule, ::DropoutLayer, Rₖ, aₖ, Rₖ₊₁) = (Rₖ .= Rₖ₊₁)
+lrp!(::AbstractLRPRule, ::ReshapingLayer, Rₖ, aₖ, Rₖ₊₁) = (Rₖ .= reshape(Rₖ₊₁, size(aₖ)))
 
 # To implement new rules, we can define two custom functions `modify_params` and `modify_denominator`.
 # If this isn't done, the following fallbacks are used by default:
@@ -125,10 +129,10 @@ Commonly used on the first layer for pixel input.
 struct ZBoxRule <: AbstractLRPRule end
 
 # The ZBoxRule requires its own implementation of relevance propagation.
-lrp(::ZBoxRule, layer::Dense, aₖ, Rₖ₊₁) = lrp_zbox(layer, aₖ, Rₖ₊₁)
-lrp(::ZBoxRule, layer::Conv, aₖ, Rₖ₊₁) = lrp_zbox(layer, aₖ, Rₖ₊₁)
+lrp!(::ZBoxRule, layer::Dense, Rₖ, aₖ, Rₖ₊₁) = lrp_zbox!(layer, Rₖ, aₖ, Rₖ₊₁)
+lrp!(::ZBoxRule, layer::Conv, Rₖ, aₖ, Rₖ₊₁) = lrp_zbox!(layer, Rₖ, aₖ, Rₖ₊₁)
 
-function lrp_zbox(layer::L, aₖ::T1, Rₖ₊₁::T2) where {L,T1,T2}
+function lrp_zbox!(layer::L, Rₖ::T1, aₖ::T1, Rₖ₊₁::T2) where {L,T1,T2}
     W, b = get_params(layer)
     l, h = fill.(extrema(aₖ), (size(aₖ),))
 
@@ -144,5 +148,6 @@ function lrp_zbox(layer::L, aₖ::T1, Rₖ₊₁::T2) where {L,T1,T2}
         s = Zygote.@ignore safedivide(Rₖ₊₁, z; eps=1e-9)
         z ⋅ s
     end
-    return aₖ .* c + l .* cₗ + h .* cₕ # Rₖ from backward pass
+    Rₖ .= aₖ .* c + l .* cₗ + h .* cₕ
+    return nothing
 end

src/precompile.jl

@@ -1,6 +1,6 @@
 using ExplainabilityMethods
 using ExplainabilityMethods: modify_params
-import ExplainabilityMethods: _modify_layer, lrp
+import ExplainabilityMethods: _modify_layer, lrp!
 using Flux
 using LinearAlgebra
 using ReferenceTests
@@ -25,7 +25,9 @@ const RULES = Dict(
     Rₖ = [17 / 90, 316 / 675] # expected output
 
     layer = Dense(W, b, relu)
-    @test lrp(rule, layer, aₖ, Rₖ₊₁) ≈ Rₖ
+    R̂ₖ = similar(aₖ) # will be inplace updated
+    @inferred lrp!(rule, layer, R̂ₖ, aₖ, Rₖ₊₁)
+    @test R̂ₖ ≈ Rₖ
 
     ## Pooling layer
     Rₖ₊₁ = Float32.([1 2; 3 4]//30)
@@ -38,7 +40,9 @@ const RULES = Dict(
     Rₖ = reshape(repeat(Rₖ, 1, 3), 3, 3, 3, 1)
 
     layer = MaxPool((2, 2); stride=(1, 1))
-    @test lrp(rule, layer, aₖ, Rₖ₊₁) ≈ Rₖ
+    R̂ₖ = similar(aₖ) # will be inplace updated
+    @inferred lrp!(rule, layer, R̂ₖ, aₖ, Rₖ₊₁)
+    @test R̂ₖ ≈ Rₖ
 end
 
 # Fixed pseudo-random numbers
@@ -69,7 +73,8 @@ layers = Dict(
             for (layername, layer) in layers
                 @testset "$layername" begin
                     Rₖ₊₁ = layer(aₖ)
-                    Rₖ = @inferred lrp(rule, layer, aₖ, Rₖ₊₁)
+                    Rₖ = similar(aₖ)
+                    @inferred lrp!(rule, layer, Rₖ, aₖ, Rₖ₊₁)
 
                     @test typeof(Rₖ) == typeof(aₖ)
                     @test size(Rₖ) == size(aₖ)
@@ -110,14 +115,17 @@ equalpairs = Dict( # these pairs of layers are all equal
                     l1, l2 = layers
                     Rₖ₊₁ = l1(aₖ)
                     @test Rₖ₊₁ == l2(aₖ)
-                    Rₖ = @inferred lrp(rule, l1, aₖ, Rₖ₊₁)
-                    @test Rₖ == lrp(rule, l2, aₖ, Rₖ₊₁)
+                    Rₖ1 = similar(aₖ)
+                    Rₖ2 = similar(aₖ)
+                    @inferred lrp!(rule, l1, Rₖ1, aₖ, Rₖ₊₁)
+                    @inferred lrp!(rule, l2, Rₖ2, aₖ, Rₖ₊₁)
+                    @test Rₖ1 == Rₖ2
 
-                    @test typeof(Rₖ) == typeof(aₖ)
-                    @test size(Rₖ) == size(aₖ)
+                    @test typeof(Rₖ1) == typeof(aₖ)
+                    @test size(Rₖ1) == size(aₖ)
 
                     @test_reference "references/rules/$rulename/$layername.jld2" Dict(
-                        "R" => Rₖ
+                        "R" => Rₖ1
                     ) by = (r, a) -> isapprox(r["R"], a["R"]; rtol=0.02)
                 end
             end
@@ -143,7 +151,8 @@ layers = Dict(
             for (layername, layer) in layers
                 @testset "$layername" begin
                     Rₖ₊₁ = layer(aₖ)
-                    Rₖ = @inferred lrp(rule, layer, aₖ, Rₖ₊₁)
+                    Rₖ = similar(aₖ)
+                    @inferred lrp!(rule, layer, Rₖ, aₖ, Rₖ₊₁)
 
                     @test typeof(Rₖ) == typeof(aₖ)
                     @test size(Rₖ) == size(aₖ)
@@ -164,7 +173,7 @@ struct TestWrapper{T}
 end
 (w::TestWrapper)(x) = w.layer(x)
 _modify_layer(r::AbstractLRPRule, w::TestWrapper) = _modify_layer(r, w.layer)
-lrp(rule::ZBoxRule, w::TestWrapper, aₖ, Rₖ₊₁) = lrp(rule, w.layer, aₖ, Rₖ₊₁)
+lrp!(rule::ZBoxRule, w::TestWrapper, Rₖ, aₖ, Rₖ₊₁) = lrp!(rule, w.layer, Rₖ, aₖ, Rₖ₊₁)
 
 layers = Dict(
     "Conv" => (Conv((3, 3), 2 => 4; init=pseudorandn), aₖ),
@@ -179,7 +188,8 @@ layers = Dict(
                 @testset "$layername" begin
                     wrapped_layer = TestWrapper(layer)
                     Rₖ₊₁ = wrapped_layer(aₖ)
-                    Rₖ = @inferred lrp(rule, wrapped_layer, aₖ, Rₖ₊₁)
+                    Rₖ = similar(aₖ)
+                    @inferred lrp!(rule, wrapped_layer, Rₖ, aₖ, Rₖ₊₁)
 
                     @test typeof(Rₖ) == typeof(aₖ)
                     @test size(Rₖ) == size(aₖ)