datasets massively slows data reads, even in memory

[lrast]

Describe the bug

Loading image data in a huggingface dataset results in very slow read speeds, approximately 1000 times longer than reading the same data from a pytorch dataset. This applies even when the dataset is loaded into RAM using a keep_in_memory=True flag.

The following script reproduces the result with random data, but it applies equally to datasets that are loaded from the hub.

Steps to reproduce the bug

The following script should reproduce the behavior

import torch
import time
from datasets import Dataset

images = torch.randint(0, 255, (1000, 3, 224, 224), dtype=torch.uint8)
labels = torch.randint(0, 200, (1000,), dtype=torch.uint8)

pt_dataset = torch.utils.data.TensorDataset(images, labels)

hf_dataset = Dataset.from_dict({'image': images, 'label':labels})
hf_dataset.set_format('torch', dtype=torch.uint8)
hf_in_memory = hf_dataset.map(lambda x: x, keep_in_memory=True)

# measure access speeds
def time_access(dataset, img_col):
    start_time = time.time()
    for i in range(1000):
        _ = dataset[i][img_col].shape
    end_time = time.time()
    return end_time - start_time


print(f"In-memory Tensor access: {time_access(pt_dataset, 0):.4f} seconds")
print(f"HF Dataset access: {time_access(hf_dataset, 'image'):.4f} seconds")
print(f"In-memory HF Dataset access: {time_access(hf_in_memory, 'image'):.4f} seconds")

Expected behavior

For me, the above script produces

In-memory Tensor access: 0.0025 seconds
HF Dataset access: 2.9317 seconds
In-memory HF Dataset access: 2.8082 seconds

I think that this difference is larger than expected.

Environment info

• datasets version: 4.0.0
• Platform: macOS-14.7.7-arm64-arm-64bit
• Python version: 3.12.11
• huggingface_hub version: 0.34.3
• PyArrow version: 18.0.0
• Pandas version: 2.2.3
• fsspec version: 2024.9.0

[lhoestq]

Hi ! you should try

from datasets import Array3D, Dataset, Features, Value

features = Features({"image": Array3D(shape=(3, 224, 224), dtype="uint8"), "label": Value("uint8")})
hf_dataset = Dataset.from_dict({'image': images, 'label':labels}, features=features)

otherwise the type of the "image" column is List(List(List(Value("uint8")))) and is less efficient.

[lrast]

Thanks! This leads to a 10x speedup:

import torch
import time
from datasets import Array3D, Dataset, Features, Value

images = torch.randint(0, 255, (1000, 3, 224, 224), dtype=torch.uint8)
labels = torch.randint(0, 200, (1000,), dtype=torch.uint8)

pt_dataset = torch.utils.data.TensorDataset(images, labels)

features = Features({"image": Array3D(shape=(3, 224, 224), dtype="uint8"), "label": Value("uint8")})
hf_dataset = Dataset.from_dict({'image': images, 'label':labels}, features=features)
hf_in_memory = hf_dataset.map(lambda x: x, keep_in_memory=True)

hf_dataset.set_format('torch', dtype=torch.uint8)
hf_in_memory.set_format('torch', dtype=torch.uint8)

# measure access speeds
def time_access(dataset, img_col):
    start_time = time.time()
    for i in range(1000):
        _ = dataset[i][img_col].shape
    end_time = time.time()
    return end_time - start_time


print(f"In-memory Tensor access: {time_access(pt_dataset, 0):.4f} seconds")
print(f"HF Dataset access: {time_access(hf_dataset, 'image'):.4f} seconds")
print(f"In-memory HF Dataset access: {time_access(hf_in_memory, 'image'):.4f} seconds")

Produces

In-memory Tensor access: 0.0026 seconds
HF Dataset access: 0.2070 seconds
In-memory HF Dataset access: 0.2112 seconds

Curious if there is a reason why this is not the default behavior for huggingface image processors?

from transformers import ViTImageProcessor
from transformers import AutoImageProcessor

from datasets import load_dataset
# Load the dataset
ds = load_dataset('ylecun/mnist', split='train[0:100]')

# Instantiate the processor, explicitly requesting NumPy arrays
processor1 = ViTImageProcessor.from_pretrained('facebook/vit-mae-base',  do_convert_rgb=True)
processor2 = AutoImageProcessor.from_pretrained("facebook/detr-resnet-50", use_fast=True)

processed1 = ds.map(lambda row: processor1(row['image']))
processed2 = ds.map(lambda row: processor2(row['image']))

print( type(processed1['pixel_values'][0]), type(processed1['pixel_values'][0]))

produces

<class 'list'> <class 'list'>

I can, of course, manually manipulate the dataset to the use the correct format, but this is fairly standard for images, and the performance implications seem large.