Skip certain step during training

[real-junjiezhang]

Bug description

I want to ignore some batch step during training, how can I write the code? Any suggestions would be appreciated.Thanks in advance.

The chatGPT answer below:

    def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
        # Define steps to skip
        steps_to_skip = {199, 302, 493, 1283}
        if self.trainer.global_step in steps_to_skip:
            return -1  # Skip training this step

Is this correct?

Thanks!

What version are you seeing the problem on?

master

How to reproduce the bug

No response

Error messages and logs

# Error messages and logs here please

Environment

Current environment

#- Lightning Component (e.g. Trainer, LightningModule, LightningApp, LightningWork, LightningFlow):
#- PyTorch Lightning Version (e.g., 1.5.0):
#- Lightning App Version (e.g., 0.5.2):
#- PyTorch Version (e.g., 2.0):
#- Python version (e.g., 3.9):
#- OS (e.g., Linux):
#- CUDA/cuDNN version:
#- GPU models and configuration:
#- How you installed Lightning(`conda`, `pip`, source):
#- Running environment of LightningApp (e.g. local, cloud):

More info

No response

[awaelchli]

@real-junjiezhang It's possible to skip a step by returning None from the training step. Would this answer your question?

[real-junjiezhang]

@real-junjiezhang It's possible to skip a step by returning None from the training step. Would this answer your question?

Thanks for your quick reply! What if currently I am using DDP to train the model? Can I write like this ?

    def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
        # Define steps to skip
        steps_to_skip = {199, 302, 493, 1283}
        if self.trainer.global_step in steps_to_skip:
            return -1  # Skip training this step

Thanks!

[awaelchli]

No, for DDP this is not supported, because the processes could fall out of sync.
In this case you need to think about how to rewrite your logic so that it doesn't have to be skipped. For example, if you have problematic data points, exclude them in your dataset beforehand.