Add LRP model checks and support unknown layers via AD fallback

[adrhill]

closes #25

Model checks

Adds the check_model(:LRP, model) function to check models for compatibility with LRP. This function is called by default in the constructor of all LRP analyzers and returns a summary using PrettyTables.jl (shown in screenshots below).

These checks can be fully skipped using the LRP analyzer keyword argument skip_checks for users who know what they are doing. Alternatively, individual custom layers and activation functions can be registered:

LRP_CONFIG.supports_layer(::MyLayer) = true               # for structs
LRP_CONFIG.supports_layer(::typeof(mylayer)) = true       # for functions

LRP_CONFIG.supports_activation(::MyActivation) = true        # for structs
LRP_CONFIG.supports_activation(::typeof(myfunction)) = true  # for functions

Custom layer support

This registration mechanism opens up the possibility to define LRP rules on custom layers.
For this purpose, type restrictions on the default automatic differentiation fallback have been removed.

By default, the following AD fallback will be applied to registered but unknown layers:

function (rule::AbstractLRPRule)(layer, aₖ, Rₖ₊₁)
    layerᵨ = modify_layer(rule, layer)
    function fwpass(a)
        z = layerᵨ(a)
        s = Zygote.dropgrad(Rₖ₊₁ ./ modify_denominator(rule, z))
        return z ⋅ s
    end
    c = gradient(fwpass, aₖ)[1]
    Rₖ = aₖ .* c
    return Rₖ
end

Where by default modify_layer(rule, layer) = layer and modify_denominator(rule, z) = stabilize_denom(d; eps=1.0f-9). These can be extended by the user for any combination of rule and custom layer.

The user can also use multiple dispatch to implement fully custom rules using this interface:

function (rule::AbstractLRPRule)(layer::MyLayer, aₖ, Rₖ₊₁)
    # ...
    return Rₖ
end

Screenshots

Unknown/unsupported activation function

Unknown layer

[codecov-commenter]

Codecov Report

Merging #26 (55e361b) into master (614bdb7) will decrease coverage by 4.77%.
The diff coverage is 80.70%.

@@            Coverage Diff             @@
##           master      #26      +/-   ##
==========================================
- Coverage   95.73%   90.95%   -4.78%     
==========================================
  Files           8        9       +1     
  Lines         164      199      +35     
==========================================
+ Hits          157      181      +24     
- Misses          7       18      +11     

Impacted Files	Coverage Δ
src/lrp_checks.jl	69.44% <69.44%> (ø)
src/flux.jl	96.29% <100.00%> (ø)
src/lrp.jl	93.10% <100.00%> (-0.23%)	⬇️
src/lrp_rules.jl	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 614bdb7...55e361b. Read the comment docs.

[adrhill]

Only needs some updated doc strings and documentation!