Migrations
/
pytorch-image-models
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Update schedulers

Browse Source
pull/1/head
Ross Wightman 6 years ago
parent b5255960d9
commit b1a5a71151
5 changed files with 66 additions and 51 deletions
Show Stats Download Patch File Download Diff File

49
scheduler/cosine_lr.py
Unescape View File

@ -21,8 +21,10 @@ class CosineLRScheduler(Scheduler):
t_mul: float = 1., t_mul: float = 1.,
lr_min: float = 0., lr_min: float = 0.,
decay_rate: float = 1., decay_rate: float = 1.,
warmup_updates=0, warmup_t=0,
warmup_lr_init=0, warmup_lr_init=0,
warmup_prefix=False,
t_in_epochs=True,
initialize=True) -> None: initialize=True) -> None:
super().__init__(optimizer, param_group_field="lr", initialize=initialize) super().__init__(optimizer, param_group_field="lr", initialize=initialize)
@ -35,32 +37,31 @@ class CosineLRScheduler(Scheduler):
self.t_mul = t_mul self.t_mul = t_mul
self.lr_min = lr_min self.lr_min = lr_min
self.decay_rate = decay_rate self.decay_rate = decay_rate
self.warmup_updates = warmup_updates self.warmup_t = warmup_t
self.warmup_lr_init = warmup_lr_init self.warmup_lr_init = warmup_lr_init
if self.warmup_updates: self.warmup_prefix = warmup_prefix
self.warmup_steps = [(v - warmup_lr_init) / self.warmup_updates for v in self.base_values] self.t_in_epochs = t_in_epochs
if self.warmup_t:
self.warmup_steps = [(v - warmup_lr_init) / self.warmup_t for v in self.base_values]
super().update_groups(self.warmup_lr_init)
else: else:
self.warmup_steps = [1 for _ in self.base_values] self.warmup_steps = [1 for _ in self.base_values]
if self.warmup_lr_init:
super().update_groups(self.warmup_lr_init)
def get_epoch_values(self, epoch: int):
# this scheduler doesn't update on epoch
return None
def get_update_values(self, num_updates: int): def _get_lr(self, t):
if num_updates < self.warmup_updates: if t < self.warmup_t:
lrs = [self.warmup_lr_init + num_updates * s for s in self.warmup_steps] lrs = [self.warmup_lr_init + t * s for s in self.warmup_steps]
else: else:
curr_updates = num_updates - self.warmup_updates if self.warmup_prefix:
t = t - self.warmup_t
if self.t_mul != 1: if self.t_mul != 1:
i = math.floor(math.log(1 - curr_updates / self.t_initial * (1 - self.t_mul), self.t_mul)) i = math.floor(math.log(1 - t / self.t_initial * (1 - self.t_mul), self.t_mul))
t_i = self.t_mul ** i * self.t_initial t_i = self.t_mul ** i * self.t_initial
t_curr = curr_updates - (1 - self.t_mul ** i) / (1 - self.t_mul) * self.t_initial t_curr = t - (1 - self.t_mul ** i) / (1 - self.t_mul) * self.t_initial
else: else:
i = curr_updates // self.t_initial i = t // self.t_initial
t_i = self.t_initial t_i = self.t_initial
t_curr = curr_updates - (self.t_initial * i) t_curr = t - (self.t_initial * i)
gamma = self.decay_rate ** i gamma = self.decay_rate ** i
lr_min = self.lr_min * gamma lr_min = self.lr_min * gamma
@ -70,3 +71,15 @@ class CosineLRScheduler(Scheduler):
lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * t_curr / t_i)) for lr_max in lr_max_values lr_min + 0.5 * (lr_max - lr_min) * (1 + math.cos(math.pi * t_curr / t_i)) for lr_max in lr_max_values
] ]
return lrs return lrs
def get_epoch_values(self, epoch: int):
if self.t_in_epochs:
return self._get_lr(epoch)
else:
return None
def get_update_values(self, num_updates: int):
if not self.t_in_epochs:
return self._get_lr(num_updates)
else:
return None

2
scheduler/scheduler.py
Unescape View File

@ -56,7 +56,7 @@ class Scheduler:
def step(self, epoch: int, metric: float = None) -> None: def step(self, epoch: int, metric: float = None) -> None:
self.metric = metric self.metric = metric
values = self.get_epoch_values(epoch) values = self.get_epoch_values(epoch + 1) # +1 to calculate for next epoch
if values is not None: if values is not None:
self.update_groups(values) self.update_groups(values)

42
scheduler/step_lr.py
Unescape View File

