|
|
|
from typing import Optional, Sequence |
|
|
|
from mmengine.hooks import Hook |
|
from mmengine.runner import Runner |
|
|
|
from mmdet.registry import HOOKS |
|
from mmdet.structures import DetDataSample |
|
|
|
|
|
@HOOKS.register_module() |
|
class MemoryProfilerHook(Hook): |
|
"""Memory profiler hook recording memory information including virtual |
|
memory, swap memory, and the memory of the current process. |
|
|
|
Args: |
|
interval (int): Checking interval (every k iterations). |
|
Default: 50. |
|
""" |
|
|
|
def __init__(self, interval: int = 50) -> None: |
|
try: |
|
from psutil import swap_memory, virtual_memory |
|
self._swap_memory = swap_memory |
|
self._virtual_memory = virtual_memory |
|
except ImportError: |
|
raise ImportError('psutil is not installed, please install it by: ' |
|
'pip install psutil') |
|
|
|
try: |
|
from memory_profiler import memory_usage |
|
self._memory_usage = memory_usage |
|
except ImportError: |
|
raise ImportError( |
|
'memory_profiler is not installed, please install it by: ' |
|
'pip install memory_profiler') |
|
|
|
self.interval = interval |
|
|
|
def _record_memory_information(self, runner: Runner) -> None: |
|
"""Regularly record memory information. |
|
|
|
Args: |
|
runner (:obj:`Runner`): The runner of the training or evaluation |
|
process. |
|
""" |
|
|
|
virtual_memory = self._virtual_memory() |
|
swap_memory = self._swap_memory() |
|
|
|
process_memory = self._memory_usage()[0] |
|
factor = 1024 * 1024 |
|
runner.logger.info( |
|
'Memory information ' |
|
'available_memory: ' |
|
f'{round(virtual_memory.available / factor)} MB, ' |
|
'used_memory: ' |
|
f'{round(virtual_memory.used / factor)} MB, ' |
|
f'memory_utilization: {virtual_memory.percent} %, ' |
|
'available_swap_memory: ' |
|
f'{round((swap_memory.total - swap_memory.used) / factor)}' |
|
' MB, ' |
|
f'used_swap_memory: {round(swap_memory.used / factor)} MB, ' |
|
f'swap_memory_utilization: {swap_memory.percent} %, ' |
|
'current_process_memory: ' |
|
f'{round(process_memory)} MB') |
|
|
|
def after_train_iter(self, |
|
runner: Runner, |
|
batch_idx: int, |
|
data_batch: Optional[dict] = None, |
|
outputs: Optional[dict] = None) -> None: |
|
"""Regularly record memory information. |
|
|
|
Args: |
|
runner (:obj:`Runner`): The runner of the training process. |
|
batch_idx (int): The index of the current batch in the train loop. |
|
data_batch (dict, optional): Data from dataloader. |
|
Defaults to None. |
|
outputs (dict, optional): Outputs from model. Defaults to None. |
|
""" |
|
if self.every_n_inner_iters(batch_idx, self.interval): |
|
self._record_memory_information(runner) |
|
|
|
def after_val_iter( |
|
self, |
|
runner: Runner, |
|
batch_idx: int, |
|
data_batch: Optional[dict] = None, |
|
outputs: Optional[Sequence[DetDataSample]] = None) -> None: |
|
"""Regularly record memory information. |
|
|
|
Args: |
|
runner (:obj:`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. |
|
outputs (Sequence[:obj:`DetDataSample`], optional): |
|
Outputs from model. Defaults to None. |
|
""" |
|
if self.every_n_inner_iters(batch_idx, self.interval): |
|
self._record_memory_information(runner) |
|
|
|
def after_test_iter( |
|
self, |
|
runner: Runner, |
|
batch_idx: int, |
|
data_batch: Optional[dict] = None, |
|
outputs: Optional[Sequence[DetDataSample]] = None) -> None: |
|
"""Regularly record memory information. |
|
|
|
Args: |
|
runner (:obj:`Runner`): The runner of the testing process. |
|
batch_idx (int): The index of the current batch in the test loop. |
|
data_batch (dict, optional): Data from dataloader. |
|
Defaults to None. |
|
outputs (Sequence[:obj:`DetDataSample`], optional): |
|
Outputs from model. Defaults to None. |
|
""" |
|
if self.every_n_inner_iters(batch_idx, self.interval): |
|
self._record_memory_information(runner) |
|
|
