Method training_forward

utils/pl_utils.py:1564–1601

Handle forward for each training case (distributed, single gpu, etc...) :param batch: :param batch_idx: :return:

(self, batch, batch_idx, opt_idx, hiddens)

        return 0, grad_norm_dic, all_log_metrics

    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, opt_idx]

        # distributed forward
        if self.use_ddp or self.use_dp:
            output = self.model(*args)
        # single GPU forward
        elif self.single_gpu:
            gpu_id = 0
            if isinstance(self.data_parallel_device_ids, list):
                gpu_id = self.data_parallel_device_ids[0]
            batch = self.transfer_batch_to_gpu(copy.copy(batch), gpu_id)
            args[0] = batch
            output = self.model.training_step(*args)
        # CPU forward
        else:
            output = self.model.training_step(*args)

        # allow any mode to define training_end
        model_ref = self.get_model()
        output_ = model_ref.training_end(output)
        if output_ is not None:
            output = output_

        # format and reduce outputs accordingly
        output = self.process_output(output, train=True)

        return output

    # ---------------
    # Utils