def on_stage_start(self, state: _State):
"""
Checks that the current stage has correct criterion
"""
criterion = state.get_attr(
key="criterion", inner_key=self.criterion_key
)
assert criterion is not None
self._criterion = criterion
def on_stage_start(self, state: _State):
"""
Checks that the current stage has correct optimizer
"""
optimizer = state.get_attr(
key="optimizer", inner_key=self.optimizer_key
)
assert optimizer is not None
self._optimizer = optimizer
def on_stage_start(self, state: _State):
scheduler = state.get_attr(key="scheduler",
inner_key=self.scheduler_key)
assert scheduler is not None
self._scheduler = scheduler
if self.mode is None:
if isinstance(scheduler, BatchScheduler):
self.mode = "batch"
else:
self.mode = "epoch"
if isinstance(scheduler, OneCycleLRWithWarmup) and \
self.mode == "batch":
scheduler.reset()
def on_stage_start(self, state: _State):
optimizer = state.get_attr(key="optimizer",
inner_key=self.optimizer_key)
assert optimizer is not None
self._optimizer = optimizer
self.init_lr = optimizer.defaults["lr"]