Document support for batches (#48)

Author: adrhill

README.md

@@ -9,7 +9,7 @@ ___
 
 Explainable AI in Julia using [Flux.jl](https://fluxml.ai).
 
-This package implements interpretability methods and visualizations for neural networks, similar to [Captum](https://github.com/pytorch/captum) for PyTorch and [iNNvestigate](https://github.com/albermax/innvestigate) for Keras models. 
+This package implements interpretability methods and visualizations for neural networks, similar to [Captum][captum-repo] and [Zennit][zennit-repo] for PyTorch and [iNNvestigate][innvestigate-repo] for Keras models. 
 
 ## Installation 
 To install this package and its dependencies, open the Julia REPL and run 
@@ -20,22 +20,31 @@ julia> ]add ExplainableAI
 ⚠️ This package is still in early development, expect breaking changes. ⚠️
 
 ## Example
-Let's use LRP to explain why an image of a cat gets classified as a cat:
+Let's use LRP to explain why an MNIST digit gets classified as a 9 using a small pre-trained LeNet5 model.
+If you want to follow along, the model can be found [here][model-bson-url].
 ```julia
 using ExplainableAI
 using Flux
-using Metalhead
+using MLDatasets
+using BSON: @load
 
 # Load model
-vgg = VGG19()
-model = strip_softmax(vgg.layers)
+@load "model.bson" model
+model = strip_softmax(model)
+
+# Load input
+x, _ = MNIST.testdata(Float32, 10)
+input = reshape(x, 28, 28, 1, :)   # reshape to WHCN format
 
 # Run XAI method
 analyzer = LRP(model)
-expl = analyze(img, analyzer)
+expl = analyze(input, analyzer)    # or: expl = analyzer(input)
 
 # Show heatmap
 heatmap(expl)
+
+# Or analyze & show heatmap directly
+heatmap(input, analyzer)
 ```
 ![][heatmap]
 
@@ -69,7 +78,7 @@ Contributions are welcome!
 > Adrian Hill acknowledges support by the Federal Ministry of Education and Research (BMBF) for the Berlin Institute for the Foundations of Learning and Data (BIFOLD) (01IS18037A).
 
 [banner-img]: https://raw.githubusercontent.com/adrhill/ExplainableAI.jl/gh-pages/assets/banner.png
-[heatmap]: https://raw.githubusercontent.com/adrhill/ExplainableAI.jl/gh-pages/assets/heatmap.png
+[heatmap]: https://raw.githubusercontent.com/adrhill/ExplainableAI.jl/gh-pages/assets/mnist9.png
 
 [docs-stab-img]: https://img.shields.io/badge/docs-stable-blue.svg
 [docs-stab-url]: https://adrhill.github.io/ExplainableAI.jl/stable
@@ -88,3 +97,9 @@ Contributions are welcome!
 
 [doi-img]: https://zenodo.org/badge/337430397.svg
 [doi-url]: https://zenodo.org/badge/latestdoi/337430397
+
+[model-bson-url]: https://github.com/adrhill/ExplainableAI.jl/blob/master/docs/src/model.bson
+
+[captum-repo]: https://github.com/pytorch/captum
+[zennit-repo]: https://github.com/chr5tphr/zennit
+[innvestigate-repo]: https://github.com/albermax/innvestigate

docs/literate/example.jl

@@ -48,7 +48,7 @@ input = reshape(x, 28, 28, 1, :);
 # ## Calling the analyzer
 # We can now select an analyzer of our choice
 # and call [`analyze`](@ref) to get an `Explanation`:
-analyzer = LRPZero(model)
+analyzer = LRP(model)
 expl = analyze(input, analyzer);
 
 # This `Explanation` bundles the following data:
@@ -79,6 +79,26 @@ heatmap(input, analyzer, 5)
 #md #     expl = analyze(img, analyzer, 5)
 #md #     ```
 
+# ## Input batches
+# ExplainableAI also supports input batches:
+batchsize = 100
+xs, _ = MNIST.testdata(Float32, 1:batchsize)
+batch = reshape(xs, 28, 28, 1, :) # reshape to WHCN format
+expl_batch = analyze(batch, analyzer);
+
+# This will once again return a single `Explanation` `expl_batch` for the entire batch.
+# Calling `heatmap` on `expl_batch` will detect the batch dimension and return a vector of heatmaps.
+#
+# Let's check if the digit at `index = 10` still matches.
+hs = heatmap(expl_batch)
+hs[index]
+
+# JuliaImages' `mosaic` can be used to return a tiled view of all heatmaps:
+mosaic(hs; nrow=10)
+
+# We can also evaluate a batch with respect to a specific output neuron, e.g. for the digit zero at index `1`:
+mosaic(heatmap(batch, analyzer, 1); nrow=10)
+
 # ## Automatic heatmap presets
 # Currently, the following analyzers are implemented:
 #
@@ -111,4 +131,7 @@ heatmap(expl; cs=ColorSchemes.jet)
 #
 heatmap(expl; reduce=:sum, normalize=:extrema, cs=ColorSchemes.inferno)
 
+# This also works with batches
+mosaic(heatmap(expl_batch; normalize=:extrema, cs=ColorSchemes.inferno); nrow=10)
+
 # For the full list of keyword arguments, refer to the [`heatmap`](@ref) documentation.