与 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

请注意，如果未向 with_format 提供 device 参数，则它将使用默认设备，即 jax.devices()[0]。

N 维数组

如果您的数据集由 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 没有内置的数据加载功能，因此你需要使用 PyTorch 并通过 DataLoader 来加载数据，或者使用 TensorFlow 并通过 tf.data.Dataset 来加载数据。引用 JAX 文档中关于此主题的说明： “JAX 专注于程序转换和基于加速器的 NumPy，因此我们不在 JAX 库中包含数据加载或数据处理功能。市面上已经有很多优秀的加载器，所以让我们直接使用它们，而不是重新发明轮子。我们将使用 PyTorch 的数据加载器，并编写一个小型垫片（shim）使其能够与 NumPy 数组配合使用。”

这正是 datasets 中 JAX 格式如此有用的原因，因为它让你能够使用 HuggingFace Hub 上的任何模型与 JAX 配合，而无需担心数据加载部分。

使用 with_format('jax')

从数据集获取 JAX 数组最简单的方法是使用 with_format('jax') 方法。假设我们想在 HuggingFace Hub 上 https://huggingface.co/datasets/ylecun/mnist 提供的 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)}

设置好格式后，我们可以使用Dataset.iter()方法将数据集以批次的形式输入 JAX 模型：

>>> for epoch in range(epochs):
...     for batch in ds["train"].iter(batch_size=32):
...         x, y = batch["image"], batch["label"]
...         ...

Update on GitHub

Datasets

与 JAX 配合使用

数据集格式

N 维数组

其他特征类型

数据加载

使用 with_format('jax')