Add ProgressMeter to NoiseAugmentation

Author: adrhill

Project.toml

@@ -7,6 +7,7 @@ version = "0.10.1"
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
@@ -16,6 +17,7 @@ XAIBase = "9b48221d-a747-4c1b-9860-46a1d8ba24a7"
 ADTypes = "1"
 DifferentiationInterface = "0.6"
 Distributions = "0.25"
+ProgressMeter = "1.10.4"
 Random = "<0.0.1, 1"
 Reexport = "1"
 Statistics = "<0.0.1, 1"

src/input_augmentation.jl

@@ -29,17 +29,20 @@ struct NoiseAugmentation{A<:AbstractXAIMethod,D<:Sampleable,R<:AbstractRNG} <:
     n::Int
     distribution::D
     rng::R
+    show_progress::Bool
 
     function NoiseAugmentation(
-        analyzer::A, n::Int, distribution::D, rng::R=GLOBAL_RNG
+        analyzer::A, n::Int, distribution::D, rng::R=GLOBAL_RNG, show_progress=true
     ) where {A<:AbstractXAIMethod,D<:Sampleable,R<:AbstractRNG}
         n < 1 && throw(ArgumentError("Number of samples `n` needs to be larger than zero."))
-        return new{A,D,R}(analyzer, n, distribution, rng)
+        return new{A,D,R}(analyzer, n, distribution, rng, show_progress)
     end
 end
-function NoiseAugmentation(analyzer, n::Int, std::T=1.0f0, rng=GLOBAL_RNG) where {T<:Real}
+function NoiseAugmentation(
+    analyzer, n::Int, std::T=1.0f0, rng=GLOBAL_RNG, show_progress=true
+) where {T<:Real}
     distribution = Normal(zero(T), std^2)
-    return NoiseAugmentation(analyzer, n, distribution, rng)
+    return NoiseAugmentation(analyzer, n, distribution, rng, show_progress)
 end
 
 function call_analyzer(input, aug::NoiseAugmentation, ns::AbstractOutputSelector; kwargs...)
@@ -48,17 +51,20 @@ function call_analyzer(input, aug::NoiseAugmentation, ns::AbstractOutputSelector
     output_indices = ns(output)
     output_selector = AugmentationSelector(output_indices)
 
+    p = Progress(aug.n; desc="Sampling NoiseAugmentation...", showspeed=aug.show_progress)
     # First augmentation
     input_aug = similar(input)
     input_aug = sample_noise!(input_aug, input, aug)
     expl_aug = aug.analyzer(input_aug, output_selector)
     sum_val = expl_aug.val
+    next!(p)
 
     # Further augmentations
     for _ in 2:(aug.n)
         input_aug = sample_noise!(input_aug, input, aug)
         expl_aug = aug.analyzer(input_aug, output_selector)
         sum_val += expl_aug.val
+        next!(p)
     end
 
     # Average explanation