2.1 `tf.data.Dataset` API 简介

文档摘要

2.1 API 简介 TensorFlow 数据集 (tf.data): API 简介的核心概念表示一系列元素，每个元素通常是一个或多个张量。可以将数据集视为一个数据管道，数据通过该管道进行转换，最终被馈送到模型中。 API 的主要优点包括：高效性：能够并行处理数据，利用多核 CPU 和 GPU 加速数据加载和转换。灵活性：支持各种数据源和转换操作，可以构建复杂的数据管道。可扩展性：能够处理大型数据集，这些数据集可能无法一次性加载到内存中。创建可以从多种来源创建，包括：内存中的数据：使用。 TFRecord 文件：使用。文本文件：使用。 Python 生成器：使用。 2.1 从内存中的数据创建是最简单的创建数据集的方法之一。

2.1 `tf.data.Dataset` API 简介

TensorFlow 数据集 (tf.data): `tf.data.Dataset` API 简介

1. `tf.data.Dataset` 的核心概念

tf.data.Dataset 表示一系列元素，每个元素通常是一个或多个张量。可以将数据集视为一个数据管道，数据通过该管道进行转换，最终被馈送到模型中。tf.data API 的主要优点包括：

高效性： 能够并行处理数据，利用多核 CPU 和 GPU 加速数据加载和转换。
灵活性： 支持各种数据源和转换操作，可以构建复杂的数据管道。
可扩展性： 能够处理大型数据集，这些数据集可能无法一次性加载到内存中。

2. 创建 `tf.data.Dataset`

tf.data.Dataset 可以从多种来源创建，包括：

内存中的数据： 使用 tf.data.Dataset.from_tensor_slices。
TFRecord 文件： 使用 tf.data.TFRecordDataset。
文本文件： 使用 tf.data.TextLineDataset。
Python 生成器： 使用 tf.data.Dataset.from_generator。

2.1 从内存中的数据创建

tf.data.Dataset.from_tensor_slices 是最简单的创建数据集的方法之一。它接受一个或多个张量，并将它们沿着第一个维度切片，创建数据集。


import tensorflow as tf
import numpy as np
# 创建一些示例数据
data = np.array([[1, 2], [3, 4], [5, 6]])
labels = np.array([0, 1, 0])
# 从数据和标签创建数据集
dataset = tf.data.Dataset.from_tensor_slices((data, labels))
# 打印数据集中的元素
for element in dataset:
    print(element)

输出:


(<tf.Tensor: shape=(2,), dtype=int64, numpy=array([1, 2])>, <tf.Tensor: shape=(), dtype=int64, numpy=0>)
(<tf.Tensor: shape=(2,), dtype=int64, numpy=array([3, 4])>, <tf.Tensor: shape=(), dtype=int64, numpy=1>)
(<tf.Tensor: shape=(2,), dtype=int64, numpy=array([5, 6])>, <tf.Tensor: shape=(), dtype=int64, numpy=0>)

图示:

2.2 从 TFRecord 文件创建

TFRecord 是一种 TensorFlow 推荐的用于存储大型数据集的二进制文件格式。 tf.data.TFRecordDataset 用于从 TFRecord 文件中读取数据。


# 假设我们有一个名为 "example.tfrecord" 的 TFRecord 文件
# 创建一个 TFRecordDataset
dataset = tf.data.TFRecordDataset("example.tfrecord")
# 要读取 TFRecord 文件，需要定义如何解析每个记录
def _parse_function(example_proto):
    # 定义特征的描述
    feature_description = {
        'feature0': tf.io.FixedLenFeature([], tf.int64, default_value=0),
        'feature1': tf.io.FixedLenFeature([], tf.float32, default_value=0.0),
        'feature2': tf.io.FixedLenFeature([10], tf.float32, default_value=tf.zeros([10], dtype=tf.float32)),
    }
    # 解析单个记录
    return tf.io.parse_single_example(example_proto, feature_description)
# 将解析函数应用于数据集
parsed_dataset = dataset.map(_parse_function)
# 打印解析后的数据集中的元素
for element in parsed_dataset:
    print(element)
    break # 只打印第一个元素

