Attend the Women in ML Symposium on December 7 Register now

性能提示

使用集合让一切井井有条 根据您的偏好保存内容并对其进行分类。

本文档提供了 TensorFlow Datasets (TFDS) 特定的性能提示。请注意,TFDS 以 tf.data.Dataset 对象的形式提供数据集,因此 tf.data 指南中的建议仍然适用。

对数据集进行基准分析

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

确保指示 batch_size= 以将结果归一化(例如100 iter/sec -> 3200 ex/sec)。这适用于任何可迭代对象(例如 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)

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

自动缓存

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

  • 数据集总大小(所有拆分)已定义且 < 250 MiB
  • shuffle_files 被停用,或仅读取单个分片

可以选择退出自动缓存,方法是在 tfds.load 中将 try_autocaching=False 传递至 tfds.ReadConfig。请参阅数据集目录文档,了解特定数据集是否将使用自动缓存。

将全部数据作为单个张量加载

如果您的数据集能够装入内存,您也可以将整个数据集作为单个张量或 NumPy 数组加载。为此,可以通过设置 batch_size=-1 来批处理单个 tf.Tensor 中的所有样本。然后使用 tfds.as_numpytf.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 以使分片成多个文件的大型数据集获得良好的乱序行为。否则,周期将以相同的顺序读取分片,因此无法真正地随机化数据。

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.ReadConfiginput_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 时,2 个工作进程中的每一个将读取 1/4 的数据。

shuffle_files=True 时,文件将在一个工作进程内(不是工作进程之间)随机打乱顺序。每个工作进程都将在周期之间读取相同比例的文件。

注:使用 tf.distribute.Strategy 时,input_context 可以使用 distribute_datasets_from_function 自动创建

在工作进程之间对数据进行自动分片 (Jax)

借助 Jax,您可以使用 tfds.split_for_jax_processtfds.even_splits API 在工作进程之间分布数据。请参阅 split 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 运算可以提高性能:

解码指南中提供了两个示例的代码。

跳过未使用的特征

如果您只使用特征的一个子集,则可以完全跳过某些特征。如果您的数据集中有许多未使用的特征,则不解码这些特征可以显著提高性能。请参阅 https://www.tensorflow.org/datasets/decode#only_decode_a_sub-set_of_the_features