Fix benchmarks

Author: adrhill

benchmark/benchmarks.jl

@@ -6,7 +6,7 @@ on_CI = haskey(ENV, "GITHUB_ACTIONS")
 
 include("../test/vgg11.jl")
 vgg11 = VGG11(; pretrain=false)
-model = flatten_model(strip_softmax(vgg19.layers))
+model = flatten_model(strip_softmax(vgg11.layers))
 img = rand(MersenneTwister(123), Float32, (224, 224, 3, 1))
 
 # Benchmark custom LRP composite
@@ -23,11 +23,13 @@ algs = Dict(
 )
 
 # Define benchmark
+contruct_analyzer(alg, model) = alg(model) # for use with @benchmarkable macro
+
 SUITE = BenchmarkGroup()
 SUITE["VGG"] = BenchmarkGroup([k for k in keys(algs)])
 for (name, alg) in algs
     SUITE["VGG"][name] = BenchmarkGroup(["construct analyzer", "analyze"])
-    SUITE["VGG"][name]["construct analyzer"] = @benchmarkable alg($(model))
+    SUITE["VGG"][name]["construct analyzer"] = @benchmarkable contruct_analyzer($(alg), $(model))
 
     analyzer = alg(model)
     SUITE["VGG"][name]["analyze"] = @benchmarkable analyze($(img), $(analyzer))
@@ -58,15 +60,15 @@ rules = Dict(
 )
 rulenames = [k for k in keys(rules)]
 
+test_rule(rule, layer, aₖ, Rₖ₊₁) = rule(layer, aₖ, Rₖ₊₁) # for use with @benchmarkable macro
+
 for (layername, (layer, aₖ)) in layers
     SUITE[layername] = BenchmarkGroup(rulenames)
+    Rₖ₊₁ = layer(aₖ)
 
-    for (rulename, ruletype) in rules
-        Rₖ₊₁ = layer(aₖ)
+    for (rulename, rule) in rules
         SUITE[layername][rulename] = BenchmarkGroup(["dispatch", "AD fallback"])
-        SUITE[layername][rulename]["dispatch"] = @benchmarkable rule($layer, $aₖ, $Rₖ₊₁)
-        SUITE[layername][rulename]["AD fallback"] = @benchmarkable rule(
-            $TestWrapper(layer), $aₖ, $Rₖ₊₁
-        )
+        SUITE[layername][rulename]["dispatch"] = @benchmarkable test_rule($(rule), $(layer), $(aₖ), $(Rₖ₊₁))
+        SUITE[layername][rulename]["AD fallback"] = @benchmarkable test_rule($(rule), $(TestWrapper(layer)), $(aₖ), $(Rₖ₊₁))
     end
 end