Streaming IterableDatasets do not work with torch DataLoaders
NONE
Edit: This behavior is true even without `.take/.set`
2023-04-09T16:38:48Z
https://github.com/huggingface/datasets/issues/5720#issuecomment-1501166850
https://api.github.com/repos/huggingface/datasets/issues/5720
IC_kwDODunzps5Zef0C
https://api.github.com/repos/huggingface/datasets/issues/comments/1501166850/reactions
2023-04-09T16:38:48Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1501166850
https://avatars.githubusercontent.com/u/29244648?v=4
https://api.github.com/users/jlehrer1/events{/privacy}
https://api.github.com/users/jlehrer1/followers
https://api.github.com/users/jlehrer1/following{/other_user}
https://api.github.com/users/jlehrer1/gists{/gist_id}
https://github.com/jlehrer1
jlehrer1
MDQ6VXNlcjI5MjQ0NjQ4
https://api.github.com/users/jlehrer1/orgs
https://api.github.com/users/jlehrer1/received_events
https://api.github.com/users/jlehrer1/repos
https://api.github.com/users/jlehrer1/starred{/owner}{/repo}
https://api.github.com/users/jlehrer1/subscriptions
User
https://api.github.com/users/jlehrer1
public
NONE
I'm experiencing the same problem that @jlehrer1. I was able to reproduce it with a very small example:

```py
from datasets import Dataset, load_dataset, load_dataset_builder
from torch.utils.data import DataLoader


def my_gen():
    for i in range(1, 4):
        yield {"a": i}

# Saving the dataset as a parquet file
dataset = Dataset.from_generator(my_gen)
train_path = "/tmp/test.parquet"
dataset.to_parquet(train_path)

# Creating a local dataset from the parquet file
data_files = {"train": [str(train_path)]}
builder = load_dataset_builder("parquet", data_files=data_files)
builder.download_and_prepare("/tmp/test_ds", file_format="parquet")

# Loading the dataset from the local directory as streaming
dataset = load_dataset("parquet", data_dir="/tmp/test_ds", split="train", streaming=True)
dataset.with_format("torch")

dl = DataLoader(dataset, batch_size=2, num_workers=1)
for row in dl:
    print(row)
```

My env info:
```
datasets                      2.11.0
torch                         2.0.0
torchvision                   0.15.1
Python 3.9.16
```

