Различия между Dataset и IterableDataset

Существует два типа объектов датасетов: Dataset и IterableDataset. Выбор типа датасета для использования или создания зависит от его размера. Как правило, IterableDataset идеально подходит для больших датасетов (речь идёт о сотнях гигабайт!) благодаря ленивому поведению и преимуществам в скорости, тогда как Dataset отлично справляется со всеми остальными задачами. На этой странице мы сравним различия между Dataset и IterableDataset, чтобы помочь вам выбрать подходящий объект датасета.

Загрузка и потоковая передача

Если у вас есть обычный Dataset, к нему можно обращаться с помощью my_dataset[0]. Это обеспечивает случайный доступ к строкам. Такие наборы данных также называют «map-style». Например, вы можете скачать ImageNet-1k следующим образом и получить доступ к любой строке:

from datasets import load_dataset

imagenet = load_dataset("timm/imagenet-1k-wds", split="train")  # downloads the full dataset
print(imagenet[0])

Однако есть одно ограничение: весь набор данных должен быть сохранён на диске или в памяти, что не позволяет работать с наборами данных, превышающими объём диска. Поскольку это может быть неудобно для больших наборов данных, существует другой тип набора данных — IterableDataset. Если у вас есть IterableDataset, вы можете обращаться к нему с помощью цикла for, чтобы загружать данные по мере итерации по набору. Таким образом, в памяти загружается лишь небольшая часть примеров, и ничего не записывается на диск.

Например, вы можете стримить датасет ImageNet-1k, не скачивая его на диск:

from datasets import load_dataset

imagenet = load_dataset("timm/imagenet-1k-wds", split="train", streaming=True)  # will start loading the data when iterated over
for example in imagenet:
    print(example)
    break

Streaming позволяет читать данные из сети, не записывая файлы на диск. Например, можно стримить датасеты, состоящие из нескольких шардов, каждый из которых занимает сотни гигабайт, как C4 или LAION-2B. Подробнее о том, как стримить датасет, читайте в Dataset Streaming Guide.

Однако это не единственное различие, поскольку «ленивое» поведение IterableDataset проявляется и при создании и обработке набора данных.

Создание наборов данных в стиле map и итерируемых наборов данных

Вы можете создать Dataset с помощью списков или словарей, при этом данные полностью преобразуются в формат Arrow, что позволяет легко обращаться к любой строке:

my_dataset = Dataset.from_dict({"col_1": [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]})
print(my_dataset[0])

Чтобы создать IterableDataset, с другой стороны, необходимо указать «ленивый» способ загрузки данных. В Python для этого обычно используют генераторы. Такие функции yield по одному примеру за раз, что означает, что вы не можете получить доступ к строке с помощью среза, как в обычном Dataset:

def my_generator(n):
    for i in range(n):
        yield {"col_1": i}

my_iterable_dataset = IterableDataset.from_generator(my_generator, gen_kwargs={"n": 10})
for example in my_iterable_dataset:
    print(example)
    break

Загрузка локальных файлов целиком и по частям

Локальные или удалённые файлы данных можно преобразовать в набор данных Arrow Dataset с помощью load_dataset():

data_files = {"train": ["path/to/data.csv"]}
my_dataset = load_dataset("csv", data_files=data_files, split="train")
print(my_dataset[0])

Однако для этого требуется этап преобразования из формата CSV в формат Arrow, что занимает время и место на диске, если ваш набор данных большой.

Чтобы сэкономить место на диске и пропустить этап конвертации, вы можете определить IterableDataset, стримя данные непосредственно из локальных файлов. Таким образом, данные будут считываться из локальных файлов по мере итерации по набору данных:

data_files = {"train": ["path/to/data.csv"]}
my_iterable_dataset = load_dataset("csv", data_files=data_files, split="train", streaming=True)
for example in my_iterable_dataset:  # this reads the CSV file progressively as you iterate over the dataset
    print(example)
    break

Поддерживается множество форматов файлов, таких как CSV, JSONL и Parquet, а также изображения и аудиофайлы. Подробнее о загрузке табличных, текстовых, изображений и аудио датасетов можно узнать в соответствующих руководствах.

Обработка данных в режиме Eager и обработка данных в режиме Lazy

