Lightning-AI/pytorch-lightning

`CombinedLoader` takes a long time when `num_workers > 0`

Open

#18,584 建立於 2023年9月19日

在 GitHub 查看
 (1 留言) (0 反應) (0 負責人)Python (3,233 fork)batch import
bughelp wantedperformancerepro neededver: 2.0.x

倉庫指標

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

描述

Bug description

I am currently using CombinedLoader (https://lightning.ai/docs/pytorch/stable/api/lightning.pytorch.utilities.combined_loader.html#lightning.pytorch.utilities.combined_loader.CombinedLoader) to combine multiple datasets. It works fine but I noticed that setting the dataloader with num_workers > 0 causes it to run extremely slow. Is there a logical explanation for this? Feel like this could be a bug otherwise. I attached a chunk of my code to show what I am doing.

What version are you seeing the problem on?

v2.0

How to reproduce the bug

import pytorch_lightning as pl
from lightning.pytorch.utilities.combined_loader import CombinedLoader

class CollectiveDataloader(pl.LightningDataModule):
    def __init__(self, datasets, num_workers=8, batch_size=10, shuffle=True):
        super().__init__()
        self.train_set = CollectiveDataset(
            datasets, num_workers, batch_size, shuffle
        ).datasets

    def train_dataloader(self):
        return CombinedLoader(self.train_set, "sequential")

class CollectiveDataset:
    def __init__(self, datasets, num_workers, batch_size, shuffle):
        # datasets is a dictionary of {dataset_name : Dataset object}
        loaded_datasets = {
            name: DataLoader(
                dataset,
                batch_size=batch_size,
                shuffle=shuffle,
                ### SETTING THIS > 0 RUNS REALLY SLOW ###
                num_workers=num_workers,
            )
            for name, dataset in datasets.items()
        }
        self.datasets = loaded_datasets

Environment

#- PyTorch Lightning Version: 2.0.7
#- PyTorch Version: 2.0.1
#- Python version: 3.10.12
#- OS: Linux
#- CUDA/cuDNN version: 12.0

cc @borda

貢獻者指南