huggingface / huggingface/datasets

Streaming IterableDatasets do not work with torch DataLoaders

Open
#5,720 10 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
22k
Forks
3.4k
Avg merge
5d 7h
Merged PRs (30d)
17

Description

### 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..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
```

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.