parallelising tf.data.Dataset.from_generator

by

Turns out I can use Dataset.map if I make the generator super lightweight (only generating meta data) and then move the actual heavy lighting into a stateless function. This way I can parallelise just the heavy lifting part with .map using a py_func.

Works; but feels a tad clumsy… Would be great to be able to just add num_parallel_calls to from_generator 🙂

def pure_numpy_and_pil_complex_calculation(metadata, label):
  # some complex pil and numpy work nothing to do with tf
  ...

dataset = tf.data.Dataset.from_generator(lightweight_generator,
                                         output_types=(tf.string,   # metadata
                                                       tf.string))  # label

def wrapped_complex_calulation(metadata, label):
  return tf.py_func(func = pure_numpy_and_pil_complex_calculation,
                    inp = (metadata, label),
                    Tout = (tf.uint8,    # (H,W,3) img
                            tf.string))  # label
dataset = dataset.map(wrapped_complex_calulation,
                      num_parallel_calls=8)

dataset = dataset.batch(64)
iter = dataset.make_one_shot_iterator()
imgs, labels = iter.get_next()

More Related Contents:

  1. Meaning of buffer_size in Dataset.map , Dataset.prefetch and Dataset.shuffle
  2. How do I split Tensorflow datasets?
  3. Split a dataset created by Tensorflow dataset API in to Train and Test?
  4. How to extract data/labels back from TensorFlow dataset
  5. tf.data with multiple inputs / outputs in Keras
  6. Parallelism isn’t reducing the time in dataset map
  7. How to Properly Combine TensorFlow’s Dataset API and Keras?
  8. How Can I Define Only the Gradient for a Tensorflow Subgraph?
  9. Numpy to TFrecords: Is there a more simple way to handle batch inputs from tfrecords?
  10. Get the value of some weights in a model trained by TensorFlow
  11. Confused by the behavior of `tf.cond`
  12. TensorFlow: slow performance when getting gradients at inputs
  13. TensorFlow, why there are 3 files after saving the model?
  14. Tensorflow doesn’t seem to see my gpu
  15. tensorflow on GPU: no known devices, despite cuda’s deviceQuery returning a “PASS” result
  16. How can I run Tensorflow on one single core?
  17. Holding variables constant during optimizer
  18. TensorFlow: Dst tensor is not initialized
  19. Is there a simpler way to handle batch inputs from tfrecords?
  20. What is the proper way to install TensorFlow on Apple M1 in 2022
  21. TensorFlow: training on my own image
  22. Split .tfrecords file into many .tfrecords files
  23. How do I convert a directory of jpeg images to TFRecords file in tensorflow?
  24. ImportError: libcublas.so.10.0: cannot open shared object file: No such file or directory
  25. Tensorflow get all variables in scope
  26. Convert between NHWC and NCHW in TensorFlow
  27. How can I convert a trained Tensorflow model to Keras?
  28. How to display custom images in TensorBoard using Keras?
  29. ImportError: libcudnn when running a TensorFlow program
  30. Choosing number of Steps per Epoch

Leave a Comment