Function batch

axlearn/common/input_tf_data.py:839–961

Returns a function that generates a tf.data.Dataset object. Note: batch(is_training=True) requires sufficient number of examples per host. When your data is too small, you should add `ds = ds.repeat()` before your batch. Args: global_batch_size: The global physical batch si

(
    global_batch_size: int,
    *,
    is_training: bool,
    pad_example_fn: PadExampleFn,
    global_logical_batch_size: Optional[int] = None,
    logical_feed_indices: Optional[Sequence[int]] = None,
    prefetch_buffer_size: Optional[int] = None,
    post_batch_processor: Optional[ConfigOr[DatasetToDatasetFn]] = None,
    repeat: Optional[int] = None,
)

def batch(
    global_batch_size: int,
    *,
    is_training: bool,
    pad_example_fn: PadExampleFn,
    global_logical_batch_size: Optional[int] = None,
    logical_feed_indices: Optional[Sequence[int]] = None,
    prefetch_buffer_size: Optional[int] = None,
    post_batch_processor: Optional[ConfigOr[DatasetToDatasetFn]] = None,
    repeat: Optional[int] = None,
) -> DatasetToDatasetFn:
    """Returns a function that generates a tf.data.Dataset object.

    Note: batch(is_training=True) requires sufficient number of examples
    per host. When your data is too small, you should add `ds = ds.repeat()`
    before your batch.

    Args:
        global_batch_size: The global physical batch size across all replicas.
            Must be divisible by the number of JAX processes and devices.
        is_training: Whether the examples are used for training.
            This parameter will be passed through to all input sources.
        pad_example_fn: Create padded examples with the given function.
        global_logical_batch_size: The global size of the logical batch, i.e. the physical batch
            subset corresponding to elements drawn from the input dataset.
                If None, assumed to be equal to the global physical batch size.
        logical_feed_indices: The JAX process indices corresponding to feeds that provide logical
            data after batching. Process indices that are not in this set will produce
            physical-only (padded) batches, with no elements drawn from the input dataset.
            If None, assumed to be the set of all JAX training processes.
        prefetch_buffer_size: Size of prefetch buffer. This allows later
            elements to be prepared while the current element is being
            processed. If not set, `tf.data.experimental.AUTOTUNE` is used.
        post_batch_processor: An optional processor (or config instantiating to a processor) that
            applies batch-wise processing functions.
        repeat: The number of times to repeat the batches from the dataset.
            If None, repeat indefinitely if is_training=True and do not repeat otherwise.
            Otherwise must be a positive integer.

    Returns:
        A DatasetToDataset fn.

    Raises:
        ValueError: If
            - global_batch_size is not divisible by the number of JAX processes, or
            - repeat is not a positive integer, or
            - global_logical_batch_size and logical_feed_indices are not both set or both unset, or
            - global_logical_batch_size is not divisible by the number of logical_feed_indices.
    """
    num_data_feeds = jax.process_count()
    if global_batch_size % num_data_feeds != 0:
        raise ValueError(
            f"global_batch_size ({global_batch_size}) must be divisible by "
            f"number of JAX processes (data feeds) ({num_data_feeds})."
        )
    per_feed_batch_size = global_batch_size // num_data_feeds

    if repeat is not None and (not isinstance(repeat, int) or repeat <= 0):