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
class EvaluationLoop(DataLoaderLoop):
"""Loops over all dataloaders for evaluation."""
def __init__(self):
super().__init__()
self.outputs: List[EPOCH_OUTPUT] = []
self.epoch_loop = EvaluationEpochLoop()
self._results = ResultCollection(training=False)
self._max_batches: Optional[Union[int, Sequence[int]]] = None
self._has_run: bool = False
@property
def num_dataloaders(self) -> int:
"""Returns the total number of dataloaders."""
# case where user does:
# return dl1, dl2
dataloaders = self.dataloaders
if dataloaders is None:
return 0
length = len(dataloaders)
if length > 0 and isinstance(dataloaders[0], (list, tuple)):
length = len(dataloaders[0])
return length
@property
def dataloaders(self) -> Sequence[DataLoader]:
"""Returns the validation or test dataloaders."""
if self.trainer.testing:
return self.trainer.test_dataloaders
return self.trainer.val_dataloaders
def connect(self, epoch_loop: EvaluationEpochLoop):
"""Connect the evaluation epoch loop with this loop."""
self.epoch_loop = epoch_loop
@property
def done(self) -> bool:
"""Returns whether all dataloaders are processed or evaluation should be skipped altogether."""
return super().done or self.skip
@property
def skip(self) -> bool:
"""Returns whether the evaluation should be skipped."""
max_batches = self.get_max_batches()
return sum(max_batches) == 0
def reset(self) -> None:
"""Resets the internal state of the loop."""
self._max_batches = self.get_max_batches()
# bookkeeping
self.outputs = []
if isinstance(self._max_batches, int):
self._max_batches = [self._max_batches] * len(self.dataloaders)
super().reset()
def on_skip(self) -> List:
return []
def on_run_start(self, *args: Any, **kwargs: Any) -> None:
"""Runs the ``on_evaluation_model_eval``, ``on_evaluation_start`` and ``on_evaluation_epoch_start``
hooks."""
void(*args, **kwargs)
# hook
self.on_evaluation_model_eval()
self.trainer.lightning_module.zero_grad()
self.on_evaluation_start()
self.on_evaluation_epoch_start()
def advance(self, *args: Any, **kwargs: Any) -> None:
"""Performs evaluation on one single dataloader."""
void(*args, **kwargs)
dataloader_idx: int = self.current_dataloader_idx
dataloader = self.trainer.accelerator.process_dataloader(
self.current_dataloader)
dataloader = self.trainer.data_connector.get_profiled_dataloader(
dataloader, dataloader_idx=dataloader_idx)
dl_max_batches = self._max_batches[dataloader_idx]
dl_outputs = self.epoch_loop.run(dataloader, dataloader_idx,
dl_max_batches, self.num_dataloaders)
# store batch level output per dataloader
self.outputs.append(dl_outputs)
if not self.trainer.sanity_checking:
# indicate the loop has run
self._has_run = True
def on_run_end(self) -> Any:
"""Runs the ``on_evaluation_epoch_end`` hook."""
outputs = self.outputs
# free memory
self.outputs = []
# with a single dataloader don't pass a 2D list
if len(outputs) > 0 and self.num_dataloaders == 1:
outputs = outputs[0]
# lightning module method
self.evaluation_epoch_end(outputs)
# hook
self.on_evaluation_epoch_end()
# log epoch metrics
eval_loop_results = self.trainer.logger_connector.update_eval_epoch_metrics(
)
# hook
self.on_evaluation_end()
# enable train mode again
self.on_evaluation_model_train()
return eval_loop_results
def get_max_batches(self) -> List[Union[int, float]]:
"""Returns the max number of batches for each dataloader."""
if self.trainer.testing:
max_batches = self.trainer.num_test_batches
else:
if self.trainer.sanity_checking:
self.trainer.num_sanity_val_batches = [
min(self.trainer.num_sanity_val_steps, val_batches)
for val_batches in self.trainer.num_val_batches
]
max_batches = self.trainer.num_sanity_val_batches
else:
max_batches = self.trainer.num_val_batches
return max_batches
def reload_evaluation_dataloaders(self) -> None:
"""Reloads dataloaders if necessary."""
if self.trainer.testing:
self.trainer.reset_test_dataloader()
elif self.trainer.val_dataloaders is None or self.trainer._should_reload_dl_epoch:
self.trainer.reset_val_dataloader()
def on_evaluation_start(self, *args: Any, **kwargs: Any) -> None:
"""Runs ``on_{validation/test}_start`` hooks."""
assert self._results is not None
self._results.to(device=self.trainer.lightning_module.device)
if self.trainer.testing:
self.trainer.call_hook("on_test_start", *args, **kwargs)
else:
self.trainer.call_hook("on_validation_start", *args, **kwargs)
def on_evaluation_model_eval(self) -> None:
"""Sets model to eval mode."""
if self.trainer.testing:
self.trainer.call_hook("on_test_model_eval")
else:
self.trainer.call_hook("on_validation_model_eval")
def on_evaluation_model_train(self) -> None:
"""Sets model to train mode."""
model_ref = self.trainer.lightning_module
if self.trainer.testing:
model_ref.on_test_model_train()
else:
model_ref.on_validation_model_train()
def on_evaluation_end(self, *args: Any, **kwargs: Any) -> None:
"""Runs ``on_{validation/test}_end`` hook."""
if self.trainer.testing:
self.trainer.call_hook("on_test_end", *args, **kwargs)
else:
self.trainer.call_hook("on_validation_end", *args, **kwargs)
# reset any `torchmetrics.Metric` and the logger connector state
self.trainer.logger_connector.reset(metrics=True)
def on_evaluation_epoch_start(self, *args: Any, **kwargs: Any) -> None:
"""Runs ``on_epoch_start`` and ``on_{validation/test}_epoch_start`` hooks."""
self.trainer.logger_connector.on_epoch_start()
self.trainer.call_hook("on_epoch_start", *args, **kwargs)
if self.trainer.testing:
self.trainer.call_hook("on_test_epoch_start", *args, **kwargs)
else:
self.trainer.call_hook("on_validation_epoch_start", *args,
**kwargs)
def evaluation_epoch_end(self, outputs: EPOCH_OUTPUT) -> None:
"""Runs ``{validation/test}_epoch_end``"""
# inform logger the batch loop has finished
self.trainer.logger_connector.epoch_end_reached()
# call the model epoch end
model = self.trainer.lightning_module
# unset dataloader_idx in model
model._current_dataloader_idx = None
if self.trainer.testing:
if is_overridden("test_epoch_end", model):
model._current_fx_name = "test_epoch_end"
model.test_epoch_end(outputs)
else:
if is_overridden("validation_epoch_end", model):
model._current_fx_name = "validation_epoch_end"
model.validation_epoch_end(outputs)
def on_evaluation_epoch_end(self) -> None:
"""Runs ``on_{validation/test}_epoch_end`` hook."""
hook_name = "on_test_epoch_end" if self.trainer.testing else "on_validation_epoch_end"
self.trainer.call_hook(hook_name)
self.trainer.call_hook("on_epoch_end")
self.trainer.logger_connector.on_epoch_end()
def teardown(self) -> None:
self._results.cpu()
self.epoch_loop.teardown()
class TrainingEpochLoop(loops.Loop):
"""
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: int, max_steps: int):
super().__init__()
self.min_steps: int = min_steps
self.max_steps: int = max_steps
self.global_step: int = 0
# manually tracking which is the last batch is necessary for iterable dataset support
self.is_last_batch: Optional[bool] = None
self.batch_progress = Progress()
self.scheduler_progress = SchedulerProgress()
self.batch_loop: Optional[TrainingBatchLoop] = None
self.val_loop: Optional["loops.EvaluationLoop"] = None
self._results = ResultCollection(training=True)
self._epoch_output: Optional[List[List[STEP_OUTPUT]]] = None
@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 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).
"""
max_steps_reached = self.max_steps is not None and self.global_step >= self.max_steps
return max_steps_reached or self.trainer.should_stop or self._num_training_batches_reached(
self.is_last_batch)
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"""
self.is_last_batch = False
# track epoch output
self._epoch_output = [[] for _ in range(
self.batch_loop.num_active_optimizers(self.total_batch_idx))]
if not self.restarting:
self.batch_progress.current.reset()
self.scheduler_progress.current.reset()
self.batch_loop.optim_progress.reset_on_epoch()
def on_run_start(self, dataloader_iter: Iterator, **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.dataloader_iter = _prepare_dataloader_iter(
dataloader_iter, 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
"""
batch_idx, (batch, is_last) = 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()
with self.trainer.profiler.profile("run_training_batch"):
batch_output = self.batch_loop.run(batch, batch_idx)
self.batch_progress.increment_processed()
self.is_last_batch = is_last
# when returning -1 from train_step, we end epoch early
if batch_output.signal == -1:
raise StopIteration
# 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_training_batches_reached(is_last):
self.update_lr_schedulers("epoch", update_plateau_schedulers=False)
batch_end_outputs = [
opt_idx_out for opt_idx_out in batch_output.training_step_output
if len(opt_idx_out)
]
processed_batch_end_outputs = self._prepare_outputs(batch_end_outputs,
batch_mode=True)
# hook
self.trainer.call_hook("on_train_batch_end",
processed_batch_end_outputs, batch,
self.batch_idx, 0)
self.trainer.call_hook("on_batch_end")
self.trainer.logger_connector.on_batch_end()
self.batch_progress.increment_completed()
# figure out what to track for epoch end
self._track_epoch_end_reduce_metrics(self._epoch_output,
batch_end_outputs)
# -----------------------------------------
# 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.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)
# progress global step according to grads progress
self._increment_accumulated_grad_global_step()
def on_run_end(self) -> List[List[STEP_OUTPUT]]:
"""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``
"""
if self.batch_progress.current.ready == 0:
# dataloader/iterator did not produce a batch
return
# inform logger the batch loop has finished
self.trainer.logger_connector.epoch_end_reached()
# prepare epoch output
processed_outputs = self._prepare_outputs(self._epoch_output,
batch_mode=False)
# get the model and call model.training_epoch_end
model = self.trainer.lightning_module
if is_overridden("training_epoch_end", model):
# run training_epoch_end
# refresh the result for custom logging at the epoch level
model._current_fx_name = "training_epoch_end"
# lightningmodule hook
training_epoch_end_output = model.training_epoch_end(
processed_outputs)
if training_epoch_end_output is not None:
raise MisconfigurationException(
"training_epoch_end expects a return of None. "
"HINT: remove the return statement in training_epoch_end")
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_training_batches_reached(self.is_last_batch):
self.update_lr_schedulers("epoch", update_plateau_schedulers=True)
epoch_output = self._epoch_output
# free memory
self._epoch_output = None
return epoch_output
def teardown(self) -> None:
self._results.cpu()
self.batch_loop.teardown()
self.val_loop.teardown()
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_training_batches_reached(self,
is_last_batch: bool = False) -> bool:
"""Checks if we are in the last batch or if there are more batches to follow.
Args:
is_last_batch: Whether the current batch is the last one
"""
return self.batch_progress.current.ready == self.trainer.num_training_batches or is_last_batch
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()
is_final_batch = self._num_training_batches_reached()
return not (accumulation_done or is_final_batch)
def _track_epoch_end_reduce_metrics(
self, epoch_output: List[List[STEP_OUTPUT]],
batch_end_outputs: STEP_OUTPUT) -> None:
"""Adds the batch outputs to the epoch outputs and prepares reduction"""
hook_overridden = is_overridden("training_epoch_end",
self.trainer.lightning_module)
if not hook_overridden:
return
# track the outputs to reduce at the end of the epoch
for opt_idx, opt_outputs in enumerate(batch_end_outputs):
# with 1 step (no tbptt) don't use a sequence at epoch end
if isinstance(opt_outputs, list) and len(opt_outputs) == 1:
opt_outputs = opt_outputs[0]
epoch_output[opt_idx].append(opt_outputs)
@staticmethod
def _prepare_outputs(
outputs: List[List[List["ResultCollection"]]], batch_mode: bool
) -> Union[List[List[List[Dict]]], List[List[Dict]], List[Dict], Dict]:
"""
Extract required information from batch or epoch end results.
Args:
outputs: A 3-dimensional list of ``ResultCollection`` objects with dimensions:
``[optimizer outs][batch outs][tbptt steps]``.
batch_mode: If True, ignore the batch output dimension.
Returns:
The cleaned outputs with ``ResultCollection`` objects converted to dictionaries.
All list dimensions of size one will be collapsed.
"""
processed_outputs = []
for opt_outputs in outputs:
# handle an edge case where an optimizer output is the empty list
if len(opt_outputs) == 0:
continue
processed_batch_outputs = []
if batch_mode:
opt_outputs = [opt_outputs]
for batch_outputs in opt_outputs:
processed_tbptt_outputs = []
if isinstance(batch_outputs, ResultCollection):
batch_outputs = [batch_outputs]
for tbptt_output in batch_outputs:
out = {}
if tbptt_output.minimize is not None:
out["loss"] = tbptt_output.minimize.detach()
out.update(tbptt_output.extra)
processed_tbptt_outputs.append(out)
# if there was only one tbptt step then we can collapse that dimension
if len(processed_tbptt_outputs) == 1:
processed_tbptt_outputs = processed_tbptt_outputs[0]
processed_batch_outputs.append(processed_tbptt_outputs)
# batch_outputs should be just one dict (or a list of dicts if using tbptt) per optimizer
if batch_mode:
processed_batch_outputs = processed_batch_outputs[0]
processed_outputs.append(processed_batch_outputs)
# if there is only one optimiser then we collapse that dimension
if len(processed_outputs) == 1:
processed_outputs = processed_outputs[0]
return processed_outputs
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
self.trainer.optimizer_connector.update_learning_rates(
interval=interval,
update_plateau_schedulers=update_plateau_schedulers,
opt_indices=[
opt_idx for opt_idx, _ in
self.batch_loop.get_active_optimizers(self.total_batch_idx)
],
)
def _increment_accumulated_grad_global_step(self) -> None:
"""Increments global step according to grads progress"""
if not self._should_accumulate():
self.global_step = self.trainer.accelerator.update_global_step(
self.batch_progress.current.ready, self.trainer.global_step)
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()
class TrainingEpochLoop(loops.Loop):
"""
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: int, max_steps: int):
super().__init__()
self.min_steps: int = min_steps
self.max_steps: int = max_steps
self.global_step: int = 0
# the total batch index across all epochs
self.total_batch_idx: int = 0
# the current split index when the batch gets split into chunks in truncated backprop through time
self.split_idx: Optional[int] = None
# the number of batches seen this run, updates immediately after batch_loop.run()
# TODO: replace by progress tracking
self.batches_seen: int = 0
self.is_last_batch: Optional[bool] = None
self.batch_progress = Progress()
self.scheduler_progress = SchedulerProgress()
self.batch_loop: Optional[TrainingBatchLoop] = None
self.val_loop: Optional["loops.EvaluationLoop"] = None
self._results = ResultCollection(training=True)
self._dataloader_idx: Optional[int] = None
self._warning_cache: WarningCache = WarningCache()
self._epoch_output: Optional[List[List[STEP_OUTPUT]]] = None
@property
def batch_idx(self) -> int:
"""Returns the current batch index (within this epoch)"""
return self.iteration_count
@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).
"""
max_steps_reached = self.max_steps is not None and self.global_step >= self.max_steps
return max_steps_reached or self.trainer.should_stop or self._num_training_batches_reached(
self.is_last_batch)
def connect(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"""
self.iteration_count = 0
self.batches_seen = 0
self.is_last_batch = False
self._dataloader_idx = 0
# track epoch output
self._epoch_output = [[] for _ in range(
self.batch_loop.num_active_optimizers(self.total_batch_idx))]
if self.restarting:
self.iteration_count = self.batches_seen = self.batch_progress.current.completed
else:
self.batch_progress.current.reset()
self.scheduler_progress.current.reset()
self.batch_loop.optim_progress.reset_on_epoch()
def on_run_start(self, *args: Any, **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()
def advance(self, dataloader_iter: Iterator, **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
"""
_, (batch, is_last) = next(dataloader_iter)
self.is_last_batch = is_last
# ------------------------------------
# TRAINING_STEP + TRAINING_STEP_END
# ------------------------------------
with self.trainer.profiler.profile("training_batch_to_device"):
batch = self.trainer.accelerator.batch_to_device(
batch, dataloader_idx=self._dataloader_idx)
self.batch_progress.increment_ready()
with self.trainer.profiler.profile("run_training_batch"):
batch_output = self.batch_loop.run(batch, self.iteration_count,
self._dataloader_idx)
self.batches_seen += 1
self.batch_progress.increment_processed()
# when returning -1 from train_step, we end epoch early
if batch_output.signal == -1:
raise StopIteration
# 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_training_batches_reached(is_last):
self.update_lr_schedulers("epoch", update_plateau_schedulers=False)
batch_end_outputs = [
opt_idx_out for opt_idx_out in batch_output.training_step_output
if len(opt_idx_out)
]
processed_batch_end_outputs = self._prepare_outputs(batch_end_outputs,
batch_mode=True)
# hook
self.trainer.call_hook("on_train_batch_end",
processed_batch_end_outputs, batch,
self.iteration_count, self._dataloader_idx)
self.trainer.call_hook("on_batch_end")
self.trainer.logger_connector.on_batch_end()
self.batch_progress.increment_completed()
# figure out what to track for epoch end
self._track_epoch_end_reduce_metrics(self._epoch_output,
batch_end_outputs)
# -----------------------------------------
# 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.iteration_count,
self.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)
self.total_batch_idx += 1
# progress global step according to grads progress
self._increment_accumulated_grad_global_step()
if self.done:
raise StopIteration
def on_run_end(self) -> List[List[STEP_OUTPUT]]:
"""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``
"""
if self.batches_seen == 0:
# dataloader/iterator did not produce a batch
return
# inform logger the batch loop has finished
self.trainer.logger_connector.epoch_end_reached()
# prepare epoch output
processed_outputs = self._prepare_outputs(self._epoch_output,
batch_mode=False)
# get the model and call model.training_epoch_end
model = self.trainer.lightning_module
if is_overridden("training_epoch_end", model):
# run training_epoch_end
# refresh the result for custom logging at the epoch level
model._current_fx_name = "training_epoch_end"
# lightningmodule hook
training_epoch_end_output = model.training_epoch_end(
processed_outputs)
if training_epoch_end_output is not None:
raise MisconfigurationException(
"training_epoch_end expects a return of None. "
"HINT: remove the return statement in training_epoch_end")
self.trainer.fit_loop.epoch_progress.increment_processed()
# call train epoch end hooks
self._on_train_epoch_end_hook(processed_outputs)
self.trainer.call_hook("on_epoch_end")
self.trainer.logger_connector.on_epoch_end()
self.update_lr_schedulers("epoch", update_plateau_schedulers=True)
epoch_output = self._epoch_output
# free memory
self._epoch_output = None
return epoch_output
def teardown(self) -> None:
self._results.cpu()
self.batch_loop.teardown()
self.val_loop.teardown()
def _run_validation(self):
# reload dataloaders
self.val_loop.reload_evaluation_dataloaders()
with torch.no_grad():
self.val_loop.run()
def _on_train_epoch_end_hook(
self, processed_epoch_output: List[List[STEP_OUTPUT]]) -> None:
"""Runs ``on_train_epoch_end hook``."""
# We cannot rely on Trainer.call_hook because the signatures might be different across
# lightning module and callback
# As a result, we need to inspect if the module accepts `outputs` in `on_train_epoch_end`
# This implementation is copied from Trainer.call_hook
hook_name = "on_train_epoch_end"
prev_fx_name = self.trainer.lightning_module._current_fx_name
self.trainer.lightning_module._current_fx_name = hook_name
# always profile hooks
with self.trainer.profiler.profile(hook_name):
# first call trainer hook
if hasattr(self.trainer, hook_name):
trainer_hook = getattr(self.trainer, hook_name)
trainer_hook(processed_epoch_output)
# next call hook in lightningModule
model_ref = self.trainer.lightning_module
if is_overridden(hook_name, model_ref):
hook_fx = getattr(model_ref, hook_name)
if is_param_in_hook_signature(hook_fx, "outputs"):
self._warning_cache.deprecation(
"The signature of `ModelHooks.on_train_epoch_end` has changed in v1.3."
" `outputs` parameter has been deprecated."
" Support for the old signature will be removed in v1.5"
)
model_ref.on_train_epoch_end(processed_epoch_output)
else:
model_ref.on_train_epoch_end()
# call the accelerator hook
if hasattr(self.trainer.accelerator, hook_name):
accelerator_hook = getattr(self.trainer.accelerator, hook_name)
accelerator_hook()
# restore current_fx when nested context
self.trainer.lightning_module._current_fx_name = prev_fx_name
def _num_training_batches_reached(self,
is_last_batch: bool = False) -> bool:
"""Checks if we are in the last batch or if there are more batches to follow."""
# TODO: Can we combine this with training_batch_loop's arg that does a similar check?
return self.batches_seen == self.trainer.num_training_batches or is_last_batch
def _track_epoch_end_reduce_metrics(
self, epoch_output: List[List[STEP_OUTPUT]],
batch_end_outputs: STEP_OUTPUT) -> None:
"""Adds the batch outputs to the epoch outputs and prepares reduction"""
hook_overridden = self._should_add_batch_output_to_epoch_output()
if not hook_overridden:
return
# track the outputs to reduce at the end of the epoch
for opt_idx, opt_outputs in enumerate(batch_end_outputs):
# with 1 step (no tbptt) don't use a sequence at epoch end
if (isinstance(opt_outputs, list) and len(opt_outputs) == 1
and not isinstance(opt_outputs[0], ResultCollection)):
opt_outputs = opt_outputs[0]
epoch_output[opt_idx].append(opt_outputs)
def _should_add_batch_output_to_epoch_output(self) -> bool:
"""
We add to the epoch outputs if
1. The model defines training_epoch_end OR
2. The model overrides on_train_epoch_end which has `outputs` in the signature
"""
# TODO: in v1.5 this only needs to check if training_epoch_end is overridden
lightning_module = self.trainer.lightning_module
if is_overridden("training_epoch_end", lightning_module):
return True
if is_overridden("on_train_epoch_end", lightning_module):
model_hook_fx = getattr(lightning_module, "on_train_epoch_end")
if is_param_in_hook_signature(model_hook_fx, "outputs"):
return True
return False
@staticmethod
def _prepare_outputs(
outputs: List[List[List["ResultCollection"]]], batch_mode: bool
) -> Union[List[List[List[Dict]]], List[List[Dict]], List[Dict], Dict]:
"""
Extract required information from batch or epoch end results.
Args:
outputs: A 3-dimensional list of ``ResultCollection`` objects with dimensions:
``[optimizer outs][batch outs][tbptt steps]``.
batch_mode: If True, ignore the batch output dimension.
Returns:
The cleaned outputs with ``ResultCollection`` objects converted to dictionaries.
All list dimensions of size one will be collapsed.
"""
processed_outputs = []
for opt_outputs in outputs:
# handle an edge case where an optimizer output is the empty list
if len(opt_outputs) == 0:
continue
processed_batch_outputs = []
if batch_mode:
opt_outputs = [opt_outputs]
for batch_outputs in opt_outputs:
processed_tbptt_outputs = []
if isinstance(batch_outputs, ResultCollection):
batch_outputs = [batch_outputs]
for tbptt_output in batch_outputs:
out = tbptt_output.extra
if tbptt_output.minimize is not None:
out["loss"] = tbptt_output.minimize.detach()
processed_tbptt_outputs.append(out)
# if there was only one tbptt step then we can collapse that dimension
if len(processed_tbptt_outputs) == 1:
processed_tbptt_outputs = processed_tbptt_outputs[0]
processed_batch_outputs.append(processed_tbptt_outputs)
# batch_outputs should be just one dict (or a list of dicts if using tbptt) per optimizer
if batch_mode:
processed_batch_outputs = processed_batch_outputs[0]
processed_outputs.append(processed_batch_outputs)
# if there is only one optimiser then we collapse that dimension
if len(processed_outputs) == 1:
processed_outputs = processed_outputs[0]
return processed_outputs
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.batch_loop.should_accumulate():
return
self.trainer.optimizer_connector.update_learning_rates(
interval=interval,
update_plateau_schedulers=update_plateau_schedulers,
opt_indices=[
opt_idx for opt_idx, _ in
self.batch_loop.get_active_optimizers(self.total_batch_idx)
],
)
def _increment_accumulated_grad_global_step(self) -> None:
"""increments global step"""
num_accumulated_batches_reached = self.batch_loop._accumulated_batches_reached(
)
num_training_batches_reached = self._num_training_batches_reached()
# progress global step according to grads progress
if num_accumulated_batches_reached or num_training_batches_reached:
self.global_step = self.trainer.accelerator.update_global_step(
self.total_batch_idx, self.trainer.global_step)
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()
