Function tfds_dataset

axlearn/common/input_tf_data.py:285–388

Returns a BuildDatasetFn for the given TFDS dataset name and split. Args: dataset_name: The tensorflow dataset name. split: The dataset split. is_training: Whether the examples are used for training. If True, examples will be read in parallel and shuffled.

(
    dataset_name: str,
    *,
    split: str,
    is_training: bool,
    train_shuffle_buffer_size: Optional[int] = None,
    train_shuffle_files: Optional[bool] = None,
    data_dir: Optional[str] = None,
    download: bool = False,
    read_config: Optional[InstantiableConfig] = None,
    decoders: Optional[InstantiableConfig] = None,
)

def tfds_dataset(
    dataset_name: str,
    *,
    split: str,
    is_training: bool,
    train_shuffle_buffer_size: Optional[int] = None,
    train_shuffle_files: Optional[bool] = None,
    data_dir: Optional[str] = None,
    download: bool = False,
    read_config: Optional[InstantiableConfig] = None,
    decoders: Optional[InstantiableConfig] = None,
) -> BuildDatasetFn:
    """Returns a BuildDatasetFn for the given TFDS dataset name and split.

    Args:
        dataset_name: The tensorflow dataset name.
        split: The dataset split.
        is_training: Whether the examples are used for training.
            If True, examples will be read in parallel and shuffled.
            Otherwise examples will be read sequentially to ensure a deterministic order.
        train_shuffle_buffer_size: The shuffle buffer size (required) when is_training=True.
            If is_training=False or shuffle_buffer_size <= 0, no shuffling is done.
        train_shuffle_files: Whether to shuffle files when is_training=True.
            If is_training=False, no shuffling is done.
            If is_training=True and train_shuffle_files is None, infer from shuffle_buffer_size > 0.
        data_dir: Used for tfds.load. If None, use the value of the environment variable
            DATA_DIR, TFDS_DATA_DIR, or "~/tensorflow_datasets" (in that order).
        download: Whether to download the examples. If false, use the data under data_dir.
        read_config: a Config that instantiates to a tfds.ReadConfig.
            If None, constructs a default one with tfds_read_config().
        decoders: An optional config instantiating a (nested) mapping of feature names to decoders.
                See: https://www.tensorflow.org/datasets/api_docs/python/tfds/decode/Decoder

    Returns:
        A BuildDatasetFn, which returns a tf.data.Dataset for the specified TFDS.

    Raises:
        ValueError: if train_shuffle_buffer_size is None when is_training=True.
    """
    if is_training:
        if train_shuffle_buffer_size is None:
            raise ValueError("train_shuffle_buffer_size is required when is_training=True")
        shuffle_files = (
            train_shuffle_buffer_size > 0 if train_shuffle_files is None else train_shuffle_files
        )
        shuffle_buffer_size = train_shuffle_buffer_size
    else:
        # Disable shuffling.
        #
        # Note that it's OK for is_training=True and train_shuffle_buffer_size=0 so that the inputs
        # are deterministic.
        shuffle_buffer_size = 0
        shuffle_files = False

    if data_dir is None:
        data_dir = get_data_dir()

    if read_config is None: