successive_halving.hpp

Go to the documentation of this file.

 
#pragma once
 
#include <algorithm>
#include <functional>
#include <limits>
#include <memory>
#include <stdexcept>
#include <vector>
 
#include "pulsatrix/hpo_sampling.hpp"
#include "pulsatrix/search_space.hpp"
 
namespace pulsatrix {
 
class ResumableTrial {
public:
    virtual ~ResumableTrial() = default;
 
    virtual double TrainForEpochs(int num_epochs) = 0;
};
 
using TrialFactory = std::function<std::unique_ptr<ResumableTrial>(const Configuration&)>;
 
struct SuccessiveHalvingResult {
    Configuration best_configuration;
    double best_metric;
    size_t total_epochs_trained;
};
 
inline SuccessiveHalvingResult RunSuccessiveHalvingOnConfigs(std::vector<Configuration> configs,
                                                              const TrialFactory& make_trial,
                                                              int initial_epoch_budget, double eta) {
    if (configs.empty()) {
        throw std::invalid_argument("RunSuccessiveHalvingOnConfigs: configs must not be empty");
    }
    if (initial_epoch_budget <= 0) {
        throw std::invalid_argument("RunSuccessiveHalvingOnConfigs: initial_epoch_budget must be positive");
    }
    if (eta <= 1.0) {
        throw std::invalid_argument("RunSuccessiveHalvingOnConfigs: eta must be > 1.0");
    }
 
    struct Candidate {
        Configuration config;
        std::unique_ptr<ResumableTrial> trial;
        double metric;
    };
    std::vector<Candidate> candidates;
    candidates.reserve(configs.size());
    for (auto& config : configs) {
        auto trial = make_trial(config);
        candidates.push_back(Candidate{std::move(config), std::move(trial),
                                        -std::numeric_limits<double>::infinity()});
    }
 
    size_t total_epochs_trained = 0;
    int budget = initial_epoch_budget;
    bool trained_at_least_once = false;
 
    while (true) {
        for (auto& c : candidates) {
            c.metric = c.trial->TrainForEpochs(budget);
        }
        total_epochs_trained += candidates.size() * static_cast<size_t>(budget);
        trained_at_least_once = true;
 
        if (candidates.size() <= 1) {
            break;
        }
        std::sort(candidates.begin(), candidates.end(),
                  [](const Candidate& a, const Candidate& b) { return a.metric > b.metric; });
 
        size_t num_survivors = std::max<size_t>(1, static_cast<size_t>(candidates.size() / eta));
        if (num_survivors >= candidates.size()) {
            break;  // eta too small to actually reduce the population further -- stop here
        }
        candidates.resize(num_survivors);
        budget = static_cast<int>(static_cast<double>(budget) * eta);
    }
    (void)trained_at_least_once;  // always true by this point -- documents the loop's own invariant
 
    auto best_it = std::max_element(candidates.begin(), candidates.end(),
                                     [](const Candidate& a, const Candidate& b) { return a.metric < b.metric; });
    return SuccessiveHalvingResult{best_it->config, best_it->metric, total_epochs_trained};
}
 
template <typename RNG>
SuccessiveHalvingResult RunSuccessiveHalving(const SearchSpace& space, const TrialFactory& make_trial,
                                              size_t num_configs, int initial_epoch_budget, double eta,
                                              RNG& rng) {
    if (num_configs == 0) {
        throw std::invalid_argument("RunSuccessiveHalving: num_configs must be positive");
    }
    std::vector<Configuration> configs;
    configs.reserve(num_configs);
    for (size_t i = 0; i < num_configs; ++i) {
        configs.push_back(RandomSample(space, rng));
    }
    return RunSuccessiveHalvingOnConfigs(std::move(configs), make_trial, initial_epoch_budget, eta);
}
 
}  // namespace pulsatrix