optimizer.h

Go to the documentation of this file.

// SPDX-License-Identifier: AGPL-3.0-only
// Copyright (C) 2026 Artificial Intelligence Techniques, S.L. and Roberto López González
 
#pragma once
 
#include <array>
#include <functional>
#include "opennn/dataset/batch.h"
#include "opennn/core/device_backend.h"
#include "opennn/network/forward_propagation.h"
#include "opennn/core/json.h"
#include "opennn/training/loss.h"
#include "opennn/training/training_context.h"
#include "opennn/core/tensor_types.h"
#include "opennn/core/thread_safe_queue.h"
#include "opennn/training/training_result.h"
 
namespace opennn
{
 
inline constexpr float GRADIENT_NORM_EPS = 1e-6f;
 
class Network;
struct Buffer;
struct BackPropagation;
 
struct OptimizerData
{
    void set(const vector<Shape>&, Device device = Device::CPU);
    void set(const vector<Shape>&, const vector<Type>& slot_types, Device device = Device::CPU);
 
    Buffer data;
    vector<TensorView> views;
};
 
class Optimizer
{
 
public:
 
    explicit Optimizer(Loss* = nullptr);
    virtual ~Optimizer();
 
    using StoppingCondition = opennn::StoppingCondition;
 
    const Loss* get_loss() const noexcept { return loss; }
 
    bool get_display() const noexcept { return display; }
 
    void set(Loss* new_loss) { loss = new_loss; }
 
    virtual void set_display(bool new_display) { display = new_display; }
 
    void set_display_period(const Index new_display_period)
    {
        throw_if(new_display_period <= 0, "Optimizer::set_display_period: period must be positive.");
        display_period = new_display_period;
    }
 
    void set_workers_number(int new_workers_number) { workers_number = max(1, new_workers_number); }
 
    void set_cuda_graph(bool enabled) { use_cuda_graph = enabled; }
 
    bool get_cuda_graph_capture_failed() const noexcept { return cuda_graph_capture_failed; }
 
    void set_shuffle(bool enabled) { shuffle_samples = enabled; }
 
    void set_batch_pool_size(int size) { batch_pool_size_override = size; }
 
    void set_maximum_epochs(const Index new_maximum_epochs) { maximum_epochs = new_maximum_epochs; }
    Index get_maximum_epochs() const noexcept { return maximum_epochs; }
    void set_maximum_time(const float new_maximum_time) { maximum_time = new_maximum_time; }
 
    void set_loss_goal(const float new_loss_goal) { training_loss_goal = new_loss_goal; }
    void set_maximum_validation_failures(const Index new_maximum_validation_failures) { maximum_validation_failures = new_maximum_validation_failures; }
    void set_validation_period(const Index n)
    {
        throw_if(n <= 0, "Optimizer::set_validation_period: period must be positive.");
        validation_period = n;
    }
 
    void set_restore_best(bool enabled) { restore_best = enabled; }
 
    void set_gradient_clip_norm(const float new_clip) { gradient_clip_norm = new_clip; }
 
    void set_joint_gradient_arena(bool enabled) noexcept
    {
        joint_gradient_arena = enabled;
    }
 
    Index get_batch_size() const noexcept { return batch_size; }
    void set_batch_size(const Index new_batch_size) { batch_size = new_batch_size; }
 
    virtual void configure_for_task(NetworkTask);
 
    virtual TrainingResult train();
 
    Index get_maximum_batch_size() const;
 
    const string& get_name() const noexcept { return name; }
 
    virtual void from_JSON(const JsonDocument&);
 
    virtual void to_JSON(JsonWriter&) const;
 
    void save(const filesystem::path&) const;
    void load(const filesystem::path&);
 
    function<void(Index, float, float, Network*)> post_epoch_callback;
    function<void(Network*)> post_batch_callback;
    function<void(Index, float)> post_best_callback;
 
protected:
 
    enum class UpdateMode { Standard, Capturable };
 
    struct TrainingSession
    {
        static constexpr int group_size = 8;
        static constexpr int slots_count = 2 * group_size;
        static constexpr int pipelines_count = 2;
 
