github
/
ColossalAI
mirror of https://github.com/hpcaitech/ColossalAI
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

[NFC] polish <colossalai/nn/lr_scheduler/poly.py> code style (#1267)

pull/1272/head
oahzxl 2022-07-12 18:18:14 +08:00 committed by GitHub
parent c92f84fcdb
commit 0cf8e8e91c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 18 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

28
colossalai/nn/lr_scheduler/poly.py
View File

@ -17,7 +17,12 @@ class PolynomialLR(_LRScheduler):
the schedule is started from the beginning or When last_epoch=-1, sets initial lr as lr.
"""
def __init__(self, optimizer, total_steps: int, end_lr: float = 0.0001, power: float = 1.0, last_epoch: int = -1,
def __init__(self,
optimizer,
total_steps: int,
end_lr: float = 0.0001,
power: float = 1.0,
last_epoch: int = -1,
**kwargs):
if end_lr < 0:
raise ValueError(f'end_lr must >= 0, got {end_lr}')
@ -30,11 +35,9 @@ class PolynomialLR(_LRScheduler):
return self._get_closed_form_lr()
def _get_closed_form_lr(self):
return [
(base_lr - self.end_lr) * ((1 - min(self.last_epoch, self.total_steps) /
self.total_steps) ** self.power) + self.end_lr
for base_lr in self.base_lrs
]
return [(base_lr - self.end_lr) *
((1 - min(self.last_epoch, self.total_steps) / self.total_steps)**self.power) + self.end_lr
for base_lr in self.base_lrs]
@LR_SCHEDULERS.register_module
@ -51,8 +54,13 @@ class PolynomialWarmupLR(WarmupScheduler):
the schedule is started from the beginning or When last_epoch=-1, sets initial lr as lr.
"""
def __init__(self, optimizer, total_steps: int, warmup_steps: int = 0, end_lr: float = 0.0001, power: float = 1.0,
last_epoch: int = -1, **kwargs):
base_scheduler = PolynomialLR(
optimizer, total_steps - warmup_steps, end_lr=end_lr, power=power)
def __init__(self,
optimizer,
total_steps: int,
warmup_steps: int = 0,
end_lr: float = 0.0001,
power: float = 1.0,
last_epoch: int = -1,
**kwargs):
base_scheduler = PolynomialLR(optimizer, total_steps - warmup_steps, end_lr=end_lr, power=power)
super().__init__(optimizer, warmup_steps, base_scheduler, last_epoch=last_epoch)