Source code for mmengine.runner.loops
# Copyright (c) OpenMMLab. All rights reserved.
import bisect
import logging
import time
from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
import torch
from torch.utils.data import BatchSampler, DataLoader, IterableDataset
from mmengine.evaluator import Evaluator
from mmengine.logging import HistoryBuffer, print_log
from mmengine.registry import LOOPS
from mmengine.structures import BaseDataElement
from mmengine.utils import is_list_of
from .amp import autocast
from .base_loop import BaseLoop
from .utils import calc_dynamic_intervals
[docs]
@LOOPS.register_module()
class EpochBasedTrainLoop(BaseLoop):
"""Loop for epoch-based training.
Args:
runner (Runner): A reference of runner.
dataloader (Dataloader or dict): A dataloader object or a dict to
build a dataloader.
max_epochs (int): Total training epochs.
val_begin (int): The epoch that begins validating.
Defaults to 1.
val_interval (int): Validation interval. Defaults to 1.
dynamic_intervals (List[Tuple[int, int]], optional): The
first element in the tuple is a milestone and the second
element is a interval. The interval is used after the
corresponding milestone. Defaults to None.
"""
def __init__(
self,
runner,
dataloader: Union[DataLoader, Dict],
max_epochs: int,
val_begin: int = 1,
val_interval: int = 1,
dynamic_intervals: Optional[List[Tuple[int, int]]] = None) -> None:
super().__init__(runner, dataloader)
self._max_epochs = int(max_epochs)
assert self._max_epochs == max_epochs, \
f'`max_epochs` should be a integer number, but get {max_epochs}.'
self._max_iters = self._max_epochs * len(self.dataloader)
self._epoch = 0
self._iter = 0
self.val_begin = val_begin
self.val_interval = val_interval
# This attribute will be updated by `EarlyStoppingHook`
# when it is enabled.
self.stop_training = False
if hasattr(self.dataloader.dataset, 'metainfo'):
self.runner.visualizer.dataset_meta = \
self.dataloader.dataset.metainfo
else:
print_log(
f'Dataset {self.dataloader.dataset.__class__.__name__} has no '
'metainfo. ``dataset_meta`` in visualizer will be '
'None.',
logger='current',
level=logging.WARNING)
self.dynamic_milestones, self.dynamic_intervals = \
calc_dynamic_intervals(
self.val_interval, dynamic_intervals)
@property
def max_epochs(self):
"""int: Total epochs to train model."""
return self._max_epochs
@property
def max_iters(self):
"""int: Total iterations to train model."""
return self._max_iters
@property
def epoch(self):
"""int: Current epoch."""
return self._epoch
@property
def iter(self):
"""int: Current iteration."""
return self._iter
[docs]
def run(self) -> torch.nn.Module:
"""Launch training."""
self.runner.call_hook('before_train')
while self._epoch < self._max_epochs and not self.stop_training:
self.run_epoch()
self._decide_current_val_interval()
if (self.runner.val_loop is not None
and self._epoch >= self.val_begin
and (self._epoch % self.val_interval == 0
or self._epoch == self._max_epochs)):
self.runner.val_loop.run()
self.runner.call_hook('after_train')
return self.runner.model
[docs]
def run_epoch(self) -> None:
"""Iterate one epoch."""
self.runner.call_hook('before_train_epoch')
self.runner.model.train()
for idx, data_batch in enumerate(self.dataloader):
self.run_iter(idx, data_batch)
self.runner.call_hook('after_train_epoch')
self._epoch += 1
[docs]
def run_iter(self, idx, data_batch: Sequence[dict]) -> None:
"""Iterate one min-batch.
Args:
data_batch (Sequence[dict]): Batch of data from dataloader.
"""
self.runner.call_hook(
'before_train_iter', batch_idx=idx, data_batch=data_batch)
# Enable gradient accumulation mode and avoid unnecessary gradient
# synchronization during gradient accumulation process.
# outputs should be a dict of loss.
outputs = self.runner.model.train_step(
data_batch, optim_wrapper=self.runner.optim_wrapper)
self.runner.call_hook(
'after_train_iter',
batch_idx=idx,
data_batch=data_batch,
outputs=outputs)
self._iter += 1
def _decide_current_val_interval(self) -> None:
"""Dynamically modify the ``val_interval``."""
step = bisect.bisect(self.dynamic_milestones, (self.epoch + 1))
self.val_interval = self.dynamic_intervals[step - 1]
class _InfiniteDataloaderIterator:
"""An infinite dataloader iterator wrapper for IterBasedTrainLoop.
It resets the dataloader to continue iterating when the iterator has
iterated over all the data. However, this approach is not efficient, as the
workers need to be restarted every time the dataloader is reset. It is
recommended to use `mmengine.dataset.InfiniteSampler` to enable the
dataloader to iterate infinitely.
"""
def __init__(self, dataloader: DataLoader) -> None:
self._dataloader = dataloader
# The iterator is created lazily so that, when resuming, the sampler
# can be fast-forwarded *before* any worker is spawned and prefetches
# data. Eagerly creating it here would make `num_workers > 0` load and
# discard the skipped data.
self._iterator: Any = None
self._epoch = 0
def __iter__(self):
return self
def __next__(self) -> Sequence[dict]:
return self._next_data()
def _ensure_iterator(self) -> None:
if self._iterator is None:
self._iterator = iter(self._dataloader)
def skip_iter(self, num_iters: int) -> None:
if num_iters <= 0:
return
sampler = self._resolve_sampler()
batch_size = self._resolve_batch_size()
if batch_size and sampler is not None:
# Fast path (multi-worker safe): advance the deterministic sampler
# stream by the number of already-consumed indices, then drop the
# iterator so it is rebuilt lazily on the next ``__next__`` and the
# workers prefetch from the resumed position without loading the
# skipped data. Relies on the ``skip(num_samples)`` contract
# documented on ``InfiniteSampler.skip``.
sampler.skip(num_iters * batch_size)
self._iterator = None
else:
print_log(
'Fast sampler-level resume is unavailable: no sampler with a '
'`skip` method was found directly or through a standard '
'`BatchSampler`, or the batch size could not be resolved. '
'Falling back to advancing the dataloader iterator one step at '
'a time, which is slower and, with `num_workers > 0`, still '
'loads and discards the skipped data.',
logger='current',
level=logging.WARNING)
for _ in range(num_iters):
self._next_data(skip_loading=True)
def _resolve_sampler(self) -> Optional[Any]:
sampler = getattr(self._dataloader, 'sampler', None)
if sampler is not None and hasattr(sampler, 'skip'):
return sampler
batch_sampler = getattr(self._dataloader, 'batch_sampler', None)
# Only unwrap PyTorch's fixed-size BatchSampler. Custom batch samplers
# may consume sampler indices differently from ``num_iters * batch_size``.
if (batch_sampler is not None
and batch_sampler.__class__ is BatchSampler):
sampler = getattr(batch_sampler, 'sampler', None)
if sampler is not None and hasattr(sampler, 'skip'):
return sampler
return None
def _resolve_batch_size(self) -> Optional[int]:
batch_size = getattr(self._dataloader, 'batch_size', None)
if batch_size is not None:
return batch_size
# ``batch_size`` is None when a custom ``batch_sampler`` is used; only a
# standard fixed-size batch sampler exposes a usable ``batch_size``.
batch_sampler = getattr(self._dataloader, 'batch_sampler', None)
return getattr(batch_sampler, 'batch_size', None)
def _next_data(self, skip_loading: bool = False) -> Any:
self._ensure_iterator()
try:
if skip_loading and self._can_skip_without_loading():
self._iterator._next_index()
return None
return next(self._iterator)
except StopIteration:
self._reset_iterator()
if skip_loading and self._can_skip_without_loading():
self._iterator._next_index()
return None
return next(self._iterator)
def _can_skip_without_loading(self) -> bool:
# ``_next_index`` is a private PyTorch API. It only advances the sampler
# of a single-process iterator without loading data. For multi-worker
# iterators it cannot be used safely (prefetch state would desync), and
# it may be absent on future/other iterator types, hence the guards.
return (getattr(self._dataloader, 'num_workers', 0) == 0
and not isinstance(self._dataloader.dataset, IterableDataset)
and hasattr(self._iterator, '_next_index'))
def _reset_iterator(self) -> None:
print_log(
'Reach the end of the dataloader, it will be '
'restarted and continue to iterate. It is '
'recommended to use '
'`mmengine.dataset.InfiniteSampler` to enable the '
'dataloader to iterate infinitely.',
logger='current',
level=logging.WARNING)
self._epoch += 1
if hasattr(self._dataloader, 'sampler') and hasattr(
self._dataloader.sampler, 'set_epoch'):
# In case the` _SingleProcessDataLoaderIter` has no sampler,
# or data loader uses `SequentialSampler` in Pytorch.
self._dataloader.sampler.set_epoch(self._epoch)
elif hasattr(self._dataloader, 'batch_sampler') and hasattr(
self._dataloader.batch_sampler.sampler, 'set_epoch'):
# In case the` _SingleProcessDataLoaderIter` has no batch
# sampler. batch sampler in pytorch warps the sampler as its
# attributes.
self._dataloader.batch_sampler.sampler.set_epoch(self._epoch)
time.sleep(2) # Prevent possible deadlock during epoch transition
self._iterator = iter(self._dataloader)
[docs]
@LOOPS.register_module()
class IterBasedTrainLoop(BaseLoop):
"""Loop for iter-based training.
Args:
runner (Runner): A reference of runner.
dataloader (Dataloader or dict): A dataloader object or a dict to
build a dataloader.
max_iters (int): Total training iterations.
val_begin (int): The iteration that begins validating.
Defaults to 1.
val_interval (int): Validation interval. Defaults to 1000.
dynamic_intervals (List[Tuple[int, int]], optional): The
first element in the tuple is a milestone and the second
element is a interval. The interval is used after the
corresponding milestone. Defaults to None.
fast_forward_on_resume (bool): Whether to skip advancing the
dataloader iterator when resuming from a checkpoint. When
`False` (default), the dataloader is advanced through the
already-trained steps to maintain training state consistency.
If the sampler supports a `skip` method (e.g.
:class:`~mmengine.dataset.InfiniteSampler`) and the batch size
can be resolved, this advance is done cheaply at the sampler
level without loading the skipped data; otherwise it falls back
to advancing the dataloader iterator step by step. When `True`,
this fast-forward is skipped to save time, but may affect
reproducibility if the dataloader has state-dependent behavior.
"""
def __init__(self,
runner,
dataloader: Union[DataLoader, Dict],
max_iters: int,
val_begin: int = 1,
val_interval: int = 1000,
dynamic_intervals: Optional[List[Tuple[int, int]]] = None,
fast_forward_on_resume: bool = False) -> None:
super().__init__(runner, dataloader)
self._max_iters = int(max_iters)
assert self._max_iters == max_iters, \
f'`max_iters` should be a integer number, but get {max_iters}'
self._max_epochs = 1 # for compatibility with EpochBasedTrainLoop
self._epoch = 0
self._iter = 0
self.val_begin = val_begin
self.val_interval = val_interval
self.fast_forward_on_resume = fast_forward_on_resume
# This attribute will be updated by `EarlyStoppingHook`
# when it is enabled.
self.stop_training = False
if hasattr(self.dataloader.dataset, 'metainfo'):
self.runner.visualizer.dataset_meta = \
self.dataloader.dataset.metainfo
else:
print_log(
f'Dataset {self.dataloader.dataset.__class__.__name__} has no '
'metainfo. ``dataset_meta`` in visualizer will be '
'None.',
logger='current',
level=logging.WARNING)
# get the iterator of the dataloader
self.dataloader_iterator = _InfiniteDataloaderIterator(self.dataloader)
self.dynamic_milestones, self.dynamic_intervals = \
calc_dynamic_intervals(
self.val_interval, dynamic_intervals)
@property
def max_epochs(self):
"""int: Total epochs to train model."""
return self._max_epochs
@property
def max_iters(self):
"""int: Total iterations to train model."""
return self._max_iters
@property
def epoch(self):
"""int: Current epoch."""
return self._epoch
@property
def iter(self):
"""int: Current iteration."""
return self._iter
[docs]
def run(self) -> None:
"""Launch training."""
self.runner.call_hook('before_train')
# In iteration-based training loop, we treat the whole training process
# as a big epoch and execute the corresponding hook.
self.runner.call_hook('before_train_epoch')
if self._iter > 0:
if not self.fast_forward_on_resume:
print_log(
f'Advance dataloader {self._iter} steps to skip data '
'that has already been trained',
logger='current',
level=logging.INFO)
self.dataloader_iterator.skip_iter(self._iter)
else:
print_log(
'Skip advancing dataloader to save time. Note that this '
'may affect reproducibility when resuming a training.',
logger='current',
level=logging.WARNING)
while self._iter < self._max_iters and not self.stop_training:
self.runner.model.train()
data_batch = next(self.dataloader_iterator)
self.run_iter(data_batch)
self._decide_current_val_interval()
if (self.runner.val_loop is not None
and self._iter >= self.val_begin
and (self._iter % self.val_interval == 0
or self._iter == self._max_iters)):
self.runner.val_loop.run()
self.runner.call_hook('after_train_epoch')
self.runner.call_hook('after_train')
return self.runner.model
[docs]
def run_iter(self, data_batch: Sequence[dict]) -> None:
"""Iterate one mini-batch.
Args:
data_batch (Sequence[dict]): Batch of data from dataloader.
"""
self.runner.call_hook(
'before_train_iter', batch_idx=self._iter, data_batch=data_batch)
# Enable gradient accumulation mode and avoid unnecessary gradient
# synchronization during gradient accumulation process.
# outputs should be a dict of loss.
outputs = self.runner.model.train_step(
data_batch, optim_wrapper=self.runner.optim_wrapper)
self.runner.call_hook(
'after_train_iter',
batch_idx=self._iter,
data_batch=data_batch,
outputs=outputs)
self._iter += 1
def _decide_current_val_interval(self) -> None:
"""Dynamically modify the ``val_interval``."""
step = bisect.bisect(self.dynamic_milestones, (self._iter + 1))
self.val_interval = self.dynamic_intervals[step - 1]
[docs]
@LOOPS.register_module()
class ValLoop(BaseLoop):
"""Loop for validation.
Args:
runner (Runner): A reference of runner.
dataloader (Dataloader or dict): A dataloader object or a dict to
build a dataloader.
evaluator (Evaluator or dict or list): Used for computing metrics.
fp16 (bool): Whether to enable fp16 validation. Defaults to
False.
"""
def __init__(self,
runner,
dataloader: Union[DataLoader, Dict],
evaluator: Union[Evaluator, Dict, List],
fp16: bool = False) -> None:
super().__init__(runner, dataloader)
if isinstance(evaluator, (dict, list)):
self.evaluator = runner.build_evaluator(evaluator) # type: ignore
else:
assert isinstance(evaluator, Evaluator), (
'evaluator must be one of dict, list or Evaluator instance, '
f'but got {type(evaluator)}.')
self.evaluator = evaluator # type: ignore
if hasattr(self.dataloader.dataset, 'metainfo'):
self.evaluator.dataset_meta = self.dataloader.dataset.metainfo
self.runner.visualizer.dataset_meta = \
self.dataloader.dataset.metainfo
else:
print_log(
f'Dataset {self.dataloader.dataset.__class__.__name__} has no '
'metainfo. ``dataset_meta`` in evaluator, metric and '
'visualizer will be None.',
logger='current',
level=logging.WARNING)
self.fp16 = fp16
self.val_loss: Dict[str, HistoryBuffer] = dict()
[docs]
def run(self) -> dict:
"""Launch validation."""
self.runner.call_hook('before_val')
self.runner.call_hook('before_val_epoch')
self.runner.model.eval()
# clear val loss
self.val_loss.clear()
for idx, data_batch in enumerate(self.dataloader):
self.run_iter(idx, data_batch)
# compute metrics
metrics = self.evaluator.evaluate(len(self.dataloader.dataset))
if self.val_loss:
loss_dict = _parse_losses(self.val_loss, 'val')
metrics.update(loss_dict)
self.runner.call_hook('after_val_epoch', metrics=metrics)
self.runner.call_hook('after_val')
return metrics
[docs]
@torch.no_grad()
def run_iter(self, idx, data_batch: Sequence[dict]):
"""Iterate one mini-batch.
Args:
data_batch (Sequence[dict]): Batch of data
from dataloader.
"""
self.runner.call_hook(
'before_val_iter', batch_idx=idx, data_batch=data_batch)
# outputs should be sequence of BaseDataElement
with autocast(enabled=self.fp16):
outputs = self.runner.model.val_step(data_batch)
outputs, self.val_loss = _update_losses(outputs, self.val_loss)
self.evaluator.process(data_samples=outputs, data_batch=data_batch)
self.runner.call_hook(
'after_val_iter',
batch_idx=idx,
data_batch=data_batch,
outputs=outputs)
[docs]
@LOOPS.register_module()
class TestLoop(BaseLoop):
"""Loop for test.
Args:
runner (Runner): A reference of runner.
dataloader (Dataloader or dict): A dataloader object or a dict to
build a dataloader.
evaluator (Evaluator or dict or list): Used for computing metrics.
fp16 (bool): Whether to enable fp16 testing. Defaults to
False.
"""
def __init__(self,
runner,
dataloader: Union[DataLoader, Dict],
evaluator: Union[Evaluator, Dict, List],
fp16: bool = False):
super().__init__(runner, dataloader)
if isinstance(evaluator, dict) or isinstance(evaluator, list):
self.evaluator = runner.build_evaluator(evaluator) # type: ignore
else:
self.evaluator = evaluator # type: ignore
if hasattr(self.dataloader.dataset, 'metainfo'):
self.evaluator.dataset_meta = self.dataloader.dataset.metainfo
self.runner.visualizer.dataset_meta = \
self.dataloader.dataset.metainfo
else:
print_log(
f'Dataset {self.dataloader.dataset.__class__.__name__} has no '
'metainfo. ``dataset_meta`` in evaluator, metric and '
'visualizer will be None.',
logger='current',
level=logging.WARNING)
self.fp16 = fp16
self.test_loss: Dict[str, HistoryBuffer] = dict()
[docs]
def run(self) -> dict:
"""Launch test."""
self.runner.call_hook('before_test')
self.runner.call_hook('before_test_epoch')
self.runner.model.eval()
# clear test loss
self.test_loss.clear()
for idx, data_batch in enumerate(self.dataloader):
self.run_iter(idx, data_batch)
# compute metrics
metrics = self.evaluator.evaluate(len(self.dataloader.dataset))
if self.test_loss:
loss_dict = _parse_losses(self.test_loss, 'test')
metrics.update(loss_dict)
self.runner.call_hook('after_test_epoch', metrics=metrics)
self.runner.call_hook('after_test')
return metrics
[docs]
@torch.no_grad()
def run_iter(self, idx, data_batch: Sequence[dict]) -> None:
"""Iterate one mini-batch.
Args:
data_batch (Sequence[dict]): Batch of data from dataloader.
"""
self.runner.call_hook(
'before_test_iter', batch_idx=idx, data_batch=data_batch)
# predictions should be sequence of BaseDataElement
with autocast(enabled=self.fp16):
outputs = self.runner.model.test_step(data_batch)
outputs, self.test_loss = _update_losses(outputs, self.test_loss)
self.evaluator.process(data_samples=outputs, data_batch=data_batch)
self.runner.call_hook(
'after_test_iter',
batch_idx=idx,
data_batch=data_batch,
outputs=outputs)
def _parse_losses(losses: Dict[str, HistoryBuffer],
stage: str) -> Dict[str, float]:
"""Parses the raw losses of the network.
Args:
losses (dict): raw losses of the network.
stage (str): The stage of loss, e.g., 'val' or 'test'.
Returns:
dict[str, float]: The key is the loss name, and the value is the
average loss.
"""
all_loss = 0
loss_dict: Dict[str, float] = dict()
for loss_name, loss_value in losses.items():
avg_loss = loss_value.mean()
loss_dict[loss_name] = avg_loss
if 'loss' in loss_name:
all_loss += avg_loss
loss_dict[f'{stage}_loss'] = all_loss
return loss_dict
def _update_losses(outputs: list, losses: dict) -> Tuple[list, dict]:
"""Update and record the losses of the network.
Args:
outputs (list): The outputs of the network.
losses (dict): The losses of the network.
Returns:
list: The updated outputs of the network.
dict: The updated losses of the network.
"""
if isinstance(outputs[-1],
BaseDataElement) and outputs[-1].keys() == ['loss']:
loss = outputs[-1].loss # type: ignore
outputs = outputs[:-1]
else:
loss = dict()
for loss_name, loss_value in loss.items():
if loss_name not in losses:
losses[loss_name] = HistoryBuffer()
if isinstance(loss_value, torch.Tensor):
losses[loss_name].update(loss_value.item())
elif is_list_of(loss_value, torch.Tensor):
for loss_value_i in loss_value:
losses[loss_name].update(loss_value_i.item())
return outputs, losses