При обработке объекта Dataset с помощью Dataset.map() весь набор данных обрабатывается немедленно и возвращается. Это похоже на то, как, например, работает pandas.

my_dataset = my_dataset.map(process_fn)  # process_fn is applied on all the examples of the dataset
print(my_dataset[0])

С другой стороны, в силу «ленивой» природы IterableDataset вызов IterableDataset.map() не применяет вашу функцию map ко всему набору данных. Вместо этого ваша функция map применяется на лету.

Поэтому вы можете объединить несколько этапов обработки в цепочку, и все они будут выполнены одновременно при начале итерации по набору данных:

my_iterable_dataset = my_iterable_dataset.map(process_fn_1)
my_iterable_dataset = my_iterable_dataset.filter(filter_fn)
my_iterable_dataset = my_iterable_dataset.map(process_fn_2)

# process_fn_1, filter_fn and process_fn_2 are applied on-the-fly when iterating over the dataset
for example in my_iterable_dataset:  
    print(example)
    break

Точное и быстрое приближённое перемешивание

При перемешивании Dataset с помощью Dataset.shuffle() выполняется точное перемешивание набора данных. Это работает так: берётся список индексов [0, 1, 2, ... len(my_dataset) - 1] и перемешивается. Затем обращение к my_dataset[0] возвращает строку и индекс, определённые первым элементом перемешанного отображения индексов:

my_dataset = my_dataset.shuffle(seed=42)
print(my_dataset[0])

Поскольку в случае IterableDataset у нас нет случайного доступа к строкам, мы не можем использовать перемешанный список индексов и обращаться к строке в произвольной позиции. Это исключает возможность использования точного перемешивания. Вместо этого в IterableDataset.shuffle() используется быстрое приближённое перемешивание. Оно использует буфер перемешивания, чтобы итеративно извлекать случайные примеры из набора данных. Поскольку набор данных по-прежнему читается итеративно, это обеспечивает высокую скорость:

my_iterable_dataset = my_iterable_dataset.shuffle(seed=42, buffer_size=100)
for example in my_iterable_dataset:
    print(example)
    break

Однако использование буфера перемешивания недостаточно для обеспечения удовлетворительного перемешивания при обучении моделей машинного обучения. Поэтому IterableDataset.shuffle() также перемешивает фрагменты набора данных, если он состоит из нескольких файлов или источников. Кроме того, буфер заполняется с использованием не более чем max_buffer_input_shards фрагментов за один раз (по умолчанию — 10), что значительно повышает качество перемешивания:

# Stream from the internet
my_iterable_dataset = load_dataset("deepmind/code_contests", split="train", streaming=True)
my_iterable_dataset.num_shards  # 39

# Stream from local files
data_files = {"train": [f"path/to/data_{i}.csv" for i in range(1024)]}
my_iterable_dataset = load_dataset("csv", data_files=data_files, split="train", streaming=True)
my_iterable_dataset.num_shards  # 1024

# From a generator function
def my_generator(n, sources):
    for source in sources:
        for example_id_for_current_source in range(n):
            yield {"example_id": f"{source}_{example_id_for_current_source}"}

gen_kwargs = {"n": 10, "sources": [f"path/to/data_{i}" for i in range(1024)]}
my_iterable_dataset = IterableDataset.from_generator(my_generator, gen_kwargs=gen_kwargs)
my_iterable_dataset.num_shards  # 1024

Если в вашем наборе данных мало файлов, вы можете повысить качество перемешивания, вызвав reshard() перед перемешиванием. Например, для наборов данных в формате Parquet это перераспределяет данные по группам строк, а не по одному файлу на каждый шард, что позволяет получить значительно больше шардов для перемешивания.

Различия в скорости

Обычные объекты Dataset основаны на Arrow, который обеспечивает быстрый случайный доступ к строкам. Благодаря отображению в память и тому, что Arrow — это формат, работающий в оперативной памяти, чтение данных с диска не требует дорогостоящих системных вызовов и десериализации. При итерации с помощью цикла for загрузка данных происходит ещё быстрее, так как итерация выполняется по непрерывным пакетам записей Arrow.

