From 6b45b354473d5775cffe2985c442d8c04d189c66 Mon Sep 17 00:00:00 2001 From: James Betker Date: Fri, 31 Jul 2020 11:21:11 -0600 Subject: [PATCH] Allow multi_step_lr_scheduler to load a new LR schedule when restoring state --- codes/models/lr_scheduler.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/codes/models/lr_scheduler.py b/codes/models/lr_scheduler.py index 63503cff..da90c223 100644 --- a/codes/models/lr_scheduler.py +++ b/codes/models/lr_scheduler.py @@ -54,6 +54,12 @@ class MultiStepLR_Restart(_LRScheduler): for group in self.optimizer.param_groups ] + # Allow this scheduler to use newly appointed milestones partially through a training run.. + def load_state_dict(self, s): + milestones_cache = self.milestones + super(MultiStepLR_Restart, self).load_state_dict(s) + self.milestones = milestones_cache + class CosineAnnealingLR_Restart(_LRScheduler): def __init__(self, optimizer, T_period, restarts=None, weights=None, eta_min=0, last_epoch=-1):