ray.data.Dataset.iter_batches#

Dataset.iter_batches(*, prefetch_batches: int = 1, batch_size: int | None = 256, batch_format: str | None = 'default', drop_last: bool = False, local_shuffle_buffer_size: int | None = None, local_shuffle_seed: int | None = None, _collate_fn: Callable[[pyarrow.Table | pandas.DataFrame | Dict[str, numpy.ndarray]], CollatedData] | None = None) → Iterable[pyarrow.Table | pandas.DataFrame | Dict[str, numpy.ndarray]][源代码]#

返回一个数据批次的可迭代对象。

此方法对模型训练很有用。

备注

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

示例

import ray

ds = ray.data.read_images("example://image-datasets/simple")

for batch in ds.iter_batches(batch_size=2, batch_format="numpy"):
    print(batch)

{'image': array([[[[...]]]], dtype=uint8)}
...
{'image': array([[[[...]]]], dtype=uint8)}

时间复杂度: O(1)

参数:

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

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

  • batch_format – 如果 "default" 或 "numpy"，批次是 Dict[str, numpy.ndarray]。如果 "pandas"，批次是 pandas.DataFrame。

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

  • local_shuffle_buffer_size – 如果不是 None ，数据会使用本地内存中的随机缓冲区进行随机打乱，这个值作为本地内存中随机缓冲区中必须包含的最小行数，以便生成一个批次。当没有更多的行可以添加到缓冲区时，缓冲区中剩余的行将被排空。

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

返回:

一个数据批次的迭代器。