Datasets 文档
使用 Datasets 与 TensorFlow
并获得增强的文档体验
开始使用
使用 Datasets 与 TensorFlow
本文档快速介绍了如何将 datasets
与 TensorFlow 结合使用,特别关注如何从我们的数据集获取 tf.Tensor
对象,以及如何将数据从 Hugging Face Dataset
对象流式传输到 Keras 方法(如 model.fit()
)。
Dataset 格式
默认情况下,datasets 返回常规 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 格式化的数据集会输出 RaggedTensor
而不是单个张量
>>> 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
类涵盖了广泛的用例 - 它通常从内存中的张量创建,或使用加载函数来读取磁盘或外部存储上的文件。可以使用 map()
方法任意转换数据集,或者可以使用 batch()
和 shuffle()
等方法来创建可用于训练的数据集。这些方法不会以任何方式修改存储的数据 - 相反,这些方法构建了一个数据管道图,该图将在数据集被迭代时执行,通常在模型训练或推理期间。这与 Hugging Face Dataset
对象的 map()
方法不同,后者会立即运行 map 函数并保存新的或更改的列。
由于整个数据预处理管道都可以在 tf.data.Dataset
中编译,因此这种方法允许大规模并行、异步数据加载和训练。但是,图编译的要求可能是一个限制,特别是对于 Hugging Face 分词器,它们通常(尚未!)可作为 TF 图的一部分进行编译。因此,我们通常建议将数据集预处理为 Hugging Face 数据集,其中可以使用任意 Python 函数,然后在之后使用 to_tf_dataset()
转换为 tf.data.Dataset
,以获得可用于训练的批处理数据集。要查看此方法的示例,请参阅 transformers
的 示例 或 notebooks。
使用 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 示例 和 notebooks 中看到这方面的示例,其中可变序列长度非常常见。
如果您发现使用 to_tf_dataset
加载速度很慢,您还可以使用 num_workers
参数。这将启动多个子进程以并行加载数据。此功能是最近推出的,并且仍然有些实验性 - 如果您在使用过程中遇到任何错误,请提交 issue!
何时使用 to_tf_dataset
细心的读者可能已经注意到,此时我们提供了两种方法来实现相同的目标 - 如果您想将数据集传递给 TensorFlow 模型,您可以将数据集转换为 Tensor
或 Tensors
的 dict
,使用 .with_format('tf')
,或者您可以使用 to_tf_dataset()
将数据集转换为 tf.data.Dataset
。这两种方法都可以传递给 model.fit()
,那么您应该选择哪一种呢?
需要认识到的关键是,当您将整个数据集转换为 Tensor
时,它是静态的并且完全加载到 RAM 中。这很简单方便,但如果以下任何一项适用,您可能应该改用 to_tf_dataset()
- 您的数据集太大,无法放入 RAM 中。
to_tf_dataset()
一次只流式传输一个批次,因此即使是非常大的数据集也可以使用此方法处理。 - 您想使用
dataset.with_transform()
或collate_fn
应用随机转换。这在多种模态中很常见,例如训练视觉模型时的图像增强,或训练掩码语言模型时的随机掩码。使用to_tf_dataset()
将在加载批次时应用这些转换,这意味着每次加载相同的样本时,它们将获得不同的增强。这通常是您想要的。 - 您的数据具有可变维度,例如 NLP 中的输入文本,它们由不同数量的 tokens 组成。当您创建具有可变维度的样本批次时,标准解决方案是将较短的样本填充到最长样本的长度。当您使用
to_tf_dataset
从数据集流式传输样本时,您可以通过collate_fn
将此填充应用于每个批次。但是,如果您想将此类数据集转换为密集Tensor
,那么您将不得不将样本填充到整个数据集中最长样本的长度!这可能会导致大量的填充,从而浪费内存并降低模型的速度。
注意事项和限制
目前,to_tf_dataset()
始终返回批处理数据集 - 我们将很快添加对非批处理数据集的支持!