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
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
|
from torch.optim.lr_scheduler import ReduceLROnPlateau, StepLR, MultiStepLR
from transformers import (get_linear_schedule_with_warmup,
get_cosine_schedule_with_warmup)
def initialize_scheduler(config, optimizer, n_train_steps):
# construct schedulers
if config.scheduler is None:
return None
elif config.scheduler == 'linear_schedule_with_warmup':
scheduler = get_linear_schedule_with_warmup(
optimizer,
num_training_steps=n_train_steps,
**config.scheduler_kwargs)
step_every_batch = True
use_metric = False
elif config.scheduler == 'cosine_schedule_with_warmup':
scheduler = get_cosine_schedule_with_warmup(
optimizer,
num_training_steps=n_train_steps,
**config.scheduler_kwargs)
step_every_batch = True
use_metric = False
elif config.scheduler == 'ReduceLROnPlateau':
assert config.scheduler_metric_name, f'scheduler metric must be specified for {config.scheduler}'
scheduler = ReduceLROnPlateau(
optimizer,
**config.scheduler_kwargs)
step_every_batch = False
use_metric = True
elif config.scheduler == 'StepLR':
scheduler = StepLR(optimizer, **config.scheduler_kwargs)
step_every_batch = False
use_metric = False
elif config.scheduler == 'MultiStepLR':
scheduler = MultiStepLR(optimizer, **config.scheduler_kwargs)
step_every_batch = False
use_metric = False
else:
raise ValueError('Scheduler not recognized.')
# add a step_every_batch field
scheduler.step_every_batch = step_every_batch
scheduler.use_metric = use_metric
return scheduler
def step_scheduler(scheduler, metric=None):
if isinstance(scheduler, ReduceLROnPlateau):
assert metric is not None
scheduler.step(metric)
else:
scheduler.step()
|