В этом документе содержатся советы по повышению производительности для наборов данных 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 ГБ)
Все наборы данных 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 МБ.
-
shuffle_files
отключен, или читается только один фрагмент
Можно отказаться от автоматического кэширования, передав try_autocaching=False
в tfds.ReadConfig
в tfds.load
. Просмотрите документацию по каталогу наборов данных, чтобы узнать, будет ли конкретный набор данных использовать автоматическое кэширование.
Загрузка полных данных в виде одного тензора
Если ваш набор данных помещается в память, вы также можете загрузить полный набор данных в виде одного массива Tensor или 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
, чтобы получить хорошее поведение при перемешивании больших наборов данных, разбитых на несколько файлов. В противном случае эпохи будут считывать осколки в одном и том же порядке, поэтому данные не будут по-настоящему рандомизированы.
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)
При обучении нескольких рабочих процессов вы можете использовать аргумент input_context
функции tfds.ReadConfig
, чтобы каждый рабочий процесс считывал подмножество данных.
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)
Это дополнение к subsplit API. Во-первых, применяется subplit API: train[:50%]
преобразуется в список файлов для чтения. Затем к этим файлам применяется ds.shard()
. Например, при использовании train[:50%]
с num_input_pipelines=2
каждый из двух рабочих будет считывать 1/4 данных.
Когда shuffle_files=True
файлы перемешиваются внутри одного рабочего процесса, но не между рабочими процессами. Каждый рабочий процесс будет читать одно и то же подмножество файлов между эпохами.
Автоматическое разделение ваших данных между работниками (Jax)
С Jax вы можете использовать API tfds.split_for_jax_process
или tfds.even_splits
для распределения ваших данных между работниками. См. руководство по раздельному 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://www.tensorflow.org/datasets/decode#only_decode_a_sub-set_of_the_features .