Repository: tensorflow
Issue: tf.data.Dataset unusable with steps_per_epoch standard training loop
Labels: stat:awaiting tensorflower,type:bug,comp:data,TF 2.10

**System information**
- Have I written custom code (as opposed to using a stock example script provided in TensorFlow): No
- OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Mint 
- TensorFlow installed from (source or binary): binary
- TensorFlow version (use command below): 2.1.0
- Python version: 3.7.4

**Describe the current behavior**

This is the underlying issue for https://github.com/tensorflow/tensorflow/issues/36153

Basically:
- `Dataset` can be used as/converted to an iterator
- Once the iterator reaches the end of the dataset it does not restart -> Yields no more samples
- `steps_per_epoch` for the `fit` method of a keras model can be used to specify the number of batches to run per epoch, more exactly: The number of times the iterator is advanced/dereferenced
- `steps_per_epoch` is required for e.g. if not the full dataset should be traversed. Either due to external requirements or to avoid using the trailing (incomplete) batch.
- `steps_per_epoch` is absolutely required to be set for `MultiWorkerMirroredStrategy`: https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras#train_the_model_with_multiworkermirroredstrategy
- Whether to recreate an iterator for each epoch out of the dataset is determined by `DatasetAdapter.should_recreate_iterator` at https://github.com/tensorflow/tensorflow/blob/v2.1.0/tensorflow/python/keras/engine/training_v2.py#L241-L242
- Note that the iterator absolutely must be recreated for all common use cases:
  - Iterate over full dataset
  - Iterate over full dataset except incomplete trailing batch
  - Iterate over a random subset of the full dataset per epoch
  - Not to be recreated for: Infinite datasets. **Maybe** If the full dataset should be consumed over multiple epochs. But why? What should happen if the dataset is exhausted after some epoch? Usually restart, right?

In the current TF (2.1.0) the implementation of `DatasetAdapter.should_recreate_iterator` is: `return self.get_size() is not None or steps_per_epoch is None`: https://github.com/tensorflow/tensorflow/blob/v2.1.0/tensorflow/python/keras/engine/data_adapter.py#L208

This is **wrong**. For datasets the size is always None (see https://github.com/tensorflow/tensorflow/issues/36531 which intends for this to be changed) and as motivated above having `steps_per_epoch` set is a common use case but the iterator should still be recreated on each epoch.

This was recently changed to ` (self._user_steps is None or cardinality.cardinality(self._dataset).numpy() == self._user_steps)` https://github.com/tensorflow/tensorflow/commit/6be131d0860559954c42685a87c63f16cebb2185#diff-f8dd40712ac721c1b363e1a1ec44c1a3R741-R747

This is also **wrong**. Again assume `_user_steps` is set (see above reasoning): First the `cardinality` might be unknown, e.g. for any TFDS dataset (and `TFRecordDataset`) it is unknown. I guess due to use of `interleave` in the dataset reader. Second even if the size was known it may not be equal to the number of steps. Common example: Skipping the last batch.

**Describe the expected behavior**

This is a design issue and hence hard to resolve.

In general it would be best to eliminate the UNKNOWN size of a dataset. But when reading data line-by-line from a file it might not be known upfront. So the user has to input the size of the dataset or specify explicitly `AUTO` which would iterate over the whole dataset once to get the number of samples. This can be costly but should not be possible in general (e.g. TFDS knows the number of samples)

I think the sanest approach would be to default to recreating the iterator on each epoch UNLESS the dataset is known to be infinite. This might still be wrong for cases I can't imagine right now but is correct for all cases I can think of. Maybe even allow the user to specify this, but this default is IMO way better than the current.

The other approach would be to recreate the iterator when it runs out of data after starting an epoch. This would partially solve the issue, but fails for:
- Omitting the trailing batch: It would yield the incomplete batch from the last epoch first, but that should be skipped.
- Using only some random samples per batch but wanting to shuffle before each batch: It would happily consume the rest of the samples and not see the samples used in an earlier epoch until it runs out of data

With the UNKNOWN size fixed, one could also fix the check at  https://github.com/tensorflow/tensorflow/commit/6be131d0860559954c42685a87c63f16cebb2185#diff-f8dd40712ac721c1b363e1a1ec44c1a3R741-R747 to check if `_user_steps` yields 1 epoch (and optional trailing batches) by using `size // steps == 1` which would be better than the previous approach (recreate iterator when out of data) as the use case with of Omitting the trailing batch is covered. But it would fail for the other.

So my suggestion would to
- (optional) avoid UNKNOWN sizes
- recreate iterators unless dataset is infinite by default
- allow the user to overwrite this explicitly, maybe via `with_options` of the dataset


**Code to reproduce the issue**
Some reduced example code based on e.g. https://www.tensorflow.org/tutorials/distribute/multi_worker_with_keras#train_the_model_with_multiworkermirroredstrategy

```
import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow as tf
from tensorflow.python.data.experimental.ops import cardinality
import numpy as np
tfds.disable_progress_bar()

# Scaling MNIST data from (0, 255] to (0., 1.]
def scale(image, label):
    image = tf.cast(image, tf.float32)
    image /= 255
    return image, label

def build_and_compile_cnn_model():
    model = tf.keras.Sequential([
        tf.keras.layers.Conv2D(32,
                               3,
                               activation='relu',
                               input_shape=(28, 28, 1)),
        tf.keras.layers.MaxPooling2D(),
        tf.keras.layers.Flatten(),
        tf.keras.layers.Dense(64, activation='relu'),
        tf.keras.layers.Dense(10, activation='softmax')
    ])
    model.compile(loss=tf.keras.losses.sparse_categorical_crossentropy,
                  optimizer=tf.keras.optimizers.SGD(learning_rate=0.001),
                  metrics=['accuracy'])
    return model

BATCH_SIZE = 64
if False:
    examples = np.ones([10*BATCH_SIZE,28,28,1])
    labels = np.ones([examples.shape[0]])
    dataset = tf.data.Dataset.from_tensor_slices((examples, labels))
    num_examples = examples.shape[0]
else:
    datasets, info = tfds.load(name='mnist', with_info=True, as_supervised=True)
    dataset = datasets['test']
    num_examples = info.splits['test'].num_examples

x = dataset.map(scale).cache().shuffle(10000).batch(BATCH_SIZE)
model = build_and_compile_cnn_model()

card = cardinality.cardinality(x)
num_batches = sum(1 for _ in x)
full_batches = num_examples // BATCH_SIZE
print("Samples: %s\nBatches: %s (%s full)\nCardinality: %s" % (num_examples, num_batches, full_batches, card))
model.fit(x=x, epochs=2, steps_per_epoch=full_batches)

```

**Other info / logs**
There is a warning:
> WARNING:tensorflow:Your input ran out of data; interrupting training. Make sure that your dataset or generator can generate at least `steps_per_epoch * epochs` batches (in this case, 312 batches). You may need to use the repeat() function when building your dataset.

It seems that adding `repeat()` and hence creating an infinite dataset is a viable option. However the docu also states
> In TF 1.X, the idiomatic way to create epochs was through the `repeat` transformation:
>     In TF 2.0, `tf.data.Dataset` objects are Python iterables which makes it possible to also create epochs through Python iteration:

So it seems that it should not be required and as per above explanation the return value of `DatasetAdapter.should_recreate_iterator` is not correct.