Использование с PyTorch

Этот документ представляет собой краткое введение к использованию datasets с PyTorch, с особым акцентом на то, как получать объекты torch.Tensor из наших наборов данных, и как использовать DataLoader PyTorch и Dataset Hugging Face с максимальной производительностью.

Формат набора данных

По умолчанию наборы данных возвращают обычные объекты Python: целые числа, числа с плавающей точкой, строки, списки и т. д.

Чтобы получать тензоры PyTorch, можно задать формат набора данных как pytorch с помощью Dataset.with_format():

>>> from datasets import Dataset
>>> data = [[1, 2],[3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': tensor([1, 2])}
>>> ds[:2]
{'data': tensor([[1, 2],
         [3, 4]])}

Объект Dataset представляет собой обёртку над таблицей Arrow, что позволяет выполнять быстрые чтения без копирования из массивов в наборе данных в тензоры PyTorch.

Чтобы загрузить данные в виде тензоров на GPU, укажите аргумент device:

>>> import torch
>>> device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
>>> ds = ds.with_format("torch", device=device)
>>> ds[0]
{'data': tensor([1, 2], device='cuda:0')}

Многомерные массивы

Если ваш набор данных состоит из N-мерных массивов, вы увидите, что по умолчанию они считаются одним и тем же тензором, если форма фиксирована:

>>> from datasets import Dataset
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]]  # fixed shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': tensor([[1, 2],
         [3, 4]])}

>>> from datasets import Dataset
>>> data = [[[1, 2],[3]],[[4, 5, 6],[7, 8]]]  # varying shape
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': [tensor([1, 2]), tensor([3])]}

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

>>> from datasets import Dataset, Features, Array2D
>>> data = [[[1, 2],[3, 4]],[[5, 6],[7, 8]]]
>>> features = Features({"data": Array2D(shape=(2, 2), dtype='int32')})
>>> ds = Dataset.from_dict({"data": data}, features=features)
>>> ds = ds.with_format("torch")
>>> ds[0]
{'data': tensor([[1, 2],
         [3, 4]])}
>>> ds[:2]
{'data': tensor([[[1, 2],
          [3, 4]],
 
         [[5, 6],
          [7, 8]]])}

Другие типы признаков

Данные ClassLabel корректно преобразуются в тензоры:

>>> from datasets import Dataset, Features, ClassLabel
>>> labels = [0, 0, 1]
>>> features = Features({"label": ClassLabel(names=["negative", "positive"])})
>>> ds = Dataset.from_dict({"label": labels}, features=features) 
>>> ds = ds.with_format("torch")  
>>> ds[:3]
{'label': tensor([0, 0, 1])}

Строковые и двоичные объекты остаются без изменений, так как PyTorch поддерживает только числа.

Также поддерживаются типы функций Image и Audio.

Чтобы использовать тип функции Image, необходимо установить дополнительный пакет vision командой pip install datasets[vision].

>>> from datasets import Dataset, Features, Audio, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features) 
>>> ds = ds.with_format("torch")
>>> ds[0]["image"].shape
torch.Size([512, 512, 4])
>>> ds[0]
{'image': tensor([[[255, 215, 106, 255],
         [255, 215, 106, 255],
         ...,
         [255, 255, 255, 255],
         [255, 255, 255, 255]]], dtype=torch.uint8)}
>>> ds[:2]["image"].shape
torch.Size([2, 512, 512, 4])
>>> ds[:2]
{'image': tensor([[[[255, 215, 106, 255],
          [255, 215, 106, 255],
          ...,
          [255, 255, 255, 255],
          [255, 255, 255, 255]]]], dtype=torch.uint8)}

Чтобы использовать тип функции Audio, необходимо установить дополнительный пакет audio командой pip install datasets[audio].

>>> from datasets import Dataset, Features, Audio, Image
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features) 
>>> ds = ds.with_format("torch")  
>>> ds[0]["audio"]["array"]
tensor([ 6.1035e-05,  1.5259e-05,  1.6785e-04,  ..., -1.5259e-05,
        -1.5259e-05,  1.5259e-05])
