tf.data：构建 TensorFlow 输入流水线

View 在 TensorFlow.org 上查看 在 Google Colab 中运行 在 GitHub 中查看源代码 下载笔记本

借助 tf.data API，您可以根据简单的可重用片段构建复杂的输入流水线。例如，图像模型的流水线可以聚合来自分布式文件系统中文件的数据，对每个图像应用随机扰动，并将随机选中的图像合并成一个批次进行训练。文本模型的流水线可能涉及从原始文本数据中提取符号，将提取的符号转换为带有查找表的嵌入向量标识符，并将不同长度的序列组合成一个批次。tf.data API 可以处理大量数据、从不同的数据格式读取数据和执行复杂的转换。

tf.data API 引入了一个 tf.data.Dataset 抽象，它表示一个元素序列，其中每个元素都由一个或多个组件组成。例如，在一个图像流水线中，一个元素可能是一个训练样本，有一对表示图像及其标签的张量组件。

您可以通过两种不同的方式创建数据集：

  • 数据源从存储在内存中或存储在一个或多个文件中的数据构造 Dataset。

  • 数据转换从一个或多个 tf.data.Dataset 对象构造数据集。

import tensorflow as tf

2022-12-14 21:17:59.484022: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2022-12-14 21:17:59.484116: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2022-12-14 21:17:59.484126: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.

import pathlib
import os
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

np.set_printoptions(precision=4)

基本机制

要创建输入流水线，您必须从数据源开始。例如，要从内存中的数据构造一个 Dataset，您可以使用 tf.data.Dataset.from_tensors() 或 tf.data.Dataset.from_tensor_slices()。或者，如果您的输入数据以推荐的 TFRecord 格式存储在文件中，则您可以使用 tf.data.TFRecordDataset()。

有了一个 Dataset 对象之后，您可以通过链接 tf.data.Dataset 对象上的方法调用将其转换成一个新的 Dataset。例如，您可以应用逐元素转换（例如 Dataset.map）和多元素转换（例如 Dataset.batch）。有关完整的转换列表，请参阅 tf.data.Dataset 文档。

Dataset 对象是一个 Python 可迭代对象。这使得利用 for 循环使用它的元素成为可能：

dataset = tf.data.Dataset.from_tensor_slices([8, 3, 0, 8, 2, 1])
dataset

<TensorSliceDataset element_spec=TensorSpec(shape=(), dtype=tf.int32, name=None)>

for elem in dataset:
  print(elem.numpy())

8
3
0
8
2
1

或者使用 iter 显式创建一个 Python 迭代器，并利用 next 来使用它的元素：

it = iter(dataset)

print(next(it).numpy())

8

或者，还可以利用 reduce 转换来使用数据集元素，从而减少所有元素以生成单个结果。以下示例演示了如何使用 reduce 转换来计算整数数据集的和。