batch_indices passed to PredictionWriter write_on_epoch_end is wrong when strategy=ddp

[ozen]

🐛 Bug

Let's say we are using ddp and there is single dataloader, the number of data points in a process is 140, and the batch size is 64.

When the PredictionWriter's write_on_epoch_end is called on that process, the sizes of predictions and batch_indices parameters are as follows:

len(predictions) == 1
len(predictions[0]) == 3
len(predictions[0][0]) == 64
len(predictions[0][1]) == 64
len(predictions[0][2]) == 12

len(batch_indices) == 1
len(batch_indices[0]) == 3
len(batch_indices[0][0]) == 12
len(batch_indices[0][1]) == 12
len(batch_indices[0][2]) == 12

Also the contents of batch_indices[0][0], batch_indices[0][1], and batch_indices[0][2] are the same. predictions are correctly different.

The source of batch_indices is trainer.predict_loop.epoch_batch_indices which is populated here:

https://github.com/PyTorchLightning/pytorch-lightning/blob/3089dc3829c6456d74ab95aef06891927519eab9/pytorch_lightning/loops/dataloader/prediction_loop.py#L91-L95

dl_batch_indices is coming from PredictionEpochLoop._all_batch_indices

https://github.com/PyTorchLightning/pytorch-lightning/blob/a002f872eaad5a0f8f2069c49295cdda3568a406/pytorch_lightning/loops/epoch/prediction_epoch_loop.py#L102-L109

Which is filled in the _store_batch_indices method:

https://github.com/PyTorchLightning/pytorch-lightning/blob/a002f872eaad5a0f8f2069c49295cdda3568a406/pytorch_lightning/loops/epoch/prediction_epoch_loop.py#L163-L171

During debugging I see batch_sampler.batch_indices is in fact the last batch indices every time this method is called.

I couldn't dig deeper but I guess the sampler yields all batches before the program reaches to _store_batch_indices.

Environment

• PyTorch Lightning Version (e.g., 1.5.0): 1.5.3
• PyTorch Version (e.g., 1.10): 1.9.0
• Python version (e.g., 3.9): 3.8
• OS (e.g., Linux): Linux (Ubuntu)
• CUDA/cuDNN version: 11.3
• GPU models and configuration: 3x RTX 3080
• How you installed PyTorch (conda, pip, source): Using NVIDIA's Docker Image
• If compiling from source, the output of torch.__config__.show():
• Any other relevant information:

[awaelchli]

@ozen
I attempted to reproduce this. Here is what I have so far based on your description:

import os

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.callbacks import BasePredictionWriter


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def predict_step(self, batch, batch_idx):
        return self(batch)

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


class Writer(BasePredictionWriter):
    def write_on_batch_end(self, trainer, pl_module, prediction, batch_indices, batch, batch_idx, dataloader_idx):
        pass

    def write_on_epoch_end(self, trainer, pl_module, predictions, batch_indices):
        print(len(batch_indices[0][0]))
        print(len(batch_indices[0][1]))
        # print(len(batch_indices[0][2]))


def run():
    predict_data = DataLoader(RandomDataset(32, 140), batch_size=64)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        max_epochs=1,
        enable_model_summary=False,
        callbacks=[Writer(write_interval="epoch")],
        accelerator="cpu",
        strategy="ddp",
        devices=2,
    )
    trainer.predict(model, predict_data)


if __name__ == "__main__":
    run()

It outputs

64
6

on each process. So far, these are the expected numbers.
Could you you please take my script and modify it to your case?

[ozen]

Thank you @awaelchli

I forgot to try this without multi-process data loading. So, adding num_workers > 0 to the data loader should do it.

Moreover, I was wrong relating this to ddp. It seems this happens when data loading is multi-process.

[awaelchli]

Yes you are right. This is happening with a pytorch dataloader and a batch sampler:

if __name__ == "__main__":
    import torch
    from torch.utils.data import DataLoader, Dataset, SequentialSampler, BatchSampler
    
    ds = range(16)
    s = SequentialSampler(ds)
    bs = BatchSampler(s, batch_size=4, drop_last=False)
    dl = DataLoader(ds, batch_sampler=bs, num_workers=2)
    for x in dl:
        print("retrieve")

With num_workers=2 we get the output:

[0, 1, 2, 3]
[4, 5, 6, 7]
[8, 9, 10, 11]
[12, 13, 14, 15]
retrieve
retrieve
retrieve
retrieve

and with num_workers=0 we get

retrieve
[0, 1, 2, 3]
retrieve
[4, 5, 6, 7]
retrieve
[8, 9, 10, 11]
retrieve
[12, 13, 14, 15]

This shows that the sampling behavior is different depending on the DataLoader setting. When multiple workers are used, the indices are all fetched before a batch gets produced. When no workers are used, it is alternating between fetching indices and fetching the batch.

This means our implementation of the BatchSamplerWrapper here is fundamentally unusable for this scenario of multiple workers. It is not possible to capture the indices this way.

https://github.com/PyTorchLightning/pytorch-lightning/blob/f7c15f256ceb7073f004d592e77dace98dcf2423/pytorch_lightning/overrides/distributed.py#L122-L133

This has to be redesigned.

[awaelchli]

This is now the repro script:

import os

import torch
from torch.utils.data import DataLoader, Dataset

from pytorch_lightning import LightningModule, Trainer
from pytorch_lightning.callbacks import BasePredictionWriter


class RandomDataset(Dataset):
    def __init__(self, size, length):
        self.len = length
        self.data = torch.randn(length, size)

    def __getitem__(self, index):
        return self.data[index]

    def __len__(self):
        return self.len


class BoringModel(LightningModule):
    def __init__(self):
        super().__init__()
        self.layer = torch.nn.Linear(32, 2)

    def forward(self, x):
        return self.layer(x)

    def predict_step(self, batch, batch_idx):
        return self(batch)

    def configure_optimizers(self):
        return torch.optim.SGD(self.layer.parameters(), lr=0.1)


class Writer(BasePredictionWriter):
    def write_on_batch_end(self, trainer, pl_module, prediction, batch_indices, batch, batch_idx, dataloader_idx):
        pass

    def write_on_epoch_end(self, trainer, pl_module, predictions, batch_indices):
        print(batch_indices[0])


def run():
    predict_data = DataLoader(RandomDataset(32, 16), batch_size=4, num_workers=2)

    model = BoringModel()
    trainer = Trainer(
        default_root_dir=os.getcwd(),
        max_epochs=1,
        enable_model_summary=False,
        enable_progress_bar=False,
        callbacks=[Writer(write_interval="epoch")],
    )
    trainer.predict(model, predict_data)


if __name__ == "__main__":
    run()

Current output:

[[12, 13, 14, 15], [12, 13, 14, 15], [12, 13, 14, 15], [12, 13, 14, 15]]

Expected output:

[[0, 1, 2, 3], [4, 5, 6, 7], [8, 9, 10, 11], [12, 13, 14, 15]]

[awaelchli]

Thanks @ozen for helping to reproduce this. Fix will be released in 1.5.5 mid of next week. You can also install from source (see readme) if you want to test the fix right now.

[ozen]

Thanks for the quick fix @awaelchli :)