CLU deterministic data from tf.data.Dataset.from_tensor_slices

[homerjed]

Is there a way to do essentially

train_ds = deterministic_data.create_dataset(
    dataset_builder,
    split=train_split,
    rng=jax.random.PRNGKey(0),
    shuffle_buffer_size=100,
    batch_dims=[jax.local_device_count(), batch_size // jax.device_count()],
    num_epochs=42,
    preprocess_fn=preprocess_fn,
    shuffle=True)

with something like

tf.data.Dataset.from_tensor_slices((np.ones((200, 32, 32, 1)), np.ones((200, 1)))),

instead of the dataset_builder in the first code block?

I would like multi-device deterministic dataloading from a from_tensor_slices tensorflow dataset, has anyone got something like this working? Cheers

[Marvin182]

create_dataset() only uses the as_dataset() function of TFDS's DatasetBuilder. You could try something like:

class DatasetBuilder:
  def as_dataset(self, split, shuffle_files, read_config, decoders):
    return tf.data.Dataset.from_tensor_slices((np.ones((200, 32, 32, 1)), np.ones((200, 1))))

[homerjed]

Hi thanks for your reply, there is another error that comes up, something deep in tensorflow, I'll look into it further.

…

On Mon, Feb 6, 2023, 15:32 Marvin ***@***.***> wrote: create_dataset() only uses the as_dataset() function of TFDS's DatasetBuilder. You could try something like: class DatasetBuilder: def as_dataset(self, split, shuffle_files, read_config, decoders): return tf.data.Dataset.from_tensor_slices((np.ones((200, 32, 32, 1)), np.ones((200, 1)))) — Reply to this email directly, view it on GitHub <#2851 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/ALI5PXDCAKBYRWCDH5CGCPTWWEDPDANCNFSM6AAAAAAUSVBLRQ> . You are receiving this because you authored the thread.Message ID: ***@***.***>