huggingface / huggingface/datasets

Slow iteration speeds when using IterableDataset.shuffle with load_dataset(data_files=..., streaming=True)

Open
#7,102 2 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 I load a dataset from a number of arrow files, as in:

```
random_dataset = load_dataset(
"arrow",
data_files={split: shard_filepaths},
streaming=True,
split=split,
)
```

I'm able to get fast iteration speeds when iterating over the dataset without shuffling.

When I shuffle the dataset, the iteration speed is reduced by ~1000x.

It's very possible the way I'm loading dataset shards is not appropriate; if so please advise!

Thanks for the help

### Steps to reproduce the bug

Here's full code to reproduce the issue:

- Generate a random dataset
- Create shards of data independently using Dataset.save_to_disk()
- The below will generate 16 shards (arrow files), of 512 examples each

```
import time
from pathlib import Path
from multiprocessing import Pool, cpu_count

import torch
from datasets import Dataset, load_dataset

split = "train"
split_save_dir = "/tmp/random_split"

def generate_random_example():
return {
'inputs': torch.randn(128).tolist(),
'indices': torch.randint(0, 10000, (2, 20000)).tolist(),
'values': torch.randn(20000).tolist(),
}

def generate_shard_dataset(examples_per_shard: int = 512):
dataset_dict = {
'inputs': [],
'indices': [],
'values': []
}

for _ in range(examples_per_shard):
example = generate_random_example()
dataset_dict['inputs'].append(example['inputs'])
dataset_dict['indices'].append(example['indices'])
dataset_dict['values'].append(example['values'])

return Dataset.from_dict(dataset_dict)

def save_shard(shard_idx, save_dir, examples_per_shard):
shard_dataset = generate_shard_dataset(examples_per_shard)
shard_write_path = Path(save_dir) / f"shard_{shard_idx}"
shard_dataset.save_to_disk(shard_write_path)
return str(Path(shard_write_path) / "data-00000-of-00001.arrow")

def generate_split_shards(save_dir, num_shards: int = 16, examples_per_shard: int = 512):
with Pool(cpu_count()) as pool:
args = [(m, save_dir, examples_per_shard) for m in range(num_shards)]
shard_filepaths = pool.starmap(save_shard, args)

return shard_filepaths

shard_filepaths = generate_split_shards(split_save_dir)
```

Load the dataset as IterableDataset:

```
random_dataset = load_dataset(
"arrow",
data_files={split: shard_filepaths},
streaming=True,
split=split,
)
random_dataset = random_dataset.with_format("numpy")
```

Observe the iterations/second when iterating over the dataset directly, and applying shuffling before iterating:

Without shuffling, this gives ~1500 iterations/second

```
start_time = time.time()
for count, item in enumerate(random_dataset):
if count > 0 and count % 100 == 0:
elapsed_time = time.time() - start_time
iterations_per_second = count / elapsed_time
print(f"Processed {count} items at an average of {iterations_per_second:.2f} iterations/second")
```

```
Processed 100 items at an average of 705.74 iterations/second
Processed 200 items at an average of 1169.68 iterations/second
Processed 300 items at an average of 1497.97 iterations/second
Processed 400 items at an average of 1739.62 iterations/second
Processed 500 items at an average of 1931.11 iterations/second`
```

When shuffling, this gives ~3 iterations/second:

```

random_dataset = random_dataset.shuffle(buffer_size=100,seed=42)

start_time = time.time()
for count, item in enumerate(random_dataset):
if count > 0 and count % 100 == 0:
elapsed_time = time.time() - start_time
iterations_per_second = count / elapsed_time
print(f"Processed {count} items at an average of {iterations_per_second:.2f} iterations/second")
```

```
Processed 100 items at an average of 3.75 iterations/second
Processed 200 items at an average of 3.93 iterations/second
```

### Expected behavior

Iterations per second should be barely affected by shuffling, especially with a small buffer size

### Environment info

Datasets version: 2.21.0
Python 3.10
Ubuntu 22.04

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.