Note that the example above doesn't fail if the number of workers used is `0`
2023-04-14T15:28:12Z
https://github.com/huggingface/datasets/issues/5720#issuecomment-1508794283
https://api.github.com/repos/huggingface/datasets/issues/5720
IC_kwDODunzps5Z7l-r
https://api.github.com/repos/huggingface/datasets/issues/comments/1508794283/reactions
2023-04-14T15:28:12Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1508794283
https://avatars.githubusercontent.com/u/895720?v=4
https://api.github.com/users/ivanprado/events{/privacy}
https://api.github.com/users/ivanprado/followers
https://api.github.com/users/ivanprado/following{/other_user}
https://api.github.com/users/ivanprado/gists{/gist_id}
https://github.com/ivanprado
ivanprado
MDQ6VXNlcjg5NTcyMA==
https://api.github.com/users/ivanprado/orgs
https://api.github.com/users/ivanprado/received_events
https://api.github.com/users/ivanprado/repos
https://api.github.com/users/ivanprado/starred{/owner}{/repo}
https://api.github.com/users/ivanprado/subscriptions
User
https://api.github.com/users/ivanprado
public
COLLABORATOR
I cannot reproduce this error, not even with your MRE @ivanprado (your env appears to be the same as Colab's, and your code runs there without issues). 
2023-04-20T17:58:50Z
https://github.com/huggingface/datasets/issues/5720#issuecomment-1516731345
https://api.github.com/repos/huggingface/datasets/issues/5720
IC_kwDODunzps5aZ3vR
https://api.github.com/repos/huggingface/datasets/issues/comments/1516731345/reactions
2023-04-20T17:58:50Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1516731345
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
NONE
@mariosasko you are right, it works on Colab. I digged deeper and found that the problem arises when the multiprocessing method is set to be `spawn`. This code reproduces the problem in Colab:

```py
from datasets import Dataset, load_dataset, load_dataset_builder
from torch.utils.data import DataLoader
import multiprocessing as mp

mp.set_start_method('spawn')

def my_gen():
    for i in range(1, 4):
        yield {"a": i}


def main():
    # Saving the dataset as a parquet file
    dataset = Dataset.from_generator(my_gen)
    train_path = "/tmp/test.parquet"
    dataset.to_parquet(train_path)

    # Creating a local dataset from the parquet file
    data_files = {"train": [str(train_path)]}
    builder = load_dataset_builder("parquet", data_files=data_files)
    builder.download_and_prepare("/tmp/test_ds", file_format="parquet")

    # Loading the dataset from the local directory as streaming
    dataset = load_dataset("parquet", data_dir="/tmp/test_ds", split="train", streaming=True)
    dataset.with_format("torch")

    dl = DataLoader(dataset, batch_size=2, num_workers=1)
    for row in dl:
        print(row)

main()
```
2023-04-21T07:23:10Z
https://github.com/huggingface/datasets/issues/5720#issuecomment-1517385577
https://api.github.com/repos/huggingface/datasets/issues/5720
IC_kwDODunzps5acXdp
https://api.github.com/repos/huggingface/datasets/issues/comments/1517385577/reactions
2023-04-21T07:23:10Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1517385577
https://avatars.githubusercontent.com/u/895720?v=4
https://api.github.com/users/ivanprado/events{/privacy}
https://api.github.com/users/ivanprado/followers
https://api.github.com/users/ivanprado/following{/other_user}
https://api.github.com/users/ivanprado/gists{/gist_id}
https://github.com/ivanprado
ivanprado
MDQ6VXNlcjg5NTcyMA==
https://api.github.com/users/ivanprado/orgs
https://api.github.com/users/ivanprado/received_events
https://api.github.com/users/ivanprado/repos
https://api.github.com/users/ivanprado/starred{/owner}{/repo}
https://api.github.com/users/ivanprado/subscriptions
User
https://api.github.com/users/ivanprado
public
NONE
So is there a way to fix this by changing the `mp` method? This is blocking any usage of the `datasets` library for me
2023-04-24T22:59:03Z
https://github.com/huggingface/datasets/issues/5720#issuecomment-1520927388
https://api.github.com/repos/huggingface/datasets/issues/5720
IC_kwDODunzps5ap4Kc
https://api.github.com/repos/huggingface/datasets/issues/comments/1520927388/reactions
2023-04-24T22:59:03Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1520927388
https://avatars.githubusercontent.com/u/29244648?v=4
https://api.github.com/users/jlehrer1/events{/privacy}
https://api.github.com/users/jlehrer1/followers
https://api.github.com/users/jlehrer1/following{/other_user}
https://api.github.com/users/jlehrer1/gists{/gist_id}
https://github.com/jlehrer1
jlehrer1
MDQ6VXNlcjI5MjQ0NjQ4
https://api.github.com/users/jlehrer1/orgs
https://api.github.com/users/jlehrer1/received_events
https://api.github.com/users/jlehrer1/repos
https://api.github.com/users/jlehrer1/starred{/owner}{/repo}
https://api.github.com/users/jlehrer1/subscriptions
User
https://api.github.com/users/jlehrer1
public
NONE
@jlehrer1 can you try adding `mp.set_start_method('fork')` at the beginning of your code? Maybe this helps you. Keep us posted. 
2023-04-25T07:06:59Z
https://github.com/huggingface/datasets/issues/5720#issuecomment-1521260440
https://api.github.com/repos/huggingface/datasets/issues/5720
IC_kwDODunzps5arJeY
https://api.github.com/repos/huggingface/datasets/issues/comments/1521260440/reactions
2023-04-25T07:06:59Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1521260440
https://avatars.githubusercontent.com/u/895720?v=4
https://api.github.com/users/ivanprado/events{/privacy}
https://api.github.com/users/ivanprado/followers
https://api.github.com/users/ivanprado/following{/other_user}
https://api.github.com/users/ivanprado/gists{/gist_id}
https://github.com/ivanprado
ivanprado
MDQ6VXNlcjg5NTcyMA==
https://api.github.com/users/ivanprado/orgs
https://api.github.com/users/ivanprado/received_events
https://api.github.com/users/ivanprado/repos
https://api.github.com/users/ivanprado/starred{/owner}{/repo}
https://api.github.com/users/ivanprado/subscriptions
User
https://api.github.com/users/ivanprado
public
NONE
I have a similar issue: 
> mp.set_start_method('fork')


Didnt work
2023-05-27T12:57:07Z
https://github.com/huggingface/datasets/issues/5720#issuecomment-1565401616
https://api.github.com/repos/huggingface/datasets/issues/5720
IC_kwDODunzps5dTiIQ
https://api.github.com/repos/huggingface/datasets/issues/comments/1565401616/reactions
2023-05-27T12:57:07Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1565401616
https://avatars.githubusercontent.com/u/58503178?v=4
https://api.github.com/users/Wonder1905/events{/privacy}
https://api.github.com/users/Wonder1905/followers
https://api.github.com/users/Wonder1905/following{/other_user}
https://api.github.com/users/Wonder1905/gists{/gist_id}
https://github.com/Wonder1905
Wonder1905
MDQ6VXNlcjU4NTAzMTc4
https://api.github.com/users/Wonder1905/orgs
https://api.github.com/users/Wonder1905/received_events
https://api.github.com/users/Wonder1905/repos
https://api.github.com/users/Wonder1905/starred{/owner}{/repo}
https://api.github.com/users/Wonder1905/subscriptions
User
https://api.github.com/users/Wonder1905
public
CONTRIBUTOR
What if I want to use GPU? spawn is a must have @ivanprado 
2024-11-15T12:35:24Z
https://github.com/huggingface/datasets/issues/5720#issuecomment-2478718781
https://api.github.com/repos/huggingface/datasets/issues/5720
IC_kwDODunzps6Tvj89
https://api.github.com/repos/huggingface/datasets/issues/comments/2478718781/reactions
2024-11-15T12:35:24Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2478718781
https://avatars.githubusercontent.com/u/11533479?v=4
https://api.github.com/users/npuichigo/events{/privacy}
https://api.github.com/users/npuichigo/followers
https://api.github.com/users/npuichigo/following{/other_user}
https://api.github.com/users/npuichigo/gists{/gist_id}
https://github.com/npuichigo
npuichigo
MDQ6VXNlcjExNTMzNDc5
https://api.github.com/users/npuichigo/orgs
https://api.github.com/users/npuichigo/received_events
https://api.github.com/users/npuichigo/repos
https://api.github.com/users/npuichigo/starred{/owner}{/repo}
https://api.github.com/users/npuichigo/subscriptions
User
https://api.github.com/users/npuichigo
public
NONE
@ivanprado you're right, this problem gets solved in case number of workers is set to 0, but this essentially destroys any level parallelism we can get.
2025-03-19T11:47:37Z
https://github.com/huggingface/datasets/issues/5720#issuecomment-2736346299
https://api.github.com/repos/huggingface/datasets/issues/5720
IC_kwDODunzps6jGVS7
https://api.github.com/repos/huggingface/datasets/issues/comments/2736346299/reactions
2025-03-19T11:47:37Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2736346299
https://avatars.githubusercontent.com/u/31623825?v=4
https://api.github.com/users/ieshaan12/events{/privacy}
https://api.github.com/users/ieshaan12/followers
https://api.github.com/users/ieshaan12/following{/other_user}
https://api.github.com/users/ieshaan12/gists{/gist_id}
https://github.com/ieshaan12
ieshaan12
MDQ6VXNlcjMxNjIzODI1
https://api.github.com/users/ieshaan12/orgs
https://api.github.com/users/ieshaan12/received_events
https://api.github.com/users/ieshaan12/repos
https://api.github.com/users/ieshaan12/starred{/owner}{/repo}
https://api.github.com/users/ieshaan12/subscriptions
User
https://api.github.com/users/ieshaan12
public
NONE
Exactly guys, agree with you. I'm just one like yours here. I'm not a datasets contributor. This issue prevented me to use this library.
2025-03-19T14:06:45Z
https://github.com/huggingface/datasets/issues/5720#issuecomment-2736775718
https://api.github.com/repos/huggingface/datasets/issues/5720
IC_kwDODunzps6jH-Im
https://api.github.com/repos/huggingface/datasets/issues/comments/2736775718/reactions
2025-03-19T14:06:45Z
https://api.github.com/repos/huggingface/datasets/issues/comments/2736775718
https://avatars.githubusercontent.com/u/895720?v=4
https://api.github.com/users/ivanprado/events{/privacy}
https://api.github.com/users/ivanprado/followers
https://api.github.com/users/ivanprado/following{/other_user}
https://api.github.com/users/ivanprado/gists{/gist_id}
https://github.com/ivanprado
ivanprado
MDQ6VXNlcjg5NTcyMA==
https://api.github.com/users/ivanprado/orgs
https://api.github.com/users/ivanprado/received_events
https://api.github.com/users/ivanprado/repos
https://api.github.com/users/ivanprado/starred{/owner}{/repo}
https://api.github.com/users/ivanprado/subscriptions
User
https://api.github.com/users/ivanprado
public
### Describe the bug

When using streaming datasets set up with train/val split using `.skip()` and `.take()`, the following error occurs when iterating over a torch dataloader:
```
  File "/Users/julian/miniconda3/envs/sims/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 363, in __iter__
    self._iterator = self._get_iterator()
  File "/Users/julian/miniconda3/envs/sims/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 314, in _get_iterator
    return _MultiProcessingDataLoaderIter(self)
  File "/Users/julian/miniconda3/envs/sims/lib/python3.9/site-packages/torch/utils/data/dataloader.py", line 927, in __init__
    w.start()
  File "/Users/julian/miniconda3/envs/sims/lib/python3.9/multiprocessing/process.py", line 121, in start
    self._popen = self._Popen(self)
  File "/Users/julian/miniconda3/envs/sims/lib/python3.9/multiprocessing/context.py", line 224, in _Popen
    return _default_context.get_context().Process._Popen(process_obj)
  File "/Users/julian/miniconda3/envs/sims/lib/python3.9/multiprocessing/context.py", line 284, in _Popen
    return Popen(process_obj)
  File "/Users/julian/miniconda3/envs/sims/lib/python3.9/multiprocessing/popen_spawn_posix.py", line 32, in __init__
    super().__init__(process_obj)
  File "/Users/julian/miniconda3/envs/sims/lib/python3.9/multiprocessing/popen_fork.py", line 19, in __init__
    self._launch(process_obj)
  File "/Users/julian/miniconda3/envs/sims/lib/python3.9/multiprocessing/popen_spawn_posix.py", line 47, in _launch
    reduction.dump(process_obj, fp)
  File "/Users/julian/miniconda3/envs/sims/lib/python3.9/multiprocessing/reduction.py", line 60, in dump
    ForkingPickler(file, protocol).dump(obj)
AttributeError: Can't pickle local object '_generate_examples_from_tables_wrapper.<locals>.wrapper'
```

To reproduce, run the code 
```
from datasets import load_dataset
data = load_dataset(args.dataset_name, split="train", streaming=True)
train_len = 5000
val_len = 100

train, val = data.take(train_len), data.skip(train_len).take(val_len)
traindata = IterableClipDataset(data, context_length=args.max_len, tokenizer=tokenizer, image_key="url", text_key="text")

traindata = DataLoader(traindata, batch_size=args.batch_size, num_workers=args.num_workers, persistent_workers=True)
```

Where the class IterableClipDataset is a simple wrapper to cast the dataset to a torch iterabledataset, defined via 
```
from torch.utils.data import Dataset, IterableDataset
from torchvision.transforms import Compose, Resize, ToTensor
from transformers import AutoTokenizer
import requests
from PIL import Image

class IterableClipDataset(IterableDataset):
    def __init__(self, dataset, context_length: int, image_transform=None, tokenizer=None, image_key="image", text_key="text"):
        self.dataset = dataset
        self.context_length = context_length
        self.image_transform = Compose([Resize((224, 224)), ToTensor()]) if image_transform is None else image_transform
        self.tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased") if tokenizer is None else tokenizer
        self.image_key = image_key
        self.text_key = text_key

    def read_image(self, url: str):
        try: # Try to read the image
            image = Image.open(requests.get(url, stream=True).raw)
        except:
            image = Image.new("RGB", (224, 224), (0, 0, 0))
        return image
    
    def process_sample(self, image, text):
        if isinstance(image, str):
            image = self.read_image(image)
        if self.image_transform is not None:
            image = self.image_transform(image)
        text = self.tokenizer.encode(
            text, add_special_tokens=True, max_length=self.context_length, truncation=True, padding="max_length"
        )
        text = torch.tensor(text, dtype=torch.long)
        return image, text

    def __iter__(self):
        for sample in self.dataset:
            image, text = sample[self.image_key], sample[self.text_key]
            yield self.process_sample(image, text)
```



### Steps to reproduce the bug

Steps to reproduce
1. Install `datasets`, `torch`, and `PIL` (if you want to reproduce exactly)
2. Run the code above

### Expected behavior

Batched data is produced from the dataloader 

### Environment info

```
datasets == 2.9.0
python == 3.9.12
torch == 1.11.0
```