>>> ds[0]["audio"]["sampling_rate"]
tensor(44100)

Загрузка данных

Как и объекты torch.utils.data.Dataset, объект Dataset можно передать напрямую в DataLoader PyTorch:

>>> import numpy as np
>>> from datasets import Dataset 
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(16)
>>> label = np.random.randint(0, 2, size=16)
>>> ds = Dataset.from_dict({"data": data, "label": label}).with_format("torch")
>>> dataloader = DataLoader(ds, batch_size=4)
>>> for batch in dataloader:
...     print(batch)                                                                                            
{'data': tensor([0.0047, 0.4979, 0.6726, 0.8105]), 'label': tensor([0, 1, 0, 1])}
{'data': tensor([0.4832, 0.2723, 0.4259, 0.2224]), 'label': tensor([0, 0, 0, 0])}
{'data': tensor([0.5837, 0.3444, 0.4658, 0.6417]), 'label': tensor([0, 1, 0, 0])}
{'data': tensor([0.7022, 0.1225, 0.7228, 0.8259]), 'label': tensor([1, 1, 1, 1])}

Оптимизация загрузки данных

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

Использование нескольких Workers

Вы можете параллелизировать загрузку данных с помощью аргумента num_workers объекта DataLoader PyTorch и получить более высокую пропускную способность.

Под капотом DataLoader запускает num_workers процессов. Каждый процесс перечитывает набор данных, переданный в DataLoader, и используется для запроса примеров. Перечитывание набора данных внутри рабочего процесса не заполняет оперативную память, поскольку оно просто снова отображает набор данных из диска в память.

>>> import numpy as np
>>> from datasets import Dataset, load_from_disk
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(10_000)
>>> Dataset.from_dict({"data": data}).save_to_disk("my_dataset")
>>> ds = load_from_disk("my_dataset").with_format("torch")
>>> dataloader = DataLoader(ds, batch_size=32, num_workers=4)

Потоковая передача данных

Потоковая загрузка набора данных с его загрузкой как IterableDataset. Это позволяет последовательно итерировать по удалённому набору данных без его скачивания на диск, а также по локальным файлам данных. Подробнее о том, какой тип набора данных лучше подходит для вашей задачи, читайте в руководстве выбор между обычным набором данных и итерируемым набором данных.

Итерируемый датасет из datasets наследуется от torch.utils.data.IterableDataset, поэтому его можно передать в torch.utils.data.DataLoader:

>>> import numpy as np
>>> from datasets import Dataset, load_dataset
>>> from torch.utils.data import DataLoader
>>> data = np.random.rand(10_000)
>>> Dataset.from_dict({"data": data}).push_to_hub("<username>/my_dataset")  # Upload to the Hugging Face Hub
>>> my_iterable_dataset = load_dataset("<username>/my_dataset", streaming=True, split="train")
>>> dataloader = DataLoader(my_iterable_dataset, batch_size=32)

Если набор данных разделён на несколько фрагментов (то есть состоит из нескольких файлов данных), то можно выполнять потоковую передачу параллельно с помощью num_workers:

>>> my_iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> my_iterable_dataset.num_shards
39
>>> dataloader = DataLoader(my_iterable_dataset, batch_size=32, num_workers=4)

В этом случае каждому воркеру выделяется подмножество списка шардов для потоковой передачи.

Точки контрольной записи и возобновление

Если вам нужен DataLoader, который можно сохранять в контрольную точку и возобновлять в середине обучения, вы можете использовать StatefulDataLoader из torchdata:

>>> from torchdata.stateful_dataloader import StatefulDataLoader
>>> my_iterable_dataset = load_dataset("deepmind/code_contests", streaming=True, split="train")
>>> dataloader = StatefulDataLoader(my_iterable_dataset, batch_size=32, num_workers=4)
>>> # save in the middle of training
>>> state_dict = dataloader.state_dict()
>>> # and resume later
>>> dataloader.load_state_dict(state_dict)