Однако как только ваш Dataset получает отображение индексов (например, через Dataset.shuffle()), скорость может снизиться в 10 раз. Это происходит потому, что появляется дополнительный шаг для получения индекса строки для чтения с использованием отображения индексов, и, что наиболее важно, вы больше не читаете непрерывные блоки данных. Чтобы восстановить скорость, вам нужно будет снова переписать весь набор данных на диске с помощью Dataset.flatten_indices(), что удаляет отображение индексов. В зависимости от размера вашего набора данных это может занять много времени:

my_dataset[0]  # fast
my_dataset = my_dataset.shuffle(seed=42)
my_dataset[0]  # up to 10x slower
my_dataset = my_dataset.flatten_indices()  # rewrite the shuffled dataset on disk as contiguous chunks of data
my_dataset[0]  # fast again

В этом случае мы рекомендуем перейти на IterableDataset и использовать его быстрый метод приблизительного перемешивания IterableDataset.shuffle(). Он перемешивает только порядок шардов и добавляет буфер перемешивания к вашему набору данных, что позволяет поддерживать оптимальную скорость работы набора данных. Вы также можете легко перемешать набор данных заново:

for example in enumerate(my_iterable_dataset):  # fast
    pass

shuffled_iterable_dataset = my_iterable_dataset.shuffle(seed=42, buffer_size=100)

for example in enumerate(shuffled_iterable_dataset):  # as fast as before
    pass

shuffled_iterable_dataset = my_iterable_dataset.shuffle(seed=1337, buffer_size=100)  # reshuffling using another seed is instantaneous

for example in enumerate(shuffled_iterable_dataset):  # still as fast as before
    pass

Если вы используете свой набор данных в нескольких эпохах, эффективный сид для перемешивания порядка шардов в буфере перемешивания равен seed + epoch. Это упрощает перемешивание набора данных между эпохами:

for epoch in range(n_epochs):
    my_iterable_dataset.set_epoch(epoch)
    for example in my_iterable_dataset:  # fast + reshuffled at each epoch using `effective_seed = seed + epoch`
        pass

Чтобы перезапустить итерацию набора данных map-стиля, достаточно просто пропустить первые примеры:

my_dataset = my_dataset.select(range(start_index, len(dataset)))

Но если вы используете DataLoader с Sampler, вам следует сохранить состояние вашего сэмплера (возможно, вы написали собственный сэмплер, который позволяет возобновлять процесс).

С другой стороны, итерируемые наборы данных не предоставляют случайный доступ к конкретному индексу примера для возобновления с него. Однако вы можете использовать IterableDataset.state_dict() и IterableDataset.load_state_dict(), чтобы возобновить работу с контрольной точки, аналогично тому, как это делается для моделей и оптимизаторов:

>>> iterable_dataset = Dataset.from_dict({"a": range(6)}).to_iterable_dataset(num_shards=3)
>>> # save in the middle of training
>>> state_dict = iterable_dataset.state_dict()
>>> # and resume later
>>> iterable_dataset.load_state_dict(state_dict)

Под капотом итерируемый набор данных отслеживает текущий шард, из которого выполняется чтение, и индекс примера в текущем шарде, сохраняя эту информацию в state_dict.

Чтобы возобновить работу с контрольной точки, набор данных пропускает все сегменты, которые были прочитаны ранее, чтобы начать с текущего сегмента. Затем он читает сегмент и пропускает примеры, пока не достигнет точного примера из контрольной точки.

Поэтому перезапуск набора данных происходит достаточно быстро, так как он не перечитывает уже обработанные фрагменты. Тем не менее возобновление работы с набором данных обычно не происходит мгновенно, поскольку необходимо начать чтение с начала текущего фрагмента и пропустить примеры, пока не будет достигнута позиция контрольной точки.

Это можно использовать вместе с StatefulDataLoader из torchdata, см. стриминг с PyTorch DataLoader.

Переключение с map-style на iterable

Если вы хотите воспользоваться «ленивым» поведением IterableDataset или их преимуществами в скорости, вы можете переключить ваш map-style Dataset на IterableDataset:

my_iterable_dataset = my_dataset.to_iterable_dataset()

Если вы хотите перемешать набор данных или использовать его с PyTorch DataLoader, мы рекомендуем создать шардированный IterableDataset:

my_iterable_dataset = my_dataset.to_iterable_dataset(num_shards=1024)
my_iterable_dataset.num_shards  # 1024