Merge pull request #31 from axelparmentier/clean-parallel-perturbed

Parallelize perturbed with ThreadsX

Author: BatyLeo

Project.toml

@@ -15,6 +15,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+ThreadsX = "ac1d9e8a-700a-412c-b207-f0111f4b6c0d"
 
 [compat]
 ChainRulesCore = "1"
@@ -23,6 +24,7 @@ Krylov = "0.8"
 LinearOperators = "2.3"
 SimpleTraits = "0.9"
 StatsBase = "0.33"
+ThreadsX = "0.1.11"
 julia = "1.7"
 
 [extras]

src/InferOpt.jl

@@ -14,6 +14,7 @@ using SparseArrays
 using Statistics
 using StatsBase: StatsBase, sample
 using Test
+using ThreadsX
 
 include("utils/probability_distribution.jl")
 include("utils/pushforward.jl")

src/fenchel_young/perturbed.jl

@@ -1,11 +1,26 @@
-function fenchel_young_F_and_first_part_of_grad(
-    perturbed::AbstractPerturbed, θ::AbstractArray{<:Real}; kwargs...
+function compute_F_and_y_samples(
+    perturbed::AbstractPerturbed{false}, θ::AbstractArray{<:Real}, Z_samples; kwargs...
 )
-    Z_samples = sample_perturbations(perturbed, θ)
     F_and_y_samples = [
         fenchel_young_F_and_first_part_of_grad(perturbed, θ, Z; kwargs...) for
         Z in Z_samples
     ]
+    return F_and_y_samples
+end
+
+function compute_F_and_y_samples(
+    perturbed::AbstractPerturbed{true}, θ::AbstractArray{<:Real}, Z_samples; kwargs...
+)
+    return ThreadsX.map(
+        Z -> fenchel_young_F_and_first_part_of_grad(perturbed, θ, Z; kwargs...), Z_samples
+    )
+end
+
+function fenchel_young_F_and_first_part_of_grad(
+    perturbed::AbstractPerturbed, θ::AbstractArray{<:Real}; kwargs...
+)
+    Z_samples = sample_perturbations(perturbed, θ)
+    F_and_y_samples = compute_F_and_y_samples(perturbed, θ, Z_samples; kwargs...)
     return mean(first, F_and_y_samples), mean(last, F_and_y_samples)
 end
 

src/perturbed/abstract_perturbed.jl

@@ -1,7 +1,8 @@
 """
-    AbstractPerturbed
+    AbstractPerturbed{B}
 
 Differentiable perturbation of a black box optimizer.
+The parameter `parallel` is a boolean value, equal to true if the perturbations are run in parallel.
 
 # Applicable functions
 
@@ -21,7 +22,7 @@ These subtypes share the following fields:
 - `rng::AbstractRNG`: random number generator
 - `seed::Union{Nothing,Int}`: random seed
 """
-abstract type AbstractPerturbed end
+abstract type AbstractPerturbed{parallel} end
 
 """
     sample_perturbations(perturbed::AbstractPerturbed, θ)
@@ -35,13 +36,31 @@ function sample_perturbations(perturbed::AbstractPerturbed, θ::AbstractArray{<:
     return Z_samples
 end
 
+function compute_atoms(
+    perturbed::AbstractPerturbed{false},
+    θ::AbstractArray{<:Real},
+    Z_samples::Vector{<:AbstractArray{<:Real}};
+    kwargs...,
+)
+    return [perturb_and_optimize(perturbed, θ, Z; kwargs...) for Z in Z_samples]
+end
+
+function compute_atoms(
+    perturbed::AbstractPerturbed{true},
+    θ::AbstractArray{<:Real},
+    Z_samples::Vector{<:AbstractArray{<:Real}};
+    kwargs...,
+)
+    return ThreadsX.map(Z -> perturb_and_optimize(perturbed, θ, Z; kwargs...), Z_samples)
+end
+
 function compute_probability_distribution(
     perturbed::AbstractPerturbed,
     θ::AbstractArray{<:Real},
     Z_samples::Vector{<:AbstractArray{<:Real}};
     kwargs...,
 )
-    atoms = [perturb_and_optimize(perturbed, θ, Z; kwargs...) for Z in Z_samples]
+    atoms = compute_atoms(perturbed, θ, Z_samples; kwargs...)
     weights = ones(length(atoms)) ./ length(atoms)
     probadist = FixedAtomsProbabilityDistribution(atoms, weights)
     return probadist

src/perturbed/additive.jl

@@ -7,12 +7,20 @@ See also: [`AbstractPerturbed`](@ref).
 
 Reference: <https://arxiv.org/abs/2002.08676>
 """
-struct PerturbedAdditive{F,R<:AbstractRNG,S<:Union{Nothing,Int}} <: AbstractPerturbed
+struct PerturbedAdditive{F,R<:AbstractRNG,S<:Union{Nothing,Int},parallel} <:
+       AbstractPerturbed{parallel}
     maximizer::F
     ε::Float64
     nb_samples::Int
     rng::R
     seed::S
+
+    function PerturbedAdditive{F,R,S,parallel}(
+        maximizer::F, ε::Float64, nb_samples::Int, rng::R, seed::S
+    ) where {F,R<:AbstractRNG,S<:Union{Nothing,Int},parallel}
+        @assert parallel isa Bool
+        return new{F,R,S,parallel}(maximizer, ε, nb_samples, rng, seed)
+    end
 end
 
 function Base.show(io::IO, perturbed::PerturbedAdditive)
@@ -28,12 +36,22 @@ end
 Shorter constructor with defaults.
 """
 function PerturbedAdditive(
-    maximizer; ε=1.0, epsilon=nothing, nb_samples=1, rng=MersenneTwister(0), seed=nothing
-)
+    maximizer::F;
+    ε=1.0,
+    epsilon=nothing,
+    nb_samples=1,
+    rng::R=MersenneTwister(0),
+    seed::S=nothing,
+    is_parallel=false,
+) where {F,R,S}
     if isnothing(epsilon)
-        return PerturbedAdditive(maximizer, float(ε), nb_samples, rng, seed)
+        return PerturbedAdditive{F,R,S,is_parallel}(
+            maximizer, float(ε), nb_samples, rng, seed
+        )
     else
-        return PerturbedAdditive(maximizer, float(epsilon), nb_samples, rng, seed)
+        return PerturbedAdditive{F,R,S,is_parallel}(
+            maximizer, float(epsilon), nb_samples, rng, seed
+        )
     end
 end
 

src/perturbed/multiplicative.jl

@@ -7,12 +7,20 @@ See also: [`AbstractPerturbed`](@ref).
 
 Reference: preprint coming soon.
 """
-struct PerturbedMultiplicative{F,R<:AbstractRNG,S<:Union{Nothing,Int}} <: AbstractPerturbed
+struct PerturbedMultiplicative{F,R<:AbstractRNG,S<:Union{Nothing,Int},parallel} <:
+       AbstractPerturbed{parallel}
     maximizer::F
     ε::Float64
     nb_samples::Int
     rng::R
     seed::S
+
+    function PerturbedMultiplicative{F,R,S,parallel}(
+        maximizer::F, ε::Float64, nb_samples::Int, rng::R, seed::S
+    ) where {F,R<:AbstractRNG,S<:Union{Nothing,Int},parallel}
+        @assert parallel isa Bool
+        return new{F,R,S,parallel}(maximizer, ε, nb_samples, rng, seed)
+    end
 end
 
 function Base.show(io::IO, perturbed::PerturbedMultiplicative)
@@ -28,12 +36,22 @@ end
 Shorter constructor with defaults.
 """
 function PerturbedMultiplicative(
-    maximizer; ε=1.0, epsilon=nothing, nb_samples=1, rng=MersenneTwister(0), seed=nothing
-)
+    maximizer::F;
+    ε=1.0,
+    epsilon=nothing,
+    nb_samples=1,
+    rng::R=MersenneTwister(0),
+    seed::S=nothing,
+    is_parallel=false,
+) where {F,R,S}
     if isnothing(epsilon)
-        return PerturbedMultiplicative(maximizer, float(ε), nb_samples, rng, seed)
+        return PerturbedMultiplicative{F,R,S,is_parallel}(
+            maximizer, float(ε), nb_samples, rng, seed
+        )
     else
-        return PerturbedMultiplicative(maximizer, float(epsilon), nb_samples, rng, seed)
+        return PerturbedMultiplicative{F,R,S,is_parallel}(
+            maximizer, float(epsilon), nb_samples, rng, seed
+        )
     end
 end
 

test/paths.jl

@@ -48,6 +48,20 @@ pipelines_imitation_y = [
         maximizer=identity,
         loss=FenchelYoungLoss(PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=5)),
     ),
+    (
+        encoder=encoder_factory(),
+        maximizer=identity,
+        loss=FenchelYoungLoss(
+            PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=5, is_parallel=true)
+        ),
+    ),
+    (
+        encoder=encoder_factory(),
+        maximizer=identity,
+        loss=FenchelYoungLoss(
+            PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=5, is_parallel=true)
+        ),
+    ),
     # Perturbed + other loss
     (
         encoder=encoder_factory(),
@@ -59,6 +73,18 @@ pipelines_imitation_y = [
         maximizer=PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=10),
         loss=Flux.Losses.mse,
     ),
+    (
+        encoder=encoder_factory(),
+        maximizer=PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=10, is_parallel=true),
+        loss=Flux.Losses.mse,
+    ),
+    (
+        encoder=encoder_factory(),
+        maximizer=PerturbedMultiplicative(
+            true_maximizer; ε=1.0, nb_samples=10, is_parallel=true
+        ),
+        loss=Flux.Losses.mse,
+    ),
     # Generic regularized + FYL
     (
         encoder=encoder_factory(),
@@ -88,6 +114,21 @@ pipelines_experience = [
             PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=10), cost
         ),
     ),
+    (
+        encoder=encoder_factory(),
+        maximizer=identity,
+        loss=Pushforward(
+            PerturbedAdditive(true_maximizer; ε=1.0, nb_samples=10, is_parallel=true), cost
+        ),
+    ),
+    (
+        encoder=encoder_factory(),
+        maximizer=identity,
+        loss=Pushforward(
+            PerturbedMultiplicative(true_maximizer; ε=1.0, nb_samples=10, is_parallel=true),
+            cost,
+        ),
+    ),
     (
         encoder=encoder_factory(),
         maximizer=identity,