ray.data.Dataset.iter_torch_batches#

Dataset.iter_torch_batches(*, prefetch_batches: int = 1, batch_size: int | None = 256, dtypes: torch.dtype | Dict[str, torch.dtype] | None = None, device: str = 'auto', collate_fn: Callable[[Dict[str, numpy.ndarray]], CollatedData] | None = None, drop_last: bool = False, local_shuffle_buffer_size: int | None = None, local_shuffle_seed: int | None = None) → Iterable[Dict[str, torch.Tensor] | CollatedData][源代码]#

返回一个表示为 Torch 张量的数据批次的可迭代对象。

这个可迭代对象生成类型为 Dict[str, torch.Tensor] 的批次。为了获得更多灵活性，可以调用 iter_batches() 并手动将数据转换为 Torch 张量。

备注

此操作将触发对此数据集执行的延迟转换。

示例

>>> 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])}

使用 collate_fn 来自定义张量批次的创建方式。

>>> 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]])

时间复杂度: O(1)

参数:

  • prefetch_batches – 要预取的批次数量，超过当前批次。如果设置为大于0，将使用一个单独的线程池来将对象获取到本地节点，格式化批次，并应用 collate_fn。默认为1。

  • batch_size – 每个批次中的行数，或 None 以将整个块用作批次（块可能包含不同数量的行）。如果 drop_last 为 False，则最后一个批次可能包含少于 batch_size 行。默认为 256。

  • dtypes – 创建的张量(s)的Torch数据类型(s)；如果为``None``，则从张量数据推断数据类型。不能与``collate_fn``参数一起使用。

  • device – 张量应放置的设备。默认为 “auto”，当数据集传递给 Ray Train 且未提供 collate_fn 时，张量会被移动到适当的设备。否则，默认为 CPU。您不能将此参数与 collate_fn 一起使用。

  • collate_fn – 一个将 Numpy 批次转换为 PyTorch 张量批次的函数。当指定此参数时，用户应在 collate_fn 外部手动处理主机到设备的数据传输。这对于在数据被批处理后进一步处理数据非常有用。潜在的用例包括沿第一个维度以外的维度进行整理，填充不同长度的序列，或一般处理不同长度的张量批次。如果未提供，则使用默认的 collate 函数，该函数仅将 numpy 数组的批次转换为 PyTorch 张量的批次。此 API 仍在实验阶段，可能会发生变化。您不能将此参数与 dtypes 或 device 一起使用。

  • drop_last – 如果最后一个批次不完整，是否丢弃它。

  • local_shuffle_buffer_size – 如果不是 None ，数据将使用本地内存中的随机洗牌缓冲区进行随机洗牌，并且此值作为本地内存中随机洗牌缓冲区中必须存在的最小行数，以便生成一个批次。当没有更多的行可以添加到缓冲区时，缓冲区中剩余的行将被排空。使用本地洗牌时，还必须指定 batch_size 。

  • local_shuffle_seed – 用于本地随机洗牌的种子。

返回:

一个遍历 Torch Tensor 批次的可迭代对象。

参见

Dataset.iter_batches()

调用此方法以手动将您的数据转换为 Torch 张量。