ray.data.Dataset.iter_torch_batches
ray.data.Dataset.iter_torch_batches#
- Dataset.iter_torch_batches(*, prefetch_batches: int = 1, batch_size: Optional[int] = 256, dtypes: Optional[Union[torch.dtype, Dict[str, torch.dtype]]] = None, device: Optional[str] = None, collate_fn: Optional[Callable[[Union[numpy.ndarray, Dict[str, numpy.ndarray]]], Any]] = None, drop_last: bool = False, local_shuffle_buffer_size: Optional[int] = None, local_shuffle_seed: Optional[int] = None, prefetch_blocks: int = 0) Iterator[TorchTensorBatchType] [source]#
Return a local batched iterator of Torch Tensors over the dataset.
This iterator will yield single-tensor batches if the underlying dataset consists of a single column; otherwise, it will yield a dictionary of column-tensors. If looking for more flexibility in the tensor conversion (e.g. casting dtypes) or the batch format, try use
iter_batches
directly, which is a lower-level API.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.shape) torch.Size([4, 1]) torch.Size([4, 1]) torch.Size([4, 1])
Time complexity: O(1)
- Parameters
prefetch_batches – The number of batches to fetch ahead of the current batch to fetch. If set to greater than 0, a separate threadpool will be used to fetch the objects to the local node, format the batches, and apply the collate_fn. Defaults to 1. You can revert back to the old prefetching behavior that uses
prefetch_blocks
by settinguse_legacy_iter_batches
to True in the datasetContext.batch_size – 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 ifdrop_last
isFalse
. Defaults to 256.dtypes – The Torch dtype(s) for the created tensor(s); if None, the dtype will be inferred from the tensor data.
device – The device on which the tensor should be placed; if None, the Torch tensor will be constructed on the CPU.
collate_fn – A function to convert a Numpy batch to a PyTorch tensor batch. 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.
drop_last – Whether to drop the last batch if it’s incomplete.
local_shuffle_buffer_size – If non-None, the data will be randomly shuffled using a local in-memory shuffle buffer, and this value will serve as the minimum number of rows that must be in the local in-memory shuffle buffer in order to yield a batch. When there are no more rows to add to the buffer, the remaining rows in the buffer will be drained. This buffer size must be greater than or equal to
batch_size
, and thereforebatch_size
must also be specified when using local shuffling.local_shuffle_seed – The seed to use for the local random shuffle.
- Returns
An iterator over Torch Tensor batches.