def batch_and_repeat(ds: Dataset, batch_size: int, shuffle: bool,
repeat: bool) -> Dataset:
ds = ds.prefetch(buffer_size=AUTOTUNE)
if shuffle:
ds = ds.shuffle(1024, seed=SEED)
if repeat:
ds = ds.repeat()
if batch_size > 0:
ds = ds.batch(batch_size, drop_remainder=False)
return ds
def batch_and_repeat(
ds: Dataset, batch_size: int, shuffle: bool, repeat: bool
) -> Dataset:
"""Helper method for to apply tensorflow shuffle, repeat and
batch (in this order)
Args:
ds (Dataset): Tensorflow Dataset
batch_size (int): Will call ds.batch(batch_size, drop_remainder=False)
if batch_size is greater zero
shuffle (int): Will call ds.shuffle(1024)
repeat (bool): Will call ds.repeat()
Returns:
Dataset
"""
ds = ds.prefetch(buffer_size=AUTOTUNE)
if shuffle:
ds = ds.shuffle(1024, seed=SEED)
if repeat:
ds = ds.repeat()
if batch_size > 0:
ds = ds.batch(batch_size, drop_remainder=False)
return ds