Это возможно благодаря IterableDataset.state_dict() и IterableDataset.load_state_dict().

Распределённый

To split your dataset across your training nodes, you can use datasets.distributed.split_dataset_by_node():

import os
from datasets.distributed import split_dataset_by_node

ds = split_dataset_by_node(ds, rank=int(os.environ["RANK"]), world_size=int(os.environ["WORLD_SIZE"]))

Это работает как для наборов данных в стиле map, так и для итерируемых наборов данных. Набор данных разбивается для узла с рангом rank в пуле узлов размером world_size.

Для наборов данных в стиле карты:

Каждому узлу назначается фрагмент данных, например, узлу с рангом 0 передается первый фрагмент набора данных.

Для итерируемых наборов данных:

Стратегия разбиения strategy может принимать значения "auto", "shards" или "examples":

• "shards": шарды равномерно распределяются по узлам
• "examples": каждый узел хранит 1 пример из world_size, пропуская остальные примеры
• "auto" (по умолчанию): использует "shards", если количество шардов в наборе данных является делителем world_size (то есть если dataset.num_shards % world_size == 0), что является наиболее оптимизированным вариантом. В противном случае использует "examples".

Это также можно объединить с torch.utils.data.DataLoader, если вы хотите, чтобы каждый узел использовал несколько рабочих процессов для загрузки данных.

Если вы перемешиваете итерируемый набор данных в распределённой конфигурации, убедитесь, что вы задали фиксированное значение seed в IterableDataset.shuffle(), чтобы на каждом узле использовался один и тот же перемешанный список шардов, и узел знал, какие шарды ему следует пропустить.

Гибридный шардированный параллелизм данных (HSDP)

Объект torch.distributed.device_mesh.DeviceMesh описывает топологию обучения в виде сетки рангов, где каждое измерение соответствует стратегии параллелизма. Для загрузки данных важны только измерения, связанные с параллелизмом по данным, а гибридное шардирование делает их легко пропустить:

• "fsdp" (как правило, внутри узла): группа FSDP, в которой каждый ранг хранит только фрагмент состояний модели (параметры, градиенты и состояния оптимизатора) и временно собирает полные параметры модуля с помощью all-gather непосредственно перед вычислениями. Обучение остаётся data-parallel: каждый ранг выполняет полный прямой и обратный проход на своей собственной выборке, поэтому в этом измерении, как и в DDP, каждому рангу нужны разные данные.
• "dp_replicate" (как правило, между узлами): дублированные копии этой шардированной модели, синхронизируемые с помощью дополнительного all-reduce между репликами. Реплики также обрабатывают разные данные.
• "tp" (внутри узла): ранги тензорного параллелизма, которые разделяют веса каждого слоя для совместного вычисления одной партии. Они не являются данными-параллельными и должны читать одни и те же примеры.

Степень параллелизма данных — это произведение размерностей шардирования (FSDP) и репликации, которое становится меньше WORLD_SIZE, как только в сетке появляется размерность, например "tp". Передайте в datasets.distributed.split_dataset_by_node() это произведение и ранг, сплющенный по этим двум размерностям:

from torch.distributed.device_mesh import init_device_mesh
from datasets.distributed import split_dataset_by_node

# Hybrid sharding spreads data parallelism over two mesh dimensions, but both
# need different data per rank, so `dp_size` is their product.
mesh = init_device_mesh("cuda", (2, 2, 2), mesh_dim_names=("dp_replicate", "fsdp", "tp"))
dp_replicate_mesh, fsdp_mesh = mesh["dp_replicate"], mesh["fsdp"]

dp_size = dp_replicate_mesh.size() * fsdp_mesh.size()
dp_rank = dp_replicate_mesh.get_local_rank() * shard.size() + fsdp_mesh.get_local_rank()

ds = split_dataset_by_node(ds, rank=dp_rank, world_size=dp_size)

Здесь dp_size равно 2 * 2 = 4: каждый ранг с параллелизмом по данным читает четверть набора данных, а два ранга каждой пары "tp" читают одни и те же примеры.