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.8282537844031334 0.5321771798936737 … 0.5140143418634737 0.749287496391157; 0.935080407055553 0.1665326583336414 … 0.20416118095868918 0.020811576749240834; … ; 0.8592653751623502 0.23819072320329948 … 0.9319849972429548 0.8403241407120906; 0.21148000071600215 0.2370467935429954 … 0.24588982844890683 0.5324104543759597;;;;], 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.578012; -14.809775; … ; -0.22507226; 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.