Fix ZBoxRule (#69)

Author: adrhill

benchmark/benchmarks.jl

@@ -9,11 +9,16 @@ on_CI = haskey(ENV, "GITHUB_ACTIONS")
 include("../test/vgg11.jl")
 vgg11 = VGG11(; pretrain=false)
 model = flatten_model(strip_softmax(vgg11.layers))
-img = rand(MersenneTwister(123), Float32, (224, 224, 3, 1))
+
+T = Float32
+img = rand(MersenneTwister(123), T, (224, 224, 3, 1))
 
 # Benchmark custom LRP composite
 function LRPCustom(model::Chain)
-    return LRP(model, [ZBoxRule(), repeat([GammaRule()], length(model.layers) - 1)...])
+    return LRP(
+        model,
+        [ZBoxRule(zero(T), oneunit(T)), repeat([GammaRule()], length(model.layers) - 1)...],
+    )
 end
 
 # Use one representative algorithm of each type
@@ -53,20 +58,20 @@ lrp!(rule::ZBoxRule, w::TestWrapper, Rₖ, aₖ, Rₖ₊₁) = lrp!(rule, w.laye
 insize = (64, 64, 3, 1)
 in_dense = 500
 out_dense = 100
-aₖ = randn(Float32, insize)
+aₖ = randn(T, insize)
 
 layers = Dict(
     "MaxPool" => (MaxPool((3, 3); pad=0), aₖ),
     "Conv" => (Conv((3, 3), 3 => 2), aₖ),
-    "Dense" => (Dense(in_dense, out_dense, relu), randn(Float32, in_dense, 1)),
+    "Dense" => (Dense(in_dense, out_dense, relu), randn(T, in_dense, 1)),
     "WrappedDense" =>
-        (TestWrapper(Dense(in_dense, out_dense, relu)), randn(Float32, in_dense, 1)),
+        (TestWrapper(Dense(in_dense, out_dense, relu)), randn(T, in_dense, 1)),
 )
 rules = Dict(
     "ZeroRule" => ZeroRule(),
     "EpsilonRule" => EpsilonRule(),
     "GammaRule" => GammaRule(),
-    "ZBoxRule" => ZBoxRule(),
+    "ZBoxRule" => ZBoxRule(zero(T), oneunit(T)),
 )
 
 SUITE["Layer"] = BenchmarkGroup([k for k in keys(layers)])

docs/literate/advanced_lrp.jl

@@ -23,7 +23,7 @@ input = reshape(x, 28, 28, 1, :);
 # 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:
 rules = [
-    ZBoxRule(),
+    ZBoxRule(0.0f0, 1.0f0),
     GammaRule(),
     GammaRule(),
     EpsilonRule(),

src/lrp_rules.jl

@@ -107,21 +107,41 @@ end
 modify_denominator(r::EpsilonRule, d) = stabilize_denom(d, r.ϵ)
 
 """
-    ZBoxRule()
+    ZBoxRule(low, high)
 
 Constructor for LRP-``z^{\\mathcal{B}}``-rule.
 Commonly used on the first layer for pixel input.
+
+The parameters `low` and `high` should be set to the lower and upper bounds of the input features,
+e.g. `0.0` and `1.0` for raw image data.
+It is also possible to provide two arrays of that match the input size.
 """
-struct ZBoxRule <: AbstractLRPRule end
+struct ZBoxRule{T} <: AbstractLRPRule
+    low::T
+    high::T
+end
 
 # The ZBoxRule requires its own implementation of relevance propagation.
-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ₖ₊₁)
+lrp!(r::ZBoxRule, layer::Dense, Rₖ, aₖ, Rₖ₊₁) = lrp_zbox!(r, layer, Rₖ, aₖ, Rₖ₊₁)
+lrp!(r::ZBoxRule, layer::Conv, Rₖ, aₖ, Rₖ₊₁) = lrp_zbox!(r, layer, Rₖ, aₖ, Rₖ₊₁)
+
+_zbox_bound(T, c::Real, in_size) = fill(convert(T, c), in_size)
+function _zbox_bound(T, A::AbstractArray, in_size)
+    size(A) != in_size && throw(
+        ArgumentError(
+            "Bounds `low`, `high` of ZBoxRule should either be scalar or match input size.",
+        ),
+    )
+    return convert.(T, A)
+end
 
-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ₖ),))
+function lrp_zbox!(r::ZBoxRule, layer::L, Rₖ::T1, aₖ::T1, Rₖ₊₁::T2) where {L,T1,T2}
+    T = eltype(aₖ)
+    in_size = size(aₖ)
+    l = _zbox_bound(T, r.low, in_size)
+    h = _zbox_bound(T, r.high, in_size)
 
+    W, b = get_params(layer)
     layer⁺ = set_params(layer, max.(0, W), max.(0, b)) # W⁺, b⁺
     layer⁻ = set_params(layer, min.(0, W), min.(0, b)) # W⁻, b⁻
 

test/references/rules/ZBoxRule/Conv.jld2

@@ -11,7 +11,7 @@ const RULES = Dict(
     "ZeroRule" => ZeroRule(),
     "EpsilonRule" => EpsilonRule(),
     "GammaRule" => GammaRule(),
-    "ZBoxRule" => ZBoxRule(),
+    "ZBoxRule" => ZBoxRule(0.0f0, 1.0f0),
 )
 
 ## Hand-written tests

test/references/rules/ZBoxRule/Dense_identity.jld2

@@ -26,7 +26,9 @@ vgg11 = VGG11(; pretrain=false)
 model = flatten_model(strip_softmax(vgg11.layers))
 
 function LRPCustom(model::Chain)
-    return LRP(model, [ZBoxRule(), repeat([GammaRule()], length(model.layers) - 1)...])
+    return LRP(
+        model, [ZBoxRule(0.0f0, 1.0f0), repeat([GammaRule()], length(model.layers) - 1)...]
+    )
 end
 
 function test_vgg11(name, method; kwargs...)