def _init_state(self,
*,
mode: str,
stage: str = None,
**kwargs) -> RunnerState:
"""
Inner method for children's classes for state specific initialization.
:return: RunnerState with all necessary parameters.
"""
additional_kwargs = {}
# transfer previous counters from old state
if self.state is not None:
additional_kwargs = {
"step": self.state.step,
"epoch": self.state.epoch + 1,
"best_metrics": self.state.best_metrics
}
return RunnerState(device=self.device,
model=self.model,
stage=self.stage,
criterion=self.criterion,
optimizer=self.optimizer,
scheduler=self.scheduler,
**kwargs,
**additional_kwargs)
def _prepare_state(self, stage: str):
migrating_params = {}
if self.state is not None:
migrating_params.update({
"step": self.state.step,
"epoch": self.state.epoch + 1
})
self.model, criterion, optimizer, scheduler, self.device = \
self._get_experiment_components(stage)
self.state = RunnerState(stage=stage,
model=self.model,
device=self.device,
criterion=criterion,
optimizer=optimizer,
scheduler=scheduler,
**self.experiment.get_state_params(stage),
**migrating_params)