Repository: tensorflow
Issue: Significant Performance Drop When Training Sequential Model Using `tf.data.Dataset.from_generator`
Labels: stat:awaiting tensorflower,comp:data,type:performance,2.17

### Issue type

Performance

### Have you reproduced the bug with TensorFlow Nightly?

Yes

### Source

source

### TensorFlow version

2.13.1
2.16.1
tf-nightly (2.18.0-dev20240613)

### Custom code

No

### OS platform and distribution

_No response_

### Mobile device

_No response_

### Python version

3.11.3

### Bazel version

_No response_

### GCC/compiler version

_No response_

### CUDA/cuDNN version

_No response_

### GPU model and memory

_No response_

### Current behavior?

I noticed a significant drop in the performance of the Sequential model in a binary-cross entropy problem when I switched to using tf.dataset.from_generator from using pd.DataFrame during the training.
Here is link to a google colab notebook which demonstrate the issue:
~~https://colab.research.google.com/drive/1AoysBY3nL6Tti1fKnCb6ch9i-f20i-z6?usp=sharing~~
https://colab.research.google.com/drive/1kdEQJpRpORiF9ornGzzqn0-2f6pUVims?usp=sharing

### Standalone code to reproduce the issue

```shell
import tensorflow as tf
import keras
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
import numpy as np
import pandas as pd

def set_seeds(seed=999):
    tf.random.set_seed(seed)
    np.random.seed(seed)
    tf.keras.utils.set_random_seed(seed)  # Sets seeds for base-python, numpy, and tf

n_features = 5

# Generate synthetic dataset and split into training and testing sets
X, y = make_classification(n_samples=1_000_000, n_features=n_features, n_classes=2, random_state=42)
X_train = pd.DataFrame(X)
y_train = pd.DataFrame(y)

X_train, X_test, y_train, y_test = train_test_split(X_train, y_train, test_size=0.2, random_state=42)

# Define a Sequential model
def define_model(n_features=n_features):
    model = keras.models.Sequential([
        keras.layers.BatchNormalization(input_shape=(n_features, )),
        keras.layers.Dense(10, activation='relu'),
        keras.layers.Dense(1, activation='sigmoid')
    ])
    return model

def compile(model):
    # Compile the model
    model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])


# Train based on pd.DataFrame
set_seeds()
model = define_model()
compile(model)
history = model.fit(X_train, y_train, epochs=10, batch_size=10_000, validation_split=0.0, verbose=False)
ax = pd.DataFrame(history.history).plot()
ax.set_ylim(0, 1.0)
ax.set_title("Training using pd.DataFrame")

# Evaluate on the test sample
print("Performance using pd.DataFrame:") 
model.evaluate(X_test, y_test, batch_size=10_000)


# Train model using tf.dataset

class Generator:
    def __call__(self):
        yield X_train.values, y_train.values

ds_train = tf.data.Dataset.from_generator(
    Generator(),
    output_signature=(
        tf.TensorSpec(shape=(None, n_features), dtype=tf.float32),
        tf.TensorSpec(shape=(None, 1), dtype=tf.int32)
    )
)

set_seeds()
model = define_model()
compile(model)
history = model.fit(ds_train, epochs=10, batch_size=10_000, verbose=False)
ax = pd.DataFrame(history.history).plot()
ax.set_ylim(0, 1.0)
ax.set_title("Training using tf.dataset")


# Evaluate on the test sample
print("Performance using tf.dataset:")
model.evaluate(X_test, y_test, batch_size=10_000)
```
```


### Relevant log output

```shell
Performance using pd.DataFrame:
20/20 [==============================] - 0s 3ms/step - loss: 0.1074 - accuracy: 0.9617

Performance using tf.dataset:
20/20 [==============================] - 0s 3ms/step - loss: 0.6245 - accuracy: 0.6626
```