Hugging Face 的 logo

数据集文档

使用 TensorFlow 处理数据集

Hugging Face's logo
加入 Hugging Face 社区

并获得增强型文档体验

开始使用

使用 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 格式的数据集会输出 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 和 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("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() 方法不同，后者会立即运行映射函数并保存新的或更改的列。

由于整个数据预处理管道可以在 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。这是一个函数，它获取数据集的多个元素并将它们组合成一个批次。当所有元素具有相同的长度时，内置的默认收集器就足够了，但是对于更复杂的任务，可能需要自定义收集器。特别是，许多任务具有具有不同序列长度的样本，这将需要一个 数据收集器 来正确地填充批次。您可以在 transformers NLP 的 示例 和 笔记本 中看到此示例，在这些示例中，可变序列长度非常常见。

如果您发现使用 to_tf_dataset 加载速度很慢，您也可以使用 num_workers 参数。这会启动多个子进程以并行加载数据。此功能是最近才有的，并且仍在实验阶段 - 如果在使用时遇到任何错误，请提交问题！

何时使用 to_tf_dataset

细心的读者可能已经注意到，我们已经提供了两种方法来实现相同的目标 - 如果您想将数据集传递给 TensorFlow 模型，您可以使用 .with_format('tf') 将数据集转换为 Tensor 或 Tensor 的 dict，或者您可以使用 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 中由不同数量的标记组成的输入文本。当您使用具有可变维度的样本创建批次时，标准解决方案是将较短的样本填充到最长的样本的长度。当您使用 to_tf_dataset 从数据集中流式传输样本时，您可以通过您的 collate_fn 将此填充应用于每个批次。但是，如果您想将这样的数据集转换为密集的 Tensor，那么您将不得不将样本填充到整个数据集中最长样本的长度！这可能导致大量的填充，从而浪费内存并降低模型的速度。

注意事项和限制

目前，to_tf_dataset() 始终返回一个批处理数据集 - 我们将很快添加对非批处理数据集的支持！

< > 在 GitHub 上更新

←流式数据 与 PyTorch 结合使用→