WIP early stopping
@@ -12,9 +12,20 @@ class EarlyStoppingPlugin(StrategyPlugin):
The state of the best model is saved after each improvement on the
given metric and is loaded back into the model before stopping the
training procedure.
+
+ .. warning::
+ The plugin checks the metric value, which is updated by the strategy
+ during the evaluation. This means that you must ensure that the
+ evaluation is called frequently enough during the training loop.
+
+ For example, if you set `patience=1`, you must also set `eval_every=1`
+ in the `BaseStrategy`, otherwise the metric won't be updated after
+ every epoch/iteration.
+
"""
def __init__(self, patience: int, val_stream_name: str,
- metric_name: str = 'Top1_Acc_Stream', mode: str = 'max'):
+ metric_name: str = 'Top1_Acc_Stream', mode: str = 'max',
+ peval_mode: str='epoch'):
"""
:param patience: Number of epochs to wait before stopping the training.
:param val_stream_name: Name of the validation stream to search in the
@@ -24,10 +35,17 @@ class EarlyStoppingPlugin(StrategyPlugin):
reported in the evaluator.
:param mode: Must be "max" or "min". max (resp. min) means that the
given metric should me maximized (resp. minimized).
+ :param peval_mode: one of {'epoch', 'iteration'}. Decides whether the
+ early stopping should happen after `patience`
+ epochs or iterations (Default='epoch').
"""
super().__init__()
self.val_stream_name = val_stream_name
self.patience = patience
+
+ assert peval_mode in {'epoch', 'iteration'}
+ self.peval_mode = peval_mode
+
self.metric_name = metric_name
self.metric_key = f'{self.metric_name}/eval_phase/' \
f'{self.val_stream_name}'
@@ -37,16 +55,26 @@ class EarlyStoppingPlugin(StrategyPlugin):
self.best_state = None # Contains the best parameters
self.best_val = None
- self.best_epoch = None
+ self.best_step = None
def before_training(self, strategy, **kwargs):
self.best_state = None
self.best_val = None
- self.best_epoch = None
+ self.best_step = None
+
+ def before_training_iteration(self, strategy, **kwargs):
+ if self.peval_mode == 'iteration':
+ self._update_best(strategy)
+ curr_step = self._get_strategy_counter(strategy)
+ if curr_step - self.best_step >= self.patience:
+ strategy.model.load_state_dict(self.best_state)
+ strategy.stop_training()
def before_training_epoch(self, strategy, **kwargs):
+ if self.peval_mode == 'epoch':
self._update_best(strategy)
- if strategy.clock.train_exp_epochs - self.best_epoch >= self.patience:
+ curr_step = self._get_strategy_counter(strategy)
+ if curr_step - self.best_step >= self.patience:
strategy.model.load_state_dict(self.best_state)
strategy.stop_training()
@@ -56,4 +84,12 @@ class EarlyStoppingPlugin(StrategyPlugin):
if self.best_val is None or self.operator(val_acc, self.best_val):
self.best_state = deepcopy(strategy.model.state_dict())
self.best_val = val_acc
- self.best_epoch = strategy.clock.train_exp_epochs
+ self.best_step = self._get_strategy_counter(strategy)
+
+ def _get_strategy_counter(self, strategy):
+ if self.peval_mode == 'epoch':
+ return strategy.clock.train_exp_epochs
+ elif self.peval_mode == 'iteration':
+ return strategy.clock.train_exp_iterations
+ else:
+ raise ValueError("Invalid `peval_mode`:", self.peval_mode)
\ No newline at end of file
