def make_tf_data(
cls,
filename: type_utils.PathLike,
buffer_size: Optional[int] = None,
) -> tf.data.Dataset:
from riegeli.tensorflow.ops import riegeli_dataset_ops as riegeli_tf # pylint: disable=g-import-not-at-top
return riegeli_tf.RiegeliDataset(filename, buffer_size=buffer_size)
def make_tf_data(cls, filename: type_utils.PathLike,
buffer_size: tf.int64) -> tf.data.Dataset:
"""Returns TensorFlow Dataset comprising given record file."""
return riegeli_tf.RiegeliDataset(filename, buffer_size=buffer_size)
def dataset_fn(self, filenames, num_epochs=1, batch_size=None):
repeat_dataset = riegeli_dataset_ops.RiegeliDataset(filenames).repeat(
num_epochs)
if batch_size:
return repeat_dataset.batch(batch_size)
return repeat_dataset
