Skip to content

Commit 114dab9

Browse files
CydralCydral
committed
Update
1 parent 9396527 commit 114dab9

4 files changed

Lines changed: 1747 additions & 0 deletions

File tree

‎dlib/dnn/lr_scheduler.h‎

Lines changed: 385 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,385 @@
1+
// Copyright (C) 2025 Cydral (cydraltechnology@gmail.com)
2+
// License: Boost Software License See LICENSE.txt for the full license.
3+
#ifndef DLIB_DNN_LR_SCHEDULER_H_
4+
#define DLIB_DNN_LR_SCHEDULER_H_
5+
6+
#include "lr_scheduler_abstract.h"
7+
#include "../serialize.h"
8+
#include <cmath>
9+
#include <algorithm>
10+
#include <string>
11+
12+
namespace dlib
13+
{
14+
15+
// ----------------------------------------------------------------------------------------
16+
17+
namespace impl
18+
{
19+
constexpr double lr_scheduler_pi = 3.14159265358979323846;
20+
}
21+
22+
// ----------------------------------------------------------------------------------------
23+
24+
enum class lr_decay_type
25+
{
26+
COSINE,
27+
LINEAR,
28+
CONSTANT,
29+
EXPONENTIAL
30+
};
31+
32+
// ----------------------------------------------------------------------------------------
33+
34+
class lr_scheduler
35+
{
36+
public:
37+
38+
lr_scheduler(
39+
) :
40+
current_step_(0),
41+
warmup_steps_(2000),
42+
hold_steps_(0),
43+
total_steps_(100000),
44+
initial_lr_(1e-7),
45+
peak_lr_(3e-4),
46+
min_lr_(1e-6),
47+
decay_type_(lr_decay_type::COSINE)
48+
{
49+
compute_decay_steps();
50+
}
51+
52+
lr_scheduler(
53+
double peak_lr,
54+
size_t warmup_steps,
55+
size_t total_steps,
56+
double min_lr = 1e-6,
57+
lr_decay_type decay_type = lr_decay_type::COSINE
58+
) :
59+
current_step_(0),
60+
warmup_steps_(warmup_steps),
61+
hold_steps_(0),
62+
total_steps_(total_steps),
63+
initial_lr_(min_lr),
64+
peak_lr_(peak_lr),
65+
min_lr_(min_lr),
66+
decay_type_(decay_type)
67+
{
68+
DLIB_CASSERT(peak_lr > 0, "peak_lr must be positive");
69+
DLIB_CASSERT(min_lr >= 0, "min_lr must be non-negative");
70+
DLIB_CASSERT(min_lr < peak_lr, "min_lr must be less than peak_lr");
71+
DLIB_CASSERT(warmup_steps < total_steps, "warmup_steps must be less than total_steps");
72+
compute_decay_steps();
73+
}
74+
75+
double get_learning_rate(
76+
) const
77+
{
78+
// Phase 1: Warmup
79+
if (current_step_ < warmup_steps_)
80+
{
81+
if (warmup_steps_ == 0)
82+
return peak_lr_;
83+
const double progress = static_cast<double>(current_step_) / warmup_steps_;
84+
return initial_lr_ + (peak_lr_ - initial_lr_) * progress;
85+
}
86+
87+
// Phase 2: Hold (optional)
88+
const size_t post_warmup = current_step_ - warmup_steps_;
89+
if (post_warmup < hold_steps_)
90+
return peak_lr_;
91+
92+
// Phase 3: Decay
93+
if (decay_steps_ == 0)
94+
return peak_lr_;
95+
96+
const size_t decay_step = post_warmup - hold_steps_;
97+
const double progress = std::min(1.0, static_cast<double>(decay_step) / decay_steps_);
98+
99+
switch (decay_type_)
100+
{
101+
case lr_decay_type::COSINE:
102+
return min_lr_ + 0.5 * (peak_lr_ - min_lr_) * (1.0 + std::cos(impl::lr_scheduler_pi * progress));
103+
104+
case lr_decay_type::LINEAR:
105+
return peak_lr_ - (peak_lr_ - min_lr_) * progress;
106+
107+
case lr_decay_type::EXPONENTIAL:
108+
return peak_lr_ * std::pow(min_lr_ / peak_lr_, progress);
109+
110+
case lr_decay_type::CONSTANT:
111+
default:
112+
return peak_lr_;
113+
}
114+
}
115+
116+
double get_learning_rate(
117+
size_t step
118+
) const
119+
{
120+
lr_scheduler temp = *this;
121+
temp.current_step_ = step;
122+
return temp.get_learning_rate();
123+
}
124+
125+
void step(
126+
size_t n = 1
127+
)
128+
{
129+
current_step_ += n;
130+
}
131+
132+
void reset(
133+
)
134+
{
135+
current_step_ = 0;
136+
}
137+
138+
void set_current_step(
139+
size_t step
140+
)
141+
{
142+
current_step_ = step;
143+
}
144+
145+
size_t get_current_step(
146+
) const { return current_step_; }
147+
148+
size_t get_warmup_steps(
149+
) const { return warmup_steps_; }
150+
151+
size_t get_hold_steps(
152+
) const { return hold_steps_; }
153+
154+
size_t get_total_steps(
155+
) const { return total_steps_; }
156+
157+
size_t get_decay_steps(
158+
) const { return decay_steps_; }
159+
160+
double get_initial_lr(
161+
) const { return initial_lr_; }
162+
163+
double get_peak_lr(
164+
) const { return peak_lr_; }
165+
166+
double get_min_lr(
167+
) const { return min_lr_; }
168+
169+
lr_decay_type get_decay_type(
170+
) const { return decay_type_; }
171+
172+
void set_peak_lr(
173+
double lr
174+
)
175+
{
176+
DLIB_CASSERT(lr > 0 && lr > min_lr_);
177+
peak_lr_ = lr;
178+
}
179+
180+
void set_min_lr(
181+
double lr
182+
)
183+
{
184+
DLIB_CASSERT(lr >= 0 && lr < peak_lr_);
185+
min_lr_ = lr;
186+
}
187+
188+
void set_initial_lr(
189+
double lr
190+
)
191+
{
192+
DLIB_CASSERT(lr >= 0 && lr <= peak_lr_);
193+
initial_lr_ = lr;
194+
}
195+
196+
void set_warmup_steps(
197+
size_t steps
198+
)
199+
{
200+
DLIB_CASSERT(steps < total_steps_);
201+
warmup_steps_ = steps;
202+
compute_decay_steps();
203+
}
204+
205+
void set_hold_steps(
206+
size_t steps
207+
)
208+
{
209+
hold_steps_ = steps;
210+
compute_decay_steps();
211+
}
212+
213+
void set_total_steps(
214+
size_t steps
215+
)
216+
{
217+
DLIB_CASSERT(steps > warmup_steps_);
218+
total_steps_ = steps;
219+
compute_decay_steps();
220+
}
221+
222+
void set_decay_type(
223+
lr_decay_type type
224+
)
225+
{
226+
decay_type_ = type;
227+
}
228+
229+
bool is_warmup_complete(
230+
) const { return current_step_ >= warmup_steps_; }
231+
232+
bool is_training_complete(
233+
) const { return current_step_ >= total_steps_; }
234+
235+
double get_warmup_progress(
236+
) const
237+
{
238+
if (warmup_steps_ == 0)
239+
return 1.0;
240+
return std::min(1.0, static_cast<double>(current_step_) / warmup_steps_);
241+
}
242+
243+
double get_total_progress(
244+
) const
245+
{
246+
if (total_steps_ == 0)
247+
return 1.0;
248+
return std::min(1.0, static_cast<double>(current_step_) / total_steps_);
249+
}
250+
251+
std::string get_phase_name(
252+
) const
253+
{
254+
if (current_step_ < warmup_steps_)
255+
return "warmup";
256+
else if (current_step_ < warmup_steps_ + hold_steps_)
257+
return "hold";
258+
else
259+
return "decay";
260+
}
261+
262+
private:
263+
264+
void compute_decay_steps(
265+
)
266+
{
267+
const size_t non_decay = warmup_steps_ + hold_steps_;
268+
decay_steps_ = (total_steps_ > non_decay) ? (total_steps_ - non_decay) : 0;
269+
}
270+
271+
size_t current_step_;
272+
size_t warmup_steps_;
273+
size_t hold_steps_;
274+
size_t total_steps_;
275+
size_t decay_steps_;
276+
double initial_lr_;
277+
double peak_lr_;
278+
double min_lr_;
279+
lr_decay_type decay_type_;
280+
};
281+
282+
// ----------------------------------------------------------------------------------------
283+
284+
inline void serialize(
285+
const lr_scheduler& item,
286+
std::ostream& out
287+
)
288+
{
289+
serialize("lr_scheduler", out);
290+
serialize(item.get_current_step(), out);
291+
serialize(item.get_warmup_steps(), out);
292+
serialize(item.get_hold_steps(), out);
293+
serialize(item.get_total_steps(), out);
294+
serialize(item.get_decay_steps(), out);
295+
serialize(item.get_initial_lr(), out);
296+
serialize(item.get_peak_lr(), out);
297+
serialize(item.get_min_lr(), out);
298+
serialize(static_cast<int>(item.get_decay_type()), out);
299+
}
300+
301+
inline void deserialize(
302+
lr_scheduler& item,
303+
std::istream& in
304+
)
305+
{
306+
std::string version;
307+
deserialize(version, in);
308+
if (version != "lr_scheduler")
309+
throw serialization_error("Unexpected version '" + version +
310+
"' found while deserializing lr_scheduler.");
311+
312+
size_t current_step, warmup_steps, hold_steps, total_steps, decay_steps;
313+
double initial_lr, peak_lr, min_lr;
314+
int decay_type_int;
315+
316+
deserialize(current_step, in);
317+
deserialize(warmup_steps, in);
318+
deserialize(hold_steps, in);
319+
deserialize(total_steps, in);
320+
deserialize(decay_steps, in);
321+
deserialize(initial_lr, in);
322+
deserialize(peak_lr, in);
323+
deserialize(min_lr, in);
324+
deserialize(decay_type_int, in);
325+
326+
item = lr_scheduler(peak_lr, warmup_steps, total_steps, min_lr,
327+
static_cast<lr_decay_type>(decay_type_int));
328+
item.set_initial_lr(initial_lr);
329+
item.set_hold_steps(hold_steps);
330+
item.set_current_step(current_step);
331+
}
332+
333+
inline std::ostream& operator<<(
334+
std::ostream& out,
335+
const lr_scheduler& item
336+
)
337+
{
338+
out << "lr_scheduler ("
339+
<< "step=" << item.get_current_step()
340+
<< ", lr=" << item.get_learning_rate()
341+
<< ", phase=" << item.get_phase_name()
342+
<< ", warmup=" << item.get_warmup_steps()
343+
<< ", total=" << item.get_total_steps()
344+
<< ", peak=" << item.get_peak_lr()
345+
<< ", min=" << item.get_min_lr()
346+
<< ")";
347+
return out;
348+
}
349+
350+
// ----------------------------------------------------------------------------------------
351+
352+
inline lr_scheduler make_transformer_scheduler(
353+
double peak_lr,
354+
size_t total_steps,
355+
double warmup_fraction = 0.02,
356+
double min_lr = 1e-6,
357+
lr_decay_type decay_type = lr_decay_type::COSINE
358+
)
359+
{
360+
DLIB_CASSERT(peak_lr > 0, "peak_lr must be positive");
361+
DLIB_CASSERT(total_steps > 0, "total_steps must be positive");
362+
DLIB_CASSERT(warmup_fraction > 0 && warmup_fraction < 1, "warmup_fraction must be in (0, 1)");
363+
DLIB_CASSERT(min_lr >= 0 && min_lr < peak_lr, "min_lr must be in [0, peak_lr)");
364+
365+
size_t warmup_steps = static_cast<size_t>(total_steps * warmup_fraction);
366+
warmup_steps = std::max(size_t(100), warmup_steps);
367+
return lr_scheduler(peak_lr, warmup_steps, total_steps, min_lr, decay_type);
368+
}
369+
370+
inline size_t estimate_total_steps(
371+
size_t dataset_size,
372+
size_t batch_size,
373+
size_t num_epochs
374+
)
375+
{
376+
DLIB_CASSERT(batch_size > 0, "batch_size must be positive");
377+
const size_t steps_per_epoch = (dataset_size + batch_size - 1) / batch_size;
378+
return steps_per_epoch * num_epochs;
379+
}
380+
381+
// ----------------------------------------------------------------------------------------
382+
383+
}
384+
385+
#endif // DLIB_DNN_LR_SCHEDULER_H_

0 commit comments

Comments
 (0)