(n_devices,
                 dataset_name,
                 split,
                 batch_size,
                 data_dir,
                 shuffle_files,
                 shuffle_buffer_size,
                 batch_shuffle_size,
                 preprocess_fun,
                 repeat=True)
Streams batches of examples from tfds, with pure-python preprocessing.
Streams batches of examples from tfds, with pure-python preprocessing.
Streams batches of examples from tfds with pure - python preprocessing .
def _tfds_stream(n_devices,
                 dataset_name,
                 split,
                 batch_size,
                 data_dir,
                 shuffle_files,
                 shuffle_buffer_size,
                 batch_shuffle_size,
                 preprocess_fun,
                 repeat=True):
  """Streams batches of examples from tfds, with pure-python preprocessing."""
  # TODO(piotrekp1): delete if switched to data_streams
  if batch_size % n_devices != 0:
    raise ValueError(f'Batch size ({batch_size}) not divisible'
                     ' by number of devices ({n_devices})')
  ds = tfds.load(
      name=dataset_name,
      split=split,
      data_dir=data_dir,
      shuffle_files=shuffle_files)
  if repeat:
    ds = ds.repeat()
  if shuffle_buffer_size is not None:
    ds = ds.shuffle(shuffle_buffer_size)
  ds = ds.batch(batch_size)
  if batch_shuffle_size is not None:
    ds = ds.shuffle(batch_shuffle_size)

  for batch in tfds.as_numpy(ds):
    if preprocess_fun is not None:
      yield preprocess_fun(batch)
    else:
      yield batch
def _tfds_stream ( n_devices , dataset_name , split , batch_size , data_dir , shuffle_files , shuffle_buffer_size , batch_shuffle_size , preprocess_fun , repeat = True ) : # TODO(piotrekp1): delete if switched to data_streams if batch_size % n_devices != 0 : raise ValueError ( f'Batch size ({batch_size}) not divisible' ' by number of devices ({n_devices})' ) ds = tfds . load ( name = dataset_name , split = split , data_dir = data_dir , shuffle_files = shuffle_files ) if repeat : ds = ds . repeat ( ) if shuffle_buffer_size is not None : ds = ds . shuffle ( shuffle_buffer_size ) ds = ds . batch ( batch_size ) if batch_shuffle_size is not None : ds = ds . shuffle ( batch_shuffle_size ) for batch in tfds . as_numpy ( ds ) : if preprocess_fun is not None : yield preprocess_fun ( batch ) else : yield batch