weight_optimizer.h

Go to the documentation of this file.

/* Brush
copyright 2020 William La Cava
license: GNU/GPL v3
 
Code below heavily inspired by heal-research/operon, Copyright 2019-2022 Heal Research
*/
#ifndef WEIGHT_OPTIMIZER_H
#define WEIGHT_OPTIMIZER_H
 
#include <ceres/dynamic_autodiff_cost_function.h>
#include <ceres/dynamic_numeric_diff_cost_function.h>
#include <ceres/loss_function.h>
#include <ceres/solver.h>
#include <ceres/tiny_solver.h>
 
#include "tiny_cost_function.h"
 
namespace Brush
{
 
struct OptimizerSummary {
    double InitialCost;
    double FinalCost;
    int Iterations;
    int FunctionEvaluations;
    int JacobianEvaluations;
    bool Success;
};
 
template<typename PT>
struct ResidualEvaluator {
    typedef float Scalar; 
    ResidualEvaluator(PT& program, Dataset const& dataset,
                      const vector<float>& class_weights = {})
        : program_(program)
        , dataset_(dataset)
        , numParameters_(program.get_weights().size())
        , y_true_(dataset.y)
        , class_weights_(class_weights)
    {}
 
    template<typename T>
    auto operator()(Eigen::DenseBase<T>& parameters, Eigen::DenseBase<T>& residuals) const noexcept -> void
    {
        return (*this)(parameters.data(), residuals.data());
    }
 
    template <typename T>
    auto operator()(T const* parameters, T* residuals) const -> bool
    {
        const T ** new_weights = &parameters; 
        auto residualMap = Eigen::Map<Array<T, Dynamic, 1>>(
            residuals, GetDataset().get_n_samples());
 
        if constexpr (PT::program_type == ProgramType::MulticlassClassifier)
        {
            using MatrixType = Array<T, Dynamic, Dynamic>;
            auto probabilities = GetProgram().template predict_with_weights<MatrixType>(
                GetDataset(), new_weights);
            for (int i = 0; i < GetDataset().get_n_samples(); ++i)
            {
                const int label = static_cast<int>(GetTarget()(i));
                const T probability = probabilities(i, label);
                const float weight = class_weights_.empty() ? 1.0f
                    : class_weights_.at(label);
                // TinySolver minimizes squared residuals, so sqrt(loss)
                // makes its objective the weighted multinomial log loss.
                residualMap(i) = sqrt(T(weight) * -log(probability));
            }
        }
        else
        {
            using ArrayType = Array<T, Dynamic, 1>;
            ArrayType y_pred = GetProgram().template predict_with_weights<ArrayType>(
                GetDataset(), new_weights);
 
            if (GetDataset().classification)
            {
            // tolerance to avoid numeric errors.
 
            // Using an eps with 7 significant digits to avoid weird behavior.
            // From some experiments, I noticed that using very small values (previous
            // was pow(10,-10) caused the optimizer to evaluate to nans), probably because
            // we were calculating log(0) due to cpp ignoring more than 7 significant digits
            // for floats. using -6. As a reference, tensorflow uses -7.
            // We need to remember that we dont want to have eps too big, otherwise
            // we are biasing the predictions away from 0 and 1
            float eps = 1e-6f;
 
            auto y = GetTarget();
 
            // cout << T(1.0) << ", " << T(eps) << endl;
 
            // clamp values and avoid log(0)
            y_pred = y_pred.min(T(1.0) - T(eps)).max(T(eps));
            
                for (int i = 0; i < y_pred.size(); ++i)
                {
                    const float weight = class_weights_.empty() ? 1.0f
                        : class_weights_.at(static_cast<int>(y(i)));
                    const T log_loss = -(T(y(i)) * log(y_pred(i))
                        + (T(1.0f) - T(y(i))) * log(T(1.0f) - y_pred(i)));
                    // See multiclass branch above: this makes the least-squares
                    // objective equal weighted binary log loss.
                    residualMap(i) = sqrt(T(weight) * log_loss);
                }
            }
            else
            {
                residualMap = y_pred - GetTarget();
            }
        }
 
        return true;
    }
 
    [[nodiscard]] auto NumParameters() const -> size_t { return numParameters_; }
    [[nodiscard]] auto NumResiduals() const -> size_t { return y_true_.get().size(); }
    inline auto GetProgram() const { return program_.get();};
    inline auto GetDataset() const { return dataset_.get();};
    inline auto GetTarget() const { return y_true_.get();};
 
private:
    std::reference_wrapper<PT> program_;
    std::reference_wrapper<Dataset const> dataset_;
    std::reference_wrapper<ArrayXf const> y_true_;
    vector<float> class_weights_;
    size_t numParameters_; // cache the number of parameters in the tree
};
 
// TODO: see this struct and try to understand how to make non-templated classes
struct WeightOptimizer
{
    template<typename PT>
    void update(PT& program, const Dataset& dataset,
                const vector<float>& class_weights = {})
    {
        if (program.get_n_weights() == 0)
            return;
            
        // fmt::print("number of weights: {}\n",program.get_n_weights());
        auto init_weights = program.get_weights();
 
        using CFType = Brush::TinyCostFunction<ResidualEvaluator<PT>> ; 
        ResidualEvaluator<PT> evaluator(program, dataset, class_weights);
        CFType cost_function(evaluator);
        ceres::TinySolver<CFType> solver;
        solver.options.max_num_iterations = 10;
 
        typename decltype(solver)::Parameters parameters = program.get_weights(); 
        solver.Solve(cost_function, &parameters);
    
        // fmt::print("Summary:\nInitial cost: {}\nFinal Cost: {}\nIterations: {}\n",
        //     solver.summary.initial_cost,
        //     solver.summary.final_cost,
        //     solver.summary.iterations
        // );
        // fmt::print("Initial weights: {}\nFinal weights: {}\n", 
        //     init_weights, 
        //     parameters
        // );
        
        // Validate optimizer results before applying
        bool costs_valid = !std::isnan(solver.summary.final_cost) && 
                          !std::isinf(solver.summary.final_cost) &&
                          !std::isnan(solver.summary.initial_cost) && 
                          !std::isinf(solver.summary.initial_cost);
        
        bool weights_valid = parameters.allFinite();
        
        if (costs_valid && weights_valid && 
            solver.summary.final_cost < solver.summary.initial_cost)
        {
            program.set_weights(parameters);
        }
 
    }
};
 
} // namespace Brush
#endif