Support buffer shuffle for to_tf_dataset
COLLABORATOR
cc @Rocketknight1 
2023-09-13T17:57:22Z
https://github.com/huggingface/datasets/issues/6236#issuecomment-1718077705
https://api.github.com/repos/huggingface/datasets/issues/6236
IC_kwDODunzps5mZ8kJ
https://api.github.com/repos/huggingface/datasets/issues/comments/1718077705/reactions
2023-09-13T17:57:22Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1718077705
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
MEMBER
Hey! You can implement this yourself, just:

1) Create the dataset with `to_tf_dataset()` with `shuffle=False`
2) Add an `unbatch()` at the end (or use batch_size=1)
3) Add a `shuffle()` to the resulting dataset with your desired buffer size
4) Add a `batch()` at the end again to re-batch your dataset.

Note that the way we construct datasets in `to_tf_dataset()`, we don't actually shuffle the entire dataset in-memory, using `tf.data.Dataset.shuffle()`! Instead, we shuffle an index array and then load from the dataset with that. This means that shuffling with `tf.data.Dataset.shuffle()` will probably be slower and use more memory than our approach - I don't think adding the option for smaller shuffle buffers will actually save you memory on this!
2023-09-14T17:11:56Z
https://github.com/huggingface/datasets/issues/6236#issuecomment-1719841395
https://api.github.com/repos/huggingface/datasets/issues/6236
IC_kwDODunzps5mgrJz
https://api.github.com/repos/huggingface/datasets/issues/comments/1719841395/reactions
2023-09-14T17:14:01Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1719841395
https://avatars.githubusercontent.com/u/12866554?v=4
https://api.github.com/users/Rocketknight1/events{/privacy}
https://api.github.com/users/Rocketknight1/followers
https://api.github.com/users/Rocketknight1/following{/other_user}
https://api.github.com/users/Rocketknight1/gists{/gist_id}
https://github.com/Rocketknight1
Rocketknight1
MDQ6VXNlcjEyODY2NTU0
https://api.github.com/users/Rocketknight1/orgs
https://api.github.com/users/Rocketknight1/received_events
https://api.github.com/users/Rocketknight1/repos
https://api.github.com/users/Rocketknight1/starred{/owner}{/repo}
https://api.github.com/users/Rocketknight1/subscriptions
User
https://api.github.com/users/Rocketknight1
public
NONE
Thanks for your reply! @Rocketknight1 
"We don't actually shuffle the entire dataset in-memory, using tf.data.Dataset.shuffle()! Instead, we shuffle an index array and then load from the dataset with that."
In such case, there will be random access to dataset data during shuffling. When the dataset is large, the performance can be X10 times slow. I have tried many ways with to_tf_dataset() trying to achieve comparable performance with tf.data.Dataset().shuffle(buffer_size).batch(). But the performance with to_tf_dataset() is still slow. 

2023-09-18T01:11:04Z
https://github.com/huggingface/datasets/issues/6236#issuecomment-1722635615
https://api.github.com/repos/huggingface/datasets/issues/6236
IC_kwDODunzps5mrVVf
https://api.github.com/repos/huggingface/datasets/issues/comments/1722635615/reactions
2023-09-18T01:11:21Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1722635615
https://avatars.githubusercontent.com/u/7635551?v=4
https://api.github.com/users/EthanRock/events{/privacy}
https://api.github.com/users/EthanRock/followers
https://api.github.com/users/EthanRock/following{/other_user}
https://api.github.com/users/EthanRock/gists{/gist_id}
https://github.com/EthanRock
EthanRock
MDQ6VXNlcjc2MzU1NTE=
https://api.github.com/users/EthanRock/orgs
https://api.github.com/users/EthanRock/received_events
https://api.github.com/users/EthanRock/repos
https://api.github.com/users/EthanRock/starred{/owner}{/repo}
https://api.github.com/users/EthanRock/subscriptions
User
https://api.github.com/users/EthanRock
public
### Feature request

I'm using to_tf_dataset to convert a large dataset to tf.data.Dataset and use Keras fit to train model. 
Currently, to_tf_dataset only supports full size shuffle, which can be very slow on large dataset. 
tf.data.Dataset support buffer shuffle by default. 
shuffle(
    buffer_size, seed=None, reshuffle_each_iteration=None, name=None
)

### Motivation

I'm very frustrated to find the loading with shuffling large dataset is very slow. It seems impossible to shuffle before training Keras with big dataset. 

### Your contribution

NA