kxngluca/DL-Art-School
1
0
Fork 0
forked from mrq/DL-Art-School
Code Pull Requests Activity

Force LR fix

Browse Source
This commit is contained in:
James Betker 2021-10-21 12:01:01 -06:00
parent 40cb25292a
commit 9a3e89ec53
1 changed files with 3 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

4
codes/trainer/lr_scheduler.py
View File

@ -76,7 +76,7 @@ class MultiStepLR_Restart(_LRScheduler):
def get_lr(self): def get_lr(self):
# Note to self: for the purposes of this trainer, "last_epoch" should read "last_step" # Note to self: for the purposes of this trainer, "last_epoch" should read "last_step"
if self.force_lr is not None: if self.force_lr is not None:
return [self.force_lr for group in self.optimizer.param_groups] return [self.force_lr for _ in self.optimizer.param_groups]
if self.last_epoch in self.restarts: if self.last_epoch in self.restarts:
if self.clear_state: if self.clear_state:
self.optimizer.state = defaultdict(dict) self.optimizer.state = defaultdict(dict)
@ -95,8 +95,10 @@ class MultiStepLR_Restart(_LRScheduler):
# Allow this scheduler to use newly appointed milestones partially through a training run.. # Allow this scheduler to use newly appointed milestones partially through a training run..
def load_state_dict(self, s): def load_state_dict(self, s):
milestones_cache = self.milestones milestones_cache = self.milestones
force_lr_cache = self.force_lr
super(MultiStepLR_Restart, self).load_state_dict(s) super(MultiStepLR_Restart, self).load_state_dict(s)
self.milestones = milestones_cache self.milestones = milestones_cache
self.force_lr = force_lr_cache
class CosineAnnealingLR_Restart(_LRScheduler): class CosineAnnealingLR_Restart(_LRScheduler):