# Copyright The PyTorch Lightning team. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. """ The lightning training loop handles everything except the actual computations of your model. To decide what will happen in your training loop, define the `training_step` function. Below are all the things lightning automates for you in the training loop. Accumulated gradients --------------------- Accumulated gradients runs K small batches of size N before doing a backwards pass. The effect is a large effective batch size of size KxN. .. code-block:: python # DEFAULT (ie: no accumulated grads) trainer = Trainer(accumulate_grad_batches=1) Force training for min or max epochs ------------------------------------ It can be useful to force training for a minimum number of epochs or limit to a max number .. code-block:: python # DEFAULT trainer = Trainer(min_epochs=1, max_epochs=1000) Force disable early stop ------------------------ To disable early stopping pass None to the early_stop_callback .. code-block:: python # DEFAULT trainer = Trainer(early_stop_callback=None) Gradient Clipping ----------------- Gradient clipping may be enabled to avoid exploding gradients. Specifically, this will `clip the gradient norm computed over all model parameters `together `_. .. code-block:: python # DEFAULT (ie: don't clip) trainer = Trainer(gradient_clip_val=0) # clip gradients with norm above 0.5 trainer = Trainer(gradient_clip_val=0.5) Inspect gradient norms ---------------------- Looking at grad norms can help you figure out where training might be going wrong. .. code-block:: python # DEFAULT (-1 doesn't track norms) trainer = Trainer(track_grad_norm=-1) # track the LP norm (P=2 here) trainer = Trainer(track_grad_norm=2) Set how much of the training set to check ----------------------------------------- If you don't want to check 100% of the training set (for debugging or if it's huge), set this flag. limit_train_batches will be overwritten by overfit_batches if `overfit_batches > 0` .. code-block:: python # DEFAULT trainer = Trainer(limit_train_batches=1.0) # check 10% only trainer = Trainer(limit_train_batches=0.1) # check 10 batches only trainer = Trainer(limit_train_batches=10) Packed sequences as inputs -------------------------- When using PackedSequence, do 2 things: 1. return either a padded tensor in dataset or a list of variable length tensors in the dataloader collate_fn (example above shows the list implementation). 2. Pack the sequence in forward or training and validation steps depending on use case. .. code-block:: python # For use in dataloader def collate_fn(batch): x = [item[0] for item in batch] y = [item[1] for item in batch] return x, y # In module def training_step(self, batch, batch_idx): x = rnn.pack_sequence(batch[0], enforce_sorted=False) y = rnn.pack_sequence(batch[1], enforce_sorted=False) Truncated Backpropagation Through Time -------------------------------------- There are times when multiple backwards passes are needed for each batch. For example, it may save memory to use Truncated Backpropagation Through Time when training RNNs. When this flag is enabled each batch is split into sequences of size truncated_bptt_steps and passed to training_step(...) separately. A default splitting function is provided, however, you can override it for more flexibility. See `tbptt_split_batch`. .. code-block:: python # DEFAULT (single backwards pass per batch) trainer = Trainer(truncated_bptt_steps=None) # (split batch into sequences of size 2) trainer = Trainer(truncated_bptt_steps=2) NaN detection and intervention ------------------------------ When the `terminate_on_nan` flag is enabled, after every forward pass during training, Lightning will check that 1. the loss you return in `training_step` is finite (not NaN and not +/-inf) 2. the model parameters have finite values. Lightning will terminate the training loop with an error message if NaN or infinite values are detected. If this happens, you should investigate numerically unstable operations in your model. .. code-block:: python # DEFAULT (won't perform the NaN check) trainer = Trainer(terminate_on_nan=False) # (NaN check each batch and terminate on NaN or infinite values) trainer = Trainer(terminate_on_nan=True) """ import subprocess from abc import ABC, abstractmethod from copy import copy from typing import Callable from typing import Union, List import numpy as np import torch import torch.distributed as torch_distrib from torch.utils.data import DataLoader from copy import deepcopy from pytorch_lightning import _logger as log from pytorch_lightning.callbacks import ModelCheckpoint from pytorch_lightning.callbacks.base import Callback from pytorch_lightning.core.lightning import LightningModule from pytorch_lightning.core.step_result import EvalResult, Result from pytorch_lightning.loggers import LightningLoggerBase from pytorch_lightning.trainer.states import TrainerState from pytorch_lightning.trainer.supporters import TensorRunningAccum, Accumulator from pytorch_lightning.utilities import rank_zero_warn, AMPType from pytorch_lightning.utilities.exceptions import MisconfigurationException from pytorch_lightning.utilities.memory import recursive_detach from pytorch_lightning.utilities.parsing import AttributeDict try: from apex import amp except ImportError: amp = None try: import torch_xla.distributed.parallel_loader as xla_pl import torch_xla.core.xla_model as xm except ImportError: XLA_AVAILABLE = False else: XLA_AVAILABLE = True try: import horovod.torch as hvd except (ModuleNotFoundError, ImportError): HOROVOD_AVAILABLE = False else: HOROVOD_AVAILABLE = True # constant which signals should be catched for graceful trainer shutdown SIGNAL_TERMINATE = ('SIGTERM', 'SIGSEGV', 'SIGINT') class TrainerTrainLoopMixin(ABC): # this is just a summary on variables used in this abstract class, # the proper values/initialisation should be done in child class max_epochs: int min_epochs: int on_gpu: bool use_ddp: bool use_dp: bool use_ddp2: bool use_horovod: bool use_single_gpu: bool use_tpu: bool data_parallel_device_ids: ... check_val_every_n_epoch: ... num_training_batches: int val_check_batch: ... disable_validation: bool fast_dev_run: ... accumulation_scheduler: ... lr_schedulers: ... early_stop_callback: ... callback_metrics: ... logger: Union[LightningLoggerBase, bool] global_step: int testing: bool log_save_interval: float global_rank: int row_log_interval: float truncated_bptt_steps: ... optimizers: ... optimizer_frequencies: ... accumulate_grad_batches: int track_grad_norm: ... model: LightningModule interrupted: bool running_loss: ... progress_bar_dict: ... reduce_lr_on_plateau_scheduler: ... profiler: ... batch_idx: int precision: ... train_dataloader: DataLoader reload_dataloaders_every_epoch: bool max_steps: int min_steps: int total_batch_idx: int terminate_on_nan: bool tpu_id: int interactive_ddp_procs: ... state: TrainerState amp_backend: AMPType on_tpu: bool # Callback system callbacks: List[Callback] on_train_start: Callable on_train_end: Callable on_batch_start: Callable on_batch_end: Callable on_train_batch_start: Callable on_train_batch_end: Callable on_epoch_start: Callable on_epoch_end: Callable on_validation_end: Callable on_keyboard_interrupt: Callable on_train_epoch_start: Callable on_train_epoch_end: Callable @abstractmethod def get_model(self) -> LightningModule: """Warning: this is just empty shell for code implemented in other class.""" @abstractmethod def is_function_implemented(self, *args, **kwargs): """Warning: this is just empty shell for code implemented in other class.""" @abstractmethod def run_evaluation(self, *args, **kwargs): """Warning: this is just empty shell for code implemented in other class.""" @abstractmethod def transfer_batch_to_gpu(self, *args): """Warning: this is just empty shell for code implemented in other class.""" @abstractmethod def transfer_batch_to_tpu(self, *args): """Warning: this is just empty shell for code implemented in other class.""" @abstractmethod def clip_gradients(self, *args): """Warning: this is just empty shell for code implemented in other class.""" @abstractmethod def detect_nan_tensors(self, *args): """Warning: this is just empty shell for code implemented in other class.""" @abstractmethod def is_overridden(self, *args): """Warning: this is just empty shell for code implemented in other class.""" @abstractmethod def add_progress_bar_metrics(self, *args): """Warning: this is just empty shell for code implemented in other class.""" @abstractmethod def log_metrics(self, *args): """Warning: this is just empty shell for code implemented in other class.""" @abstractmethod def process_output(self, *args): """Warning: this is just empty shell for code implemented in other class.""" @abstractmethod def reset_train_dataloader(self, *args): """Warning: this is just empty shell for code implemented in other class.""" @abstractmethod def reset_val_dataloader(self, model): """Warning: this is just empty shell for code implemented in other class.""" @abstractmethod def has_arg(self, *args): """Warning: this is just empty shell for code implemented in other class.""" def train(self): # add signal handlers for process kills # def _signal_kill_handler(*args): # return TrainerTrainLoopMixin.run_training_teardown(self) # # orig_signal_handlers = {} # for sig_name in SIGNAL_TERMINATE: # orig_signal_handlers[sig_name] = signal.signal(getattr(signal, sig_name), # _signal_kill_handler) # get model model = self.get_model() # enable train mode model.train() # enable gradients torch.set_grad_enabled(True) # load data # if reload_dataloaders_every_epoch, this is moved to the epoch loop if not self.reload_dataloaders_every_epoch: self.reset_train_dataloader(model) if model.val_dataloader is not None: self.reset_val_dataloader(model) # Train start events with self.profiler.profile('on_train_start'): # callbacks self.on_train_start() # model hooks model.on_train_start() try: # run all epochs for epoch in range(self.current_epoch, self.max_epochs): # reset train dataloader if self.reload_dataloaders_every_epoch: self.reset_train_dataloader(model) # set seed for distributed sampler (enables shuffling for each epoch) if (self.use_ddp or self.use_horovod or self.on_tpu) \ and hasattr(self.train_dataloader, 'sampler') \ and hasattr(self.train_dataloader.sampler, 'set_epoch'): self.train_dataloader.sampler.set_epoch(epoch) # update training progress in trainer and model model.current_epoch = epoch self.current_epoch = epoch # changing gradient according accumulation_scheduler self.accumulation_scheduler.on_epoch_start(self, self.get_model()) # stores accumulated grad fractions per batch self.batch_loss_value = TensorRunningAccum( window_length=self.accumulate_grad_batches ) # ----------------- # RUN TNG EPOCH # ----------------- self.run_training_epoch() if self.max_steps and self.max_steps <= self.global_step: self.run_training_teardown() return # update LR schedulers self.update_learning_rates(interval='epoch') # early stopping met_min_epochs = epoch >= self.min_epochs - 1 met_min_steps = self.global_step >= self.min_steps if self.min_steps else True if self.should_stop: if (met_min_epochs and met_min_steps): self.run_training_teardown() return else: log.info('Trainer was signaled to stop but required minimum epochs' f' ({self.min_epochs}) or minimum steps ({self.min_steps}) has' ' not been met. Training will continue...') self.run_training_teardown() except KeyboardInterrupt: rank_zero_warn('Detected KeyboardInterrupt, attempting graceful shutdown...') # user could press ctrl+c many times... only shutdown once if not self.interrupted: self.interrupted = True self.state = TrainerState.INTERRUPTED self.on_keyboard_interrupt() self.run_training_teardown() def prepare_train_loop_dataloader(self, train_dataloader): # on TPU we have to wrap it under the ParallelLoader if self.use_tpu: device = xm.xla_device(self.tpu_id) train_dataloader = xla_pl.ParallelLoader(train_dataloader, [device]) train_dataloader = train_dataloader.per_device_loader(device) return train_dataloader def run_on_epoch_start_hook(self, model): # Epoch start events with self.profiler.profile('on_epoch_start'): # callbacks self.on_epoch_start() # model hooks if self.is_function_implemented('on_epoch_start'): model.on_epoch_start() # Epoch start events with self.profiler.profile('on_train_epoch_start'): # callbacks self.on_train_epoch_start() # model hooks if self.is_function_implemented('on_train_epoch_start'): model.on_train_epoch_start() def run_training_epoch(self): # get model model = self.get_model() # Epoch start events self.run_on_epoch_start_hook(model) # modify dataloader if needed (ddp, etc...) train_dataloader = self.prepare_train_loop_dataloader(self.train_dataloader) # bookkeeping num_optimizers = len(self._get_optimizers_iterable()) epoch_output = [[] for _ in range(num_optimizers)] should_check_val = False # structured result accumulators for callbacks early_stopping_accumulator = Accumulator() checkpoint_accumulator = Accumulator() # run epoch for batch_idx, (batch, is_last_batch) in self.profiler.profile_iterable( enumerate(_with_is_last(train_dataloader)), "get_train_batch" ): # stop epoch if we limited the number of training batches if batch_idx >= self.num_training_batches: break self.batch_idx = batch_idx model.global_step = self.global_step # ------------------------------------ # TRAINING_STEP + TRAINING_STEP_END # ------------------------------------ batch_output = self.run_training_batch(batch, batch_idx) # only track outputs when user implements training_epoch_end # otherwise we will build up unnecessary memory epoch_end_outputs = self.process_train_step_outputs( batch_output.training_step_output_for_epoch_end, early_stopping_accumulator, checkpoint_accumulator ) # track the outputs to reduce at the end of the epoch for opt_idx, opt_outputs in enumerate(epoch_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], Result): opt_outputs = opt_outputs[0] epoch_output[opt_idx].append(opt_outputs) # when returning -1 from train_step, we end epoch early self.should_stop = batch_output.signal == -1 # ----------------------------------------- # VALIDATE IF NEEDED + CHECKPOINT CALLBACK # ----------------------------------------- should_check_val = self.should_check_val(batch_idx, is_last_batch) if should_check_val: self.run_evaluation(test_mode=False) # ----------------------------------------- # SAVE LOGGERS (ie: Tensorboard, etc...) # ----------------------------------------- self.save_loggers_in_training_loop(batch_idx) # ----------------------------------------- # SAVE METRICS TO LOGGERS # ----------------------------------------- self.save_train_loop_metrics_to_loggers(batch_idx, batch_output) # update LR schedulers monitor_metrics = deepcopy(self.callback_metrics) monitor_metrics.update(batch_output.batch_log_metrics) self.update_train_loop_lr_schedulers(monitor_metrics=monitor_metrics) # progress global step according to grads progress self.increment_accumulated_grad_global_step() # max steps reached, end training if self.max_steps is not None and self.max_steps == self.global_step: break # end epoch early # stop when the flag is changed or we've gone past the amount # requested in the batches if self.should_stop: break # let ddp devices catch up when using horovod self.sync_horovod() # process epoch outputs self.run_training_epoch_end(epoch_output, checkpoint_accumulator, early_stopping_accumulator, num_optimizers) # checkpoint callback self.check_checkpoint_callback(should_check_val) # epoch end hook self.run_on_epoch_end_hook(model) def process_train_step_outputs(self, all_train_step_outputs, early_stopping_accumulator, checkpoint_accumulator): """ Figure out what needs to be tracked/logged at the end of the epoch """ # the training step outputs a list per optimizer. The list contains the outputs at each time step # when no TBPTT is used, then the list has 1 item per batch # when TBPTT IS used, then the list has n items (1 per time step) epoch_end_outputs = [] for optimizer_idx_outputs in all_train_step_outputs: # extract one representative sample from each time step (1 if no tbptt) and 0th optimizer sample_output = optimizer_idx_outputs[-1] # pull out callback info if available (ie: Results object) if isinstance(sample_output, dict) and 'early_stop_on' in sample_output: early_stopping_accumulator.accumulate(sample_output['early_stop_on']) if isinstance(sample_output, dict) and 'checkpoint_on' in sample_output: checkpoint_accumulator.accumulate(sample_output['checkpoint_on']) # decide if we need to reduce at the end of the epoch automatically auto_reduce_tng_result = isinstance(sample_output, Result) and sample_output.should_reduce_on_epoch_end # only track when a) it needs to be autoreduced OR b) the user wants to manually reduce on epoch end if self.is_overridden('training_epoch_end', model=self.get_model()) or auto_reduce_tng_result: epoch_end_outputs.append(optimizer_idx_outputs) return epoch_end_outputs def check_checkpoint_callback(self, should_check_val): # when no val loop is present or fast-dev-run still need to call checkpoints # TODO bake this logic into the checkpoint callback should_activate = not self.is_overridden('validation_step') and not should_check_val if should_activate: checkpoint_callbacks = [c for c in self.callbacks if isinstance(c, ModelCheckpoint)] [c.on_validation_end(self, self.get_model()) for c in checkpoint_callbacks] def update_train_loop_lr_schedulers(self, monitor_metrics=None): if ((self.batch_idx + 1) % self.accumulate_grad_batches == 0 or (self.batch_idx + 1) == self.num_training_batches): # update lr self.update_learning_rates(interval='step', monitor_metrics=monitor_metrics) def run_on_epoch_end_hook(self, model): with self.profiler.profile('on_epoch_end'): # callbacks self.on_epoch_end() # model hooks if self.is_function_implemented('on_epoch_end'): model.on_epoch_end() with self.profiler.profile('on_train_epoch_end'): # callbacks self.on_train_epoch_end() # model hooks if self.is_function_implemented('on_train_epoch_end'): model.on_train_epoch_end() def run_training_epoch_end(self, epoch_output, checkpoint_accumulator, early_stopping_accumulator, num_optimizers): # epoch output is a list. Each item in that list has all the outputs per optimizer # epoch_output[optimizer_idx][training_step_idx][tbptt_index] # remember that not using truncated backprop is equivalent with truncated back prop of len(1) model = self.get_model() epoch_log_metrics = {} epoch_callback_metrics = {} epoch_progress_bar_metrics = {} # ----------------------- # Calculate epoch callback values if given # ----------------------- if checkpoint_accumulator.num_values > 0: epoch_callback_metrics['checkpoint_on'] = checkpoint_accumulator.mean() if early_stopping_accumulator.num_values > 0: epoch_callback_metrics['early_stop_on'] = early_stopping_accumulator.mean() # ------------------------ # determine if using a result obj # ------------------------ # [optimizer_idx][training_step_idx][tbptt_index] opt_idx_outputs = epoch_output[0] try: sample_obj = opt_idx_outputs[0][0] if isinstance(opt_idx_outputs[0], list) else opt_idx_outputs[0] is_result_obj = len(epoch_output) > 0 and isinstance(sample_obj, Result) except IndexError as e: is_result_obj = False # -------------------------- # EPOCH END STEP IF DEFINED # -------------------------- if self.is_overridden('training_epoch_end', model=model): self.global_step += 1 if is_result_obj: # with result object gather across time and training steps so each opt idx has a single result obj epoch_output = self.__gather_result_across_time_and_optimizers(epoch_output) if num_optimizers == 1: epoch_output = epoch_output[0] # run training_epoch_end # a list with a result per optimizer index epoch_output = model.training_epoch_end(epoch_output) if isinstance(epoch_output, Result): epoch_log_metrics = epoch_output.epoch_log_metrics epoch_progress_bar_metrics = epoch_output.epoch_pbar_metrics else: _processed_outputs = self.process_output(epoch_output) epoch_progress_bar_metrics = _processed_outputs[1] epoch_log_metrics = _processed_outputs[2] epoch_callback_metrics = _processed_outputs[3] # -------------------------- # Structured Result (auto epoch end) # -------------------------- elif is_result_obj: epoch_log_metrics, epoch_progress_bar_metrics = self.__auto_reduce_results_on_epoch_end(epoch_output) # -------------------------- # track results # -------------------------- # add the metrics to the loggers if epoch_log_metrics and len(epoch_log_metrics) > 0: self.log_metrics(epoch_log_metrics, {}) # add metrics to callbacks self.callback_metrics.update(epoch_callback_metrics) # add metrics to progress_bar if len(epoch_progress_bar_metrics) > 0: self.add_progress_bar_metrics(epoch_progress_bar_metrics) def __auto_reduce_results_on_epoch_end(self, epoch_output): epoch_log_metrics = {} epoch_progress_bar_metrics = {} for opt_outputs in epoch_output: # reduce across time first time_reduced_outputs = [] for train_step_idx in range(len(opt_outputs)): tbptt_outs = opt_outputs[train_step_idx] tbptt_outs = tbptt_outs[0].__class__.reduce_across_time(tbptt_outs) time_reduced_outputs.append(tbptt_outs) # reduce across training steps opt_outputs = time_reduced_outputs[0].__class__.reduce_on_epoch_end(time_reduced_outputs) opt_outputs.minimize = opt_outputs.minimize.mean() epoch_log_metrics.update(opt_outputs.epoch_log_metrics) epoch_progress_bar_metrics.update(opt_outputs.epoch_pbar_metrics) return epoch_log_metrics, epoch_progress_bar_metrics def __gather_result_across_time_and_optimizers(self, epoch_output): """ Gather results into a single padded tensor per metric where each tensor is gathered across time and across time steps. Returns: a list where each element is a Result with the tensors gathered """ gathered_epoch_outputs = [] for opt_outputs in epoch_output: # gather across time first time_gathered_outputs = [] for train_step_idx in range(len(opt_outputs)): tbptt_outs = opt_outputs[train_step_idx] tbptt_outs = tbptt_outs[0].__class__.gather(tbptt_outs) time_gathered_outputs.append(tbptt_outs) # gather across training steps # each metric has dimensions (training_steps, seq_len) (seq_len=1 when no tbptt is used) gathered_opt_output = time_gathered_outputs[0].__class__.padded_gather(time_gathered_outputs) gathered_epoch_outputs.append(gathered_opt_output) return gathered_epoch_outputs def sync_horovod(self): if self.use_horovod: hvd.join(hvd.local_rank() if self.on_gpu else -1) def increment_accumulated_grad_global_step(self): # progress global step according to grads progress if ((self.batch_idx + 1) % self.accumulate_grad_batches == 0 or (self.batch_idx + 1) == self.num_training_batches): self.global_step += 1 self.total_batch_idx += 1 def save_train_loop_metrics_to_loggers(self, batch_idx, batch_output): # when metrics should be logged should_log_metrics = (batch_idx + 1) % self.row_log_interval == 0 or self.should_stop if should_log_metrics or self.fast_dev_run: # logs user requested information to logger metrics = batch_output.batch_log_metrics grad_norm_dic = batch_output.grad_norm_dic if len(metrics) > 0 or len(grad_norm_dic) > 0: self.log_metrics(metrics, grad_norm_dic) def save_loggers_in_training_loop(self, batch_idx): # when loggers should save to disk should_save_log = (batch_idx + 1) % self.log_save_interval == 0 or self.should_stop if should_save_log or self.fast_dev_run: if self.is_global_zero and self.logger is not None: self.logger.save() def should_check_val(self, batch_idx, is_last_batch): # decide if we should run validation is_val_check_batch = (batch_idx + 1) % self.val_check_batch == 0 can_check_epoch = (self.current_epoch + 1) % self.check_val_every_n_epoch == 0 can_check_val = self.enable_validation and can_check_epoch should_check_val = is_val_check_batch or self.should_stop is_last_batch_for_infinite_dataset = (is_last_batch and self.val_check_batch == float('inf')) should_check_val = can_check_val and (should_check_val or is_last_batch_for_infinite_dataset) return should_check_val def run_training_batch(self, batch, batch_idx): # track grad norms grad_norm_dic = {} # track all metrics for callbacks batch_callback_metrics = [] # track metrics to log batch_log_metrics = [] using_results_obj = False # track all outputs across time and num of optimizers batch_outputs = [[] for i in range(len(self._get_optimizers_iterable()))] if batch is None: return AttributeDict(signal=0, grad_norm_dic=grad_norm_dic) # Batch start events # TODO: deprecate 1.0 with self.profiler.profile('on_batch_start'): # callbacks self.on_batch_start() # hooks if self.is_function_implemented('on_batch_start'): response = self.get_model().on_batch_start(batch) if response == -1: return AttributeDict(signal=-1, grad_norm_dic=grad_norm_dic) with self.profiler.profile('on_train_batch_start'): # forward support for multiple loaders dataloader_idx = 0 self.on_train_batch_start(batch, batch_idx, dataloader_idx) # hooks if self.is_function_implemented('on_train_batch_start'): response = self.get_model().on_train_batch_start(batch, batch_idx, dataloader_idx) if response == -1: return AttributeDict(signal=-1, grad_norm_dic=grad_norm_dic) splits = [batch] if self.truncated_bptt_steps is not None: model_ref = self.get_model() with self.profiler.profile('tbptt_split_batch'): splits = model_ref.tbptt_split_batch(batch, self.truncated_bptt_steps) self.hiddens = None for split_idx, split_batch in enumerate(splits): self.split_idx = split_idx for opt_idx, optimizer in self._get_optimizers_iterable(): # make sure only the gradients of the current optimizer's parameters are calculated # in the training step to prevent dangling gradients in multiple-optimizer setup. if len(self.optimizers) > 1: for param in self.get_model().parameters(): param.requires_grad = False for group in optimizer.param_groups: for param in group['params']: param.requires_grad = True # ------------------- # calculate loss (train step + train step end) # ------------------- opt_closure_result = self.optimizer_closure( split_batch, batch_idx, opt_idx, optimizer, self.hiddens ) using_results_obj = isinstance(opt_closure_result.training_step_output, Result) # ------------------------------ # POST forward bookkeeping # ------------------------------ batch_callback_metrics.append(opt_closure_result.training_step_output.callback_metrics) # add metrics to loggers if using_results_obj: metrics_to_log = opt_closure_result.training_step_output.batch_log_metrics step_pbar_metrics = opt_closure_result.training_step_output.batch_pbar_metrics else: metrics_to_log = opt_closure_result.training_step_output.log_metrics step_pbar_metrics = opt_closure_result.training_step_output.pbar_on_batch_end # track metrics batch_log_metrics.append(metrics_to_log) if len(step_pbar_metrics) > 0: self.add_progress_bar_metrics(step_pbar_metrics) # track hiddens self.hiddens = opt_closure_result.hiddens if using_results_obj: opt_closure_result.training_step_output_for_epoch_end.drop_hiddens() # check if loss or model weights are nan if self.terminate_on_nan: self.detect_nan_tensors(opt_closure_result.loss) # track total loss for logging (avoid mem leaks) self.batch_loss_value.append(opt_closure_result.loss) # track all the outputs across all steps batch_outputs[opt_idx].append(opt_closure_result.training_step_output_for_epoch_end) # ------------------------------ # BACKWARD PASS # ------------------------------ # gradient update with accumulated gradients if ((self.batch_idx + 1) % self.accumulate_grad_batches == 0 or (self.batch_idx + 1) == self.num_training_batches): # backward grad_norm_dic = self.run_batch_backward_pass(split_batch, batch_idx, opt_idx, optimizer) # calculate running loss for display self.running_loss.append(self.batch_loss_value.mean() * self.accumulate_grad_batches) # reset for next set of accumulated grads self.batch_loss_value.reset() # Batch end events with self.profiler.profile('on_batch_end'): # callbacks self.on_batch_end() # model hooks if self.is_function_implemented('on_batch_end'): self.get_model().on_batch_end() with self.profiler.profile('on_train_batch_end'): # forward support for multiple loaders dataloader_idx = 0 self.on_train_batch_end(batch, batch_idx, dataloader_idx) # model hooks if self.is_function_implemented('on_train_batch_end'): self.get_model().on_train_batch_end(batch, batch_idx, dataloader_idx) # collapse all metrics into one dict batch_log_metrics = {k: v for d in batch_log_metrics for k, v in d.items()} # track all metrics for callbacks if not using_results_obj: self.callback_metrics.update({k: v for d in batch_callback_metrics for k, v in d.items()}) result = AttributeDict( signal=0, grad_norm_dic=grad_norm_dic, batch_log_metrics=batch_log_metrics, training_step_output_for_epoch_end=batch_outputs ) return result def run_batch_backward_pass(self, split_batch, batch_idx, opt_idx, optimizer): # ------------------ # GRAD NORMS # ------------------ # track gradient norms when requested grad_norm_dic = {} if batch_idx % self.row_log_interval == 0: if float(self.track_grad_norm) > 0: model = self.get_model() grad_norm_dic = model.grad_norm( self.track_grad_norm) # ------------------ # CLIP GRADS # ------------------ if self.amp_backend == AMPType.NATIVE and not self.use_tpu: self.scaler.unscale_(optimizer) self.clip_gradients(optimizer) # ------------------ # .STEP + ZERO_GRAD # ------------------ self.call_optimizer_step(optimizer, opt_idx, batch_idx, split_batch) return grad_norm_dic def call_optimizer_step(self, optimizer, opt_idx, batch_idx, split_batch): # calls .step(), .zero_grad() # override function to modify this behavior model = self.get_model() with self.profiler.profile('optimizer_step'): lambda_closure = lambda: self.optimizer_closure( split_batch, batch_idx, opt_idx, optimizer, self.hiddens, ).loss # apply TPU optimizer if self.use_tpu and XLA_AVAILABLE: model.optimizer_step(self.current_epoch, batch_idx, optimizer, opt_idx, lambda_closure, on_tpu=True) # for LBFGS do something a bit different elif isinstance(optimizer, torch.optim.LBFGS): # native amp + lbfgs is a no go right now if self.amp_backend == AMPType.NATIVE: raise MisconfigurationException( 'native PyTorch amp and lbfgs are not compatible.' ' To request, please file a Github issue in PyTorch and tag @mcarilli') model.optimizer_step(self.current_epoch, batch_idx, optimizer, opt_idx, lambda_closure, using_lbfgs=True) # when using 16-bit else: native_amp = self.amp_backend == AMPType.NATIVE model.optimizer_step(self.current_epoch, batch_idx, optimizer, opt_idx, lambda_closure, using_native_amp=native_amp) # in native 16-bit we need to update scaler after optimizer step if self.amp_backend == AMPType.NATIVE and not self.use_tpu: self.scaler.update() # model hook model.on_before_zero_grad(optimizer) # clear gradients model.optimizer_zero_grad(self.current_epoch, batch_idx, optimizer, opt_idx) def optimizer_closure(self, split_batch, batch_idx, opt_idx, optimizer, hiddens): """ wrap the forward step in a closure so second order methods work """ # --------------------------- # FORWARD (TRAINING STEP + TRAIN STEP END) # --------------------------- with self.profiler.profile('model_forward'): if self.amp_backend == AMPType.NATIVE and not self.use_tpu: with torch.cuda.amp.autocast(): training_step_output = self.training_forward(split_batch, batch_idx, opt_idx, hiddens) else: training_step_output = self.training_forward(split_batch, batch_idx, opt_idx, hiddens) # ---------------------------- # PROCESS THE RESULT # ---------------------------- # format and reduce outputs accordingly training_step_output_for_epoch_end = training_step_output is_result_obj = isinstance(training_step_output, Result) # track batch size for weighted average if is_result_obj: training_step_output.track_batch_size(len(split_batch)) # don't allow EvalResult in the training_step if isinstance(training_step_output, EvalResult): raise MisconfigurationException('training_step cannot return EvalResult, ' 'use a dict or TrainResult instead') # handle regular dicts if not is_result_obj: training_step_output = self.process_output(training_step_output, train=True) training_step_output = AttributeDict( batch_loss=training_step_output[0], pbar_on_batch_end=training_step_output[1], log_metrics=training_step_output[2], callback_metrics=training_step_output[3], hiddens=training_step_output[4], ) # if the user decides to finally reduce things in epoch_end, save raw output without graphs if isinstance(training_step_output_for_epoch_end, torch.Tensor): training_step_output_for_epoch_end = training_step_output_for_epoch_end.detach() elif is_result_obj: training_step_output_for_epoch_end = copy(training_step_output) training_step_output_for_epoch_end.detach() else: training_step_output_for_epoch_end = recursive_detach(training_step_output_for_epoch_end) # accumulate loss # (if accumulate_grad_batches = 1 no effect) closure_loss = training_step_output.minimize if is_result_obj else training_step_output.batch_loss closure_loss = closure_loss / self.accumulate_grad_batches # the loss will get scaled for amp. avoid any modifications to it untouched_loss = closure_loss.detach().clone() # backward pass model_ref = self.get_model() with self.profiler.profile('model_backward'): # scale loss for 16 bit if self.precision == 16 and not self.on_tpu: closure_loss = model_ref.amp_scale_loss(closure_loss, optimizer, opt_idx, amp_backend=self.amp_backend) # enter amp context if self.amp_backend == AMPType.APEX: self.dev_debugger.track_event('AMP', str(AMPType.APEX)) context = closure_loss closure_loss = closure_loss.__enter__() # do backward pass model_ref.backward(self, closure_loss, optimizer, opt_idx) # exit amp context if self.precision == 16 and self.amp_backend == AMPType.APEX and not self.on_tpu: a, b, c = None, None, None error = context.__exit__(a, b, c) if error: rank_zero_warn(a, b, c) raise Exception('apex unscale error') # once backward has been applied, release graph closure_loss = closure_loss.detach() if is_result_obj: training_step_output.detach() else: training_step_output.batch_loss = training_step_output.batch_loss.detach() if self.use_horovod: # Synchronize Horovod to ensure gradient manipulations (e.g., loss scaling) are valid optimizer.synchronize() # insert after step hook if self.is_function_implemented('on_after_backward'): model_ref = self.get_model() with self.profiler.profile('on_after_backward'): model_ref.on_after_backward() # when in dev debugging track the losses self.dev_debugger.track_train_loss_history(batch_idx, untouched_loss.detach()) result = AttributeDict( loss=untouched_loss, training_step_output=training_step_output, training_step_output_for_epoch_end=training_step_output_for_epoch_end, hiddens=training_step_output.hiddens, ) return result def _get_optimizers_iterable(self): if not self.optimizer_frequencies: # call training_step once per optimizer return list(enumerate(self.optimizers)) optimizer_freq_cumsum = np.cumsum(self.optimizer_frequencies) optimizers_loop_length = optimizer_freq_cumsum[-1] current_place_in_loop = self.total_batch_idx % optimizers_loop_length # find optimzier index by looking for the first {item > current_place} in the cumsum list opt_idx = np.argmax(optimizer_freq_cumsum > current_place_in_loop) return [(opt_idx, self.optimizers[opt_idx])] # @atexit.register def run_training_teardown(self): if hasattr(self, '_teardown_already_run') and self._teardown_already_run: return self._teardown_already_run = True # Train end events with self.profiler.profile('on_train_end'): # callbacks self.on_train_end() # model hooks if self.is_function_implemented('on_train_end'): self.get_model().on_train_end() if self.logger is not None: self.logger.finalize("success") # summarize profile results if self.global_rank == 0: self.profiler.describe() if self.global_rank == 0: for proc in self.interactive_ddp_procs: subprocess.Popen.kill(proc) # clean up dist group if self.use_ddp or self.use_ddp2: torch_distrib.destroy_process_group() # clear mem if self.on_gpu: model = self.get_model() model.cpu() torch.cuda.empty_cache() def training_forward(self, batch, batch_idx, opt_idx, hiddens): """ Handle forward for each training case (distributed, single gpu, etc...) :param batch: :param batch_idx: :return: """ # --------------- # FORWARD # --------------- # enable not needing to add opt_idx to training_step args = [batch, batch_idx] if len(self.optimizers) > 1: if self.has_arg('training_step', 'optimizer_idx'): args.append(opt_idx) else: num_opts = len(self.optimizers) raise ValueError( f'Your LightningModule defines {num_opts} optimizers but ' f'training_step is missing the "optimizer_idx" argument.' ) # pass hiddens if using tbptt if self.truncated_bptt_steps is not None: args.append(hiddens) # distributed forward if self.use_ddp or self.use_ddp2 or self.use_dp: output = self.model(*args) # Horovod elif self.use_horovod and self.on_gpu: batch = self.transfer_batch_to_gpu(batch, hvd.local_rank()) args[0] = batch output = self.model.training_step(*args) # single GPU forward elif self.use_single_gpu: gpu_id = 0 if isinstance(self.data_parallel_device_ids, list): gpu_id = self.data_parallel_device_ids[0] # Don't copy the batch since there is a single gpu that the batch could # be referenced from and if there are multiple optimizers the batch will # wind up copying it to the same device repeatedly. batch = self.transfer_batch_to_gpu(batch, gpu_id) args[0] = batch output = self.model.training_step(*args) # TPU support elif self.use_tpu: batch = self.transfer_batch_to_tpu(batch, self.tpu_id) args[0] = batch output = self.model.training_step(*args) # CPU forward else: output = self.model.training_step(*args) is_result_obj = isinstance(output, Result) # allow any mode to define training_step_end # do something will all the dp outputs (like softmax) if self.is_overridden('training_step_end'): model_ref = self.get_model() with self.profiler.profile('training_step_end'): # TODO: modify when using result obj output = model_ref.training_step_end(output) elif is_result_obj and (self.use_dp or self.use_ddp2): output.dp_reduce() # allow any mode to define training_end # TODO: remove in 1.0.0 if self.is_overridden('training_end'): model_ref = self.get_model() with self.profiler.profile('training_end'): output = model_ref.training_end(output) rank_zero_warn('`training_end` was deprecated in 0.7.0 and will be removed 1.0.0.' ' Use training_epoch_end instead', DeprecationWarning) return output def update_learning_rates(self, interval: str, monitor_metrics=None): """Update learning rates. Args: interval: either 'epoch' or 'step'. monitor_metrics: dict of possible values to monitor """ if not self.lr_schedulers: return for scheduler_idx, lr_scheduler in enumerate(self.lr_schedulers): current_idx = self.batch_idx if interval == 'step' else self.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: # If instance of ReduceLROnPlateau, we need to pass validation loss if lr_scheduler['reduce_on_plateau']: monitor_key = lr_scheduler['monitor'] if monitor_metrics is not None: monitor_val = monitor_metrics.get(monitor_key) else: monitor_val = self.callback_metrics.get(monitor_key) if monitor_val is None: avail_metrics = ','.join(list(self.callback_metrics.keys())) 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' ) if self.dev_debugger.enabled: old_lr = lr_scheduler['scheduler'].optimizer.param_groups[0]['lr'] # update LR lr_scheduler['scheduler'].step(monitor_val) if self.dev_debugger.enabled: new_lr = lr_scheduler['scheduler'].optimizer.param_groups[0]['lr'] self.dev_debugger.track_lr_schedulers_update( self.batch_idx, interval, scheduler_idx, old_lr, new_lr, monitor_key, ) else: if self.dev_debugger.enabled: old_lr = lr_scheduler['scheduler'].optimizer.param_groups[0]['lr'] # update LR lr_scheduler['scheduler'].step() if self.dev_debugger.enabled: new_lr = lr_scheduler['scheduler'].optimizer.param_groups[0]['lr'] self.dev_debugger.track_lr_schedulers_update( self.batch_idx, interval, scheduler_idx, old_lr, new_lr ) def _with_is_last(iterable): """Pass through values from the given iterable with an added boolean indicating if this is the last item. See `https://stackoverflow.com/a/1630350 `_""" it = iter(iterable) last = next(it) for val in it: # yield last and has next yield last, False last = val # yield last, no longer has next yield last, True