Tree Parzen within Hyperband

[AlexanderNenninger]

Is your feature request related to a problem? Please describe.

Yes, but not of this package. Hyperopt.jl's API is a bit of a mess, unfortunately. That's why I've implemented BOHB with this package's TreeParzen optimizer as a model provider.

The Trial API in combination with the ask-and-tell API of this package made the implementation particularly elegant.

Describe the solution you'd like Integrate code of BOHB into this package.

Describe alternatives you've considered

I considered keeping the code in a private Julia repo, since I don't want to maintain a public package myself. But this makes the code less useful for the wider Julia community.

Additional context

Prototype code:

using TreeParzen

"""
    hyperband(f_budget, bmin, bmax, space, η=1.3, config=Config())

Treeparzen estimator within hyperband. We ask treeparzen for recommendations, then we evaluate
the most promising recommendations using successive_halfing.
"""
function hyperband(f_budget, bmin, bmax, space, η=1.3, config=Config())
    # Ensure space is ok
    TreeParzen.Graph.checkspace(space)
    @assert bmin < bmax
    @assert check_signature(f_budget) "Expected signature `f_budget(;hyperparams..., budget)`"
    @assert η > 1

    smax = floor(Int, log(η, bmax / bmin))
    trial_history = TreeParzen.Trials.Trial[]
    for s in smax:-1:0
        n = ceil((smax + 1) / (s + 1) / η^s)
        trials = [TreeParzen.ask(space, trial_history, config) for _ in 1:n]
        successive_halfing!(trials, f_budget, bmax / η^s, bmax, η)
        append!(trial_history, trials)
    end
    trial_history
end

"""
    successive_halfing!(trials::Vector{TreeParzen.Trials.Trial}, f_budget::Function, bmin, bmax, η)

Takes a vector of trials and evaluates f_budget with different budgets on each element,
increasing evaluation time for promising candidates, while discarding the rest. Writes the budget
and loss into trials. 
"""
function successive_halfing!(trials::Vector{TreeParzen.Trials.Trial}, f_budget::Function, bmin, bmax, η)
    budget = bmin
    trials_view = view(trials, :)
    while budget <= bmax && length(trials) >= 1
        (trial -> trial.hyperparams[:budget] = budget).(trials_view)

        # We can assume this line is the most expensive one. But with the current structure, 
        # it is trivially parallelizable using Distributed or Dagger.
        (trial -> tell!(trial, f_budget(; trial.hyperparams...))).(trials_view)
        sort!(trials_view, by=trial -> trial.loss)
        trials_view = @view trials_view[1:floor(Int, length(trials) / η)]
        budget = η * budget
    end
end

"""
    check_signature(f::Function)

Determin whether the f can take the keyword argument budget. 
"""
function check_signature(f::Function)
    methods = Base.methods(f)
    any(any(Base.kwarg_decl(method) .== :budget) for method in methods)
end

# Usage

space = Dict(:x => HP.Uniform(:x, -5.0, 5.0))

# This function becomes less noisy with increasing budget.
f(; x, budget) = (x - 1)^2 + randn()^2 / budget
trials = hyperband(f, 1, 10, space, 2)
t = trials[argmin(t.loss for t in trials)]