[Tensorflow] Support tf.dataset.repeat() to avoid duplicating and dropping samples in one epoch with shuffle?

[chongxiaoc]

Current implementation of make_petastorm_dataset for tensorflow doesn't support multiple iterations:https://github.com/uber/petastorm/blob/7f37e8dde6ff1b13f055d22a6289e2de8bb5d473/petastorm/tf_utils.py#L370

It is recommended to set reader's num_epochs > 1 to support multiple iterations.

This will cause possible duplication and drop and samples in one epoch when using together with tf.dataset.shuffle. Let's say:

• My training rows are [0,1,2,3,4,5,6,7]
• I plan to train 10 epochs, so I set reader.num_epochs = 10.
• I want to shuffle my samples as well in training, so I add shuffle to the dataset as below:

dataset = make_petastorm_dataset(reader)
dataset = dataset.shuffle(8)

• During training, due the implmentation of shuffle in tf: API docs (TF API page):

This dataset fills a buffer with buffer_size elements, then randomly samples elements from this buffer, replacing the selected elements with new elements. For perfect shuffling, a buffer size greater than or equal to the full size of the dataset is required.

For instance, if your dataset contains 10,000 elements but buffer_size is set to 1,000, then shuffle will initially select a random element from only the first 1,000 elements in the buffer. Once an element is selected, its space in the buffer is replaced by the next (i.e. 1,001-st) element, maintaining the 1,000 element buf

• A training epoch will be like:

  • Randomly pick 7, and refill with 0 (next available in reader) again: [0,1,2,3,4,5,6,0]
  • Now it is highly possible 0 will be selected twice in next few iterations, even before other values are selected once.
  • In the end of an epoch, it is highly possible 0 is duplicated, while some values are dropped.

• This example is reproducible in our real applications, the reason is that by using tf.dataset.shuffle, it will always refill shuffle buffer with next available element from reader (since we set reader to be multiple epochs)

• A work-around to fix this:

  • Set reader num_epochs to be 1.
  • Enable multiple iterations in make_petastorm_dataset
  • Using shuffle and repeat from tf.dataset as:

dataset = make_petastorm_dataset(reader)
dataset = dataset.shuffle(8)
dataset = dataset.repeat(num_epochs)

• In this case, repeat isolates different epochs, so we are not seeing sample drop and duplication in a single epoch. (verified on our applications).

@selitvin What do you think?