Make LRP rules type inferable (#33)

* introduce `lrp` for rule calls
* help type inference and specialization in rule calls  
* test type stability with `@inferred`
* add precompilation of LRP rules

Author: adrhill

benchmark/benchmarks.jl

@@ -1,7 +1,7 @@
 using BenchmarkTools
 using Flux
 using ExplainabilityMethods
-import ExplainabilityMethods: _modify_layer
+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)
-(rule::ZBoxRule)(w::TestWrapper, aₖ, Rₖ₊₁) = rule(w.layer, aₖ, Rₖ₊₁)
+lrp(rule::ZBoxRule, w::TestWrapper, aₖ, Rₖ₊₁) = lrp(rule, w.layer, aₖ, Rₖ₊₁)
 
 # generate input for conv layers
 insize = (64, 64, 3, 1)
@@ -66,15 +66,13 @@ rules = Dict(
     "ZBoxRule" => ZBoxRule(),
 )
 
-test_rule(rule, layer, aₖ, Rₖ₊₁) = rule(layer, aₖ, Rₖ₊₁) # for use with @benchmarkable macro
-
 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ₖ₊₁ = layer(aₖ)
     for (rulename, rule) in rules
-        SUITE["Layer"][layername][rulename] = @benchmarkable test_rule(
+        SUITE["Layer"][layername][rulename] = @benchmarkable lrp(
             $(rule), $(layer), $(aₖ), $(Rₖ₊₁)
         )
     end

src/ExplainabilityMethods.jl

@@ -20,6 +20,8 @@ include("lrp_checks.jl")
 include("lrp_rules.jl")
 include("lrp.jl")
 include("heatmap.jl")
+include("precompile.jl")
+_precompile_()
 
 export analyze
 
@@ -32,7 +34,7 @@ export LRP, LRPZero, LRPEpsilon, LRPGamma
 export AbstractLRPRule
 export LRP_CONFIG
 export ZeroRule, EpsilonRule, GammaRule, ZBoxRule
-export modify_params, modify_denominator
+export lrp, modify_params, modify_denominator
 export check_model
 
 # heatmapping

src/lrp.jl

@@ -63,7 +63,7 @@ function (analyzer::LRP)(input, ns::AbstractNeuronSelector; layerwise_relevances
 
     # Backward pass through layers, applying LRP rules
     for (i, rule) in Iterators.reverse(enumerate(analyzer.rules))
-        rels[i] .= rule(layers[i], acts[i], rels[i + 1]) # Rₖ = rule(layer, aₖ, Rₖ₊₁)
+        rels[i] .= lrp(rule, layers[i], acts[i], rels[i + 1])
     end
 
     if layerwise_relevances

src/lrp_rules.jl

@@ -22,30 +22,38 @@ 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.
-(rule::AbstractLRPRule)(layer, aₖ, Rₖ₊₁) = lrp_autodiff(rule, layer, aₖ, Rₖ₊₁)
+function lrp(rule::R, layer::L, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule,L}
+    return lrp_autodiff(rule, layer, aₖ, Rₖ₊₁)
+end
 
-function lrp_autodiff(rule, layer, aₖ, Rₖ₊₁)
+function lrp_autodiff(
+    rule::R, layer::L, aₖ::T1, Rₖ₊₁::T2
+) where {R<:AbstractLRPRule,L,T1,T2}
     layerᵨ = _modify_layer(rule, layer)
-    function fwpass(a)
-        z = layerᵨ(a)
-        s = Zygote.dropgrad(Rₖ₊₁ ./ modify_denominator(rule, z))
-        return z ⋅ s
-    end
-    return aₖ .* gradient(fwpass, aₖ)[1] # Rₖ
+    c::T1 = only(
+        gradient(aₖ) do a
+            z::T2 = layerᵨ(a)
+            s = Zygote.@ignore Rₖ₊₁ ./ modify_denominator(rule, z)
+            z ⋅ s
+        end,
+    )
+    return aₖ .* c # Rₖ
 end
 
 # For linear layer types such as Dense layers, using autodiff is overkill.
-(rule::AbstractLRPRule)(layer::Dense, aₖ, Rₖ₊₁) = lrp_dense(rule, layer, aₖ, Rₖ₊₁)
+function lrp(rule::R, layer::Dense, aₖ, Rₖ₊₁) where {R<:AbstractLRPRule}
+    return lrp_dense(rule, layer, aₖ, Rₖ₊₁)
+end
 
-function lrp_dense(rule, l, aₖ, Rₖ₊₁)
+function lrp_dense(rule::R, l, 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]
 end
 
 # Other special cases that are dispatched on layer type:
-(::AbstractLRPRule)(::DropoutLayer, aₖ, Rₖ₊₁) = Rₖ₊₁
-(::AbstractLRPRule)(::ReshapingLayer, aₖ, Rₖ₊₁) = reshape(Rₖ₊₁, size(aₖ))
+lrp(::AbstractLRPRule, ::DropoutLayer, aₖ, Rₖ₊₁) = Rₖ₊₁
+lrp(::AbstractLRPRule, ::ReshapingLayer, aₖ, 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:
@@ -65,7 +73,7 @@ modify_denominator(::AbstractLRPRule, d) = stabilize_denom(d; eps=1.0f-9) # gene
 
 # This helper function applies `modify_params`:
 _modify_layer(::AbstractLRPRule, layer) = layer # skip layers without modify_params
-function _modify_layer(rule::AbstractLRPRule, layer::Union{Dense,Conv})
+function _modify_layer(rule::R, layer::L) where {R<:AbstractLRPRule,L<:Union{Dense,Conv}}
     return set_params(layer, modify_params(rule, get_params(layer)...)...)
 end
 
@@ -117,26 +125,24 @@ Commonly used on the first layer for pixel input.
 struct ZBoxRule <: AbstractLRPRule end
 
 # The ZBoxRule requires its own implementation of relevance propagation.
-(rule::ZBoxRule)(layer::Dense, aₖ, Rₖ₊₁) = lrp_zbox(layer, aₖ, Rₖ₊₁)
-(rule::ZBoxRule)(layer::Conv, aₖ, Rₖ₊₁) = lrp_zbox(layer, aₖ, Rₖ₊₁)
+lrp(::ZBoxRule, layer::Dense, aₖ, Rₖ₊₁) = lrp_zbox(layer, aₖ, Rₖ₊₁)
+lrp(::ZBoxRule, layer::Conv, aₖ, Rₖ₊₁) = lrp_zbox(layer, aₖ, Rₖ₊₁)
 
-function lrp_zbox(layer, aₖ, Rₖ₊₁)
+function lrp_zbox(layer::L, aₖ::T1, Rₖ₊₁::T2) where {L,T1,T2}
     W, b = get_params(layer)
     l, h = fill.(extrema(aₖ), (size(aₖ),))
 
     layer⁺ = set_params(layer, max.(0, W), max.(0, b)) # W⁺, b⁺
     layer⁻ = set_params(layer, min.(0, W), min.(0, b)) # W⁻, b⁻
 
-    # Forward pass
-    function fwpass(a, l, h)
-        f = layer(a)
-        f⁺ = layer⁺(l)
-        f⁻ = layer⁻(h)
+    c::T1, cₗ::T1, cₕ::T1 = gradient(aₖ, l, h) do a, l, h
+        f::T2 = layer(a)
+        f⁺::T2 = layer⁺(l)
+        f⁻::T2 = layer⁻(h)
 
         z = f - f⁺ - f⁻
-        s = Zygote.dropgrad(safedivide(Rₖ₊₁, z; eps=1e-9))
-        return z ⋅ s
+        s = Zygote.@ignore safedivide(Rₖ₊₁, z; eps=1e-9)
+        z ⋅ s
     end
-    c, cₗ, cₕ = gradient(fwpass, aₖ, l, h) # w.r.t. three inputs
     return aₖ .* c + l .* cₗ + h .* cₕ # Rₖ from backward pass
 end

src/precompile.jl

@@ -0,0 +1,44 @@
+macro warnpcfail(ex::Expr)
+    modl = __module__
+    file = __source__.file === nothing ? "?" : String(__source__.file)
+    line = __source__.line
+    quote
+        $(esc(ex)) || @warn """precompile directive $($(Expr(:quote, ex)))
+        failed. Please report an issue in $($modl) (after checking for duplicates) or remove this directive.""" _file =
+            $file _line = $line
+    end
+end
+
+function _precompile_()
+    eltypes = (Float32,)
+    ruletypes = (ZeroRule, EpsilonRule, GammaRule, ZBoxRule)
+    layertypes = (
+        Dense,
+        Conv,
+        MaxPool,
+        AdaptiveMaxPool,
+        GlobalMaxPool,
+        MeanPool,
+        AdaptiveMeanPool,
+        GlobalMeanPool,
+        DepthwiseConv,
+        ConvTranspose,
+        CrossCor,
+        Dropout,
+        AlphaDropout,
+        typeof(Flux.flatten),
+    )
+
+    for R in ruletypes
+        for T in eltypes
+            AT = Array{T}
+            @warnpcfail precompile(modify_denominator, (R, AT))
+            @warnpcfail precompile(modify_params, (R, AT, AT))
+
+            for L in layertypes
+                @warnpcfail precompile(_modify_layer, (R, L))
+                @warnpcfail precompile(lrp, (R, L, AT, AT))
+            end
+        end
+    end
+end

test/test_neuron_selection.jl

@@ -1,8 +1,8 @@
 using ExplainabilityMethods: MaxActivationNS, IndexNS
 
 A = [-2.1694243, 2.4023275, 0.99464744, -0.1514646, 1.0307171]
-ns1 = MaxActivationNS()
-ns2 = IndexNS(4)
+ns1 = @inferred MaxActivationNS()
+ns2 = @inferred IndexNS(4)
 
 @test ns1(A) == 2
 @test ns2(A) == 4

test/test_rules.jl

@@ -1,6 +1,6 @@
 using ExplainabilityMethods
 using ExplainabilityMethods: modify_params
-import ExplainabilityMethods: _modify_layer
+import ExplainabilityMethods: _modify_layer, lrp
 using Flux
 using LinearAlgebra
 using ReferenceTests
@@ -25,7 +25,7 @@ const RULES = Dict(
     Rₖ = [17 / 90, 316 / 675] # expected output
 
     layer = Dense(W, b, relu)
-    @test rule(layer, aₖ, Rₖ₊₁) ≈ Rₖ
+    @test lrp(rule, layer, aₖ, Rₖ₊₁) ≈ Rₖ
 
     ## Pooling layer
     Rₖ₊₁ = Float32.([1 2; 3 4]//30)
@@ -38,7 +38,7 @@ const RULES = Dict(
     Rₖ = reshape(repeat(Rₖ, 1, 3), 3, 3, 3, 1)
 
     layer = MaxPool((2, 2); stride=(1, 1))
-    @test rule(layer, aₖ, Rₖ₊₁) ≈ Rₖ
+    @test lrp(rule, layer, aₖ, Rₖ₊₁) ≈ Rₖ
 end
 
 # Fixed pseudo-random numbers
@@ -48,7 +48,7 @@ pseudorandn(dims...) = randn(MersenneTwister(123), T, dims...)
 ## Test individual rules
 @testset "modify_params" begin
     W, b = [1.0 -1.0; 2.0 0.0], [-1.0, 1.0]
-    ρW, ρb = modify_params(GammaRule(; γ=0.42), W, b)
+    ρW, ρb = @inferred modify_params(GammaRule(; γ=0.42), W, b)
     @test ρW ≈ [1.42 -1.0; 2.84 0.0]
     @test ρb ≈ [-1.0, 1.42]
 end
@@ -69,7 +69,7 @@ layers = Dict(
             for (layername, layer) in layers
                 @testset "$layername" begin
                     Rₖ₊₁ = layer(aₖ)
-                    Rₖ = rule(layer, aₖ, Rₖ₊₁)
+                    Rₖ = @inferred lrp(rule, layer, aₖ, Rₖ₊₁)
 
                     @test typeof(Rₖ) == typeof(aₖ)
                     @test size(Rₖ) == size(aₖ)
@@ -110,8 +110,8 @@ equalpairs = Dict( # these pairs of layers are all equal
                     l1, l2 = layers
                     Rₖ₊₁ = l1(aₖ)
                     @test Rₖ₊₁ == l2(aₖ)
-                    Rₖ = rule(l1, aₖ, Rₖ₊₁)
-                    @test Rₖ == rule(l2, aₖ, Rₖ₊₁)
+                    Rₖ = @inferred lrp(rule, l1, aₖ, Rₖ₊₁)
+                    @test Rₖ == lrp(rule, l2, aₖ, Rₖ₊₁)
 
                     @test typeof(Rₖ) == typeof(aₖ)
                     @test size(Rₖ) == size(aₖ)
@@ -143,7 +143,7 @@ layers = Dict(
             for (layername, layer) in layers
                 @testset "$layername" begin
                     Rₖ₊₁ = layer(aₖ)
-                    Rₖ = rule(layer, aₖ, Rₖ₊₁)
+                    Rₖ = @inferred lrp(rule, layer, aₖ, Rₖ₊₁)
 
                     @test typeof(Rₖ) == typeof(aₖ)
                     @test size(Rₖ) == size(aₖ)
@@ -158,13 +158,13 @@ layers = Dict(
 end
 
 ## Test custom layers & default AD fallback using the ZeroRule
-## Compare with references of non-wrapped layers
+# Compare with references of non-wrapped layers
 struct TestWrapper{T}
     layer::T
 end
 (w::TestWrapper)(x) = w.layer(x)
 _modify_layer(r::AbstractLRPRule, w::TestWrapper) = _modify_layer(r, w.layer)
-(rule::ZBoxRule)(w::TestWrapper, aₖ, Rₖ₊₁) = rule(w.layer, aₖ, Rₖ₊₁)
+lrp(rule::ZBoxRule, w::TestWrapper, aₖ, Rₖ₊₁) = lrp(rule, w.layer, aₖ, Rₖ₊₁)
 
 layers = Dict(
     "Conv" => (Conv((3, 3), 2 => 4; init=pseudorandn), aₖ),
@@ -179,7 +179,7 @@ layers = Dict(
                 @testset "$layername" begin
                     wrapped_layer = TestWrapper(layer)
                     Rₖ₊₁ = wrapped_layer(aₖ)
-                    Rₖ = rule(wrapped_layer, aₖ, Rₖ₊₁)
+                    Rₖ = @inferred lrp(rule, wrapped_layer, aₖ, Rₖ₊₁)
 
                     @test typeof(Rₖ) == typeof(aₖ)
                     @test size(Rₖ) == size(aₖ)