Skip to main content

Data Loading for Distributed Training in TensorFlow

This page provides a practical guide to managing Datasets and DataGenerators for distributed TensorFlow model training. It addresses the challenges introduced in the main data loading page.

This page covers:

The documentation concludes with a complete example of optimised data loading, as well as a practical implementation on Jean Zay using a Jupyter Notebook.

Preliminary remark

This documentation is limited to TensorFlow releases 2.x and above.

Datasets​

In TensorFlow, database handling is managed using the tf.data.Dataset object.

Predefined Datasets in TensorFlow​

TensorFlow provides a set of ready-to-use datasets in the tensorflow-datasets module. Some of these are already available in Jean Zay's shared DSDIR space, in the tensorflow_datasets directory.

When loading, it is possible to distinguish between training and validation data using the split option.

For example, for the ImageNet2012 dataset:

import tensorflow as tf
import tensorflow_datasets as tfds

# load imagenet2012 dataset from DSDIR
tfds_dir = os.environ['DSDIR'] + '/tensorflow_datasets/'

# load data for training
dataset = tfds.load('imagenet2012',
                    data_dir=tfds_dir,
                    split='train')

# load data for testing
dataset_test = tfds.load('imagenet2012',
                    data_dir=tfds_dir,
                    split='test')

Each loading function then offers dataset-specific features (data quality, extracting a subset of the data, etc.). Please consult the official documentation for more details.

Remarks
  • Datasets in the tensorflow-datasets collection are in TFRecord format.
  • The $DSDIR/tensorflow_datasets directory can be enriched upon request to IDRIS support (assist@idris.fr).
important

Some functions allow downloading datasets online using the download=True argument. We remind you that Jean Zay compute nodes do not have internet access, so such operations must be performed in advance from a login node or a pre/post-processing node.

Custom Datasets​

There are several ways to generate a tf.data.Dataset object from input files. The most common are:

  • tf.data.Dataset.list_files() for reading standard formats:
import tensorflow as tf

# locate Places365 dataset in DSDIR
places365_dir = os.environ['DSDIR']+'/Places365-Standard/data_256/'

# extract images from places beginning with `a`
dataset = tf.data.Dataset.list_files(places365_dir + 'a/*.jpg')

By default, the list_files function organises files randomly. To disable this feature, you can specify the argument shuffle=False.

Remark

The list_files function applies the Python glob function to its input parameter. If you already have the list of paths to the input files, it is preferable to extract the data using the tf.data.Dataset.from_tensor_slices(glob_list) function. To randomly shuffle the files, you must first apply a function such as random.Random(seed).shuffle(glob_list) to the list of paths. In the case of a .shard() function call (see below), the seed value must be the same on each worker.

  • tf.data.TFRecordDataset() for reading files in TFRecord format:
import tensorflow as tf

# locate imagenet dataset in DSDIR
imagenet_path = os.environ['DSDIR']+'/imagenet/'

# extract training data from dataset
dataset = tf.data.TFRecordDataset(imagenet_path + 'train*')
Remarks
  • You can create your own TFRecord format files from a dataset by following this tutorial. The DSDIR directory can also be enriched with TFRecord-format Datasets upon request to IDRIS support [assist@idris.fr].
  • It is not possible to determine the size of a Dataset loaded from a TFRecord using the len(dataset) command (this returns an error). You must retrieve this information from the source database used to create the TFRecord.
  • If your Dataset is partitioned across multiple TFRecord files, it is recommended to use the interleave() function when reading the Dataset. Partitioned TFRecords are then read in parallel, which is faster. This also avoids redundant reads between different processes when using data parallelism. 
    dataset = tf.data.Dataset.list_files(pattern)
    dataset = dataset.interleave(tf.data.TFRecordDataset,
                                 num_parallel_calls=tf.data.AUTOTUNE, deterministic=False)
    The num_parallel_calls=tf.data.AUTOTUNE option also enables multithreading. The number of CPUs used is automatically determined based on the available CPUs. The deterministic=False option improves processing performance when the order of operations is not important.

Transformations​

Map Function​

Input data can be transformed using the tf.data.Dataset.map function. The transformation function is defined by the user and then passed to the map function. For example, to resize images from the ImageNet2012 dataset:

# define resizing function
def img_resize(x) :
    return tf.image.resize(x,[300,300])

# apply transformation to dataset
dataset = dataset.map(img_resize,
                      num_parallel_calls=tf.data.AUTOTUNE,
                      deterministic=False)
