From 12046e5b30047afbacb8fd57fbca0d979d0adfaf Mon Sep 17 00:00:00 2001 From: Adrian Hill Date: Thu, 10 Oct 2024 16:58:44 +0200 Subject: [PATCH] Update benchmarks and test them with PkgJogger (#175) --- benchmark/Project.toml | 3 +- benchmark/bench_jogger.jl | 47 +++++++++++++++++++ benchmark/benchmarks.jl | 98 ++------------------------------------- test/Project.toml | 2 + test/runtests.jl | 4 ++ test/test_batches.jl | 1 + test/test_benchmarks.jl | 4 ++ 7 files changed, 63 insertions(+), 96 deletions(-) create mode 100644 benchmark/bench_jogger.jl create mode 100644 test/test_benchmarks.jl diff --git a/benchmark/Project.toml b/benchmark/Project.toml index 6bb7d4e..d2a943a 100644 --- a/benchmark/Project.toml +++ b/benchmark/Project.toml @@ -2,9 +2,8 @@ BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" ExplainableAI = "4f1bc3e1-d60d-4ed0-9367-9bdff9846d3b" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" -LoopVectorization = "bdcacae8-1622-11e9-2a5c-532679323890" PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d" -Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc" +PkgJogger = "10150987-6cc1-4b76-abee-b1c1cbd91c01" [compat] BenchmarkTools = "1" diff --git a/benchmark/bench_jogger.jl b/benchmark/bench_jogger.jl new file mode 100644 index 0000000..e786975 --- /dev/null +++ b/benchmark/bench_jogger.jl @@ -0,0 +1,47 @@ +using BenchmarkTools +using Flux +using ExplainableAI + +on_CI = haskey(ENV, "GITHUB_ACTIONS") + +T = Float32 +input_size = (32, 32, 3, 1) +input = rand(T, input_size) + +model = Chain( + Chain( + Conv((3, 3), 3 => 8, relu; pad=1), + Conv((3, 3), 8 => 8, relu; pad=1), + MaxPool((2, 2)), + Conv((3, 3), 8 => 16, relu; pad=1), + Conv((3, 3), 16 => 16, relu; pad=1), + MaxPool((2, 2)), + ), + Chain( + Flux.flatten, + Dense(1024 => 512, relu), # 102_764_544 parameters + Dropout(0.5), + Dense(512 => 100, relu), + ), +) +Flux.testmode!(model, true) + +# Use one representative algorithm of each type +METHODS = Dict( + "Gradient" => Gradient, + "InputTimesGradient" => InputTimesGradient, + "SmoothGrad" => model -> SmoothGrad(model, 5), + "IntegratedGradients" => model -> IntegratedGradients(model, 5), +) + +# Define benchmark +construct(method, model) = method(model) # for use with @benchmarkable macro + +suite = BenchmarkGroup() +suite["CNN"] = BenchmarkGroup([k for k in keys(METHODS)]) +for (name, method) in METHODS + analyzer = method(model) + suite["CNN"][name] = BenchmarkGroup(["construct analyzer", "analyze"]) + suite["CNN"][name]["constructor"] = @benchmarkable construct($(method), $(model)) + suite["CNN"][name]["analyze"] = @benchmarkable analyze($(input), $(analyzer)) +end diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index a601058..37cdbf3 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -1,95 +1,5 @@ -using BenchmarkTools -using LoopVectorization -using Tullio -using Flux +using PkgJogger using ExplainableAI -using ExplainableAI: lrp!, modify_layer - -on_CI = haskey(ENV, "GITHUB_ACTIONS") - -T = Float32 -input_size = (32, 32, 3, 1) -input = rand(T, input_size) - -model = Chain( - Chain( - Conv((3, 3), 3 => 8, relu; pad=1), - Conv((3, 3), 8 => 8, relu; pad=1), - MaxPool((2, 2)), - Conv((3, 3), 8 => 16, relu; pad=1), - Conv((3, 3), 16 => 16, relu; pad=1), - MaxPool((2, 2)), - ), - Chain( - Flux.flatten, - Dense(1024 => 512, relu), # 102_764_544 parameters - Dropout(0.5), - Dense(512 => 100, relu), - ), -) -Flux.testmode!(model, true) - -# Use one representative algorithm of each type -algs = Dict( - "Gradient" => Gradient, - "InputTimesGradient" => InputTimesGradient, - "LRP" => LRP, - "LREpsilonPlusFlat" => model -> LRP(model, EpsilonPlusFlat()), - "SmoothGrad" => model -> SmoothGrad(model, 5), - "IntegratedGradients" => model -> IntegratedGradients(model, 5), -) - -# Define benchmark -_alg(alg, model) = alg(model) # for use with @benchmarkable macro - -SUITE = BenchmarkGroup() -SUITE["CNN"] = BenchmarkGroup([k for k in keys(algs)]) -for (name, alg) in algs - analyzer = alg(model) - SUITE["CNN"][name] = BenchmarkGroup(["construct analyzer", "analyze"]) - SUITE["CNN"][name]["construct analyzer"] = @benchmarkable _alg($(alg), $(model)) - SUITE["CNN"][name]["analyze"] = @benchmarkable analyze($(input), $(analyzer)) -end - -# generate input for conv layers -insize = (32, 32, 3, 1) -in_dense = 64 -out_dense = 10 -aᵏ = rand(T, insize) - -layers = Dict( - "Conv" => (Conv((3, 3), 3 => 2), aᵏ), - "Dense" => (Dense(in_dense, out_dense, relu), randn(T, in_dense, 1)), -) -rules = Dict( - "ZeroRule" => ZeroRule(), - "EpsilonRule" => EpsilonRule(), - "GammaRule" => GammaRule(), - "WSquareRule" => WSquareRule(), - "FlatRule" => FlatRule(), - "AlphaBetaRule" => AlphaBetaRule(), - "ZPlusRule" => ZPlusRule(), - "ZBoxRule" => ZBoxRule(zero(T), oneunit(T)), -) - -layernames = String.(keys(layers)) -rulenames = String.(keys(rules)) - -SUITE["modify layer"] = BenchmarkGroup(rulenames) -SUITE["apply rule"] = BenchmarkGroup(rulenames) -for rname in rulenames - SUITE["modify layer"][rname] = BenchmarkGroup(layernames) - SUITE["apply rule"][rname] = BenchmarkGroup(layernames) -end - -for (lname, (layer, aᵏ)) in layers - Rᵏ = similar(aᵏ) - Rᵏ⁺¹ = layer(aᵏ) - for (rname, rule) in rules - modified_layer = modify_layer(rule, layer) - SUITE["modify layer"][rname][lname] = @benchmarkable modify_layer($(rule), $(layer)) - SUITE["apply rule"][rname][lname] = @benchmarkable lrp!( - $(Rᵏ), $(rule), $(layer), $(modified_layer), $(aᵏ), $(Rᵏ⁺¹) - ) - end -end +# Use PkgJogger.@jog to create the JogExplainableAI module +@jog ExplainableAI +SUITE = JogExplainableAI.suite() diff --git a/test/Project.toml b/test/Project.toml index 08f0277..a616a45 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,11 +1,13 @@ [deps] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" JuliaFormatter = "98e50ef6-434e-11e9-1051-2b60c6c9e899" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +PkgJogger = "10150987-6cc1-4b76-abee-b1c1cbd91c01" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" ReferenceTests = "324d217c-45ce-50fc-942e-d289b448e8cf" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" diff --git a/test/runtests.jl b/test/runtests.jl index f6b79ee..8bd5819 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -33,4 +33,8 @@ using JET @info "Testing analyzers on batches..." include("test_batches.jl") end + @testset "Benchmark correctness" begin + @info "Testing whether benchmarks are up-to-date..." + include("test_benchmarks.jl") + end end diff --git a/test/test_batches.jl b/test/test_batches.jl index c7ba69a..832d86e 100644 --- a/test/test_batches.jl +++ b/test/test_batches.jl @@ -3,6 +3,7 @@ using Test using Flux using Random +using StableRNGs: StableRNG using Distributions: Laplace pseudorand(dims...) = rand(StableRNG(123), Float32, dims...) diff --git a/test/test_benchmarks.jl b/test/test_benchmarks.jl new file mode 100644 index 0000000..45ee42e --- /dev/null +++ b/test/test_benchmarks.jl @@ -0,0 +1,4 @@ +using PkgJogger +using ExplainableAI + +PkgJogger.@test_benchmarks ExplainableAI