Shortcuts

mmengine.hooks.hook 源代码

# Copyright (c) OpenMMLab. All rights reserved.
from typing import Dict, Optional, Sequence, Union

from mmengine import is_method_overridden

DATA_BATCH = Optional[Union[dict, tuple, list]]


[文档]class Hook: """Base hook class. All hooks should inherit from this class. """ priority = 'NORMAL' stages = ('before_run', 'after_load_checkpoint', 'before_train', 'before_train_epoch', 'before_train_iter', 'after_train_iter', 'after_train_epoch', 'before_val', 'before_val_epoch', 'before_val_iter', 'after_val_iter', 'after_val_epoch', 'after_val', 'before_save_checkpoint', 'after_train', 'before_test', 'before_test_epoch', 'before_test_iter', 'after_test_iter', 'after_test_epoch', 'after_test', 'after_run')
[文档] def before_run(self, runner) -> None: """All subclasses should override this method, if they need any operations before the training validation or testing process. Args: runner (Runner): The runner of the training, validation or testing process. """
[文档] def after_run(self, runner) -> None: """All subclasses should override this method, if they need any operations before the training validation or testing process. Args: runner (Runner): The runner of the training, validation or testing process. """
[文档] def before_train(self, runner) -> None: """All subclasses should override this method, if they need any operations before train. Args: runner (Runner): The runner of the training process. """
[文档] def after_train(self, runner) -> None: """All subclasses should override this method, if they need any operations after train. Args: runner (Runner): The runner of the training process. """
[文档] def before_val(self, runner) -> None: """All subclasses should override this method, if they need any operations before validation. Args: runner (Runner): The runner of the validation process. """
[文档] def after_val(self, runner) -> None: """All subclasses should override this method, if they need any operations after validation. Args: runner (Runner): The runner of the validation process. """
[文档] def before_test(self, runner) -> None: """All subclasses should override this method, if they need any operations before testing. Args: runner (Runner): The runner of the testing process. """
[文档] def after_test(self, runner) -> None: """All subclasses should override this method, if they need any operations after testing. Args: runner (Runner): The runner of the testing process. """
[文档] def before_save_checkpoint(self, runner, checkpoint: dict) -> None: """All subclasses should override this method, if they need any operations before saving the checkpoint. Args: runner (Runner): The runner of the training, validation or testing process. checkpoint (dict): Model's checkpoint. """
[文档] def after_load_checkpoint(self, runner, checkpoint: dict) -> None: """All subclasses should override this method, if they need any operations after loading the checkpoint. Args: runner (Runner): The runner of the training, validation or testing process. checkpoint (dict): Model's checkpoint. """
[文档] def before_train_epoch(self, runner) -> None: """All subclasses should override this method, if they need any operations before each training epoch. Args: runner (Runner): The runner of the training process. """ self._before_epoch(runner, mode='train')
[文档] def before_val_epoch(self, runner) -> None: """All subclasses should override this method, if they need any operations before each validation epoch. Args: runner (Runner): The runner of the validation process. """ self._before_epoch(runner, mode='val')
[文档] def before_test_epoch(self, runner) -> None: """All subclasses should override this method, if they need any operations before each test epoch. Args: runner (Runner): The runner of the testing process. """ self._before_epoch(runner, mode='test')
[文档] def after_train_epoch(self, runner) -> None: """All subclasses should override this method, if they need any operations after each training epoch. Args: runner (Runner): The runner of the training process. """ self._after_epoch(runner, mode='train')
[文档] def after_val_epoch(self, runner, metrics: Optional[Dict[str, float]] = None) -> None: """All subclasses should override this method, if they need any operations after each validation epoch. Args: runner (Runner): The runner of the validation process. metrics (Dict[str, float], optional): Evaluation results of all metrics on validation dataset. The keys are the names of the metrics, and the values are corresponding results. """ self._after_epoch(runner, mode='val')
[文档] def after_test_epoch(self, runner, metrics: Optional[Dict[str, float]] = None) -> None: """All subclasses should override this method, if they need any operations after each test epoch. Args: runner (Runner): The runner of the testing process. metrics (Dict[str, float], optional): Evaluation results of all metrics on test dataset. The keys are the names of the metrics, and the values are corresponding results. """ self._after_epoch(runner, mode='test')
[文档] def before_train_iter(self, runner, batch_idx: int, data_batch: DATA_BATCH = None) -> None: """All subclasses should override this method, if they need any operations before each training iteration. Args: runner (Runner): The runner of the training process. batch_idx (int): The index of the current batch in the train loop. data_batch (dict or tuple or list, optional): Data from dataloader. """ self._before_iter( runner, batch_idx=batch_idx, data_batch=data_batch, mode='train')
[文档] def before_val_iter(self, runner, batch_idx: int, data_batch: DATA_BATCH = None) -> None: """All subclasses should override this method, if they need any operations before each validation iteration. Args: runner (Runner): The runner of the validation process. batch_idx (int): The index of the current batch in the val loop. data_batch (dict, optional): Data from dataloader. Defaults to None. """ self._before_iter( runner, batch_idx=batch_idx, data_batch=data_batch, mode='val')
[文档] def before_test_iter(self, runner, batch_idx: int, data_batch: DATA_BATCH = None) -> None: """All subclasses should override this method, if they need any operations before each test iteration. Args: runner (Runner): The runner of the testing process. batch_idx (int): The index of the current batch in the test loop. data_batch (dict or tuple or list, optional): Data from dataloader. Defaults to None. """ self._before_iter( runner, batch_idx=batch_idx, data_batch=data_batch, mode='test')
[文档] def after_train_iter(self, runner, batch_idx: int, data_batch: DATA_BATCH = None, outputs: Optional[dict] = None) -> None: """All subclasses should override this method, if they need any operations after each training iteration. Args: runner (Runner): The runner of the training process. batch_idx (int): The index of the current batch in the train loop. data_batch (dict tuple or list, optional): Data from dataloader. outputs (dict, optional): Outputs from model. """ self._after_iter( runner, batch_idx=batch_idx, data_batch=data_batch, outputs=outputs, mode='train')
[文档] def after_val_iter(self, runner, batch_idx: int, data_batch: DATA_BATCH = None, outputs: Optional[Sequence] = None) -> None: """All subclasses should override this method, if they need any operations after each validation iteration. Args: runner (Runner): The runner of the validation process. batch_idx (int): The index of the current batch in the val loop. data_batch (dict or tuple or list, optional): Data from dataloader. outputs (Sequence, optional): Outputs from model. """ self._after_iter( runner, batch_idx=batch_idx, data_batch=data_batch, outputs=outputs, mode='val')
[文档] def after_test_iter(self, runner, batch_idx: int, data_batch: DATA_BATCH = None, outputs: Optional[Sequence] = None) -> None: """All subclasses should override this method, if they need any operations after each test iteration. Args: runner (Runner): The runner of the training process. batch_idx (int): The index of the current batch in the test loop. data_batch (dict or tuple or list, optional): Data from dataloader. outputs (Sequence, optional): Outputs from model. """ self._after_iter( runner, batch_idx=batch_idx, data_batch=data_batch, outputs=outputs, mode='test')
def _before_epoch(self, runner, mode: str = 'train') -> None: """All subclasses should override this method, if they need any operations before each epoch. Args: runner (Runner): The runner of the training, validation or testing process. mode (str): Current mode of runner. Defaults to 'train'. """ def _after_epoch(self, runner, mode: str = 'train') -> None: """All subclasses should override this method, if they need any operations after each epoch. Args: runner (Runner): The runner of the training, validation or testing process. mode (str): Current mode of runner. Defaults to 'train'. """ def _before_iter(self, runner, batch_idx: int, data_batch: DATA_BATCH = None, mode: str = 'train') -> None: """All subclasses should override this method, if they need any operations before each iter. Args: runner (Runner): The runner of the training, validation or testing process. batch_idx (int): The index of the current batch in the loop. data_batch (dict or tuple or list, optional): Data from dataloader. mode (str): Current mode of runner. Defaults to 'train'. """ def _after_iter(self, runner, batch_idx: int, data_batch: DATA_BATCH = None, outputs: Optional[Union[Sequence, dict]] = None, mode: str = 'train') -> None: """All subclasses should override this method, if they need any operations after each epoch. Args: runner (Runner): The runner of the training, validation or testing process. batch_idx (int): The index of the current batch in the loop. data_batch (dict or tuple or list, optional): Data from dataloader. outputs (dict or Sequence, optional): Outputs from model. mode (str): Current mode of runner. Defaults to 'train'. """
[文档] def every_n_epochs(self, runner, n: int, start: int = 0) -> bool: """Test whether current epoch can be evenly divided by n. Args: runner (Runner): The runner of the training, validation or testing process. n (int): Whether current epoch can be evenly divided by n. start (int): Starting from `start` to check the logic for every n epochs. Defaults to 0. Returns: bool: Whether current epoch can be evenly divided by n. """ dividend = runner.epoch + 1 - start return dividend % n == 0 if dividend >= 0 and n > 0 else False
[文档] def every_n_inner_iters(self, batch_idx: int, n: int) -> bool: """Test whether current inner iteration can be evenly divided by n. Args: batch_idx (int): Current batch index of the training, validation or testing loop. n (int): Whether current inner iteration can be evenly divided by n. Returns: bool: Whether current inner iteration can be evenly divided by n. """ return (batch_idx + 1) % n == 0 if n > 0 else False
[文档] def every_n_train_iters(self, runner, n: int, start: int = 0) -> bool: """Test whether current training iteration can be evenly divided by n. Args: runner (Runner): The runner of the training, validation or testing process. n (int): Whether current iteration can be evenly divided by n. start (int): Starting from `start` to check the logic for every n iterations. Defaults to 0. Returns: bool: Return True if the current iteration can be evenly divided by n, otherwise False. """ dividend = runner.iter + 1 - start return dividend % n == 0 if dividend >= 0 and n > 0 else False
[文档] def end_of_epoch(self, dataloader, batch_idx: int) -> bool: """Check whether the current iteration reaches the last iteration of the dataloader. Args: dataloader (Dataloader): The dataloader of the training, validation or testing process. batch_idx (int): The index of the current batch in the loop. Returns: bool: Whether reaches the end of current epoch or not. """ return batch_idx + 1 == len(dataloader)
[文档] def is_last_train_epoch(self, runner) -> bool: """Test whether current epoch is the last train epoch. Args: runner (Runner): The runner of the training process. Returns: bool: Whether reaches the end of training epoch. """ return runner.epoch + 1 == runner.max_epochs
[文档] def is_last_train_iter(self, runner) -> bool: """Test whether current iteration is the last train iteration. Args: runner (Runner): The runner of the training process. Returns: bool: Whether current iteration is the last train iteration. """ return runner.iter + 1 == runner.max_iters
[文档] def get_triggered_stages(self) -> list: """Get all triggered stages with method name of the hook. Returns: list: List of triggered stages. """ trigger_stages = set() for stage in Hook.stages: if is_method_overridden(stage, Hook, self): trigger_stages.add(stage) # some methods will be triggered in multi stages # use this dict to map method to stages. method_stages_map = { '_before_epoch': ['before_train_epoch', 'before_val_epoch', 'before_test_epoch'], '_after_epoch': ['after_train_epoch', 'after_val_epoch', 'after_test_epoch'], '_before_iter': ['before_train_iter', 'before_val_iter', 'before_test_iter'], '_after_iter': ['after_train_iter', 'after_val_iter', 'after_test_iter'], } for method, map_stages in method_stages_map.items(): if is_method_overridden(method, Hook, self): trigger_stages.update(map_stages) return list(trigger_stages)