iter_tf_batches#

Dataset.iter_tf_batches(*, prefetch_batches: int = 1, batch_size: int | None = 256, dtypes: tf.dtypes.DType | Dict[str, tf.dtypes.DType] | None = None, drop_last: bool = False, local_shuffle_buffer_size: int | None = None, local_shuffle_seed: int | None = None) → Iterable[tf.Tensor | Dict[str, tf.Tensor]][source]#

Return an iterable over batches of data represented as TensorFlow tensors.

This iterable yields batches of type Dict[str, tf.Tensor]. For more flexibility, call iter_batches() and manually convert your data to TensorFlow tensors.

Tip

If you don’t need the additional flexibility provided by this method, consider using to_tf() instead. It’s easier to use.

Note

This operation will trigger execution of the lazy transformations performed on this dataset.

Examples

>>> import ray
>>> ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
>>> tf_dataset = ds.iter_tf_batches(batch_size=2)
>>> for batch in tf_dataset:
...     print(batch["sepal length (cm)"], batch["target"])
...     break
tf.Tensor([5.1 4.9], shape=(2,), dtype=float64) tf.Tensor([0 0], shape=(2,), dtype=int64)

Time complexity: O(1)

Parameters:

prefetch_batches (int) – The number of batches to fetch ahead of the current batch. Ray Data prepares batches in background threads while your code consumes the current one. If greater than 0, a thread pool formats the batches. The pool has prefetch_batches threads, up to a limit of 4. To change the limit, set the RAY_DATA_MAX_FORMAT_THREADPOOL_NUM_WORKERS environment variable. If 0, Ray Data doesn’t prefetch blocks and formats each batch without the thread pool. Defaults to 1.
batch_size (int | None) – The number of rows in each batch, or None to use entire blocks as batches (blocks may contain different numbers of rows). The final batch may include fewer than batch_size rows if drop_last is False. Defaults to 256.
dtypes (tf.dtypes.DType | Dict[str, tf.dtypes.DType] | None) – The TensorFlow dtype(s) for the created tensor(s); if None, the dtype is inferred from the tensor data.
drop_last (bool) – Whether to drop the last batch if it’s incomplete.
local_shuffle_buffer_size (int | None) – If not None, Ray Data randomly shuffles the data using a local in-memory shuffle buffer, and this value serves as the minimum number of rows that must be in the buffer to yield a batch. When there are no more rows to add to the buffer, Ray Data drains the remaining rows. A larger buffer makes the shuffle more random but might increase the latency to the first batch. If this value is smaller than batch_size, Ray Data rounds it up to batch_size. Local shuffling raises a ValueError if batch_size is None.
local_shuffle_seed (int | None) – The seed to use for the local random shuffle.

Returns:

An iterable over TensorFlow Tensor batches.

Return type:

Iterable[tf.Tensor | Dict[str, tf.Tensor]]