|
|
from ..core import * |
|
|
from ..callback import * |
|
|
from ..basic_train import Learner, LearnerCallback |
|
|
|
|
|
__all__ = ['GeneralScheduler', 'TrainingPhase'] |
|
|
|
|
|
@dataclass |
|
|
class TrainingPhase(): |
|
|
"Schedule hyper-parameters for a phase of `length` iterations." |
|
|
length:int |
|
|
|
|
|
def __post_init__(self): self.scheds = dict() |
|
|
def schedule_hp(self, name, vals, anneal=None): |
|
|
"Adds a schedule for `name` between `vals` using `anneal`." |
|
|
self.scheds[name] = Scheduler(vals, self.length, anneal) |
|
|
return self |
|
|
|
|
|
class GeneralScheduler(LearnerCallback): |
|
|
"Schedule multiple `TrainingPhase` for a `Learner`." |
|
|
def __init__(self, learn:Learner, phases:Collection[TrainingPhase], start_epoch:int=None): |
|
|
super().__init__(learn) |
|
|
self.phases,self.start_epoch = phases,start_epoch |
|
|
|
|
|
def on_train_begin(self, epoch:int, **kwargs:Any)->None: |
|
|
"Initialize the schedulers for training." |
|
|
res = {'epoch':self.start_epoch} if self.start_epoch is not None else None |
|
|
self.start_epoch = ifnone(self.start_epoch, epoch) |
|
|
self.scheds = [p.scheds for p in self.phases] |
|
|
self.opt = self.learn.opt |
|
|
for k,v in self.scheds[0].items(): |
|
|
v.restart() |
|
|
self.opt.set_stat(k, v.start) |
|
|
self.idx_s = 0 |
|
|
return res |
|
|
|
|
|
def jump_to_epoch(self, epoch:int)->None: |
|
|
for _ in range(len(self.learn.data.train_dl) * epoch): |
|
|
self.on_batch_end(True) |
|
|
|
|
|
def on_batch_end(self, train, **kwargs:Any)->None: |
|
|
"Take a step in lr,mom sched, start next stepper when the current one is complete." |
|
|
if train: |
|
|
if self.idx_s >= len(self.scheds): return {'stop_training': True, 'stop_epoch': True} |
|
|
sched = self.scheds[self.idx_s] |
|
|
for k,v in sched.items(): self.opt.set_stat(k, v.step()) |
|
|
if list(sched.values())[0].is_done: self.idx_s += 1 |