        struct GraphPipeline
        {
            array<unique_ptr<Batch>, group_size> slots;
            device::GraphExecHandle exec;
            device::CudaEvent fork_event;
            array<device::CudaEvent, group_size> copy_done_events;
        };
 
        Batch* fixed_batch() const { return pipelines[0].slots[0].get(); }
 
        bool has_graph_batches() const
        {
            return pipelines[0].slots[1] || pipelines[1].slots[0];
        }
 
        void disable_cuda_graph_capture()
        {
            for (GraphPipeline& pipeline : pipelines)
                pipeline.exec.reset();
            cuda_graph_capture_allowed = false;
        }
 
        struct TailContext
        {
            unique_ptr<Batch> batch;
            unique_ptr<TrainingContext> context;
            device::GraphExecHandle exec;
            Index size = 0;
            bool capture_failed = false;
        };
 
        array<GraphPipeline, pipelines_count> pipelines;
        TailContext tail;
        Buffer device_metrics{Device::CUDA};
        array<device::CudaEvent, 4> throttle_events;
        size_t throttle_cursor = 0;
        bool cuda_graph_capture_allowed = false;
    };
 
    void set_names();
    void prepare_training_artifacts();
 
    // Training always runs at least one epoch, even when the limit is zero.
    Index get_epoch_budget() const noexcept { return max<Index>(maximum_epochs, 1); }
 
    struct EpochReport
    {
        Index completed_epochs = 0;
        float seconds = 0.0f;
        float training_error = 0.0f;
        float loss = 0.0f;
        optional<float> loss_decrease;
        Index validation_failures = 0;
    };
 
    optional<StoppingCondition> find_stopping_condition(const EpochReport&,
                                                        float minimum_loss_decrease = -MAX) const;
 
    void close_training_result(TrainingResult&, const EpochReport&, bool has_validation) const;
 
    struct BestModelSnapshot
    {
        float validation_error = MAX;
        Index epoch = -1;
        vector<float> parameters;
        vector<float> states;
    };
 
    void update_best_parameters(Network*,
                                float,
                                Index,
                                Index&,
                                BestModelSnapshot&);
 
    void restore_best_parameters(Network*,
                                 TrainingResult&,
                                 const BestModelSnapshot&);
 
    void write_common_json(JsonWriter&) const;
    void read_common_json(const Json*);
 
    void setup_device_training();
    void teardown_device_training();
 
    void warmup_device_training(TrainingContext&,
                                ThreadSafeQueue<Batch*>&,
                                const vector<vector<Index>>&,
                                const FeatureSelection&,
                                TrainingSession&,
                                OptimizerData&,
                                ForwardPropagation* validation_forward_propagation = nullptr,
                                ThreadSafeQueue<Batch*>* validation_empty_queue = nullptr,
                                const vector<vector<Index>>* validation_batches = nullptr);
 
    void prefetch_batch(Batch&);
 
    void sync_device(bool on_gpu, bool has_recurrent_layers, TrainingSession&);
 
    static void clip_gradient_norm(BackPropagation&, float);
 
    bool should_display(Index epoch) const noexcept { return display && epoch % display_period == 0; }
 
    void display_epoch_results(Index epoch, float training_error, float training_accuracy,
                               float validation_error, float validation_accuracy,
                               bool has_validation, bool validation_fresh, bool is_token_cross_entropy,
                               float elapsed_time) const;
 
    void setup_batch_pools(BatchPools& pools,
                           Dataset& dataset,
                           Network& network,
                           Index training_batch_size,
                           Index validation_batch_size,
                           bool has_validation,
                           TrainingSession& training_session);
 
    struct WorkerProfileCounters;
    struct GraphEpochContext;
    struct TrainingEpochContext;
 
    unique_ptr<BatchPrefetchSession> start_batch_prefetch(
        ThreadSafeQueue<Batch*>&,
        const vector<vector<Index>>&,
        const FeatureSelection&,
        FillMode,
        WorkerProfileCounters* profile_counters = nullptr);
 
    int get_batch_pool_size(const Network&) const;
 
    struct EpochLoopContext;
    Loss::EvaluationResult run_epoch_loop(EpochLoopContext&);
 