注意: 上述代码假设你已经创建了一个名为 "example.tfrecord" 的 TFRecord 文件，并且文件包含具有 'feature0', 'feature1', 和 'feature2' 等特征的记录。创建 TFRecord 文件涉及到序列化数据，超出本文章的范围。

2.3 从文本文件创建

tf.data.TextLineDataset 用于从文本文件中逐行读取数据。


# 假设我们有一个名为 "example.txt" 的文本文件
# 创建一个 TextLineDataset
dataset = tf.data.TextLineDataset("example.txt")
# 打印数据集中的元素
for element in dataset:
    print(element)
    break # 只打印第一个元素

注意: 上述代码假设你已经创建了一个名为 "example.txt" 的文本文件。

2.4 从 Python 生成器创建

tf.data.Dataset.from_generator 允许使用 Python 生成器创建数据集。这对于从自定义数据源读取数据非常有用。


def generator():
    for i in range(5):
        yield i, i**2
dataset = tf.data.Dataset.from_generator(
    generator,
    output_signature=(
        tf.TensorSpec(shape=(), dtype=tf.int32),
        tf.TensorSpec(shape=(), dtype=tf.int32)
    )
)
# 打印数据集中的元素
for element in dataset:
    print(element)

输出:


(<tf.Tensor: shape=(), dtype=int32, numpy=0>, <tf.Tensor: shape=(), dtype=int32, numpy=0>)
(<tf.Tensor: shape=(), dtype=int32, numpy=1>, <tf.Tensor: shape=(), dtype=int32, numpy=1>)
(<tf.Tensor: shape=(), dtype=int32, numpy=2>, <tf.Tensor: shape=(), dtype=int32, numpy=4>)
(<tf.Tensor: shape=(), dtype=int32, numpy=3>, <tf.Tensor: shape=(), dtype=int32, numpy=9>)
(<tf.Tensor: shape=(), dtype=int32, numpy=4>, <tf.Tensor: shape=(), dtype=int32, numpy=16>)

图示:

3. 数据集转换

tf.data.Dataset API 提供了多种转换操作，用于预处理和增强数据。一些常用的转换包括：

map： 将一个函数应用于数据集中的每个元素。
batch： 将连续的元素组合成批次。
shuffle： 随机打乱数据集中的元素。
filter： 根据条件过滤数据集中的元素。
repeat： 重复数据集若干次。
prefetch： 在训练期间预取数据，以提高性能。

3.1 `map` 转换

map 转换用于将一个函数应用于数据集中的每个元素。


# 创建一个示例数据集
dataset = tf.data.Dataset.range(5)
# 定义一个函数，将每个元素乘以 2
def multiply_by_two(x):
    return x * 2
# 使用 map 转换将函数应用于数据集
mapped_dataset = dataset.map(multiply_by_two)
# 打印转换后的数据集中的元素
for element in mapped_dataset:
    print(element)

输出:


tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(2, shape=(), dtype=int64)
tf.Tensor(4, shape=(), dtype=int64)
tf.Tensor(6, shape=(), dtype=int64)
tf.Tensor(8, shape=(), dtype=int64)

图示:

3.2 `batch` 转换

batch 转换用于将连续的元素组合成批次。


# 创建一个示例数据集
dataset = tf.data.Dataset.range(8)
# 使用 batch 转换将元素组合成大小为 3 的批次
batched_dataset = dataset.batch(3)
# 打印批处理后的数据集中的元素
for element in batched_dataset:
    print(element)

输出:


tf.Tensor([0 1 2], shape=(3,), dtype=int64)
tf.Tensor([3 4 5], shape=(3,), dtype=int64)
tf.Tensor([6 7], shape=(2,), dtype=int64)

图示:

3.3 `shuffle` 转换

shuffle 转换用于随机打乱数据集中的元素。


# 创建一个示例数据集
dataset = tf.data.Dataset.range(8)
# 使用 shuffle 转换随机打乱数据集中的元素
# buffer_size 参数指定用于打乱元素的缓冲区大小
shuffled_dataset = dataset.shuffle(buffer_size=8)
# 打印打乱后的数据集中的元素
for element in shuffled_dataset:
    print(element)

