Re-shuffle data each iteration about flax HOT 7 CLOSED

reshuffle_each_iteration is False by default? I guess this is a great example of why implicit random seeds are evil.

In the past I have also noticed that creating iterators in tf.data tends to come with significant overhead so I prefer to create one infinite iterator and use a manual for step in range(steps_per_epoch) to loop for an epoch. This also avoids some other issues like partial batches.

from flax.

Just checked, it actually does default to True:

reshuffle_each_iteration: (Optional.) A boolean, which if true indicates that the dataset should be pseudorandomly reshuffled each time it is iterated over. (Defaults to True.)

Perhaps the issue is that the nlp_seq example does not shuffle the data?

from flax.

Marking as PR welcome. A successful PR would also include (a link to) new TensorBoard logs verifying the metrics (see README.md in examples/nlp_seq for expected numbers).

from flax.

It looks like dataset.shuffle was just missing. I've made the change. I'll train on a V100 similar to the README and submit a PR with updated logs and code. I've checked other example folders. Others don't miss the shuffle to the training dataset. Please correct if I'm wrong.

from flax.

Awesome @gan3sh500 ! Indeed if you can reproduce the same results on a V100 and share a link to tensorboard.dev we'll merge your PR.

from flax.

Sorry for slow turnaround. I didn't get the V100 but I now have a 2080Ti locally to run and have ran the same batch_size and iterations. The results are 0.23% lower than the one shown in the repo. Please see for yourself. TB Dev #321

from flax.

It looks like @bohnetbd felt this wasn't worth adding due to the added complexity (given that this is an educational example). So I am closing this.

from flax.