    struct FullBatchContext
    {
        Network* network = nullptr;
        Index training_samples_number = 0;
        Index validation_samples_number = 0;
        unique_ptr<Batch> training_batch;
        unique_ptr<Batch> validation_batch;
        unique_ptr<ForwardPropagation> training_forward_propagation;
        unique_ptr<ForwardPropagation> validation_forward_propagation;
    };
 
    struct FullBatchStep
    {
        float training_error = 0.0f;
        float displayed_error = 0.0f;
        float loss = 0.0f;
    };
 
    struct FullBatchHooks
    {
        function<void()> setup_state;
        function<FullBatchStep()> train_step;
        function<float()> validation_error;
        function<void()> display_extra;
        function<void()> post_step;
        float minimum_loss_decrease = -MAX;
    };
 
    void prepare_full_batch_training(FullBatchContext&, const char* banner);
    TrainingResult train_full_batch(FullBatchContext&, const FullBatchHooks&);
 
    virtual string get_display_name() const { return name; }
    virtual void setup_optimizer_data(OptimizerData&, Index, Device) {}
    virtual void update_parameters(BackPropagation&, OptimizerData&,
                                   UpdateMode = UpdateMode::Standard)
    { 
        throw runtime_error("train() requires a mini-batch optimizer (SGD or Adam)."); 
    }
 
    virtual bool supports_cuda_graph() const noexcept { return false; }
 
    bool network_has_active_dropout() const;
 
    bool can_use_cuda_graph() const
    {
        return use_cuda_graph
            && supports_cuda_graph()
            && device::is_cuda_build()
            && loss
            && loss->supports_device_epoch_metrics()
            && !network_has_active_dropout();
    }
    virtual void on_epoch_begin(Index, OptimizerData&) {}
 
    bool uses_joint_gradient_arena() const noexcept;
 
    Loss::EvaluationResult run_graph_epoch(TrainingSession&,
                                           OptimizerData&,
                                           ForwardPropagation&,
                                           BackPropagation&,
                                           ThreadSafeQueue<Batch*>&,
                                           const vector<vector<Index>>&,
                                           const FeatureSelection&);
 
    Loss::EvaluationResult train_epoch(TrainingContext&,
                                       ThreadSafeQueue<Batch*>&,
                                       const vector<vector<Index>>&,
                                       const FeatureSelection&,
                                       TrainingSession&,
                                       OptimizerData&);
 
    Loss::EvaluationResult evaluate_epoch(ForwardPropagation&,
                                          ThreadSafeQueue<Batch*>&,
                                          const vector<vector<Index>>&,
                                          const FeatureSelection&,
                                          TrainingSession&);
 
    Loss* loss = nullptr;
 
    float training_loss_goal = 0.0f;
 
    Index maximum_validation_failures = numeric_limits<Index>::max();
    Index validation_period = 1;
 
    float gradient_clip_norm = 0.0f;
 
    bool joint_gradient_arena = false;
 
    bool restore_best = true;
 
    Index maximum_epochs = 10000;
 
    Index batch_size = 0;
 
    float maximum_time = 360000.0f;
 
    Index display_period = 10;
 
    bool display = true;
 
    bool shuffle_samples = true;
 
    int batch_pool_size_override = 0;
 
    // How many threads prepare batches ahead of the device. Two was a bare
    // literal that nothing outside the tests ever changed, so it is measured
    // here rather than left to look arbitrary: on the CNN family, whose fill
    // decodes and resizes an image per sample and is the most expensive one in
    // the suite, six workers were slower than two in every run and in both
    // orders -- medians 378 against 406 samples/s, 7% down. The fills are not
    // waiting on anything the extra threads could overlap with, so the workers
    // only take the cores the main thread wants.
    //
    // OPENNN_BATCH_WORKERS overrides it, which is how the above was run. The
    // batch pool sizes itself from this, so raising it widens the queue too.
    static const int default_workers_number;
 
    int workers_number = default_workers_number;
 
    bool use_cuda_graph = false;
    bool cuda_graph_capture_failed = false;
 
    string name;
};
 
unique_ptr<Optimizer> create_optimizer(const string& name);
 
}