Add default composites (#87)

* rename `RuleMap` to `TypeRule`

* add `FirstLayerTypeRule` and `LastLayerTypeRule`

* Add default composites

* Document default composites

* Add compat entries for test dependencies

Author: adrhill

Project.toml

@@ -27,14 +27,3 @@ PrettyTables = "1"
 Tullio = "0.3"
 Zygote = "0.6"
 julia = "1.6"
-
-[extras]
-JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
-LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890"
-Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-ReferenceTests = "324d217c-45ce-50fc-942e-d289b448e8cf"
-Suppressor = "fd094767-a336-5f1f-9728-57cf17d0bbfb"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[targets]
-test = ["JLD2", "LoopVectorization", "Random", "ReferenceTests", "Suppressor", "Test"]

docs/literate/advanced_lrp.jl

@@ -19,7 +19,8 @@ index = 10
 x, _ = MNIST(Float32, :test)[10]
 input = reshape(x, 28, 28, 1, :);
 
-# ## Custom LRP composites
+# ## LRP composites
+# ### Custom composites
 # When creating an LRP-analyzer, we can assign individual rules to each layer.
 # The array of rules has to match the length of the Flux chain:
 rules = [
@@ -49,11 +50,11 @@ heatmap(input, analyzer)
 # To obtain the same set of rules as in the previous example, we can define
 composite = Composite(
     ZeroRule(),                         # default rule
-    GlobalRuleMap(
+    GlobalTypeRule(
         Conv => GammaRule(),       # apply GammaRule on all convolutional layers
         MaxPool => EpsilonRule(),  # apply EpsilonRule on all pooling-layers
     ),
-    FirstRule(ZBoxRule(0.0f0, 1.0f0)),  # apply ZBoxRule on the first layer
+    FirstLayerRule(ZBoxRule(0.0f0, 1.0f0)),  # apply ZBoxRule on the first layer
 )
 
 analyzer = LRP(model, composite)        # construct LRP analyzer from composite
@@ -63,24 +64,31 @@ heatmap(input, analyzer)
 analyzer.rules # show rules
 
 # ### Composite primitives
-# The following sets of primitives can used to construct a [`Composite`](@ref).
+# The following [Composite primitives](@ref composite_primitive_api) can used to construct a [`Composite`](@ref).
 #
 # To apply a single rule, use:
 # * [`LayerRule`](@ref) to apply a rule to the `n`-th layer of a model
-# * [`GlobalRule`](@ref) to apply a rule to all layers of a model
-# * [`RangeRule`](@ref) to apply a rule to a positional range of layers of a model
-# * [`FirstRule`](@ref) to apply a rule to the first layer of a model
-# * [`LastRule`](@ref) to apply a rule to the last layer of a model
+# * [`GlobalRule`](@ref) to apply a rule to all layers
+# * [`RangeRule`](@ref) to apply a rule to a positional range of layers
+# * [`FirstLayerRule`](@ref) to apply a rule to the first layer
+# * [`LastLayerRule`](@ref) to apply a rule to the last layer
 #
-# To apply a set of rules to multiple layers, use:
-# * [`GlobalRuleMap`](@ref) to apply a dictionary that maps layer types to LRP-rules
-# * [`RangeRuleMap`](@ref) for a `RuleMap` on generalized ranges
-# * [`FirstNRuleMap`](@ref) for a `RuleMap` on the first `n` layers of a model
-# * [`LastNRuleMap`](@ref) for a `RuleMap` on the last `n` layers
+# To apply a set of rules to layers based on their type, use:
+# * [`GlobalTypeRule`](@ref) to apply a dictionary that maps layer types to LRP-rules
+# * [`RangeTypeRule`](@ref) for a `TypeRule` on generalized ranges
+# * [`FirstLayerTypeRule`](@ref) for a `TypeRule` on the first layer of a model
+# * [`LastLayerTypeRule`](@ref) for a `TypeRule` on the last layer
+# * [`FirstNTypeRule`](@ref) for a `TypeRule` on the first `n` layers
+# * [`LastNTypeRule`](@ref) for a `TypeRule` on the last `n` layers
 #
 # Primitives are called sequentially in the order the `Composite` was created with
 # and overwrite rules specified by previous primitives.
 
+# ### Default composites
+# A list of implemented default composites can be found under
+# [Default composites](@ref default_composite_api) in the API reference, e.g. [`EpsilonPlusFlat`](@ref):
+EpsilonPlusFlat()
+
 # ## Custom LRP rules
 # Let's define a rule that modifies the weights and biases of our layer on the forward pass.
 # The rule has to be of type `AbstractLRPRule`.

docs/src/api.md

@@ -50,21 +50,33 @@ LRP_CONFIG.supports_activation
 Composite
 ```
 
+### [Composite primitives](@id composite_primitive_api)
 Composite primitives that apply a single rule:
 ```@docs
 LayerRule
 GlobalRule
 RangeRule
-FirstRule
-LastRule
+FirstLayerRule
+LastLayerRule
 ```
 
 Composite primitives that apply a set of rules to multiple layers:
 ```@docs
-GlobalRuleMap
-RangeRuleMap
-FirstNRuleMap
-LastNRuleMap
+GlobalTypeRule
+RangeTypeRule
+FirstLayerTypeRule
+LastLayerTypeRule
+FirstNTypeRule
+LastNTypeRule
+```
+
+### [Default composites](@id default_composite_api)
+```@docs
+EpsilonGammaBox
+EpsilonPlus
+EpsilonAlpha2Beta1
+EpsilonPlusFlat
+EpsilonAlpha2Beta1Flat
 ```
 
 # Utilities

src/ExplainableAI.jl

@@ -53,8 +53,13 @@ export check_model
 
 # LRP composites
 export Composite, AbstractCompositePrimitive
-export LayerRule, GlobalRule, RangeRule, FirstRule, LastRule
-export GlobalRuleMap, RangeRuleMap, FirstNRuleMap, LastNRuleMap
+export LayerRule, GlobalRule, RangeRule, FirstLayerRule, LastLayerRule
+export GlobalTypeRule, RangeTypeRule, FirstLayerTypeRule, LastLayerTypeRule
+export FirstNTypeRule, LastNTypeRule
+# Default composites
+export EpsilonGammaBox, EpsilonPlus, EpsilonAlpha2Beta1, EpsilonPlusFlat
+export EpsilonAlpha2Beta1Flat
+# Useful type unions
 export ConvLayer, PoolingLayer, DropoutLayer, ReshapingLayer
 
 # heatmapping