iter_torch_batches#

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

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

Note

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

Examples

>>> import ray
>>> for batch in ray.data.range(
...     12,
... ).iter_torch_batches(batch_size=4):
...     print(batch)
{'id': tensor([0, 1, 2, 3])}
{'id': tensor([4, 5, 6, 7])}
{'id': tensor([ 8,  9, 10, 11])}

Use the collate_fn to customize how the tensor batch is created.

>>> from typing import Any, Dict
>>> import torch
>>> import numpy as np
>>> import ray
>>> def collate_fn(batch: Dict[str, np.ndarray]) -> Any:
...     return torch.stack(
...         [torch.as_tensor(array) for array in batch.values()],
...         axis=1
...     )
>>> dataset = ray.data.from_items([
...     {"col_1": 1, "col_2": 2},
...     {"col_1": 3, "col_2": 4}])
>>> for batch in dataset.iter_torch_batches(collate_fn=collate_fn):
...     print(batch)
tensor([[1, 2],
        [3, 4]])

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 and applies collate_fn. 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 and collates 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 number of rows). The final batch may include fewer than batch_size rows if drop_last is False. Defaults to 256.
dtypes (torch.dtype | Dict[str, torch.dtype] | None) – The Torch dtype(s) for the created tensor(s); if None, the dtype is inferred from the tensor data. You can’t use this parameter with collate_fn.
device (str | torch.device | int | Literal['auto']) – The device on which the tensor should be placed. Defaults to “auto” which moves the tensors to the appropriate device when the Dataset is passed to Ray Train, and to CPU otherwise. When used together with collate_fn, the device transfer only applies if the collate_fn output is a TensorBatchType; for other output types, you must handle the device transfer manually.
collate_fn (Callable[[Dict[str, ndarray]], CollatedData] | CollateFn | None) – A function to convert a Numpy batch to a PyTorch tensor batch. If the output of collate_fn is a TensorBatchType, it is automatically moved to the target device (see device); otherwise, you must handle the device transfer manually. This is useful for further processing the data after it has been batched. Potential use cases include collating along a dimension other than the first, padding sequences of various lengths, or generally handling batches of different length tensors. If not provided, the default collate function is used which simply converts the batch of numpy arrays to a batch of PyTorch tensors. This API is still experimental and is subject to change. You can’t use this parameter in conjunction with dtypes.
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.
pin_memory (bool) – [Alpha] Pin memory if True and the collate_fn output is a TensorBatchType. It is recommended to use this flag to pin memory instead of manually pinning memory in the collate_fn.

Returns:

An iterable over Torch Tensor batches.

Return type:

Iterable[Dict[str, torch.Tensor] | CollatedData]