Method run_training_batch

utils/commons/trainer.py:369–456

(self, batch_idx, batch)

        task_ref.on_train_end()

    def run_training_batch(self, batch_idx, batch):
        if batch is None:
            return {}
        all_progress_bar_metrics = []
        all_log_metrics = []
        task_ref = self.get_task_ref()
        for opt_idx, optimizer in enumerate(self.optimizers):
            if optimizer is None:
                continue
            # make sure only the gradients of the current optimizer's paramaters are calculated
            # in the training step to prevent dangling gradients in multiple-optimizer setup.
            if len(self.optimizers) > 1:
                for k, param in task_ref.named_parameters():
                    param.requires_grad = False
                for group in optimizer.param_groups:
                    for param in group['params']:
                        param.requires_grad = True

            # forward pass
            with Timer("forward_training_step", enable=self.debug):
                with autocast(enabled=self.amp):
                    if self.on_gpu:
                        batch = move_to_cuda(copy.copy(batch), self.root_gpu)
                    args = [batch, batch_idx, opt_idx]
                    if self.use_ddp:
                        output = self.task(*args)
                    else:
                        output = task_ref.training_step(*args)
                    loss = output['loss']
                    if loss is None:
                        continue
                    progress_bar_metrics = output['progress_bar']
                    log_metrics = output['tb_log']
                    # accumulate loss
                    loss = loss / self.accumulate_grad_batches
                
            # backward pass
            with Timer("backward_training_step", enable=self.debug):
                if loss.requires_grad:
                    if self.amp:
                        self.amp_scalar.scale(loss).backward()
                    else:
                        loss.backward()

                # track progress bar metrics
                all_log_metrics.append(log_metrics)
                all_progress_bar_metrics.append(progress_bar_metrics)

                if loss is None:
                    continue

            # nan grads
            with Timer("checkNan_training_step", enable=self.debug):
                has_nan_grad = False
                nan_params_names = []
                if self.print_nan_grads:
                    for name, param in task_ref.named_parameters():
                        if (param.grad is not None) and torch.isnan(param.grad.float()).any():