File size: 1,227 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 27 28 29 30 31 32 33 | #include "neuroflow/scheduler.hpp"
#include <algorithm>
#include <cmath>
namespace neuroflow {
CosineScheduler::CosineScheduler(float lr_max, size_t total_steps,
float lr_min_ratio, float warmup_ratio)
: lr_max_(lr_max),
lr_min_(lr_max * lr_min_ratio),
warmup_steps_(std::max(size_t(1), static_cast<size_t>(static_cast<float>(total_steps) * warmup_ratio))),
total_steps_(total_steps) {}
float CosineScheduler::get_lr(size_t step) const {
if (total_steps_ == 0) return lr_max_;
if (step >= total_steps_) return lr_min_;
if (step < warmup_steps_) {
return lr_max_ * static_cast<float>(step + 1) / static_cast<float>(warmup_steps_);
}
float progress = static_cast<float>(step - warmup_steps_)
/ static_cast<float>(total_steps_ - warmup_steps_);
return lr_min_ + 0.5f * (lr_max_ - lr_min_)
* (1.0f + std::cos(3.14159265358979f * progress));
}
CosineScheduler::Phase CosineScheduler::get_phase(size_t step) const {
if (step < warmup_steps_) return Phase::WARMUP;
if (step < total_steps_) return Phase::COSINE;
return Phase::DONE;
}
} // namespace neuroflow |