注意: 由于 shuffle 是随机的，因此每次运行代码时输出顺序可能会不同。 buffer_size 参数非常重要，它决定了shuffle的程度。如果 buffer_size 小于数据集的大小，则shuffle效果会降低。

3.4 `filter` 转换

filter 转换用于根据条件过滤数据集中的元素。


# 创建一个示例数据集
dataset = tf.data.Dataset.range(10)
# 定义一个函数，用于过滤偶数
def is_even(x):
    return tf.equal(x % 2, 0)
# 使用 filter 转换过滤数据集中的元素
filtered_dataset = dataset.filter(is_even)
# 打印过滤后的数据集中的元素
for element in filtered_dataset:
    print(element)

输出:


tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(2, shape=(), dtype=int64)
tf.Tensor(4, shape=(), dtype=int64)
tf.Tensor(6, shape=(), dtype=int64)
tf.Tensor(8, shape=(), dtype=int64)

图示:

3.5 `repeat` 转换

repeat 转换用于重复数据集若干次。


# 创建一个示例数据集
dataset = tf.data.Dataset.range(3)
# 使用 repeat 转换重复数据集 2 次
repeated_dataset = dataset.repeat(2)
# 打印重复后的数据集中的元素
for element in repeated_dataset:
    print(element)

输出:


tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(2, shape=(), dtype=int64)
tf.Tensor(0, shape=(), dtype=int64)
tf.Tensor(1, shape=(), dtype=int64)
tf.Tensor(2, shape=(), dtype=int64)

3.6 `prefetch` 转换

prefetch 转换用于在训练期间预取数据，以提高性能。它允许数据准备与模型训练并行进行。


# 创建一个示例数据集
dataset = tf.data.Dataset.range(1000)
# 使用 prefetch 转换预取数据
# tf.data.AUTOTUNE 允许 TensorFlow 动态调整预取缓冲区的大小
prefetched_dataset = dataset.prefetch(tf.data.AUTOTUNE)
# 在训练循环中使用 prefetched_dataset
# ...

注意: prefetch 转换通常是数据管道的最后一个转换，以确保数据在需要时可用。

4. 构建数据管道

可以将多个转换操作链接在一起，构建复杂的数据管道。


# 创建一个示例数据集
data = np.array([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10]])
labels = np.array([0, 1, 0, 1, 0])
dataset = tf.data.Dataset.from_tensor_slices((data, labels))
# 定义一个函数，用于缩放数据
def scale_data(data, label):
    return data / 255.0, label
# 构建数据管道
dataset = (
    dataset
    .shuffle(buffer_size=5)  # 打乱数据
    .map(scale_data)       # 缩放数据
    .batch(2)             # 组合成批次
    .prefetch(tf.data.AUTOTUNE) # 预取数据
)
# 在训练循环中使用数据集
for data, labels in dataset:
    print("Data:", data)
    print("Labels:", labels)
    break # 只打印第一个批次

图示:

5. 总结

tf.data.Dataset API 是 TensorFlow 中构建高效数据管道的关键工具。通过使用各种数据集创建方法和转换操作，可以轻松地加载、预处理和增强数据，使其能够无缝地输入到 TensorFlow 模型中。掌握 tf.data.Dataset API 对于构建高性能的 TensorFlow 应用至关重要。本文只是一个介绍， tf.data 库还有许多高级特性，例如使用 tf.data.Dataset.zip 合并多个数据集，使用 tf.data.Dataset.interleave 交错读取多个数据集等。深入研究官方文档和示例可以帮助你充分利用 tf.data API 的强大功能。

作者与出处

原作者：灏天文库

来源：灏天文库

原文： https://aiknowledge.cn/collection/219-tensorflow%E6%B7%B1%E5%BA%A6%E5%AD%A6%E4%B9%A0%E5%AE%9E%E6%88%98%E6%95%99%E7%A8%8B

整理：灏天文库整理

由灏天文库平台收录，内容或由平台用户上传，仅供学习交流

2.1 `tf.data.Dataset` API 简介

文档摘要