TFDS 现在支持 Croissant 🥐 格式！阅读文档了解更多信息。

性能提示

本文档提供 TensorFlow 数据集 (TFDS) 特定的性能提示。请注意，TFDS 将数据集作为 tf.data.Dataset 对象提供，因此 tf.data 指南中的建议仍然适用。

基准数据集

使用 tfds.benchmark(ds) 对任何 tf.data.Dataset 对象进行基准测试。

确保指示 batch_size= 以规范化结果（例如，100 次迭代/秒 -> 3200 个示例/秒）。这适用于任何可迭代对象（例如，tfds.benchmark(tfds.as_numpy(ds))）。

ds = tfds.load('mnist', split='train').batch(32).prefetch()
# Display some benchmark statistics
tfds.benchmark(ds, batch_size=32)
# Second iteration is much faster, due to auto-caching
tfds.benchmark(ds, batch_size=32)

小型数据集（小于 1 GB）

所有 TFDS 数据集都将数据存储在磁盘上的 TFRecord 格式中。对于小型数据集（例如，MNIST、CIFAR-10/-100），从 .tfrecord 读取可能会增加大量开销。

由于这些数据集适合内存，因此可以通过缓存或预加载数据集来显著提高性能。请注意，TFDS 会自动缓存小型数据集（下一节将详细介绍）。

缓存数据集

以下是一个数据管道的示例，该管道在对图像进行归一化后显式缓存数据集。

def normalize_img(image, label):
  """Normalizes images: `uint8` -> `float32`."""
  return tf.cast(image, tf.float32) / 255., label


ds, ds_info = tfds.load(
    'mnist',
    split='train',
    as_supervised=True,  # returns `(img, label)` instead of dict(image=, ...)
    with_info=True,
)
# Applying normalization before `ds.cache()` to re-use it.
# Note: Random transformations (e.g. images augmentations) should be applied
# after both `ds.cache()` (to avoid caching randomness) and `ds.batch()` (for
# vectorization [1]).
ds = ds.map(normalize_img, num_parallel_calls=tf.data.AUTOTUNE)
ds = ds.cache()
# For true randomness, we set the shuffle buffer to the full dataset size.
ds = ds.shuffle(ds_info.splits['train'].num_examples)
# Batch after shuffling to get unique batches at each epoch.
ds = ds.batch(128)
ds = ds.prefetch(tf.data.experimental.AUTOTUNE)

[1] 向量化映射

在迭代此数据集时，由于缓存，第二次迭代将比第一次迭代快得多。

自动缓存

默认情况下，TFDS 会自动缓存（使用 ds.cache()）满足以下约束的数据集

数据集总大小（所有拆分）已定义且 < 250 MiB
shuffle_files 已禁用，或者只读取单个分片

可以通过将 try_autocaching=False 传递给 tfds.ReadConfig（在 tfds.load 中）来选择退出自动缓存。查看数据集目录文档，了解特定数据集是否会使用自动缓存。

将完整数据加载为单个张量

如果数据集适合内存，您还可以将完整数据集加载为单个张量或 NumPy 数组。可以通过将 batch_size=-1 设置为将所有示例批处理到单个 tf.Tensor 中来实现。然后使用 tfds.as_numpy 将 tf.Tensor 转换为 np.array。

(img_train, label_train), (img_test, label_test) = tfds.as_numpy(tfds.load(
    'mnist',
    split=['train', 'test'],
    batch_size=-1,
    as_supervised=True,
))

大型数据集

大型数据集被分片（拆分为多个文件），通常不适合内存，因此不应缓存。

混洗和训练

在训练期间，对数据进行充分混洗非常重要 - 混洗不充分的数据会导致训练精度降低。

除了使用 ds.shuffle 混洗记录外，还应设置 shuffle_files=True，以便对拆分为多个文件的大型数据集获得良好的混洗行为。否则，每个 epoch 会按相同的顺序读取分片，因此数据不会真正随机化。

ds = tfds.load('imagenet2012', split='train', shuffle_files=True)

此外，当 shuffle_files=True 时，TFDS 会禁用 options.deterministic，这可能会略微提高性能。要获得确定性混洗，可以使用 tfds.ReadConfig 选择退出此功能：通过设置 read_config.shuffle_seed 或覆盖 read_config.options.deterministic 来实现。

在工作进程之间自动分片数据（TF）

在多个工作进程上进行训练时，可以使用 tfds.ReadConfig 的 input_context 参数，以便每个工作进程读取数据的子集。

input_context = tf.distribute.InputContext(
    input_pipeline_id=1,  # Worker id
    num_input_pipelines=4,  # Total number of workers
)
read_config = tfds.ReadConfig(
    input_context=input_context,
)
ds = tfds.load('dataset', split='train', read_config=read_config)

这与子分割 API 相辅相成。首先，应用子分割 API：train[:50%] 被转换为要读取的文件列表。然后，在这些文件上应用 ds.shard() 操作。例如，当使用 train[:50%] 和 num_input_pipelines=2 时，两个工作器中的每一个都将读取 1/4 的数据。

当 shuffle_files=True 时，文件在一个工作器内进行混洗，但在工作器之间不进行混洗。每个工作器在 epoch 之间将读取相同的文件子集。

在工作器之间自动分片您的数据（Jax）

使用 Jax，您可以使用 tfds.split_for_jax_process 或 tfds.even_splits API 在工作器之间分配您的数据。请参阅分割 API 指南。

split = tfds.split_for_jax_process('train', drop_remainder=True)
ds = tfds.load('my_dataset', split=split)

tfds.split_for_jax_process 是以下内容的简单别名

# The current `process_index` loads only `1 / process_count` of the data.
splits = tfds.even_splits('train', n=jax.process_count(), drop_remainder=True)
split = splits[jax.process_index()]

更快的图像解码

默认情况下，TFDS 会自动解码图像。但是，在某些情况下，使用 tfds.decode.SkipDecoding 跳过图像解码并手动应用 tf.io.decode_image 操作可能会更高效

在过滤示例（使用 tf.data.Dataset.filter）时，在过滤示例后解码图像。
在裁剪图像时，使用融合的 tf.image.decode_and_crop_jpeg 操作。

这两个示例的代码都可以在解码指南中找到。

跳过未使用的特征

如果您只使用特征的一个子集，则可以完全跳过某些特征。如果您的数据集包含许多未使用的特征，则不解码它们可以显着提高性能。请参阅 https://tensorflowcn.cn/datasets/decode#only_decode_a_sub-set_of_the_features

tf.data 使用了我所有的 RAM！

如果您在 RAM 上受到限制，或者如果您在使用 tf.data 时并行加载许多数据集，以下是一些可以提供帮助的选项

覆盖缓冲区大小

builder.as_dataset(
  read_config=tfds.ReadConfig(
    ...
    override_buffer_size=1024,  # Save quite a bit of RAM.
  ),
  ...
)

这将覆盖传递给 TFRecordDataset（或等效项）的 buffer_size：https://tensorflowcn.cn/api_docs/python/tf/data/TFRecordDataset#args

使用 tf.data.Dataset.with_options 停止神奇行为

https://tensorflowcn.cn/api_docs/python/tf/data/Dataset#with_options

options = tf.data.Options()

# Stop magic stuff that eats up RAM:
options.autotune.enabled = False
options.experimental_distribute.auto_shard_policy = (
  tf.data.experimental.AutoShardPolicy.OFF)
options.experimental_optimization.inject_prefetch = False

data = data.with_options(options)