ray.train.tensorflow.prepare_dataset_shard#

ray.train.tensorflow.prepare_dataset_shard(tf_dataset_shard: tf.data.Dataset)[source]#

A utility function that overrides default config for Tensorflow Dataset.

This should be used on a TensorFlow Dataset created by calling iter_tf_batches() on a ray.data.Dataset returned by ray.train.get_dataset_shard() since the dataset has already been sharded across the workers.

Parameters:

tf_dataset_shard (tf.data.Dataset) – A TensorFlow Dataset.

Returns:

  • autosharding turned off

  • prefetching turned on with autotune enabled

Return type:

A TensorFlow Dataset with

PublicAPI (beta): This API is in beta and may change before becoming stable.