Example Implementations

The following examples demonstrate the implementation of XAI methods using the XAIBase.jl interface. To evaluate our methods, we load a small, pre-trained LeNet5 model and the MNIST dataset:

using Flux
using BSON

model = BSON.load("model.bson", @__MODULE__)[:model] # load pre-trained LeNet-5 model

Chain(
  Conv((5, 5), 1 => 6, relu),           # 156 parameters
  MaxPool((2, 2)),
  Conv((5, 5), 6 => 16, relu),          # 2_416 parameters
  MaxPool((2, 2)),
  Flux.flatten,
  Dense(256 => 120, relu),              # 30_840 parameters
  Dense(120 => 84, relu),               # 10_164 parameters
  Dense(84 => 10),                      # 850 parameters
)                   # Total: 10 arrays, 44_426 parameters, 174.344 KiB.

using MLDatasets
using ImageCore, ImageIO, ImageShow

index = 10
x, y = MNIST(Float32, :test)[10]

# By convention in Flux.jl, the input needs to be resized to WHCN format
# by adding a color channel and batch dimensions.
input = reshape(x, 28, 28, 1, :);

convert2image(MNIST, x)

Example 1: Random explanation

To get started, we implement a nonsensical method that returns a random explanation in the shape of the input.

using XAIBase
import XAIBase: call_analyzer

struct RandomAnalyzer{M} <: AbstractXAIMethod
    model::M
end

function call_analyzer(input, method::RandomAnalyzer, output_selector::AbstractOutputSelector; kwargs...)
    output = method.model(input)
    output_selection = output_selector(output)

    val = rand(size(input)...)
    return Explanation(val, input, output, output_selection, :RandomAnalyzer, :sensitivity, nothing)
end

call_analyzer (generic function with 2 methods)

We can directly use XAIBase's analyze function to compute the random explanation:

analyzer = RandomAnalyzer(model)
expl = analyze(input, analyzer)

Explanation{Array{Float64, 4}, Array{Float32, 4}, Matrix{Float32}, Vector{CartesianIndex{2}}, Nothing}([0.7273925748723502 0.690652989164023 … 0.06244825968041412 0.7485865198891324; 0.39566743212040767 0.63849859266577 … 0.4814987142217383 0.42002066847604447; … ; 0.7394244019133709 0.9682783414945164 … 0.18283448357566745 0.8207014650011124; 0.86085989925645 0.5320034672864322 … 0.35427691037156517 0.17150778975221426;;;;], Float32[0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0; … ; 0.0 0.0 … 0.0 0.0; 0.0 0.0 … 0.0 0.0;;;;], Float32[-17.578014; -14.809775; … ; -0.22507292; 21.765005;;], CartesianIndex{2}[CartesianIndex(10, 1)], :RandomAnalyzer, :sensitivity, nothing)

Using either VisionHeatmaps.jl or TextHeatmaps.jl, which provide package extensions on XAIBase's Explanation type, we can visualize the explanations:

using VisionHeatmaps # load heatmapping functionality

heatmap(expl.val)

As expected, the explanation is just noise.

Example 2: Input sensitivity

In this second example, we naively reimplement the Gradient analyzer from ExplainableAI.jl.

using XAIBase
import XAIBase: call_analyzer

using Zygote: gradient

struct MyGradient{M} <: AbstractXAIMethod
    model::M
end

function call_analyzer(input, method::MyGradient, output_selector::AbstractOutputSelector; kwargs...)
    output = method.model(input)
    output_selection = output_selector(output)

    grad = gradient((x) -> only(method.model(x)[output_selection]), input)
    val = only(grad)
    return Explanation(val, input, output, output_selection, :MyGradient, :sensitivity, nothing)
end

call_analyzer (generic function with 3 methods)

Once again, we can directly use XAIBase's analyze and VisionHeatmaps' heatmap functions

using VisionHeatmaps

analyzer = MyGradient(model)
expl = analyze(input, analyzer)
heatmap(expl.val)

heatmap(expl.val, colorscheme=:twilight, reduce=:norm, rangescale=:centered)

and make use of all the features provided by the Julia-XAI ecosystem.