Datasets documentation
Использование с JAX
Использование с JAX
Этот документ представляет собой краткое введение к использованию datasets с JAX, с особым акцентом на то, как получать объекты jax.Array из наших наборов данных и как использовать их для обучения моделей JAX.
Для воспроизведения приведённого выше кода необходимы
jaxиjaxlib, поэтому убедитесь, что вы установили их командойpip install datasets[jax].
Формат набора данных
По умолчанию наборы данных возвращают обычные объекты Python: целые числа, числа с плавающей точкой, строки, списки и т. д., при этом строковые и двоичные объекты остаются неизменными, поскольку JAX поддерживает только числа.
Чтобы получать массивы JAX (похожие на numpy), можно задать формат набора данных как jax:
>>> from datasets import Dataset
>>> data = [[1, 2], [3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("jax")
>>> ds[0]
{'data': DeviceArray([1, 2], dtype=int32)}
>>> ds[:2]
{'data': DeviceArray([
[1, 2],
[3, 4]], dtype=int32)}Объект Dataset представляет собой обёртку над таблицей Arrow, которая позволяет быстро читать данные из массивов набора в массивы JAX.
Обратите внимание, что точно такая же процедура применяется к объектам DatasetDict, так что
при установке формата DatasetDict в jax все Dataset внутри него
будут отформатированы как jax:
>>> from datasets import DatasetDict
>>> data = {"train": {"data": [[1, 2], [3, 4]]}, "test": {"data": [[5, 6], [7, 8]]}}
>>> dds = DatasetDict.from_dict(data)
>>> dds = dds.with_format("jax")
>>> dds["train"][:2]
{'data': DeviceArray([
[1, 2],
[3, 4]], dtype=int32)}Ещё один момент, который нужно учитывать: форматирование применяется только тогда, когда вы фактически обращаетесь к данным. Поэтому, если вы хотите получить массив JAX из набора данных, вам сначала нужно обратиться к данным, иначе формат останется прежним.
Наконец, чтобы загрузить данные на выбранное вами устройство, можно указать аргумент device, но обратите внимание, что jaxlib.xla_extension.Device не поддерживается, так как он не сериализуется ни с помощью pickle, ни с помощью dill, поэтому вместо него нужно использовать его строковый идентификатор:
>>> import jax
>>> from datasets import Dataset
>>> data = [[1, 2], [3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> device = str(jax.devices()[0]) # Not casting to `str` before passing it to `with_format` will raise a `ValueError`
>>> ds = ds.with_format("jax", device=device)
>>> ds[0]
{'data': DeviceArray([1, 2], dtype=int32)}
>>> ds[0]["data"].device()
TFRT_CPU_0
>>> assert ds[0]["data"].device() == jax.devices()[0]
TrueОбратите внимание, что если аргумент device не передан в with_format, будет использовано устройство по умолчанию, а именно jax.devices()[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("jax")
>>> ds[0]
{'data': Array([[1, 2],
[3, 4]], dtype=int32)}>>> 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("jax")
>>> ds[0]
{'data': [Array([1, 2], dtype=int32), Array([3], dtype=int32)]}Однако такая логика часто требует медленных сравнений форм и копирования данных.
Чтобы избежать этого, необходимо явно использовать тип фичи 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("jax")
>>> ds[0]
{'data': Array([[1, 2],
[3, 4]], dtype=int32)}
>>> ds[:2]
{'data': Array([[[1, 2],
[3, 4]],
[[5, 6],
[7, 8]]], dtype=int32)}Другие типы признаков
Данные 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("jax")
>>> ds[:3]
{'label': DeviceArray([0, 0, 1], dtype=int32)}Строковые и двоичные объекты остаются без изменений, поскольку JAX поддерживает только числа.
Также поддерживаются типы функций Image и Audio.
Чтобы использовать тип функции Image, необходимо установить дополнительный пакет
visionкомандойpip install datasets[vision].
>>> from datasets import Dataset, Features, Image
>>> images = ["path/to/image.png"] * 10
>>> features = Features({"image": Image()})
>>> ds = Dataset.from_dict({"image": images}, features=features)
>>> ds = ds.with_format("jax")
>>> ds[0]["image"].shape
(512, 512, 3)
>>> ds[0]
{'image': DeviceArray([[[ 255, 255, 255],
[ 255, 255, 255],
...,
[ 255, 255, 255],
[ 255, 255, 255]]], dtype=uint8)}
>>> ds[:2]["image"].shape
(2, 512, 512, 3)
>>> ds[:2]
{'image': DeviceArray([[[[ 255, 255, 255],
[ 255, 255, 255],
...,
[ 255, 255, 255],
[ 255, 255, 255]]]], dtype=uint8)}Чтобы использовать тип функции Audio, необходимо установить дополнительный пакет
audioкомандойpip install datasets[audio].
>>> from datasets import Dataset, Features, Audio
>>> audio = ["path/to/audio.wav"] * 10
>>> features = Features({"audio": Audio()})
>>> ds = Dataset.from_dict({"audio": audio}, features=features)
>>> ds = ds.with_format("jax")
>>> ds[0]["audio"]["array"]
DeviceArray([-0.059021 , -0.03894043, -0.00735474, ..., 0.0133667 ,
0.01809692, 0.00268555], dtype=float32)
>>> ds[0]["audio"]["sampling_rate"]
DeviceArray(44100, dtype=int32, weak_type=True)Загрузка данных
В JAX нет встроенных возможностей для загрузки данных, поэтому для загрузки данных с помощью DataLoader вам потребуется использовать библиотеку, например PyTorch, или TensorFlow с использованием tf.data.Dataset. Как указано в документации JAX по этой теме:
«JAX сфокусирован на преобразованиях программ и NumPy с поддержкой ускорителей, поэтому мы не включаем в библиотеку JAX загрузку или обработку данных. Уже существует множество отличных загрузчиков данных, поэтому давайте просто будем использовать их, а не изобретать велосипед. Мы возьмем загрузчик данных PyTorch и создадим небольшой шим, чтобы он работал с массивами NumPy».
Именно поэтому форматирование JAX в datasets так полезно: оно позволяет использовать любую модель из HuggingFace Hub с JAX, не беспокоясь о загрузке данных.
Использование с with_format('jax')
Самый простой способ получить массивы JAX из набора данных — использовать метод with_format('jax'). Допустим, мы хотим обучить нейронную сеть на наборе данных MNIST, доступном
на HuggingFace Hub по адресу https://huggingface.co/datasets/ylecun/mnist.
>>> from datasets import load_dataset
>>> ds = load_dataset("ylecun/mnist")
>>> ds = ds.with_format("jax")
>>> ds["train"][0]
{'image': DeviceArray([[ 0, 0, 0, ...],
[ 0, 0, 0, ...],
...,
[ 0, 0, 0, ...],
[ 0, 0, 0, ...]], dtype=uint8),
'label': DeviceArray(5, dtype=int32)}После того как формат задан, мы можем передавать датасет в модель JAX партиями, используя метод Dataset.iter():
>>> for epoch in range(epochs):
... for batch in ds["train"].iter(batch_size=32):
... x, y = batch["image"], batch["label"]
... ...