Remarks
  • The num_parallel_calls=tf.data.AUTOTUNE option enables multithreading for the transformation step. The number of CPUs used is automatically determined based on the available CPUs.
  • The deterministic=False option improves processing performance when the order of operations is not important.

Cache Function​

Some transformation operations are deterministic (data reading, normalisation, resizing, etc.), meaning they generate the same transformed data each time, while others are related to Data Augmentation and are random. Deterministic transformations generally involve heavy processing that generates low-volume data. This transformed data can be stored in a memory cache to avoid repeating the transformation unnecessarily at each epoch. To do this, use the cache() function:

# define random transformation
def random_trans(x) :
    return ...

# define deterministic transformation
def determ_trans(x) :
    return ...

# apply deterministic transformation to dataset and store them in cache
dataset = dataset.map(determ_trans, num_parallel_calls=tf.data.AUTOTUNE).cache()

# apply random transformation to dataset at each epoch
dataset = dataset.map(random_trans, num_parallel_calls=tf.data.AUTOTUNE)

By default, the cache memory used is the RAM of the CPU node. It is also possible to store the transformed data in an external file. To do this, simply specify the path to the storage file when calling the cache() function: .cache(path/to/file).

Remark

The cache() function is useful if the CPU's RAM can hold the entire dataset. The time savings with .cache(path/to/file) to a storage file are limited, except in extreme cases where loading the data and its deterministic transformations would be a major bottleneck in the training loop.

Configuration and Optimisation of Preprocessing Steps​

Random Processing of Input Data​

The shuffle() function enables random shuffling of the data at each epoch during training. This function takes the buffer_size argument, which sets the size of the memory buffer used for this operation. Ideally, the memory buffer should be able to hold the entire dataset to avoid any bias in the random processing. In practice, a dataset is often too large, and a smaller buffer size (1,000 or 10,000 pointers) allows you to overcome the memory constraint.

Data Distribution Across Multiple Processes for Distributed Training​

To distribute data across multiple processes (or workers) for distributed training, use the shard() function. This function takes as input arguments the number of workers and the global worker index.

These values can be retrieved from the Slurm compute configuration as follows:

import tensorflow as tf
import idr_tf # IDRIS package available in all TensorFlow modules to interface with Slurm

# load dataset
dataset = ...

# get number of processes/workers
num_workers = idr_tf.size    # or hvd.size() if Horovod is used
worker_index = idr_tf.rank   # or hvd.rank() if Horovod is used

# distribute dataset
dataset = dataset.shard(num_workers,worker_index)

Here, we use the idr_tf library to retrieve information about the Slurm compute environment. This library, developed by IDRIS, is available in all TensorFlow modules on Jean Zay.

If you use the tf.distribute.MultiWorkerMirroredStrategy distribution strategy, Autosharding is configured by default. Therefore, you should not use .shard(). If you wish to disable Autosharding, please refer to the TensorFlow documentation.

Remarks
  • When calling the shard function, each worker only reads a portion of the dataset.
  • In a distributed configuration, it is important to perform the shuffling step after the distribution step. Otherwise, the entire dataset will be read by each worker, reducing performance. Generally, the shard function should be applied as early as possible.

Optimising Resource Usage During Training​

The batch size is defined using the batch() function. An optimal batch size ensures efficient use of compute resources, i.e., maximising GPU memory usage and evenly distributing the workload across GPUs.

The drop_remainder=True option allows you to ignore the last batch if its size is smaller than the requested batch size. The batch function is always called after the shuffle function, to ensure unique batches at each epoch.

Overlapping Data Transfers and Computation​

It is possible to optimise batch transfers from CPU to GPU by overlapping data transfers with computation. The prefetch() function allows you to preload the next batches to be processed during training. The number of preloaded batches is controlled by the buffer_size argument. For automatic buffer_size definition, you can set buffer_size=tf.data.AUTOTUNE.

Remarks
  • The .repeat() function allows you to repeat a Dataset indefinitely in a DataGenerator, or to repeat it for a given number of epochs using .repeat(num_epoch).
    It is optional: omitting this function is equivalent to .repeat(1). It can be useful if you want to set a number of training iterations rather than a number of epochs. It can be called anywhere in the chain; the result will be the same.
  • If the dataset contains corrupted data that generates an error, you can ignore errors using the following command, to be applied after .map():
    dataset = dataset.apply(tf.data.experimental.ignore_errors())

