Lightning-AI/pytorch-lightning

Returning None from training_step with multi GPU DDP training

Open

#5,243 建立於 2020年12月23日

在 GitHub 查看
 (26 留言) (4 反應) (1 負責人)Python (3,233 fork)batch import
distributedfeaturehelp wantedpriority: 1

倉庫指標

Star
 (26,687 star)
PR 合併指標
 (平均合併 9天 15小時) (30 天內合併 3 個 PR)

描述

🐛 Bug

Returning None from training_step with multi GPU DDP training freezes the training without exception

To Reproduce

Starting multi-gpu training with a None-returning training_step function.

Example training_step function:

    def training_step(self, batch, batch_idx):
        data, target = batch
        model_outputs = self.forward(images)
        loss = calc_loss(model_outputs, target)

        if torch.isnan(loss) or random.random() < .05:
            return None

        return loss

Example trainer:

 trainer = Trainer(
    gpus=2,
    distributed_backend="ddp"
)

Expected behavior

To continue training with skipping the current batch as pointed out at here.

Environment

No specific environment is needed to reproduce this bug.

Additional context

This issue was mentioned here: #4956 but not with specifics.

Note: By the time this issue being investigated, a help for a workaround would be great!

貢獻者指南