File size: 574 Bytes
26d5b81
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
#ifndef NEUROFLOW_SCHEDULER_HPP
#define NEUROFLOW_SCHEDULER_HPP

#include <cstddef>

namespace neuroflow {

class CosineScheduler {
public:
    enum class Phase { WARMUP, COSINE, DONE };

    float lr_max_;
    float lr_min_;
    size_t warmup_steps_;
    size_t total_steps_;

    CosineScheduler(float lr_max, size_t total_steps,
                    float lr_min_ratio = 0.1f, float warmup_ratio = 0.01f);

    float get_lr(size_t step) const;
    Phase get_phase(size_t step) const;
};

} // namespace neuroflow

#endif // NEUROFLOW_SCHEDULER_HPP