def on_stage_start(self, runner: IRunner): """Checks that the current stage has correct criterion.""" criterion = runner.get_attr( key="criterion", inner_key=self.criterion_key ) assert criterion is not None self._criterion = criterion
def on_stage_start(self, runner: IRunner) -> None: """Stage start hook. Args: runner: current runner """ self.reduced_metric = self.reduced_metric or runner.main_metric scheduler = runner.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() assert self.mode is not None
def on_stage_start(self, runner: IRunner) -> None: """Checks that the current stage has correct optimizer. Args: runner(IRunner): current runner """ self._optimizer = runner.get_attr(key="optimizer", inner_key=self.optimizer_key) assert self._optimizer is not None
def on_stage_start(self, runner: IRunner) -> None: """Stage start hook. Args: runner (IRunner): current runner """ optimizer = runner.get_attr(key="optimizer", inner_key=self.optimizer_key) assert optimizer is not None self._optimizer = optimizer self.init_lr = optimizer.defaults["lr"]
def on_stage_start(self, runner: IRunner) -> None: """Checks that the current stage has correct optimizer. Args: runner(IRunner): current runner """ from torch.cuda.amp import GradScaler self._optimizer = runner.get_attr(key="optimizer", inner_key=self.optimizer_key) self.scaler = GradScaler() assert self._optimizer is not None
def on_stage_start(self, runner: IRunner) -> None: """Checks that the current stage has correct optimizer. Args: runner(IRunner): current runner """ self._optimizer = runner.get_attr(key="optimizer", inner_key=self.optimizer_key) # device based optimization step if runner.device.type == "xla": self._optimizer_step_fn = self._optimizer_step_tpu else: self._optimizer_step_fn = self._optimizer_step assert self._optimizer is not None