Datasets documentation
使用 Datasets 与 TensorFlow
使用 Datasets 与 TensorFlow
本文档简要介绍了如何使用 datasets 与 TensorFlow 配合，重点说明如何从我们的数据集中获取 tf.Tensor 对象，以及如何将数据从 Hugging Face Dataset 对象流式传输到 Keras 方法（例如 model.fit()）。
数据集格式
默认情况下，数据集返回的是普通的 Python 对象：整数、浮点数、字符串、列表等。
若要获取 TensorFlow 张量，可以将数据集的格式设置为 tf：
>>> from datasets import Dataset
>>> data = [[1, 2],[3, 4]]
>>> ds = Dataset.from_dict({"data": data})
>>> ds = ds.with_format("tf")
>>> ds[0]
{'data': <tf.Tensor: shape=(2,), dtype=int64, numpy=array([1, 2])>}
>>> ds[:2]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
array([[1, 2],
[3, 4]])>}Dataset 对象是 Arrow 表的封装，它允许从数据集的数组中快速读取为 TensorFlow 张量。
这在将数据集转换为Tensor对象的字典，或编写生成器以从中加载 TF 样本时可能很有用。如果您希望将整个数据集转换为Tensor，只需查询完整的数据集即可：
>>> ds[:]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
array([[1, 2],
[3, 4]])>}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("tf")
>>> ds[0]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
array([[1, 2],
[3, 4]])>}否则，TensorFlow 格式的 dataset 会输出一个 RaggedTensor，而不是单个 tensor：
>>> 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': <tf.RaggedTensor [[1, 2], [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("tf")
>>> ds[0]
{'data': <tf.Tensor: shape=(2, 2), dtype=int64, numpy=
array([[1, 2],
[3, 4]])>}
>>> ds[:2]
{'data': <tf.Tensor: shape=(2, 2, 2), dtype=int64, numpy=
array([[[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("tf")
>>> ds[:3]
{'label': <tf.Tensor: shape=(3,), dtype=int64, numpy=array([0, 0, 1])>}字符串和二进制对象也受支持：
>>> from datasets import Dataset, Features
>>> text = ["foo", "bar"]
>>> data = [0, 1]
>>> ds = Dataset.from_dict({"text": text, "data": data})
>>> ds = ds.with_format("tf")
>>> ds[:2]
{'text': <tf.Tensor: shape=(2,), dtype=string, numpy=array([b'foo', b'bar'], dtype=object)>,
'data': <tf.Tensor: shape=(2,), dtype=int64, numpy=array([0, 1])>}您也可以显式地格式化某些列，而将其他列保持未格式化状态：
>>> ds = ds.with_format("tf", columns=["data"], output_all_columns=True)
>>> ds[:2]
{'data': <tf.Tensor: shape=(2,), dtype=int64, numpy=array([0, 1])>,
'text': ['foo', 'bar']}字符串和二进制对象保持不变，因为 PyTorch 仅支持数字。
要使用 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("tf")
>>> ds[0]
{'image': <tf.Tensor: shape=(512, 512, 4), dtype=uint8, numpy=
array([[[255, 215, 106, 255],
[255, 215, 106, 255],
...,
[255, 255, 255, 255],
[255, 255, 255, 255]]], dtype=uint8)>}
>>> ds[:2]
{'image': <tf.Tensor: shape=(2, 512, 512, 4), dtype=uint8, numpy=
array([[[[255, 215, 106, 255],
[255, 215, 106, 255],
...,
[255, 255, 255, 255],
[255, 255, 255, 255]]]], dtype=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("tf")
>>> ds[0]["audio"]["array"]
<tf.Tensor: shape=(202311,), dtype=float32, numpy=
array([ 6.1035156e-05, 1.5258789e-05, 1.6784668e-04, ...,
-1.5258789e-05, -1.5258789e-05, 1.5258789e-05], dtype=float32)>
>>> ds[0]["audio"]["sampling_rate"]
<tf.Tensor: shape=(), dtype=int32, numpy=44100>数据加载
虽然你可以通过对数据集进行索引来加载单个样本和批次，但如果你想使用 Keras 方法，如 fit() 和 predict()，这种方式将不起作用。你可以编写一个生成器函数，从数据集中打乱并加载批次，然后在该生成器上调用 fit()，但这听起来像是很多不必要的工作。相反，如果你想从数据集中实时流式传输数据，我们建议使用 to_tf_dataset() 方法将你的数据集转换为 tf.data.Dataset。
tf.data.Dataset 类涵盖了广泛的应用场景——它通常由内存中的 Tensor 创建，或者使用加载函数来读取磁盘或外部存储上的文件。数据集可以通过 map() 方法进行任意转换，也可以使用 batch() 和 shuffle() 等方法来创建可直接用于训练的数据集。这些方法不会以任何方式修改存储的数据——相反，这些方法会构建一个数据管道图，该图将在迭代数据集时执行，通常是在模型训练或推理期间。这与 Hugging Face Dataset 对象的 map() 方法不同，后者会立即运行映射函数并保存新增或更改的列。
由于整个数据预处理流水线都可以编译为 tf.data.Dataset，这种方法能够实现大规模并行、异步的数据加载和训练。然而，图编译的要求可能成为一种限制，尤其是对于 Hugging Face 分词器，它们通常（目前还）无法作为 TF 图的一部分进行编译。因此，我们通常建议先将数据集作为 Hugging Face 数据集进行预处理，在此过程中可以使用任意 Python 函数，然后再使用 to_tf_dataset() 将其转换为 tf.data.Dataset，以获得可直接用于训练的批处理数据集。要了解这种方法的示例，请参阅 transformers 的 示例 或 笔记本。
使用 to_tf_dataset()
使用 to_tf_dataset() 非常简单。一旦数据集完成预处理并准备就绪，只需按以下方式调用即可：
>>> from datasets import Dataset
>>> data = {"inputs": [[1, 2],[3, 4]], "labels": [0, 1]}
>>> ds = Dataset.from_dict(data)
>>> tf_ds = ds.to_tf_dataset(
columns=["inputs"],
label_cols=["labels"],
batch_size=2,
shuffle=True
)此处返回的 tf_ds 对象现已完全准备好用于训练，可以直接传递给 model.fit()。请注意，您在创建数据集时已设置了批处理大小，因此在调用 fit() 时无需再指定它：
>>> model.fit(tf_ds, epochs=2)有关参数的完整说明，请参阅 to_tf_dataset() 文档。在许多情况下，您还需要在调用中添加 collate_fn。这是一个函数，它接收数据集中的多个元素并将它们组合成一个批次。当所有元素长度相同时，内置的默认 collator 就足够了，但对于更复杂的任务，可能需要自定义 collator。特别是，许多任务的样本具有不同的序列长度，这将需要一个能够正确填充批次的 data collator。您可以在 transformers NLP examples 和 notebooks 中看到此类示例，其中可变序列长度非常常见。
如果您发现使用 to_tf_dataset 加载数据较慢，也可以使用 num_workers 参数。该参数会启动多个子进程以并行加载数据。此功能较新，仍具有一定实验性——如果您在使用过程中遇到任何 bug，请提交 issue！
何时使用 to_tf_dataset
细心的读者可能已经注意到，我们提供了两种方法来实现相同的目标——如果你想将数据集传递给 TensorFlow 模型，既可以使用 .with_format('tf') 将数据集转换为 Tensor 或 Tensors 的 dict，也可以使用 to_tf_dataset() 将数据集转换为 tf.data.Dataset。这两种方式都可以传递给 model.fit()，那么你应该选择哪一种呢？
需要明确的一点是，当你将整个数据集转换为 Tensor 时，它是静态的，并且会完全加载到 RAM 中。这种方式简单且方便，但如果符合以下任一情况，你最好改用 to_tf_dataset()：
- 您的数据集太大，无法装入内存。
to_tf_dataset()每次仅流式传输一个批次，因此即使是非常大的数据集也能通过此方法进行处理。 - 你希望使用
dataset.with_transform()或collate_fn来应用随机变换。这在多种模态中很常见，例如训练视觉模型时的图像增强，或训练掩码语言模型时的随机掩码。使用to_tf_dataset()会在加载批次时应用这些变换，这意味着每次加载相同的样本时，它们都会获得不同的增强。这通常是你想要的效果。 - 您的数据具有可变维度，例如 NLP 中输入文本由数量不等的 token 组成。当使用具有可变维度的样本创建批次时，标准解决方案是将较短的样本填充到最长样本的长度。当您使用
to_tf_dataset从数据集中流式传输样本时，可以通过collate_fn对每个批次应用此填充。但是，如果您希望将此类数据集转换为稠密的Tensor，则必须将样本填充到 整个数据集中最长样本的长度！这可能会导致大量的填充，从而浪费内存并降低模型的速度。
注意事项和局限性
目前，to_tf_dataset() 始终返回批处理数据集 - 我们将很快添加对非批处理数据集的支持！