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/onecycle.py code style (#1269)

pull/1298/head
Super Daniel 2022-07-12 18:56:34 +08:00 committed by Frank Lee
parent 0e06f62160
commit 52d145a342
1 changed files with 8 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
colossalai/nn/lr_scheduler/onecycle.py
View File

@ -68,7 +68,9 @@ class OneCycleLR(_OneCycleLR):
https://arxiv.org/abs/1708.07120
"""
def __init__(self, optimizer, total_steps: int,
def __init__(self,
optimizer,
total_steps: int,
pct_start=0.3,
anneal_strategy='cos',
cycle_momentum=True,
@ -76,9 +78,12 @@ class OneCycleLR(_OneCycleLR):
max_momentum=0.95,
div_factor=25.0,
final_div_factor=10000.0,
last_epoch=-1, **kwargs):
last_epoch=-1,
**kwargs):
max_lrs = list(map(lambda group: group['lr'], optimizer.param_groups))
super().__init__(optimizer, max_lrs, total_steps=total_steps,
super().__init__(optimizer,
max_lrs,
total_steps=total_steps,
pct_start=pct_start,
anneal_strategy=anneal_strategy,
cycle_momentum=cycle_momentum,