Example Implementations

The following examples demonstrate the implementation of feature-attribution 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.359 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 attribution

To get started, we implement a nonsensical method that returns a random attribution 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 Attribution(val, input, output, output_selection, NormPooling())
end

call_analyzer (generic function with 2 methods)

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

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

Attribution{Array{Float64, 4}, Array{Float32, 4}, Matrix{Float32}, Vector{CartesianIndex{2}}, NormPooling, Nothing}([0.6796528119017861 0.6040942809515998 … 0.5271368640565308 0.67155018924836; 0.4019314629512123 0.10906397603016005 … 0.09414176538390184 0.23489586880900049; … ; 0.2509689141597298 0.532704095003402 … 0.8369979388102061 0.9768284100560938; 0.8619197044307623 0.2502307760895385 … 0.05779316327863526 0.3439083451500261;;;;], 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)], NormPooling(), nothing)

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

using VisionHeatmaps # load heatmapping functionality

heatmap(attr.val)

As expected, the attribution 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 Attribution(val, input, output, output_selection, NormPooling())
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)
attr = analyze(input, analyzer)
heatmap(attr.val)

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