def __init__(self, min_steps: Optional[int] = 0, max_steps: int = -1) -> None:
super().__init__()
if max_steps is None:
rank_zero_deprecation(
"Setting `max_steps = None` is deprecated in v1.5 and will no longer be supported in v1.7."
" Use `max_steps = -1` instead."
)
max_steps = -1
elif max_steps < -1:
raise MisconfigurationException(
f"`max_steps` must be a non-negative integer or -1 (infinite steps). You passed in {max_steps}."
)
self.min_steps = min_steps
self.max_steps = max_steps
self.global_step: int = 0
self.batch_progress = BatchProgress()
self.scheduler_progress = SchedulerProgress()
self.batch_loop: Optional[TrainingBatchLoop] = None
self.val_loop: Optional["loops.EvaluationLoop"] = None
self._results = ResultCollection(training=True)
self._outputs: _OUTPUTS_TYPE = []
self._warning_cache = WarningCache()
self._dataloader_iter: Optional[Iterator] = None
# caches the loaded dataloader state until dataloader objects are available
self._dataloader_state_dict: Dict[str, Any] = {}
def __init__(self) -> None:
super().__init__()
self.batch_progress = BatchProgress()
self._outputs: EPOCH_OUTPUT = []
self._dl_max_batches = 0
self._data_fetcher: Optional[AbstractDataFetcher] = None
self._dataloader_state_dict: Dict[str, Any] = {}
def __init__(self) -> None:
super().__init__()
self.outputs: EPOCH_OUTPUT = []
self.batch_progress = BatchProgress()
self._dl_max_batches: Optional[int] = None
self._num_dataloaders: Optional[int] = None
self._dataloader_iter: Optional[Iterator] = None
self._data_fetcher: Optional[DataFetcher] = None
self._dataloader_state_dict: Dict[str, Any] = None
def __init__(self,
min_steps: Optional[int] = None,
max_steps: int = -1) -> None:
super().__init__()
if max_steps < -1:
raise MisconfigurationException(
f"`max_steps` must be a non-negative integer or -1 (infinite steps). You passed in {max_steps}."
)
self.min_steps = min_steps
self.max_steps = max_steps
self.batch_progress = BatchProgress()
self.scheduler_progress = SchedulerProgress()
self.batch_loop = TrainingBatchLoop()
self.val_loop = loops.EvaluationLoop(verbose=False)
self._results = _ResultCollection(training=True)
self._outputs: _OUTPUTS_TYPE = []
self._warning_cache = WarningCache()
# caches the loaded dataloader state until dataloader objects are available
self._dataloader_state_dict: Dict[str, Any] = {}
self._batches_that_stepped: int = 0
def __init__(self) -> None:
super().__init__()
self.accumulated_loss: Optional[Tensor] = None
self.batch_outputs: Optional[List[List[STEP_OUTPUT]]] = None
self.running_loss: TensorRunningAccum = TensorRunningAccum(window_length=20)
self.batch_idx: int = 0
self.split_idx: Optional[int] = None
self.progress = BatchProgress()
self.optim_progress = OptimizationProgress()
self._warning_cache: WarningCache = WarningCache()
self._hiddens: Optional[Tensor] = None
self._optimizer_freq_cumsum: Optional[int] = None
self._remaining_splits: Optional[List[Any]] = None
self._skip_backward: bool = False
class EvaluationEpochLoop(Loop):
"""This is the loop performing the evaluation.
It mainly loops over the given dataloader and runs the validation or test step (depending on the trainer's current
state).
"""
def __init__(self) -> None:
super().__init__()
self.batch_progress = BatchProgress()
self._outputs: EPOCH_OUTPUT = []
self._dl_max_batches = 0
self._dataloader_iter: Optional[Iterator] = None
self._data_fetcher: Optional[AbstractDataFetcher] = None
self._dataloader_state_dict: Dict[str, Any] = {}
@property
def done(self) -> bool:
"""Returns ``True`` if the current iteration count reaches the number of dataloader batches."""
return self.batch_progress.current.completed >= self._dl_max_batches
def reset(self) -> None:
"""Resets the loop's internal state."""
self._dl_max_batches = 0
self._data_fetcher = None
self._outputs = []
if not self.restarting:
self.batch_progress.reset_on_run()
else:
self.batch_progress.reset_on_restart()
# when restarting, if we are running `validate` or `test` twice, since there's no concept of `max_epochs` we
# need to reset the current state when the loop has finished running
if self.done and self.trainer.state.fn != TrainerFn.FITTING:
self.batch_progress.reset_on_run()
def on_run_start( # type: ignore[override]
self, data_fetcher: AbstractDataFetcher,
dataloader_idx: Optional[int], dl_max_batches: int) -> None:
"""Adds the passed arguments to the loop's state if necessary.
Args:
data_fetcher: the current data_fetcher wrapping the dataloader
dataloader_idx: index of the current dataloader
dl_max_batches: maximum number of batches the dataloader can produce
"""
void(dataloader_idx)
self._dl_max_batches = dl_max_batches
self._data_fetcher = data_fetcher
self._reload_dataloader_state_dict(data_fetcher)
self._dataloader_iter = iter(data_fetcher)
def advance( # type: ignore[override]
self, data_fetcher: AbstractDataFetcher,
dataloader_idx: Optional[int], dl_max_batches: int) -> None:
"""Calls the evaluation step with the corresponding hooks and updates the logger connector.
Args:
data_fetcher: iterator over the dataloader
dataloader_idx: index of the current dataloader
dl_max_batches: maximum number of batches the dataloader can produce
Raises:
StopIteration: If the current batch is None
"""
void(dl_max_batches)
assert self._dataloader_iter is not None
batch, self.batch_progress.is_last_batch = next(self._dataloader_iter)
if batch is None:
raise StopIteration
# configure step_kwargs
# TODO: each loop should construct its own kwargs, so we avoid the dataloader_idx reference here
kwargs = self._build_kwargs(batch, self.batch_progress.current.ready,
dataloader_idx)
self.batch_progress.increment_ready()
# hook
self._on_evaluation_batch_start(**kwargs)
self.batch_progress.increment_started()
# lightning module methods
output = self._evaluation_step(**kwargs)
output = self._evaluation_step_end(output)
self.batch_progress.increment_processed()
# track loss history
self._on_evaluation_batch_end(output, **kwargs)
self.batch_progress.increment_completed()
# log batch metrics
self.trainer.logger_connector.update_eval_step_metrics()
# track epoch level outputs
if self._should_track_batch_outputs_for_epoch_end(
) and output is not None:
self._outputs.append(output)
if self.trainer.move_metrics_to_cpu:
# the evaluation step output is not moved as they are not considered "metrics"
assert self.trainer._results is not None
self.trainer._results.cpu()
if not self.batch_progress.is_last_batch:
# if fault tolerant is enabled and process has been notified, exit.
self.trainer._exit_gracefully_on_signal()
def on_run_end(self) -> EPOCH_OUTPUT:
"""Returns the outputs of the whole run."""
outputs, self._outputs = self._outputs, [] # free memory
self._dataloader_iter = None
self._data_fetcher = None
return outputs
def teardown(self) -> None:
# in case the model changes
self._should_track_batch_outputs_for_epoch_end.cache_clear()
def on_save_checkpoint(self) -> Dict:
state_dict = super().on_save_checkpoint()
if (self._data_fetcher is None
or self._num_completed_batches_reached() # did not finish
# TODO: fault-tolerance requires a minimum number of batches so probably should be > 0
or self.batch_progress.current.ready == 0 # did not start
):
return state_dict
# TODO: this should use `pytorch_lightning/trainer/supporters.py::CombinedLoader._state_dict_fn`
state_to_save = "state" if self._has_completed() else "previous_state"
state: Optional[MergedIteratorState] = getattr(
self._data_fetcher.dataloader_iter, state_to_save, None)
if state:
state_dict[
"dataloader_state_dict"] = _collect_states_on_rank_zero_over_collection(
asdict(state))
return state_dict
def on_load_checkpoint(self, state_dict: Dict) -> None:
# cache the dataloader state dict until the dataloader objects are available
# dataset states are collected across all ranks
dataloader_state_dict = state_dict.get("dataloader_state_dict", None)
if not _fault_tolerant_training() or not dataloader_state_dict:
return
self._dataloader_state_dict = dataloader_state_dict[
self.trainer.global_rank]
def _reload_dataloader_state_dict(
self, data_fetcher: AbstractDataFetcher) -> None:
if self.trainer.sanity_checking or not self._dataloader_state_dict:
return
dataloader = data_fetcher.dataloader
if isinstance(dataloader, CombinedLoader):
raise MisconfigurationException(
"Reloading support hasn't been implemented for `CombinedLoader`. You can request it by opening an issue"
" in `https://github.com/PyTorchLightning/pytorch-lightning/issues`."
)
assert dataloader is not None
_reload_dataloader_state_dict(dataloader, self._dataloader_state_dict)
self._dataloader_state_dict = {}
def _num_completed_batches_reached(self) -> bool:
epoch_finished_on_completed = self.batch_progress.current.completed == self._dl_max_batches
dataloader_consumed_successfully = self.batch_progress.is_last_batch and self._has_completed(
)
return epoch_finished_on_completed or dataloader_consumed_successfully
def _has_completed(self) -> bool:
return self.batch_progress.current.ready == self.batch_progress.current.completed
def _evaluation_step(self, **kwargs: Any) -> Optional[STEP_OUTPUT]:
"""The evaluation step (validation_step or test_step depending on the trainer's state).
Args:
batch: The current batch to run through the step.
batch_idx: The index of the current batch
dataloader_idx: the index of the dataloader producing the current batch
Returns:
the outputs of the step
"""
if self.trainer.testing:
output = self.trainer._call_strategy_hook("test_step",
*kwargs.values())
else:
output = self.trainer._call_strategy_hook("validation_step",
*kwargs.values())
return output
def _evaluation_step_end(self, *args: Any,
**kwargs: Any) -> Optional[STEP_OUTPUT]:
"""Calls the `{validation/test}_step_end` hook."""
hook_name = "test_step_end" if self.trainer.testing else "validation_step_end"
model_output = self.trainer._call_lightning_module_hook(
hook_name, *args, **kwargs)
strategy_output = self.trainer._call_strategy_hook(
hook_name, *args, **kwargs)
output = strategy_output if model_output is None else model_output
return output
def _on_evaluation_batch_start(self, **kwargs: Any) -> None:
"""Calls the ``on_{validation/test}_batch_start`` hook.
Args:
batch: The current batch to run through the step
batch_idx: The index of the current batch
dataloader_idx: The index of the dataloader producing the current batch
Raises:
AssertionError: If the number of dataloaders is None (has not yet been set).
"""
self.trainer.logger_connector.on_batch_start(**kwargs)
kwargs.setdefault("dataloader_idx",
0) # TODO: the argument should be keyword for these
if self.trainer.testing:
self.trainer._call_callback_hooks("on_test_batch_start",
*kwargs.values())
self.trainer._call_lightning_module_hook("on_test_batch_start",
*kwargs.values())
else:
self.trainer._call_callback_hooks("on_validation_batch_start",
*kwargs.values())
self.trainer._call_lightning_module_hook(
"on_validation_batch_start", *kwargs.values())
def _on_evaluation_batch_end(self, output: Optional[STEP_OUTPUT],
**kwargs: Any) -> None:
"""The ``on_{validation/test}_batch_end`` hook.
Args:
output: The output of the performed step
batch: The input batch for the step
batch_idx: The index of the current batch
dataloader_idx: Index of the dataloader producing the current batch
"""
kwargs.setdefault("dataloader_idx",
0) # TODO: the argument should be keyword for these
hook_name = "on_test_batch_end" if self.trainer.testing else "on_validation_batch_end"
self.trainer._call_callback_hooks(hook_name, output, *kwargs.values())
self.trainer._call_lightning_module_hook(hook_name, output,
*kwargs.values())
self.trainer.logger_connector.on_batch_end()
def _build_kwargs(
self, batch: Any, batch_idx: int,
dataloader_idx: Optional[int]) -> Dict[str, Union[Any, int]]:
"""Helper function to build the arguments for the current step.
Args:
batch: The current batch to run through the step
batch_idx: the index of the current batch
dataloader_idx: the index of the dataloader producing the current batch
Returns:
the keyword arguments to pass to the step function
"""
# make dataloader_idx arg in validation_step optional
step_kwargs = OrderedDict([("batch", batch), ("batch_idx", batch_idx)])
if dataloader_idx is not None:
step_kwargs["dataloader_idx"] = dataloader_idx
return step_kwargs
@lru_cache(1)
def _should_track_batch_outputs_for_epoch_end(self) -> bool:
"""Whether the batch outputs should be stored for later usage."""
model = self.trainer.lightning_module
if self.trainer.testing:
return is_overridden("test_epoch_end", model)
return is_overridden("validation_epoch_end", model)
class TrainingEpochLoop(loops.Loop[_OUTPUTS_TYPE]):
"""Runs over all batches in a dataloader (one epoch).
Args:
min_steps: The minimum number of steps (batches) to process
max_steps: The maximum number of steps (batches) to process
"""
def __init__(self, min_steps: Optional[int] = 0, max_steps: int = -1) -> None:
super().__init__()
if max_steps is None:
rank_zero_deprecation(
"Setting `max_steps = None` is deprecated in v1.5 and will no longer be supported in v1.7."
" Use `max_steps = -1` instead."
)
max_steps = -1
elif max_steps < -1:
raise MisconfigurationException(
f"`max_steps` must be a non-negative integer or -1 (infinite steps). You passed in {max_steps}."
)
self.min_steps = min_steps
self.max_steps = max_steps
self.global_step: int = 0
self.batch_progress = BatchProgress()
self.scheduler_progress = SchedulerProgress()
self.batch_loop: Optional[TrainingBatchLoop] = None
self.val_loop: Optional["loops.EvaluationLoop"] = None
self._results = ResultCollection(training=True)
self._outputs: _OUTPUTS_TYPE = []
self._warning_cache = WarningCache()
self._dataloader_iter: Optional[Iterator] = None
# caches the loaded dataloader state until dataloader objects are available
self._dataloader_state_dict: Dict[str, Any] = {}
@property
def total_batch_idx(self) -> int:
"""Returns the current batch index (across epochs)"""
# use `ready` instead of `completed` in case this is accessed after `completed` has been increased
# but before the next `ready` increase
return self.batch_progress.total.ready - 1
@property
def batch_idx(self) -> int:
"""Returns the current batch index (within this epoch)"""
# use `ready` instead of `completed` in case this is accessed after `completed` has been increased
# but before the next `ready` increase
return self.batch_progress.current.ready - 1
@property
def _is_training_done(self) -> bool:
max_steps_reached = _is_max_limit_reached(self.global_step, self.max_steps)
return max_steps_reached or self._num_ready_batches_reached()
@property
def _is_validation_done(self) -> bool:
# when we are restarting we want to check whether the val loop has finished
return not self.restarting or self.val_loop.done
@property
def done(self) -> bool:
"""Returns whether the training should be stopped.
The criteria are that the number of steps reached the max steps, the last batch is reached or the trainer
signals to stop (e.g. by early stopping).
"""
return (self._is_training_done and self._is_validation_done) or self.trainer.should_stop
def connect(
self,
batch_loop: TrainingBatchLoop = None,
val_loop: Optional["loops.EvaluationLoop"] = None,
) -> None:
"""Optionally connect a custom batch or validation loop to this training epoch loop."""
if batch_loop is not None:
self.batch_loop = batch_loop
if val_loop is not None:
self.val_loop = val_loop
def reset(self) -> None:
"""Resets the internal state of the loop for a new run."""
assert self.batch_loop is not None
assert self.batch_loop.optimizer_loop is not None
if self.restarting:
self.batch_progress.reset_on_restart()
self.scheduler_progress.reset_on_restart()
self.batch_loop.optimizer_loop.optim_progress.reset_on_restart()
else:
self.batch_progress.reset_on_run()
self.scheduler_progress.reset_on_run()
self.batch_loop.optimizer_loop.optim_progress.reset_on_run()
self._outputs = []
def on_run_start(self, data_fetcher: AbstractDataFetcher, **kwargs: Any) -> None:
# hook
self.trainer.logger_connector.on_epoch_start()
self.trainer.call_hook("on_epoch_start")
self.trainer.call_hook("on_train_epoch_start")
self.trainer.fit_loop.epoch_progress.increment_started()
self._reload_dataloader_state_dict(data_fetcher)
self._dataloader_iter = _update_dataloader_iter(data_fetcher, self.batch_idx + 1)
def advance(self, *args: Any, **kwargs: Any) -> None:
"""Runs a single training batch.
Args:
dataloader_iter: the iterator over the dataloader producing the new batch
Raises:
StopIteration: When the epoch is canceled by the user returning -1
"""
if self.restarting and self._should_check_val_fx(self.batch_idx, self.batch_progress.is_last_batch):
# skip training and run validation in `on_advance_end`
return
batch_idx, (batch, self.batch_progress.is_last_batch) = next(self._dataloader_iter)
if not self.trainer._data_connector.train_data_fetcher.store_on_device:
with self.trainer.profiler.profile("training_batch_to_device"):
batch = self.trainer.accelerator.batch_to_device(batch)
self.batch_progress.increment_ready()
# cache the batch size value to avoid extracting it again after the batch loop runs as the value will be
# different if tbptt is enabled
batch_size = self.trainer.logger_connector.on_batch_start(batch_idx, batch)
if batch is None:
self._warning_cache.warn("train_dataloader yielded None. If this was on purpose, ignore this warning...")
batch_output = []
else:
# hook
response = self.trainer.call_hook("on_batch_start")
if response == -1:
self.batch_progress.increment_processed()
raise StopIteration
# TODO: Update this in v1.7 (deprecation: #9816)
model_fx = self.trainer.lightning_module.on_train_batch_start
extra_kwargs = (
{"dataloader_idx": 0}
if callable(model_fx) and is_param_in_hook_signature(model_fx, "dataloader_idx", explicit=True)
else {}
)
# hook
response = self.trainer.call_hook("on_train_batch_start", batch, batch_idx, **extra_kwargs)
if response == -1:
self.batch_progress.increment_processed()
raise StopIteration
self.batch_progress.increment_started()
with self.trainer.profiler.profile("run_training_batch"):
batch_output = self.batch_loop.run(batch, batch_idx)
self.trainer._results.batch_size = batch_size
self.batch_progress.increment_processed()
# update non-plateau LR schedulers
# update epoch-interval ones only when we are at the end of training epoch
self.update_lr_schedulers("step", update_plateau_schedulers=False)
if self._num_ready_batches_reached():
self.update_lr_schedulers("epoch", update_plateau_schedulers=False)
batch_end_outputs = self._prepare_outputs_training_batch_end(
batch_output,
automatic=self.trainer.lightning_module.trainer.lightning_module.automatic_optimization,
num_optimizers=len(self.trainer.optimizers),
)
# TODO: Update this in v1.7 (deprecation: #9816)
model_fx = self.trainer.lightning_module.on_train_batch_end
extra_kwargs = (
{"dataloader_idx": 0}
if callable(model_fx) and is_param_in_hook_signature(model_fx, "dataloader_idx", explicit=True)
else {}
)
self.trainer.call_hook("on_train_batch_end", batch_end_outputs, batch, batch_idx, **extra_kwargs)
self.trainer.call_hook("on_batch_end")
self.trainer.logger_connector.on_batch_end()
self.batch_progress.increment_completed()
if is_overridden("training_epoch_end", self.trainer.lightning_module):
self._outputs.append(batch_output)
# -----------------------------------------
# SAVE METRICS TO LOGGERS AND PROGRESS_BAR
# -----------------------------------------
self.trainer.logger_connector.update_train_step_metrics()
def on_advance_end(self):
"""Runs validation and Checkpointing if necessary.
Raises:
StopIteration: if :attr:`done` evaluates to ``True`` to finish this epoch
"""
# -----------------------------------------
# VALIDATE IF NEEDED + CHECKPOINT CALLBACK
# -----------------------------------------
should_check_val = self._should_check_val_fx(self.batch_idx, self.batch_progress.is_last_batch)
if should_check_val:
self.trainer.validating = True
self._run_validation()
self.trainer.training = True
# -----------------------------------------
# SAVE LOGGERS (ie: Tensorboard, etc...)
# -----------------------------------------
self._save_loggers_on_train_batch_end()
# update plateau LR scheduler after metrics are logged
self.update_lr_schedulers("step", update_plateau_schedulers=True)
if not self._should_accumulate():
# progress global step according to grads progress
self.global_step += 1
# if training finished, try to exit in `on_run_end` instead as we should have enough time
# TODO: @tchaton verify this assumption is True.
if not self._is_training_done:
# if fault tolerant is enabled and process has been notified, exit.
self.trainer._exit_gracefully_on_signal()
def on_run_end(self) -> None:
"""Calls the on_epoch_end hook.
Returns:
The output of each training step for each optimizer
Raises:
MisconfigurationException: ``train_epoch_end`` does not return ``None``
"""
# inform logger the batch loop has finished
self.trainer.logger_connector.epoch_end_reached()
# get the model and call model.training_epoch_end
model = self.trainer.lightning_module
if is_overridden("training_epoch_end", model) and self._outputs:
epoch_end_outputs = self._prepare_outputs_training_epoch_end(
self._outputs,
automatic=model.automatic_optimization,
num_optimizers=len(self.trainer.optimizers),
)
# run lightning module hook training_epoch_end
# refresh the result for custom logging at the epoch level
model._current_fx_name = "training_epoch_end"
epoch_end_outputs = model.training_epoch_end(epoch_end_outputs)
if epoch_end_outputs is not None:
raise MisconfigurationException(
"`training_epoch_end` expects a return of None. "
"HINT: remove the return statement in `training_epoch_end`."
)
# free memory
self._outputs = []
self.trainer.fit_loop.epoch_progress.increment_processed()
# call train epoch end hooks
self.trainer.call_hook("on_train_epoch_end")
self.trainer.call_hook("on_epoch_end")
self.trainer.logger_connector.on_epoch_end()
if self._num_ready_batches_reached():
self.update_lr_schedulers("epoch", update_plateau_schedulers=True)
# if fault tolerant is enabled and process has been notified, exit.
self.trainer._exit_gracefully_on_signal()
def teardown(self) -> None:
self._results.cpu()
self.batch_loop.teardown()
self.val_loop.teardown()
def on_save_checkpoint(self) -> Dict:
state_dict = super().on_save_checkpoint()
if (
self.trainer.train_dataloader is None
or self._num_completed_batches_reached() # did not finish
# TODO: fault-tolerance requires a minimum number of batches so probably should be > 0
or self.batch_progress.current.ready == 0 # did not start
):
return state_dict
state_dict["dataloader_state_dict"] = self.trainer.train_dataloader.state_dict(
has_completed=self._has_completed()
)
return state_dict
def on_load_checkpoint(self, state_dict: Dict) -> None:
# cache the dataloader state dict until the dataloader objects are available
self._dataloader_state_dict = state_dict.get("dataloader_state_dict")
def _run_validation(self):
# reload dataloaders
self.val_loop._reload_evaluation_dataloaders()
with torch.no_grad():
self.val_loop.run()
def _accumulated_batches_reached(self) -> bool:
"""Determine if accumulation will be finished by the end of the current batch."""
return self.batch_progress.current.ready % self.trainer.accumulate_grad_batches == 0
def _num_ready_batches_reached(self) -> bool:
"""Checks if we are in the last batch or if there are more batches to follow."""
epoch_finished_on_ready = self.batch_progress.current.ready == self.trainer.num_training_batches
return epoch_finished_on_ready or self.batch_progress.is_last_batch
def _num_completed_batches_reached(self) -> bool:
epoch_finished_on_completed = self.batch_progress.current.completed == self.trainer.num_training_batches
dataloader_consumed_successfully = self.batch_progress.is_last_batch and self._has_completed()
return epoch_finished_on_completed or dataloader_consumed_successfully
def _has_completed(self) -> bool:
return self.batch_progress.current.ready == self.batch_progress.current.completed
def _should_accumulate(self) -> bool:
"""Checks if the optimizer step should be performed or gradients should be accumulated for the current
step."""
accumulation_done = self._accumulated_batches_reached()
# Lightning steps on the final batch
is_final_batch = self._num_ready_batches_reached()
# but the TTP might not
ttp_accumulates_on_final_batch = (
self.trainer.training_type_plugin.handles_gradient_accumulation or not is_final_batch
)
return not accumulation_done and ttp_accumulates_on_final_batch
@staticmethod
def _prepare_outputs_training_batch_end(
batch_output: _BATCH_OUTPUTS_TYPE,
automatic: bool,
num_optimizers: int,
) -> Union[List[List[Dict[str, Any]]], List[Dict[str, Any]]]:
"""Processes the outputs from the batch loop into the format passed to the ``training_batch_end`` hook.
``(tbptt_steps, n_opt) -> (n_opt, tbptt_steps)``. The optimizer dimension might have been squeezed.
"""
if not batch_output:
return []
# convert optimizer dicts to list
if automatic:
batch_output = apply_to_collection(
batch_output, dtype=dict, function=_convert_optim_dict, num_optimizers=num_optimizers
)
array = np.array(batch_output, dtype=object)
if array.ndim == 1:
array = np.expand_dims(array, 1)
array = array.transpose((1, 0))
array = array.squeeze()
array = array.tolist()
array = _recursive_unpad(array)
return array
@staticmethod
def _prepare_outputs_training_epoch_end(
batch_outputs: _OUTPUTS_TYPE,
automatic: bool,
num_optimizers: int,
) -> Union[List[List[List[Dict[str, Any]]]], List[List[Dict[str, Any]]], List[Dict[str, Any]]]:
"""Processes the outputs from the batch loop into the format passed to the ``training_epoch_end`` hook.
``(n_batches, tbptt_steps, n_opt) -> (n_opt, n_batches, tbptt_steps)``.
All single-element dimensions might have been squeezed.
This processing is necessary because the format of the inputs to the ``training_epoch_end`` hook does not
match the loop structure and because empty dimensions are squeezed. This could break with loop customization.
"""
# `batch_outputs` (plural) is the same as `epoch_end_output` (singular)
if not batch_outputs:
return []
# convert optimizer dicts to list
if automatic:
batch_outputs = apply_to_collection(
batch_outputs, dtype=dict, function=_convert_optim_dict, num_optimizers=num_optimizers
)
array = _recursive_pad(batch_outputs)
if array.ndim == 2:
array = np.expand_dims(array, 2)
array = array.transpose((2, 0, 1))
array = array.squeeze()
array = array.tolist()
array = _recursive_unpad(array)
# in case we squeezed from 1-element array to a 0-dim array
array = array if isinstance(array, list) else [array]
# remove residual empty lists
array = [item for item in array if not isinstance(item, list) or len(item)]
return array
def update_lr_schedulers(self, interval: str, update_plateau_schedulers: bool) -> None:
"""updates the lr schedulers based on the given interval."""
if interval == "step" and self._should_accumulate():
return
active_optimizers = _get_active_optimizers(
self.trainer.optimizers, self.trainer.optimizer_frequencies, self.total_batch_idx
)
self._update_learning_rates(
interval=interval,
update_plateau_schedulers=update_plateau_schedulers,
opt_indices=[opt_idx for opt_idx, _ in active_optimizers],
)
def _update_learning_rates(
self, interval: str, update_plateau_schedulers: bool, opt_indices: Optional[List[int]] = None
) -> None:
"""Update learning rates.
Args:
interval: either 'epoch' or 'step'.
update_plateau_schedulers: control whether ``ReduceLROnPlateau`` or non-plateau schedulers get updated.
This is used so non-plateau schedulers can be updated before running validation. Checkpoints are
commonly saved during validation, however, on-plateau schedulers might monitor a validation metric
so they have to be updated separately.
opt_indices: indices of the optimizers to update.
"""
if not self.trainer.lr_schedulers or not self.trainer.lightning_module.automatic_optimization:
return
if opt_indices is None:
opt_indices = []
for lr_scheduler in self.trainer.lr_schedulers:
if isinstance(lr_scheduler["opt_idx"], int) and lr_scheduler["opt_idx"] not in opt_indices:
continue
if update_plateau_schedulers ^ lr_scheduler["reduce_on_plateau"]:
continue
current_idx = self.batch_idx if interval == "step" else self.trainer.current_epoch
current_idx += 1 # account for both batch and epoch starts from 0
# Take step if call to update_learning_rates matches the interval key and
# the current step modulo the schedulers frequency is zero
if lr_scheduler["interval"] == interval and current_idx % lr_scheduler["frequency"] == 0:
monitor_val = None
if lr_scheduler["reduce_on_plateau"]:
# If instance of ReduceLROnPlateau, we need a monitor
monitor_key = lr_scheduler["monitor"]
monitor_val = self._get_monitor_value(monitor_key)
if monitor_val is None:
if lr_scheduler.get("strict", True):
avail_metrics = list(self.trainer.callback_metrics)
raise MisconfigurationException(
f"ReduceLROnPlateau conditioned on metric {monitor_key}"
f" which is not available. Available metrics are: {avail_metrics}."
" Condition can be set using `monitor` key in lr scheduler dict"
)
rank_zero_warn(
f"ReduceLROnPlateau conditioned on metric {monitor_key}"
" which is not available but strict is set to `False`."
" Skipping learning rate update.",
RuntimeWarning,
)
continue
self.scheduler_progress.increment_ready()
# update LR
if lr_scheduler["reduce_on_plateau"]:
lr_scheduler["scheduler"].step(monitor_val)
else:
lr_scheduler["scheduler"].step()
self.scheduler_progress.increment_completed()
def _get_monitor_value(self, key: str) -> Any:
# this is a separate method to aid in testing
return self.trainer.callback_metrics.get(key)
def _should_check_val_fx(self, batch_idx: int, is_last_batch: bool) -> bool:
"""Decide if we should run validation."""
if not self.trainer.enable_validation:
return False
is_val_check_epoch = (self.trainer.current_epoch + 1) % self.trainer.check_val_every_n_epoch == 0
if not is_val_check_epoch:
return False
# val_check_batch is inf for iterable datasets with no length defined
is_infinite_dataset = self.trainer.val_check_batch == float("inf")
if is_last_batch and is_infinite_dataset:
return True
if self.trainer.should_stop:
return True
# TODO(@awaelchli): let training/eval loop handle logic around limit_*_batches and val_check_batch
is_val_check_batch = is_last_batch
if isinstance(self.trainer.limit_train_batches, int) and is_infinite_dataset:
is_val_check_batch = (batch_idx + 1) % self.trainer.limit_train_batches == 0
elif self.trainer.val_check_batch != float("inf"):
is_val_check_batch = (batch_idx + 1) % self.trainer.val_check_batch == 0
return is_val_check_batch
def _save_loggers_on_train_batch_end(self) -> None:
"""Flushes loggers to disk."""
# when loggers should save to disk
should_flush_logs = self.trainer.logger_connector.should_flush_logs
if should_flush_logs and self.trainer.is_global_zero and self.trainer.logger is not None:
self.trainer.logger.save()
def _reload_dataloader_state_dict(self, data_fetcher: AbstractDataFetcher):
if self._dataloader_state_dict:
data_fetcher.dataloader.load_state_dict(self._dataloader_state_dict)
self._dataloader_state_dict = None
def test_epoch_loop_progress_increment_sequence():
"""Test sequences for incrementing batches reads and epochs."""
batch = BatchProgress(total=Tracker(started=None))
epoch = EpochProgress(batch=batch)
loop = EpochLoopProgress(epoch=epoch)
batch.increment_ready()
assert batch.total == Tracker(ready=1, started=None)
assert batch.current == Tracker(ready=1)
batch.increment_started()
assert batch.total == Tracker(ready=1, started=None)
assert batch.current == Tracker(ready=1)
batch.increment_processed()
assert batch.total == Tracker(ready=1, started=None, processed=1)
assert batch.current == Tracker(ready=1, processed=1)
batch.increment_completed()
assert batch.total == Tracker(ready=1, started=None, processed=1, completed=1)
assert batch.current == Tracker(ready=1, processed=1, completed=1)
assert epoch.total == Tracker()
assert epoch.current == Tracker()
loop.increment_epoch_completed()
assert batch.total == Tracker(ready=1, started=None, processed=1, completed=1)
assert batch.current == Tracker()
assert epoch.total == Tracker(completed=1)
assert epoch.current == Tracker()
batch.increment_ready()
assert batch.total == Tracker(ready=2, started=None, processed=1, completed=1)
assert batch.current == Tracker(ready=1)
assert epoch.total == Tracker(completed=1)
assert epoch.current == Tracker()
loop.reset_on_epoch()
assert batch.total == Tracker(ready=2, started=None, processed=1, completed=1)
assert batch.current == Tracker()
assert epoch.total == Tracker(completed=1)
assert epoch.current == Tracker()
class EvaluationEpochLoop(Loop):
"""This is the loop performing the evaluation.
It mainly loops over the given dataloader and runs the validation or test step (depending on the trainer's current
state).
"""
def __init__(self) -> None:
super().__init__()
self.batch_progress = BatchProgress()
self._outputs: EPOCH_OUTPUT = []
self._dl_max_batches = 0
self._data_fetcher: Optional[AbstractDataFetcher] = None
self._dataloader_state_dict: Dict[str, Any] = {}
@property
def done(self) -> bool:
"""Returns ``True`` if the current iteration count reaches the number of dataloader batches."""
return self.batch_progress.current.completed >= self._dl_max_batches
def reset(self) -> None:
"""Resets the loop's internal state."""
self._dl_max_batches = 0
self._data_fetcher = None
self._outputs = []
if not self.restarting:
self.batch_progress.reset_on_run()
else:
self.batch_progress.reset_on_restart()
# when restarting, if we are running `validate` or `test` twice, since there's no concept of `max_epochs` we
# need to reset the current state when the loop has finished running
if self.done and self.trainer.state.fn != TrainerFn.FITTING:
self.batch_progress.reset_on_run()
def on_run_start( # type: ignore[override]
self, data_fetcher: AbstractDataFetcher, dl_max_batches: int,
kwargs: OrderedDict) -> None:
"""Adds the passed arguments to the loop's state if necessary.
Args:
data_fetcher: the current data_fetcher wrapping the dataloader
dl_max_batches: maximum number of batches the dataloader can produce
kwargs: the kwargs passed down to the hooks.
"""
void(kwargs)
self._dl_max_batches = dl_max_batches
self._reload_dataloader_state_dict(data_fetcher)
# creates the iterator inside the fetcher but returns `self`
self._data_fetcher = iter(data_fetcher)
# add the previous `fetched` value to properly track `is_last_batch` with no prefetching
data_fetcher.fetched += self.batch_progress.current.ready
stage = self.trainer.state.stage
assert stage is not None
stage = stage.dataloader_prefix
self._profiler_fetch_action = (
f"[{self.__class__.__name__}].{stage}_dataloader_idx_{kwargs.get('dataloader_idx', 0)}_next"
)
data_fetcher._start_profiler = self._on_before_fetch
data_fetcher._stop_profiler = self._on_after_fetch
def _on_before_fetch(self) -> None:
self.trainer.profiler.start(self._profiler_fetch_action)
def _on_after_fetch(self) -> None:
self.trainer.profiler.stop(self._profiler_fetch_action)
def advance( # type: ignore[override]
self,
data_fetcher: AbstractDataFetcher,
dl_max_batches: int,
kwargs: OrderedDict,
) -> None:
"""Calls the evaluation step with the corresponding hooks and updates the logger connector.
Args:
data_fetcher: iterator over the dataloader
dl_max_batches: maximum number of batches the dataloader can produce
kwargs: the kwargs passed down to the hooks.
Raises:
StopIteration: If the current batch is None
"""
void(dl_max_batches)
if not isinstance(data_fetcher, DataLoaderIterDataFetcher):
batch_idx = self.batch_progress.current.ready
batch = next(data_fetcher)
else:
batch_idx, batch = next(data_fetcher)
self.batch_progress.is_last_batch = data_fetcher.done
# configure step_kwargs
kwargs = self._build_kwargs(kwargs, batch, batch_idx)
self.batch_progress.increment_ready()
# hook
self._on_evaluation_batch_start(**kwargs)
self.batch_progress.increment_started()
# lightning module methods
output = self._evaluation_step(**kwargs)
output = self._evaluation_step_end(output)
self.batch_progress.increment_processed()
# track loss history
self._on_evaluation_batch_end(output, **kwargs)
self.batch_progress.increment_completed()
# log batch metrics
self.trainer._logger_connector.update_eval_step_metrics()
# track epoch level outputs
if self._should_track_batch_outputs_for_epoch_end(
) and output is not None:
self._outputs.append(output)
if self.trainer.move_metrics_to_cpu:
# the evaluation step output is not moved as they are not considered "metrics"
assert self.trainer._results is not None
self.trainer._results.cpu()
if not self.batch_progress.is_last_batch:
# if fault tolerant is enabled and process has been notified, exit.
self.trainer._exit_gracefully_on_signal()
def on_run_end(self) -> EPOCH_OUTPUT:
"""Returns the outputs of the whole run."""
outputs, self._outputs = self._outputs, [] # free memory
self._data_fetcher = None
return outputs
def teardown(self) -> None:
# in case the model changes
self._should_track_batch_outputs_for_epoch_end.cache_clear()
def on_save_checkpoint(self) -> Dict:
state_dict = super().on_save_checkpoint()
if (self.trainer is not None
and self.trainer.state._fault_tolerant_mode.is_enabled
and self._data_fetcher is not None
and not self._num_completed_batches_reached() # did not finish
and self.batch_progress.current.ready # did start
):
state = CombinedLoader._state_dict_fn(
self._data_fetcher.dataloader_iter, self._has_completed())
if state:
state_dict[
"dataloader_state_dict"] = _collect_states_on_rank_zero_over_collection(
state)
return state_dict
def on_load_checkpoint(self, state_dict: Dict) -> None:
# cache the dataloader state dict until the dataloader objects are available
# dataset states are collected across all ranks
dataloader_state_dict = state_dict.get("dataloader_state_dict", None)
if not _fault_tolerant_training() or not dataloader_state_dict:
return
self._dataloader_state_dict = dataloader_state_dict[
self.trainer.global_rank]
def _reload_dataloader_state_dict(
self, data_fetcher: AbstractDataFetcher) -> None:
if self.trainer.sanity_checking or not self._dataloader_state_dict:
return
dataloader = data_fetcher.dataloader
if isinstance(dataloader, CombinedLoader):
raise MisconfigurationException(
"Reloading support hasn't been implemented for `CombinedLoader`. You can request it by opening an issue"
" in `https://github.com/PyTorchLightning/pytorch-lightning/issues`."
)
assert isinstance(dataloader, DataLoader)
_reload_dataloader_state_dict(dataloader, self._dataloader_state_dict)
self._dataloader_state_dict = {}
def _num_completed_batches_reached(self) -> bool:
epoch_finished_on_completed = self.batch_progress.current.completed == self._dl_max_batches
dataloader_consumed_successfully = self.batch_progress.is_last_batch and self._has_completed(
)
return epoch_finished_on_completed or dataloader_consumed_successfully
def _has_completed(self) -> bool:
return self.batch_progress.current.ready == self.batch_progress.current.completed
def _evaluation_step(self, **kwargs: Any) -> Optional[STEP_OUTPUT]:
"""The evaluation step (validation_step or test_step depending on the trainer's state).
Args:
batch: The current batch to run through the step.
batch_idx: The index of the current batch
dataloader_idx: the index of the dataloader producing the current batch
Returns:
the outputs of the step
"""
hook_name = "test_step" if self.trainer.testing else "validation_step"
output = self.trainer._call_strategy_hook(hook_name, *kwargs.values())
return output
def _evaluation_step_end(self, *args: Any,
**kwargs: Any) -> Optional[STEP_OUTPUT]:
"""Calls the `{validation/test}_step_end` hook."""
hook_name = "test_step_end" if self.trainer.testing else "validation_step_end"
model_output = self.trainer._call_lightning_module_hook(
hook_name, *args, **kwargs)
strategy_output = self.trainer._call_strategy_hook(
hook_name, *args, **kwargs)
output = strategy_output if model_output is None else model_output
return output
def _on_evaluation_batch_start(self, **kwargs: Any) -> None:
"""Calls the ``on_{validation/test}_batch_start`` hook.
Args:
batch: The current batch to run through the step
batch_idx: The index of the current batch
dataloader_idx: The index of the dataloader producing the current batch
Raises:
AssertionError: If the number of dataloaders is None (has not yet been set).
"""
self.trainer._logger_connector.on_batch_start(**kwargs)
kwargs.setdefault("dataloader_idx",
0) # TODO: the argument should be keyword for these
hook_name = "on_test_batch_start" if self.trainer.testing else "on_validation_batch_start"
self.trainer._call_callback_hooks(hook_name, *kwargs.values())
self.trainer._call_lightning_module_hook(hook_name, *kwargs.values())
def _on_evaluation_batch_end(self, output: Optional[STEP_OUTPUT],
**kwargs: Any) -> None:
"""The ``on_{validation/test}_batch_end`` hook.
Args:
output: The output of the performed step
batch: The input batch for the step
batch_idx: The index of the current batch
dataloader_idx: Index of the dataloader producing the current batch
"""
kwargs.setdefault("dataloader_idx",
0) # TODO: the argument should be keyword for these
hook_name = "on_test_batch_end" if self.trainer.testing else "on_validation_batch_end"
self.trainer._call_callback_hooks(hook_name, output, *kwargs.values())
self.trainer._call_lightning_module_hook(hook_name, output,
*kwargs.values())
self.trainer._logger_connector.on_batch_end()
def _build_kwargs(self, kwargs: OrderedDict, batch: Any,
batch_idx: int) -> OrderedDict:
"""Helper method to build the arguments for the current step.
Args:
kwargs: The kwargs passed down to the hooks.
batch: The current batch to run through the step.
Returns:
The kwargs passed down to the hooks.
"""
kwargs.update(batch=batch, batch_idx=batch_idx)
# `dataloader_idx` should be last so we need to push these to the front
kwargs.move_to_end("batch_idx", last=False)
kwargs.move_to_end("batch", last=False)
return kwargs
@lru_cache(1)
def _should_track_batch_outputs_for_epoch_end(self) -> bool:
"""Whether the batch outputs should be stored for later usage."""
model = self.trainer.lightning_module
if self.trainer.testing:
return is_overridden("test_epoch_end", model)
return is_overridden("validation_epoch_end", model)
class TrainingEpochLoop(loops.Loop[_OUTPUTS_TYPE]):
"""Runs over all batches in a dataloader (one epoch).
Args:
min_steps: The minimum number of steps (batches) to process
max_steps: The maximum number of steps (batches) to process
"""
def __init__(self,
min_steps: Optional[int] = None,
max_steps: int = -1) -> None:
super().__init__()
if max_steps is None:
rank_zero_deprecation(
"Setting `max_steps = None` is deprecated in v1.5 and will no longer be supported in v1.7."
" Use `max_steps = -1` instead.")
max_steps = -1
elif max_steps < -1:
raise MisconfigurationException(
f"`max_steps` must be a non-negative integer or -1 (infinite steps). You passed in {max_steps}."
)
self.min_steps = min_steps
self.max_steps = max_steps
self.batch_progress = BatchProgress()
self.scheduler_progress = SchedulerProgress()
self.batch_loop = TrainingBatchLoop()
self.val_loop = loops.EvaluationLoop(verbose=False)
self._results = _ResultCollection(training=True)
self._outputs: _OUTPUTS_TYPE = []
self._warning_cache = WarningCache()
# caches the loaded dataloader state until dataloader objects are available
self._dataloader_state_dict: Dict[str, Any] = {}
self._batches_that_stepped: int = 0
@property
def total_batch_idx(self) -> int:
"""Returns the current batch index (across epochs)"""
# use `ready` instead of `completed` in case this is accessed after `completed` has been increased
# but before the next `ready` increase
return self.batch_progress.total.ready - 1
@property
def batch_idx(self) -> int:
"""Returns the current batch index (within this epoch)"""
# use `ready` instead of `completed` in case this is accessed after `completed` has been increased
# but before the next `ready` increase
return self.batch_progress.current.ready - 1
@property
def global_step(self) -> int:
lightning_module = self.trainer.lightning_module
if lightning_module is None or lightning_module.automatic_optimization:
return self.batch_loop.optimizer_loop.optim_progress.optimizer_steps
return self.batch_loop.manual_loop.optim_step_progress.total.completed
@property
def _is_training_done(self) -> bool:
max_steps_reached = _is_max_limit_reached(self.global_step,
self.max_steps)
return max_steps_reached or self._num_ready_batches_reached()
@property
def _is_validation_done(self) -> bool:
# when we are restarting we want to check whether the val loop has finished
return not self.restarting or self.val_loop.done
@property
def done(self) -> bool:
"""Evaluates when to leave the loop."""
return (self._is_training_done
and self._is_validation_done) or self.trainer.should_stop
def connect( # type: ignore[override]
self,
batch_loop: Optional[TrainingBatchLoop] = None,
val_loop: Optional["loops.EvaluationLoop"] = None,
) -> None:
"""Optionally connect a custom batch or validation loop to this training epoch loop."""
if batch_loop is not None:
self.batch_loop = batch_loop
if val_loop is not None:
self.val_loop = val_loop
def reset(self) -> None:
"""Resets the internal state of the loop for a new run."""
if self.restarting:
self.batch_progress.reset_on_restart()
self.scheduler_progress.reset_on_restart()
self.batch_loop.optimizer_loop.optim_progress.reset_on_restart()
trainer = self.trainer
if not trainer.state._fault_tolerant_mode.is_enabled and trainer.num_training_batches != float(
"inf"):
expected_steps = math.ceil(trainer.num_training_batches /
trainer.accumulate_grad_batches)
if self.global_step % expected_steps != 0:
rank_zero_warn(
"You're resuming from a checkpoint that ended before the epoch ended. This can cause unreliable"
" results if further training is done. Consider using an end-of-epoch checkpoint or enabling"
" fault-tolerant training:"
" https://pytorch-lightning.readthedocs.io/en/stable/advanced/fault_tolerant_training.html"
)
else:
self.batch_progress.reset_on_run()
self.scheduler_progress.reset_on_run()
self.batch_loop.optimizer_loop.optim_progress.reset_on_run()
# when the epoch starts, the total val batch progress should be reset as it's supposed to count the batches
# seen per epoch, this is useful for tracking when validation is run multiple times per epoch
self.val_loop.epoch_loop.batch_progress.total.reset()
self._outputs = []
def on_run_start(
self, data_fetcher: AbstractDataFetcher
) -> None: # type: ignore[override]
self._reload_dataloader_state_dict(data_fetcher)
_ = iter(data_fetcher) # creates the iterator inside the fetcher
# add the previous `fetched` value to properly track `is_last_batch` with no prefetching
data_fetcher.fetched += self.batch_progress.current.ready
def advance(
self, data_fetcher: AbstractDataFetcher
) -> None: # type: ignore[override]
"""Runs a single training batch.
Raises:
StopIteration: When the epoch is canceled by the user returning -1
"""
if self.restarting and self._should_check_val_fx(
self.batch_idx, self.batch_progress.is_last_batch):
# skip training and run validation in `on_advance_end`
return
# we are going to train first so the val loop does not need to restart
self.val_loop.restarting = False
if not isinstance(data_fetcher, DataLoaderIterDataFetcher):
batch_idx = self.batch_idx + 1
batch = next(data_fetcher)
else:
batch_idx, batch = next(data_fetcher)
self.batch_progress.is_last_batch = data_fetcher.done
self.batch_progress.increment_ready()
self.trainer._logger_connector.on_batch_start(batch, batch_idx)
if batch is None:
self._warning_cache.warn(
"train_dataloader yielded None. If this was on purpose, ignore this warning..."
)
batch_output = []
else:
# hook
self.trainer._call_callback_hooks("on_batch_start")
# TODO: Update this in v1.7 (deprecation: #9816)
model_fx = self.trainer.lightning_module.on_train_batch_start
extra_kwargs = ({
"dataloader_idx": 0
} if callable(model_fx) and is_param_in_hook_signature(
model_fx, "dataloader_idx", explicit=True) else {})
# hook
self.trainer._call_callback_hooks("on_train_batch_start", batch,
batch_idx, **extra_kwargs)
response = self.trainer._call_lightning_module_hook(
"on_train_batch_start", batch, batch_idx, **extra_kwargs)
self.trainer._call_strategy_hook("on_train_batch_start", batch,
batch_idx, **extra_kwargs)
if response == -1:
self.batch_progress.increment_processed()
raise StopIteration
self.batch_progress.increment_started()
with self.trainer.profiler.profile("run_training_batch"):
batch_output = self.batch_loop.run(batch, batch_idx)
self.batch_progress.increment_processed()
# update non-plateau LR schedulers
# update epoch-interval ones only when we are at the end of training epoch
self.update_lr_schedulers("step", update_plateau_schedulers=False)
if self._num_ready_batches_reached():
self.update_lr_schedulers("epoch", update_plateau_schedulers=False)
batch_end_outputs = self._prepare_outputs_training_batch_end(
batch_output,
lightning_module=self.trainer.lightning_module,
num_optimizers=len(self.trainer.optimizers),
)
# TODO: Update this in v1.7 (deprecation: #9816)
model_fx = self.trainer.lightning_module.on_train_batch_end
extra_kwargs = ({
"dataloader_idx": 0
} if callable(model_fx) and is_param_in_hook_signature(
model_fx, "dataloader_idx", explicit=True) else {})
self.trainer._call_callback_hooks("on_train_batch_end",
batch_end_outputs, batch, batch_idx,
**extra_kwargs)
self.trainer._call_lightning_module_hook("on_train_batch_end",
batch_end_outputs, batch,
batch_idx, **extra_kwargs)
self.trainer._call_callback_hooks("on_batch_end")
self.trainer._logger_connector.on_batch_end()
self.batch_progress.increment_completed()
if is_overridden("training_epoch_end", self.trainer.lightning_module):
self._outputs.append(batch_output)
# -----------------------------------------
# SAVE METRICS TO LOGGERS AND PROGRESS_BAR
# -----------------------------------------
self.trainer._logger_connector.update_train_step_metrics()
def on_advance_end(self) -> None:
# -----------------------------------------
# VALIDATE IF NEEDED
# -----------------------------------------
should_check_val = self._should_check_val_fx(
self.batch_idx, self.batch_progress.is_last_batch)
if should_check_val:
self.trainer.validating = True
self._run_validation()
self.trainer.training = True
# update plateau LR scheduler after metrics are logged
self.update_lr_schedulers("step", update_plateau_schedulers=True)
if not self._should_accumulate():
# this is increased once per batch disregarding multiple optimizers or tbptt on purpose for loggers
self._batches_that_stepped += 1
# this will save based on the `batches_that_stepped` value
self._save_loggers_on_train_batch_end()
# if training finished, defer exit to the parent. this assumes there will be enough time in between
# which might not be the case depending on what's in the `*_epoch_end` hooks
if not self._is_training_done:
# if fault tolerant is enabled and process has been notified, exit.
self.trainer._exit_gracefully_on_signal()
def on_run_end(self) -> _OUTPUTS_TYPE:
outputs, self._outputs = self._outputs, []
return outputs
def teardown(self) -> None:
self._results.cpu()
self.batch_loop.teardown()
self.val_loop.teardown()
def on_save_checkpoint(self) -> Dict:
state_dict = super().on_save_checkpoint()
if (self.trainer is not None
and self.trainer.state._fault_tolerant_mode.is_enabled
and self.trainer.train_dataloader is not None
and not self._num_completed_batches_reached() # did not finish
# TODO: fault-tolerance requires a minimum number of batches so probably should be > 0
and self.batch_progress.current.ready # did start
):
loader: CombinedLoader = self.trainer.train_dataloader
state = loader.state_dict(has_completed=self._has_completed())
if state:
state_dict[
"dataloader_state_dict"] = _collect_states_on_rank_zero_over_collection(
state)
return state_dict
def on_load_checkpoint(self, state_dict: Dict) -> None:
# cache the dataloader state dict until the dataloader objects are available
self._dataloader_state_dict = state_dict.get("dataloader_state_dict")
def _run_validation(self) -> None:
# reload dataloaders
self.val_loop._reload_evaluation_dataloaders()
with torch.no_grad():
self.val_loop.run()
def _accumulated_batches_reached(self) -> bool:
"""Determine if accumulation will be finished by the end of the current batch."""
return self.batch_progress.current.ready % self.trainer.accumulate_grad_batches == 0
def _num_ready_batches_reached(self) -> bool:
"""Checks if we are in the last batch or if there are more batches to follow."""
epoch_finished_on_ready = self.batch_progress.current.ready == self.trainer.num_training_batches
return epoch_finished_on_ready or self.batch_progress.is_last_batch
def _num_completed_batches_reached(self) -> bool:
epoch_finished_on_completed = self.batch_progress.current.completed == self.trainer.num_training_batches
dataloader_consumed_successfully = self.batch_progress.is_last_batch and self._has_completed(
)
return epoch_finished_on_completed or dataloader_consumed_successfully
def _has_completed(self) -> bool:
return self.batch_progress.current.ready == self.batch_progress.current.completed
def _should_accumulate(self) -> bool:
"""Checks if the optimizer step should be performed or gradients should be accumulated for the current
step."""
accumulation_done = self._accumulated_batches_reached()
# Lightning steps on the final batch
is_final_batch = self._num_ready_batches_reached()
# but the strategy might not
strategy_accumulates_on_final_batch = self.trainer.strategy.handles_gradient_accumulation or not is_final_batch
return not accumulation_done and strategy_accumulates_on_final_batch
@staticmethod
def _prepare_outputs_training_batch_end(
batch_output: _BATCH_OUTPUTS_TYPE,
lightning_module: "pl.LightningModule",
num_optimizers: int,
) -> Union[List[List[Dict[str, Any]]], List[Dict[str, Any]]]:
"""Processes the outputs from the batch loop into the format passed to the ``on_train_batch_end`` hook."""
if not batch_output:
return []
# convert optimizer dicts to list
if lightning_module.automatic_optimization:
batch_output = apply_to_collection(batch_output,
dtype=dict,
function=_convert_optim_dict,
num_optimizers=num_optimizers)
array = np.array(batch_output, dtype=object)
# TODO: remove in v1.8
if (num_optimizers > 1 and lightning_module.truncated_bptt_steps > 0
and
not _v1_8_output_format(lightning_module.on_train_batch_end)):
rank_zero_deprecation(
"You are training with multiple optimizers AND truncated backpropagation through time enabled."
" The current format of the `on_train_batch_end(outputs, ...)` is a 2d list with sizes"
" (n_optimizers, tbptt_steps), however, this has been deprecated and will change in version v1.8 to"
" (tbptt_steps, n_optimizers). You can update your code by adding the following parameter to your"
" hook signature: `on_train_batch_end(outputs, ..., new_format=True)`."
)
# (tbptt_steps, n_opt) -> (n_opt, tbptt_steps)
if array.ndim == 1:
array = np.expand_dims(array, 1)
array = array.transpose((1, 0))
# squeeze all single-element dimensions
array = array.squeeze()
array = array.tolist()
array = _recursive_unpad(array)
return array
@staticmethod
def _prepare_outputs_training_epoch_end(
batch_outputs: _OUTPUTS_TYPE,
lightning_module: "pl.LightningModule",
num_optimizers: int,
) -> Union[List[List[List[Dict[str, Any]]]], List[List[Dict[str, Any]]],
List[Dict[str, Any]]]:
"""Processes the outputs from the batch loop into the format passed to the ``training_epoch_end`` hook."""
# `batch_outputs` (plural) is the same as `epoch_end_output` (singular)
if not batch_outputs:
return []
# convert optimizer dicts to list
if lightning_module.automatic_optimization:
batch_outputs = apply_to_collection(batch_outputs,
dtype=dict,
function=_convert_optim_dict,
num_optimizers=num_optimizers)
array = _recursive_pad(batch_outputs)
# TODO: remove in v1.8
if (num_optimizers > 1 and lightning_module.truncated_bptt_steps > 0
and
not _v1_8_output_format(lightning_module.on_train_epoch_end)):
rank_zero_deprecation(
"You are training with multiple optimizers AND truncated backpropagation through time enabled."
" The current format of the `training_epoch_end(outputs)` is a 3d list with sizes"
" (n_optimizers, n_batches, tbptt_steps), however, this has been deprecated and will change in version"
" v1.8 to (n_batches, tbptt_steps, n_optimizers). You can update your code by adding the following"
" parameter to your hook signature: `training_epoch_end(outputs, new_format=True)`."
)
# (n_batches, tbptt_steps, n_opt) -> (n_opt, n_batches, tbptt_steps)
if array.ndim == 2:
array = np.expand_dims(array, 2)
array = array.transpose((2, 0, 1))
# squeeze all single-element dimensions
array = array.squeeze()
array = array.tolist()
array = _recursive_unpad(array)
# in case we squeezed from 1-element array to a 0-dim array
array = array if isinstance(array, list) else [array]
# remove residual empty lists
array = [
item for item in array if not isinstance(item, list) or len(item)
]
return array
def update_lr_schedulers(self, interval: str,
update_plateau_schedulers: bool) -> None:
"""updates the lr schedulers based on the given interval."""
if interval == "step" and self._should_accumulate():
return
active_optimizers = _get_active_optimizers(
self.trainer.optimizers, self.trainer.optimizer_frequencies,
self.total_batch_idx)
self._update_learning_rates(
interval=interval,
update_plateau_schedulers=update_plateau_schedulers,
opt_indices=[opt_idx for opt_idx, _ in active_optimizers],
)
def _update_learning_rates(
self,
interval: str,
update_plateau_schedulers: bool,
opt_indices: Optional[List[int]] = None) -> None:
"""Update learning rates.
Args:
interval: either 'epoch' or 'step'.
update_plateau_schedulers: control whether ``ReduceLROnPlateau`` or non-plateau schedulers get updated.
This is used so non-plateau schedulers can be updated before running validation. Checkpoints are
commonly saved during validation, however, on-plateau schedulers might monitor a validation metric
so they have to be updated separately.
opt_indices: indices of the optimizers to update.
"""
if not self.trainer.lr_scheduler_configs or not self.trainer.lightning_module.automatic_optimization:
return
if opt_indices is None:
opt_indices = []
for config in self.trainer.lr_scheduler_configs:
if config.opt_idx not in opt_indices:
continue
if update_plateau_schedulers ^ config.reduce_on_plateau:
continue
current_idx = self.batch_idx if interval == "step" else self.trainer.current_epoch
current_idx += 1 # account for both batch and epoch starts from 0
# Take step if call to update_learning_rates matches the interval key and
# the current step modulo the schedulers frequency is zero
if config.interval == interval and current_idx % config.frequency == 0:
monitor_val = None
if config.reduce_on_plateau:
# If instance of ReduceLROnPlateau, we need a monitor
monitor_key = config.monitor
monitor_val = self._get_monitor_value(monitor_key)
if monitor_val is None:
if config.strict:
avail_metrics = list(self.trainer.callback_metrics)
raise MisconfigurationException(
f"ReduceLROnPlateau conditioned on metric {monitor_key}"
f" which is not available. Available metrics are: {avail_metrics}."
" Condition can be set using `monitor` key in lr scheduler dict"
)
rank_zero_warn(
f"ReduceLROnPlateau conditioned on metric {monitor_key}"
" which is not available but strict is set to `False`."
" Skipping learning rate update.",
category=RuntimeWarning,
)
continue
self.scheduler_progress.increment_ready()
# update LR
self.trainer._call_lightning_module_hook(
"lr_scheduler_step",
config.scheduler,
config.opt_idx,
monitor_val,
)
self.scheduler_progress.increment_completed()
def _get_monitor_value(self, key: str) -> Any:
# this is a separate method to aid in testing
return self.trainer.callback_metrics.get(key)
def _should_check_val_epoch(self):
return (self.trainer.enable_validation
and (self.trainer.current_epoch + 1) %
self.trainer.check_val_every_n_epoch == 0)
def _should_check_val_fx(self, batch_idx: int,
is_last_batch: bool) -> bool:
"""Decide if we should run validation."""
if not self._should_check_val_epoch():
return False
# val_check_batch is inf for iterable datasets with no length defined
is_infinite_dataset = self.trainer.val_check_batch == float("inf")
if is_last_batch and is_infinite_dataset:
return True
if self.trainer.should_stop:
return True
# TODO(@awaelchli): let training/eval loop handle logic around limit_*_batches and val_check_batch
is_val_check_batch = is_last_batch
if isinstance(self.trainer.limit_train_batches,
int) and is_infinite_dataset:
is_val_check_batch = (batch_idx +
1) % self.trainer.limit_train_batches == 0
elif self.trainer.val_check_batch != float("inf"):
is_val_check_batch = (batch_idx +
1) % self.trainer.val_check_batch == 0
return is_val_check_batch
def _save_loggers_on_train_batch_end(self) -> None:
"""Flushes loggers to disk."""
# this assumes that `batches_that_stepped` was increased before
should_flush = self._batches_that_stepped % self.trainer.flush_logs_every_n_steps == 0
if should_flush or self.trainer.should_stop:
for logger in self.trainer.loggers:
logger.save()
def _reload_dataloader_state_dict(
self, data_fetcher: AbstractDataFetcher) -> None:
if self._dataloader_state_dict:
data_fetcher.dataloader.load_state_dict(
self._dataloader_state_dict)
self._dataloader_state_dict = None
class EvaluationEpochLoop(Loop):
"""This is the loop performing the evaluation.
It mainly loops over the given dataloader and runs the validation or test step (depending on the trainer's current
state).
"""
def __init__(self) -> None:
super().__init__()
self.outputs: EPOCH_OUTPUT = []
self.batch_progress = BatchProgress()
self._dl_max_batches: Optional[int] = None
self._num_dataloaders: Optional[int] = None
self._dataloader_iter: Optional[Iterator] = None
self._data_fetcher: Optional[DataFetcher] = None
self._dataloader_state_dict: Dict[str, Any] = None
@property
def done(self) -> bool:
"""Returns ``True`` if the current iteration count reaches the number of dataloader batches."""
return self.batch_progress.current.completed >= self._dl_max_batches
def connect(self, **kwargs: "Loop") -> None:
raise NotImplementedError(f"{self.__class__.__name__} does not connect any child loops.")
def reset(self) -> None:
"""Resets the loop's internal state."""
self._dl_max_batches = None
self._num_dataloaders = None
self._data_fetcher = None
self.outputs = []
if not self.restarting:
self.batch_progress.reset_on_run()
else:
self.batch_progress.reset_on_restart()
def on_run_start(
self, data_fetcher: AbstractDataFetcher, dataloader_idx: int, dl_max_batches: int, num_dataloaders: int
) -> None:
"""Adds the passed arguments to the loop's state if necessary.
Args:
data_fetcher: the current data_fetcher wrapping the dataloader
dataloader_idx: index of the current dataloader
dl_max_batches: maximum number of batches the dataloader can produce
num_dataloaders: the total number of dataloaders
"""
void(dataloader_idx)
self._dl_max_batches = dl_max_batches
self._num_dataloaders = num_dataloaders
self._data_fetcher = data_fetcher
self._reload_dataloader_state_dict(data_fetcher)
self._dataloader_iter = _update_dataloader_iter(data_fetcher, self.batch_progress.current.ready)
def advance(
self, data_fetcher: AbstractDataFetcher, dataloader_idx: int, dl_max_batches: int, num_dataloaders: int
) -> None:
"""Calls the evaluation step with the corresponding hooks and updates the logger connector.
Args:
data_fetcher: iterator over the dataloader
dataloader_idx: index of the current dataloader
dl_max_batches: maximum number of batches the dataloader can produce
num_dataloaders: the total number of dataloaders
Raises:
StopIteration: If the current batch is None
"""
void(data_fetcher, dl_max_batches, num_dataloaders)
batch_idx, (batch, self.batch_progress.is_last_batch) = next(self._dataloader_iter)
if batch is None:
raise StopIteration
if not self.trainer._data_connector.evaluation_data_fetcher.store_on_device:
with self.trainer.profiler.profile("evaluation_batch_to_device"):
batch = self.trainer.accelerator.batch_to_device(batch, dataloader_idx=dataloader_idx)
self.batch_progress.increment_ready()
# hook
self._on_evaluation_batch_start(batch, batch_idx, dataloader_idx)
self.batch_progress.increment_started()
# lightning module methods
with self.trainer.profiler.profile("evaluation_step_and_end"):
output = self._evaluation_step(batch, batch_idx, dataloader_idx)
output = self._evaluation_step_end(output)
self.batch_progress.increment_processed()
# track loss history
self._on_evaluation_batch_end(output, batch, batch_idx, dataloader_idx)
self.batch_progress.increment_completed()
# log batch metrics
self.trainer.logger_connector.update_eval_step_metrics()
# track epoch level outputs
if self._should_track_batch_outputs_for_epoch_end():
output = recursive_detach(output, to_cpu=self.trainer.move_metrics_to_cpu)
if output is not None:
self.outputs.append(output)
if not self.batch_progress.is_last_batch:
# if fault tolerant is enabled and process has been notified, exit.
self.trainer._exit_gracefully_on_signal()
def on_run_end(self) -> EPOCH_OUTPUT:
"""Returns the outputs of the whole run."""
outputs = self.outputs
# free memory
self.outputs = []
self._dataloader_iter = None
self._data_fetcher = None
return outputs
def teardown(self) -> None:
# in case the model changes
self._should_track_batch_outputs_for_epoch_end.cache_clear()
def on_save_checkpoint(self) -> Dict:
state_dict = super().on_save_checkpoint()
if (
self._data_fetcher is None
or self._num_completed_batches_reached() # did not finish
# TODO: fault-tolerance requires a minimum number of batches so probably should be > 0
or self.batch_progress.current.ready == 0 # did not start
):
return state_dict
# TODO: this should use `pytorch_lightning/trainer/supporters.py::CombinedLoader._state_dict_fn`
state_to_save = "state" if self._has_completed() else "previous_state"
state: Optional[MergedIteratorState] = getattr(self._data_fetcher.dataloader_iter, state_to_save, None)
if state:
state_dict["dataloader_state_dict"] = asdict(state)
return state_dict
def on_load_checkpoint(self, state_dict: Dict) -> None:
# cache the dataloader state dict until the dataloader objects are available
self._dataloader_state_dict = state_dict.get("dataloader_state_dict")
def _reload_dataloader_state_dict(self, data_fetcher: AbstractDataFetcher):
if not self.trainer.sanity_checking and self._dataloader_state_dict:
reload_dataloader_state_dict(data_fetcher.dataloader, self._dataloader_state_dict)
self._dataloader_state_dict = None
def _num_completed_batches_reached(self) -> bool:
epoch_finished_on_completed = self.batch_progress.current.completed == self._dl_max_batches
dataloader_consumed_successfully = self.batch_progress.is_last_batch and self._has_completed()
return epoch_finished_on_completed or dataloader_consumed_successfully
def _has_completed(self) -> bool:
return self.batch_progress.current.ready == self.batch_progress.current.completed
def _evaluation_step(self, batch: Any, batch_idx: int, dataloader_idx: int) -> Optional[STEP_OUTPUT]:
"""The evaluation step (validation_step or test_step depending on the trainer's state).
Args:
batch: The current batch to run through the step.
batch_idx: The index of the current batch
dataloader_idx: the index of the dataloader producing the current batch
Returns:
the outputs of the step
"""
# configure step_kwargs
step_kwargs = self._build_kwargs(batch, batch_idx, dataloader_idx)
if self.trainer.testing:
self.trainer.lightning_module._current_fx_name = "test_step"
with self.trainer.profiler.profile("test_step"):
output = self.trainer.accelerator.test_step(step_kwargs)
else:
self.trainer.lightning_module._current_fx_name = "validation_step"
with self.trainer.profiler.profile("validation_step"):
output = self.trainer.accelerator.validation_step(step_kwargs)
return output
def _evaluation_step_end(self, *args: Any, **kwargs: Any) -> Optional[STEP_OUTPUT]:
"""Calls the `{validation/test}_step_end` hook."""
hook_name = "test_step_end" if self.trainer.testing else "validation_step_end"
output = self.trainer.call_hook(hook_name, *args, **kwargs)
return output
def _on_evaluation_batch_start(self, batch: Any, batch_idx: int, dataloader_idx: int) -> None:
"""Calls the ``on_{validation/test}_batch_start`` hook.
Args:
batch: The current batch to run through the step
batch_idx: The index of the current batch
dataloader_idx: The index of the dataloader producing the current batch
Raises:
AssertionError: If the number of dataloaders is None (has not yet been set).
"""
self.trainer.logger_connector.on_batch_start(batch_idx, batch)
assert self._num_dataloaders is not None
self.trainer.logger_connector.on_evaluation_batch_start(dataloader_idx, self._num_dataloaders)
if self.trainer.testing:
self.trainer.call_hook("on_test_batch_start", batch, batch_idx, dataloader_idx)
else:
self.trainer.call_hook("on_validation_batch_start", batch, batch_idx, dataloader_idx)
def _on_evaluation_batch_end(
self, output: Optional[STEP_OUTPUT], batch: Any, batch_idx: int, dataloader_idx: int
) -> None:
"""The ``on_{validation/test}_batch_end`` hook.
Args:
output: The output of the performed step
batch: The input batch for the step
batch_idx: The index of the current batch
dataloader_idx: Index of the dataloader producing the current batch
"""
hook_name = "on_test_batch_end" if self.trainer.testing else "on_validation_batch_end"
self.trainer.call_hook(hook_name, output, batch, batch_idx, dataloader_idx)
self.trainer.logger_connector.on_batch_end()
def _build_kwargs(self, batch: Any, batch_idx: int, dataloader_idx: int) -> Dict[str, Union[Any, int]]:
"""Helper function to build the arguments for the current step.
Args:
batch: The current batch to run through the step
batch_idx: the index of the current batch
dataloader_idx: the index of the dataloader producing the current batch
Returns:
the keyword arguments to pass to the step function
"""
# make dataloader_idx arg in validation_step optional
step_kwargs = OrderedDict([("batch", batch), ("batch_idx", batch_idx)])
multiple_val_loaders = not self.trainer.testing and self._num_dataloaders > 1
multiple_test_loaders = self.trainer.testing and self._num_dataloaders > 1
if multiple_test_loaders or multiple_val_loaders:
step_kwargs["dataloader_idx"] = dataloader_idx
return step_kwargs
@lru_cache(1)
def _should_track_batch_outputs_for_epoch_end(self) -> bool:
"""Whether the batch outputs should be stored for later usage."""
model = self.trainer.lightning_module
if self.trainer.testing:
return is_overridden("test_epoch_end", model)
return is_overridden("validation_epoch_end", model)