Complete Example of Optimised Data Loading​

Here is a complete example of optimised ImageNet dataset loading for distributed training on Jean Zay:

import tensorflow as tf
import idr_tf # IDRIS package available in all TensorFlow modules to interface with Slurm
import os
import glob
import random

IMG_WIDTH=320
IMG_HEIGHT=320
def decode_img(file_path):
    # parse label
    label = tf.strings.split(file_path, sep='/')[-2]
    # read input file
    img = tf.io.read_file(file_path)
    # decode jpeg format (channel=3 for RGB, channel=0 for Grayscale)
    img = tf.image.decode_jpeg(img, channels=3)
    # convert to [0,1] format for TensorFlow compatibility
    img = tf.image.convert_image_dtype(img, tf.float32)
    # resize image
    return label, tf.image.resize(img, [IMG_WIDTH, IMG_HEIGHT])

# Create a generator
rng = tf.random.Generator.from_seed(123, alg='philox')

def randomized_preprocessing(label, img):
    # randomly adjust image contrast - Data Augmentation
    contrast_factor = random.random() + 1.0
    img = tf.image.adjust_contrast(img,contrast_factor=contrast_factor)
    img = tf.image.stateless_random_flip_left_right(img,rng.make_seeds(2)[0])
    return label, img

# configuration
num_epochs = 3
batch_size = 64
shuffling_buffer_size = 5000
num_parallel_calls = tf.data.AUTOTUNE
prefetch_factor = tf.data.experimental.AUTOTUNE

# locate Places365 dataset in DSDIR and list them
places365_path = glob.glob(os.environ['DSDIR']+"/Places365-Standard/data_256/**/*.jpg", recursive=True)
random.Random(123).shuffle(places365_path)

# create a dataset object from path
dataset = tf.data.Dataset.from_tensor_slices(places365_path)
# distribute dataset
num_workers = idr_tf.size
worker_index = idr_tf.rank
dataset = dataset.shard(num_workers,worker_index).shuffle(shuffling_buffer_size)

# deterministic transformation
dataset = dataset.map(decode_img, num_parallel_calls=num_parallel_calls, deterministic=False)
# cache function only relevant for small to medium datasets
dataset = dataset.cache()
# random transformations
dataset = dataset.map(randomized_preprocessing, num_parallel_calls=num_parallel_calls, deterministic=False)

dataset = dataset.batch(batch_size, drop_remainder=True).prefetch(prefetch_factor)

## Repeat dataset num_epochs times
#dataset = dataset.repeat(num_epochs)
#for label, img in dataset:
#    train_step(label, img)
## equivalent to:
for epoch in range(num_epochs):
    for label, img in dataset:
        train_step(label, img)

Here is an example of the loading chain when loading a single TFRecord file:

dataset = tf.data.TFRecordDataset(input_file)
dataset = dataset.shard(num_workers, worker_index)
dataset = dataset.shuffle(shuffle_buffer_size)
dataset = dataset.map(parser_fn, num_parallel_calls=num_map_threads)
dataset = dataset.batch(batch_size, drop_remainder=True).prefetch(prefetch_factor)
dataset = dataset.repeat(num_epochs)

Finally, an example of the loading chain when loading a Dataset partitioned across multiple TFRecords:

dataset = Dataset.list_files(pattern)
dataset = dataset.shard(num_workers, worker_index)
dataset = dataset.shuffle(shuffle_buffer_size)
dataset = dataset.interleave(tf.data.TFRecordDataset,
                             cycle_length=num_readers, block_length=1)
dataset = dataset.map(parser_fn, num_parallel_calls=num_map_threads)
dataset = dataset.batch(batch_size, drop_remainder=True).prefetch(prefetch_factor)
dataset = dataset.repeat(num_epochs)

Practical Implementation on Jean Zay​

To put the above documentation into practice and see the benefits of each feature offered by the TensorFlow Dataset, you can retrieve the notebook_data_preprocessing_tensorflow-eng.ipynb notebook from DSDIR. For example, to copy it to your WORK directory:

cp $DSDIR/examples_IA/Tensorflow_parallel/notebook_data_preprocessing_tensorflow-eng.ipynb $WORK

You can then run the Notebook from a Jean Zay login node (see our JupyterHub access documentation for more information on using Notebooks on Jean Zay).

Creating TFRecord files and using them are also covered.

Official Documentation​