@ -10,39 +10,41 @@ class StepLRScheduler(Scheduler):
def __init__(self, def __init__(self,
optimizer: torch.optim.Optimizer, optimizer: torch.optim.Optimizer,
decay_epochs: int, decay_t: int,
decay_rate: float = 1., decay_rate: float = 1.,
warmup_updates=0, warmup_t=0,
warmup_lr_init=0, warmup_lr_init=0,
t_in_epochs=True,
initialize=True) -> None: initialize=True) -> None:
super().__init__(optimizer, param_group_field="lr", initialize=initialize) super().__init__(optimizer, param_group_field="lr", initialize=initialize)
self.decay_epochs = decay_epochs self.decay_t = decay_t
self.decay_rate = decay_rate self.decay_rate = decay_rate
self.warmup_updates = warmup_updates self.warmup_t = warmup_t
self.warmup_lr_init = warmup_lr_init self.warmup_lr_init = warmup_lr_init
self.t_in_epochs = t_in_epochs
if self.warmup_updates: if self.warmup_t:
self.warmup_active = warmup_updates > 0 # this state updates with num_updates self.warmup_steps = [(v - warmup_lr_init) / self.warmup_t for v in self.base_values]
self.warmup_steps = [(v - warmup_lr_init) / self.warmup_updates for v in self.base_values]
super().update_groups(self.warmup_lr_init) super().update_groups(self.warmup_lr_init)
else: else:
self.warmup_steps = [1 for _ in self.base_values] self.warmup_steps = [1 for _ in self.base_values]
def get_epoch_values(self, epoch: int): def _get_lr(self, t):
if not self.warmup_active: if t < self.warmup_t:
lrs = [v * (self.decay_rate ** ((epoch + 1) // self.decay_epochs)) lrs = [self.warmup_lr_init + t * s for s in self.warmup_steps]
for v in self.base_values]
else: else:
lrs = None # no epoch updates while warming up lrs = [v * (self.decay_rate ** (t // self.decay_t))
for v in self.base_values]
return lrs return lrs
def get_update_values(self, num_updates: int): def get_epoch_values(self, epoch: int):
if num_updates < self.warmup_updates: if self.t_in_epochs:
lrs = [self.warmup_lr_init + num_updates * s for s in self.warmup_steps] return self._get_lr(epoch)
else: else:
self.warmup_active = False # warmup cancelled by first update past warmup_update count return None
lrs = None # no change on update afte warmup stage
return lrs
def get_update_values(self, num_updates: int):
if not self.t_in_epochs:
return self._get_lr(num_updates)
else:
return None

2
scheduler/tanh_lr.py
Unescape View File

@ -27,7 +27,7 @@ class TanhLRScheduler(Scheduler):
warmup_lr_init=0, warmup_lr_init=0,
warmup_prefix=False, warmup_prefix=False,
cycle_limit=0, cycle_limit=0,
t_in_epochs=False, t_in_epochs=True,
initialize=True) -> None: initialize=True) -> None:
super().__init__(optimizer, param_group_field="lr", initialize=initialize) super().__init__(optimizer, param_group_field="lr", initialize=initialize)

22
train.py
Unescape View File

@ -162,7 +162,7 @@ def main():
if args.opt.lower() == 'sgd': if args.opt.lower() == 'sgd':
optimizer = optim.SGD( optimizer = optim.SGD(
model.parameters(), lr=args.lr, model.parameters(), lr=args.lr,
momentum=args.momentum, weight_decay=args.weight_decay, nesterov=False) momentum=args.momentum, weight_decay=args.weight_decay, nesterov=True)
elif args.opt.lower() == 'adam': elif args.opt.lower() == 'adam':
optimizer = optim.Adam( optimizer = optim.Adam(
model.parameters(), lr=args.lr, weight_decay=args.weight_decay, eps=args.opt_eps) model.parameters(), lr=args.lr, weight_decay=args.weight_decay, eps=args.opt_eps)
@ -183,32 +183,32 @@ def main():
if optimizer_state is not None: if optimizer_state is not None:
optimizer.load_state_dict(optimizer_state) optimizer.load_state_dict(optimizer_state)
updates_per_epoch = len(loader_train)
if args.sched == 'cosine': if args.sched == 'cosine':
lr_scheduler = scheduler.CosineLRScheduler( lr_scheduler = scheduler.CosineLRScheduler(
optimizer, optimizer,
t_initial=100 * updates_per_epoch, t_initial=130,
t_mul=1.0, t_mul=1.0,
lr_min=0, lr_min=0,
decay_rate=0.5, decay_rate=args.decay_rate,
warmup_lr_init=1e-4, warmup_lr_init=1e-4,
warmup_updates=1 * updates_per_epoch warmup_t=3,
t_in_epochs=True,
) )
elif args.sched == 'tanh': elif args.sched == 'tanh':
lr_scheduler = scheduler.TanhLRScheduler( lr_scheduler = scheduler.TanhLRScheduler(
optimizer, optimizer,
t_initial=80 * updates_per_epoch, t_initial=130,
t_mul=1.0, t_mul=1.0,
lr_min=1e-5, lr_min=1e-6,
decay_rate=0.5,
warmup_lr_init=.001, warmup_lr_init=.001,
warmup_t=5 * updates_per_epoch, warmup_t=3,
cycle_limit=1 cycle_limit=1,
t_in_epochs=True,
) )
else: else:
lr_scheduler = scheduler.StepLRScheduler( lr_scheduler = scheduler.StepLRScheduler(
optimizer, optimizer,
decay_epochs=args.decay_epochs, decay_t=args.decay_epochs,
decay_rate=args.decay_rate, decay_rate=args.decay_rate,
) )

Write Preview
Loading…
Cancel
Save